From fb3717bfff3d7f0270e2e877a6767eb0031068b2 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 29 Sep 2026 01:13:07 -0700 Subject: [PATCH 1/2] refactor(lora): take one plain PEFT adapter directory per component instead of component subfolders --- docs/fsdp_model_roles.md | 14 +++-- docs/user-guide/cli-reference.md | 6 +-- miles/backends/fsdp_utils/actor.py | 6 +-- miles/backends/fsdp_utils/model_loader.py | 23 +++----- miles/utils/arguments.py | 52 +++++++++++++------ .../fsdp_utils/_reference_models_worker.py | 21 ++++---- .../fsdp_utils/test_reference_models.py | 31 ++++++----- tests/fast/utils/test_reference_arguments.py | 48 +++++++++++------ 8 files changed, 122 insertions(+), 79 deletions(-) diff --git a/docs/fsdp_model_roles.md b/docs/fsdp_model_roles.md index 02da0fbbd..ece375088 100644 --- a/docs/fsdp_model_roles.md +++ b/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. diff --git a/docs/user-guide/cli-reference.md b/docs/user-guide/cli-reference.md index 5f0713e1d..1576d44d7 100644 --- a/docs/user-guide/cli-reference.md +++ b/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`. | diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index deb386272..a8d21e8f7 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/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. diff --git a/miles/backends/fsdp_utils/model_loader.py b/miles/backends/fsdp_utils/model_loader.py index b3e17abf9..ca61e39ec 100644 --- a/miles/backends/fsdp_utils/model_loader.py +++ b/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) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index cd67711f3..a02d7b144 100644 --- a/miles/utils/arguments.py +++ b/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: diff --git a/tests/fast-gpu/backends/fsdp_utils/_reference_models_worker.py b/tests/fast-gpu/backends/fsdp_utils/_reference_models_worker.py index 7c4166a43..f84fe1a57 100644 --- a/tests/fast-gpu/backends/fsdp_utils/_reference_models_worker.py +++ b/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,11 +48,11 @@ 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, @@ -59,7 +60,7 @@ def check_roles(mesh, adapter_path, reference_has_adapter, cpu_offload): 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") diff --git a/tests/fast/backends/fsdp_utils/test_reference_models.py b/tests/fast/backends/fsdp_utils/test_reference_models.py index d5eff247a..993a905a1 100644 --- a/tests/fast/backends/fsdp_utils/test_reference_models.py +++ b/tests/fast/backends/fsdp_utils/test_reference_models.py @@ -105,13 +105,20 @@ def make_model_loader_args(): ) +def component_adapter_paths(path): + """One plain PEFT adapter directory per component, in --update-weight-target-module order.""" + return [str(path / component) for component in make_model_loader_args().update_weight_target_modules] + + def save_component_lora_adapters(path): - for index, component in enumerate(make_model_loader_args().update_weight_target_modules): + adapter_paths = component_adapter_paths(path) + for index, adapter_path in enumerate(adapter_paths): model = get_peft_model(TinyComponent(), LoraConfig(r=2, lora_alpha=4, target_modules=["proj"])) with torch.no_grad(): model.base_model.model.block.proj.lora_A.default.weight.fill_(0.15 + 0.1 * index) model.base_model.model.block.proj.lora_B.default.weight.fill_(0.25 + 0.1 * index) - model.save_pretrained(path / component) + model.save_pretrained(adapter_path) + return adapter_paths def load_actor_forward_method(): @@ -191,7 +198,7 @@ def test_loaded_reference_and_teacher_are_independent_and_frozen( tmp_path, unsharded_model_loader, actor_mode, reference_has_adapter ): parallel, calls = unsharded_model_loader - save_component_lora_adapters(tmp_path) + adapter_paths = save_component_lora_adapters(tmp_path) args, backend, config = make_model_loader_args(), TinyBackend(), TinyPipelineConfig() actor_has_adapter = actor_mode != "full" args.use_lora = actor_has_adapter @@ -205,7 +212,7 @@ def test_loaded_reference_and_teacher_are_independent_and_frozen( config, parallel, checkpoint_path="actor", - lora_adapter_path=str(tmp_path) if actor_mode == "saved_lora" else None, + lora_adapter_paths=adapter_paths if actor_mode == "saved_lora" else None, trainable=True, ) reference = model_loader.load_fsdp_models( @@ -214,7 +221,7 @@ def test_loaded_reference_and_teacher_are_independent_and_frozen( config, parallel, checkpoint_path="reference", - lora_adapter_path=str(tmp_path) if reference_has_adapter else None, + lora_adapter_paths=adapter_paths if reference_has_adapter else None, cpu_offload=True, ) teacher = model_loader.load_fsdp_models( @@ -223,7 +230,7 @@ def test_loaded_reference_and_teacher_are_independent_and_frozen( config, parallel, checkpoint_path="teacher", - lora_adapter_path=str(tmp_path), + lora_adapter_paths=adapter_paths, ) assert args.hf_checkpoint == "actor" and args.use_lora is actor_has_adapter assert calls == [("shared-fsdp-mesh", flag) for flag in (False, False, True, True, False, False)] @@ -273,14 +280,14 @@ def test_loaded_reference_and_teacher_are_independent_and_frozen( def test_lora_base_forward_restores_adapters_and_actor_gradients(tmp_path, unsharded_model_loader): parallel, _ = unsharded_model_loader - save_component_lora_adapters(tmp_path) + adapter_paths = save_component_lora_adapters(tmp_path) actor = model_loader.load_fsdp_models( make_model_loader_args(), TinyBackend(), TinyPipelineConfig(), parallel, checkpoint_path="actor", - lora_adapter_path=str(tmp_path), + lora_adapter_paths=adapter_paths, trainable=True, ) component = "transformer_2" @@ -357,14 +364,14 @@ def test_nft_gets_ema_as_old_and_lora_base_as_kl_reference(tmp_path, unsharded_m # actor adapters moved away from their EMA copy: # old_pred = actor with EMA weights (pi_old), ref_pred = actor with adapters disabled, all three differ parallel, _ = unsharded_model_loader - save_component_lora_adapters(tmp_path) + adapter_paths = save_component_lora_adapters(tmp_path) actor = model_loader.load_fsdp_models( make_model_loader_args(), TinyBackend(), TinyPipelineConfig(), parallel, checkpoint_path="actor", - lora_adapter_path=str(tmp_path), + lora_adapter_paths=adapter_paths, trainable=True, ) component = "transformer" @@ -391,14 +398,14 @@ def test_nft_gets_ema_as_old_and_lora_base_as_kl_reference(tmp_path, unsharded_m def test_ipc_publication_uses_loaded_adapter_config(tmp_path, unsharded_model_loader): parallel, _ = unsharded_model_loader - save_component_lora_adapters(tmp_path) + adapter_paths = save_component_lora_adapters(tmp_path) models = model_loader.load_fsdp_models( make_model_loader_args(), TinyBackend(), TinyPipelineConfig(), parallel, checkpoint_path="actor", - lora_adapter_path=str(tmp_path), + lora_adapter_paths=adapter_paths, trainable=True, ) path = Path(__file__).resolve().parents[4] / "miles/backends/fsdp_utils/diffusion_update_weight_utils.py" diff --git a/tests/fast/utils/test_reference_arguments.py b/tests/fast/utils/test_reference_arguments.py index 5b8afff48..7246253c8 100644 --- a/tests/fast/utils/test_reference_arguments.py +++ b/tests/fast/utils/test_reference_arguments.py @@ -4,6 +4,7 @@ actor: LoRA ---> lora_base: reuse actor with its adapter disabled teacher ---> read only by custom prepare/loss hooks, so it needs one actor LoRA enabled --> adapter config --> checkpoint A/B support + uniform IPC scaling + CLI r/alpha match + adapter paths ---> one plain PEFT directory per component, in order; any other count is rejected Validation rejects incomplete role sources and unused options without changing args. """ @@ -25,13 +26,13 @@ def _args(**overrides): ref_mode="none", ref_load=None, teacher_load=None, - lora_adapter_path=None, + lora_adapter_paths=None, lora_ipc_weight_sync=False, lora_rank=2, lora_alpha=4, update_weight_target_modules=["transformer", "transformer_2"], - ref_lora_adapter_path=None, - teacher_lora_adapter_path=None, + ref_lora_adapter_paths=None, + teacher_lora_adapter_paths=None, ref_cpu_offload=False, teacher_cpu_offload=False, use_lora=False, @@ -53,12 +54,12 @@ def _args(**overrides): def test_frozen_models_do_not_inherit_actor_training_mode(use_lora, role_has_adapter): args = _args( use_lora=use_lora, - lora_adapter_path="actor-adapter" if use_lora else None, + lora_adapter_paths=["actor-high", "actor-low"] if use_lora else None, ref_mode="ref", ref_load="reference-base", teacher_load="teacher-base", - ref_lora_adapter_path="reference-adapter" if role_has_adapter else None, - teacher_lora_adapter_path="teacher-adapter" if role_has_adapter else None, + ref_lora_adapter_paths=["reference-high", "reference-low"] if role_has_adapter else None, + teacher_lora_adapter_paths=["teacher-high", "teacher-low"] if role_has_adapter else None, ref_cpu_offload=True, teacher_cpu_offload=True, custom_loss_function_path="my_loss.teacher_loss", @@ -69,7 +70,7 @@ def test_frozen_models_do_not_inherit_actor_training_mode(use_lora, role_has_ada def test_lora_base_needs_no_separate_reference_checkpoint(): - args = _args(use_lora=True, lora_adapter_path="actor-adapter", ref_mode="lora_base") + args = _args(use_lora=True, lora_adapter_paths=["actor-high", "actor-low"], ref_mode="lora_base") validate_reference_model_args(args) assert args.ref_load is None @@ -79,8 +80,13 @@ def test_lora_base_needs_no_separate_reference_checkpoint(): [ ({"ref_mode": "ref"}, "--ref-mode ref requires --ref-load"), ({"ref_load": "reference"}, "--ref-load requires --ref-mode ref"), - ({"ref_lora_adapter_path": "adapter"}, "--ref-lora-adapter-path requires --ref-load"), - ({"teacher_lora_adapter_path": "adapter"}, "--teacher-lora-adapter-path requires --teacher-load"), + ({"ref_lora_adapter_paths": ["high", "low"]}, "--ref-lora-adapter-path requires --ref-load"), + ({"teacher_lora_adapter_paths": ["high", "low"]}, "--teacher-lora-adapter-path requires --teacher-load"), + # two components, one adapter --> the adapter/component pairing is ambiguous + ( + {"ref_mode": "ref", "ref_load": "reference", "ref_lora_adapter_paths": ["only-one"]}, + "needs one adapter per", + ), ({"ref_cpu_offload": True}, "--ref-cpu-offload requires --ref-load"), ({"teacher_cpu_offload": True}, "--teacher-cpu-offload requires --teacher-load"), ({"ref_mode": "lora_base"}, "--ref-mode lora_base requires --use-lora"), @@ -97,9 +103,10 @@ def test_incomplete_model_roles_are_rejected(overrides, message): def test_loaded_adapter_matching_actor_cli_passes_without_changing_args(tmp_path): # adapter r=2 alpha=4 == --lora-rank 2 --lora-alpha 4 -> accepted, args untouched - args = _args(use_lora=True, lora_adapter_path=str(tmp_path), lora_ipc_weight_sync=True) - for component in args.update_weight_target_modules: - LoraConfig(r=2, lora_alpha=4, target_modules=["proj"]).save_pretrained(tmp_path / component) + adapter_paths = [str(tmp_path / "high"), str(tmp_path / "low")] + for adapter_path in adapter_paths: + LoraConfig(r=2, lora_alpha=4, target_modules=["proj"]).save_pretrained(adapter_path) + args = _args(use_lora=True, lora_adapter_paths=adapter_paths, lora_ipc_weight_sync=True) original = vars(args).copy() validate_actor_lora_adapter(args) assert vars(args) == original @@ -109,13 +116,22 @@ def test_loaded_adapter_matching_actor_cli_passes_without_changing_args(tmp_path def test_loaded_adapter_rejects_mismatched_actor_cli(tmp_path, cli): # adapter r=2 alpha=4 vs CLI 64 -> rejected instead of silently ignoring the CLI values LoraConfig(r=2, lora_alpha=4, target_modules=["proj"]).save_pretrained(tmp_path) - args = _args(use_lora=True, lora_adapter_path=str(tmp_path), update_weight_target_modules=["transformer"], **cli) + args = _args( + use_lora=True, lora_adapter_paths=[str(tmp_path)], update_weight_target_modules=["transformer"], **cli + ) with pytest.raises(ValueError, match="set --lora-rank 2 --lora-alpha 4"): validate_actor_lora_adapter(args) +def test_actor_adapter_count_must_match_the_components(): + # components [transformer, transformer_2], one adapter --> rejected before any config is read + args = _args(use_lora=True, lora_adapter_paths=["unread-adapter-checkpoint"]) + with pytest.raises(ValueError, match="--lora-adapter-path needs one adapter per"): + validate_actor_lora_adapter(args) + + def test_actor_adapter_requires_lora_training_before_loading_config(): - args = _args(lora_adapter_path="unread-adapter-checkpoint") + args = _args(lora_adapter_paths=["unread-adapter-checkpoint"]) with pytest.raises(ValueError, match="--lora-adapter-path requires --use-lora"): validate_actor_lora_adapter(args) @@ -131,7 +147,7 @@ def test_loaded_adapter_rejects_unsupported_export_state(tmp_path, config_option LoraConfig(r=2, lora_alpha=4, target_modules=["proj"], **config_options).save_pretrained(tmp_path) args = _args( use_lora=True, - lora_adapter_path=str(tmp_path), + lora_adapter_paths=[str(tmp_path)], lora_ipc_weight_sync=ipc, update_weight_target_modules=["transformer"], ) @@ -144,5 +160,5 @@ def test_layer_specific_scaling_is_allowed_without_ipc(tmp_path): r=2, lora_alpha=4, use_rslora=True, rank_pattern={"proj": 4}, alpha_pattern={"proj": 8} ).save_pretrained(tmp_path) validate_actor_lora_adapter( - _args(use_lora=True, lora_adapter_path=str(tmp_path), update_weight_target_modules=["transformer"]) + _args(use_lora=True, lora_adapter_paths=[str(tmp_path)], update_weight_target_modules=["transformer"]) ) From 6665786554a35f5926712c2037decbdc8d4f8690 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 29 Sep 2026 15:48:24 -0700 Subject: [PATCH 2/2] test(fsdp): keep the caller's PYTHONPATH behind the repository root in the reference worker --- .../fast-gpu/backends/fsdp_utils/test_reference_models.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/fast-gpu/backends/fsdp_utils/test_reference_models.py b/tests/fast-gpu/backends/fsdp_utils/test_reference_models.py index 6d2f8c846..95f97ea64 100644 --- a/tests/fast-gpu/backends/fsdp_utils/test_reference_models.py +++ b/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,