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

workspace/README.md

@@ -1,15 +1,28 @@
-# Training Configuration Guide
+# Training Guide
 
-This document explains the key configuration parameters that can be modified for Wall-X training.
+This document explains the key configuration parameters and memory requirements for Wall-X training.
 
 ## Quick Start Checklist
-1. **Update run.sh**: Set `code_dir` and `config_path` to your actual paths
-2. **Configure GPUs**: Set `CUDA_VISIBLE_DEVICES` for your available GPUs
-3. **Update config paths**: Replace all `/path/to/` placeholders in `config_qact.yml` with actual paths
-4. **Configure robot**: Set `dof_config` and `agent_pos_config` for your robot
-5. **Set dataset**: Choose appropriate `repo_id` for your dataset
-6. **Adjust batch size**: Set `batch_size_per_gpu` based on GPU memory
-7. **Run training**: Execute `bash ./workspace/lerobot_example/run.sh`
+
+### 🚀 **Step 1: Download Pre-trained Model**
+Choose one of the available models:
+- **WALL-OSS-FLOW**: https://huggingface.co/x-square-robot/wall-oss-flow
+- **WALL-OSS-FAST**: https://huggingface.co/x-square-robot/wall-oss-fast
+
+### ⚙️ **Step 2: Configure Environment**
+- Update `run.sh`: Set `code_dir` and `config_path` to your actual paths
+- Set `CUDA_VISIBLE_DEVICES` for your available GPUs
+
+### 📝 **Step 3: Update Configuration Files**
+- Replace all `/path/to/` placeholders in `config_qact.yml` with actual paths
+- Configure robot settings: `dof_config` and `agent_pos_config`
+- Set dataset: Choose appropriate `repo_id`
+- Adjust `batch_size_per_gpu` based on your GPU memory
+
+### ▶️ **Step 4: Start Training**
+```bash
+bash ./workspace/lerobot_example/run.sh
+```
 
 ## Enable FAST tokenizer
 To fine-tune using the FAST tokenizer, please download the repository and update the `action_tokenizer_path`. Make sure to set `use_fast_tokenizer` to `true`:
@@ -38,6 +51,15 @@ action_tokenizer_path: "/path/to/fast/"             # Must set if use_fast_token
 - `num_training_steps`: Total training steps
 - `num_epoch`: Number of training epochs
 
+### Training Optimization Settings
+- `FSDP2`: Enable FSDP2 for distributed training (default: True) - **Recommended for multi-GPU**
+- `torch_compile`: Enable PyTorch compilation optimization (default: False)
+
+**⚠️ Important Note on torch_compile:**
+- **Benefits**: Enabling `torch_compile` can significantly improve training efficiency
+- **Requirements**: Requires that the data input shape is always consistent throughout training
+- **Caution**: If you don't have sufficient understanding of torch compile, please **DO NOT** enable it as it may cause unexpected issues with dynamic input shapes
+
 ## Robot Configuration (Modify for Your Robot)
 
 ### DOF Configuration
@@ -67,6 +89,20 @@ Keep `agent_pos_config` consistent with `dof_config`.
 - `resume.ckpt`: Path to checkpoint for resuming training
 - `resume.load_ckpt_only`: Only load model weights, not optimizer state
 
-## Performance Settings (Optional)
-- `profile`: Enable PyTorch profiling (true/false)
-- `padding_side`: Token padding side (left/right)
+## Memory Usage
+
+Below are the memory consumption benchmarks for different training configurations using the `lerobot/aloha_mobile_cabinet` dataset:
+
+| Dataset | Batch Size | FSDP2 | Torch Compile | Num GPUs | Max Allocated Memory |
+|---------|------------|--------|---------------|----------|---------------------|
+| lerobot/aloha_mobile_cabinet | 1 | ❌ | ❌ | 1 | 40.11G |
+| lerobot/aloha_mobile_cabinet | 1 | ❌ | ❌ | 8 | 48.02G |
+| lerobot/aloha_mobile_cabinet | 1 | ✅ | ❌ | 2 | 43.70G |
+| lerobot/aloha_mobile_cabinet | 1 | ✅ | ❌ | 8 | 24.96G |
+| lerobot/aloha_mobile_cabinet | 1 | ✅ | ✅ | 8 | 24.21G |
+
+
+**Hardware Recommendations:**
+
+- For single GPU training: Ensure at least 48GB VRAM (e.g., RTX 6000 Ada, A6000)
+- For multi-GPU training: Enable FSDP2 for optimal memory distribution

workspace/lerobot_example/config_qact.yml

@@ -28,6 +28,10 @@ batch_size_per_gpu: 8
 padding_side: left
 epoch_save_interval: 10
 
+# Training optimization settings
+FSDP2: True
+torch_compile: False
+
 # Robot configuration - Define degrees of freedom for each component
 dof_config:
   follow_left_ee_cartesian_pos: 3      # Left end-effector Cartesian position

workspace/lerobot_example/run.sh

@@ -1,5 +1,4 @@
 #!/bin/bash
-# export CUDA_VISIBLE_DEVICES=4,5,6,7
 export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
 NUM_GPUS=$(echo $CUDA_VISIBLE_DEVICES | tr ',' '\n' | wc -l)