From 4f3e7cca22c88c66fa9b5b3ec31e00f49516a947 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Wed, 29 Jul 2026 04:43:44 +0000 Subject: [PATCH 1/5] fix: keep lora_A/lora_B paired in LoRA IPC weight buckets sglang merges LoRA adapters per IPC payload, so splitting lora_A and lora_B across update_weight_buffer_size buckets silently skips layers. Group tensors by layer prefix before bucketing and add a GPU test. --- .../diffusion_update_weight_utils.py | 113 ++++++++++++------ tests/fast-gpu/test_lora_weight_sync.py | 52 +++++++- 2 files changed, 127 insertions(+), 38 deletions(-) diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index 18dcbf61..a53237a4 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/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,59 @@ 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 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]) + layer_groups.append(tensors) + return layer_groups, unmapped_keys, num_lora_keys + + +def bucket_lora_layer_groups( + layer_groups: Sequence[Sequence[tuple[str, torch.Tensor]]], + buffer_size: int, +) -> list[list[tuple[str, torch.Tensor]]]: + """Bucket LoRA tensors by whole layer groups; never split lora_A/lora_B across buckets.""" + buckets: list[list[tuple[str, torch.Tensor]]] = [] + bucket: list[tuple[str, torch.Tensor]] = [] + bucket_size = 0 + + for group in layer_groups: + group_size = sum(_tensor_nbytes(tensor) for _, tensor in group) + if bucket and bucket_size + group_size >= buffer_size: + buckets.append(bucket) + bucket = [] + bucket_size = 0 + bucket.extend(group) + bucket_size += group_size + + if bucket: + buckets.append(bucket) + return buckets + + class DiffusionUpdateWeight(abc.ABC): """Base updater used by diffusion training actors.""" @@ -413,57 +470,41 @@ 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(): - 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: + layer_groups, unmapped_keys, num_lora_keys = collect_lora_layer_groups(model.state_dict()) + prepared_groups = [ + [(sgld_name, self._prepare_lora_param(param)) for sgld_name, param in group] + for group in layer_groups + ] + buckets = bucket_lora_layer_groups(prepared_groups, self.args.update_weight_buffer_size) + for bucket in buckets: 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: + 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, + len(buckets), len(unmapped_keys), ) if sample_layers: diff --git a/tests/fast-gpu/test_lora_weight_sync.py b/tests/fast-gpu/test_lora_weight_sync.py index 0646191c..fb2d3e5f 100644 --- a/tests/fast-gpu/test_lora_weight_sync.py +++ b/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 From 24a62f11c186342d092e4954bfadec13a0d300b4 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Wed, 29 Jul 2026 04:48:18 +0000 Subject: [PATCH 2/5] style: apply black formatting for LoRA IPC bucketing change --- miles/backends/fsdp_utils/diffusion_update_weight_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index a53237a4..e968169c 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -484,8 +484,7 @@ def update_weights(self) -> None: for target_module, model in self.models.items(): layer_groups, unmapped_keys, num_lora_keys = collect_lora_layer_groups(model.state_dict()) prepared_groups = [ - [(sgld_name, self._prepare_lora_param(param)) for sgld_name, param in group] - for group in layer_groups + [(sgld_name, self._prepare_lora_param(param)) for sgld_name, param in group] for group in layer_groups ] buckets = bucket_lora_layer_groups(prepared_groups, self.args.update_weight_buffer_size) for bucket in buckets: From 8778115c69fc0cd05bb7f45cd66b20f46a6a0480 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Wed, 29 Jul 2026 07:06:43 +0000 Subject: [PATCH 3/5] fix: materialize LoRA IPC tensors lazily per bucket Avoid preparing all lora_A/lora_B params on GPU before bucketing, which spiked peak memory by the full adapter size. Bucket using raw state-dict refs first, then _prepare_lora_param only for the bucket being flushed. --- .../fsdp_utils/diffusion_update_weight_utils.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index e968169c..a4d3848f 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -483,13 +483,11 @@ 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()) - prepared_groups = [ - [(sgld_name, self._prepare_lora_param(param)) for sgld_name, param in group] for group in layer_groups - ] - buckets = bucket_lora_layer_groups(prepared_groups, self.args.update_weight_buffer_size) - for bucket in buckets: + raw_buckets = bucket_lora_layer_groups(layer_groups, self.args.update_weight_buffer_size) + for raw_bucket in raw_buckets: + prepared_bucket = [(sgld_name, self._prepare_lora_param(param)) for sgld_name, param in raw_bucket] self.wait_and_update_bucket_weights( - bucket, + prepared_bucket, target_module, weight_update_mode=LORA_IPC_WEIGHT_UPDATE_MODE, ) @@ -503,7 +501,7 @@ def update_weights(self) -> None: target_module, num_lora_keys, num_layers, - len(buckets), + len(raw_buckets), len(unmapped_keys), ) if sample_layers: From be3ff70e1c47f5922c07f5c1e0bf017ddabe8c94 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Wed, 29 Jul 2026 07:14:57 +0000 Subject: [PATCH 4/5] fix: prepare LoRA IPC tensors only when filling each bucket Stream layer groups into buckets and call _prepare_lora_param on append, instead of building a fully prepared bucket list before flush. Keeps peak GPU memory bounded to the current bucket while preserving A/B pairing. --- .../diffusion_update_weight_utils.py | 30 +++++++++++++++---- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index a4d3848f..7bde38b3 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -483,14 +483,34 @@ 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()) - raw_buckets = bucket_lora_layer_groups(layer_groups, self.args.update_weight_buffer_size) - for raw_bucket in raw_buckets: - prepared_bucket = [(sgld_name, self._prepare_lora_param(param)) for sgld_name, param in raw_bucket] + bucket: list[tuple[str, torch.Tensor]] = [] + bucket_size = 0 + 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, + ) + num_buckets += 1 + bucket = [] + bucket_size = 0 + + 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( - prepared_bucket, + bucket, target_module, weight_update_mode=LORA_IPC_WEIGHT_UPDATE_MODE, ) + num_buckets += 1 if self.weight_version <= 2 and dist.is_initialized() and dist.get_rank() == 0: _, num_layers, sample_layers, _ = PeftLoRAKeyMapper.summarize_mapping(model.state_dict()) @@ -501,7 +521,7 @@ def update_weights(self) -> None: target_module, num_lora_keys, num_layers, - len(raw_buckets), + num_buckets, len(unmapped_keys), ) if sample_layers: From c46d6fffe7528047dae8efc8daa9615be0ec4973 Mon Sep 17 00:00:00 2001 From: niehen6174 Date: Wed, 29 Jul 2026 08:02:07 +0000 Subject: [PATCH 5/5] refactor: drop unused bucket helper and assert LoRA A/B pairs Remove bucket_lora_layer_groups now that bucketing is inlined in update_weights. Validate each collected layer group matches [layer.lora_A, layer.lora_B] before IPC sync. --- .../diffusion_update_weight_utils.py | 30 +++++-------------- 1 file changed, 7 insertions(+), 23 deletions(-) diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index 7bde38b3..9142ce0f 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -107,6 +107,12 @@ 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]: @@ -129,33 +135,11 @@ def collect_lora_layer_groups( 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 -def bucket_lora_layer_groups( - layer_groups: Sequence[Sequence[tuple[str, torch.Tensor]]], - buffer_size: int, -) -> list[list[tuple[str, torch.Tensor]]]: - """Bucket LoRA tensors by whole layer groups; never split lora_A/lora_B across buckets.""" - buckets: list[list[tuple[str, torch.Tensor]]] = [] - bucket: list[tuple[str, torch.Tensor]] = [] - bucket_size = 0 - - for group in layer_groups: - group_size = sum(_tensor_nbytes(tensor) for _, tensor in group) - if bucket and bucket_size + group_size >= buffer_size: - buckets.append(bucket) - bucket = [] - bucket_size = 0 - bucket.extend(group) - bucket_size += group_size - - if bucket: - buckets.append(bucket) - return buckets - - class DiffusionUpdateWeight(abc.ABC): """Base updater used by diffusion training actors."""