Repository navigation
Add the Kimi K3 backend primitives - #94
Merged
Merged
Conversation
Three hooks the Kimi K3 plugin in miles relies on, ported from the K3 review fork (yueming-yuan/k3-megatron-review@ebe3b9d7f): - TransformerConfig.gated_activation_func: an activation that consumes the packed [gate | linear] GLU input itself (K3's situ), used by MLP and TEGroupedMLP in place of activation_func(gate) * linear. - TransformerConfig.moe_latent_use_norm: RMSNorm the combined routed output in latent space before fc2_latent_proj (K3's routed_expert_norm). - Parameter.sum_gradients_across_tp_domain: sum (not average) TP-replicated partial gradients in finalize_model_grads, for K3's expert-shared LoRA factors whose gradient is a partial sum over the TP shards.
Zhichenzzz
approved these changes
Sep 22, 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
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.
Three hooks the Kimi K3 plugin in miles (radixark/miles#1825) relies on, ported from the K3 review fork (yueming-yuan/k3-megatron-review@ebe3b9d7f) onto
miles-main. They were never opened against this repo; the K3 image carried the fork instead.TransformerConfig.gated_activation_func: an activation that consumes the packed[gate | linear]GLU input itself (K3's situ).MLPandTEGroupedMLPcall it in place ofactivation_func(gate) * linearwhen set;Nonekeeps today's behavior.TransformerConfig.moe_latent_use_norm: RMSNorm the combined routed output in latent space beforefc2_latent_proj(K3'srouted_expert_norm, a real checkpoint tensor).Parameter.sum_gradients_across_tp_domain: sum, not average, TP-replicated partial gradients infinalize_model_grads, for K3's expert-shared LoRA factors whose gradient is a partial sum over the TP shards.Without the norm hook a K3 checkpoint loads with the latent norm silently dropped;
--check-weight-update-equalon the 4-layer K3 run caught it asrouted_expert_norm.weightnever reaching the engine.Validation
yueming/k3-lora-on-sglang-miles) with--check-weight-update-equal: base weights compare equal after the sync, two rollouts train.sum_gradients_across_tp_domain(the smoke run is TP=1); no unit test — the fork's only test covered an unrelated SGD change that is already onmiles-main.