feat: LoRA IPC weight sync for colocated diffusion GRPO#33
Merged
Conversation
Sync only lora_A/lora_B tensors to rollout via weight_update_mode=lora_merge, avoiding full merged DiT weight gather and IPC transfer each step.
Select DiffusionUpdateWeightFromTensorLoRAIPC when enabled, and pass lora_merge metadata plus target modules to rollout weight updates.
Rockdu
reviewed
Jul 15, 2026
Fill per-model LoRA target modules during arg validation so rollout IPC sync does not re-resolve config at engine startup. Add CPU CI tests for PeftLoRAKeyMapper and server-args wiring.
Collaborator
|
Nice, merged! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Add an opt-in LoRA IPC weight sync path for colocated diffusion GRPO training. When
--lora-ipc-weight-syncis enabled, miles pushes onlylora_A/lora_Btensors to rollout via IPC withweight_update_mode=lora_merge, instead of all-gathering full base weights.This reduces weight-sync cost from near-100% model parameters to ~1–5% LoRA parameters, while keeping rollout behavior equivalent for correctly mapped models.
Modifications
miles/utils/arguments.py--lora-ipc-weight-syncmiles/backends/fsdp_utils/diffusion_update_weight_utils.pyPeftLoRAKeyMapperfor PEFT → sglang-d LoRA key mappingDiffusionUpdateWeightFromTensorLoRAIPCweight_update_mode,lora_alpha, andlora_rankmiles/backends/fsdp_utils/actor.pyDiffusionUpdateWeightFromTensorLoRAIPCwhenuse_lora and lora_ipc_weight_syncmiles/backends/sglang_diffusion_utils/sglang_diffusion_engine.pyweight_update_mode,lora_alpha,lora_rankin/update_weights_from_tensorlora_target_modulesto rollout server args when LoRA IPC is enabled--lora-ipc-weight-syncin:scripts/run-diffusion-grpo-sd3-ocr-sglang.shscripts/run-diffusion-grpo-ocr-2gpu-flowgrpo-aligned.shRollout-side LoRA IPC merge is implemented in sglang-d.
Performance
Benchmarks were run on 2-GPU colocate OCR GRPO with 5 rollout cycles, comparing merge baseline vs LoRA IPC.
Weight sync time (
perf/update_weights_time)Rollout-side verification:
Train/inference consistency on Qwen-Image
We compared LoRA IPC vs the existing merged-weight sync path on the same Qwen-Image OCR colocate setup (10 rollout steps). The plot below shows
train/log_prob_mean_abs_diffover training steps:Both paths remain on the same order of magnitude (~1e-5 to ~5e-5). Step 1 is bit-identical, and later steps stay numerically close. This suggests LoRA IPC does not materially affect train/rollout log-prob consistency on Qwen-Image.