Chapter 10
Which parameters move, and how to fit it in memory
4 sections · about 3 minutes
10.1 Parameter groups
| Module | Approximate size | Recommendation |
|---|---|---|
ActionDiT + the MoT action pathway + proprio_encoder | about 0.7 B | Full fine-tune; learning rate 1e-5 |
| Video expert (Wan2.2 DiT) | about 5 B | LoRA adapters (rank 64 on attention and feed-forward projections), learning rate 5e-5, updated only by L_WM |
| VAE, T5 | — | Frozen, as in pretraining |
| Value head (only if using PPO with a critic) | tiny | Learning rate 1e-4 |
The actor learning rate is the one verifiable reference point from the VLA-RL literature scaled up slightly: SimpleVLA-RL trains a 7-billion-parameter autoregressive policy at 5e-6, and a 0.7-billion-parameter flow head tolerates a somewhat larger step. (πRL and RLinf-VLA defer their learning rates to appendices that were not verified for this document.) WAM-RL fully fine-tunes its world model; LoRA is recommended here for three reasons: it fits in memory without ZeRO sharding, it gives you the frozen reference model for free, and it limits how far the features can drift, which is what the KL term is trying to achieve anyway. Full fine-tuning with the existing train_zero1.sh DeepSpeed machinery is a later ablation.
10.2 Why the compiled inference path cannot be used in the learner
infer_action with compile_action_infer=true compiles the cache prefill and the per-step denoise with torch.compile(mode="reduce-overhead"), which captures CUDA graphs. The code even has to clone the cache outputs because the graph owns its buffers and overwrites them on replay. That path runs under @torch.no_grad() and cannot produce gradients. The learner therefore needs a second implementation of the same computation that is uncompiled (or compiled in the default mode), accepts a batch dimension greater than one, and runs with gradients enabled. Write it as a thin grad-capable twin of _denoise_action_with_video_cache and prefill_video_cache_tensor, and verify in Chapter 13 that it produces the same numbers as the compiled path.
Rollout workers, on the other hand, should keep using the compiled path, with the Flow-SDE noise injection added between compiled denoising calls (the noise addition is a cheap elementwise op that can live outside the compiled region).
10.3 Memory
The actor update backpropagates through each stochastic denoising step (one with πRL's sampler, K_sde with the variant), each of which is a 30-layer action expert forward attending into a 30-layer cache. Enable gradient checkpointing (mot_checkpoint_mixed_attn: true in the model config turns it on for both experts). With a minibatch of 64 macro-steps, bf16, and checkpointing, the actor phase fits comfortably on an 80 GB GPU; the world-model phase is the same cost as ordinary Fast-WAM training at batch 16 and already fits under the existing trainer's settings.
10.4 Synchronising weights to the workers
After each iteration, the workers need the new actor weights (and, if the video expert is being updated, the new LoRA weights). The compiled CUDA graphs in the workers were captured against the old parameter tensors; after loading new weights, delete the _prefill_video_cache_compiled and _denoise_action_with_video_cache_compiled attributes on the model so that infer_action recaptures them on the next call. Loading into the existing parameter tensors in place (param.data.copy_) rather than replacing them may let the graphs survive, but deletion and recapture is the safe default.