fix(test): unblock Kimi-Linear HF reference and gate MiniMax routed parity - #3644
Closed
yuhezhang-ai wants to merge 4 commits into
Closed
Conversation
19 tasks
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 14:50
6f904f5 to
cde5b6e
Compare
Contributor
Author
|
Restructured: the Kimi-Linear fixes are independent of the MiniMax work, so this stacked PR is replaced by a standalone PR based directly on #3567's branch (same head branch, now containing only the Kimi commits). |
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
1 participant
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 does this PR do?
Unblocks the Kimi-Linear vanilla-HF reference (AMINT-288) and closes the MiniMax M2.7 cross-framework residual with measured router evidence (AMINT-286). Stacked on #3643; includes a merge of #3635 (GLM parity diagnostics), whose router/shape machinery this PR extends.
With this stack,
kimi_linear_48b_a3b_hellaswagpasses all six phases on the standard profile for the first time (pipeline 64182587: source HF vs AM mean KL 0.0021 / cosine 0.9994; export reload 0.0023; resume exact-state), andminimax_m2.7_hellaswag_loragates its shared 128-token prefix with full-2K metrics and router evidence kept diagnostic.Changelog
Kimi-Linear reference compatibility (the checkpoint's remote code targets transformers 4.57.1 / fla-core 0.4.0; the CI image runs transformers 5.12.1 / fla-core 0.4.2; no in-tree implementation exists):
fix(transformers): alias the relocatedOutputRecorderback intotransformers.utils.genericinapply_cache_compatibility_patches(production compat patcher; same class object; no-op when present).auto_mapwhen AutoModel's_CUSTOM_CONFIG_REGISTRATIONShijacks AutoConfig to a nemo-owned class, and drop nemo component configs forwarded from config-based recipes.input_embedskwarg.fused_kda_gate(g, A, head_k_dim, g_bias=...)calls onto fla-core 0.4.2's renamed/reshaped API.MiniMax M2.7 residual (router precision + measured tie sensitivity):
fix(moe): score sigmoid routing in fp32 by default, mirroring the softmax branch and the HF references'sigmoid(logits.float()); final weights keep the existing cast to input dtype. Production behavior change for sigmoid-router MoE models (MiniMax, GLM, DeepSeek families): bf16 sigmoid quantized scores at ~2e-3, flipping 7-13% of top-k selections per layer against the fp32 reference on M2.7's real bias lattices.mlp.gate.e_score_correction_biasnever matched in-treemlp.e_score_correction_bias, silently casting the reference's router bias to bf16 (215-242/255 adjacent expert pairs collide exactly in bf16 at bias magnitude 8.6).MiniMaxM2TopKRoutercapture;model_familyderived fromconfig.model_type; grouplessn_expert_groups=0captures normalized) and enable router + [128, 512] shape diagnostics on the recipe.Validation
🤖 Generated with Claude Code