add_axk1 (#46867)
* add_axk1
* fix TP bug of load balancing bias term
* change default value of n_routed_experts and max_position_embeddings
* address review comments
* change model type to snake case, and rebase
* refactor AXK1 to inherit from deepseek_v3 and modernize config
- inherit from deepseek_v3 (AXK1Attention / AXK1TopkRouter / AXK1MoE / AXK1Experts /
AXK1DecoderLayer / AXK1Model / AXK1ForCausalLM), keeping the A.X-K1-specific
post_mlp_layernorm (applied to the MoE block output) as the only override
- config: rewrite in @strict style inheriting DeepseekV3Config (class-level typed
attributes, automatic docstring), moved into modular_axk1.py; drop pretraining_tp /
moe_layer_freq / the rope-float override; inherit base_model_ep_plan
- experts AXK1NaiveMoe(MixtralExperts) -> AXK1Experts(DeepseekV3Experts) (same stacked
layout / keys); routing now lives in the router
- _keys_to_ignore_on_load_unexpected = ["inv_freq"]
- add docs (model_doc/axk1.md + toctree)
- bump copyrights to 2026
Verified on the light checkpoint: loads with the only newly-initialized key being the
(zero) e_score_correction_bias the checkpoint omits, and greedy generation is coherent
("대한민국의 수도는" -> " 서울이다.").
* address second review round: youtu-style tests, docs cleanup, inherit decoder layer
* use Expectations class and add batched logits integration test
* fixup repo
* idk why but go
* Add random small model test
* test fixes
* oops that was from my local
* add integration test
* fixup integration tests
* style
---------
Co-authored-by: minsangkim <kmswin1@sktelecom.com>
Co-authored-by: vasqu <antonprogamer@gmail.com>
