From cdf2994062545be6c6cee2be220346adff055f51 Mon Sep 17 00:00:00 2001 From: niehen6174 <1639206518@qq.com> Date: Tue, 1 Sep 2026 03:22:32 +0000 Subject: [PATCH 1/3] refactor(h3): drop trainer-side LoRA collector sglang-d lora_merge now maps PEFT/diffusers names and applies the FFN swap, so H3 no longer needs a dedicated IPC grouper. --- miles/backends/fsdp_utils/configs/h3.py | 9 +- .../configs/train_pipeline_config.py | 3 - .../diffusion_update_weight_utils.py | 33 +--- .../fsdp_utils/h3_weight_key_mapper.py | 182 ------------------ 4 files changed, 12 insertions(+), 215 deletions(-) delete mode 100644 miles/backends/fsdp_utils/h3_weight_key_mapper.py diff --git a/miles/backends/fsdp_utils/configs/h3.py b/miles/backends/fsdp_utils/configs/h3.py index a07824b16..7ac726fb1 100644 --- a/miles/backends/fsdp_utils/configs/h3.py +++ b/miles/backends/fsdp_utils/configs/h3.py @@ -21,7 +21,6 @@ class H3TrainPipelineConfig(TrainPipelineConfig): supports_cfg_training = False sde_timestep_divisor = 1000.0 optimizer_state_allowed_missing = ["audio"] - lora_layer_group_collector_path = "miles.backends.fsdp_utils.h3_weight_key_mapper.collect_h3_lora_layer_groups" lora_target_modules = [ "attn.to_q", @@ -34,10 +33,10 @@ class H3TrainPipelineConfig(TrainPipelineConfig): @classmethod def validate_args(cls, args: Namespace) -> None: - # sglang's H3 DiT renames modules and fuses Q/K/V, so weights only reach the - # rollout through the LoRA IPC path's layer grouper; any other sync mode would - # push names the engine drops with a warning, silently training nothing. - # SFT (--train-only) has no rollout engine and therefore no sync constraint. + # H3's rollout DiT fuses Q/K/V and rewrites FFN layout. Those transforms + # live in sglang-d's lora_merge IPC path; train-side merge / full-weight + # sync would push dense names the engine drops. SFT (--train-only) has + # no rollout engine and therefore no sync constraint. if not args.train_only and not (args.use_lora and args.lora_ipc_weight_sync): raise ValueError("H3 training requires --use-lora with --lora-ipc-weight-sync") diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index dec07c3f9..56420a62b 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -113,9 +113,6 @@ def apply_rollout_sampling_params( """ sde_timestep_divisor = 1.0 - # LoRA IPC layer grouper for families whose rollout module names or tensor layout - # differ from the trained diffusers ones; None keeps the generic PEFT grouping. - lora_layer_group_collector_path: str | None = None def configure(self, args) -> None: # noqa: B027 optional no-op hook, not abstract """Bind the request constants a family needs at train time; default binds none.""" diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index ff45c2136..0f68bf4d3 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -459,31 +459,14 @@ def _prepare_lora_param(self, param: torch.Tensor) -> torch.Tensor: def _collect_layer_groups( self, model: torch.nn.Module ) -> tuple[list[list[tuple[str, torch.Tensor]]], list[str], int]: - """Group this model's LoRA tensors into rollout layer names, per model family.""" - from miles.utils.misc import load_function - - collector_path = None - if self.args.train_pipeline_config_path: - collector_path = load_function(self.args.train_pipeline_config_path).lora_layer_group_collector_path - if collector_path is None: - return collect_lora_layer_groups(model.state_dict()) - - # A family whose rollout fuses several projections into one layer (H3's - # qkv_proj) combines adapters here, so DTensor shards must resolve first. - lora_state = { - name: self._prepare_lora_param(param) - for name, param in model.state_dict().items() - if PeftLoRAKeyMapper.is_lora_key(name) - } - layer_groups, unmapped_keys, num_lora_keys = load_function(collector_path)(lora_state) - if unmapped_keys: - # The rollout only warns about a name it cannot resolve, which would - # leave that adapter frozen at its checkpoint value. - raise ValueError( - f"{collector_path} could not map {len(unmapped_keys)} adapter modules to " - f"rollout layer names (first 5: {unmapped_keys[:5]})" - ) - return layer_groups, unmapped_keys, num_lora_keys + """Group PEFT LoRA tensors so each layer's A/B pair stays in one IPC bucket. + + Names stay PEFT/diffusers-shaped (``transformer_blocks.0.attn.to_q.lora_A``). + sglang-d's ``lora_merge`` path applies ``param_names_mapping`` and the + disk-load FFN swap, so fused families such as H3 do not need a trainer-side + collector. + """ + return collect_lora_layer_groups(model.state_dict()) def update_weights(self) -> None: self.weight_version += 1 diff --git a/miles/backends/fsdp_utils/h3_weight_key_mapper.py b/miles/backends/fsdp_utils/h3_weight_key_mapper.py deleted file mode 100644 index 764e0af76..000000000 --- a/miles/backends/fsdp_utils/h3_weight_key_mapper.py +++ /dev/null @@ -1,182 +0,0 @@ -"""Map diffusers MiniMax H3 LoRA names to sglang H3 DiT layer names. - -Training uses diffusers ``MiniMaxH3Transformer3DModel`` (separate Q/K/V). -Rollout uses sglang ``MiniMaxH3DiTModel`` (fused ``qkv_proj``). LoRA IPC sync -must therefore rename modules and stack the Q/K/V adapters before the push. - -Names that resolve to no sglang layer are skipped with a warning on the rollout -side, so an incomplete map silently freezes those adapters at their checkpoint -values — anything unrecognized is reported as unmapped instead of guessed. -""" - -from __future__ import annotations - -import re -from collections.abc import Mapping - -import torch - -_QKV_RE = re.compile( - r"^(?P(?:token_refiner\.)?refiner_blocks\.(?P\d+)|transformer_blocks\.(?P\d+))" - r"\.attn\.to_(?Pq|k|v)\.weight$" -) -# After PEFT strip, refiner path is token_refiner.refiner_blocks -> token_refiner.blocks -_QKV_RE_SGL = re.compile( - r"^(?P(?:token_refiner\.)?blocks\.(?P\d+)|blocks\.(?P\d+))" - r"\.attn\.to_(?Pq|k|v)\.weight$" -) - - -def _qkv_group_key(name: str) -> tuple[str, str] | None: - for regex in (_QKV_RE, _QKV_RE_SGL): - m = regex.match(name) - if m is None: - continue - prefix = m.group("prefix") - if prefix.startswith("token_refiner."): - block_idx = m.group("idx") - sgld_prefix = f"token_refiner.blocks.{block_idx}" - elif prefix.startswith("refiner_blocks."): - block_idx = m.group("idx") - sgld_prefix = f"token_refiner.blocks.{block_idx}" - else: - block_idx = m.group("idx2") - sgld_prefix = f"blocks.{block_idx}" - return sgld_prefix, m.group("which") - return None - - -def _swap_gated_ffn_halves(tensor: torch.Tensor) -> torch.Tensor: - """Reorder a gated FFN input projection from diffusers' halves to sglang's. - - diffusers' GEGLU splits the fused projection as ``[up, gate]`` and computes - ``up * gelu(gate)``; sglang's ``mlp.fc1`` splits it as ``[gate, up]`` and - computes ``silu(gate) * up``. Same weights, opposite halves. - """ - rows = tensor.shape[0] - if rows % 2: - raise ValueError(f"H3 gated FFN projection must have an even row count, got {rows}") - half = rows // 2 - return torch.cat([tensor[half:], tensor[:half]], dim=0) - - -_LORA_AB_RE = re.compile(r"\.lora_([AB])(?:\.[^.]+)?(?:\.weight)?$") -_PEFT_PREFIX = "base_model.model." - -# LoRA-able H3 submodules other than Q/K/V, as (diffusers suffix, sglang suffix). -# Kept as an explicit whitelist: an unrecognized module must surface as unmapped -# rather than reach the rollout under a guessed name, where it would be skipped -# with only a warning and silently freeze that adapter. -_LORA_MODULE_SUFFIXES: tuple[tuple[str, str], ...] = ( - (".attn.to_out.0", ".attn.out_proj"), - (".ff.net.0.proj", ".mlp.fc1"), - (".ff.net.2", ".mlp.fc2"), -) - -_BLOCK_PREFIX_REPLACEMENTS: tuple[tuple[str, str], ...] = ( - (r"^token_refiner\.refiner_blocks\.", "token_refiner.blocks."), - (r"^refiner_blocks\.", "token_refiner.blocks."), - (r"^transformer_blocks\.", "blocks."), -) - - -def _strip_peft_prefix(name: str) -> str: - return name[len(_PEFT_PREFIX) :] if name.startswith(_PEFT_PREFIX) else name - - -def _normalize_block_prefix(module_path: str) -> str: - out = module_path - for pattern, repl in _BLOCK_PREFIX_REPLACEMENTS: - out = re.sub(pattern, repl, out) - return out - - -def _stack_qkv_lora(triple: dict[str, dict[str, torch.Tensor]], layer: str) -> tuple[torch.Tensor, torch.Tensor]: - """Stack per-projection LoRA into the 3D layout sglang's fused qkv expects. - - ``MergedColumnParallelLinearWithLoRA`` multiplies a 3D ``B @ A`` batchwise and - flattens the result, so stacking along a leading axis yields a delta ordered - ``[q_all, k_all, v_all]`` — exactly how sglang stores ``qkv_proj.weight``. - Note this differs from the dense path, which must instead emit the head-major - grouped layout because it goes through the checkpoint weight loader. - """ - missing = {"q", "k", "v"} - set(triple) - if missing: - raise ValueError(f"H3 LoRA IPC incomplete QKV for {layer}: missing {sorted(missing)}") - order = ("q", "k", "v") - a_shapes = {triple[w]["A"].shape for w in order} - b_shapes = {triple[w]["B"].shape for w in order} - if len(a_shapes) != 1 or len(b_shapes) != 1: - raise ValueError( - f"H3 LoRA IPC expects MHA-shaped Q/K/V adapters for {layer}, " - f"got A={sorted(a_shapes)} B={sorted(b_shapes)}" - ) - lora_a = torch.stack([triple[w]["A"] for w in order], dim=0) - lora_b = torch.stack([triple[w]["B"] for w in order], dim=0) - return lora_a, lora_b - - -def collect_h3_lora_layer_groups( - state_dict: Mapping[str, torch.Tensor], -) -> tuple[list[list[tuple[str, torch.Tensor]]], list[str], int]: - """Group PEFT LoRA tensors into sglang H3 layer names for IPC weight sync. - - Returns ``(layer_groups, unmapped_keys, num_lora_keys)`` to match - ``collect_lora_layer_groups``; each group holds one layer's A/B pair so they - always land in the same IPC bucket. - """ - per_module: dict[str, dict[str, torch.Tensor]] = {} - unmapped: list[str] = [] - num_lora_keys = 0 - - for name, tensor in state_dict.items(): - if ".lora_A" not in name and ".lora_B" not in name: - continue - stripped = _strip_peft_prefix(name) - match = _LORA_AB_RE.search(stripped) - if match is None: - unmapped.append(name) - continue - per_module.setdefault(stripped[: match.start()], {})[match.group(1)] = tensor - num_lora_keys += 1 - - qkv_pending: dict[str, dict[str, dict[str, torch.Tensor]]] = {} - simple: dict[str, dict[str, torch.Tensor]] = {} - - for module_path, ab in per_module.items(): - if "A" not in ab or "B" not in ab: - unmapped.append(module_path) - continue - - probe = f"{module_path}.weight" - qkv = _qkv_group_key(probe) - if qkv is not None: - sgld_prefix, which = qkv - qkv_pending.setdefault(f"{sgld_prefix}.attn.qkv_proj", {})[which] = ab - continue - - normalized = _normalize_block_prefix(module_path) - for diffusers_suffix, sglang_suffix in _LORA_MODULE_SUFFIXES: - if normalized.endswith(diffusers_suffix): - layer = normalized[: -len(diffusers_suffix)] + sglang_suffix - if diffusers_suffix == ".ff.net.0.proj": - # Gated FFN: diffusers stores [up, gate], sglang fc1 wants - # [gate, up]. Only B is row-indexed by output, so A is untouched. - ab = {"A": ab["A"], "B": _swap_gated_ffn_halves(ab["B"])} - simple[layer] = ab - break - else: - unmapped.append(module_path) - - groups: list[list[tuple[str, torch.Tensor]]] = [] - for layer in sorted(simple): - ab = simple[layer] - groups.append([(f"{layer}.lora_A", ab["A"]), (f"{layer}.lora_B", ab["B"])]) - for layer in sorted(qkv_pending): - lora_a, lora_b = _stack_qkv_lora(qkv_pending[layer], layer) - groups.append([(f"{layer}.lora_A", lora_a), (f"{layer}.lora_B", lora_b)]) - - return groups, unmapped, num_lora_keys - - -__all__ = ["collect_h3_lora_layer_groups"] From 4036bcadca0cf96651cdbd7b516ff168058c3eef Mon Sep 17 00:00:00 2001 From: niehen6174 <1639206518@qq.com> Date: Tue, 1 Sep 2026 03:22:35 +0000 Subject: [PATCH 2/3] test(h3): assert IPC LoRA keys stay PEFT/diffusers-shaped --- .../fsdp_utils/test_peft_lora_key_mapper.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/fast/backends/fsdp_utils/test_peft_lora_key_mapper.py b/tests/fast/backends/fsdp_utils/test_peft_lora_key_mapper.py index d74aac0fa..384a52c82 100644 --- a/tests/fast/backends/fsdp_utils/test_peft_lora_key_mapper.py +++ b/tests/fast/backends/fsdp_utils/test_peft_lora_key_mapper.py @@ -66,3 +66,22 @@ def test_summarize_mapping_reports_unmapped_lora_keys(self): assert num_layers == 1 assert sample_layers == ["transformer_blocks.0.attn.to_q"] assert unmapped == ["base_model.model.weird.lora_A.default.weight.extra"] + + def test_h3_peft_keys_stay_diffusers_shaped(self): + """H3 IPC no longer pre-fuses names; sglang-d maps transformer_blocks → blocks.""" + state_dict = { + "base_model.model.transformer_blocks.0.attn.to_q.lora_A.default.weight": torch.zeros(4, 8), + "base_model.model.transformer_blocks.0.attn.to_q.lora_B.default.weight": torch.zeros(8, 4), + "base_model.model.transformer_blocks.0.ff.net.0.proj.lora_A.default.weight": torch.zeros(4, 8), + "base_model.model.transformer_blocks.0.ff.net.0.proj.lora_B.default.weight": torch.zeros(8, 4), + "base_model.model.token_refiner.refiner_blocks.1.attn.to_k.lora_A.default.weight": torch.zeros(4, 8), + "base_model.model.token_refiner.refiner_blocks.1.attn.to_k.lora_B.default.weight": torch.zeros(8, 4), + } + assert PeftLoRAKeyMapper.collect_sgld_names(state_dict) == { + "transformer_blocks.0.attn.to_q.lora_A", + "transformer_blocks.0.attn.to_q.lora_B", + "transformer_blocks.0.ff.net.0.proj.lora_A", + "transformer_blocks.0.ff.net.0.proj.lora_B", + "token_refiner.refiner_blocks.1.attn.to_k.lora_A", + "token_refiner.refiner_blocks.1.attn.to_k.lora_B", + } From 618c7b09826d9e8f176e649679a2f18419311575 Mon Sep 17 00:00:00 2001 From: niehen6174 <1639206518@qq.com> Date: Tue, 1 Sep 2026 03:22:36 +0000 Subject: [PATCH 3/3] docs(h3): note that sglang-d maps LoRA names on IPC --- docs/advanced/lora.md | 6 ++++-- docs/models/h3/h3.md | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/advanced/lora.md b/docs/advanced/lora.md index 6c74397c8..a90a8d939 100644 --- a/docs/advanced/lora.md +++ b/docs/advanced/lora.md @@ -83,8 +83,10 @@ sglang-d expects (e.g. `transformer_blocks.0.attn.to_q.weight`). 1. `collect_lora_layer_groups()` groups state-dict entries by layer prefix so **lora_A and lora_B for the same layer always stay together**. -2. `PeftLoRAKeyMapper.to_sgld_name()` maps PEFT keys to sglang-d names - (e.g. `transformer_blocks.0.attn.to_q.lora_A`). +2. `PeftLoRAKeyMapper.to_sgld_name()` strips PEFT wrappers + (e.g. `transformer_blocks.0.attn.to_q.lora_A`). Fused families such as H3 + keep these diffusers names; sglang-d `lora_merge` applies + `param_names_mapping` and the disk-load FFN swap. 3. FSDP shard all-gather → pack into buckets capped by **`--update-weight-buffer-size`** (recipes use 2 GB) → CUDA IPC. 4. Rollout engine receives `weight_update_mode="lora_merge"` with diff --git a/docs/models/h3/h3.md b/docs/models/h3/h3.md index bac9b6395..5f89683b3 100644 --- a/docs/models/h3/h3.md +++ b/docs/models/h3/h3.md @@ -25,7 +25,7 @@ From `miles/backends/fsdp_utils/configs/h3.py`: |---|---|---| | CFG training | Off (asserted) | CFG is distilled into the checkpoint; the forward is unguided | | Weight sync | `--use-lora --lora-ipc-weight-sync` (asserted) | sgl-d's H3 DiT renames modules and fuses Q/K/V; other sync modes push names the engine drops with a warning, silently training nothing | -| LoRA targets | `attn.to_{q,k,v}`, `attn.to_out.0`, `ff.net.0.proj`, `ff.net.2` | Grouped for the fused engine layers by `h3_weight_key_mapper.collect_h3_lora_layer_groups` | +| LoRA targets | `attn.to_{q,k,v}`, `attn.to_out.0`, `ff.net.0.proj`, `ff.net.2` | Trainer pushes PEFT/diffusers names; sglang-d `lora_merge` maps and fuses QKV / swaps FFN | | Optimizer state | `audio` allowed missing | The audio branch is rolled out but never trained | | Sample micro-batch | 1 (asserted) | One packed sequence per forward | | Forced sampling params | `task=t2va`, `short_edge=768`, `conditions=[]` | sgl-d accepts only these for H3, so none of them is exposed as an argument |