fix(config): use explicit KDA beta-gate exclusion paths

#8
by iCavan - opened

This replaces the relative attention.b_proj exclusion with the 63 complete KDA beta-gate paths from the published checkpoint index, for example model.layers.0.attention.b_proj.

The BF16/FP8 policy is unchanged. Full checkpoint paths work with the existing checkpoint-name conversion in Transformers and remain valid SGLang module prefixes. This avoids a special conversion rule for relative attention.* names without broadening the exclusion to the FP8 MLA q_b_proj or kv_b_proj layers. All other config fields, model code, and weights are unchanged.

I compared the full checkpoint at 93e96ab91e80 with the candidate configuration on the same four GPUs, separately within each engine:

Engine Check Result
Transformers native 8 fixed-prefix steps, cached and cacheless, all 157,184 logits per step; greedy generation Bit-exact logits in both modes (max absolute difference 0), identical generated tokens
Transformers native Checkpoint dtype and loaded-module audit Identical audits: 639 FP8 Linear, 73 BF16 Linear, and 68 FP8 expert modules; no loading or dtype errors
SGLang native English, repeated English without flushing, and Chinese after flushing; 8 generated tokens each Identical text, token IDs, and all returned logprob fields, including top-20 and selected-token scores

Transformers used a3a371bf4711, eager attention, the PyTorch KDA fallback, and trust_remote_code=False. Its baseline retains the relative-name conversion rule; the candidate removes that rule and uses these full paths.

SGLang used 8cd3c494a3aa, TP4/EP4, Triton FP8 GEMM/MoE, FlashInfer KDA prefill and Triton KDA decode, with CUDA graphs disabled. It used model_impl="sglang" and trust_remote_code=True to read the official config class. Speculative decoding/MTP was disabled.

These are baseline-versus-candidate checks within each engine, not a claim of numerical equivalence between engines. The repeated short English request did not hit the prefix cache.

Thanks for taking a look!

A little more context on the main motivation: this follows the Transformers reviewer’s recommendation during the native Ling 3.1 integration. The preference was to express these exclusions as full checkpoint module paths, rather than add another name-mapping rule in Transformers just to handle the relative attention.b_proj entry.

Transformers already converts full paths such as model.layers.0.attention.b_proj to model.layers.0.self_attn.b_proj through its existing checkpoint-name conversion. The relative entry starts directly with attention., so supporting it in the proposed integration would otherwise need an additional start-of-name mapping. Using full paths in the checkpoint config lets us reuse the established conversion and remove that extra special case.

So this is primarily about keeping the upstream Transformers integration simpler and easier to maintain, not changing the intended quantization policy. The same 63 KDA beta gates remain BF16, and the real-checkpoint comparisons above confirmed unchanged results within both Transformers and SGLang under the tested settings.

Thanks for considering this checkpoint-side change!

inclusionAI org

Engine coverage. Only an unreleased Transformers commit (trust_remote_code=False) and an SGLang commit were tested. Please also check a released Transformers version with trust_remote_code=True, since that’s what most users run today. Also check whichever other engines the model card advertises, such as vLLM. A bare attention.b_proj entry may be matched differently by each engine’s ignore-list logic, and some apply their own name remapping to it.

Thanks for requesting the additional compatibility checks. I tested released Transformers and vLLM against the full checkpoint, keeping the weights and model code identical and changing only config.json as proposed in this PR.

Engine Before PR #8: current main configuration With PR #8
Transformers 4.56.2, trust_remote_code=True Loads and generates successfully, with the expected FP8/BF16 module selection. Same module precision and identical Chinese/English generation. Full-vocabulary logits are bit-exact, with max absolute difference = 0, across eight cached and eight cacheless decoding steps.
vLLM 0.32.0, native model implementation The relative attention.b_proj exclusion is not remapped correctly. All 63 KDA beta projections are included in fused FP8 qkvb_proj modules, and generation is visibly corrupted—for example, TrastrasterM0 hardest Award. All 63 KDA beta projections correctly remain separate BF16 modules, while MLA q_b_proj / kv_b_proj remain FP8. Chinese/English generation completes with coherent output and finite returned logprobs.

The vLLM results used TP4/EP4, Triton MoE/MLA backends, eager execution, and two prompts with eight generated tokens each. MTP and CUDA graphs were not enabled. These are focused compatibility checks, not a full model-quality benchmark or a claim of numerical equivalence between engines.

The tested main configuration already breaks correct generation in vLLM 0.32.0; this PR fixes that existing regression rather than introducing a new quantization policy. The fix also preserves the tested Transformers behavior exactly.

Given the impact on users loading the published configuration, could we prioritize merging this fix? Thanks again for the careful review.

michael-qiu changed pull request status to merged

Sign up or log in to comment