fix(h3): drop trainer-side LoRA collector; keep PEFT names on IPC

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

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 |

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")
 

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."""

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

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",
+        }