Starrick Liu
421db17d53
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
117 lines
3.2 KiB
YAML
117 lines
3.2 KiB
YAML
# Training Configuration for Wall-X Robotic Multi-Modal Learning
|
|
# This configuration supports multi-modal learning with vision, language, and action data
|
|
|
|
# Model and paths configuration
|
|
log_name: "robotic_training"
|
|
log_project: "vla_training"
|
|
model_type: qwen2_5
|
|
pretrained_wallx_path: "/path/to/wallx_model/" # Must set
|
|
save_path: "/path/to/workspace/" # Must set
|
|
use_fast_tokenizer: False # True: train FAST, False: train Flow
|
|
action_tokenizer_path: "/path/to/fast/" # Must set if use_fast_tokenizer is true
|
|
|
|
# Torch Profile
|
|
profile: False
|
|
profile_save_path: /path/to/profile/
|
|
profile_wait_iters: 10
|
|
profile_warmup_iters: 5
|
|
profile_active_iters: 2
|
|
|
|
# Training hyperparameters
|
|
num_warmup_steps: 100
|
|
num_training_steps: 64000000
|
|
learning_rate: 0.00009
|
|
min_lr: 0.00005
|
|
num_epoch: 100
|
|
gradient_accumulation_steps: 32
|
|
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
|
|
follow_left_ee_rotation: 3 # Left end-effector rotation
|
|
follow_left_gripper: 1 # Left gripper control
|
|
follow_right_ee_cartesian_pos: 3 # Right end-effector Cartesian position
|
|
follow_right_ee_rotation: 3 # Right end-effector rotation
|
|
follow_right_gripper: 1 # Right gripper control
|
|
head_actions: 2 # Head/camera movement
|
|
height: 1 # Mobile base height control
|
|
car_pose: 3 # Mobile base pose (x, y, theta)
|
|
|
|
# Agent proprioception configuration (typically matches DOF config)
|
|
agent_pos_config:
|
|
follow_left_ee_cartesian_pos: 3
|
|
follow_left_ee_rotation: 3
|
|
follow_left_gripper: 1
|
|
follow_right_ee_cartesian_pos: 3
|
|
follow_right_ee_rotation: 3
|
|
follow_right_gripper: 1
|
|
head_actions: 2
|
|
height: 1
|
|
car_pose: 3
|
|
|
|
# # Checkpoint resuming configuration
|
|
# resume:
|
|
# ckpt: "/path/to/resume_model/"
|
|
# load_ckpt_only: true
|
|
|
|
# Data configuration
|
|
data:
|
|
use_lerobot: true
|
|
|
|
# LeRobot dataset configuration
|
|
lerobot_config:
|
|
repo_id: "lerobot/aloha_mobile_cabinet"
|
|
root: null
|
|
episodes: null
|
|
image_transforms: null
|
|
delta_timestamps: null
|
|
tolerance_s: 1e-4
|
|
revision: null
|
|
force_cache_sync: false
|
|
download_videos: true
|
|
video_backend: null
|
|
|
|
action_horizon: 32
|
|
train_test_split: 0.95
|
|
|
|
# Action keys for observation and prediction
|
|
obs_action_keys:
|
|
- follow_left_ee_cartesian_pos
|
|
- follow_left_ee_rotation
|
|
- follow_left_gripper
|
|
- follow_right_ee_cartesian_pos
|
|
- follow_right_ee_rotation
|
|
- follow_right_gripper
|
|
- head_actions
|
|
- height
|
|
- car_pose
|
|
|
|
predict_action_keys:
|
|
- follow_left_ee_cartesian_pos
|
|
- follow_left_ee_rotation
|
|
- follow_left_gripper
|
|
- follow_right_ee_cartesian_pos
|
|
- follow_right_ee_rotation
|
|
- follow_right_gripper
|
|
- head_actions
|
|
- height
|
|
- car_pose
|
|
|
|
# Image resolution configuration for different camera views
|
|
resolution:
|
|
face_view: 256
|
|
left_wrist_view: 256
|
|
right_wrist_view: 256
|
|
move1_view: 256
|
|
move2_view: 256
|
|
top_view: 256
|
|
wall_view: 256
|
|
multi_modal: 256
|