diff --git a/nemo_automodel/_transformers/utils.py b/nemo_automodel/_transformers/utils.py index 8ee86322b7..7436455c5a 100644 --- a/nemo_automodel/_transformers/utils.py +++ b/nemo_automodel/_transformers/utils.py @@ -261,9 +261,22 @@ def _to_legacy_cache(self): DynamicCache.to_legacy_cache = _to_legacy_cache + # OutputRecorder moved from transformers.utils.generic to + # transformers.utils.output_capturing in transformers v5.x. Pre-v5 + # remote-code models (e.g. Kimi-Linear and MiniMax-M2 checkpoints) import + # it from the old location for their auxiliary router-logit recorders and + # otherwise fail at module import. Alias the relocated class back; v5.x + # re-exports it from transformers.modeling_utils. + import transformers.modeling_utils as mu + import transformers.utils.generic as generic_utils + + if not hasattr(generic_utils, "OutputRecorder"): + _output_recorder = getattr(mu, "OutputRecorder", None) + if _output_recorder is not None: + generic_utils.OutputRecorder = _output_recorder + # _tied_weights_keys changed from list to dict in transformers v5.x. # Patch post_init to auto-convert list -> dict for remote-code models. - import transformers.modeling_utils as mu if not getattr(mu.PreTrainedModel.post_init, "_nemo_tied_keys_patched", False): _orig_post_init = mu.PreTrainedModel.post_init diff --git a/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_llm.py b/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_llm.py index ab019310e0..63ecc5d122 100644 --- a/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_llm.py +++ b/tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_llm.py @@ -343,6 +343,75 @@ def _load_hf_fp8_dequantized_config( return config +def _is_nemo_owned_config(config) -> bool: + """Return True when a config object is an AutoModel component config.""" + return type(config).__module__.startswith("nemo_automodel") + + +def _replace_nemo_owned_reference_config( + config, + pretrained_model_name_or_path: str | Path, + *, + trust_remote_code: bool, + revision: str | None = None, + token: str | bool | None = None, +): + """Re-resolve a vanilla-reference config hijacked by AutoModel registrations. + + nemo_automodel registers its component config classes into Transformers' + ``CONFIG_MAPPING`` (``_CUSTOM_CONFIG_REGISTRATIONS``), and a locally + registered ``model_type`` wins over checkpoint remote code even with + ``trust_remote_code=True``. Inside the harness process, AutoConfig then + resolves checkpoints such as Kimi-Linear to an AutoModel-owned class while + the model class still comes from the checkpoint's ``auto_map``, and + ``from_pretrained`` rejects the pair with a ``config_class`` mismatch + (AMINT-288). Resolve the checkpoint's own config class from its + ``auto_map`` instead, preserving a load-time FP8 ``dequantize`` request. + + Args: + config: Config resolved for the vanilla reference load. + pretrained_model_name_or_path: Checkpoint the reference loads from. + trust_remote_code: Whether the reference load trusts remote code. + revision: Optional checkpoint revision. + token: Optional Hub token. + + Returns: + Tuple of the faithful config and whether a replacement happened. + """ + if not trust_remote_code or not _is_nemo_owned_config(config): + return config, False + auto_map = getattr(config, "auto_map", None) or {} + class_reference = auto_map.get("AutoConfig") + if not class_reference: + return config, False + + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + load_kwargs: dict[str, str | bool] = { + "local_files_only": os.environ.get("HF_HUB_OFFLINE", "0") == "1", + } + if revision is not None: + load_kwargs["revision"] = revision + if token is not None: + load_kwargs["token"] = token + config_cls = get_class_from_dynamic_module(class_reference, pretrained_model_name_or_path, **load_kwargs) + replacement = config_cls.from_pretrained(pretrained_model_name_or_path, **load_kwargs) + + original_quantization = getattr(config, "quantization_config", None) + requested_dequantize = ( + original_quantization.get("dequantize") + if isinstance(original_quantization, dict) + else getattr(original_quantization, "dequantize", None) + ) + if requested_dequantize: + replacement_quantization = getattr(replacement, "quantization_config", None) + if isinstance(replacement_quantization, dict): + replacement.quantization_config = {**replacement_quantization, "dequantize": True} + elif replacement_quantization is not None: + replacement_quantization.dequantize = True + return replacement, True + + def _dequantize_hf_fp8_weights_in_place(model, output_dtype: torch.dtype) -> int: """Dequantize native per-tensor HF FP8 modules without their runtime kernel. @@ -447,7 +516,7 @@ def _peft_adapter_load_kwargs(hf_kwargs: dict[str, object]) -> dict[str, object] def _patch_remote_masking_api_compatibility() -> None: - """Allow remote model code to pass masking kwargs removed by Transformers.""" + """Adapt remote model code to masking kwargs removed or renamed by Transformers.""" import transformers.masking_utils as masking_utils for function_name in ("create_causal_mask", "create_sliding_window_causal_mask"): @@ -455,22 +524,100 @@ def _patch_remote_masking_api_compatibility() -> None: if getattr(mask_function, "_nemo_removed_kwargs_patched", False): continue parameters = inspect.signature(mask_function).parameters.values() - accepts_cache_position = any( - parameter.name == "cache_position" or parameter.kind is inspect.Parameter.VAR_KEYWORD - for parameter in parameters + parameter_names = {parameter.name for parameter in parameters} + accepts_var_keyword = any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters) + drop_cache_position = "cache_position" not in parameter_names and not accepts_var_keyword + # Transformers v5.x renamed ``input_embeds`` to ``inputs_embeds``; + # pre-v5 remote code (e.g. Kimi-Linear) still passes the old keyword. + rename_input_embeds = ( + "input_embeds" not in parameter_names and "inputs_embeds" in parameter_names and not accepts_var_keyword ) - if accepts_cache_position: + if not drop_cache_position and not rename_input_embeds: continue @wraps(mask_function) - def compatible_mask_function(*args, _mask_function=mask_function, **kwargs): - kwargs.pop("cache_position", None) + def compatible_mask_function( + *args, + _mask_function=mask_function, + _drop_cache_position=drop_cache_position, + _rename_input_embeds=rename_input_embeds, + **kwargs, + ): + if _drop_cache_position: + kwargs.pop("cache_position", None) + if _rename_input_embeds and "input_embeds" in kwargs: + kwargs["inputs_embeds"] = kwargs.pop("input_embeds") return _mask_function(*args, **kwargs) compatible_mask_function._nemo_removed_kwargs_patched = True # type: ignore[attr-defined] setattr(masking_utils, function_name, compatible_mask_function) +def _patch_remote_fla_api_compatibility() -> None: + """Adapt remote model code to the renamed fla-core KDA gate API. + + Kimi-Linear's pre-0.4.2 remote code calls + ``fused_kda_gate(g, A_log, head_k_dim, g_bias=...)`` with a flat + ``g`` of shape ``[..., heads * head_k_dim]`` that the old kernel reshaped + internally. fla-core 0.4.2 renamed the API to + ``fused_kda_gate(g, A_log, dt_bias=None, lower_bound=None, ...)`` and + expects ``g`` pre-reshaped to ``[..., heads, head_k_dim]``. Translate + legacy calls when the installed function no longer accepts the old form; + an installed fla that still accepts ``g_bias`` is left untouched. + """ + try: + import fla.ops.kda as kda_ops + import fla.ops.kda.gate as kda_gate + except ImportError: + return + + gate_function = kda_gate.fused_kda_gate + if getattr(gate_function, "_nemo_legacy_kda_gate_patched", False): + return + parameter_names = set(inspect.signature(gate_function).parameters) + if "g_bias" in parameter_names or "dt_bias" not in parameter_names: + return + + @wraps(gate_function) + def compatible_fused_kda_gate(g, A_log, *args, _gate_function=gate_function, **kwargs): + """Translate a legacy KDA gate call onto the renamed fla-core API. + + Args: + g: Gate projection. Legacy callers pass a flat Tensor of shape + [..., heads * head_k_dim] together with a positional + ``head_k_dim``; new-style callers pass [..., heads, head_k_dim]. + A_log: Per-head log-decay Tensor of shape [heads]. + *args: A leading int is the legacy positional ``head_k_dim``; a + following tensor is the legacy positional ``g_bias``. + **kwargs: Legacy ``g_bias``/``beta``/``threshold`` keywords are + translated or rejected; everything else passes through. + + Returns: + Gate Tensor of shape [..., heads, head_k_dim] from the new API. + """ + if args and isinstance(args[0], int): + head_k_dim = args[0] + remaining = list(args[1:]) + if remaining: + kwargs.setdefault("g_bias", remaining.pop(0)) + if remaining: + raise TypeError("Unexpected extra positional arguments for legacy fused_kda_gate call") + g_bias = kwargs.pop("g_bias", None) + beta = kwargs.pop("beta", 1.0) + threshold = kwargs.pop("threshold", 20.0) + if beta != 1.0 or threshold != 20.0: + raise TypeError( + "Legacy fused_kda_gate beta/threshold overrides are not supported by the installed fla API" + ) + return _gate_function(g.view(*g.shape[:-1], -1, head_k_dim), A_log, dt_bias=g_bias, **kwargs) + return _gate_function(g, A_log, *args, **kwargs) + + compatible_fused_kda_gate._nemo_legacy_kda_gate_patched = True # type: ignore[attr-defined] + kda_gate.fused_kda_gate = compatible_fused_kda_gate + if getattr(kda_ops, "fused_kda_gate", None) is gate_function: + kda_ops.fused_kda_gate = compatible_fused_kda_gate + + def _rss_gb() -> float: """Current RSS in GB from /proc/self/statm.""" page_size = os.sysconf("SC_PAGE_SIZE") @@ -985,6 +1132,13 @@ def _hf_source_load_kwargs( "trust_remote_code", } hf_kwargs = {k: v for k, v in model_kwargs.items() if k in hf_allowed_keys} + recipe_config = hf_kwargs.get("config") + if recipe_config is not None and _is_nemo_owned_config(recipe_config): + # AutoModel component configs are never valid for a vanilla-HF + # reference; remote and in-tree model classes both reject them with a + # config_class mismatch. Let the reference resolve the checkpoint's + # own config instead. + del hf_kwargs["config"] hf_kwargs["torch_dtype"] = source_dtype hf_kwargs["trust_remote_code"] = trust_remote_code hf_kwargs["local_files_only"] = os.environ.get("HF_HUB_OFFLINE", "0") == "1" @@ -1429,6 +1583,7 @@ def _prepare_source_load_reference_rank0( apply_cache_compatibility_patches() _patch_remote_masking_api_compatibility() + _patch_remote_fla_api_compatibility() model_kwargs = _model_kwargs_from_config(cfg.model) original_pretrained_path = _model_pretrained_path(cfg.model, model_kwargs) @@ -1489,6 +1644,17 @@ def _prepare_source_load_reference_rank0( revision=hf_kwargs.get("revision"), token=hf_kwargs.get("token"), ) + hf_config, replaced_reference_config = _replace_nemo_owned_reference_config( + hf_config, + original_pretrained_path, + trust_remote_code=hf_kwargs["trust_remote_code"], + revision=hf_kwargs.get("revision"), + token=hf_kwargs.get("token"), + ) + if replaced_reference_config: + # Pass the faithful config explicitly so from_pretrained's internal + # AutoConfig resolution cannot re-select the AutoModel-owned class. + hf_kwargs["config"] = hf_config model_load_context = _hf_model_load_context( trust_remote_code=trust_remote_code, @@ -1928,6 +2094,7 @@ def _run_vanilla_hf_reload( # still carry Transformers-v4 list-form ``_tied_weights_keys``. apply_cache_compatibility_patches() _patch_remote_masking_api_compatibility() + _patch_remote_fla_api_compatibility() _, ckpt_step_dir, consolidated_dir = _checkpoint_paths(cfg) is_peft = hasattr(cfg, "peft") model_kwargs = _model_kwargs_from_config(cfg.model) @@ -2007,6 +2174,17 @@ def _run_vanilla_hf_reload( revision=model_kwargs.get("revision"), token=model_kwargs.get("token"), ) + hf_config, replaced_reference_config = _replace_nemo_owned_reference_config( + hf_config, + config_path, + trust_remote_code=trust_remote_code, + revision=model_kwargs.get("revision"), + token=model_kwargs.get("token"), + ) + if replaced_reference_config: + # Pass the faithful config explicitly so from_pretrained's internal + # AutoConfig resolution cannot re-select the AutoModel-owned class. + hf_kwargs["config"] = hf_config # Load the reference model straight onto the target GPU. Materialising a # 14B checkpoint on CPU and then ``.to(device)`` costs ~50-225s, and that # rank-0-only stall trips the NCCL watchdog while the other ranks idle at diff --git a/tests/unit_tests/_transformers/test_transformers_utils.py b/tests/unit_tests/_transformers/test_transformers_utils.py index a549d032ba..59ebb95480 100644 --- a/tests/unit_tests/_transformers/test_transformers_utils.py +++ b/tests/unit_tests/_transformers/test_transformers_utils.py @@ -327,6 +327,23 @@ def test_patches_sliding_window_cache(self): assert hasattr(cu, "SlidingWindowCache") + def test_restores_output_recorder_in_utils_generic(self, monkeypatch): + """Pre-v5 remote code imports OutputRecorder from transformers.utils.generic. + + Transformers v5.x moved the class to transformers.utils.output_capturing, + which breaks remote-code checkpoints (Kimi-Linear, MiniMax-M2) at module + import. The patch must alias the relocated class back to the old location. + """ + import transformers.modeling_utils as modeling_utils + import transformers.utils.generic as generic_utils + + monkeypatch.delattr(generic_utils, "OutputRecorder", raising=False) + apply_cache_compatibility_patches() + + assert generic_utils.OutputRecorder is modeling_utils.OutputRecorder + # The exact import form used by the remote modeling files must work. + from transformers.utils.generic import OutputRecorder # noqa: F401 + def test_patches_cache_get_usable_length(self): """Cache.get_usable_length should exist after patching.""" apply_cache_compatibility_patches() diff --git a/tests/unit_tests/ci_tests/test_checkpoint_robustness_hf_kwargs.py b/tests/unit_tests/ci_tests/test_checkpoint_robustness_hf_kwargs.py index eebdc92302..b0604fc132 100644 --- a/tests/unit_tests/ci_tests/test_checkpoint_robustness_hf_kwargs.py +++ b/tests/unit_tests/ci_tests/test_checkpoint_robustness_hf_kwargs.py @@ -19,7 +19,7 @@ import pytest import torch -from transformers import AutoModelForCausalLM +from transformers import AutoModelForCausalLM, PretrainedConfig from tests.functional_tests.checkpoint_robustness.test_checkpoint_robustness_biencoder import ( _extract_custom_args as _extract_biencoder_custom_args, @@ -49,6 +49,7 @@ _LogitParityPolicy, _model_pretrained_path, _normalize_peft_no_split_modules, + _patch_remote_fla_api_compatibility, _patch_remote_masking_api_compatibility, _peft_adapter_load_kwargs, _post_load_dequant_max_memory, @@ -56,6 +57,7 @@ _raise_distributed_failure, _record_deferred_failure, _repeatability_policy, + _replace_nemo_owned_reference_config, _resolve_hf_attn_implementation, _resolve_hf_model_class, _run_process_isolated_checkpoint_phase, @@ -309,11 +311,13 @@ def create_mask(config, inputs_embeds, attention_mask, past_key_values, position _patch_remote_masking_api_compatibility() for function_name in ("create_causal_mask", "create_sliding_window_causal_mask"): + # Pre-v5 remote code passes the removed ``cache_position`` and the + # renamed ``input_embeds`` keywords (e.g. Kimi-Linear). result = getattr(masking_utils, function_name)( - "config", - "inputs", - "attention", - "cache", + config="config", + input_embeds="inputs", + attention_mask="attention", + past_key_values="cache", position_ids="positions", cache_position="removed-argument", ) @@ -328,7 +332,7 @@ def create_mask(config, inputs_embeds, attention_mask, past_key_values, position def test_remote_masking_api_compatibility_preserves_supported_api(monkeypatch): import transformers.masking_utils as masking_utils - def create_mask(config, inputs_embeds, attention_mask, past_key_values, cache_position=None): + def create_mask(config, input_embeds, attention_mask, past_key_values, cache_position=None): return cache_position monkeypatch.setattr(masking_utils, "create_causal_mask", create_mask) @@ -1707,3 +1711,178 @@ def test_record_deferred_failure_preserves_all_comparison_failures(): _record_deferred_failure(failures, "Phase 4 HF reload parity", "HF parity failed") assert failures == ["Phase 4 HF reload parity:\nHF parity failed"] + + + +class _FakeNemoOwnedConfig(PretrainedConfig): + """Stands in for an AutoModel component config registered into AutoConfig.""" + + model_type = "stub_reference" + + +# The helper discriminates on the class's owning package, not its bases. +_FakeNemoOwnedConfig.__module__ = "nemo_automodel.components.models.test_only.config" + +_STUB_AUTO_MAP = {"AutoConfig": "configuration_stub.StubReferenceConfig"} + + +def _write_stub_remote_checkpoint(tmp_path, *, quantization_config=None): + """Write a minimal remote-code checkpoint dir exposing its own config class.""" + (tmp_path / "configuration_stub.py").write_text( + "from transformers import PretrainedConfig\n" + "\n" + "\n" + "class StubReferenceConfig(PretrainedConfig):\n" + ' model_type = "stub_reference"\n' + ) + config = {"model_type": "stub_reference", "auto_map": _STUB_AUTO_MAP, "hidden_size": 8} + if quantization_config is not None: + config["quantization_config"] = quantization_config + (tmp_path / "config.json").write_text(json.dumps(config)) + return tmp_path + + +def test_replace_nemo_owned_reference_config_resolves_checkpoint_auto_map(tmp_path): + ckpt = _write_stub_remote_checkpoint(tmp_path) + hijacked = _FakeNemoOwnedConfig(auto_map=dict(_STUB_AUTO_MAP)) + + replaced, did_replace = _replace_nemo_owned_reference_config(hijacked, ckpt, trust_remote_code=True) + + assert did_replace is True + assert type(replaced).__name__ == "StubReferenceConfig" + assert type(replaced).__module__.startswith("transformers_modules") + + +def test_replace_nemo_owned_reference_config_preserves_dequantize_request(tmp_path): + ckpt = _write_stub_remote_checkpoint(tmp_path, quantization_config={"quant_method": "fp8"}) + hijacked = _FakeNemoOwnedConfig( + auto_map=dict(_STUB_AUTO_MAP), + quantization_config={"quant_method": "fp8", "dequantize": True}, + ) + + replaced, did_replace = _replace_nemo_owned_reference_config(hijacked, ckpt, trust_remote_code=True) + + assert did_replace is True + assert replaced.quantization_config["dequantize"] is True + + +def test_replace_nemo_owned_reference_config_noop_cases(tmp_path): + from transformers import AutoConfig + + hf_config = AutoConfig.for_model("llama", vocab_size=64, hidden_size=32, num_hidden_layers=1) + assert _replace_nemo_owned_reference_config(hf_config, tmp_path, trust_remote_code=True) == (hf_config, False) + + no_auto_map = _FakeNemoOwnedConfig() + assert _replace_nemo_owned_reference_config(no_auto_map, tmp_path, trust_remote_code=True) == (no_auto_map, False) + + untrusted = _FakeNemoOwnedConfig(auto_map=dict(_STUB_AUTO_MAP)) + assert _replace_nemo_owned_reference_config(untrusted, tmp_path, trust_remote_code=False) == (untrusted, False) + + +def test_hf_source_load_kwargs_drops_nemo_owned_recipe_config(): + with patch( + "transformers.PretrainedConfig.get_config_dict", + return_value=({"model_type": "unknown_remote_model"}, {}), + ): + hf_kwargs = _hf_source_load_kwargs( + {"config": _FakeNemoOwnedConfig(), "attn_implementation": "eager"}, + pretrained_model_name_or_path="model-path", + source_dtype=torch.bfloat16, + trust_remote_code=True, + experts_implementation=None, + hf_model_cls=AutoModelForCausalLM, + device=torch.device("cpu"), + hf_device_map_auto=False, + ) + + assert "config" not in hf_kwargs + + +def test_hf_source_load_kwargs_keeps_hf_recipe_config(): + from transformers import AutoConfig + + hf_config = AutoConfig.for_model("llama", vocab_size=64, hidden_size=32, num_hidden_layers=1) + with patch( + "transformers.PretrainedConfig.get_config_dict", + return_value=({"model_type": "unknown_remote_model"}, {}), + ): + hf_kwargs = _hf_source_load_kwargs( + {"config": hf_config, "attn_implementation": "eager"}, + pretrained_model_name_or_path="model-path", + source_dtype=torch.bfloat16, + trust_remote_code=True, + experts_implementation=None, + hf_model_cls=AutoModelForCausalLM, + device=torch.device("cpu"), + hf_device_map_auto=False, + ) + + assert hf_kwargs["config"] is hf_config + + +def _install_fake_fla(monkeypatch, gate_function): + """Register a minimal fake fla.ops.kda[.gate] module tree exposing gate_function.""" + import sys + from types import ModuleType + + fla = ModuleType("fla") + ops = ModuleType("fla.ops") + kda = ModuleType("fla.ops.kda") + gate = ModuleType("fla.ops.kda.gate") + gate.fused_kda_gate = gate_function + kda.fused_kda_gate = gate_function + kda.gate = gate + ops.kda = kda + fla.ops = ops + for name, module in (("fla", fla), ("fla.ops", ops), ("fla.ops.kda", kda), ("fla.ops.kda.gate", gate)): + monkeypatch.setitem(sys.modules, name, module) + return kda, gate + + +def test_remote_fla_api_compatibility_translates_legacy_kda_gate(monkeypatch): + """Legacy fused_kda_gate(g, A, head_k_dim, g_bias=...) calls must reach the renamed API. + + The fake installed API mirrors fla-core 0.4.2: g arrives pre-reshaped to + [..., heads, head_k_dim] and the bias keyword is dt_bias. + """ + calls = {} + + def fused_kda_gate(g, A_log, dt_bias=None, lower_bound=None, output_dtype=torch.float32): + calls["g_shape"] = tuple(g.shape) + calls["dt_bias"] = dt_bias + return g + + kda, gate = _install_fake_fla(monkeypatch, fused_kda_gate) + _patch_remote_fla_api_compatibility() + + # Legacy call: flat g of shape [batch, sequence, heads * head_k_dim]. + g = torch.zeros(2, 3, 8) + bias = torch.ones(8) + gate.fused_kda_gate(g, torch.zeros(2), 4, g_bias=bias) + assert calls["g_shape"] == (2, 3, 2, 4) + assert calls["dt_bias"] is bias + + # New-style calls pass through untouched. + gate.fused_kda_gate(torch.zeros(2, 3, 2, 4), torch.zeros(2), dt_bias=None) + assert calls["g_shape"] == (2, 3, 2, 4) + assert calls["dt_bias"] is None + + # The package-level re-export is patched consistently, and re-patching no-ops. + assert kda.fused_kda_gate is gate.fused_kda_gate + patched = gate.fused_kda_gate + _patch_remote_fla_api_compatibility() + assert gate.fused_kda_gate is patched + + with pytest.raises(TypeError, match="beta/threshold"): + gate.fused_kda_gate(g, torch.zeros(2), 4, g_bias=bias, beta=2.0) + + +def test_remote_fla_api_compatibility_preserves_legacy_capable_api(monkeypatch): + def fused_kda_gate(g, A, head_k_dim, g_bias=None, beta=1.0, threshold=20.0): + return g + + kda, gate = _install_fake_fla(monkeypatch, fused_kda_gate) + _patch_remote_fla_api_compatibility() + + assert gate.fused_kda_gate is fused_kda_gate + assert kda.fused_kda_gate is fused_kda_gate