feat: Major optimization and robustness improvements (#31)

This release introduces significant performance optimizations, memory efficiency
improvements, and enhanced system robustness:

🚀 Performance Optimizations:
- Add three new fused CUDA kernels (rope_index, rot_pos_emb, get_window_index)
  for accelerated multimodal preprocessing
- Implement FSDP2 support for distributed training with improved memory efficiency
- Add Torch.compile integration for additional performance gains
- Optimize memory usage: reduce peak allocation from 48GB to 24GB on 8-GPU setup

🔧 System Robustness:
- Fix missing token position inputs in prediction pipeline
- Add type-robust negation operations in RoPE CUDA kernels (half/bfloat16 support)
- Fix dataset root parameter initialization in LeRobot data loader
- Enhanced error handling and input validation across fusion operators

📚 Documentation & Usability:
- Add comprehensive memory usage benchmarks and hardware recommendations
- Update citation format with proper arXiv reference
- Improve training configuration documentation with quick start guide
- Add detailed API documentation for new fusion operators

🛠️ Technical Details:
- Version bump to 1.0.1
- New CUDA kernels: rope_index.cu, rot_pos.cu, window_index.cu
- FSDP2 state dict loading with distribute_tensor support
- Enhanced multimodal RoPE with 3D position encoding
- Window attention optimization for Vision Transformers

Breaking Changes: None - all changes are backward compatible

train_qact.py

@@ -3,6 +3,7 @@ import json
 import time
 import yaml
 import wandb
+import accelerate
 from argparse import ArgumentParser
 from accelerate import (
     Accelerator,
@@ -38,9 +39,32 @@ def setup_accelerator(config):
     ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
     accelerator_dataloader_config = DataLoaderConfiguration(dispatch_batches=False)
 
+    if config.get("FSDP2", False):
+        # Use Fully Sharded Data Parallel (FSDP) version 2
+        fsdp_plugin = accelerate.utils.dataclasses.FullyShardedDataParallelPlugin(
+            fsdp_version=2, reshard_after_forward=True
+        )
+        print("[INFO] Using FSDP version 2 for distributed training")
+    else:
+        fsdp_plugin = None
+
+    if config.get("torch_compile", False):
+        # Use Torch Dynamo for compilation
+        dynamo_plugin = accelerate.utils.TorchDynamoPlugin(
+            backend="inductor",
+            mode="default",
+            fullgraph=False,
+            dynamic=False,
+        )
+        print("[INFO] Using Torch Dynamo for compilation")
+    else:
+        dynamo_plugin = None
+
     accelerator = Accelerator(
         kwargs_handlers=[ddp_kwargs],
         mixed_precision="bf16",
+        fsdp_plugin=fsdp_plugin,
+        dynamo_plugin=dynamo_plugin,
         dataloader_config=accelerator_dataloader_config,
         gradient_accumulation_steps=config.get("gradient_accumulation_steps", 1),
     )