fix: keep lora_A/lora_B paired in LoRA IPC weight buckets

miles/backends/fsdp_utils/diffusion_update_weight_utils.py

@@ -65,6 +65,10 @@ def to_sgld_name(cls, name: str) -> str | None:
         ab = match.group(1)
         return f"{layer_prefix}.lora_{ab}"
 
+    @classmethod
+    def layer_prefix(cls, sgld_name: str) -> str:
+        return sgld_name.rsplit(".lora_", 1)[0]
+
     @classmethod
     def collect_sgld_names(cls, state_dict: Mapping[str, torch.Tensor]) -> set[str]:
         names: set[str] = set()
@@ -99,6 +103,43 @@ def summarize_mapping(
         return len(sgld_names), len(layer_prefixes), sample, unmapped
 
 
+def _tensor_nbytes(tensor: torch.Tensor) -> int:
+    return tensor.numel() * tensor.element_size()
+
+
+def _assert_lora_ab_pair(layer_prefix: str, tensors: list[tuple[str, torch.Tensor]]) -> None:
+    names = [name for name, _ in tensors]
+    expected = [f"{layer_prefix}.lora_A", f"{layer_prefix}.lora_B"]
+    assert names == expected, f"LoRA layer {layer_prefix!r} expected {expected}, got {names}"
+
+
+def collect_lora_layer_groups(
+    state_dict: Mapping[str, torch.Tensor],
+) -> tuple[list[list[tuple[str, torch.Tensor]]], list[str], int]:
+    """Group LoRA state-dict entries by layer prefix so lora_A/lora_B stay together."""
+    groups_by_layer: dict[str, list[tuple[str, torch.Tensor]]] = {}
+    unmapped_keys: list[str] = []
+    num_lora_keys = 0
+
+    for name, param in 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
+        layer_prefix = PeftLoRAKeyMapper.layer_prefix(sgld_name)
+        groups_by_layer.setdefault(layer_prefix, []).append((sgld_name, param))
+        num_lora_keys += 1
+
+    layer_groups: list[list[tuple[str, torch.Tensor]]] = []
+    for layer_prefix in sorted(groups_by_layer):
+        tensors = sorted(groups_by_layer[layer_prefix], key=lambda item: item[0])
+        _assert_lora_ab_pair(layer_prefix, tensors)
+        layer_groups.append(tensors)
+    return layer_groups, unmapped_keys, num_lora_keys
+
+
 class DiffusionUpdateWeight(abc.ABC):
     """Base updater used by diffusion training actors."""
 
@@ -413,57 +454,58 @@ def _verify_weight_sync(self, pairs: list[tuple[str, torch.Tensor]], target_modu
 class DiffusionUpdateWeightFromTensorLoRAIPC(DiffusionUpdateWeightFromTensor):
     """Push only lora_A/lora_B tensors; rollout merges locally via weight_update_mode=lora_merge."""
 
+    def _prepare_lora_param(self, param: torch.Tensor) -> torch.Tensor:
+        param = param.cuda()
+        if isinstance(param, DTensor):
+            param = param.redistribute(
+                placements=[Replicate()] * param.device_mesh.ndim,
+                async_op=True,
+            ).to_local()
+        return param
+
     def update_weights(self) -> None:
         self.weight_version += 1
         for target_module, model in self.models.items():
+            layer_groups, unmapped_keys, num_lora_keys = collect_lora_layer_groups(model.state_dict())
             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:
+            num_buckets = 0
+            buffer_size = self.args.update_weight_buffer_size
+
+            for group in layer_groups:
+                group_size = sum(_tensor_nbytes(param) for _, param in group)
+                if bucket and bucket_size + group_size >= buffer_size:
                     self.wait_and_update_bucket_weights(
                         bucket,
                         target_module,
                         weight_update_mode=LORA_IPC_WEIGHT_UPDATE_MODE,
                     )
-                    bucket, bucket_size = [], 0
+                    num_buckets += 1
+                    bucket = []
+                    bucket_size = 0
 
-                bucket.append((sgld_name, param))
-                bucket_size += sz
-                num_lora_keys += 1
+                for sgld_name, param in group:
+                    bucket.append((sgld_name, self._prepare_lora_param(param)))
+                bucket_size += group_size
 
             if bucket:
                 self.wait_and_update_bucket_weights(
                     bucket,
                     target_module,
                     weight_update_mode=LORA_IPC_WEIGHT_UPDATE_MODE,
                 )
+                num_buckets += 1
 
-            if self.weight_version <= 2 and dist.get_rank() == 0:
+            if self.weight_version <= 2 and dist.is_initialized() 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)",
+                    "LoRA IPC weight sync v%s [%s]: pushed %d lora tensors, "
+                    "%d layer prefixes in %d buckets (unmapped=%d)",
                     self.weight_version,
                     target_module,
                     num_lora_keys,
                     num_layers,
+                    num_buckets,
                     len(unmapped_keys),
                 )
                 if sample_layers:

tests/fast-gpu/test_lora_weight_sync.py

@@ -12,7 +12,11 @@
 import torch
 from peft import LoraConfig, get_peft_model
 
-from miles.backends.fsdp_utils.diffusion_update_weight_utils import DiffusionUpdateWeightFromTensorLoRA
+from miles.backends.fsdp_utils.diffusion_update_weight_utils import (
+    DiffusionUpdateWeightFromTensorLoRA,
+    DiffusionUpdateWeightFromTensorLoRAIPC,
+    PeftLoRAKeyMapper,
+)
 
 
 class _TinyBlock(torch.nn.Module):
@@ -29,7 +33,18 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.buckets: list[list[tuple[str, torch.Tensor]]] = []
 
-    def wait_and_update_bucket_weights(self, bucket, target_module):
+    def wait_and_update_bucket_weights(self, bucket, target_module, weight_update_mode=None):
+        self.buckets.append([(name, tensor.clone()) for name, tensor in bucket])
+
+
+class _CaptureLoRAIPCUpdater(DiffusionUpdateWeightFromTensorLoRAIPC):
+    """Capture LoRA IPC buckets without rollout engines."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.buckets: list[list[tuple[str, torch.Tensor]]] = []
+
+    def wait_and_update_bucket_weights(self, bucket, target_module, weight_update_mode=None):
         self.buckets.append([(name, tensor.clone()) for name, tensor in bucket])
 
 
@@ -50,6 +65,23 @@ def _run_update(peft_model, buffer_size):
     return updater.buckets
 
 
+def _run_lora_ipc_update(peft_model, buffer_size):
+    updater = _CaptureLoRAIPCUpdater(Namespace(update_weight_buffer_size=buffer_size), {"transformer": peft_model})
+    updater.update_weights()
+    return updater.buckets
+
+
+def _assert_buckets_have_complete_ab_pairs(buckets):
+    for bucket in buckets:
+        by_layer: dict[str, set[str]] = {}
+        for name, _ in bucket:
+            prefix = PeftLoRAKeyMapper.layer_prefix(name)
+            ab = "A" if ".lora_A" in name else "B"
+            by_layer.setdefault(prefix, set()).add(ab)
+        for prefix, abs_ in by_layer.items():
+            assert abs_ == {"A", "B"}, f"incomplete LoRA pair for {prefix} in bucket: {abs_}"
+
+
 def test_lora_merge_and_name_mapping():
     peft_model = _make_peft_model()
     synced = {name: tensor for bucket in _run_update(peft_model, 1 << 30) for name, tensor in bucket}
@@ -70,6 +102,22 @@ def test_bucket_flush_respects_buffer_size():
     assert sum(len(bucket) for bucket in buckets) == 3
 
 
+def test_lora_ipc_bucket_keeps_ab_together():
+    peft_model = _make_peft_model()
+    lora_layer = peft_model.base_model.model.proj
+    pair_size = (
+        lora_layer.lora_A["default"].weight.numel() * lora_layer.lora_A["default"].weight.element_size()
+        + lora_layer.lora_B["default"].weight.numel() * lora_layer.lora_B["default"].weight.element_size()
+    )
+    buckets = _run_lora_ipc_update(peft_model, buffer_size=pair_size)
+    _assert_buckets_have_complete_ab_pairs(buckets)
+    synced = {name: tensor for bucket in buckets for name, tensor in bucket}
+    assert set(synced) == {"proj.lora_A", "proj.lora_B"}
+    A, B = lora_layer.lora_A["default"].weight, lora_layer.lora_B["default"].weight
+    torch.testing.assert_close(synced["proj.lora_A"], A)
+    torch.testing.assert_close(synced["proj.lora_B"], B)
+
+
 if __name__ == "__main__":
     import sys