huangyanpei/VLA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Major optimization and robustness improvements (#31)

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py
+10 -8
View File
@@ -560,15 +560,17 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel):
`torch.Tensor`: hidden_states.
"""
hidden_states = self.patch_embed(hidden_states)
rotary_pos_emb = self.rot_pos_emb(grid_thw)
window_index, cu_window_seqlens = self.get_window_index(grid_thw)
window_index = window_index.to(hidden_states.device)
cu_window_seqlens = torch.tensor(
cu_window_seqlens,
device=hidden_states.device,
dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
rotary_pos_emb = ops.rot_pos_emb(
self.rotary_pos_emb.inv_freq, grid_thw, self.spatial_merge_size
)
window_index, cu_window_seqlens = ops.get_window_index(
grid_thw=grid_thw,
window_size=self.window_size,
spatial_merge_size=self.spatial_merge_size,
patch_size=self.patch_size,
spatial_merge_unit=self.spatial_merge_unit,
)
cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)
seq_len, _ = hidden_states.size()
hidden_states = hidden_states.reshape(
wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py
+35 -12
View File
@@ -1244,12 +1244,17 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
or self.rope_deltas is None
or (past_key_values is None or past_key_values.get_seq_length() == 0)
):
position_ids, rope_deltas = self.get_rope_index(
input_ids,
image_grid_thw,
video_grid_thw,
second_per_grid_ts,
attention_mask,
position_ids, rope_deltas = ops.get_rope_index(
input_ids=input_ids,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
second_per_grid_ts=second_per_grid_ts,
attention_mask=attention_mask,
spatial_merge_size=self.config.vision_config.spatial_merge_size,
image_token_id=self.config.image_token_id,
video_token_id=self.config.video_token_id,
vision_start_token_id=self.config.vision_start_token_id,
tokens_per_second=self.config.vision_config.tokens_per_second,
)
self.rope_deltas = rope_deltas
# Use previously calculated rope deltas to get correct position IDs
@@ -1720,12 +1725,17 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
or self.rope_deltas is None
or (past_key_values is None or past_key_values.get_seq_length() == 0)
):
position_ids, rope_deltas = self.get_rope_index(
input_ids,
image_grid_thw,
video_grid_thw,
second_per_grid_ts,
attention_mask,
position_ids, rope_deltas = ops.get_rope_index(
input_ids=input_ids,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
second_per_grid_ts=second_per_grid_ts,
attention_mask=attention_mask,
spatial_merge_size=self.config.vision_config.spatial_merge_size,
image_token_id=self.config.image_token_id,
video_token_id=self.config.video_token_id,
vision_start_token_id=self.config.vision_start_token_id,
tokens_per_second=self.config.vision_config.tokens_per_second,
)
self.rope_deltas = rope_deltas
# Use previously calculated rope deltas to get correct position IDs
@@ -1882,6 +1892,17 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
)
dof_mask = dof_mask.to(inputs_embeds.device).to(inputs_embeds.dtype)
# Calculate token distribution across MoE expert groups
group_size = torch.zeros(
self.config.num_experts, dtype=torch.long, device="cpu"
)
for i in range(self.config.num_experts):
group_size[i] = (moe_token_types == i).sum()
# Calculate start and end indices for each expert group
start_indices = torch.cumsum(group_size, dim=0) - group_size
end_indices = torch.cumsum(group_size, dim=0)
def step(timestep, noisy_action):
"""
Single denoising step for diffusion process.
@@ -1915,6 +1936,8 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
past_key_values=past_key_values,
inputs_embeds=temp_inputs_embeds,
moe_token_types=moe_token_types,
start_indices=start_indices,
end_indices=end_indices,
use_cache=True,
output_attentions=False,
output_hidden_states=False,