feat: LoRA IPC weight sync for colocated diffusion GRPO

miles/backends/fsdp_utils/actor.py

@@ -29,7 +29,11 @@
 )
 from . import checkpoint
 from .arguments import deterministic_capable_flash_fns
-from .diffusion_update_weight_utils import DiffusionUpdateWeightFromTensor, DiffusionUpdateWeightFromTensorLoRA
+from .diffusion_update_weight_utils import (
+    DiffusionUpdateWeightFromTensor,
+    DiffusionUpdateWeightFromTensorLoRA,
+    DiffusionUpdateWeightFromTensorLoRAIPC,
+)
 from .lr_scheduler import get_lr_scheduler
 from .parallel import create_fsdp_parallel_state
 
@@ -186,6 +190,8 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
         # sglang-d now supports /update_weights_from_tensor (PR #20464).
         if self.args.debug_train_only:
             self.weight_updater = None
+        elif self.args.use_lora and self.args.lora_ipc_weight_sync:
+            self.weight_updater = DiffusionUpdateWeightFromTensorLoRAIPC(self.args, self.models)
         elif self.args.use_lora:
             self.weight_updater = DiffusionUpdateWeightFromTensorLoRA(self.args, self.models)
         else:

miles/backends/fsdp_utils/diffusion_update_weight_utils.py

@@ -1,8 +1,9 @@
 import abc
 import logging
 import os
+import re
 from argparse import Namespace
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 
 import ray
 import torch
@@ -33,6 +34,70 @@
 
 logger = logging.getLogger(__name__)
 
+LORA_IPC_WEIGHT_UPDATE_MODE = "lora_merge"
+
+
+class PeftLoRAKeyMapper:
+    """Map PEFT LoRA state-dict keys to sglang-d tensor names for IPC sync."""
+
+    _LORA_KEY_RE = re.compile(r"\.lora_([AB])(?:\.[^.]+)?(?:\.weight)?$")
+    _PEFT_PREFIX = "base_model.model."
+
+    @classmethod
+    def is_lora_key(cls, name: str) -> bool:
+        return ".lora_A" in name or ".lora_B" in name
+
+    @classmethod
+    def to_sgld_name(cls, name: str) -> str | None:
+        """Map a PEFT state-dict key to sglang-d LoRA tensor name."""
+        if not cls.is_lora_key(name):
+            return None
+
+        stripped = name
+        if stripped.startswith(cls._PEFT_PREFIX):
+            stripped = stripped[len(cls._PEFT_PREFIX) :]
+
+        match = cls._LORA_KEY_RE.search(stripped)
+        if match is None:
+            return None
+
+        layer_prefix = stripped[: match.start()]
+        ab = match.group(1)
+        return f"{layer_prefix}.lora_{ab}"
+
+    @classmethod
+    def collect_sgld_names(cls, state_dict: Mapping[str, torch.Tensor]) -> set[str]:
+        names: set[str] = set()
+        for key in state_dict:
+            sgld_name = cls.to_sgld_name(key)
+            if sgld_name is not None:
+                names.add(sgld_name)
+        return names
+
+    @classmethod
+    def collect_layer_prefixes(cls, state_dict: Mapping[str, torch.Tensor]) -> set[str]:
+        return {name.rsplit(".lora_", 1)[0] for name in cls.collect_sgld_names(state_dict)}
+
+    @classmethod
+    def summarize_mapping(
+        cls,
+        state_dict: Mapping[str, torch.Tensor],
+    ) -> tuple[int, int, list[str], list[str]]:
+        """Return (num_tensors, num_layers, sample_layer_prefixes, unmapped_peft_keys)."""
+        sgld_names: list[str] = []
+        unmapped: list[str] = []
+        for key in state_dict:
+            if not cls.is_lora_key(key):
+                continue
+            sgld_name = cls.to_sgld_name(key)
+            if sgld_name is None:
+                unmapped.append(key)
+            else:
+                sgld_names.append(sgld_name)
+        layer_prefixes = {name.rsplit(".lora_", 1)[0] for name in sgld_names}
+        sample = sorted(layer_prefixes)[:5]
+        return len(sgld_names), len(layer_prefixes), sample, unmapped
+
 
 class DiffusionUpdateWeight(abc.ABC):
     """Base updater used by diffusion training actors."""
@@ -88,12 +153,23 @@ def _update_component_weights(self, target_module: str, model: torch.nn.Module)
             self.wait_and_update_bucket_weights(bucket, target_module)
             del bucket
 
-    def wait_and_update_bucket_weights(self, bucket, target_module: str):
+    def wait_and_update_bucket_weights(self, bucket, target_module: str, weight_update_mode=None):
         bucket = [(name, param.wait()) if hasattr(param, "wait") else (name, param) for name, param in bucket]
-        self.update_bucket_weights(bucket, target_module, weight_version=self.weight_version)
+        self.update_bucket_weights(
+            bucket,
+            target_module,
+            weight_version=self.weight_version,
+            weight_update_mode=weight_update_mode,
+        )
 
     @abc.abstractmethod
-    def update_bucket_weights(self, named_tensors, target_module: str, weight_version=None) -> None:
+    def update_bucket_weights(
+        self,
+        named_tensors,
+        target_module: str,
+        weight_version=None,
+        weight_update_mode: str | None = None,
+    ) -> None:
         pass
 
 
@@ -123,7 +199,13 @@ def connect_rollout_engines(
                 # Calculate TP rank within this SGLang engine group.
                 self.tp_rank = dist.get_rank() - start_rank
 
-    def update_bucket_weights(self, named_tensors, target_module: str, weight_version=None) -> None:
+    def update_bucket_weights(
+        self,
+        named_tensors,
+        target_module: str,
+        weight_version=None,
+        weight_update_mode: str | None = None,
+    ) -> None:
         monkey_patch_torch_reductions()
         logger.info("Using flattened tensor bucket (diffusion updater, module=%s)", target_module)
         named_tensors_by_dtypes = {}
@@ -173,6 +255,10 @@ def update_bucket_weights(self, named_tensors, target_module: str, weight_versio
                     "target_modules": [target_module],
                     "weight_version": str(weight_version),
                 }
+                if weight_update_mode is not None:
+                    kwargs["weight_update_mode"] = weight_update_mode
+                    kwargs["lora_alpha"] = self.args.lora_alpha
+                    kwargs["lora_rank"] = self.args.lora_rank
                 ref = self._ipc_engine.update_weights_from_tensor.remote(**kwargs)
                 ray.get(ref)
 
@@ -322,3 +408,78 @@ def _verify_weight_sync(self, pairs: list[tuple[str, torch.Tensor]], target_modu
         all_equal = all(s == first_sum for _, s in engine_sums)
         pretty = "  ".join(f"eng{idx}={s[:16] if isinstance(s, str) else s}" for idx, s in engine_sums)
         logger.warning(f"[weight_sync verify v{self.weight_version} cross-engine] " f"all_equal={all_equal}  {pretty}")
+
+
+class DiffusionUpdateWeightFromTensorLoRAIPC(DiffusionUpdateWeightFromTensor):
+    """Push only lora_A/lora_B tensors; rollout merges locally via weight_update_mode=lora_merge."""
+
+    def update_weights(self) -> None:
+        self.weight_version += 1
+        for target_module, model in self.models.items():
+            bucket: list[tuple[str, torch.Tensor]] = []
+            bucket_size = 0
+            num_lora_keys = 0
+            unmapped_keys: list[str] = []
+
+            for name, param in model.state_dict().items():
+                if not PeftLoRAKeyMapper.is_lora_key(name):
+                    continue
+                sgld_name = PeftLoRAKeyMapper.to_sgld_name(name)
+                if sgld_name is None:
+                    unmapped_keys.append(name)
+                    continue
+
+                param = param.cuda()
+                if isinstance(param, DTensor):
+                    param = param.redistribute(
+                        placements=[Replicate()] * param.device_mesh.ndim,
+                        async_op=True,
+                    ).to_local()
+
+                sz = param.numel() * param.element_size()
+                if bucket and bucket_size + sz >= self.args.update_weight_buffer_size:
+                    self.wait_and_update_bucket_weights(
+                        bucket,
+                        target_module,
+                        weight_update_mode=LORA_IPC_WEIGHT_UPDATE_MODE,
+                    )
+                    bucket, bucket_size = [], 0
+
+                bucket.append((sgld_name, param))
+                bucket_size += sz
+                num_lora_keys += 1
+
+            if bucket:
+                self.wait_and_update_bucket_weights(
+                    bucket,
+                    target_module,
+                    weight_update_mode=LORA_IPC_WEIGHT_UPDATE_MODE,
+                )
+
+            if self.weight_version <= 2 and dist.get_rank() == 0:
+                _, num_layers, sample_layers, _ = PeftLoRAKeyMapper.summarize_mapping(model.state_dict())
+                logger.info(
+                    "LoRA IPC weight sync v%s [%s]: pushed %d lora tensors, " "%d layer prefixes (unmapped=%d)",
+                    self.weight_version,
+                    target_module,
+                    num_lora_keys,
+                    num_layers,
+                    len(unmapped_keys),
+                )
+                if sample_layers:
+                    logger.info(
+                        "LoRA IPC [%s] sample layer prefixes: %s",
+                        target_module,
+                        sample_layers,
+                    )
+                if unmapped_keys:
+                    logger.warning(
+                        "LoRA IPC unmapped PEFT keys [%s] (first 5): %s",
+                        target_module,
+                        unmapped_keys[:5],
+                    )
+                if num_lora_keys == 0:
+                    logger.error(
+                        "LoRA IPC [%s]: no lora tensors found in training state_dict",
+                        target_module,
+                    )

miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py

@@ -229,6 +229,9 @@ def update_weights_from_tensor(
         load_format: str | None = None,
         target_modules: list[str] | None = None,
         weight_version: str | None = None,
+        weight_update_mode: str | None = None,
+        lora_alpha: int | None = None,
+        lora_rank: int | None = None,
     ):
         """
         Update model weights from tensor data. The HTTP server will only post meta data, and the real weights will be copied directly from GPUs.
@@ -244,6 +247,12 @@ def update_weights_from_tensor(
             payload["target_modules"] = target_modules
         if weight_version is not None:
             payload["weight_version"] = weight_version
+        if weight_update_mode is not None:
+            payload["weight_update_mode"] = weight_update_mode
+        if lora_alpha is not None:
+            payload["lora_alpha"] = lora_alpha
+        if lora_rank is not None:
+            payload["lora_rank"] = lora_rank
         return self._make_request(
             "update_weights_from_tensor",
             payload,
@@ -291,7 +300,8 @@ def simulate_crash(self):
 
 
 def _compute_server_args(args, host, port, nccl_port):
-    # Only fields SGL-D's ServerArgs accepts; GPU pinning is done in _init_normal via CUDA_VISIBLE_DEVICES.
+    # Only set fields SGL-D's ServerArgs actually accepts. GPU pinning is done
+    # in `_init_normal` via CUDA_VISIBLE_DEVICES — SGL-D has no base_gpu_id arg.
     kwargs = {
         "model_path": args.diffusion_model,
         "trust_remote_code": True,
@@ -316,6 +326,8 @@ def _compute_server_args(args, host, port, nccl_port):
         if hasattr(args, f"sglang_{attr.name}") and attr.name not in kwargs:
             kwargs[attr.name] = getattr(args, f"sglang_{attr.name}")
 
+    if getattr(args, "use_lora", False) and getattr(args, "lora_ipc_weight_sync", False):
+        kwargs["lora_target_modules"] = args.lora_target_modules
     # dit_precision / vae_precision are PipelineConfig fields, not ServerArgs, so forward them explicitly (only when changed from the class default, to avoid clobbering a subclass override).
     from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
 

miles/utils/arguments.py

@@ -1039,6 +1039,15 @@ def add_debug_arguments(parser):
                 default=None,
                 help="Override LoRA target modules. Default: per-model from TrainPipelineConfig.",
             )
+            parser.add_argument(
+                "--lora-ipc-weight-sync",
+                action="store_true",
+                default=False,
+                help=(
+                    "Sync only lora_A/lora_B to rollout via IPC with weight_update_mode=lora_merge "
+                    "(requires matching sglang-d LoRAPipeline support)."
+                ),
+            )
             parser.add_argument(
                 "--diffusion-init-lora-weight",
                 type=str,
@@ -1360,6 +1369,17 @@ def miles_validate_args(args):
         if args.model_backend_path is None:
             args.model_backend_path = cfg_cls.model_backend_path
         cfg_cls.validate_args(args)
+        if args.use_lora and args.lora_target_modules is None:
+            args.lora_target_modules = list(cfg_cls.lora_target_modules)
+
+    if getattr(args, "lora_ipc_weight_sync", False):
+        if not args.use_lora:
+            raise ValueError("--lora-ipc-weight-sync requires --use-lora")
+        if not args.lora_target_modules:
+            raise ValueError(
+                "--lora-ipc-weight-sync requires LoRA target modules; "
+                "set --diffusion-model (for per-model defaults) or --lora-target-modules."
+            )
 
     if args.dump_details is not None:
         args.save_debug_rollout_data = f"{args.dump_details}/rollout_data/{{rollout_id}}.pt"

scripts/run-diffusion-grpo-ocr-2gpu-flowgrpo-aligned.sh

@@ -63,6 +63,7 @@ hf download --repo-type dataset rockdu/miles-diffusion-datasets \
   --num-gpus-per-node 2 \
   --colocate \
   --use-lora \
+  --lora-ipc-weight-sync \
   --lora-rank 64 \
   --lora-alpha 128 \
   --diffusion-init-lora-weight gaussian \

scripts/run-diffusion-grpo-sd3-ocr-sglang.sh

@@ -111,6 +111,7 @@ python -u "${ROOT_DIR}/train_diffusion.py" \
   --use-miles-router \
   --sglang-server-concurrency 8 \
   --use-lora \
+  --lora-ipc-weight-sync \
   --lora-rank 32 \
   --lora-alpha 64 \
   --diffusion-init-lora-weight gaussian \

tests/fast/backends/fsdp_utils/test_peft_lora_key_mapper.py

@@ -0,0 +1,68 @@
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[])
+
+import pytest
+import torch
+
+from miles.backends.fsdp_utils.diffusion_update_weight_utils import PeftLoRAKeyMapper
+
+_QWEN_A = "base_model.model.transformer_blocks.0.attn.to_q.lora_A.default.weight"
+_QWEN_B = "base_model.model.transformer_blocks.0.attn.to_q.lora_B.default.weight"
+_SD3_A = "base_model.model.transformer_blocks.1.attn.add_k_proj.lora_A.weight"
+_MLP_B = "base_model.model.transformer_blocks.2.img_mlp.net.0.proj.lora_B.default.weight"
+
+
+class TestPeftLoRAKeyMapper:
+    @pytest.mark.parametrize(
+        "key,expected",
+        [
+            (_QWEN_A, "transformer_blocks.0.attn.to_q.lora_A"),
+            (_QWEN_B, "transformer_blocks.0.attn.to_q.lora_B"),
+            (_SD3_A, "transformer_blocks.1.attn.add_k_proj.lora_A"),
+            (_MLP_B, "transformer_blocks.2.img_mlp.net.0.proj.lora_B"),
+            ("transformer_blocks.0.attn.to_v.lora_A.default.weight", "transformer_blocks.0.attn.to_v.lora_A"),
+        ],
+    )
+    def test_to_sgld_name_maps_peft_keys(self, key, expected):
+        assert PeftLoRAKeyMapper.to_sgld_name(key) == expected
+
+    @pytest.mark.parametrize(
+        "key",
+        [
+            "transformer_blocks.0.attn.to_q.weight",
+            "base_model.model.norm.weight",
+            "lora_A.default.weight",
+        ],
+    )
+    def test_to_sgld_name_returns_none_for_non_lora_keys(self, key):
+        assert PeftLoRAKeyMapper.to_sgld_name(key) is None
+
+    def test_is_lora_key(self):
+        assert PeftLoRAKeyMapper.is_lora_key(_QWEN_A)
+        assert PeftLoRAKeyMapper.is_lora_key(_QWEN_B)
+        assert not PeftLoRAKeyMapper.is_lora_key("transformer_blocks.0.attn.to_q.weight")
+
+    def test_collect_sgld_names_and_layer_prefixes(self):
+        state_dict = {
+            _QWEN_A: torch.zeros(4, 8),
+            _QWEN_B: torch.zeros(8, 4),
+            "base_model.model.norm.weight": torch.zeros(8),
+        }
+        sgld_names = PeftLoRAKeyMapper.collect_sgld_names(state_dict)
+        assert sgld_names == {
+            "transformer_blocks.0.attn.to_q.lora_A",
+            "transformer_blocks.0.attn.to_q.lora_B",
+        }
+        assert PeftLoRAKeyMapper.collect_layer_prefixes(state_dict) == {"transformer_blocks.0.attn.to_q"}
+
+    def test_summarize_mapping_reports_unmapped_lora_keys(self):
+        state_dict = {
+            _QWEN_A: torch.zeros(4, 8),
+            "base_model.model.weird.lora_A.default.weight.extra": torch.zeros(4, 8),
+        }
+        num_tensors, num_layers, sample_layers, unmapped = PeftLoRAKeyMapper.summarize_mapping(state_dict)
+        assert num_tensors == 1
+        assert num_layers == 1
+        assert sample_layers == ["transformer_blocks.0.attn.to_q"]
+        assert unmapped == ["base_model.model.weird.lora_A.default.weight.extra"]