feat(model): add Kimi K3 language backbone support

.github/workflows/link-check.yml

@@ -44,6 +44,8 @@ jobs:
         uses: lycheeverse/lychee-action@v2
         with:
           fail: true
+          # TODO(kimi-k3): drop the two kimi_k3 / kimi-k3 excludes below once this PR is
+          # merged. They point at files added by the same PR, so blob/main 404s until then.
           args: >-
             --no-progress
             --max-retries 3
@@ -56,6 +58,8 @@ jobs:
             --exclude 'github\.com/NVIDIA-NeMo/Megatron-Bridge/stargazers'
             --exclude 'sched\.com'
             --exclude 'llama\.com'
+            --exclude 'Megatron-Bridge/blob/main/src/megatron/bridge/models/kimi/kimi_k3_'
+            --exclude 'Megatron-Bridge/blob/main/examples/model_verification_cards/kimi-k3/'
             "docs/fern/**/*.mdx"
             "docs/**/*.md"
         env:

docs/models/README.md

@@ -20,7 +20,7 @@ Megatron Bridge conversion, training recipe links, and model-specific notes.
 | **Gemma** | [Gemma (deprecated)](gemma/gemma.md), [Gemma 2 (deprecated)](gemma/gemma2.md), [Gemma 3](gemma/gemma3.md), [Gemma 3 VL](gemma/gemma3-vl.md), [Gemma 4 VL](gemma/gemma4-vl.md) |
 | **GLM** | [GLM 4.5](glm/glm45.md), [GLM-4.5V](glm/glm-45v.md), [GLM-4.7 / 4.7-Flash](glm/glm47.md), [GLM-5 / 5.1](glm/glm5.md) |
 | **GPT-OSS** | [GPT OSS](gpt_oss/gpt-oss.md) |
-| **Kimi** | [Kimi K2](kimi/kimi-k2.md), [Kimi-K2.5-VL](kimi/kimi-k25-vl.md) |
+| **Kimi** | [Kimi K2](kimi/kimi-k2.md), [Kimi-K2.5-VL](kimi/kimi-k25-vl.md), [Kimi K3](kimi/kimi-k3.md) |
 | **Llama** | [Llama 2 (deprecated)](llama/llama2.md), [Llama 3](llama/llama3.md) |
 | **MiniMax** | [MiniMax-M2 / M2.5 / M2.7](minimax/minimax-m2.md), [MiniMax-M3](minimax/minimax-m3.md) |
 | **Mistral** | [Mistral 7B / Small 3 24B (deprecated)](mistral/mistral.md), [Ministral 3](mistral/ministral3.md) |

docs/models/kimi/index.md

@@ -7,9 +7,11 @@ Kimi model documentation is organized by model variant.
 
 kimi-k2.md
 kimi-k25-vl.md
+kimi-k3.md
 ```
 
 | Variant | Guide |
 |---------|-------|
 | Kimi K2 | [kimi-k2.md](kimi-k2.md) |
 | Kimi-K2.5-VL | [kimi-k25-vl.md](kimi-k25-vl.md) |
+| Kimi K3 | [kimi-k3.md](kimi-k3.md) |

docs/models/kimi/kimi-k3.md

@@ -0,0 +1,76 @@
+# Kimi K3
+
+[Kimi K3](https://huggingface.co/moonshotai/Kimi-K3) is a large sparse MoE model from Moonshot AI. Megatron Bridge supports the **language backbone** of the published multimodal checkpoint through the `KimiK3Bridge`.
+
+```{note}
+Support for this model is in progress. Conversion (HF → Megatron) and Megatron greedy inference are verified; strict full-checkpoint export, exact round-trip parity, and every training workflow are not. See [Known Limitations](#known-limitations) and the machine-readable [verification card](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/model_verification_cards/kimi-k3/card.yaml) before relying on this path.
+```
+
+## Supported Variants
+
+Megatron Bridge supports checkpoints with the `KimiK3ForConditionalGeneration` architecture and the `kimi_k3` model type:
+
+| Variant | HF Path |
+|---------|---------|
+| Kimi-K3 | [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) |
+
+Requires `transformers >= 4.56.2` and `--trust-remote-code`.
+
+## Architecture Notes
+
+K3 uses a heterogeneous attention schedule rather than a single attention type:
+
+- **KDA (Kimi Delta Attention)** on the layers listed in the HF config's `linear_attn_config.kda_layers` — a gated delta-rule linear-attention block with short depthwise convolutions over Q/K/V, a low-rank forget gate, and a per-head beta projection. In the published 93-layer checkpoint, 69 layers are KDA.
+- **No-RoPE MLA** on the remaining layers (`full_attn_layers`, 24 layers in the published checkpoint).
+
+Both layer lists hold **1-indexed** global layer numbers and together partition `1..num_hidden_layers`.
+
+Other notable properties:
+
+- Latent MoE (`moe_latent_size`) with shared experts, grouped GEMM, and all-to-all token dispatch.
+- AttnRes residual banks with a configurable block size (`attn_res_block_size`), which the pipeline payload carries between stages.
+- Published routed-expert weights are MXFP4 (`uint8`-packed E2M1 values with UE8M0 scales) and are dequantized to BF16 on import.
+- The published KDA `A_log` tensors carry 96 active entries plus 32 zero-only padding entries. Import validates the padding is all-zero and removes it; export restores it.
+- Export of the language backbone passes the checkpoint's `vision_tower.*` and `mm_projector.*` tensors through unchanged.
+
+## Conversion
+
+```bash
+# HF → Megatron
+./scripts/conversion/convert.sh import \
+    --hf-model moonshotai/Kimi-K3 \
+    --megatron-path /workspace/kimi-k3 \
+    --torch-dtype bfloat16 \
+    --tp 2 --pp 3 --ep 8 --etp 2 \
+    --distributed-timeout-minutes 180 \
+    --trust-remote-code
+```
+
+The full checkpoint needs a multi-node allocation — import was validated on 48 GB200 GPUs at TP2/PP3/EP8/ETP2.
+
+For a fast local iteration loop, build a truncated proxy checkpoint with
+[`examples/conversion/create_hf_toy_model.py`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/create_hf_toy_model.py), which truncates the heterogeneous layer schedule and downloads only the safetensor shards the selected layers need.
+
+## Inference
+
+Megatron greedy generation was validated on 24 GB300 GPUs at TP1/PP3/EP8/ETP1. See the verification card for the exact command and the recorded deterministic completion.
+
+## Training
+
+No training recipe ships for K3 yet. Pretraining, SFT, and PEFT configs, checkpoint-resume validation, and performance tuning are pending.
+
+## Known Limitations
+
+- Strict full-checkpoint Megatron → HF export, HF reload, and exact round-trip parity are unverified.
+- Full HF/Megatron forward-logit correlation is unverified. A four-layer proxy reached cosine similarity `0.9998` and Pearson correlation `0.9998`.
+- Virtual pipeline parallelism (VPP) is not supported.
+- KDA layers do not support context parallelism (`CP > 1`).
+- Only the language backbone is covered. Native K3 vision/video modeling and multimodal inference are not implemented; export only preserves the published vision and projector tensors unchanged.
+- The model has not been performance-tuned. Reported timings are sanity checks, not optimized throughput results.
+
+## Related Implementation
+
+- Bridge: [`src/megatron/bridge/models/kimi/kimi_k3_bridge.py`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/kimi/kimi_k3_bridge.py)
+- Provider: [`src/megatron/bridge/models/kimi/kimi_k3_provider.py`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/kimi/kimi_k3_provider.py)
+- Layer spec and KDA/MLA modules: [`src/megatron/bridge/models/kimi/kimi_k3_layers.py`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/kimi/kimi_k3_layers.py)
+- Verification card: [`examples/model_verification_cards/kimi-k3/card.yaml`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/model_verification_cards/kimi-k3/card.yaml)

examples/conversion/create_hf_toy_model.py

@@ -87,19 +87,81 @@ def _parse_args() -> argparse.Namespace:
     parser.add_argument("output_dir", type=Path, help="Directory for the truncated checkpoint.")
     parser.add_argument("--num-hidden-layers", type=int, required=True, help="Number of leading layers to retain.")
     parser.add_argument("--revision", help="Optional Hugging Face Hub revision when source is a model ID.")
+    parser.add_argument("--cache-dir", type=Path, help="Optional Hugging Face Hub cache directory.")
     parser.add_argument("--overwrite", action="store_true", help="Replace output_dir if it already exists.")
     return parser.parse_args()
 
 
-def _resolve_source(source: str, *, revision: str | None) -> Path:
-    source_path = Path(source).expanduser()
-    if source_path.is_dir():
-        return source_path.resolve()
+def _select_hub_files(
+    repo_files: list[str],
+    index: dict[str, object] | None,
+    *,
+    num_hidden_layers: int,
+) -> list[str]:
+    """Select metadata and weight shards needed for a shallow Hub checkpoint."""
+    selected = {filename for filename in repo_files if "/" not in filename and _should_copy_metadata(Path(filename))}
+    if "config.json" in repo_files:
+        selected.add("config.json")
+    if index is None:
+        safetensors_files = [
+            filename for filename in repo_files if "/" not in filename and filename.endswith(".safetensors")
+        ]
+        if len(safetensors_files) != 1:
+            raise ValueError(f"Expected one unsharded safetensors file, found {len(safetensors_files)}")
+        selected.update(safetensors_files)
+        return sorted(selected)
+
+    selected.add(_SAFETENSORS_INDEX_NAME)
+    weight_map = cast(dict[str, str], index["weight_map"])
+    selected.update(
+        shard
+        for tensor_name, shard in weight_map.items()
+        if _retains_tensor(tensor_name, num_hidden_layers=num_hidden_layers)
+    )
+    return sorted(selected)
 
-    from huggingface_hub import snapshot_download
 
-    LOGGER.info("Downloading Hugging Face checkpoint %s", source)
-    return Path(snapshot_download(repo_id=source, revision=revision))
+def _resolve_source(
+    source: str,
+    *,
+    revision: str | None,
+    num_hidden_layers: int,
+    cache_dir: Path | None,
+) -> tuple[Path, bool]:
+    source_path = Path(source).expanduser()
+    if source_path.is_dir():
+        return source_path.resolve(), False
+
+    from huggingface_hub import HfApi, hf_hub_download, snapshot_download
+
+    repo_files = HfApi().list_repo_files(source, revision=revision)
+    index = None
+    if _SAFETENSORS_INDEX_NAME in repo_files:
+        index_path = Path(
+            hf_hub_download(
+                repo_id=source,
+                filename=_SAFETENSORS_INDEX_NAME,
+                revision=revision,
+                cache_dir=cache_dir,
+            )
+        )
+        index = cast(dict[str, object], json.loads(index_path.read_text()))
+    selected_files = _select_hub_files(repo_files, index, num_hidden_layers=num_hidden_layers)
+    selected_shards = sum(filename.endswith(".safetensors") for filename in selected_files)
+    LOGGER.info(
+        "Downloading %d weight shard(s) plus metadata from Hugging Face checkpoint %s",
+        selected_shards,
+        source,
+    )
+    snapshot_path = Path(
+        snapshot_download(
+            repo_id=source,
+            revision=revision,
+            allow_patterns=selected_files,
+            cache_dir=cache_dir,
+        )
+    )
+    return snapshot_path, index is not None
 
 
 def _prepare_output(source_dir: Path, output_dir: Path, *, overwrite: bool) -> Path:
@@ -127,27 +189,56 @@ def _copy_metadata_files(source_dir: Path, output_dir: Path) -> None:
             shutil.copy2(source_path, output_dir / source_path.name, follow_symlinks=True)
 
 
+def _language_config(config: dict[str, object]) -> dict[str, object]:
+    """Return the config dictionary that owns the transformer layers."""
+    if "num_hidden_layers" in config:
+        return config
+
+    text_config = config.get("text_config")
+    if isinstance(text_config, dict) and "num_hidden_layers" in text_config:
+        return cast(dict[str, object], text_config)
+
+    raise ValueError("config.json does not contain num_hidden_layers at the top level or under text_config")
+
+
 def _truncate_config(source_dir: Path, output_dir: Path, *, num_hidden_layers: int) -> int:
     config_path = source_dir / "config.json"
     if not config_path.is_file():
         raise FileNotFoundError(f"Hugging Face config not found: {config_path}")
 
     config = json.loads(config_path.read_text())
-    if "num_hidden_layers" not in config:
-        raise ValueError("config.json does not contain a top-level num_hidden_layers field")
-    original_num_hidden_layers = int(config["num_hidden_layers"])
+    transformer_config = _language_config(config)
+    original_num_hidden_layers = int(transformer_config["num_hidden_layers"])
     if not 0 < num_hidden_layers <= original_num_hidden_layers:
         raise ValueError(
             f"num_hidden_layers must be between 1 and {original_num_hidden_layers}, got {num_hidden_layers}"
         )
 
-    config["num_hidden_layers"] = num_hidden_layers
-    if "max_window_layers" in config:
-        config["max_window_layers"] = min(int(config["max_window_layers"]), num_hidden_layers)
-    if isinstance(config.get("layer_types"), list):
-        config["layer_types"] = config["layer_types"][:num_hidden_layers]
-    if isinstance(config.get("mlp_only_layers"), list):
-        config["mlp_only_layers"] = [layer for layer in config["mlp_only_layers"] if layer < num_hidden_layers]
+    transformer_config["num_hidden_layers"] = num_hidden_layers
+    if "max_window_layers" in transformer_config:
+        transformer_config["max_window_layers"] = min(int(transformer_config["max_window_layers"]), num_hidden_layers)
+    if isinstance(transformer_config.get("layer_types"), list):
+        transformer_config["layer_types"] = transformer_config["layer_types"][:num_hidden_layers]
+    if isinstance(transformer_config.get("mlp_only_layers"), list):
+        transformer_config["mlp_only_layers"] = [
+            layer for layer in transformer_config["mlp_only_layers"] if layer < num_hidden_layers
+        ]
+    linear_attn_config = transformer_config.get("linear_attn_config")
+    if isinstance(linear_attn_config, dict):
+        # Unlike the exclusive bounds used above, these lists hold 1-indexed global layer
+        # *numbers*, not 0-indexed offsets, so the retained range is [1, num_hidden_layers]
+        # and the bound is inclusive. Kimi K3's published config is the reference case:
+        # across 93 layers, `kda_layers` spans 1..91 and `full_attn_layers` spans 4..93,
+        # and the two partition 1..93 exactly. The consumer side agrees — the layer spec
+        # tests `layer_number in config.kimi_kda_layers` with MCore's 1-indexed
+        # `layer_number` (see `KimiK3Attention.__init__`).
+        for layer_list_name in ("kda_layers", "full_attn_layers"):
+            if isinstance(linear_attn_config.get(layer_list_name), list):
+                linear_attn_config[layer_list_name] = [
+                    layer_number
+                    for layer_number in linear_attn_config[layer_list_name]
+                    if layer_number <= num_hidden_layers
+                ]
 
     (output_dir / "config.json").write_text(json.dumps(config, indent=2) + "\n")
     return original_num_hidden_layers
@@ -279,6 +370,7 @@ def truncate_checkpoint(
     *,
     num_hidden_layers: int,
     revision: str | None = None,
+    cache_dir: Path | None = None,
     overwrite: bool = False,
 ) -> Path:
     """Create a shallow checkpoint that retains pretrained weights from leading layers.
@@ -288,19 +380,25 @@ def truncate_checkpoint(
         output_dir: Directory where the truncated checkpoint is written.
         num_hidden_layers: Number of leading transformer layers to retain.
         revision: Optional Hub revision used when ``source`` is a model ID.
+        cache_dir: Optional Hugging Face Hub cache directory.
         overwrite: Whether to replace an existing output directory.
 
     Returns:
         Resolved path to the truncated checkpoint.
     """
-    source_dir = _resolve_source(source, revision=revision)
+    source_dir, selectively_downloaded = _resolve_source(
+        source,
+        revision=revision,
+        num_hidden_layers=num_hidden_layers,
+        cache_dir=cache_dir,
+    )
     output_dir = _prepare_output(source_dir, output_dir, overwrite=overwrite)
     _copy_metadata_files(source_dir, output_dir)
     original_num_hidden_layers = _truncate_config(source_dir, output_dir, num_hidden_layers=num_hidden_layers)
     tensor_count, tensor_bytes, removed_tensor_count = _rewrite_checkpoint(
         source_dir, output_dir, num_hidden_layers=num_hidden_layers
     )
-    if num_hidden_layers < original_num_hidden_layers and removed_tensor_count == 0:
+    if not selectively_downloaded and num_hidden_layers < original_num_hidden_layers and removed_tensor_count == 0:
         raise ValueError("No layer tensors were removed; expected tensor names containing layers.<index>")
     LOGGER.info(
         "Created %s from %s: layers=%d/%d, tensors=%d, tensor_bytes=%d",
@@ -323,6 +421,7 @@ def main() -> None:
         args.output_dir,
         num_hidden_layers=args.num_hidden_layers,
         revision=args.revision,
+        cache_dir=args.cache_dir,
         overwrite=args.overwrite,
     )