🚀 Fine-tune Kimi K3 with NeMo AutoModel — native 2.8T MoE support at scale

[akoumpa]

Fine-tune Moonshot AI's 2.8T-parameter multimodal MoE with native PyTorch distributed training

🤗 Model · 🔥 Recipe · 📖 Model coverage · 🛠️ Implementation

NeMo AutoModel now supports fine-tuning moonshotai/Kimi-K3, Moonshot AI's open-weight, native multimodal Mixture-of-Experts model.

The integration loads the released Hugging Face checkpoint through native model, configuration, tokenizer, and state-dict implementations. Scaling is configured with FSDP2, Expert Parallelism, Pipeline Parallelism, and Context Parallelism.

2.8T	104B	896 → 16	1M
total parameters	activated parameters	experts → selected/token	context window

Why Kimi K3 is interesting

• Hybrid attention: 69 Kimi Delta Attention layers and 24 Gated MLA layers balance long-context efficiency with expressive full attention.
• Attention Residuals: a separate residual stream is accumulated across transformer blocks, requiring architecture-aware pipeline partitioning.
• Extreme MoE sparsity: only 16 of 896 experts are selected for each token.
• Native multimodality: the architecture supports both text and visual inputs.
• Quantization-aware weights: the released checkpoint uses MXFP4 weights with MXFP8 activations.

For the complete architecture and evaluation results, see Moonshot AI's Kimi K3 repository and technical report.

What's ready in NeMo AutoModel

✅ Native KimiK3ForCausalLM and KimiK3ForConditionalGeneration implementations
✅ Native tokenizer, encoding, and MXFP4 checkpoint loading
✅ FSDP2, Expert Parallelism, Pipeline Parallelism, and Context Parallelism
✅ K3-aware pipeline stages that preserve complete Attention Residual blocks
✅ Residual-accumulator handoff between pipeline stages
✅ A full-parameter HellaSwag SFT recipe

Kimi K3 currently supports DP, PP, CP, and EP. Tensor Parallelism is not yet supported.

Native loading—without remote model code

The text recipe makes two choices explicit:

model:
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
  pretrained_model_name_or_path: moonshotai/Kimi-K3

  # Optional but explicit: use AutoModel's in-tree implementation rather
  # than executing Python supplied by the model repository.
  trust_remote_code: false

  # Select the text-only backbone. The checkpoint's top-level architecture
  # is KimiK3ForConditionalGeneration, which also builds the vision stack.
  config:
    architectures: [KimiK3ForCausalLM]

For multimodal Kimi K3, use NeMoAutoModelForImageTextToText and leave the checkpoint's KimiK3ForConditionalGeneration architecture unchanged.

Scale-out recipe and validation

The provided recipe uses FSDP2 with EP32 × PP8:

distributed:
  strategy: fsdp2
  tp_size: 1
  cp_size: 1
  pp_size: 8
  ep_size: 32
  activation_checkpointing: true

  pipeline:
    pp_schedule: 1f1b
    pp_microbatch_size: 1

Note

Kimi K3 is a genuinely large model. The included full-parameter, 93-layer recipe was validated at 256-GPU scale. This is a scale-out recipe, not a single-node quick start.

Model	Full 93-layer Kimi K3
Objective	Full-parameter HellaSwag SFT
Parallelism	EP32 × PP8
Run length	100 steps
Training loss	2.0011 → 1.4307
Context Parallelism	CP1/CP2 parity validated

_{Full-parameter Kimi K3 fine-tuning on HellaSwag.}

Use Kimi K3 from another Python library

Third-party training systems can construct the distributed model directly with DistributedSetup. With PP_SIZE=1, the result is an already-sharded PyTorch module with the normal forward interface. The host library can also use setup.mesh_context to access the resolved device and MoE meshes.

Multi-node Python example

import os

import torch

from nemo_automodel import NeMoAutoModelForCausalLM
from nemo_automodel.components.distributed import (
    DistributedSetup,
    FSDP2Config,
    MoEParallelizerConfig,
    ParallelismSizes,
    PipelineConfig,
    initialize_distributed,
)
from nemo_automodel.components.loss.masked_ce import MaskedCrossEntropy
from nemo_automodel.components.models.common import BackendConfig


pp_size = int(os.getenv("PP_SIZE", "1"))
cp_size = int(os.getenv("CP_SIZE", "1"))
ep_size = int(os.getenv("EP_SIZE", "1"))
local_batch_size = int(os.getenv("LOCAL_BATCH_SIZE", "8"))

# Reuses the process group if the host library initialized it already.
dist = initialize_distributed("nccl", timeout_minutes=120)

pipeline = None
if pp_size > 1:
    pipeline = PipelineConfig(
        pp_schedule="1f1b",
        pp_microbatch_size=1,
        pp_batch_size=local_batch_size,
        patch_inner_model=False,
        patch_causal_lm_model=False,
        scale_grads_in_schedule=False,
        loss_fn=MaskedCrossEntropy(),
    )

setup = DistributedSetup.build(
    strategy=FSDP2Config(sequence_parallel=False),
    parallelism_sizes=ParallelismSizes(
        tp_size=1,  # Kimi K3 currently supports DP/PP/CP/EP, but not TP.
        pp_size=pp_size,
        cp_size=cp_size,
        ep_size=ep_size,
    ),
    pipeline_config=pipeline,
    moe_parallel_config=(
        MoEParallelizerConfig(
            reshard_after_forward=False,
            wrap_outer_model=False,
        )
        if ep_size > 1
        else None
    ),
    activation_checkpointing=True,
    world_size=dist.world_size,
    timeout_minutes=120,
)

model = NeMoAutoModelForCausalLM.from_pretrained(
    "moonshotai/Kimi-K3",
    distributed_setup=setup,
    torch_dtype=torch.bfloat16,
    trust_remote_code=False,
    config={"architectures": ["KimiK3ForCausalLM"]},
    kda_mode="chunk",
    use_liger_kernel=False,
    use_sdpa_patching=False,
    backend=BackendConfig(
        attn="eager",
        linear="torch",
        rms_norm="torch_fp32",
        experts="torch_mm",
        dispatcher="torch",
        rope_fusion=False,
        enable_fsdp_optimizations=True,
    ),
    # Direct library users load through the loader. AutoModel recipes instead
    # coordinate base loading through their Checkpointer.
    load_base_model=True,
)

if dist.is_main:
    mesh = setup.mesh_context
    print(
        f"Loaded Kimi K3: DP={mesh.dp_size}, PP={mesh.pp_size}, "
        f"CP={mesh.cp_size}, EP={mesh.ep_size}"
    )

# Pass `model` and `setup.mesh_context` to the host library.
# The model is already distributed; do not wrap it again with DDP or FSDP.

Launch it with any system that provides the standard torchrun environment:

PP_SIZE=8 CP_SIZE=1 EP_SIZE=32 LOCAL_BATCH_SIZE=8 \
torchrun \
  --nnodes=32 \
  --nproc-per-node=8 \
  --node-rank="$NODE_RANK" \
  --master-addr="$MASTER_ADDR" \
  --master-port=29500 \
  load_kimi_k3.py

Data parallelism is inferred as:

DP = WORLD_SIZE / (TP × PP × CP)

Example topologies on 256 GPUs:

Use case	DP	PP	CP	EP	Returned interface
Third-party trainer + EP	256	1	1	32	Sharded PyTorch module
Validated full-model recipe	32	8	1	32	AutoPipeline
Long-context PP + CP	16	8	2	32	AutoPipeline

EP must divide DP × CP.

Important

With PP_SIZE>1, AutoModel returns an AutoPipeline. The host library must drive model.step(...) or integrate the stage modules exposed through model.parts; a generic Hugging Face-style trainer cannot treat it as a normal single-process model.

Start here

• 🔥 Kimi K3 HellaSwag recipe
• 📖 Kimi K3 model-coverage page
• 🛠️ NeMo AutoModel implementation
• 🤗 Kimi K3 checkpoint
• 📈 Validated training run

This first recipe focuses on text SFT. We would especially love feedback on the next workloads you want to see: multimodal fine-tuning, long-context training, additional datasets, or further kernel optimization.

Huge thanks to @HuiyingLi for bringing Kimi K3 support to NeMo AutoModel!