Skip to main content
Miles-diffusion trains LoRA adapters on the FSDP actor and syncs them to sglang-diffusion rollout engines each iteration. The rollout engine has no PEFT layers β€” weights arrive either as merged base tensors or as raw lora_A / lora_B pairs for local merge. For colocated LoRA training, prefer IPC merge β€” push only lora_A / lora_B and let the rollout engine merge locally:
--lora-rank / --lora-alpha vary by recipe (e.g. SD3 uses 32/64; some others use 64/128). Without --lora-ipc-weight-sync, LoRA still trains but merges on the train side and pushes full merged weights (Β§3). IPC merge requires --colocate.

​
2. Key flags

--lora-ipc-weight-sync requires both --use-lora and --colocate. Without colocation, CUDA IPC handles cannot cross the train/rollout process boundary.
LoRA target modules default from the model family’s TrainPipelineConfig.lora_target_modules. For SD3, see SD3 model guide.

​
3. Three weight-sync strategies

Selection logic in miles/backends/fsdp_utils/actor.py: Implementation: miles/backends/fsdp_utils/diffusion_update_weight_utils.py.

​
Full-weight sync (no LoRA)

FSDP shards are all-gathered into FlattenedTensorBucket objects, serialized via CUDA IPC, and sent to the rollout engine’s update_weights_from_tensor(load_format="flattened_bucket").

​
Train-side merge (LoRA, no IPC)

For each base layer with adapters, the updater computes:
on the fly, strips PEFT key prefixes, and pushes standard weight names that sglang-d expects (e.g. transformer_blocks.0.attn.to_q.weight). DiffusionUpdateWeightFromTensorLoRAIPC:
  1. collect_lora_layer_groups() groups state-dict entries by layer prefix so lora_A and lora_B for the same layer always stay together.
  2. PeftLoRAKeyMapper.to_sgld_name() maps PEFT keys to sglang-d names (e.g. transformer_blocks.0.attn.to_q.lora_A).
  3. FSDP shard all-gather β†’ pack into buckets capped by --update-weight-buffer-size (recipes use 2 GB) β†’ CUDA IPC.
  4. Rollout engine receives weight_update_mode="lora_merge" with lora_alpha and lora_rank.
Bucket packing: the IPC updater iterates layer groups, not individual tensors. When adding the next group would exceed --update-weight-buffer-size, the current bucket is flushed first; the whole group (both lora_A and lora_B) then starts the next bucket. Pairs are never split across buckets. Constant: LORA_IPC_WEIGHT_UPDATE_MODE = "lora_merge". Rollout-side merge precision is controlled by environment variable SGLANG_DIFFUSION_LORA_MERGE_FP32:
  • "1" when --diffusion-forward-dtype fp32
  • "0" otherwise (fp16 merge)
Set automatically in RolloutManager when spawning engines. On the first few syncs, rank 0 logs lines like:
After FSDP all-gather, serialized buckets are collected on the gather-src rank only; that rank calls the rollout engine. If sync stalls or VRAM grows across rollouts, check trainer logs for LoRA IPC weight sync lines and Ray worker stderr under ~/.ray/session_latest/logs/.

​
4. Internals

​
5. Limitations

  • Colocate only β€” disaggregated train/rollout is not supported for LoRA IPC.
  • Single adapter per run β€” one set of --lora-* flags per job.
  • FSDP backend β€” LoRA weight sync is implemented for the FSDP diffusion actor; Megatron LLM LoRA (in upstream Miles) uses a separate path.