diff --git a/examples/llm_finetune/kimi/k3_hellaswag.yaml b/examples/llm_finetune/kimi/k3_hellaswag.yaml index 770e2d88b7..58b609a0aa 100644 --- a/examples/llm_finetune/kimi/k3_hellaswag.yaml +++ b/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: diff --git a/examples/llm_finetune/kimi/k3_hellaswag_lora.yaml b/examples/llm_finetune/kimi/k3_hellaswag_lora.yaml new file mode 100644 index 0000000000..673f3242c9 --- /dev/null +++ b/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" diff --git a/tests/unit_tests/recipes/llm/test_k3_lora_recipe.py b/tests/unit_tests/recipes/llm/test_k3_lora_recipe.py new file mode 100644 index 0000000000..1eab695d85 --- /dev/null +++ b/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", + }