Skip to content

Add the Kimi K3 backend primitives - #94

Merged
yueming-yuan merged 1 commit into
miles-mainfrom
yueming/k3-backend-primitives
Sep 23, 2026
Merged

yueming-yuan merged 1 commit into
miles-mainfrom
yueming/k3-backend-primitives

Conversation

@yueming-yuan

Copy link
Copy Markdown
Collaborator

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). MLP and TEGroupedMLP call it in place of activation_func(gate) * linear when set; None keeps today's behavior.
  • TransformerConfig.moe_latent_use_norm: RMSNorm the combined routed output in latent space before fc2_latent_proj (K3's routed_expert_norm, a real checkpoint tensor).
  • 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.

Without the norm hook a K3 checkpoint loads with the latent norm silently dropped; --check-weight-update-equal on the 4-layer K3 run caught it as routed_expert_norm.weight never reaching the engine.

Validation

  • 4-layer Kimi K3 LoRA RL smoke run (colocated, 8xH200, Kimi-K3 full-model and LoRA RL support miles#1825 + sgl-project/sglang yueming/k3-lora-on-sglang-miles) with --check-weight-update-equal: base weights compare equal after the sync, two rollouts train.
  • Not executed: a TP>1 run exercising 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 on miles-main.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants