Fix lerobot fine-tuning out of the box: version-robust checkpoint loading, optional flash-attn, episode passthrough

[Liyux3]

Summary

Following the workspace README to fine-tune wall-oss-0.5 on a custom lerobot dataset currently breaks in several environment-dependent ways. This PR makes that path work out of the box:

1. Version-robust checkpoint loading (modeling_qwen2_5_vl_act.py)

   • The released checkpoint stores fused projections (qkv_proj_experts.*, gate_up_proj). With transformers versions that build split modules (e.g. 4.51.x, which is needed anyway since AttentionInterface does not exist in the pinned 4.49.0), load_state_dict(strict=False) silently skips ~700 tensors: all LM attention/MLP expert weights and every vision-tower MLP stay randomly initialized while training appears to run normally. This PR slices fused tensors into the split naming, with split sizes taken from the model's own state_dict (HF concat order q;k;v / gate;up), so it is a no-op on environments where names already match.
   • Checkpoint tensors whose shapes disagree with a customized dof_config (e.g. the 26-dim multi-embodiment action projections vs a custom single-arm config) previously raised RuntimeError even with strict=False; they now re-initialize with a clear message, which is the expected behavior when adapting to a new embodiment per the workspace README.
   • A load report (missing/unexpected counts + first keys) is printed so silent loading issues are visible.

2. Optional flash-attn (joint_attention.py): flash_attn is imported at module level but is not listed in requirements.txt, so a fresh install crashes at import. It is now optional; only the flash_attention_2 attention class needs it, sdpa/eager paths never call it.

3. Episode list passthrough (load_lerobot_dataset.py): lerobot_config.episodes exists in the config template but was ignored; the loader always used a sequential train/test split. Honoring it allows holding out evaluation episodes.

4. use_selective_recompute passthrough (qwen_vl_act_trainer.py): the wall-oss branch dropped this flag when calling from_pretrained, silently disabling selective recompute.

5. Opt-in train_action_preprocessor (qwen_vl_act_trainer.py): with freeze_vlm: true and a customized dof_config, the freshly re-initialized action in/out projections were frozen at random init. The new flag keeps them trainable alongside the experts; default false preserves current behavior.

Validation

Fine-tuned wall-oss-0.5 on a custom 7-DoF single-arm lerobot dataset (3 cameras, joint space + gripper), transformers 4.51.3 / torch 2.10 / python 3.12 / RTX 4090 (sm_89):

• Load report after this PR: 280 fused tensors adapted, 4 missing / 0 unexpected (tied lm_head, propri_proj which is absent from the released checkpoint, and the two intentionally re-initialized action projections)
• Single 24GB GPU fine-tune runs healthily with freeze_vlm + selective recompute, bf16, bs1

Note

requirements.txt pins transformers==4.49.0, but wall_x/model/vla_mixin.py imports AttentionInterface, which does not exist in 4.49; the pin may want a bump to >=4.50 in a separate change.