SIGN IN SIGN UP

fix(test): unblock the Kimi-Linear vanilla-HF parity reference (#3659)

* fix(transformers): alias relocated OutputRecorder for pre-v5 remote code

Transformers v5.x moved OutputRecorder from transformers.utils.generic to
transformers.utils.output_capturing without a back-compat alias. Pre-v5
remote-code checkpoints (Kimi-Linear-48B-A3B, MiniMax-M2.*) import it from
the old location for their auxiliary router-logit recorders, so the whole
modeling module fails at import (AMINT-288): with no in-tree kimi_linear in
Transformers, the checkpoint-robustness vanilla-HF source and export-reload
phases cannot construct a reference at all.

Alias the relocated class (re-exported by transformers.modeling_utils in
v5.x) back into transformers.utils.generic inside
apply_cache_compatibility_patches, which the checkpoint-robustness harness
already applies before every vanilla-HF reference load and which owns the
existing v5-removal shims (SlidingWindowCache, Cache.get_usable_length).
The alias binds the identical class object and is a no-op on Transformers
versions that still export it from utils.generic.

Verified against the cached Kimi-Linear remote code under transformers
5.12.1: module import now proceeds past OutputRecorder to the model's
fla-core dependency (present in the CI image). Kimi's remote code uses no
rope-init APIs, so it does not hit the separate v5 ROPE_INIT_FUNCTIONS
'default'-key removal that still breaks MiniMax-M2 remote code (whose
reference intentionally stays on the repaired in-tree path).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* fix(test): resolve vanilla-HF reference configs past AutoModel registrations

nemo_automodel registers component config classes into Transformers'
CONFIG_MAPPING (_CUSTOM_CONFIG_REGISTRATIONS), and a locally registered
model_type beats checkpoint remote code even under trust_remote_code=True.
Inside the checkpoint-robustness process this hijacks the vanilla-HF
reference for remote-code-only models: AutoConfig resolves Kimi-Linear to
the AutoModel-owned config class while the model class comes from the
checkpoint's auto_map, and from_pretrained rejects the pair with a
config_class mismatch in both the source-load and export-reload phases
(AMINT-288, pipeline 64093591 after the OutputRecorder shim).

Re-resolve the checkpoint's own config class from its auto_map whenever the
reference config resolves to a nemo_automodel-owned class, pass the faithful
config explicitly so from_pretrained's internal resolution cannot re-select
the registered class, preserve a load-time FP8 dequantize request across the
replacement, and drop AutoModel component configs that config-based recipes
forward through the reference load kwargs. All paths are no-ops for models
whose configs already resolve to Transformers-owned classes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* fix(test): rename legacy input_embeds kwarg for remote masking calls

With the OutputRecorder shim and reference-config resolution in place, the
Kimi-Linear vanilla-HF reference now loads and fails one step later: its
pre-v5 remote code calls create_causal_mask(input_embeds=...), which
Transformers v5.x renamed to inputs_embeds (pipeline 64160712, job
408783019). Extend _patch_remote_masking_api_compatibility to rename the
legacy keyword when the installed function only accepts the new name, next
to the existing removed-cache_position handling. Functions that accept the
legacy API are still left unwrapped.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* fix(test): translate legacy fla fused_kda_gate calls for remote code

With the masking rename fixed, Kimi-Linear's vanilla-HF reference reaches
its KDA layers and fails on the fifth compatibility break (pipeline
64176083, job 408883841): the checkpoint's remote code calls the pre-0.4.2
fla API fused_kda_gate(g, A_log, head_k_dim, g_bias=...) with a flat gate
projection, while fla-core 0.4.2 renamed the keyword to dt_bias, dropped
head_k_dim, and expects g pre-reshaped to [..., heads, head_k_dim].

Add _patch_remote_fla_api_compatibility next to the masking patch at both
vanilla-HF reference sites: translate legacy calls (reshape + keyword
rename), reject unsupported non-default beta/threshold overrides loudly,
pass new-style calls through untouched, and leave an installed fla that
still accepts g_bias unpatched.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

---------

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Y
Yuhe Zhang committed
504967bb6ffe20d02b4f70e37aab3c21b1cde198
Parent: f432e6d
Committed by GitHub <noreply@github.com> on 8/26/2026, 12:09:44 AM