refactor(lora): one plain PEFT adapter directory per component

docs/fsdp_model_roles.md

@@ -36,7 +36,8 @@ Custom prepare/loss hooks can access `ctx.reference_models` and `ctx.teacher_mod
 by component name. The teacher has no built-in consumer, so `--teacher-load` requires
 one of these hooks. Teacher inference should run inside `torch.no_grad()`.
 
-Each source may load a pretrained PEFT LoRA adapter:
+Each source may load pretrained PEFT LoRA adapters, one plain PEFT directory (what
+`save_pretrained` writes) per `--update-weight-target-module` entry, in the same order:
 
 ```bash
 --use-lora --lora-adapter-path /adapters/student --lora-rank 64 --lora-alpha 128 \
@@ -46,11 +47,16 @@ Each source may load a pretrained PEFT LoRA adapter:
 --teacher-lora-adapter-path /adapters/teacher
 ```
 
+Wan2.2, for example, takes two directories, one per expert; any other number is rejected:
+
+```bash
+--update-weight-target-module transformer,transformer_2 \
+--use-lora --lora-adapter-path /adapters/high_noise /adapters/low_noise
+```
+
 Without a role's adapter path, that frozen role loads ordinary full weights, even
 when the actor uses LoRA. Full weights and base-only weights use the same loader.
-An adapter source is a PEFT checkpoint; for multiple components it contains a PEFT
-checkpoint under each component subdirectory. `--lora-rank` and `--lora-alpha` must
-match the actor adapter's `r` and `lora_alpha`. Actor adapter weights are loaded
+`--lora-rank` and `--lora-alpha` must match the actor adapter's `r` and `lora_alpha`. Actor adapter weights are loaded
 before a training checkpoint is restored. Independent reference/teacher weights
 are always reconstructed from their configured sources, so retain those sources
 and their adapter checkpoints when resuming training.

docs/user-guide/cli-reference.md

@@ -263,10 +263,10 @@ See [Dtype Control](../advanced/dtype-control.md).
 | `--diffusion-kl-beta` | float | `0.0` | KL to the `--ref-mode` model; under `nft` adds `kl_beta * mean((v_θ − v_ref)²)` per pair. |
 | `--ref-mode` | enum | – | KL reference: `none` / `lora_base` (actor with adapters disabled) / `ref` (the model from `--ref-load`). Auto: `lora_base` when KL > 0. |
 | `--ref-load` | str | – | HF pipeline checkpoint of an independent frozen reference model. Requires `--ref-mode ref`. |
-| `--ref-lora-adapter-path` | str | – | PEFT adapter applied to `--ref-load`. |
+| `--ref-lora-adapter-path` | str+ | – | PEFT adapter directories applied to `--ref-load`, one per `--update-weight-target-module` entry, in order. |
 | `--ref-cpu-offload` | flag | off | Keep the reference shards in pinned CPU memory between forwards (FSDP `CPUOffloadPolicy`). |
 | `--teacher-load` | str | – | HF pipeline checkpoint of a frozen teacher model, read only by custom prepare/loss hooks as `ctx.teacher_models`, so it requires `--custom-prepare-train-batch-path` or `--custom-loss-function-path`. |
-| `--teacher-lora-adapter-path` | str | – | PEFT adapter applied to `--teacher-load`. |
+| `--teacher-lora-adapter-path` | str+ | – | PEFT adapter directories applied to `--teacher-load`, one per `--update-weight-target-module` entry, in order. |
 | `--teacher-cpu-offload` | flag | off | Keep the teacher shards in pinned CPU memory between forwards. |
 | `--custom-prepare-train-batch-path` | str | – | Builds DiT inputs. |
 | `--custom-loss-function-path` | str | – | Loss **formula** only — the DiT forward stays in the actor. |
@@ -325,7 +325,7 @@ Every one takes a dotted path.
 |---|---|---|---|
 | `--use-lora` | flag | off | |
 | `--lora-rank` / `--lora-alpha` | int | `64` / `64` | |
-| `--lora-adapter-path` | str | – | Resume the actor from a saved PEFT adapter; `--lora-rank` / `--lora-alpha` must match its config. |
+| `--lora-adapter-path` | str+ | – | PEFT adapter directories the actor continues from, one per `--update-weight-target-module` entry, in order; `--lora-rank` / `--lora-alpha` must match their configs. |
 | `--lora-target-modules` | str+ | – | Defaults per model family. |
 | `--lora-init-weights` | str | `gaussian` | `kaiming-uniform` maps to PEFT's default; other PEFT schemes pass through. |
 | `--lora-ipc-weight-sync` | flag | off | Push only `lora_A`/`lora_B`; the engine merges locally. Requires `--use-lora`. |

miles/backends/fsdp_utils/actor.py

@@ -108,7 +108,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
             self.train_pipeline_config,
             self.parallel_state,
             checkpoint_path=args.hf_checkpoint,
-            lora_adapter_path=args.lora_adapter_path,
+            lora_adapter_paths=args.lora_adapter_paths,
             trainable=True,
             cpu_offload=args.fsdp_cpu_offload,
         )
@@ -120,7 +120,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
                 self.train_pipeline_config,
                 self.parallel_state,
                 checkpoint_path=args.ref_load,
-                lora_adapter_path=args.ref_lora_adapter_path,
+                lora_adapter_paths=args.ref_lora_adapter_paths,
                 cpu_offload=args.ref_cpu_offload,
             )
         self.teacher_models = {}
@@ -131,7 +131,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
                 self.train_pipeline_config,
                 self.parallel_state,
                 checkpoint_path=args.teacher_load,
-                lora_adapter_path=args.teacher_lora_adapter_path,
+                lora_adapter_paths=args.teacher_lora_adapter_paths,
                 cpu_offload=args.teacher_cpu_offload,
             )
         # A natively CPU-offloaded model never sleeps: its parameters stay on the host and its buffers on the GPU.

miles/backends/fsdp_utils/model_loader.py

@@ -23,14 +23,14 @@ def load_fsdp_models(
     parallel_state,
     *,
     checkpoint_path: str,
-    lora_adapter_path: str | None = None,
+    lora_adapter_paths: list[str] | None = None,
     trainable: bool = False,
     cpu_offload: bool = False,
 ) -> dict[str, torch.nn.Module]:
     materialize_weights = dist.get_rank() == 0
     master_dtype = parse_dtype_from_str(args.fsdp_master_dtype)
     models = {}
-    for component in args.update_weight_target_modules:
+    for component_index, component in enumerate(args.update_weight_target_modules):
         with model_init_context(materialize_weights=materialize_weights):
             model = model_backend.load_component(
                 component,
@@ -42,10 +42,8 @@ def load_fsdp_models(
             model_backend.set_attention_backend(model, args.fsdp_attention_backend)
         if trainable and args.gradient_checkpointing:
             model_backend.enable_gradient_checkpointing(model)
-        if lora_adapter_path is not None:
-            model = load_lora_adapter(
-                model, lora_adapter_path, trainable=trainable, subfolder=lora_adapter_subfolder(args, component)
-            )
+        if lora_adapter_paths is not None:
+            model = load_lora_adapter(model, lora_adapter_paths[component_index], trainable=trainable)
         elif trainable and args.use_lora:
             model = apply_lora(model, args, train_pipeline_config)
         model.train(trainable)
@@ -107,19 +105,12 @@ def model_init_context(*, materialize_weights: bool):
         yield
 
 
-def lora_adapter_subfolder(args: Namespace, component: str) -> str | None:
-    """A multi-component adapter directory keeps each component's adapter in a subfolder named after it."""
-    return component if len(args.update_weight_target_modules) > 1 else None
-
-
-def load_lora_adapter(model, adapter_path: str, *, trainable: bool, subfolder: str | None = None):
+def load_lora_adapter(model, adapter_path: str, *, trainable: bool):
     from peft import PeftConfig, PeftModel, get_peft_model
 
     if dist.get_rank() == 0:
-        return PeftModel.from_pretrained(
-            model, adapter_path, subfolder=subfolder, is_trainable=trainable, autocast_adapter_dtype=False
-        )
-    config = PeftConfig.from_pretrained(adapter_path, subfolder=subfolder)
+        return PeftModel.from_pretrained(model, adapter_path, is_trainable=trainable, autocast_adapter_dtype=False)
+    config = PeftConfig.from_pretrained(adapter_path)
     config.inference_mode = not trainable
     return get_peft_model(model, config, low_cpu_mem_usage=True, autocast_adapter_dtype=False)
 

miles/utils/arguments.py

@@ -980,13 +980,21 @@ def add_algo_arguments(parser):
                 "--teacher-load", type=str, default=None, help="HF pipeline checkpoint for the frozen teacher model."
             )
             parser.add_argument(
-                "--ref-lora-adapter-path", type=str, default=None, help="Pretrained LoRA adapter for --ref-load."
+                "--ref-lora-adapter-path",
+                dest="ref_lora_adapter_paths",
+                nargs="+",
+                default=None,
+                help="PEFT adapter directories for --ref-load, one per --update-weight-target-module entry, in order.",
             )
             parser.add_argument(
                 "--teacher-lora-adapter-path",
-                type=str,
+                dest="teacher_lora_adapter_paths",
+                nargs="+",
                 default=None,
-                help="Pretrained LoRA adapter for --teacher-load.",
+                help=(
+                    "PEFT adapter directories for --teacher-load, "
+                    "one per --update-weight-target-module entry, in order."
+                ),
             )
             parser.add_argument(
                 "--ref-cpu-offload",
@@ -1080,9 +1088,13 @@ def add_lora_arguments(parser):
             )
             parser.add_argument(
                 "--lora-adapter-path",
-                type=str,
+                dest="lora_adapter_paths",
+                nargs="+",
                 default=None,
-                help="Pretrained actor LoRA adapter to continue training.",
+                help=(
+                    "PEFT adapter directories the actor continues training from, "
+                    "one per --update-weight-target-module entry, in order."
+                ),
             )
             parser.add_argument("--lora-rank", type=int, default=64)
             parser.add_argument("--lora-alpha", type=int, default=64)
@@ -1605,10 +1617,19 @@ def validate_reference_model_args(args) -> None:
         raise ValueError("--ref-mode ref requires --ref-load")
     if args.ref_load is not None and args.ref_mode != "ref":
         raise ValueError("--ref-load requires --ref-mode ref")
-    if args.ref_lora_adapter_path is not None and args.ref_load is None:
+    if args.ref_lora_adapter_paths is not None and args.ref_load is None:
         raise ValueError("--ref-lora-adapter-path requires --ref-load")
-    if args.teacher_lora_adapter_path is not None and args.teacher_load is None:
+    if args.teacher_lora_adapter_paths is not None and args.teacher_load is None:
         raise ValueError("--teacher-lora-adapter-path requires --teacher-load")
+    for flag, adapter_paths in (
+        ("--ref-lora-adapter-path", args.ref_lora_adapter_paths),
+        ("--teacher-lora-adapter-path", args.teacher_lora_adapter_paths),
+    ):
+        if adapter_paths is not None and len(adapter_paths) != len(args.update_weight_target_modules):
+            raise ValueError(
+                f"{flag} needs one adapter per --update-weight-target-module entry "
+                f"{args.update_weight_target_modules}, got {len(adapter_paths)}"
+            )
     if args.ref_cpu_offload and args.ref_load is None:
         raise ValueError("--ref-cpu-offload requires --ref-load")
     if args.teacher_cpu_offload and args.teacher_load is None:
@@ -1624,18 +1645,19 @@ def validate_reference_model_args(args) -> None:
 
 
 def validate_actor_lora_adapter(args) -> None:
-    if args.lora_adapter_path is None:
+    if args.lora_adapter_paths is None:
         return
     if not args.use_lora:
         raise ValueError("--lora-adapter-path requires --use-lora")
+    if len(args.lora_adapter_paths) != len(args.update_weight_target_modules):
+        raise ValueError(
+            "--lora-adapter-path needs one adapter per --update-weight-target-module entry "
+            f"{args.update_weight_target_modules}, got {len(args.lora_adapter_paths)}"
+        )
     from peft import PeftConfig
 
-    from miles.backends.fsdp_utils.model_loader import lora_adapter_subfolder
-
-    for component in args.update_weight_target_modules:
-        adapter_config = PeftConfig.from_pretrained(
-            args.lora_adapter_path, subfolder=lora_adapter_subfolder(args, component)
-        )
+    for component, adapter_path in zip(args.update_weight_target_modules, args.lora_adapter_paths, strict=True):
+        adapter_config = PeftConfig.from_pretrained(adapter_path)
         if (
             adapter_config.use_dora
             or adapter_config.bias != "none"
@@ -1662,7 +1684,6 @@ def validate_actor_lora_adapter(args) -> None:
 
 
 def miles_validate_args(args):
-    validate_reference_model_args(args)
     args.eval_datasets = _resolve_eval_datasets(args)
 
     if args.eval_interval is not None:
@@ -1699,6 +1720,7 @@ def miles_validate_args(args):
         )
     if len(set(args.update_weight_target_modules)) != len(args.update_weight_target_modules):
         raise ValueError(f"--update-weight-target-module has duplicates: {args.update_weight_target_module!r}")
+    validate_reference_model_args(args)
     validate_actor_lora_adapter(args)
 
     if args.wandb_log_image_interval < 1:

tests/fast-gpu/backends/fsdp_utils/_reference_models_worker.py

@@ -4,7 +4,7 @@
        |                         |                          |
        +--> dense oracle         +--> full / LoRA ref       +--> LoRA actor
                                         frozen                  gradients + AdamW
-    role checkpoint_path --> base component --> optional pretrained adapter
+    role checkpoint_path --> base component --> optional pretrained adapter (one PEFT directory per component)
 
 Native CPU offload and GPU-resident references must produce the same outputs.
 LoRA-base evaluation must leave adapter execution and adapter gradients intact.
@@ -22,6 +22,7 @@
 from tests.fast.backends.fsdp_utils.test_reference_models import (
     TinyBackend,
     TinyPipelineConfig,
+    component_adapter_paths,
     load_actor_forward_method,
     make_actor_forward_harness,
     make_model_loader_args,
@@ -38,7 +39,7 @@ def load_dense_model(checkpoint_path, component, adapter_path, *, trainable=Fals
         component, checkpoint_path=checkpoint_path, master_dtype=torch.float32, materialize_weights=True
     )
     if adapter_path is not None:
-        model = PeftModel.from_pretrained(model, adapter_path, subfolder=component, is_trainable=trainable)
+        model = PeftModel.from_pretrained(model, adapter_path, is_trainable=trainable)
     model.train(trainable)
     return model.cuda()
 
@@ -47,19 +48,19 @@ def local_tensor(tensor):
     return tensor.to_local() if isinstance(tensor, DTensor) else tensor
 
 
-def check_roles(mesh, adapter_path, reference_has_adapter, cpu_offload):
+def check_roles(mesh, adapter_paths, reference_has_adapter, cpu_offload):
     args, backend, config = make_model_loader_args(), TinyBackend(), TinyPipelineConfig()
     parallel = SimpleNamespace(get_mesh=lambda name: mesh, get_optional_mesh=lambda name: None)
     actor = load_fsdp_models(
-        args, backend, config, parallel, checkpoint_path="actor", lora_adapter_path=adapter_path, trainable=True
+        args, backend, config, parallel, checkpoint_path="actor", lora_adapter_paths=adapter_paths, trainable=True
     )
     reference = load_fsdp_models(
         args,
         backend,
         config,
         parallel,
         checkpoint_path="reference",
-        lora_adapter_path=adapter_path if reference_has_adapter else None,
+        lora_adapter_paths=adapter_paths if reference_has_adapter else None,
         cpu_offload=cpu_offload,
     )
     teacher = load_fsdp_models(
@@ -68,11 +69,11 @@ def check_roles(mesh, adapter_path, reference_has_adapter, cpu_offload):
         config,
         parallel,
         checkpoint_path="teacher",
-        lora_adapter_path=adapter_path,
+        lora_adapter_paths=adapter_paths,
         cpu_offload=cpu_offload,
     )
     inputs = torch.arange(6, dtype=torch.float32, device="cuda").reshape(2, 3) / 10
-    for component in args.update_weight_target_modules:
+    for component, adapter_path in zip(args.update_weight_target_modules, adapter_paths, strict=True):
         oracle = load_dense_model("actor", component, adapter_path, trainable=True)
         reference_oracle = load_dense_model("reference", component, adapter_path if reference_has_adapter else None)
         teacher_oracle = load_dense_model("teacher", component, adapter_path)
@@ -135,12 +136,12 @@ def main():
     torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
     dist.init_process_group("nccl", device_id=torch.device("cuda", torch.cuda.current_device()))
     mesh = init_device_mesh("cuda", (dist.get_world_size(),))
-    adapter_path = Path(sys.argv[1])
+    adapter_root = Path(sys.argv[1])
     if dist.get_rank() == 0:
-        save_component_lora_adapters(adapter_path)
+        save_component_lora_adapters(adapter_root)
     dist.barrier()
     for reference_has_adapter, cpu_offload in itertools.product((False, True), repeat=2):
-        check_roles(mesh, str(adapter_path), reference_has_adapter, cpu_offload)
+        check_roles(mesh, component_adapter_paths(adapter_root), reference_has_adapter, cpu_offload)
     dist.barrier()
     if dist.get_rank() == 0:
         print("OK")

tests/fast-gpu/backends/fsdp_utils/test_reference_models.py

@@ -35,7 +35,13 @@ def test_independent_fsdp_references(tmp_path):
             str(tmp_path),
         ],
         # The worker imports the CPU test's harness from the `tests` package at the repository root.
-        env={**os.environ, "PYTHONPATH": str(Path(__file__).resolve().parents[4]), "PYTHONUNBUFFERED": "1"},
+        env={
+            **os.environ,
+            "PYTHONPATH": os.pathsep.join(
+                filter(None, (str(Path(__file__).resolve().parents[4]), os.environ.get("PYTHONPATH")))
+            ),
+            "PYTHONUNBUFFERED": "1",
+        },
         capture_output=True,
         text=True,
         timeout=300,