fix(test): unblock the Kimi-Linear vanilla-HF parity reference - #3659
Merged
Conversation
yuhezhang-ai
marked this pull request as ready for review
August 25, 2026 15:49
yuhezhang-ai
requested review from
a team,
HuiyingLi,
akoumpa,
athitten and
snowmanwwg
as code owners
August 25, 2026 20:34
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>
…rations 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>
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>
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>
yuhezhang-ai
force-pushed
the
yuhez/fix/amint288-kimi-outputrecorder-shim
branch
from
August 25, 2026 20:56
cde5b6e to
61f2845
Compare
Contributor
Author
|
/ok to test 61f2845 |
akoumpa
approved these changes
Aug 25, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Kimi-Linear's checkpoint remote code targets a pre-v5 Transformers API, so the checkpoint-robustness harness could never load its vanilla-HF reference under the CI stack (Transformers 5.12): phases 0 and 3 crashed before producing a single parity number (AMINT-288). This PR makes the reference loadable through four independent compatibility fixes, none of which changes model math.
Replaces the stacked #3644 (closed): the Kimi work is independent of the MiniMax parity work, so it now sits directly on #3567's branch.
Changelog
apply_cache_compatibility_patchesaliasesOutputRecorderback intotransformers.utils.genericfor pre-v5 remote code (library-side shim, benefits any stale remote-code model).auto_mapwhen AutoModel'sCONFIG_MAPPINGregistrations hijackAutoConfig(preserves a load-time FP8dequantizerequest).input_embedskwarg in remotecreate_causal_maskcalls (renamed toinputs_embedsin 5.x,cache_positiondropped).flafused_kda_gate(g_flat, A, head_k_dim, g_bias=...)calls to the fla-core 0.4.2 signature (reshape +dt_bias; rejects non-defaultbeta/threshold).Validation
Final proof on this exact head (61f2845, rebased onto post-test(checkpoint): expand parity metrics and phase coverage #3567 main): nemo-ci pipeline 64546415 — GREEN, all six phases (source 0.00212 / cos 0.99943, export reload 0.00226 / 0.99941, AM reload 0.00232, every self-repeat exact).
Scoped nemo-ci pipeline 64182587, exact job
kimi_linear_48b_a3b_hellaswag: all six phases pass on the standard profile — first-ever Kimi HF parity numbers: source mean KL 0.00212 / p95 0.00849 / cos 0.99943; export reload 0.00234 / 0.00901 / 0.99941; AM reload 0.00230; resume exact.Unit tests: 149 passed (
test_checkpoint_robustness_hf_kwargs.py,test_transformers_utils.py).Note: this PR and #3674 add functions in adjacent regions of the same harness file; whichever merges second needs a trivial context-only rebase.
🤖 Generated with Claude Code