Skip to main content
The FSDP trainer owns one trainable actor, optional independent frozen reference and teacher models, and an optional EMA optimizer. Each model contains the same configured components; Wan2.2 can therefore have both transformer and transformer_2 in every role.
Only actor parameters belong to AdamW. EMA stores averages for trainable actor parameters in a separate torch.optim.Optimizer; frozen base parameters and all buffers are excluded from averaging.

​
Model sources

Use an independent reference with:
Both sources use the same family backend, component names, input preparation, and device mesh. This provides a separate model instance, not cross-family distillation. The existing loss receives the reference prediction, and custom prepare/loss hooks can access ctx.reference_models by component name. A third source, the teacher (--teacher-load, with --teacher-lora-adapter-path for its adapters), is reserved for a DMD2-style loss that miles does not ship yet. Until then it has no built-in consumer: --teacher-load requires a custom prepare/loss hook, which reads it from ctx.teacher_models, and teacher inference should run inside torch.no_grad(). Each source may load pretrained PEFT LoRA adapters, one plain PEFT directory (what save_pretrained writes) per --update-weight-target-module entry, in the same order:
Wan2.2, for example, takes two directories, one per expert; any other number is rejected:
Without a role’s adapter path, that frozen role loads ordinary full weights, even when the actor uses LoRA. Full weights and base-only weights use the same loader. --lora-rank and --lora-alpha must match the actor adapter’s r and lora_alpha. Actor adapter weights are loaded before a training checkpoint is restored. Independent reference/teacher weights are always reconstructed from their configured sources, so retain those sources and their adapter checkpoints when resuming training. --ref-mode lora_base keeps the existing low-memory alternative: temporarily disable the actor adapter for a gradient-free base forward. It does not construct another base model. FSDP’s lazy init runs at load time, while the adapters are trainable, so the gradient dtype does not depend on which forward runs first. FSDP is resharded before the adapters are re-enabled, so PEFT restores adapter gradients on the shards rather than on gathered copies.

​
EMA and checkpointing

--use-ema creates the EMA optimizer. Under --loss-type nft it is π_old: the rollout samples with it (--rollout-weights ema) and the loss trains against its prediction. --ref-mode only selects the KL reference (none, lora_base, or ref; it defaults to lora_base when --diffusion-kl-beta > 0, else none), so NFT can use EMA and a reference together. EMA updates once after a rollout’s training completes, and its decay follows the actor’s optimizer step count, as in DiffusionNFT. Publishing or retrying a weight update never advances EMA. EMA evaluation and publication swap the EMA values into the actor’s shards and swap them back afterward, dropping FSDP’s gathered copies on both sides so no forward reads stale weights; the actor’s backward re-gathers its restored shards. EMA state lives on the same device as its actor parameter. Sleep moves it to pinned CPU memory together with AdamW state, and wake moves both back. Rollout publication pushes the current actor weights unless --rollout-weights ema asks for the EMA weights.
EMA uses the existing distributed optimizer-state checkpoint wrapper, including parameter names and DTensor sharding. --no-save-optim skips saving both AdamW and LR-scheduler state; --no-load-optim skips loading AdamW state only, so a saved LR-scheduler state is still restored. EMA remains independently saved and restored. Loading a legacy checkpoint without EMA seeds EMA from the loaded actor and logs that its averaging history starts over. Disabling EMA simply omits this state from loading.

​
Sleep/wake and native CPU offload

Reference and teacher CPU offload are opt-in. Sleep covers actor, reference, teacher, their buffers, AdamW state, and EMA state. Sleep clears completed-step gradients instead of transferring them. Transfers use pinned CPU storage and enqueue asynchronous copies before synchronization; Parameter objects and optimizer bindings are preserved. A natively CPU-offloaded model never sleeps, as in miles: its parameter shards stay on CPU and its buffers stay on GPU. Sleep retains the model’s own CPU tensors, without a second actor backup. Reference and teacher weights never change, so waking keeps their pinned host copies and each later sleep points the parameters back at them instead of copying them off the GPU. The helper lets FSDP rebuild its shard views and padding through Module._apply; uneven shards remain pinned. This requires a small FSDP-internal adaptation of its pinning flag.

​
Focused validation

CPU CI covers EMA averaging, DCP resume, model-source selection, loaded LoRA, and argument validation.