diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index 18dcbf61..9142ce0f 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,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,41 +454,39 @@ 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( @@ -455,15 +494,18 @@ def update_weights(self) -> None: 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: 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