feat(kimi-k3): add expert LoRA and tune full SFT

examples/llm_finetune/kimi/k3_hellaswag.yaml

@@ -52,7 +52,7 @@ model:
     linear: torch
     rms_norm: torch_fp32
     experts: torch_mm
-    dispatcher: torch
+    dispatcher: hybridep
     rope_fusion: false
     fake_balanced_gate: false
     enable_hf_state_dict_adapter: true
@@ -80,7 +80,7 @@ distributed:
     patch_inner_model: false
     patch_causal_lm_model: false
   moe:
-    reshard_after_forward: false
+    reshard_after_forward: true
     wrap_outer_model: false
 
 loss_fn:

examples/llm_finetune/kimi/k3_hellaswag_lora.yaml

@@ -0,0 +1,162 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Ordinary LoRA SFT for the 93-layer Kimi K3 text model on 128 GPUs:
+#   32 GB200 nodes x 4 GPUs, PP4, DP32, EP32, 100 steps.
+# LoRA targets all supported linear modules. Dense modules use rank 32; routed
+# experts use rank 2 because moe_rank_scaling divides the rank by the 16 active
+# experts per token. HybridEP and torch_mm are the validated K3 expert path.
+
+recipe: TrainFinetuneRecipeForNextTokenPrediction
+
+step_scheduler:
+  gc_every_steps: 10
+  global_batch_size: 256
+  local_batch_size: 8
+  ckpt_every_steps: 100
+  val_every_steps: 100
+  max_steps: 100
+  num_epochs: 1
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 120
+
+rng:
+  _target_: nemo_automodel.components.training.rng.StatefulRNG
+  seed: 1234
+  ranked: true
+
+model:
+  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
+  pretrained_model_name_or_path: moonshotai/Kimi-K3
+  torch_dtype: bfloat16
+  trust_remote_code: false
+  config:
+    architectures: [KimiK3ForCausalLM]
+  kda_mode: chunk
+  use_liger_kernel: false
+  use_sdpa_patching: false
+  backend:
+    _target_: nemo_automodel.components.models.common.BackendConfig
+    attn: te
+    linear: torch
+    rms_norm: torch_fp32
+    experts: torch_mm
+    dispatcher: hybridep
+    rope_fusion: false
+    fake_balanced_gate: false
+    enable_hf_state_dict_adapter: true
+    enable_fsdp_optimizations: true
+
+peft:
+  _target_: nemo_automodel.components._peft.lora.PeftConfig
+  target_modules: ["*"]
+  dim: 32
+  alpha: 32
+  dropout: 0.0
+  use_memory_efficient_lora: false
+  use_triton: false
+  moe_rank_scaling: true
+
+checkpoint:
+  enabled: true
+  checkpoint_dir: checkpoints/kimi_k3_hellaswag_lora
+  model_save_format: safetensors
+  save_consolidated: final
+  dequantize_base_checkpoint: true
+  max_recent_checkpoints: 1
+
+distributed:
+  strategy: fsdp2
+  tp_size: 1
+  cp_size: 1
+  pp_size: 4
+  ep_size: 32
+  sequence_parallel: false
+  activation_checkpointing: selective
+  defer_fsdp_grad_sync: true
+  pipeline:
+    pp_schedule: 1f1b
+    pp_microbatch_size: 1
+    scale_grads_in_schedule: false
+    patch_inner_model: false
+    patch_causal_lm_model: false
+  moe:
+    ignore_router_for_ac: true
+    reshard_after_forward: true
+    wrap_outer_model: false
+
+loss_fn:
+  _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy
+
+dataset:
+  _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag
+  path_or_dataset: Rowan/hellaswag
+  split: train
+  num_samples_limit: 25600
+  pad_to_max_length: false
+  tokenizer:
+    _target_: nemo_automodel.NeMoAutoTokenizer.from_pretrained
+    pretrained_model_name_or_path: moonshotai/Kimi-K3
+    trust_remote_code: false
+
+packed_sequence:
+  packed_sequence_size: 0
+
+dataloader:
+  _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  collate_fn: nemo_automodel.components.datasets.utils.default_collater
+  shuffle: true
+  num_workers: 0
+  drop_last: true
+
+validation_dataset:
+  _target_: nemo_automodel.components.datasets.llm.hellaswag.HellaSwag
+  path_or_dataset: Rowan/hellaswag
+  split: validation
+  num_samples_limit: 256
+  pad_to_max_length: false
+  tokenizer:
+    _target_: nemo_automodel.NeMoAutoTokenizer.from_pretrained
+    pretrained_model_name_or_path: moonshotai/Kimi-K3
+    trust_remote_code: false
+
+validation_dataloader:
+  _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  collate_fn: nemo_automodel.components.datasets.utils.default_collater
+  batch_size: 8
+  num_workers: 0
+
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 1.0e-4
+  weight_decay: 0.0
+  betas: [0.9, 0.95]
+  eps: 1.0e-8
+
+wandb:
+  enable: false
+  project: kimi-k3-sft
+  name: kimi_k3_hellaswag_lora_ep32_pp4
+  dir: wandb
+
+clip_grad_norm:
+  max_norm: 1.0
+
+ci:
+  recipe_owner: huiyingl
+  nodes: 32
+  cluster_tag: gb200
+  time: "01:00:00"

tests/unit_tests/recipes/llm/test_k3_lora_recipe.py

@@ -0,0 +1,83 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Configuration contracts for the Kimi K3 SFT examples."""
+
+from pathlib import Path
+
+import yaml
+
+from nemo_automodel.components._peft.lora import PeftConfig
+from nemo_automodel.components.config.loader import load_yaml_config
+from nemo_automodel.components.models.common import BackendConfig
+from nemo_automodel.recipes._dist_utils import parse_distributed_section
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+CONFIG_PATH = REPO_ROOT / "examples/llm_finetune/kimi/k3_hellaswag_lora.yaml"
+SFT_CONFIG_PATH = REPO_ROOT / "examples/llm_finetune/kimi/k3_hellaswag.yaml"
+
+
+def test_k3_full_sft_recipe_uses_hybridep_with_resharding() -> None:
+    """The full-SFT recipe keeps HybridEP dispatch and reshard-after-forward enabled."""
+    raw_config = yaml.safe_load(SFT_CONFIG_PATH.read_text(encoding="utf-8"))
+    config = load_yaml_config(SFT_CONFIG_PATH)
+    distributed = parse_distributed_section(raw_config["distributed"])
+
+    backend = config.model.backend.instantiate()
+
+    assert isinstance(backend, BackendConfig)
+    assert backend.experts == "torch_mm"
+    assert backend.dispatcher == "hybridep"
+    assert distributed["moe_parallel_config"].reshard_after_forward is True
+
+
+def test_k3_lora_recipe_declares_expert_lora_scaling_contract() -> None:
+    """The recipe keeps the validated K3 LoRA backend, PEFT, and topology settings."""
+    raw_config = yaml.safe_load(CONFIG_PATH.read_text(encoding="utf-8"))
+    config = load_yaml_config(CONFIG_PATH)
+    distributed = parse_distributed_section(raw_config["distributed"])
+
+    backend = config.model.backend.instantiate()
+    peft = config.peft.instantiate()
+
+    assert isinstance(backend, BackendConfig)
+    assert backend.attn == "te"
+    assert backend.experts == "torch_mm"
+    assert backend.dispatcher == "hybridep"
+
+    assert isinstance(peft, PeftConfig)
+    assert peft.target_modules == ["*"]
+    assert peft.dim == 32
+    assert peft.alpha == 32
+    assert peft.use_memory_efficient_lora is False
+    assert peft.use_triton is False
+    assert peft.moe_rank_scaling is True
+
+    assert config.model.pretrained_model_name_or_path == "moonshotai/Kimi-K3"
+    assert config.step_scheduler.max_steps == 100
+    assert config.checkpoint.enabled is True
+    assert config.checkpoint.save_consolidated == "final"
+
+    assert raw_config["distributed"]["pp_size"] == 4
+    assert raw_config["distributed"]["ep_size"] == 32
+    assert distributed["activation_checkpointing"] == "selective"
+    assert distributed["moe_parallel_config"].ignore_router_for_ac is True
+    assert distributed["moe_parallel_config"].reshard_after_forward is True
+
+    assert raw_config["ci"] == {
+        "recipe_owner": "huiyingl",
+        "nodes": 32,
+        "cluster_tag": "gb200",
+        "time": "01:00:00",
+    }