diff --git a/docs/guides/lora.md b/docs/guides/lora.md index b738a33ede7..eb27954b64f 100644 --- a/docs/guides/lora.md +++ b/docs/guides/lora.md @@ -49,6 +49,7 @@ The table below maps equivalent fields and highlights the differences. | Optimized kernels | `use_triton` | — | | Adapter dtype | _(follows base layer)_ | `lora_dtype` | | Experimental A2A comm | — | `a2a_experimental` | +| Warm start from a donor adapter | `restore_from` (adapter directory) | `restore_from` (Megatron `iter_*` checkpoint) | The effective learning-rate multiplier for the adapter is `alpha / dim` on both backends. @@ -71,6 +72,7 @@ policy: dropout_position: "post" # Dropout position: "pre" or "post" lora_A_init: "xavier" # Initialization method: "xavier" or "uniform" use_triton: true # Use Triton-optimized kernels (DTensor v2 path) + restore_from: null # Warm start from a donor adapter checkpoint (see below) ``` ### DTensor Parameter Details @@ -85,6 +87,7 @@ policy: - **`dropout_position`** (str): Apply dropout before (`"pre"`) or after (`"post"`) LoRA. - **`lora_A_init`** (str): Initialization method for the LoRA A matrix (`"xavier"` or `"uniform"`). The B matrix is always initialized to zero. - **`use_triton`** (bool): Use Triton-optimized kernels for better performance. Used for DTensor v2 only. **Note**: [Automodel does not support Triton for TP > 1](https://github.com/NVIDIA-NeMo/Automodel/blob/b2db55eee98dfe81a8bfe5e23ac4e57afd8ab261/nemo_automodel/recipes/llm/train_ft.py#L199). Set to `false` when `tensor_parallel_size > 1` to avoid compatibility issues. +- **`restore_from`** (str, optional): Path to a donor PEFT adapter checkpoint (a directory containing `adapter_model.safetensors` + `adapter_config.json`, e.g. a previous run's `step_*/policy/weights`) whose adapter weights initialize this run's LoRA modules. See [Warm-Starting from a LoRA Checkpoint](#warm-starting-from-a-lora-checkpoint). ## Megatron Configuration @@ -105,6 +108,7 @@ policy: lora_B_init_method: "zero" # Initialization method for lora B: "zero" a2a_experimental: false # Enables the experimental All-to-All (A2A) communication strategy lora_dtype: None # Adapter weights dtype + restore_from: null # Warm start from a donor PEFT checkpoint (see below) ``` ### Megatron Parameter Details @@ -126,6 +130,23 @@ policy: - **`lora_B_init_method`** (str): Initialization method for the low-rank matrix B. Defaults to `"zero"`. - **`a2a_experimental`** (bool): Enables the experimental All-to-All (A2A) communication strategy. Defaults to `False`. - **`lora_dtype`** (torch.dtype): Adapter weights dtype. By default it follows `orig_linear`'s dtype, but for quantized weights (e.g. 4-bit) it must be specified explicitly. +- **`restore_from`** (str, optional): Path to a donor Megatron-Bridge PEFT checkpoint (an `iter_XXXXXXX` directory, or a checkpoint root resolving to one) whose adapter weights initialize this run's LoRA modules. See [Warm-Starting from a LoRA Checkpoint](#warm-starting-from-a-lora-checkpoint). + +## Warm-Starting from a LoRA Checkpoint + +`restore_from` initializes this run's adapters from a donor PEFT checkpoint, e.g. to carry an +SFT LoRA into GRPO. The accepted path format differs by backend: + +- **DTensor (Automodel)**: `policy.dtensor_cfg.lora_cfg.restore_from` — a directory containing + `adapter_model.safetensors` + `adapter_config.json` (a previous run's `step_*/policy/weights` + or its `model/` subdirectory). +- **Megatron Core**: `policy.megatron_cfg.peft.restore_from` — a native Megatron `iter_XXXXXXX` + directory, or a checkpoint root containing one. + +The donor's `dim`/`alpha` must match this run's; setup fails loudly otherwise. Optimizer, RNG, +and train state always start fresh. `restore_from` requires `enabled: true`. When resuming from a +NeMo RL training checkpoint, the resumed weights take precedence for the policy; the KL reference +policy still anchors to the warm-started (donor) adapters, not to the bare base model. ## Usage by Algorithm diff --git a/examples/configs/distillation_math_megatron.yaml b/examples/configs/distillation_math_megatron.yaml index 834e01ccd2b..1ecc6ff1c02 100644 --- a/examples/configs/distillation_math_megatron.yaml +++ b/examples/configs/distillation_math_megatron.yaml @@ -102,6 +102,7 @@ policy: &POLICY_BASE lora_B_init_method: "zero" a2a_experimental: false lora_dtype: null + restore_from: null # Warm start: path to a native Megatron PEFT checkpoint (iter_XXXXXXX dir or its root) to initialize this run's LoRA weights from, e.g. a previous SFT run optimizer: optimizer: "adam" diff --git a/examples/configs/dpo.yaml b/examples/configs/dpo.yaml index 94478cb8317..444a0239056 100755 --- a/examples/configs/dpo.yaml +++ b/examples/configs/dpo.yaml @@ -85,6 +85,7 @@ policy: dropout_position: "post" # Where to apply dropout: "pre" (before LoRA) or "post" (after LoRA) lora_A_init: "xavier" # Initialization method for LoRA A matrix: "xavier" or "uniform" use_triton: true # Use Triton-optimized kernels for LoRA (faster but requires flash-attn). Disable when tensor_parallel_size > 1 + restore_from: null # Warm start: path to a PEFT adapter checkpoint (adapter_model.safetensors + adapter_config.json) to initialize this run's LoRA weights from, e.g. a previous SFT run's step_*/policy/weights dynamic_batching: enabled: false @@ -198,6 +199,7 @@ policy: lora_B_init_method: "zero" a2a_experimental: false lora_dtype: null + restore_from: null # Warm start: path to a native Megatron PEFT checkpoint (iter_XXXXXXX dir or its root) to initialize this run's LoRA weights from, e.g. a previous SFT run optimizer: optimizer: "adam" diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 53c8c13abce..08c8f6653eb 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -172,6 +172,7 @@ policy: dropout_position: "post" # Where to apply dropout: "pre" (before LoRA) or "post" (after LoRA) lora_A_init: "xavier" # Initialization method for LoRA A matrix: "xavier" or "uniform" use_triton: true # Use Triton-optimized kernels for LoRA (faster but requires flash-attn). Disable when tensor_parallel_size > 1 + restore_from: null # Warm start: path to a PEFT adapter checkpoint (adapter_model.safetensors + adapter_config.json) to initialize this run's LoRA weights from, e.g. a previous SFT run's step_*/policy/weights megatron_cfg: enabled: false @@ -265,6 +266,7 @@ policy: lora_B_init_method: "zero" a2a_experimental: false lora_dtype: None + restore_from: null # Warm start: path to a native Megatron PEFT checkpoint (iter_XXXXXXX dir or its root) to initialize this run's LoRA weights from, e.g. a previous SFT run optimizer: optimizer: "adam" diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml index 15700751845..fc4bea59032 100644 --- a/examples/configs/grpo_math_1B_megatron.yaml +++ b/examples/configs/grpo_math_1B_megatron.yaml @@ -148,6 +148,7 @@ policy: lora_B_init_method: "zero" a2a_experimental: false lora_dtype: null + restore_from: null # Warm start: path to a native Megatron PEFT checkpoint (iter_XXXXXXX dir or its root) to initialize this run's LoRA weights from, e.g. a previous SFT run optimizer: optimizer: "adam" diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index e0fea886a4e..4b42482982c 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -134,6 +134,7 @@ policy: dropout_position: "post" lora_A_init: "xavier" use_triton: true + restore_from: null # Warm start: path to a PEFT adapter checkpoint (adapter_model.safetensors + adapter_config.json) to initialize this run's LoRA weights from, e.g. a previous SFT run's step_*/policy/weights megatron_cfg: enabled: false @@ -334,6 +335,7 @@ value: dropout_position: "post" lora_A_init: "xavier" use_triton: true + restore_from: null # Warm start: path to a PEFT adapter checkpoint (adapter_model.safetensors + adapter_config.json) to initialize this run's LoRA weights from, e.g. a previous SFT run's step_*/policy/weights megatron_cfg: enabled: false diff --git a/examples/configs/sft.yaml b/examples/configs/sft.yaml index 435aa2996fa..df9e0fd3f09 100644 --- a/examples/configs/sft.yaml +++ b/examples/configs/sft.yaml @@ -76,6 +76,7 @@ policy: dropout_position: "post" # Where to apply dropout: "pre" (before LoRA) or "post" (after LoRA) lora_A_init: "xavier" # Initialization method for LoRA A matrix: "xavier" or "uniform" use_triton: true # Use Triton-optimized kernels for LoRA (faster but requires flash-attn). Disable when tensor_parallel_size > 1 + restore_from: null # Warm start: path to a PEFT adapter checkpoint (adapter_model.safetensors + adapter_config.json) to initialize this run's LoRA weights from, e.g. a previous SFT run's step_*/policy/weights dynamic_batching: enabled: false @@ -184,6 +185,7 @@ policy: lora_B_init_method: "zero" a2a_experimental: false lora_dtype: None + restore_from: null # Warm start: path to a native Megatron PEFT checkpoint (iter_XXXXXXX dir or its root) to initialize this run's LoRA weights from, e.g. a previous SFT run optimizer: diff --git a/nemo_rl/models/automodel/setup.py b/nemo_rl/models/automodel/setup.py index 1fb4db4329f..2131acaf5d8 100644 --- a/nemo_rl/models/automodel/setup.py +++ b/nemo_rl/models/automodel/setup.py @@ -16,13 +16,17 @@ import importlib import inspect +import json import os +import shutil +import tempfile from functools import partial from typing import Any, Optional, Union import torch from hydra.utils import get_class from nemo_automodel import NeMoAutoModelForSequenceClassification +from safetensors import safe_open try: from nemo_automodel import NeMoAutoModelForTokenClassification @@ -51,6 +55,12 @@ class NeMoAutoModelForTokenClassification( from nemo_automodel._transformers.auto_tokenizer import NeMoAutoTokenizer from nemo_automodel._transformers.registry import ModelRegistry from nemo_automodel.components._peft.lora import PeftConfig +from nemo_automodel.components.checkpoint.checkpointing import ( + _maybe_adapt_state_dict_to_hf, +) +from nemo_automodel.components.checkpoint.stateful_wrappers import ( + _rename_dora_keys_to_hf, +) from nemo_automodel.components.config.loader import _resolve_target from nemo_automodel.components.distributed.config import FSDP2Config from nemo_automodel.components.distributed.mesh_utils import create_device_mesh @@ -71,7 +81,7 @@ class NeMoAutoModelForTokenClassification( ModelAndOptimizerState, RuntimeConfig, ) -from nemo_rl.models.policy import PolicyConfig, TokenizerConfig +from nemo_rl.models.policy import LoRAConfig, PolicyConfig, TokenizerConfig from nemo_rl.models.policy.utils import configure_dynamo_cache, resolve_model_class STRING_TO_DTYPE = { @@ -586,6 +596,186 @@ def _noop(*args, **kwargs) -> None: # pragma: no cover _model_init._restore_loaded_model_dtype = _noop +def _resolve_lora_adapter_dir(restore_from: str) -> str: + """Resolve a ``lora_cfg.restore_from`` path to the directory holding the adapter files. + + Accepts a NeMo RL checkpoint weights directory (``step_*/policy/weights``), + its ``model`` subdirectory, or any directory directly containing + ``adapter_model.safetensors`` + ``adapter_config.json`` (HF PEFT layout). + """ + for candidate in (restore_from, os.path.join(restore_from, "model")): + if os.path.isfile(os.path.join(candidate, "adapter_model.safetensors")): + return candidate + raise FileNotFoundError( + f"dtensor_cfg.lora_cfg.restore_from={restore_from!r}: no " + "adapter_model.safetensors found there or in its 'model' subdirectory. " + "restore_from must point to a PEFT adapter checkpoint (a directory " + "containing adapter_model.safetensors + adapter_config.json, e.g. a " + "previous run's step_*/policy/weights directory)." + ) + + +def _validate_lora_adapter_config( + adapter_dir: str, lora_cfg: LoRAConfig, model_name: str +) -> None: + """Fail closed if the donor adapter's metadata is incompatible with this run. + + Checks the standard HF PEFT ``adapter_config.json`` fields: the adapter must + be a LoRA adapter with the same rank (r) and scaling (lora_alpha) as this + run's ``lora_cfg``, and must have been trained on the same base model. + """ + config_path = os.path.join(adapter_dir, "adapter_config.json") + if not os.path.isfile(config_path): + raise FileNotFoundError( + f"dtensor_cfg.lora_cfg.restore_from: {config_path} not found. The " + "donor checkpoint must carry an adapter_config.json so its " + "provenance can be validated." + ) + with open(config_path) as f: + adapter_config = json.load(f) + if adapter_config.get("peft_type") != "LORA": + raise ValueError( + f"dtensor_cfg.lora_cfg.restore_from: {config_path} has " + f"peft_type={adapter_config.get('peft_type')!r}; only 'LORA' " + "adapters can be warm-started from." + ) + for config_key, lora_key in (("r", "dim"), ("lora_alpha", "alpha")): + if config_key not in adapter_config: + raise ValueError( + f"dtensor_cfg.lora_cfg.restore_from: {config_path} has no " + f"{config_key!r} field; cannot verify compatibility with this " + "run's lora_cfg." + ) + if int(adapter_config[config_key]) != int(lora_cfg[lora_key]): + raise ValueError( + f"dtensor_cfg.lora_cfg.restore_from: donor adapter " + f"{config_key}={adapter_config[config_key]} does not match this " + f"run's lora_cfg.{lora_key}={lora_cfg[lora_key]}. Warm starting " + "requires the same LoRA rank and scaling; train a new adapter " + "instead." + ) + donor_base = adapter_config.get("base_model_name_or_path") + if donor_base and donor_base != "N/A" and donor_base != model_name: + raise ValueError( + f"dtensor_cfg.lora_cfg.restore_from: donor adapter was trained on " + f"base model {donor_base!r} but this run uses model_name=" + f"{model_name!r}." + ) + + +def _validate_lora_adapter_keys( + adapter_dir: str, model: torch.nn.Module, *, moe_mesh: Any = None +) -> None: + """Fail closed if the donor adapter tensors don't cover this model's LoRA params. + + The underlying PEFT load is unconditionally non-strict (a key mismatch is + only a warning), so this check is the only guarantee that the donor covers + every LoRA parameter. Expected keys are converted through the same native + to HF state-dict adapter path used when saving, so the comparison also works + for custom model implementations and expert-parallel models. + """ + with safe_open( + os.path.join(adapter_dir, "adapter_model.safetensors"), framework="pt" + ) as f: + donor_keys = set(f.keys()) + if not donor_keys: + raise ValueError( + f"dtensor_cfg.lora_cfg.restore_from: {adapter_dir}/" + "adapter_model.safetensors contains no tensors." + ) + # HF PEFT exports prefix keys with "base_model.model."; the loader strips + # the prefix when present, so accept either form here. + prefix = "base_model.model." + normalized_donor_keys = { + key[len(prefix) :] if key.startswith(prefix) else key for key in donor_keys + } + + expected_state_dict = { + f"{prefix}{name.replace('_checkpoint_wrapped_module.', '')}": ( + param.full_tensor().detach().cpu() + if hasattr(param, "full_tensor") + else param.detach().cpu() + ) + for name, param in model.named_parameters() + if "lora_" in name + } + _rename_dora_keys_to_hf(expected_state_dict) + expected_state_dict = _maybe_adapt_state_dict_to_hf( + model, + expected_state_dict, + quantization=False, + device_mesh=moe_mesh, + ) + normalized_expected_keys = { + key[len(prefix) :] if key.startswith(prefix) else key + for key in expected_state_dict + } + missing = sorted(normalized_expected_keys - normalized_donor_keys) + unexpected = sorted(normalized_donor_keys - normalized_expected_keys) + if missing or unexpected: + raise ValueError( + "dtensor_cfg.lora_cfg.restore_from: donor adapter key mismatch " + f"against this run's LoRA parameters. Missing from donor: " + f"{missing[:5]}{' ...' if len(missing) > 5 else ''}; unexpected in " + f"donor: {unexpected[:5]}{' ...' if len(unexpected) > 5 else ''}. " + "The donor adapter must target the same modules as this run's " + "lora_cfg." + ) + + +def _load_initial_lora_adapter( + model: torch.nn.Module, + checkpoint_manager: Any, + restore_from: str, + lora_cfg: LoRAConfig, + model_name: str, +) -> None: + """Warm-start the model's LoRA adapters from a donor PEFT adapter checkpoint. + + The load goes through the Automodel checkpointer's PEFT path (each rank + reads adapter_model.safetensors, then + ``set_model_state_dict(broadcast_from_rank0=True)`` places it with + DTensor/EP awareness), the same machinery used to resume NeMo RL PEFT + checkpoints. Optimizer state is not loaded: warm starts begin with a + fresh optimizer. + """ + adapter_dir = os.path.abspath(_resolve_lora_adapter_dir(restore_from)) + _validate_lora_adapter_config(adapter_dir, lora_cfg, model_name) + _validate_lora_adapter_keys( + adapter_dir, + model, + moe_mesh=getattr(checkpoint_manager, "moe_mesh", None), + ) + staging_dir = None + load_dir = adapter_dir + if os.path.basename(adapter_dir.rstrip(os.sep)) != "model": + # Automodel's checkpointer selects its PEFT safetensors read by a + # substring test on the path ("/model" in path); a bare adapter + # directory would fall through to the DCP/HF-storage-reader branch + # instead. Expose it under a temporary "model" path component. + staging_dir = tempfile.mkdtemp(prefix="nrl_lora_warm_start_") + load_dir = os.path.join(staging_dir, "model") + os.symlink(adapter_dir, load_dir) + assert checkpoint_manager.checkpointer is not None, ( + "Checkpointer must be initialized before warm starting LoRA adapters." + ) + checkpoint_manager.update_checkpointer_config( + config_updates={ + "model_save_format": "safetensors", + "is_peft": True, + # The donor checkpoint is already dequantized. + "dequantize_base_checkpoint": False, + }, + checkpoint_root=os.path.dirname(load_dir.rstrip(os.sep)), + ) + try: + checkpoint_manager.checkpointer.load_model(model=model, model_path=load_dir) + finally: + if staging_dir is not None: + shutil.rmtree(staging_dir, ignore_errors=True) + print(f"Warm-started LoRA adapters from {adapter_dir}") + + def setup_model_and_optimizer( config: PolicyConfig, tokenizer: AutoTokenizer, @@ -676,6 +866,12 @@ def setup_model_and_optimizer( lora_cfg = config["dtensor_cfg"].get("lora_cfg", None) peft_config = None lora_enabled = lora_cfg is not None and lora_cfg["enabled"] + if not lora_enabled and (lora_cfg or {}).get("restore_from"): + raise ValueError( + "dtensor_cfg.lora_cfg.restore_from is set but " + "dtensor_cfg.lora_cfg.enabled is False. Enable LoRA to warm start " + "from an adapter checkpoint." + ) if lora_enabled: if tp_size > 1: assert not lora_cfg["use_triton"], ( @@ -872,6 +1068,20 @@ def setup_model_and_optimizer( optimizer_path=optimizer_path, scheduler=scheduler, ) + elif lora_enabled and lora_cfg.get("restore_from") is not None: + # Warm start: base weights were already loaded by from_pretrained above; + # restore only the donor adapter weights. This runs before the worker + # captures the KL reference state dict, so the reference policy is the + # warm-started initial policy. When this setup call is part of a + # deferred resume (weights_path temporarily None), the resumed + # checkpoint is loaded afterwards and overwrites these weights. + _load_initial_lora_adapter( + model=model, + checkpoint_manager=checkpoint_manager, + restore_from=lora_cfg["restore_from"], + lora_cfg=lora_cfg, + model_name=config["model_name"], + ) else: print( "No weights path provided. Loaded base HF weights via from_pretrained (default policy init)" diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 399633cea52..e7e306dc322 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -24,6 +24,7 @@ from typing import Any, Callable, Optional, TypeVar import torch +import yaml from megatron.bridge import AutoBridge from megatron.bridge.models.model_provider import ModelProviderMixin, get_model from megatron.bridge.peft.lora import LoRA @@ -55,14 +56,16 @@ _create_peft_pre_wrap_hook, _update_model_config_funcs, ) -from megatron.bridge.training.state import GlobalState +from megatron.bridge.training.state import GlobalState, TrainState from megatron.bridge.training.tokenizers.tokenizer import build_tokenizer from megatron.bridge.training.utils.pg_utils import get_pg_collection from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size from megatron.core import parallel_state from megatron.core.inference.shards import build_inference_pg_collection +from megatron.core.num_microbatches_calculator import update_num_microbatches from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.rerun_state_machine import RerunMode, get_rerun_state_machine from megatron.core.transformer import MegatronModule from megatron.core.transformer.enums import AttnBackend, InferenceCudaGraphScope from megatron.core.transformer.module import Float16Module @@ -249,7 +252,11 @@ def _sync_distrib_opt(distrib_opt): router_replay_enabled, validate_router_replay_config, ) -from nemo_rl.models.policy import MegatronConfig, PolicyConfig +from nemo_rl.models.policy import ( + MegatronConfig, + MegatronPeftConfig, + PolicyConfig, +) from nemo_rl.models.policy.utils import ( configure_dynamo_cache, get_megatron_checkpoint_dir, @@ -452,7 +459,11 @@ def _get_hf_config_overrides_hash(overrides: dict[str, Any]) -> str: return hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:12] -def _resolve_iter_dir_from_root(path: str, not_found_msg: str) -> str: +def _resolve_iter_dir_from_root( + path: str, + not_found_msg: str, + config_key: str = "pretrained_checkpoint.path", +) -> str: """Resolve the latest iteration directory under ``path``. Checks ``latest_checkpointed_iteration.txt`` first; falls back to scanning @@ -468,7 +479,7 @@ def _resolve_iter_dir_from_root(path: str, not_found_msg: str) -> str: return os.path.join(path, f"iter_{int(iteration_str):07d}") except ValueError: raise ValueError( - f"pretrained_checkpoint.path={path!r}: " + f"{config_key}={path!r}: " f"latest_checkpointed_iteration.txt contains unexpected value " f"{iteration_str!r}; expected an integer or 'release'." ) @@ -485,6 +496,177 @@ def _resolve_iter_dir_from_root(path: str, not_found_msg: str) -> str: return os.path.join(path, iter_subdirs[-1]) +def _resolve_peft_restore_dir(restore_from: str) -> str: + """Resolve a ``megatron_cfg.peft.restore_from`` path to an iteration directory. + + Accepts either a specific iteration directory (containing run_config.yaml) + or a checkpoint root with a latest_checkpointed_iteration.txt tracker file / + iter_* subdirectories. + """ + if not os.path.isdir(restore_from): + raise FileNotFoundError( + f"megatron_cfg.peft.restore_from={restore_from!r} does not exist or is " + "not a directory. It must point to a native Megatron-Bridge PEFT " + "checkpoint (an iter_XXXXXXX directory or a checkpoint root " + "containing one)." + ) + if os.path.exists(os.path.join(restore_from, "run_config.yaml")): + return restore_from + resolved = _resolve_iter_dir_from_root( + restore_from, + f"megatron_cfg.peft.restore_from={restore_from!r} does not contain " + "run_config.yaml, latest_checkpointed_iteration.txt, or any iter_* " + "subdirectories. It must point to a native Megatron-Bridge PEFT " + "checkpoint (an iter_XXXXXXX directory or a checkpoint root containing " + "one).", + config_key="megatron_cfg.peft.restore_from", + ) + if not os.path.exists(os.path.join(resolved, "run_config.yaml")): + raise FileNotFoundError( + f"megatron_cfg.peft.restore_from={restore_from!r}: resolved to " + f"iteration directory {resolved!r} but it does not contain " + "run_config.yaml." + ) + return resolved + + +def _validate_peft_restore_config( + restore_dir: str, peft_cfg: MegatronPeftConfig +) -> None: + """Fail closed if the donor checkpoint's PEFT config is incompatible. + + Adapter tensors are loadable only when the rank (dim) matches, and only + functionally equivalent when the scaling (alpha) matches, so both must + agree with this run's ``megatron_cfg.peft`` before any weights are loaded. + The targeted module sets must match as well: the adapter-only distributed + load only raises on donor keys missing from the checkpoint, while + unrequested donor keys are discarded silently (DCP's ASSUME_OK_UNEXPECTED + skips the mismatch check), so a superset donor would otherwise load with + no diagnostics and leave this run's extra adapters at fresh init. + """ + run_config_path = os.path.join(restore_dir, "run_config.yaml") + with open(run_config_path) as f: + run_config = yaml.safe_load(f) + saved_peft = run_config.get("peft") if isinstance(run_config, dict) else None + if not isinstance(saved_peft, dict): + raise ValueError( + f"megatron_cfg.peft.restore_from={restore_dir!r}: {run_config_path} " + "has no 'peft' section. restore_from requires a checkpoint saved " + "with PEFT enabled." + ) + for key in ("dim", "alpha"): + if key not in saved_peft: + raise ValueError( + f"megatron_cfg.peft.restore_from={restore_dir!r}: the donor " + f"checkpoint's run_config.yaml peft section has no {key!r} key; " + "cannot verify compatibility with this run's peft config." + ) + if int(saved_peft[key]) != int(peft_cfg[key]): + raise ValueError( + f"megatron_cfg.peft.restore_from={restore_dir!r}: donor " + f"checkpoint peft.{key}={saved_peft[key]} does not match this " + f"run's peft.{key}={peft_cfg[key]}. Warm starting requires the " + "same LoRA rank and scaling; train a new adapter instead." + ) + for key in ("target_modules", "exclude_modules"): + if key not in saved_peft: + raise ValueError( + f"megatron_cfg.peft.restore_from={restore_dir!r}: the donor " + f"checkpoint's run_config.yaml peft section has no {key!r} key; " + "cannot verify compatibility with this run's peft config." + ) + if set(saved_peft[key]) != set(peft_cfg[key]): + raise ValueError( + f"megatron_cfg.peft.restore_from={restore_dir!r}: donor " + f"checkpoint peft.{key}={saved_peft[key]} does not match this " + f"run's peft.{key}={peft_cfg[key]}. Warm starting requires the " + "donor to target the same modules; train a new adapter instead." + ) + # These megatron-bridge LoRA fields change the adapter shape/key layout + # for MoE expert layers. NeMo RL never sets them (a run always uses the + # bridge defaults), but a native Megatron-Bridge donor checkpoint may + # have; a mismatch would restore onto a different adapter layout. + moe_shaping_defaults = { + "normalize_moe_lora": False, + "share_expert_adapters": True, + "experts_shared_outer_loras": False, + } + for key, default in moe_shaping_defaults.items(): + if key in saved_peft and saved_peft[key] != peft_cfg.get(key, default): + raise ValueError( + f"megatron_cfg.peft.restore_from={restore_dir!r}: donor " + f"checkpoint peft.{key}={saved_peft[key]} does not match this " + f"run's peft.{key}={peft_cfg.get(key, default)}. Warm starting " + "requires the same MoE adapter layout; train a new adapter " + "instead." + ) + + +def _create_peft_warm_start_hook( + megatron_cfg: ConfigContainer, + state: GlobalState, + restore_dir: str, +) -> Callable[[list[MegatronModule]], list[MegatronModule]]: + """Create a pre-wrap hook that warm-starts LoRA adapters from a donor checkpoint. + + The hook must run immediately after the PEFT pre-wrap hook (which loads the + pretrained base weights and attaches fresh adapters). It routes the donor + load through Megatron-Bridge's PEFT-resume path: loading from + ``checkpoint.load`` with PEFT configured and ``finetune=False`` filters the + generated sharded state dict to adapter keys and drops to a non-strict + load, so an adapter-only donor checkpoint restores cleanly onto the fresh + base with distributed resharding handled by the torch_dist loader. + """ + + def peft_warm_start_hook(model: list[MegatronModule]) -> list[MegatronModule]: + ckpt_cfg = megatron_cfg.checkpoint + rerun_state_machine = get_rerun_state_machine() + original_rerun_mode = rerun_state_machine.get_mode() + original_load = ckpt_cfg.load + original_finetune = ckpt_cfg.finetune + original_load_optim = ckpt_cfg.load_optim + original_load_rng = ckpt_cfg.load_rng + ckpt_cfg.load = restore_dir + # finetune=False is required for the adapter-only filter, but + # optimizer/RNG state must not come from the donor run. + ckpt_cfg.finetune = False + ckpt_cfg.load_optim = False + ckpt_cfg.load_rng = False + # Bridge restores rerun metadata independently of load_rng/load_optim. + # A disabled machine ignores donor rerun state, and setup is serialized, + # so temporarily disabling it makes this adapter load weights-only. + rerun_state_machine.set_mode(RerunMode.DISABLED) + try: + _load_checkpoint_from_path( + load_dir=restore_dir, + state=state, + model=model, + optimizer=None, + opt_param_scheduler=None, + checkpointing_context={}, + skip_load_to_model_and_opt=False, + ignore_ckpt_step=True, + ) + finally: + ckpt_cfg.load = original_load + ckpt_cfg.finetune = original_finetune + ckpt_cfg.load_optim = original_load_optim + ckpt_cfg.load_rng = original_load_rng + rerun_state_machine.set_mode(original_rerun_mode) + # finetune=False also pulled the donor run's train state (step, consumed + # samples) and fed it to the global microbatch calculator. This is a + # warm start, not a resume: reset both so the new run starts at step 0. + state.train_state = TrainState() + update_num_microbatches(consumed_samples=0, verbose=False) + print( + f"Warm-started PEFT adapters from {restore_dir} " + "(optimizer, RNG, and train state start fresh)." + ) + return model + + return peft_warm_start_hook + + def validate_model_paths(config: PolicyConfig) -> tuple[str, str, bool]: """Validate and setup model paths. @@ -1719,6 +1901,11 @@ def setup_model_and_optimizer( pre_wrap_hook = [] use_peft = policy_cfg["megatron_cfg"].get("peft", {}).get("enabled", False) + if not use_peft and policy_cfg["megatron_cfg"].get("peft", {}).get("restore_from"): + raise ValueError( + "megatron_cfg.peft.restore_from is set but megatron_cfg.peft.enabled " + "is False. Enable PEFT to warm start from an adapter checkpoint." + ) draft_enabled = "draft" in policy_cfg and policy_cfg["draft"]["enabled"] resume_checkpoint_exists = ( megatron_cfg.checkpoint.load is not None @@ -1815,8 +2002,15 @@ def apply_freeze(megatron_model): a2a_experimental=peft_cfg["a2a_experimental"], lora_dtype=peft_cfg["lora_dtype"], ) + # Resolve and validate the warm-start donor checkpoint up front so a + # bad path or mismatched donor fails before any model construction. + peft_restore_dir = None + if peft_cfg.get("restore_from") is not None: + peft_restore_dir = _resolve_peft_restore_dir(peft_cfg["restore_from"]) + _validate_peft_restore_config(peft_restore_dir, peft_cfg) else: peft = None + peft_restore_dir = None megatron_cfg.peft = peft @@ -1868,6 +2062,14 @@ def composed_peft_hook(model: list[MegatronModule]) -> list[MegatronModule]: pre_wrap_hook.extend([composed_peft_hook]) + # Warm start the adapters from the donor checkpoint after the base + # weights are loaded and fresh adapters are attached. Skipped when + # resuming: the resume checkpoint already carries this run's adapters. + if peft_restore_dir is not None and not resume_checkpoint_exists: + pre_wrap_hook.append( + _create_peft_warm_start_hook(megatron_cfg, state, peft_restore_dir) + ) + if draft_enabled: draft_pre_wrap_hook = _create_draft_pre_wrap_hook( policy_cfg, @@ -2140,6 +2342,21 @@ def composed_peft_hook(model: list[MegatronModule]) -> list[MegatronModule]: ref_pre_wrap_hooks.extend([composed_peft_hook]) + # Anchor the reference policy to the warm-started initial policy: with + # restore_from set, the policy starts from the donor adapters, so the + # reference must include them too (zero-init adapters would anchor KL + # to the bare base model). Unlike the policy model this runs on resumes + # as well — the reference must stay anchored to the initial policy. + peft_restore_from = config["megatron_cfg"].get("peft", {}).get("restore_from") + if peft_restore_from is not None: + ref_pre_wrap_hooks.append( + _create_peft_warm_start_hook( + ref_megatron_cfg, + ref_state, + _resolve_peft_restore_dir(peft_restore_from), + ) + ) + try: reference_model = get_model( megatron_cfg.model, diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 7d14674a392..e6611f4260f 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -103,6 +103,13 @@ class LoRAConfig(TypedDict): dropout_position: Literal["pre", "post"] lora_A_init: str use_triton: NotRequired[bool] + # Warm start: path to a PEFT adapter checkpoint (a directory containing + # adapter_model.safetensors + adapter_config.json, e.g. a previous run's + # step_*/policy/weights or step_*/policy/weights/model directory) whose + # adapter weights initialize this run's LoRA modules. Ignored when resuming + # from a NeMo RL training checkpoint (resumed weights take precedence for + # the policy; the reference policy still anchors to the restored adapters). + restore_from: NotRequired[str | None] class AutomodelBackendConfig(TypedDict): @@ -241,6 +248,14 @@ class MegatronPeftConfig(TypedDict): lora_B_init_method: str a2a_experimental: bool lora_dtype: str | None + # Warm start: path to a native Megatron-Bridge PEFT checkpoint (an + # iter_XXXXXXX directory, or a checkpoint root resolving to one) whose + # adapter weights initialize this run's LoRA modules. The donor checkpoint + # must have been saved with a matching peft configuration (dim and alpha + # are validated against its run_config.yaml). Ignored when resuming from a + # NeMo RL training checkpoint (resumed weights take precedence for the + # policy; the reference policy still anchors to the restored adapters). + restore_from: NotRequired[str | None] class MegatronOptimizerConfig(TypedDict): diff --git a/tests/unit/models/automodel/test_automodel_setup.py b/tests/unit/models/automodel/test_automodel_setup.py index 9c26f741e15..9d45bbc11bf 100644 --- a/tests/unit/models/automodel/test_automodel_setup.py +++ b/tests/unit/models/automodel/test_automodel_setup.py @@ -15,7 +15,8 @@ """Unit tests for automodel setup utilities.""" import os -from unittest.mock import MagicMock, Mock, patch +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock, create_autospec, patch import pytest @@ -26,7 +27,9 @@ pytest.skip("nemo_automodel not available", allow_module_level=True) import torch +from nemo_automodel.components.checkpoint.checkpointing import Checkpointer +from nemo_rl.models.automodel.checkpoint import AutomodelCheckpointManager from nemo_rl.models.automodel.config import DistributedContext from nemo_rl.models.automodel.setup import ( ModelAndOptimizerState, @@ -40,6 +43,7 @@ ) +@pytest.mark.automodel def test_token_classification_backport_still_required(): with pytest.raises(ImportError): from nemo_automodel import NeMoAutoModelForTokenClassification # noqa: F401 @@ -956,6 +960,32 @@ def test_setup_model_and_optimizer_basic( # Verify config= is NOT passed (avoids duplicate arg for custom models) assert "config" not in call_kwargs + @patch("nemo_rl.models.automodel.setup.torch.distributed.get_rank") + def test_restore_from_without_lora_enabled_raises( + self, + mock_get_rank, + mock_config, + mock_runtime_config, + mock_distributed_context, + mock_checkpoint_manager, + mock_tokenizer, + ): + """restore_from with LoRA disabled must fail loudly, not silently no-op.""" + mock_get_rank.return_value = 0 + mock_config["dtensor_cfg"]["lora_cfg"] = { + "enabled": False, + "restore_from": "/donor/step_5/policy/weights", + } + + with pytest.raises(ValueError, match="lora_cfg.restore_from is set"): + setup_model_and_optimizer( + config=mock_config, + tokenizer=mock_tokenizer, + runtime_config=mock_runtime_config, + distributed_context=mock_distributed_context, + checkpoint_manager=mock_checkpoint_manager, + ) + @patch("nemo_rl.models.automodel.setup.torch.optim.lr_scheduler.LambdaLR") @patch("nemo_rl.models.automodel.setup.torch.distributed.get_rank") @patch("nemo_rl.models.automodel.setup.get_class") @@ -2129,6 +2159,7 @@ def test_does_not_use_hf_auto_tokenizer(self, mock_nemo_auto_tokenizer): mock_nemo_auto_tokenizer.from_pretrained.assert_called_once() +@pytest.mark.automodel class TestMaybeSetForceHf: """Tests for _maybe_set_force_hf adapter compatibility check.""" @@ -2290,3 +2321,590 @@ def test_automodel_dtype_restore_workaround_still_needed(monkeypatch): "_disable_automodel_checkpoint_dtype_restore() workaround in setup.py is obsolete - " "remove it and this test." ) + + +class _TinyLoraModel(torch.nn.Module): + """Tiny module with one LoRA-adapted linear for warm-start tests.""" + + def __init__(self): + from nemo_automodel.components._peft.lora import LinearLoRA + + super().__init__() + self.layer = LinearLoRA(torch.nn.Linear(4, 8), dim=2, alpha=4) + + +class _IdentityStateDictAdapter: + """Minimal custom adapter that preserves state-dict keys.""" + + @staticmethod + def to_hf(state_dict, **kwargs): + return state_dict + + +class _TinyQwen3MoeLoraModel(torch.nn.Module): + """Minimal native Qwen3-MoE LoRA key layout for adapter-key tests.""" + + def __init__(self, *, ep_size=1): + from nemo_automodel.components.models.qwen3_moe.state_dict_adapter import ( + Qwen3MoeStateDictAdapter, + ) + + super().__init__() + experts = torch.nn.Module() + experts.ep_size = ep_size + experts.register_parameter( + "lora_gate_and_up_A", torch.nn.Parameter(torch.empty(2, 4, 2)) + ) + experts.register_parameter( + "lora_gate_and_up_B", torch.nn.Parameter(torch.empty(2, 2, 6)) + ) + experts.register_parameter( + "lora_down_A", torch.nn.Parameter(torch.empty(2, 3, 2)) + ) + experts.register_parameter( + "lora_down_B", torch.nn.Parameter(torch.empty(2, 2, 4)) + ) + mlp = torch.nn.Module() + mlp.add_module("experts", experts) + layer = torch.nn.Module() + layer.add_module("mlp", mlp) + inner_model = torch.nn.Module() + inner_model.add_module("layers", torch.nn.ModuleList([layer])) + self.add_module("model", inner_model) + self.state_dict_adapter = Qwen3MoeStateDictAdapter( + config=SimpleNamespace(), + moe_config=SimpleNamespace(n_routed_experts=2), + backend=SimpleNamespace(), + ) + + +_QWEN3_MOE_HF_LORA_KEYS = ( + "base_model.model.model.layers.0.mlp.experts.base_layer.lora_B.weight", + "base_model.model.model.layers.0.mlp.experts.base_layer.lora_A.weight", + "base_model.model.model.layers.0.mlp.experts.lora_B.weight", + "base_model.model.model.layers.0.mlp.experts.lora_A.weight", +) + + +def _write_adapter_checkpoint( + adapter_dir, + *, + dim=2, + alpha=4, + base_model_name_or_path="tiny-model", + keys=("layer.lora_A.weight", "layer.lora_B.weight"), + peft_type="LORA", + fill=1.0, +): + """Write a minimal HF-PEFT-style adapter checkpoint (config + safetensors). + + Tensors are filled with a non-zero value so a test cannot confuse + "loaded the donor" with "left the zero-init adapters alone". + """ + import json as _json + + from safetensors.torch import save_file + + os.makedirs(adapter_dir, exist_ok=True) + tensors = {} + for key in keys: + shape = (2, 4) if "lora_A" in key else (8, 2) + tensors[key] = torch.full(shape, fill) + save_file(tensors, os.path.join(adapter_dir, "adapter_model.safetensors")) + with open(os.path.join(adapter_dir, "adapter_config.json"), "w") as f: + _json.dump( + { + "peft_type": peft_type, + "r": dim, + "lora_alpha": alpha, + "base_model_name_or_path": base_model_name_or_path, + "target_modules": ["layer"], + }, + f, + ) + + +def _lora_cfg(dim=2, alpha=4): + return { + "enabled": True, + "target_modules": [], + "exclude_modules": [], + "match_all_linear": True, + "dim": dim, + "alpha": alpha, + "dropout": 0.0, + "dropout_position": "post", + "lora_A_init": "xavier", + } + + +@pytest.mark.automodel +class TestResolveLoraAdapterDir: + def test_direct_adapter_dir(self, tmp_path): + from nemo_rl.models.automodel.setup import _resolve_lora_adapter_dir + + _write_adapter_checkpoint(tmp_path / "adapter") + assert _resolve_lora_adapter_dir(str(tmp_path / "adapter")) == str( + tmp_path / "adapter" + ) + + def test_weights_dir_with_model_subdir(self, tmp_path): + from nemo_rl.models.automodel.setup import _resolve_lora_adapter_dir + + weights_dir = tmp_path / "step_5" / "policy" / "weights" + _write_adapter_checkpoint(weights_dir / "model") + assert _resolve_lora_adapter_dir(str(weights_dir)) == str(weights_dir / "model") + + def test_missing_adapter_file_raises(self, tmp_path): + from nemo_rl.models.automodel.setup import _resolve_lora_adapter_dir + + with pytest.raises(FileNotFoundError, match="adapter_model.safetensors"): + _resolve_lora_adapter_dir(str(tmp_path)) + + +@pytest.mark.automodel +class TestValidateLoraAdapterConfig: + def test_matching_config_passes(self, tmp_path): + from nemo_rl.models.automodel.setup import _validate_lora_adapter_config + + adapter_dir = tmp_path / "adapter" + _write_adapter_checkpoint(adapter_dir) + _validate_lora_adapter_config( + str(adapter_dir), _lora_cfg(), "tiny-model" + ) # should not raise + + def test_peft_type_mismatch_raises(self, tmp_path): + from nemo_rl.models.automodel.setup import _validate_lora_adapter_config + + adapter_dir = tmp_path / "adapter" + _write_adapter_checkpoint(adapter_dir, peft_type="IA3") + with pytest.raises(ValueError, match="peft_type"): + _validate_lora_adapter_config(str(adapter_dir), _lora_cfg(), "tiny-model") + + def test_rank_mismatch_raises(self, tmp_path): + from nemo_rl.models.automodel.setup import _validate_lora_adapter_config + + adapter_dir = tmp_path / "adapter" + _write_adapter_checkpoint(adapter_dir, dim=8) + with pytest.raises(ValueError, match="r=8"): + _validate_lora_adapter_config(str(adapter_dir), _lora_cfg(), "tiny-model") + + def test_alpha_mismatch_raises(self, tmp_path): + from nemo_rl.models.automodel.setup import _validate_lora_adapter_config + + adapter_dir = tmp_path / "adapter" + _write_adapter_checkpoint(adapter_dir, alpha=16) + with pytest.raises(ValueError, match="lora_alpha"): + _validate_lora_adapter_config(str(adapter_dir), _lora_cfg(), "tiny-model") + + def test_base_model_mismatch_raises(self, tmp_path): + from nemo_rl.models.automodel.setup import _validate_lora_adapter_config + + adapter_dir = tmp_path / "adapter" + _write_adapter_checkpoint(adapter_dir, base_model_name_or_path="other-model") + with pytest.raises(ValueError, match="other-model"): + _validate_lora_adapter_config(str(adapter_dir), _lora_cfg(), "tiny-model") + + def test_unknown_base_model_does_not_raise(self, tmp_path): + from nemo_rl.models.automodel.setup import _validate_lora_adapter_config + + adapter_dir = tmp_path / "adapter" + _write_adapter_checkpoint(adapter_dir, base_model_name_or_path="N/A") + _validate_lora_adapter_config(str(adapter_dir), _lora_cfg(), "tiny-model") + + def test_missing_config_file_raises(self, tmp_path): + from nemo_rl.models.automodel.setup import _validate_lora_adapter_config + + adapter_dir = tmp_path / "adapter" + adapter_dir.mkdir() + with pytest.raises(FileNotFoundError, match="adapter_config.json"): + _validate_lora_adapter_config(str(adapter_dir), _lora_cfg(), "tiny-model") + + +@pytest.mark.automodel +class TestValidateLoraAdapterKeys: + def test_prefixed_keys_match(self, tmp_path): + from nemo_rl.models.automodel.setup import _validate_lora_adapter_keys + + adapter_dir = tmp_path / "adapter" + _write_adapter_checkpoint( + adapter_dir, + keys=( + "base_model.model.layer.lora_A.weight", + "base_model.model.layer.lora_B.weight", + ), + ) + _validate_lora_adapter_keys(str(adapter_dir), _TinyLoraModel()) + + def test_missing_key_raises(self, tmp_path): + from nemo_rl.models.automodel.setup import _validate_lora_adapter_keys + + adapter_dir = tmp_path / "adapter" + _write_adapter_checkpoint( + adapter_dir, keys=("base_model.model.layer.lora_A.weight",) + ) + with pytest.raises(ValueError, match="Missing from donor"): + _validate_lora_adapter_keys(str(adapter_dir), _TinyLoraModel()) + + def test_unexpected_key_raises(self, tmp_path): + from nemo_rl.models.automodel.setup import _validate_lora_adapter_keys + + adapter_dir = tmp_path / "adapter" + _write_adapter_checkpoint( + adapter_dir, + keys=( + "layer.lora_A.weight", + "layer.lora_B.weight", + "layer2.lora_A.weight", + ), + ) + with pytest.raises(ValueError, match="unexpected in donor"): + _validate_lora_adapter_keys(str(adapter_dir), _TinyLoraModel()) + + def test_state_dict_adapter_model_still_validated(self, tmp_path): + """A custom state_dict_adapter must not disable the coverage check. + + Regression test: the escape hatch used to be gated on + ``state_dict_adapter``, which 26/33 Automodel architectures set + (including plain LlamaForCausalLM) -- so the check silently passed + while the non-strict PEFT load left donor-uncovered adapters at fresh + init. The validator must instead compare keys after the adapter's HF + conversion. + """ + from nemo_rl.models.automodel.setup import _validate_lora_adapter_keys + + adapter_dir = tmp_path / "adapter" + _write_adapter_checkpoint( + adapter_dir, keys=("base_model.model.layer.lora_A.weight",) + ) + model = _TinyLoraModel() + model.state_dict_adapter = _IdentityStateDictAdapter() + with pytest.raises(ValueError, match="Missing from donor"): + _validate_lora_adapter_keys(str(adapter_dir), model) + + def test_qwen3_moe_hf_keys_match_at_ep1(self, tmp_path): + from nemo_rl.models.automodel.setup import _validate_lora_adapter_keys + + adapter_dir = tmp_path / "adapter" + _write_adapter_checkpoint(adapter_dir, keys=_QWEN3_MOE_HF_LORA_KEYS) + _validate_lora_adapter_keys(str(adapter_dir), _TinyQwen3MoeLoraModel(ep_size=1)) + + def test_expert_parallel_model_with_incomplete_donor_raises(self, tmp_path): + from nemo_rl.models.automodel.setup import _validate_lora_adapter_keys + + adapter_dir = tmp_path / "adapter" + _write_adapter_checkpoint(adapter_dir, keys=_QWEN3_MOE_HF_LORA_KEYS[:-1]) + with pytest.raises(ValueError, match="Missing from donor"): + _validate_lora_adapter_keys( + str(adapter_dir), _TinyQwen3MoeLoraModel(ep_size=2) + ) + + +@pytest.mark.automodel +class TestLoadInitialLoraAdapter: + def test_loads_through_checkpointer(self, tmp_path): + from nemo_rl.models.automodel.setup import _load_initial_lora_adapter + + weights_dir = tmp_path / "step_5" / "policy" / "weights" + adapter_dir = weights_dir / "model" + _write_adapter_checkpoint( + adapter_dir, + keys=( + "base_model.model.layer.lora_A.weight", + "base_model.model.layer.lora_B.weight", + ), + ) + model = _TinyLoraModel() + # autospec binds the mocks to the real signatures, so a wrong kwarg + # name in setup.py fails the test instead of silently recording a call. + manager = create_autospec(AutomodelCheckpointManager, instance=True) + manager.checkpointer = create_autospec(Checkpointer, instance=True) + _load_initial_lora_adapter( + model=model, + checkpoint_manager=manager, + restore_from=str(weights_dir), + lora_cfg=_lora_cfg(), + model_name="tiny-model", + ) + manager.update_checkpointer_config.assert_called_once() + config_updates = manager.update_checkpointer_config.call_args.kwargs[ + "config_updates" + ] + assert config_updates["is_peft"] is True + manager.checkpointer.load_model.assert_called_once_with( + model=model, model_path=str(adapter_dir) + ) + + def test_invalid_donor_fails_before_load(self, tmp_path): + from nemo_rl.models.automodel.setup import _load_initial_lora_adapter + + adapter_dir = tmp_path / "adapter" + # dim=8 donor vs dim=2 run -> validation must fail before any load. + _write_adapter_checkpoint(adapter_dir, dim=8) + manager = MagicMock() + with pytest.raises(ValueError, match="r=8"): + _load_initial_lora_adapter( + model=_TinyLoraModel(), + checkpoint_manager=manager, + restore_from=str(adapter_dir), + lora_cfg=_lora_cfg(), + model_name="tiny-model", + ) + manager.checkpointer.load_model.assert_not_called() + + def test_bare_adapter_dir_loaded_via_model_path(self, tmp_path): + """A bare adapter dir must still take the checkpointer's PEFT branch. + + Automodel selects its PEFT safetensors read by a substring test on the + path ("/model" in path), so the bare layout is staged under a + temporary "model" path component before loading. + """ + from nemo_rl.models.automodel.setup import _load_initial_lora_adapter + + adapter_dir = tmp_path / "adapter" + _write_adapter_checkpoint( + adapter_dir, + keys=( + "base_model.model.layer.lora_A.weight", + "base_model.model.layer.lora_B.weight", + ), + ) + model = _TinyLoraModel() + manager = create_autospec(AutomodelCheckpointManager, instance=True) + manager.checkpointer = create_autospec(Checkpointer, instance=True) + + seen = {} + + def fake_load_model(*, model, model_path, **kwargs): + # Inspected mid-call: the staging symlink is cleaned up after load. + seen["model_path"] = model_path + seen["resolves"] = os.path.isfile( + os.path.join(model_path, "adapter_model.safetensors") + ) + + manager.checkpointer.load_model.side_effect = fake_load_model + _load_initial_lora_adapter( + model=model, + checkpoint_manager=manager, + restore_from=str(adapter_dir), + lora_cfg=_lora_cfg(), + model_name="tiny-model", + ) + assert os.path.basename(seen["model_path"]) == "model" + # The staged path still resolved to the donor's adapter file. + assert seen["resolves"] + # The staging directory was cleaned up after the load. + assert not os.path.exists(os.path.dirname(seen["model_path"])) + + @pytest.mark.parametrize("relative_path", ["adapter", "model"]) + def test_relative_adapter_path_is_canonicalized( + self, tmp_path, monkeypatch, relative_path + ): + from nemo_rl.models.automodel.setup import _load_initial_lora_adapter + + adapter_dir = tmp_path / relative_path + _write_adapter_checkpoint( + adapter_dir, + keys=( + "base_model.model.layer.lora_A.weight", + "base_model.model.layer.lora_B.weight", + ), + ) + monkeypatch.chdir(tmp_path) + manager = create_autospec(AutomodelCheckpointManager, instance=True) + manager.checkpointer = create_autospec(Checkpointer, instance=True) + seen = {} + + def fake_load_model(*, model, model_path, **kwargs): + seen["model_path"] = model_path + seen["resolves"] = os.path.isfile( + os.path.join(model_path, "adapter_model.safetensors") + ) + + manager.checkpointer.load_model.side_effect = fake_load_model + _load_initial_lora_adapter( + model=_TinyLoraModel(), + checkpoint_manager=manager, + restore_from=relative_path, + lora_cfg=_lora_cfg(), + model_name="tiny-model", + ) + + assert os.path.isabs(seen["model_path"]) + assert seen["resolves"] + + +@pytest.fixture +def _init_gloo_pg(): + """Single-process gloo PG so the real Automodel Checkpointer can run on CPU.""" + if not torch.distributed.is_initialized(): + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29517") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + torch.distributed.init_process_group(backend="gloo", rank=0, world_size=1) + yield + + +class _TwoLinearModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList( + [torch.nn.Linear(4, 4), torch.nn.Linear(4, 1)] + ) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +@pytest.mark.automodel +class TestLoadInitialLoraAdapterEndToEnd: + """Warm start must actually put the donor's weights into the model.""" + + def test_donor_adapter_weights_land_in_model(self, _init_gloo_pg, tmp_path): + from nemo_automodel.components._peft.lora import ( + PeftConfig, + apply_lora_to_linear_modules, + ) + + from nemo_rl.models.automodel.setup import _load_initial_lora_adapter + + peft_config = PeftConfig( + target_modules=[], + match_all_linear=True, + dim=2, + alpha=4, + dropout=0.0, + dropout_position="post", + lora_A_init="xavier", + use_triton=False, + ) + + # Donor: distinctive non-zero adapter weights, so "loaded the donor" + # cannot be confused with "left the zero-init adapters alone". + donor = _TwoLinearModel() + apply_lora_to_linear_modules(donor, peft_config) + for name, param in donor.named_parameters(): + if "lora_" in name: + torch.nn.init.normal_(param, mean=3.0, std=1.0) + donor_lora = { + k: v.clone() for k, v in donor.state_dict().items() if "lora_" in k + } + assert donor_lora + + mesh = torch.distributed.device_mesh.init_device_mesh( + "cpu", (1,), mesh_dim_names=("dp",) + ) + manager = AutomodelCheckpointManager(dp_mesh=mesh, tp_mesh=mesh) + manager.init_checkpointer( + config_updates={"model_save_format": "safetensors", "is_peft": True} + ) + weights_path = str(tmp_path / "step_5" / "policy" / "weights") + manager.save_checkpoint( + model=donor, + weights_path=weights_path, + checkpointing_cfg={ + "enabled": True, + "model_save_format": "safetensors", + "is_peft": True, + }, + lora_enabled=True, + peft_config=peft_config, + ) + + # Fresh run: adapters start at zero, like a cold LoRA init. + model = _TwoLinearModel() + apply_lora_to_linear_modules(model, peft_config) + for name, param in model.named_parameters(): + if "lora_" in name: + param.data.zero_() + + _load_initial_lora_adapter( + model=model, + checkpoint_manager=manager, + restore_from=weights_path, + lora_cfg=_lora_cfg(), + model_name="tiny-model", + ) + + loaded = {k: v for k, v in model.state_dict().items() if "lora_" in k} + assert set(loaded) == set(donor_lora) + for key, expected in donor_lora.items(): + assert torch.allclose(loaded[key], expected), f"{key} was not warm-started" + assert not torch.allclose(loaded[key], torch.zeros_like(loaded[key])) + + def test_donor_covering_fewer_modules_raises(self, _init_gloo_pg, tmp_path): + """A donor targeting fewer modules than the run must fail closed. + + Regression test for the state_dict_adapter escape hatch: the PEFT load + is unconditionally non-strict, so without key validation this donor + would load partially (one layer warm-started, the other silently left + at fresh init) and still print a success message. + """ + from nemo_automodel.components._peft.lora import ( + PeftConfig, + apply_lora_to_linear_modules, + ) + + from nemo_rl.models.automodel.setup import _load_initial_lora_adapter + + donor_peft_config = PeftConfig( + target_modules=["*layers.0*"], + match_all_linear=False, + dim=2, + alpha=4, + dropout=0.0, + dropout_position="post", + lora_A_init="xavier", + use_triton=False, + ) + donor = _TwoLinearModel() + apply_lora_to_linear_modules(donor, donor_peft_config) + donor_lora_names = [n for n, _ in donor.named_parameters() if "lora_" in n] + # Sanity: the donor adapter covers layers.0 only, not layers.1. + assert donor_lora_names + assert all("layers.0" in n for n in donor_lora_names) + + mesh = torch.distributed.device_mesh.init_device_mesh( + "cpu", (1,), mesh_dim_names=("dp",) + ) + manager = AutomodelCheckpointManager(dp_mesh=mesh, tp_mesh=mesh) + manager.init_checkpointer( + config_updates={"model_save_format": "safetensors", "is_peft": True} + ) + weights_path = str(tmp_path / "step_5" / "policy" / "weights") + manager.save_checkpoint( + model=donor, + weights_path=weights_path, + checkpointing_cfg={ + "enabled": True, + "model_save_format": "safetensors", + "is_peft": True, + }, + lora_enabled=True, + peft_config=donor_peft_config, + ) + + run_peft_config = PeftConfig( + target_modules=[], + match_all_linear=True, + dim=2, + alpha=4, + dropout=0.0, + dropout_position="post", + lora_A_init="xavier", + use_triton=False, + ) + model = _TwoLinearModel() + apply_lora_to_linear_modules(model, run_peft_config) + + with pytest.raises(ValueError, match="Missing from donor"): + _load_initial_lora_adapter( + model=model, + checkpoint_manager=manager, + restore_from=weights_path, + lora_cfg=_lora_cfg(), + model_name="tiny-model", + ) diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 8622fac910f..79d79c0b9c0 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -28,10 +28,11 @@ from dataclasses import dataclass, field from types import SimpleNamespace from typing import Any -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import pytest import torch +import yaml @dataclass @@ -3625,3 +3626,659 @@ def test_megatron_internals_have_not_drifted(self): f"DistributedOptimizer no longer references {name!r}; " "_force_sync_optimizer_fp32_from_model's level-1 sync is now a silent no-op." ) + + +@pytest.mark.mcore +class TestPeftWarmStart: + """Tests for the megatron_cfg.peft.restore_from warm-start path.""" + + @staticmethod + def _peft_cfg(**overrides): + """A complete megatron_cfg.peft config for an enabled LoRA run.""" + cfg = { + "enabled": True, + "target_modules": ["linear_qkv"], + "exclude_modules": [], + "dim": 8, + "alpha": 32, + "dropout": 0.0, + "dropout_position": "pre", + "lora_A_init_method": "xavier", + "lora_B_init_method": "zero", + "a2a_experimental": False, + "lora_dtype": None, + } + cfg.update(overrides) + return cfg + + def _make_donor_iter_dir(self, tmp_path, peft_section=None): + """Create a donor iteration directory with a run_config.yaml.""" + iter_dir = tmp_path / "donor" / "iter_0000005" + iter_dir.mkdir(parents=True) + run_config = {"peft": peft_section} if peft_section is not None else {} + with open(iter_dir / "run_config.yaml", "w") as f: + yaml.dump(run_config, f) + return iter_dir + + def test_resolve_iter_dir_directly(self, tmp_path): + from nemo_rl.models.megatron.setup import _resolve_peft_restore_dir + + iter_dir = self._make_donor_iter_dir(tmp_path, {"dim": 8, "alpha": 32}) + assert _resolve_peft_restore_dir(str(iter_dir)) == str(iter_dir) + + def test_resolve_root_with_tracker_file(self, tmp_path): + from nemo_rl.models.megatron.setup import _resolve_peft_restore_dir + + iter_dir = self._make_donor_iter_dir(tmp_path, {"dim": 8, "alpha": 32}) + with open(tmp_path / "donor" / "latest_checkpointed_iteration.txt", "w") as f: + f.write("5") + assert _resolve_peft_restore_dir(str(tmp_path / "donor")) == str(iter_dir) + + def test_resolve_root_with_iter_subdirs_fallback(self, tmp_path): + from nemo_rl.models.megatron.setup import _resolve_peft_restore_dir + + iter_dir = self._make_donor_iter_dir(tmp_path, {"dim": 8, "alpha": 32}) + # No tracker file: falls back to scanning iter_* subdirectories. + assert _resolve_peft_restore_dir(str(tmp_path / "donor")) == str(iter_dir) + + def test_resolve_missing_path_raises(self, tmp_path): + from nemo_rl.models.megatron.setup import _resolve_peft_restore_dir + + with pytest.raises(FileNotFoundError, match="does not exist"): + _resolve_peft_restore_dir(str(tmp_path / "nonexistent")) + + def test_resolve_root_without_iterations_raises(self, tmp_path): + from nemo_rl.models.megatron.setup import _resolve_peft_restore_dir + + empty_root = tmp_path / "donor" + empty_root.mkdir() + with pytest.raises(FileNotFoundError, match="iter_"): + _resolve_peft_restore_dir(str(empty_root)) + + def test_resolve_iter_dir_missing_run_config_raises(self, tmp_path): + from nemo_rl.models.megatron.setup import _resolve_peft_restore_dir + + iter_dir = tmp_path / "donor" / "iter_0000005" + iter_dir.mkdir(parents=True) + with pytest.raises(FileNotFoundError, match="run_config.yaml"): + _resolve_peft_restore_dir(str(tmp_path / "donor")) + + def test_validate_config_match_passes(self, tmp_path): + from nemo_rl.models.megatron.setup import _validate_peft_restore_config + + iter_dir = self._make_donor_iter_dir(tmp_path, self._peft_cfg()) + _validate_peft_restore_config( + str(iter_dir), self._peft_cfg() + ) # should not raise + + def test_validate_config_dim_mismatch_raises(self, tmp_path): + from nemo_rl.models.megatron.setup import _validate_peft_restore_config + + iter_dir = self._make_donor_iter_dir(tmp_path, self._peft_cfg(dim=16)) + with pytest.raises(ValueError, match="dim"): + _validate_peft_restore_config(str(iter_dir), self._peft_cfg()) + + def test_validate_config_alpha_mismatch_raises(self, tmp_path): + from nemo_rl.models.megatron.setup import _validate_peft_restore_config + + iter_dir = self._make_donor_iter_dir(tmp_path, self._peft_cfg(alpha=64)) + with pytest.raises(ValueError, match="alpha"): + _validate_peft_restore_config(str(iter_dir), self._peft_cfg()) + + def test_validate_config_missing_peft_section_raises(self, tmp_path): + from nemo_rl.models.megatron.setup import _validate_peft_restore_config + + iter_dir = self._make_donor_iter_dir(tmp_path, peft_section=None) + with pytest.raises(ValueError, match="no 'peft' section"): + _validate_peft_restore_config(str(iter_dir), self._peft_cfg()) + + def test_validate_config_missing_dim_key_raises(self, tmp_path): + from nemo_rl.models.megatron.setup import _validate_peft_restore_config + + donor_section = self._peft_cfg() + del donor_section["dim"] + iter_dir = self._make_donor_iter_dir(tmp_path, donor_section) + with pytest.raises(ValueError, match="'dim'"): + _validate_peft_restore_config(str(iter_dir), self._peft_cfg()) + + def test_validate_config_target_modules_mismatch_raises(self, tmp_path): + """A superset donor must fail closed, not load silently. + + The adapter-only distributed load discards unrequested donor keys + without any diagnostics (DCP ASSUME_OK_UNEXPECTED skips the mismatch + check), so this comparison is the only thing standing between a + superset donor and a partial warm start. + """ + from nemo_rl.models.megatron.setup import _validate_peft_restore_config + + iter_dir = self._make_donor_iter_dir( + tmp_path, + self._peft_cfg( + target_modules=[ + "linear_qkv", + "linear_proj", + "linear_fc1", + "linear_fc2", + ] + ), + ) + with pytest.raises(ValueError, match="target_modules"): + _validate_peft_restore_config(str(iter_dir), self._peft_cfg()) + + def test_validate_config_exclude_modules_mismatch_raises(self, tmp_path): + from nemo_rl.models.megatron.setup import _validate_peft_restore_config + + iter_dir = self._make_donor_iter_dir( + tmp_path, self._peft_cfg(exclude_modules=["lm_head"]) + ) + with pytest.raises(ValueError, match="exclude_modules"): + _validate_peft_restore_config(str(iter_dir), self._peft_cfg()) + + def test_validate_config_missing_target_modules_key_raises(self, tmp_path): + from nemo_rl.models.megatron.setup import _validate_peft_restore_config + + donor_section = self._peft_cfg() + del donor_section["target_modules"] + iter_dir = self._make_donor_iter_dir(tmp_path, donor_section) + with pytest.raises(ValueError, match="'target_modules'"): + _validate_peft_restore_config(str(iter_dir), self._peft_cfg()) + + def test_validate_config_module_order_does_not_matter(self, tmp_path): + from nemo_rl.models.megatron.setup import _validate_peft_restore_config + + iter_dir = self._make_donor_iter_dir( + tmp_path, self._peft_cfg(target_modules=["linear_qkv", "linear_proj"]) + ) + _validate_peft_restore_config( + str(iter_dir), + self._peft_cfg(target_modules=["linear_proj", "linear_qkv"]), + ) # should not raise + + def test_validate_config_moe_shaping_mismatch_raises(self, tmp_path): + from nemo_rl.models.megatron.setup import _validate_peft_restore_config + + donor_section = self._peft_cfg() + donor_section["normalize_moe_lora"] = True # run uses the default False + iter_dir = self._make_donor_iter_dir(tmp_path, donor_section) + with pytest.raises(ValueError, match="normalize_moe_lora"): + _validate_peft_restore_config(str(iter_dir), self._peft_cfg()) + + def test_validate_config_absent_moe_shaping_keys_pass(self, tmp_path): + """Donors saved before these bridge fields existed are not blocked.""" + from nemo_rl.models.megatron.setup import _validate_peft_restore_config + + iter_dir = self._make_donor_iter_dir(tmp_path, self._peft_cfg()) + _validate_peft_restore_config( + str(iter_dir), self._peft_cfg() + ) # should not raise + + def test_warm_start_hook_loads_adapter_only_and_restores_cfg(self, tmp_path): + from megatron.core.rerun_state_machine import RerunMode + + from nemo_rl.models.megatron.setup import _create_peft_warm_start_hook + + ckpt_cfg = SimpleNamespace( + load="/runs/current/policy/weights", + finetune=True, # left over from the PEFT base-weights hook + load_optim=True, + load_rng=True, + ) + # GlobalState.cfg's setter installs a signal handler based on + # cfg.train, so the namespace needs a train section too. + megatron_cfg = SimpleNamespace( + checkpoint=ckpt_cfg, + train=SimpleNamespace(exit_signal_handler=False), + ) + + from megatron.bridge.training.state import GlobalState + + state = GlobalState() + # In production state.cfg IS megatron_cfg (setup.py assigns it), and + # _load_checkpoint_from_path reads its config off state.cfg. + state.cfg = megatron_cfg + state.train_state.step = 123 + state.train_state.consumed_train_samples = 456 + + class FakeRerunStateMachine: + def __init__(self): + self.mode = RerunMode.VALIDATE_RESULTS + + def get_mode(self): + return self.mode + + def set_mode(self, mode): + self.mode = mode + + rerun_state_machine = FakeRerunStateMachine() + + captured = {} + + def fake_load(**kwargs): + # Capture the checkpoint config as seen mid-load. + captured["load"] = ckpt_cfg.load + captured["finetune"] = ckpt_cfg.finetune + captured["load_optim"] = ckpt_cfg.load_optim + captured["load_rng"] = ckpt_cfg.load_rng + # The loader reads its config off state.cfg: assert the mutation + # is visible there, not just on the local SimpleNamespace. + captured["state_cfg_finetune"] = kwargs["state"].cfg.checkpoint.finetune + captured["load_dir"] = kwargs["load_dir"] + captured["optimizer"] = kwargs["optimizer"] + captured["opt_param_scheduler"] = kwargs["opt_param_scheduler"] + captured["skip_load_to_model_and_opt"] = kwargs[ + "skip_load_to_model_and_opt" + ] + captured["ignore_ckpt_step"] = kwargs["ignore_ckpt_step"] + captured["checkpointing_context"] = kwargs["checkpointing_context"] + captured["rerun_mode"] = rerun_state_machine.get_mode() + return 0, 0 + + model = [MagicMock()] + with ( + patch( + "nemo_rl.models.megatron.setup._load_checkpoint_from_path", + autospec=True, + side_effect=fake_load, + ), + patch( + "nemo_rl.models.megatron.setup.update_num_microbatches" + ) as mock_update_microbatches, + patch( + "nemo_rl.models.megatron.setup.get_rerun_state_machine", + return_value=rerun_state_machine, + ), + ): + hook = _create_peft_warm_start_hook( + megatron_cfg, state, "/donor/iter_0000005" + ) + result = hook(model) + + assert result is model + # The load was routed through the PEFT-resume path: checkpoint.load + # pointed at the donor with finetune=False (adapter-only, non-strict + # load) but optimizer/RNG state disabled. + assert captured["load_dir"] == "/donor/iter_0000005" + assert captured["load"] == "/donor/iter_0000005" + assert captured["finetune"] is False + assert captured["state_cfg_finetune"] is False + assert captured["load_optim"] is False + assert captured["load_rng"] is False + assert captured["optimizer"] is None + assert captured["opt_param_scheduler"] is None + assert captured["skip_load_to_model_and_opt"] is False + assert captured["ignore_ckpt_step"] is True + # An empty context isolates the donor load from the run's own + # dataloader-state directory. + assert captured["checkpointing_context"] == {} + assert captured["rerun_mode"] is RerunMode.DISABLED + # Config restored after the load. + assert ckpt_cfg.load == "/runs/current/policy/weights" + assert ckpt_cfg.finetune is True + assert ckpt_cfg.load_optim is True + assert ckpt_cfg.load_rng is True + assert rerun_state_machine.get_mode() is RerunMode.VALIDATE_RESULTS + # Train state reset so the run starts at step 0. + assert state.train_state.step == 0 + assert state.train_state.consumed_train_samples == 0 + mock_update_microbatches.assert_called_once_with( + consumed_samples=0, verbose=False + ) + + def test_warm_start_hook_restores_cfg_on_load_failure(self): + from megatron.core.rerun_state_machine import RerunMode + + from nemo_rl.models.megatron.setup import _create_peft_warm_start_hook + + ckpt_cfg = SimpleNamespace( + load="/runs/current/policy/weights", + finetune=True, + load_optim=True, + load_rng=True, + ) + megatron_cfg = SimpleNamespace(checkpoint=ckpt_cfg) + + from megatron.bridge.training.state import GlobalState + + state = GlobalState() + rerun_state_machine = MagicMock() + rerun_state_machine.get_mode.return_value = RerunMode.VALIDATE_RESULTS + + with ( + patch( + "nemo_rl.models.megatron.setup._load_checkpoint_from_path", + side_effect=RuntimeError("boom"), + ), + patch("nemo_rl.models.megatron.setup.update_num_microbatches"), + patch( + "nemo_rl.models.megatron.setup.get_rerun_state_machine", + return_value=rerun_state_machine, + ), + ): + hook = _create_peft_warm_start_hook( + megatron_cfg, state, "/donor/iter_0000005" + ) + with pytest.raises(RuntimeError, match="boom"): + hook([MagicMock()]) + + # Even on failure the run's own checkpoint config is restored. + assert ckpt_cfg.load == "/runs/current/policy/weights" + assert ckpt_cfg.finetune is True + assert ckpt_cfg.load_optim is True + assert ckpt_cfg.load_rng is True + assert rerun_state_machine.set_mode.call_args_list == [ + call(RerunMode.DISABLED), + call(RerunMode.VALIDATE_RESULTS), + ] + + def test_restore_from_without_peft_enabled_raises(self): + """restore_from with PEFT disabled must fail loudly, not silently no-op.""" + import nemo_rl.models.megatron.setup as setup_mod + + mock_state = MagicMock() + mock_state.start_time = 0.0 + + megatron_cfg = MagicMock() + megatron_cfg.ft = None + megatron_cfg.model.vocab_size = 32000 + megatron_cfg.model.make_vocab_size_divisible_by = 128 + megatron_cfg.model.tensor_model_parallel_size = 1 + + policy_cfg = { + "megatron_cfg": { + "peft": {"enabled": False, "restore_from": "/donor"}, + } + } + + with ( + patch.object(setup_mod, "GlobalState", return_value=mock_state), + patch.object(setup_mod, "_patch_bridge_signal_handler_for_worker_threads"), + patch.object(setup_mod, "initialize_megatron"), + patch.object(setup_mod, "set_jit_fusion_options"), + patch.object(setup_mod, "init_checkpointing_context"), + patch.object(setup_mod, "build_tokenizer"), + patch("torch.distributed.barrier"), + patch("torch.distributed.all_reduce"), + patch("torch.tensor") as mock_tensor, + ): + mock_tensor_instance = MagicMock() + mock_tensor_instance.item.return_value = 0.0 + mock_tensor.return_value = mock_tensor_instance + + with pytest.raises(ValueError, match="peft.restore_from is set"): + setup_mod.setup_model_and_optimizer( + policy_cfg=policy_cfg, + megatron_cfg=megatron_cfg, + ) + + def test_bridge_peft_resume_filters_adapters_and_loads_non_strict(self): + """Exercise both Bridge behaviors required by the warm-start hook.""" + checkpointing = pytest.importorskip( + "megatron.bridge.training.checkpointing", + reason="requires the mcore extra (Megatron-Bridge)", + ) + from megatron.bridge.training.state import TrainState + + load_dir = "/donor/iter_0000005" + peft = MagicMock() + cfg = SimpleNamespace( + peft=peft, + checkpoint=SimpleNamespace( + load=load_dir, + pretrained_checkpoint="/base-model", + finetune=False, + load_rng=False, + load_optim=False, + ckpt_format="torch_dist", + fully_parallel_save=True, + stage_precision_aware_optimizer_state_on_cpu=False, + load_main_params_from_ckpt=False, + ), + model=SimpleNamespace( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + fp16=False, + bf16=False, + ), + optimizer=SimpleNamespace(use_distributed_optimizer=False), + rng=SimpleNamespace(data_parallel_random_init=False), + ddp=SimpleNamespace(use_megatron_fsdp=False), + ) + state = SimpleNamespace( + cfg=cfg, + train_state=TrainState(), + wandb_logger=None, + mlflow_logger=None, + comet_logger=None, + ) + model = [MagicMock()] + pg_collection = MagicMock() + pg_collection.tp.rank.return_value = 0 + pg_collection.tp.size.return_value = 1 + pg_collection.pp.rank.return_value = 0 + pg_collection.pp.size.return_value = 1 + pg_collection.dp_cp = MagicMock() + + full_state_dict = { + "model": { + "decoder.layers.0.linear.weight": torch.ones(2, 2), + "decoder.layers.0.linear.adapter.lora_A": torch.ones(1, 2), + }, + "checkpoint_version": 3.0, + } + filtered_state_dict = { + "model": { + "decoder.layers.0.linear.adapter.lora_A": torch.ones(1, 2), + }, + "checkpoint_version": 3.0, + } + + def fake_load_base(*args, rank0, **kwargs): + loaded = ( + {"checkpoint_version": 3.0} if rank0 else kwargs["sharded_state_dict"] + ) + return loaded, load_dir, False, None + + run_config = { + "model": { + "tensor_model_parallel_size": 1, + "pipeline_model_parallel_size": 1, + }, + "checkpoint": { + "save_rng": False, + "save_optim": False, + "fully_parallel_save": True, + }, + } + + with ( + patch.object( + checkpointing, + "_load_base_checkpoint", + side_effect=fake_load_base, + ), + patch.object( + checkpointing, + "generate_state_dict", + return_value=full_state_dict, + ), + patch.object( + checkpointing, + "apply_peft_adapter_filter_to_state_dict", + return_value=filtered_state_dict, + ) as mock_filter, + patch.object( + checkpointing.dist_checkpointing, + "load_content_metadata", + return_value={}, + ), + patch.object(checkpointing, "read_run_config", return_value=run_config), + patch.object(checkpointing, "file_exists", return_value=True), + patch.object(checkpointing, "read_train_state", return_value=TrainState()), + patch.object(checkpointing, "update_num_microbatches"), + patch.object(checkpointing, "set_checkpoint_version"), + patch.object(checkpointing, "get_checkpoint_version", return_value=3.0), + patch.object( + checkpointing, + "_get_model_glu_interleave_sizes", + return_value=(None, None), + ), + patch.object(checkpointing, "_load_model_state_dict") as mock_model_load, + patch.object(checkpointing, "is_hf_checkpoint_dir", return_value=False), + patch.object(checkpointing, "unwrap_model", return_value=model), + patch.object(checkpointing.wandb_utils, "on_load_checkpoint_success"), + patch.object(checkpointing.mlflow_utils, "on_load_checkpoint_success"), + patch.object(checkpointing.comet_utils, "on_load_checkpoint_success"), + patch("torch.distributed.is_initialized", return_value=False), + patch("torch.cuda.empty_cache"), + ): + checkpointing._load_checkpoint_from_path( + load_dir=load_dir, + state=state, + model=model, + optimizer=None, + opt_param_scheduler=None, + strict=True, + checkpointing_context={}, + skip_load_to_model_and_opt=False, + ignore_ckpt_step=True, + pg_collection=pg_collection, + ) + + mock_filter.assert_called_once_with(full_state_dict, peft) + mock_model_load.assert_called_once_with( + model[0], filtered_state_dict["model"], False + ) + + def _run_policy_setup(self, tmp_path, *, resume_exists): + """Run setup_model_and_optimizer with PEFT warm start configured. + + Returns the _create_peft_warm_start_hook mock so the caller can assert + whether the hook was composed into the policy's pre-wrap hooks. + """ + import nemo_rl.models.megatron.setup as setup_mod + + donor_iter_dir = self._make_donor_iter_dir(tmp_path, self._peft_cfg()) + + mock_state = MagicMock() + mock_state.start_time = 0.0 + + megatron_cfg = MagicMock() + megatron_cfg.ft = None + megatron_cfg.model.vocab_size = 32000 + megatron_cfg.model.make_vocab_size_divisible_by = 128 + megatron_cfg.model.tensor_model_parallel_size = 1 + megatron_cfg.ddp.overlap_param_gather = False + megatron_cfg.checkpoint.load = ( + "/runs/current/policy/weights" if resume_exists else None + ) + megatron_cfg.checkpoint.pretrained_checkpoint = None + + policy_cfg = { + "megatron_cfg": { + "freeze_moe_router": False, + "peft": self._peft_cfg(restore_from=str(donor_iter_dir)), + } + } + + mock_model_chunk = MagicMock() + mock_optimizer = MagicMock() + mock_scheduler = MagicMock() + mock_tensor_instance = MagicMock() + mock_tensor_instance.item.return_value = 0.0 + + with ( + patch.object(setup_mod, "GlobalState", return_value=mock_state), + patch.object(setup_mod, "_patch_bridge_signal_handler_for_worker_threads"), + patch.object(setup_mod, "ProcessGroupCollection"), + patch.object(setup_mod, "initialize_megatron"), + patch.object(setup_mod, "set_jit_fusion_options"), + patch.object(setup_mod, "init_checkpointing_context"), + patch.object(setup_mod, "build_tokenizer"), + patch.object(setup_mod, "get_model", return_value=[mock_model_chunk]), + patch.object( + setup_mod, + "setup_optimizer", + return_value=(mock_optimizer, mock_scheduler), + ), + patch.object( + setup_mod, + "_create_peft_pre_wrap_hook", + return_value=lambda model: model, + ), + patch.object(setup_mod, "_create_peft_warm_start_hook") as mock_hook, + patch.object(setup_mod, "checkpoint_exists", return_value=resume_exists), + patch.object(setup_mod, "load_checkpoint"), + patch.object(setup_mod, "get_attached_draft_model", return_value=None), + patch("torch.distributed.barrier"), + patch("torch.distributed.all_reduce"), + patch("torch.tensor", return_value=mock_tensor_instance), + ): + setup_mod.setup_model_and_optimizer( + policy_cfg=policy_cfg, + megatron_cfg=megatron_cfg, + ) + return mock_hook + + def test_policy_warm_start_hook_appended_on_fresh_run(self, tmp_path): + """No resume checkpoint -> the policy warm-start hook is composed.""" + mock_hook = self._run_policy_setup(tmp_path, resume_exists=False) + mock_hook.assert_called_once() + + def test_policy_warm_start_hook_skipped_on_resume(self, tmp_path): + """Resume checkpoint already carries this run's adapters -> no hook.""" + mock_hook = self._run_policy_setup(tmp_path, resume_exists=True) + mock_hook.assert_not_called() + + def test_reference_warm_start_hook_appended_unconditionally(self, tmp_path): + """The reference model warm-starts even when the policy resumes. + + This asymmetry is the KL-anchoring story: on a resume the policy's + adapters come from the resume checkpoint, but the reference must stay + anchored to the initial (donor warm-started) policy. A regression that + drops the reference hook produces a silently wrong KL anchor rather + than a crash. + """ + import nemo_rl.models.megatron.setup as setup_mod + + donor_iter_dir = self._make_donor_iter_dir(tmp_path, self._peft_cfg()) + + megatron_cfg = MagicMock() + megatron_cfg.dist.use_torch_fsdp2 = False + + mock_model = MagicMock() + mock_model.state_dict.return_value = { + "layer1.weight": torch.tensor([1.0, 2.0]), + } + + config = { + "megatron_cfg": { + "freeze_moe_router": False, + "peft": self._peft_cfg(restore_from=str(donor_iter_dir)), + } + } + + with ( + patch.object(setup_mod, "ProcessGroupCollection"), + patch.object(setup_mod, "init_checkpointing_context"), + patch.object(setup_mod, "GlobalState", return_value=MagicMock()), + patch.object(setup_mod, "get_model", return_value=[mock_model]), + patch.object(setup_mod, "checkpoint_exists", return_value=False), + patch.object(setup_mod, "clear_global_router_replay_instances"), + patch.object(setup_mod, "load_checkpoint"), + patch.object( + setup_mod, + "_create_peft_pre_wrap_hook", + return_value=lambda model: model, + ), + patch.object(setup_mod, "_create_peft_warm_start_hook") as mock_hook, + patch.object(setup_mod, "HAVE_FSDP2", False), + ): + setup_mod.setup_reference_model_state( + config=config, + megatron_cfg=megatron_cfg, + pretrained_path="/path/to/pretrained", + ) + + mock_hook.assert_called_once() + # The hook receives the resolved donor iteration directory. + assert mock_hook.call_args.args[2] == str(donor_iter_dir) diff --git a/tests/unit/reference_configs/dpo.yaml b/tests/unit/reference_configs/dpo.yaml index 3f9bd65b83f..d7357b62cc1 100755 --- a/tests/unit/reference_configs/dpo.yaml +++ b/tests/unit/reference_configs/dpo.yaml @@ -80,6 +80,7 @@ policy: dropout_position: "post" # Where to apply dropout: "pre" (before LoRA) or "post" (after LoRA) lora_A_init: "xavier" # Initialization method for LoRA A matrix: "xavier" or "uniform" use_triton: true # Use Triton-optimized kernels for LoRA (faster but requires flash-attn). Disable when tensor_parallel_size > 1 + restore_from: null # Warm start: path to a PEFT adapter checkpoint (adapter_model.safetensors + adapter_config.json) to initialize this run's LoRA weights from, e.g. a previous SFT run's step_*/policy/weights dynamic_batching: enabled: false @@ -173,6 +174,7 @@ policy: lora_B_init_method: "zero" a2a_experimental: false lora_dtype: null + restore_from: null # Warm start: path to a native Megatron PEFT checkpoint (iter_XXXXXXX dir or its root) to initialize this run's LoRA weights from, e.g. a previous SFT run optimizer: optimizer: "adam" diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index c51186f9aea..ceb0e5d82aa 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -178,6 +178,7 @@ policy: dropout_position: "post" # Where to apply dropout: "pre" (before LoRA) or "post" (after LoRA) lora_A_init: "xavier" # Initialization method for LoRA A matrix: "xavier" or "uniform" use_triton: true # Use Triton-optimized kernels for LoRA (faster but requires flash-attn). Disable when tensor_parallel_size > 1 + restore_from: null # Warm start: path to a PEFT adapter checkpoint (adapter_model.safetensors + adapter_config.json) to initialize this run's LoRA weights from, e.g. a previous SFT run's step_*/policy/weights megatron_cfg: enabled: false @@ -253,6 +254,7 @@ policy: lora_B_init_method: "zero" a2a_experimental: false lora_dtype: None + restore_from: null # Warm start: path to a native Megatron PEFT checkpoint (iter_XXXXXXX dir or its root) to initialize this run's LoRA weights from, e.g. a previous SFT run optimizer: optimizer: "adam" diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 71a9251a559..f9197246b0d 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -130,6 +130,7 @@ policy: dropout_position: "post" lora_A_init: "xavier" use_triton: true + restore_from: null # Warm start: path to a PEFT adapter checkpoint (adapter_model.safetensors + adapter_config.json) to initialize this run's LoRA weights from, e.g. a previous SFT run's step_*/policy/weights megatron_cfg: enabled: true @@ -310,6 +311,7 @@ value: dropout_position: "post" lora_A_init: "xavier" use_triton: true + restore_from: null # Warm start: path to a PEFT adapter checkpoint (adapter_model.safetensors + adapter_config.json) to initialize this run's LoRA weights from, e.g. a previous SFT run's step_*/policy/weights megatron_cfg: enabled: true diff --git a/tests/unit/reference_configs/sft.yaml b/tests/unit/reference_configs/sft.yaml index 3924de619e7..6d5707cd0e9 100644 --- a/tests/unit/reference_configs/sft.yaml +++ b/tests/unit/reference_configs/sft.yaml @@ -71,6 +71,7 @@ policy: dropout_position: "post" # Where to apply dropout: "pre" (before LoRA) or "post" (after LoRA) lora_A_init: "xavier" # Initialization method for LoRA A matrix: "xavier" or "uniform" use_triton: true # Use Triton-optimized kernels for LoRA (faster but requires flash-attn). Disable when tensor_parallel_size > 1 + restore_from: null # Warm start: path to a PEFT adapter checkpoint (adapter_model.safetensors + adapter_config.json) to initialize this run's LoRA weights from, e.g. a previous SFT run's step_*/policy/weights dynamic_batching: enabled: false @@ -159,6 +160,7 @@ policy: lora_B_init_method: "zero" a2a_experimental: false lora_dtype: None + restore_from: null # Warm start: path to a native Megatron PEFT checkpoint (iter_XXXXXXX dir or its root) to initialize this run's LoRA weights from, e.g. a previous SFT run optimizer: