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
This commit is contained in:
Starrick Liu
2025-09-17 23:09:20 +08:00
committed by GitHub
parent 86f70a3b08
commit 421db17d53
24 changed files with 1862 additions and 76 deletions
+7 -5
View File
@@ -68,13 +68,15 @@ if __name__ == "__main__":
wrapper = VQAWrapper(model_path=MODEL_PATH_FOR_MODULE_TEST)
try:
test_image_url = "https://www.ilankelman.org/stopsigns/australia.jpg"
test_question = "What is written on the sign?"
test_question = "To move the red block in the plate with same color, what should you do next? Think step by step."
# img = Image.open("/path/to/your/local/image.jpg").convert("RGB")
import requests
# Local Image
img = Image.open("/path/to/wall-x/assets/cot_example_frame.png").convert("RGB")
# Internet Image
# import requests
# test_image_url = "https://www.ilankelman.org/stopsigns/australia.jpg"
# img = Image.open(requests.get(test_image_url, stream=True).raw).convert("RGB")
img = Image.open(requests.get(test_image_url, stream=True).raw).convert("RGB")
answer = wrapper.generate(img, test_question)
print("model answer:", answer)