transformers

b6d5084f - Add AXK2 from SKT (#47528)

Add AXK2 from SKT (#47528) * Add A.X K2, the SKT's flagship LLM * Add fused gate * fix code from bug of the attention_mask is None case * delete redundant comment * change several default config values * refactor AXK2 to inherit from deepseek_v3/v32 and address review feedback - indexer: inherit DeepseekV32Indexer; key cache moved out of the module into the shared Cache (past_key_values.update_indexer backed by DynamicIndexedLayer/StaticIndexedLayer, dispatched via config.layer_types=deepseek_sparse_attention); stateful prefill reset removed - MoE: router inherits DeepseekV32TopkRouter with a non-grouped top-k forward (released model ships n_group=None); routing now lives inside the router; experts inherit the stacked-3D implementation (@use_experts_implementation, expert-parallel friendly) - attention: inherit DeepseekV3Attention; single canonical path — q LoRA always on, interleaved RoPE imported from deepseek_v3, output gate always applied via linear_gate - fused checkpoint layout is handled at load time by the new SplitFusedAttentionGate weight converter (splits the block-diagonal q_b_proj into q_b_proj + linear_gate, re-fused on save); the split is lossless and removes all runtime fused/unfused branches - SGA no longer forces eager attention: the indexer top-k is folded into an additive sparse mask consumed by eager and SDPA; _supports_sdpa=T flash/flex disabled (same as deepseek_v32) - config: mlp_layer_types / layer_types are computed in the config; rope_interleave / gated_norm / attention_output_gate toggles removed (single released architecture) * change model type to snake case, and rebase * refactor AXK2 to inherit from deepseek_v3.2 and address review feedback - config: rewrite in @strict style inheriting DeepseekV32Config (class-level typed attributes, __post_init__ for derived fields, automatic docstring); drop pretraining_tp / moe_layer_freq / the rope-float override; attribute_map = {"num_local_experts": "n_routed_experts"}; inherit base_model_ep_plan - inherit the DSA stack from deepseek_v3.2 (AXK2Indexer / AXK2Attention / AXK2DecoderLayer / AXK2PreTrainedModel / AXK2Model / AXK2ForCausalLM), overriding only the A.X-K2 deltas - gated RMSNorm reuses CLIPMLP (bias-free, SiLU, gated_norm_rank); gate weights bridged W_down/W_up -> mlp.fc1/fc2 by the converter - indexer key LayerNorm eps set explicitly to 1e-5; indexer reads the pre-norm q - attention output gate renamed linear_gate -> g_proj; validate q_lora_rank is set - rename fused-checkpoint conversion op SplitFusedAttentionGate -> SplitFusedMLAGate (+ FuseMLAGate reverse) - _keys_to_ignore_on_load_unexpected = ["inv_freq"]; drop the legacy "AXK2" -> "axk2" conversion alias (model unreleased; aligned on model_type="axk2") - add docs (model_doc/axk2.md + toctree) and tests/models/axk2 - bump copyrights to 2026 Verified on the release-candidate light checkpoint: loads with no missing/unexpected keys and greedy generation is identical to the pre-refactor implementation ("대한민국의 수도는" -> " 서울입니다. 서울은 대한민국의 정치, 경제, 문화의 중심지로"). * address second review round: tests, MiniMax router, config cleanup - tests: add tests/models/axk2 (AXK2ModelTest(CausalLMModelTester) fast tests + slow integration) and SplitFusedMLAGate / FuseMLAGate unit tests in tests/utils/test_core_model_loading.py - router: inherit MiniMaxM2TopKRouter (non-grouped sigmoid) instead of DeepseekV32TopkRouter, keeping e_score_correction_bias on the router and fp32 routing (forward math unchanged) - config: drop n_group / topk_group / first_k_dense_replace via AttributeError(); mlp_layer_types is derived from the legacy first_k_dense_replace / moe_layer_freq kwargs in __post_init__ - attention: drop the runtime q_lora_rank check (q_lora_rank: int is enforced by @strict at config construction) - decoder: post_attention_layernorm written as a ternary - docs: drop dtype/device_map from the examples (auto under the hood) * fix converter round-trip + skip DSA-incompatible generation tests - SplitFusedMLAGate: pass an already-unfused q_b_proj through unchanged (only the fused vLLM layout is split), so save/load round-trips work and g_proj loads from its own key; add a passthrough unit test in tests/utils/test_core_model_loading.py - skip static-cache / compile generation tests (the SGA indexer mask uses a dynamic cache length), mirroring glm_moe_dsa Full fast suite on GPU: 113 passed, 2173 subtests passed, 146 skipped (deepspeed deselected — this env lacks mpi4py; they run in CI). * address final review: config super() trick, op refactor, test cleanup - config: call PreTrainedConfig.__post_init__/validate_architecture directly instead of super() so only A.X-K2's own derivations run (drops the duplicated DeepSeek-V3.2 __post_init__ block); validate_architecture now requires a positive q_lora_rank - router (AXK2TopkRouter): unpack `_, topk_indices = torch.topk(...)` and comment the only deltas over MiniMax-M2 (optional top-k norm + routed scaling) - core_model_loading: SplitFusedMLAGate/FuseMLAGate follow PermuteForRope (self.config attribute + separate _split/_fuse application methods) - unquote AXK2Config type annotations - tests: drop n_group/topk_group/first_k_dense_replace from the tester (derive mlp_layer_types); remove vestigial pipeline_model_mapping/fx_compatible/ test_torchscript; move AXK2_CHECKPOINT to an integration-test class attribute; wrap expected text in Expectations; replace test_sdpa_matches_eager_generation with test_logits; fold the fused-gate op test into TestConvertAndLoadStateDict via a dummy-model load path (no standalone TestCase). Re-checked the padding skips on GPU — they still fail (DSA selection flips past index_topk), kept skipped. * self-contained integration tests pushed, ready for your local CI; logits Expectations recorded on A100 so a GPU-specific key may need adding * use a small randomized checkpoint for integration tests * address review: g_proj TP plan entry, TP-shard-safe gate converter, adopt expand_kv, cleanup - add `layers.*.self_attn.g_proj: colwise` to the TP plan (the output gate emits one scalar per attention-output column, so it shards with the heads) - fix SplitFusedMLAGate/FuseMLAGate to derive the head count from the tensor shape: under tensor parallelism the conversion runs on each rank's head-aligned row shard, and the config-derived head count crashed on it; add a row-shard unit test - verify TP end to end: 2-GPU DistributedConfig(tp_size=2) smoke test vs a single-GPU reference — max |logit diff| 1.2e-7, greedy tokens identical - adopt the `expand_kv` split from #47460; bit-identical logits and greedy tokens on a seeded tiny model before/after - regenerate from modular: picks up yarn_apply_mscale and the causal_mask_mapping API from the rebased parents (generated files were stale) - drop the two converter-workaround comments flagged in review - tester: mlp_layer_types = ["dense", "sparse"] (shared tester default is 2 layers) - spell out abbreviated variable names in the MLA-gate loading tests * fix output_capturing dynamo NameError + point integration tests at hf-internal-testing/tiny-axk2 * support fused attention gate (attn_gate_fused) + grouped routing for A.X-K2 releases The released A.X-K2 checkpoints (e.g. A.X-K2-Preview) ship the attention output gate fused into `q_b_proj` (vLLM block-diagonal layout) and use DeepSeek-V3-style grouped expert routing. - attn_gate_fused config flag: keep `q_b_proj` fused and split the activation in the forward (matches the reference/vLLM kernel), instead of splitting the weight at load. Required for the fp8 releases, whose 128-block scales cannot be split along the per-head query/gate boundary; the fused matrix is 128-aligned and loads directly. The unfused SplitFusedMLAGate path stays for bf16 checkpoints. - filter_checkpoint_conversions hook: a model can drop registered conversions that don't apply to its config (the mapping is keyed by model_type and can't see config). A.X-K2 drops the SplitFusedMLAGate converter when attn_gate_fused, so fp8 loading doesn't trip the merge-only scale-folding / many-to-many arity guard. - n_group/topk_group grouped routing (0.0 group masking, matching the SKT reference), with validate_architecture pairing/divisibility checks; conversion mapping stacks per-expert projections. - tests: grouped router vs reference, fused-gate layout + save/load roundtrip, hub-layout roundtrip. Verified bit-identical to the SKT reference implementation on the Preview architecture (grouped + fused gate + SGA sparse); Light non-grouped parity unchanged. Native fp8 GEMM needs a compute-capability >= 8.9 GPU to exercise (A100 force-dequantizes fp8 to bf16), so the fp8 inference path should be confirmed on CI hardware. * keep the attention output gate fused (q_gate_proj) instead of splitting at load Per review: the split-at-load path doesn't make sense given the released checkpoints are fused (and fp8, which can't be split), so drop it and always use the fused variation under a clearer name. - AXK2Attention keeps the query up-projection and output gate fused in a single `q_gate_proj` (doubled input -> per-head [query | gate]) and splits the activation in the forward, matching the vLLM-style released layout. Renamed from `q_b_proj` since it also carries the gate. - drop SplitFusedMLAGate/FuseMLAGate and the two-path `attn_gate_fused` flag entirely; the checkpoint's `q_b_proj` is a plain 1->1 rename to `q_gate_proj` on load (fp8-safe: no weight split, and the prefix renames the `weight_scale_inv` companion too). Reverses to `q_b_proj` on save. - AXK2Attention is standalone (nn.Module) rather than inheriting DeepseekV32Attention, so there is no leftover `q_b_proj` from the parent's conditional q-projection. - TP: the per-head activation split uses `-1` (not the global head count) so colwise sharding of `q_gate_proj` (each rank holds whole heads) stays correct. Verified: fused forward is bit-identical to the SKT reference on the Preview architecture (grouped routing + fused gate + SGA); fp8 checkpoint loads and matches bf16 within quant error; 2-GPU tp_plan="auto" matches the single-GPU reference (max |logit diff| 1.2e-7). Light parity unchanged. * fix repo and fixup modular + clean tests * rm this * repo fixes * add integration test (depends on upstream main axk1 merge) * fix * try fixing --------- Co-authored-by: minsangkim <kmswin1@sktelecom.com>