Add Moonlight-V4-1B variant config, tokenizer, model card and training helpers

README.md

@@ -0,0 +1,186 @@
+---
+license: mit
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+  - deepseek_v4
+  - mixture-of-experts
+  - moe
+  - hybrid-attention
+  - compressed-sparse-attention
+  - hyper-connections
+  - mhc
+  - architecture-config
+  - untrained
+  - from-scratch
+  - nemo-automodel
+  - moonlight
+---
+
+# Moonlight-V4-1B-h16d256
+
+**An untrained, ~1B-parameter DeepSeek-V4-architecture configuration** (no weights): DeepSeek-V4-architecture 1B config with 16 x 256 attention heads (CSA ratio 4 + lightning indexer).
+It is the small end of a Moonlight-style down-scaling of DeepSeek-V4, sized to pre-train on two 48 GB GPUs, with
+attention dimensions chosen so that the TileLang sparse-attention kernel fits GPUs with 99 KB of shared memory
+(Ada / consumer class). This is the V4-faithful variant: 7 Compressed Sparse Attention (CSA) layers with ratio-4 overlapped compression and a lightning indexer, 6 Heavily Compressed Attention (HCA) layers (ratio 128) and 2 pure sliding-window layers.
+
+The sibling repo replaces the ratio-4 CSA layers by ratio-8 compression without an indexer, which is cheaper to train at short context and avoids the indexer kernel entirely. Sibling: [`akoumpa/Moonlight-V4-1B-h16d256-r8`](https://huggingface.co/akoumpa/Moonlight-V4-1B-h16d256-r8).
+
+| | total | non-embedding | activated / token | activated non-embedding |
+| --- | ---: | ---: | ---: | ---: |
+| parameters | **999,670,879** (999.7M) | 734.9M | 539.6M | 274.8M |
+
+## Lineage
+
+| model | architecture | size | status |
+| --- | --- | --- | --- |
+| [Moonlight-16B-A3B](https://huggingface.co/moonshotai/Moonlight-16B-A3B) (Moonshot) | DeepSeek-V3 | 16B / 3B active | released, trained with Muon |
+| Moonlight-V4-16B-A3B | DeepSeek-V4 at Moonlight's width/depth/experts | 16.5B / 3.0B active | proposed config |
+| Moonlight-V4-1B (8 heads x 512) | same, shrunk | 1.0B / 0.54B active | config; eager attention only |
+| **Moonlight-V4-1B-h16d256** | same, attention re-shaped for the kernels | 999.7M / 539.6M active | **this repo** |
+
+DeepSeek-V4 (Flash: 284B/13B, Pro: 1.6T/49B) keeps `head_dim` = 512 and 64 or 128 query heads; Moonlight's 16 heads
+were kept for the 16B analogue, and here 16 heads x 256 dims give the same 4096-wide query space as the 8 x 512
+1B config with half the shared-KV width.
+
+## Architecture
+
+| component | setting |
+| --- | --- |
+| hidden size / layers | 1024 / 15 |
+| attention schedule (`compress_ratios`) | `[0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4]`: 2 sliding-window, 7 compressed (ratio 4, overlapped, with indexer), 6 HCA (ratio 128) |
+| attention | shared-KV MQA (`num_key_value_heads` 1): 16 query heads x `head_dim` 256 (last 64 dims RoPE), the same 256-dim entry is key and value |
+| query path | `q_lora_rank` 256 -> 16 x 256; per-head RMSNorm before RoPE |
+| output projection | grouped low-rank: `o_groups` 2 (8 heads per group) x `o_lora_rank` 1024 -> hidden |
+| sliding window / attention sinks | 128 tokens on every layer / one learnable sink logit per head |
+| Lightning indexer | 64 heads x 128 dims, `index_topk` 1024 (a power of two, as the indexer kernels require); at <= 4K tokens this is >= the 1024 pooled entries, i.e. dense |
+| MoE (every layer) | 32 routed experts x 384 (SwiGLU, clamp 10.0), top-6, 1 shared expert; `sqrtsoftplus` scoring, aux-loss-free bias (`noaux_tc`), `routed_scaling_factor` 2.436 |
+| hash-routed layers | first 1 MoE layer(s) route by a fixed token-id table (`tid2eid`) |
+| residual stream | manifold-constrained hyper-connections, `hc_mult` 4, 20 Sinkhorn iterations |
+| positions | 4096 tokens, plain RoPE (theta 10000 on sliding layers, 160000 on compressed layers, no YaRN) |
+| MTP | none (`num_nextn_predict_layers` 0) |
+| vocabulary | 129280 (DeepSeek-V4 tokenizer, BOS 0, EOS 1) |
+| norm eps / init | 1e-06 / 0.02 |
+
+Per-layer parameters: attention 7.9M (sliding) / 11.6M (ratio 4) / 8.4M (HCA); MoE 39.0M
+total, 8.3M activated (each expert 1.2M); mHC mixers 196,662. Embedding and
+head are 132.4M each. KV cache per sequence at 4K / 32K tokens (FP8 non-RoPE dims, bf16 RoPE dims): 3.3 MiB / 22.1 MiB.
+Core-attention FLOPs per generated token at 4K context: 0.27 GF, against 0.81 GF of linear layers.
+`python count_params.py config.json` reproduces these numbers (the script also reproduces the published sizes of
+Moonlight, DeepSeek-V3, V4-Flash and V4-Pro).
+
+Design rules: `head_dim` and `index_topk` are powers of two and `index_n_heads` divides 128 (kernel requirements);
+`o_lora_rank` 1024 keeps DeepSeek-V4's per-group output projection shape; `routed_scaling_factor` follows Moonlight's
+recipe (expected 1/||p||_2 of renormalised top-k scores) applied to sqrt-softplus with 32 experts / top-6.
+
+## Files
+
+| file | purpose |
+| --- | --- |
+| `config.json` | Hugging Face config (same key set as `deepseek-ai/DeepSeek-V4-Flash`) |
+| `inference_config.json` | the same model in the key format of DeepSeek's reference `inference/model.py` |
+| `tokenizer.json`, `tokenizer_config.json` | DeepSeek-V4 tokenizer (MIT), copied from `deepseek-ai/DeepSeek-V4-Flash` |
+| `count_params.py` | parameter / KV-cache / FLOP calculator for DeepSeek-V3- and V4-style configs |
+| `training/pretrain.yaml`, `training/train.py` | NeMo Automodel recipe (FSDP2, 2 GPUs) and launcher that seeds the hash table and mHC mixers |
+| `training/prepare_data.py`, `training/finite_nanogpt.py`, `training/init_utils.py` | data shards from parquet text, bounded validation dataset, from-scratch initialisers |
+
+## Loading
+
+### transformers (>= 5.8, native `deepseek_v4`)
+
+```python
+from transformers import AutoTokenizer, DeepseekV4Config, DeepseekV4ForCausalLM
+cfg = DeepseekV4Config.from_pretrained("akoumpa/Moonlight-V4-1B-h16d256")          # legacy keys (compress_ratios, num_hash_layers, ...) are folded in
+tok = AutoTokenizer.from_pretrained("akoumpa/Moonlight-V4-1B-h16d256")
+model = DeepseekV4ForCausalLM(cfg)                             # random init, 999.7M parameters
+```
+
+The transformers implementation is **inference-oriented**: with no KV cache it appends compressed entries to the
+key axis without a causal mask and gathers per-query top-k entries into an `S x k` key axis, so do not train with it.
+Use it for architecture inspection and, once you have weights, for generation.
+
+### NeMo Automodel (native training implementation)
+
+```python
+from nemo_automodel.components.models.common import BackendConfig
+from nemo_automodel.components.models.deepseek_v4.config import DeepseekV4Config
+from nemo_automodel.components.models.deepseek_v4.model import DeepseekV4ForCausalLM
+cfg = DeepseekV4Config.from_pretrained("akoumpa/Moonlight-V4-1B-h16d256")
+backend = BackendConfig(attn="eager", linear="torch", rms_norm="torch_fp32", rope_fusion=False,
+                        dispatcher="torch", experts="torch_mm", enable_hf_state_dict_adapter=False)
+model = DeepseekV4ForCausalLM(cfg, backend=backend)
+model.initialize_weights(dtype=torch.bfloat16)
+```
+
+Or in a recipe: `NeMoAutoModelForCausalLM.from_config` with `config: DeepseekV4Config.from_pretrained("akoumpa/Moonlight-V4-1B-h16d256")`
+(see `training/pretrain.yaml`).
+
+### DeepSeek reference code
+
+`inference_config.json` drops into the `inference/` folder of the DeepSeek-V4 release (`ModelArgs` keys, `n_mtp_layers` 0).
+
+## Training from scratch
+
+The released implementations only ever load checkpoints, so two things must be initialised by hand; `training/init_utils.py`
+does both and `training/train.py` calls it after the recipe's own weight init on a fresh start:
+
+1. **Hash-routing table.** `tid2eid` is created as zeros (every token to expert 0). `fill_hash_tables` writes a balanced
+   token-id hash: 32 experts, 6 distinct experts per token, equal load over the vocabulary.
+2. **mHC mixers.** NeMo Automodel leaves the `fn` / `base` / `scale` tensors uninitialised; `init_hyper_connections` applies
+   transformers' rule (normal(0, 0.02) projection, zero bias, unit gates).
+
+Recipe notes that cost time to find (all encoded in `training/`):
+
+- Automodel wraps DeepSeek-V4's fp32 tensors (attention sinks, compressor position biases, mHC mixers, `lm_head`) as their
+  own FSDP2 units whose forward returns the parameter; if they reshard after forward, attention reads a freed tensor.
+  `train.py` calls `set_reshard_after_forward(False)` on those units.
+- Use the logits-based `MaskedCrossEntropy`: the fused linear cross-entropy rejects the fp32 `lm_head` x bf16 hidden states.
+- For iterable datasets the recipe passes no batch size to the DataLoader; set `dataloader.batch_size` explicitly.
+- `NanogptDataset` is an infinite stream; validation uses `FiniteNanogptDataset`.
+- The indexer's top-k has no gradient path (Automodel freezes its parameters). With `index_topk` 1024 every query sees all causal compressed entries up to 4K tokens, i.e. DeepSeek's own dense warm-up regime; sparse training at longer contexts needs an indexer distillation loss.
+- To train beyond 4K, add DeepSeek-V4's YaRN block (`rope_scaling`: factor 16, `original_max_position_embeddings` 65536)
+  and raise `max_position_embeddings` and `index_topk`.
+
+Measured on two RTX 5880 Ada (48 GB, sm_89) GPUs, bf16, `torch_mm` experts, chunked cross-entropy, single-GPU forward+backward
+(TileLang sparse attention with Sinkhorn and indexer on torch; the eager path for the same model reaches 4.2k tok/s at
+B=2 x 2048 and runs out of memory at B=4):
+
+| micro-batch | fwd+bwd time, throughput, peak memory |
+| --- | --- |
+| 2 x 2048 | 740 ms, 5.5k tok/s, 15.1 GiB |
+| 4 x 2048 | 1227 ms, 6.7k tok/s, 27.9 GiB |
+| 2 x 4096 | 1421 ms, 5.8k tok/s, 33.1 GiB |
+
+Forward-time breakdown at B=2 x 2048: attention 59 ms, indexer 82 ms, MoE 39 ms, mHC mixers 35 ms, compressor 12 ms (247 ms forward). With the recipe's FSDP2 data parallelism over 2 GPUs and
+AdamW, a 32-sequence x 2048-token global batch is a reasonable starting point (`training/pretrain.yaml`).
+
+## TileLang kernels
+
+Automodel's vendored Miles/TileLang kernels (sparse attention, indexer) and DeepSeek's TileKernels Sinkhorn were written
+for Hopper's 227 KB of shared memory. On a 99 KB-per-block GPU:
+
+| kernel | shape rule | shared memory | Ada (99 KB) |
+| --- | --- | --- | --- |
+| sparse attention fwd/bwd | `head_dim` power of two; heads padded to 16, chunked by 16, multiples of 64 above 64 | 147 KB at `head_dim` 512, fits at **256** | runs with this config; parity with the torch reference verified |
+| lightning indexer fwd | `index_n_heads` <= 64, multiple of 8, divides 128; `index_topk` power of two (bwd) | 224 KB (`block_N` 256 x 128 fp32) | does not fit; run the indexer on torch |
+| Sinkhorn (mHC) | needs the `tile_kernels` package | small | needs TileKernels installed or the torch fallback |
+
+Automodel currently switches all three together (`backend.attn: tilelang`); the measurements above used per-kernel
+selection (only sparse attention on TileLang). On Hopper GPUs the default `head_dim` 512 kernels fit and the 16 x 256
+choice is optional.
+
+## Limitations and notes
+
+- No trained weights are provided; all numbers are architecture-derived or short from-scratch measurements.
+- The name follows the Moonlight / DeepSeek-V4 lineage for clarity; this repository is not affiliated with Moonshot AI
+  or DeepSeek.
+- The tokenizer files are DeepSeek's (MIT licence, `deepseek-ai/DeepSeek-V4-Flash`).
+
+## References
+
+- Liu et al., *Muon is Scalable for LLM Training* (Moonlight), [arXiv:2502.16982](https://arxiv.org/abs/2502.16982)
+- DeepSeek-AI, *DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence*, [arXiv:2606.19348](https://arxiv.org/abs/2606.19348)
+- Xie et al., *Manifold-Constrained Hyper-Connections* (mHC), 2026
+- Roller et al., *Hash Layers for Large Sparse Models*, NeurIPS 2021
+- DeepSeek-AI, *DeepSeek-V3.2* (DeepSeek Sparse Attention, lightning indexer), 2025
+- Kernels: [Miles](https://github.com/yueming-yuan/miles) (sparse attention / indexer, vendored in NeMo Automodel), [TileKernels](https://github.com/deepseek-ai/TileKernels) (Sinkhorn), [TileLang](https://github.com/tile-ai/tilelang)

config.json

@@ -0,0 +1,65 @@
+{
+  "architectures": [
+    "DeepseekV4ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "eos_token_id": 1,
+  "hc_eps": 1e-06,
+  "hc_mult": 4,
+  "hc_sinkhorn_iters": 20,
+  "head_dim": 256,
+  "hidden_act": "silu",
+  "hidden_size": 1024,
+  "index_head_dim": 128,
+  "index_n_heads": 64,
+  "index_topk": 1024,
+  "initializer_range": 0.02,
+  "max_position_embeddings": 4096,
+  "model_type": "deepseek_v4",
+  "moe_intermediate_size": 384,
+  "n_routed_experts": 32,
+  "n_shared_experts": 1,
+  "norm_topk_prob": true,
+  "num_attention_heads": 16,
+  "num_experts_per_tok": 6,
+  "num_hidden_layers": 15,
+  "num_hash_layers": 1,
+  "num_key_value_heads": 1,
+  "num_nextn_predict_layers": 0,
+  "o_groups": 2,
+  "o_lora_rank": 1024,
+  "q_lora_rank": 256,
+  "qk_rope_head_dim": 64,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 10000,
+  "routed_scaling_factor": 2.436,
+  "scoring_func": "sqrtsoftplus",
+  "sliding_window": 128,
+  "swiglu_limit": 10.0,
+  "tie_word_embeddings": false,
+  "topk_method": "noaux_tc",
+  "torch_dtype": "bfloat16",
+  "use_cache": true,
+  "vocab_size": 129280,
+  "compress_rope_theta": 160000,
+  "compress_ratios": [
+    0,
+    0,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4
+  ]
+}

count_params.py

@@ -0,0 +1,316 @@
+#!/usr/bin/env python3
+"""
+Parameter / KV-cache / attention-FLOP calculator for
+  * DeepSeek-V3-style configs (Moonlight-16B-A3B, DeepSeek-V3)
+  * DeepSeek-V4-style configs (DeepSeek-V4-Flash / -Pro, Moonlight-V4)
+
+Formulas follow the tensor shapes of the released checkpoints:
+  * Moonlight: modeling_deepseek.py (DeepseekV3) shapes read from model.safetensors headers
+  * DeepSeek-V4: inference/model.py (ModelArgs / Block / Attention / Compressor / Indexer / MoE)
+
+Published numbers used for validation (all reproduced by this script):
+  Moonlight-16B-A3B : 15.29B total / 2.24B activated non-embedding; 16B / 3B incl. embedding
+  DeepSeek-V4-Flash : 284B total / 13B activated
+  DeepSeek-V4-Pro   : 1.6T total / 49B activated
+
+Usage:
+  python3 count_params.py                 # validation + Moonlight-V4 report
+  python3 count_params.py path/to/config.json [...]   # report for arbitrary HF-style configs
+"""
+import json
+import math
+import random
+import sys
+from pathlib import Path
+
+HERE = Path(__file__).resolve().parent
+
+# ----------------------------------------------------------------------------------------------
+# DeepSeek-V3 architecture (Moonlight, DeepSeek-V3)
+# ----------------------------------------------------------------------------------------------
+def count_v3(c):
+    V, D, L = c["vocab_size"], c["hidden_size"], c["num_hidden_layers"]
+    H = c["num_attention_heads"]
+    qk_nope, qk_rope, v_hd = c["qk_nope_head_dim"], c["qk_rope_head_dim"], c["v_head_dim"]
+    kv_r, q_r = c["kv_lora_rank"], c.get("q_lora_rank")
+    E, K, I = c["n_routed_experts"], c["num_experts_per_tok"], c["moe_intermediate_size"]
+    n_sh, I_dense = c["n_shared_experts"], c["intermediate_size"]
+    k_dense, moe_freq = c["first_k_dense_replace"], c.get("moe_layer_freq", 1)
+    n_mtp = c.get("num_nextn_predict_layers", 0)
+
+    # attention (MLA)
+    if q_r:
+        attn = D * q_r + q_r + q_r * H * (qk_nope + qk_rope)          # q_a_proj, q_a_layernorm, q_b_proj
+    else:
+        attn = D * H * (qk_nope + qk_rope)                             # q_proj
+    attn += D * (kv_r + qk_rope) + kv_r                                # kv_a_proj_with_mqa, kv_a_layernorm
+    attn += kv_r * H * (qk_nope + v_hd)                                # kv_b_proj
+    attn += H * v_hd * D                                               # o_proj
+    norms = 2 * D
+
+    dense_mlp = 3 * D * I_dense
+    gate = E * D + E                                                   # gate.weight + e_score_correction_bias
+    expert = 3 * D * I
+    moe_total = gate + E * expert + n_sh * expert
+    moe_active = gate + K * expert + n_sh * expert
+
+    def is_moe(l):
+        return l >= k_dense and (l - k_dense) % moe_freq == 0 if moe_freq > 1 else l >= k_dense
+
+    n_moe = sum(1 for l in range(L) if is_moe(l))
+    n_dense = L - n_moe
+    body_total = L * (attn + norms) + n_dense * dense_mlp + n_moe * moe_total + D
+    body_active = L * (attn + norms) + n_dense * dense_mlp + n_moe * moe_active + D
+    emb = V * D
+    head = V * D
+    # MTP module (DeepSeek-V3 style): eh_proj (2D x D) + enorm + hnorm + shared_head.norm + one full layer
+    mtp_total = n_mtp * (2 * D * D + 3 * D + attn + norms + moe_total)
+    mtp_active = n_mtp * (2 * D * D + 3 * D + attn + norms + moe_active)
+    return dict(
+        family="deepseek_v3", layers=L, n_dense=n_dense, n_moe=n_moe,
+        attn_per_layer=attn, dense_mlp=dense_mlp, moe_total_per_layer=moe_total,
+        moe_active_per_layer=moe_active, expert=expert,
+        embed=emb, head=head, body_total=body_total, body_active=body_active,
+        total_nonembed=body_total, active_nonembed=body_active,
+        total_incl_embed=body_total + emb + head, active_incl_embed=body_active + emb + head,
+        mtp_total=mtp_total, mtp_active=mtp_active,
+    )
+
+
+# ----------------------------------------------------------------------------------------------
+# DeepSeek-V4 architecture (Flash, Pro, Moonlight-V4)
+# ----------------------------------------------------------------------------------------------
+def v4_attention_params(c, ratio):
+    """Parameters of one Attention module (incl. compressor / indexer) for a layer with the given compress ratio."""
+    D, H, hd = c["hidden_size"], c["num_attention_heads"], c["head_dim"]
+    Qr, Or, G = c["q_lora_rank"], c["o_lora_rank"], c["o_groups"]
+    IH, Ihd = c["index_n_heads"], c["index_head_dim"]
+    p = {}
+    p["attn_sink"] = H
+    p["wq_a"] = D * Qr
+    p["q_norm"] = Qr
+    p["wq_b"] = Qr * H * hd
+    p["wkv"] = D * hd
+    p["kv_norm"] = hd
+    p["wo_a"] = (H * hd // G) * (G * Or)      # == H*hd*Or, applied group-wise as [G, Or, H*hd/G]
+    p["wo_b"] = G * Or * D
+    if ratio:
+        coff = 2 if ratio == 4 else 1            # overlapped compression only for the CSA ratio (4)
+        p["compressor"] = ratio * coff * hd + 2 * D * coff * hd + hd     # ape + wkv + wgate + norm
+        if ratio == 4:                           # CSA: lightning indexer on top of compressed keys
+            idx = Qr * IH * Ihd + D * IH         # wq_b + weights_proj
+            idx += 4 * 2 * Ihd + 2 * D * 2 * Ihd + Ihd   # indexer.compressor (ratio 4, overlap)
+            p["indexer"] = idx
+    return p
+
+
+def v4_hc_params(c):
+    hc, D = c["hc_mult"], c["hidden_size"]
+    mix = (2 + hc) * hc
+    return 2 * (mix * hc * D + mix + 3)         # (fn + base + scale) for attn and for ffn
+
+
+def count_v4(c):
+    V, D, L = c["vocab_size"], c["hidden_size"], c["num_hidden_layers"]
+    E, K, I = c["n_routed_experts"], c["num_experts_per_tok"], c["moe_intermediate_size"]
+    n_sh, n_hash = c["n_shared_experts"], c["num_hash_layers"]
+    hc = c["hc_mult"]
+    n_mtp = c.get("num_nextn_predict_layers", 0)
+    ratios = list(c["compress_ratios"])
+    assert len(ratios) >= L + n_mtp, f"compress_ratios needs >= {L + n_mtp} entries, got {len(ratios)}"
+
+    expert = 3 * D * I
+    gate_w, gate_b = E * D, E
+    moe_total = lambda hash_layer: gate_w + (0 if hash_layer else gate_b) + E * expert + n_sh * expert
+    moe_active = lambda hash_layer: gate_w + (0 if hash_layer else gate_b) + K * expert + n_sh * expert
+    hc_per_layer = v4_hc_params(c)
+    norms = 2 * D
+
+    per_layer = []
+    for l in range(L):
+        ap = v4_attention_params(c, ratios[l])
+        attn = sum(ap.values())
+        is_hash = l < n_hash
+        per_layer.append(dict(
+            layer=l, ratio=ratios[l], hash=is_hash, attn=attn, attn_parts=ap, hc=hc_per_layer,
+            total=attn + hc_per_layer + norms + moe_total(is_hash),
+            active=attn + hc_per_layer + norms + moe_active(is_hash),
+        ))
+    head_hc = hc * hc * D + hc + 1
+    body_total = sum(x["total"] for x in per_layer) + D + head_hc
+    body_active = sum(x["active"] for x in per_layer) + D + head_hc
+    emb, head = V * D, V * D
+
+    mtp_total = mtp_active = 0
+    for i in range(n_mtp):
+        ap = sum(v4_attention_params(c, ratios[L + i]).values())
+        extra = 2 * D * D + 3 * D + head_hc                 # e_proj, h_proj, enorm, hnorm, norm, hc_head_*
+        mtp_total += ap + hc_per_layer + norms + moe_total(False) + extra
+        mtp_active += ap + hc_per_layer + norms + moe_active(False) + extra
+
+    kinds = {0: "SWA", 4: "CSA", 128: "HCA"}
+    layer_kinds = {k: sum(1 for l in range(L) if kinds.get(ratios[l]) == k) for k in kinds.values()}
+    return dict(
+        family="deepseek_v4", layers=L, layer_kinds=layer_kinds, n_hash=n_hash,
+        per_layer=per_layer, expert=expert, hc_per_layer=hc_per_layer,
+        moe_total_per_layer=moe_total(False), moe_active_per_layer=moe_active(False),
+        embed=emb, head=head, body_total=body_total, body_active=body_active,
+        total_nonembed=body_total, active_nonembed=body_active,
+        total_incl_embed=body_total + emb + head, active_incl_embed=body_active + emb + head,
+        mtp_total=mtp_total, mtp_active=mtp_active,
+        tid2eid_entries=n_hash * V * K,                    # non-trainable hash-routing lookup tables
+    )
+
+
+def count(c):
+    return count_v4(c) if c.get("model_type") == "deepseek_v4" or "compress_ratios" in c else count_v3(c)
+
+
+# ----------------------------------------------------------------------------------------------
+# KV cache and attention FLOPs per token as a function of context length
+# ----------------------------------------------------------------------------------------------
+def kv_cache_bytes_v3(c, L, kv_bytes=2):
+    """MLA cache: one (kv_lora_rank + rope) latent per token per layer."""
+    per_tok = (c["kv_lora_rank"] + c["qk_rope_head_dim"]) * kv_bytes
+    return c["num_hidden_layers"] * per_tok * L
+
+
+def kv_cache_bytes_v4(c, L, mixed=True):
+    """V4 cache. Per entry: head_dim dims; mixed storage = fp8 for non-rope dims + bf16 for the 64 rope dims
+    (paper sec. 2.3.4). CSA layers additionally cache indexer keys (index_head_dim, FP4)."""
+    hd, rd = c["head_dim"], c["qk_rope_head_dim"]
+    win = c["sliding_window"]
+    entry = (hd - rd) * 1 + rd * 2 if mixed else hd * 2
+    idx_entry = c["index_head_dim"] * (0.5 if mixed else 2)
+    total = 0
+    for r in c["compress_ratios"][: c["num_hidden_layers"]]:
+        n = min(win, L)
+        if r:
+            n += L // r
+        total += n * entry
+        if r == 4:
+            total += (L // r) * idx_entry
+    return total
+
+
+def attn_flops_v3(c, L):
+    """Core attention FLOPs for ONE query token attending to L cached tokens, all layers.
+    Naive (non-absorbed) MLA: QK over (nope+rope) dims and PV over v dims, per head."""
+    H = c["num_attention_heads"]
+    qk = c["qk_nope_head_dim"] + c["qk_rope_head_dim"]
+    return c["num_hidden_layers"] * 2 * H * (qk + c["v_head_dim"]) * L
+
+
+def attn_flops_v4(c, L):
+    """Core attention (+ lightning indexer) FLOPs for ONE query token at context L, all layers.
+    KV entries serve as both key and value: 2*H*hd per entry for QK and 2*H*hd for PV."""
+    H, hd, win = c["num_attention_heads"], c["head_dim"], c["sliding_window"]
+    IH, Ihd, topk = c["index_n_heads"], c["index_head_dim"], c["index_topk"]
+    total = 0
+    for r in c["compress_ratios"][: c["num_hidden_layers"]]:
+        n = min(win, L)
+        if r == 4:
+            n += min(L // r, topk)
+            total += 2 * IH * Ihd * (L // r)          # indexer scores over all compressed keys
+        elif r:
+            n += L // r
+        total += 4 * H * hd * n
+    return total
+
+
+# ----------------------------------------------------------------------------------------------
+# Routed scaling factor, Moonlight's recipe (appendix C, fig. 6) generalised to any scoring function
+# ----------------------------------------------------------------------------------------------
+def gate_scaling_factor(num_experts, topk, score="sigmoid", iters=200_000, seed=0):
+    rng = random.Random(seed)
+    if score == "sigmoid":
+        f = lambda x: 1.0 / (1.0 + math.exp(-x))
+    elif score == "sqrtsoftplus":
+        f = lambda x: math.sqrt(math.log1p(math.exp(x)))
+    else:
+        raise ValueError(score)
+    acc = 0.0
+    for _ in range(iters):
+        p = sorted((f(rng.gauss(0, 1)) for _ in range(num_experts)), reverse=True)[:topk]
+        s = sum(p)
+        acc += 1.0 / math.sqrt(sum((x / s) ** 2 for x in p))
+    return acc / iters
+
+
+# ----------------------------------------------------------------------------------------------
+def fmt(n):
+    if n >= 1e12:
+        return f"{n/1e12:.3f}T"
+    if n >= 1e9:
+        return f"{n/1e9:.3f}B"
+    if n >= 1e6:
+        return f"{n/1e6:.2f}M"
+    return f"{n:,}"
+
+
+def gb(x):
+    return f"{x/2**30:.3f} GiB" if x >= 2**30 else f"{x/2**20:.1f} MiB"
+
+
+def report(name, c):
+    r = count(c)
+    print(f"== {name} ({r['family']}) ==")
+    print(f"  layers={r['layers']}" + (f"  kinds={r['layer_kinds']}  hash_layers={r['n_hash']}" if r["family"] == "deepseek_v4"
+                                       else f"  dense={r['n_dense']} moe={r['n_moe']}"))
+    print(f"  embed={fmt(r['embed'])}  head={fmt(r['head'])}  expert={fmt(r['expert'])}")
+    print(f"  total   non-embedding: {fmt(r['total_nonembed'])}   incl. embed+head: {fmt(r['total_incl_embed'])}")
+    print(f"  active  non-embedding: {fmt(r['active_nonembed'])}   incl. embed+head: {fmt(r['active_incl_embed'])}")
+    if r["mtp_total"]:
+        print(f"  MTP module(s): total {fmt(r['mtp_total'])}, active {fmt(r['mtp_active'])}"
+              f"  -> grand total incl. MTP {fmt(r['total_incl_embed'] + r['mtp_total'])}")
+    if r["family"] == "deepseek_v4":
+        for kind in ("SWA", "CSA", "HCA"):
+            ex = next((x for x in r["per_layer"] if {0: "SWA", 4: "CSA", 128: "HCA"}[x["ratio"]] == kind), None)
+            if ex:
+                parts = ", ".join(f"{k}={fmt(v)}" for k, v in ex["attn_parts"].items())
+                print(f"  attention/{kind} layer: {fmt(ex['attn'])}  [{parts}]")
+        print(f"  mHC per layer: {fmt(r['hc_per_layer'])}   tid2eid entries (non-trainable): {fmt(r['tid2eid_entries'])}")
+    else:
+        print(f"  attention per layer: {fmt(r['attn_per_layer'])}   dense MLP: {fmt(r['dense_mlp'])}")
+    print(f"  MoE per layer: total {fmt(r['moe_total_per_layer'])}, active {fmt(r['moe_active_per_layer'])}")
+    return r
+
+
+def main(argv):
+    if argv:
+        for p in argv:
+            report(Path(p).stem, json.load(open(p)))
+        return
+    ref = HERE / "reference"
+    moonlight = json.load(open(ref / "moonlight_16b_a3b_config.json"))
+    flash = json.load(open(ref / "deepseek_v4_flash_config.json"))
+    pro = json.load(open(ref / "deepseek_v4_pro_config.json"))
+    mv4 = json.load(open(HERE / "config.json"))
+
+    print("#" * 100 + "\n# Validation against published numbers\n" + "#" * 100)
+    r_ml = report("Moonlight-16B-A3B  (published: 15.29B/2.24B non-embed, 16B/3B incl. embed)", moonlight)
+    r_fl = report("DeepSeek-V4-Flash  (published: 284B total, 13B active)", flash)
+    r_pr = report("DeepSeek-V4-Pro    (published: 1.6T total, 49B active)", pro)
+    print("\n" + "#" * 100 + "\n# Moonlight-V4 (this repo)\n" + "#" * 100)
+    r_m4 = report("Moonlight-V4-16B-A3B", mv4)
+
+    print("\n#### Routed scaling factor via Moonlight's recipe (E[1/||p||_2] over renormalised top-k scores)")
+    for (E, K, s) in [(64, 6, "sigmoid"), (64, 6, "sqrtsoftplus"), (256, 6, "sqrtsoftplus"), (384, 6, "sqrtsoftplus"), (256, 8, "sigmoid")]:
+        print(f"  experts={E:4d} topk={K} score={s:12s}: {gate_scaling_factor(E, K, s, iters=20000):.3f}")
+
+    print("\n#### KV cache per sequence (Moonlight: bf16 MLA latent; V4: fp8 non-rope + bf16 rope + fp4 indexer keys)")
+    for L in (8192, 65536, 1048576):
+        print(f"  L={L:>8}: Moonlight {gb(kv_cache_bytes_v3(moonlight, L)):>12} | Moonlight-V4 {gb(kv_cache_bytes_v4(mv4, L)):>12}"
+              f" (bf16-only {gb(kv_cache_bytes_v4(mv4, L, mixed=False)):>12}) | V4-Flash {gb(kv_cache_bytes_v4(flash, L)):>12}")
+
+    print("\n#### Per-token FLOPs at context L (2*active params incl. head  +  core attention [+ indexer])")
+    for L in (8192, 65536, 1048576):
+        lin_ml = 2 * (r_ml["active_nonembed"] + r_ml["head"])
+        lin_m4 = 2 * (r_m4["active_nonembed"] + r_m4["head"])
+        a_ml, a_m4 = attn_flops_v3(moonlight, L), attn_flops_v4(mv4, L)
+        print(f"  L={L:>8}: Moonlight linear {lin_ml/1e9:6.2f} GF + attn {a_ml/1e9:8.2f} GF = {(lin_ml+a_ml)/1e9:8.2f} GF"
+              f" | Moonlight-V4 linear {lin_m4/1e9:6.2f} GF + attn {a_m4/1e9:6.2f} GF = {(lin_m4+a_m4)/1e9:6.2f} GF")
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])

inference_config.json

@@ -0,0 +1,50 @@
+{
+  "vocab_size": 129280,
+  "dim": 1024,
+  "moe_inter_dim": 384,
+  "n_layers": 15,
+  "n_hash_layers": 1,
+  "n_mtp_layers": 0,
+  "n_heads": 16,
+  "n_routed_experts": 32,
+  "n_shared_experts": 1,
+  "n_activated_experts": 6,
+  "score_func": "sqrtsoftplus",
+  "route_scale": 2.436,
+  "swiglu_limit": 10.0,
+  "q_lora_rank": 256,
+  "head_dim": 256,
+  "rope_head_dim": 64,
+  "o_groups": 2,
+  "o_lora_rank": 1024,
+  "window_size": 128,
+  "original_seq_len": 0,
+  "rope_theta": 10000,
+  "rope_factor": 1,
+  "beta_fast": 32,
+  "beta_slow": 1,
+  "index_n_heads": 64,
+  "index_head_dim": 128,
+  "index_topk": 1024,
+  "hc_mult": 4,
+  "hc_sinkhorn_iters": 20,
+  "dtype": "bf16",
+  "compress_rope_theta": 160000,
+  "compress_ratios": [
+    0,
+    0,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4,
+    128,
+    4
+  ]
+}

tokenizer_config.json

@@ -0,0 +1,34 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<｜begin▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "clean_up_tokenization_spaces": false,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "legacy": true,
+  "model_max_length": 1048576,
+  "pad_token": {
+    "__type": "AddedToken",
+    "content": "<｜end▁of▁sentence｜>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sp_model_kwargs": {},
+  "unk_token": null,
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}

training/finite_nanogpt.py

@@ -0,0 +1,22 @@
+"""Bounded variant of NeMo Automodel's NanogptDataset for validation.
+
+``NanogptDataset`` is an infinite stream (its file iterator restarts the shard list forever, which is what
+training wants), so a validation epoch over it never terminates. This subclass stops after ``max_samples``
+samples per iterator (i.e. per data-parallel rank / dataloader worker).
+"""
+from nemo_automodel.components.datasets.llm.nanogpt_dataset import NanogptDataset
+
+
+class FiniteNanogptDataset(NanogptDataset):
+    def __init__(self, *args, max_samples: int = 64, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.max_samples = int(max_samples)
+
+    def __iter__(self):
+        for i, sample in enumerate(super().__iter__()):
+            if i >= self.max_samples:
+                return
+            yield sample
+
+    def __len__(self) -> int:  # type: ignore[override]
+        return self.max_samples

training/init_utils.py

@@ -0,0 +1,90 @@
+"""Helpers for training DeepSeek-V4-architecture models from scratch.
+
+DeepSeek-V4 routes the first ``num_hash_layers`` MoE layers with a fixed token-id -> expert-id table
+(``tid2eid``, Roller et al. 2021 "Hash Layers"). Both NeMo Automodel (``DeepseekV4HashGate``) and
+transformers (``DeepseekV4HashRouter``) zero-initialise that table, which sends every token to expert 0
+when a model is built from a config instead of a checkpoint. Call :func:`fill_hash_tables` after the model
+weights have been initialised (and after FSDP wrapping; the table is a plain buffer).
+"""
+from __future__ import annotations
+
+import torch
+
+
+def balanced_hash_tid2eid(vocab_size: int, n_experts: int, topk: int, seed: int = 0,
+                          device: torch.device | str | None = None) -> torch.Tensor:
+    """Return a ``[vocab_size, topk]`` int64 table mapping each token id to ``topk`` distinct experts.
+
+    Token ids are randomly permuted (seeded) and mapped to a base expert; the remaining ``topk - 1`` experts are
+    spaced ``n_experts // topk`` apart, so every expert receives the same number of (token, slot) assignments
+    up to rounding. This is a uniform-over-token-ids hash; the released V4 tables are not published, so a
+    frequency-balanced variant (weighting by corpus token counts) is left to the user.
+    """
+    assert n_experts >= topk, "need at least topk experts"
+    stride = n_experts // topk
+    g = torch.Generator().manual_seed(seed)
+    base = torch.randperm(vocab_size, generator=g) % n_experts
+    offsets = torch.arange(topk) * stride
+    table = (base.unsqueeze(1) + offsets.unsqueeze(0)) % n_experts
+    return table.to(dtype=torch.int64, device=device)
+
+
+def fill_hash_tables(model: torch.nn.Module, seed: int = 0) -> int:
+    """Fill every ``tid2eid`` buffer in ``model`` (Automodel or transformers V4) with a balanced hash.
+
+    Returns the number of tables filled. Idempotent for a given seed.
+    """
+    n = 0
+    for module in model.modules():
+        table = getattr(module, "tid2eid", None)
+        if not isinstance(table, torch.Tensor):
+            continue
+        n_experts = getattr(module, "n_experts", None) or getattr(module, "num_experts", None)
+        if n_experts is None:
+            raise AttributeError(f"{type(module).__name__} has tid2eid but no n_experts/num_experts attribute")
+        vocab_size, topk = table.shape
+        with torch.no_grad():
+            table.copy_(balanced_hash_tid2eid(vocab_size, int(n_experts), topk, seed + n, table.device))
+        n += 1
+    return n
+
+
+def check_hash_tables(model: torch.nn.Module) -> dict:
+    """Return per-table expert-load statistics (max/mean assignment count) to confirm balance."""
+    stats = {}
+    for name, buf in model.named_buffers():
+        if name.endswith("tid2eid"):
+            counts = torch.bincount(buf.flatten().cpu())
+            stats[name] = {"min": int(counts.min()), "max": int(counts.max()), "mean": float(counts.float().mean())}
+    return stats
+
+
+def init_hyper_connections(model: torch.nn.Module, std: float = 0.02) -> int:
+    """Initialise the mHC mixer parameters (``fn`` / ``base`` / ``scale`` and the head's ``hc_fn`` /
+    ``hc_base`` / ``hc_scale``) for from-scratch training.
+
+    NeMo Automodel constructs these with ``torch.empty`` and its ``init_weights`` leaves them untouched
+    (it only ever loads them from a checkpoint). transformers' ``_init_weights`` uses
+    normal(0, initializer_range) for the projection, zeros for the static bias and ones for the gates, and that
+    is what this function applies. Works on plain tensors and on FSDP2 DTensors (operates on the local shard).
+    Returns the number of tensors initialised.
+    """
+    n = 0
+
+    def _local(p):
+        return p.to_local() if hasattr(p, "to_local") else p
+
+    with torch.no_grad():
+        for name, p in model.named_parameters():
+            leaf = name.rsplit(".", 1)[-1]
+            parent = name.rsplit(".", 2)[-2] if name.count(".") >= 1 else ""
+            if parent in ("attn_hc", "ffn_hc") and leaf in ("fn", "base", "scale") or leaf in ("hc_fn", "hc_base", "hc_scale"):
+                t = _local(p)
+                if leaf in ("fn", "hc_fn"):
+                    torch.nn.init.normal_(t, mean=0.0, std=std)
+                elif leaf in ("base", "hc_base"):
+                    torch.nn.init.zeros_(t)
+                else:
+                    torch.nn.init.ones_(t)
+                n += 1
+    return n

training/prepare_data.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+"""Tokenise local FineWeb parquet files into NeMo Automodel ``NanogptDataset`` shards.
+
+Shard format (nemo_automodel/components/datasets/llm/nanogpt_dataset.py, "new" format):
+  int32[256] header: [MAGIC=278895051, VERSION=1, num_tokens, bytes_per_token], then uint32 tokens.
+  A sibling ``.bos.idx`` file holds int32 positions of BOS tokens (used when ``align_to_bos=True``).
+Every document is prefixed with the tokenizer's BOS token (id 0 for the DeepSeek-V4 tokenizer), no EOS is added,
+matching Automodel's tools/nanogpt_data_processor.py.
+
+Example (200M train tokens + 4M validation tokens, ~0.8 GB on disk):
+  python3 prepare_data.py --parquet "/path/to/fineweb/*.parquet" \
+      --tokenizer akoumpa/Moonlight-V4-1B-h16d256 --out /path/to/data \
+      --train-tokens 200M --val-tokens 4M --shard-tokens 50M
+"""
+import argparse
+import glob
+import os
+import time
+
+import numpy as np
+import pyarrow.parquet as pq
+from transformers import AutoTokenizer
+
+MAGIC, VERSION, HEADER_SIZE = 278895051, 1, 256
+
+
+def parse_tokens(s: str) -> int:
+    s = s.strip().upper()
+    mult = {"K": 10**3, "M": 10**6, "B": 10**9}.get(s[-1], 1)
+    return int(float(s[:-1]) * mult) if s[-1] in "KMB" else int(s)
+
+
+class ShardWriter:
+    def __init__(self, path: str, bos_id: int):
+        self.path, self.bos_id, self.n = path, bos_id, 0
+        self.fp = open(path, "wb")
+        self.idx = open(path[:-4] + ".bos.idx", "wb")
+        self.fp.write(np.zeros(HEADER_SIZE, dtype=np.int32).tobytes())
+
+    def write(self, toks: np.ndarray):
+        pos = self.n
+        self.fp.write(toks.astype(np.uint32).tobytes())
+        self.idx.write((pos + np.flatnonzero(toks == self.bos_id)).astype(np.int32).tobytes())
+        self.n += toks.size
+
+    def close(self):
+        header = np.array([MAGIC, VERSION, self.n, 4] + [0] * (HEADER_SIZE - 4), dtype=np.int32)
+        self.fp.seek(0); self.fp.write(header.tobytes()); self.fp.close(); self.idx.close()
+
+
+def doc_stream(files, text_col="text", batch_rows=1024):
+    for f in files:
+        pf = pq.ParquetFile(f)
+        for batch in pf.iter_batches(batch_size=batch_rows, columns=[text_col]):
+            yield batch.column(text_col).to_pylist()
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--parquet", required=True, help="glob of parquet files with a 'text' column")
+    ap.add_argument("--tokenizer", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--train-tokens", default="200M", type=parse_tokens)
+    ap.add_argument("--val-tokens", default="4M", type=parse_tokens)
+    ap.add_argument("--shard-tokens", default="50M", type=parse_tokens)
+    ap.add_argument("--max-doc-tokens", default=32768, type=int)
+    args = ap.parse_args()
+
+    files = sorted(glob.glob(os.path.expanduser(args.parquet)))
+    assert files, f"no parquet files match {args.parquet}"
+    out = os.path.expanduser(args.out)
+    os.makedirs(out, exist_ok=True)
+    tok = AutoTokenizer.from_pretrained(args.tokenizer)
+    bos = tok.bos_token_id
+    assert bos is not None
+
+    # plan: validation shard first, then train shards
+    plan = [("val", args.val_tokens, args.val_tokens)] + [("train", args.train_tokens, args.shard_tokens)]
+    stream = doc_stream(files)
+    t0, total, n_docs = time.time(), 0, 0
+    for split, budget, shard_size in plan:
+        written, shard_id, writer = 0, 0, None
+        while written < budget:
+            if writer is None:
+                writer = ShardWriter(os.path.join(out, f"fineweb_{split}_{shard_id:04d}.bin"), bos)
+            try:
+                texts = next(stream)
+            except StopIteration:
+                break
+            ids = tok(texts, add_special_tokens=False)["input_ids"]
+            buf = np.concatenate([np.array([bos] + d[: args.max_doc_tokens - 1], dtype=np.uint32) for d in ids])
+            writer.write(buf)
+            written += buf.size; total += buf.size; n_docs += len(texts)
+            if writer.n >= shard_size:
+                writer.close(); writer = None; shard_id += 1
+            if n_docs % (1024 * 20) == 0:
+                print(f"[{split}] {written/1e6:8.1f}M / {budget/1e6:.0f}M tokens, {n_docs} docs, {total/(time.time()-t0)/1e6:.2f} Mtok/s", flush=True)
+        if writer is not None:
+            writer.close()
+        print(f"[{split}] done: {written/1e6:.1f}M tokens in {shard_id + 1} shard(s)")
+    print(f"total {total/1e6:.1f}M tokens, {n_docs} docs, {time.time()-t0:.0f}s -> {out}")
+    print(open(os.path.join(out, "README.txt"), "w").write(
+        f"tokenizer={args.tokenizer} bos={bos} train_tokens={args.train_tokens} val_tokens={args.val_tokens} source={args.parquet}\n") and "")
+
+
+if __name__ == "__main__":
+    main()

training/pretrain.yaml

@@ -0,0 +1,113 @@
+# akoumpa/Moonlight-V4-1B-h16d256: from-scratch pre-training recipe for NeMo Automodel (native deepseek_v4 model, FSDP2 over 2 GPUs).
+#
+#   torchrun --nproc-per-node 2 train.py --config pretrain.yaml        # any key can be overridden: --step_scheduler.max_steps 100
+#
+# train.py seeds the hash-routing table and the mHC mixers on a fresh start (see README, "Training from scratch").
+# Replace the /path/to/... placeholders; data shards come from prepare_data.py.
+
+recipe: TrainFinetuneRecipeForNextTokenPrediction
+
+step_scheduler:
+  global_batch_size: 32          # sequences per optimizer step (32 x 2048 = 65k tokens)
+  local_batch_size: 2            # sequences per GPU per micro-step; keep equal to dataloader.batch_size below (4 OOMs on 48 GB)
+  ckpt_every_steps: 250
+  val_every_steps: 100
+  gc_every_steps: 10
+  num_epochs: 1
+  max_steps: 3000                # ~200M tokens = one pass over the prepared shards; ~10.6 s/step measured -> ~9 h
+
+dist_env:
+  backend: nccl
+  timeout_minutes: 15
+
+rng:
+  _target_: nemo_automodel.components.training.rng.StatefulRNG
+  seed: 1234
+  ranked: true
+
+model:
+  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_config
+  config:
+    _target_: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config.from_pretrained
+    pretrained_model_name_or_path: akoumpa/Moonlight-V4-1B-h16d256
+  backend:
+    _target_: nemo_automodel.components.models.common.BackendConfig
+    attn: eager                  # see README section "TileLang kernels" before switching to 'tilelang'
+    linear: torch
+    rms_norm: torch_fp32
+    rope_fusion: false
+    dispatcher: torch
+    experts: torch_mm
+    enable_hf_state_dict_adapter: false
+
+distributed:
+  strategy: fsdp2
+  dp_size: 2
+  tp_size: 1
+  cp_size: 1
+  pp_size: 1
+  ep_size: 1
+  sequence_parallel: false
+  activation_checkpointing: false
+  moe:
+    # Automodel wraps DSV4's fp32 parameter holders (attention sinks, compressor ape, mHC mixers) as their own
+    # FSDP2 units; resharding them right after their forward frees the tensor before attention uses it.
+    # The shipped DeepSeek-V4 fine-tune recipes disable resharding for the same reason.
+    reshard_after_forward: false
+    wrap_outer_model: false
+
+checkpoint:
+  enabled: true
+  checkpoint_dir: /path/to/checkpoints/Moonlight-V4-1B-h16d256
+  model_save_format: torch_save
+  save_consolidated: false
+
+loss_fn:
+  # Logits-based loss (as in Automodel's own DeepSeek-V4 recipes). FusedLinearCrossEntropy cannot be used here:
+  # the model keeps lm_head in fp32 while hidden states are bf16 and the fused Triton kernel requires one dtype.
+  _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy
+
+dataset:
+  _target_: nemo_automodel.components.datasets.llm.nanogpt_dataset.NanogptDataset
+  file_pattern: /path/to/data/fineweb_train_*.bin
+  seq_len: 2048
+  shuffle_files: true
+  align_to_bos: false
+  bos_token: 0
+
+dataloader:
+  _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  # The recipe passes no batch size to the DataLoader for IterableDatasets (it only uses local_batch_size for the
+  # accumulation count), so without this line every micro-batch is a single sequence.
+  batch_size: 2
+  shuffle: false
+  collate_fn: nemo_automodel.components.datasets.utils.default_collater
+
+validation_dataset:
+  # NanogptDataset is an infinite stream; FiniteNanogptDataset (finite_nanogpt.py, importable because train_1b.py
+  # puts this directory on sys.path) stops after max_samples sequences per rank.
+  _target_: finite_nanogpt.FiniteNanogptDataset
+  file_pattern: /path/to/data/fineweb_val_0000.bin
+  max_samples: 64                                                             # 64 x 2048 tokens per rank per validation
+  seq_len: 2048
+  shuffle_files: false
+  align_to_bos: false
+  bos_token: 0
+
+validation_dataloader:
+  _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  batch_size: 2
+  shuffle: false
+  collate_fn: nemo_automodel.components.datasets.utils.default_collater
+
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 4.0e-4
+  betas: [0.9, 0.95]
+  eps: 1.0e-8
+  weight_decay: 0.1
+
+lr_scheduler:
+  lr_decay_style: cosine
+  lr_warmup_steps: 200
+  min_lr: 4.0e-5

training/train.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+"""Launcher for pretrain.yaml (NeMo Automodel recipe).
+
+Thin wrapper around NeMo Automodel's TrainFinetuneRecipeForNextTokenPrediction that, on a fresh start
+(no checkpoint to resume from), seeds two things Automodel's from-config path leaves at zero /
+uninitialised because it only ever loads them from released checkpoints:
+  * the hash-routing table ``tid2eid`` of the first ``num_hash_layers`` MoE layers (balanced hash), and
+  * the mHC mixer parameters (``attn_hc``/``ffn_hc``/``hc_head``: normal(0, 0.02) projection, zero bias, unit gates).
+
+Usage:  torchrun --nproc-per-node 2 train.py --config pretrain_1b.yaml [--dotted.key value ...]
+"""
+import os
+import pathlib
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from init_utils import check_hash_tables, fill_hash_tables, init_hyper_connections  # noqa: E402
+
+import torch  # noqa: E402
+from torch.distributed.fsdp import FSDPModule  # noqa: E402
+
+import nemo_automodel  # noqa: E402
+from nemo_automodel.components.config._arg_parser import parse_args_and_load_config  # noqa: E402
+from nemo_automodel.recipes.llm.train_ft import TrainFinetuneRecipeForNextTokenPrediction  # noqa: E402
+
+
+def will_resume(cfg) -> bool:
+    """Mirror the recipe's auto-resume rule: explicit restore_from, or a checkpoint_dir with a training log."""
+    if cfg.get("checkpoint.restore_from", None):
+        return True
+    ckpt_dir = cfg.get("checkpoint.checkpoint_dir", None)
+    return bool(ckpt_dir) and (pathlib.Path(ckpt_dir) / "training.jsonl").exists()
+
+
+def keep_fp32_units_unsharded(model) -> int:
+    """Automodel wraps DSV4's fp32 tensors (attention sinks, compressor position biases, mHC mixers, lm_head)
+    as separate FSDP2 units whose forward *returns the parameter itself*. If such a unit reshards right after
+    its forward, the returned tensor is freed before attention consumes it (``setStorage ... storage of size 0``).
+    Keep those (tiny) units unsharded between forward and backward."""
+    n = n_resharding = 0
+    for m in model.modules():
+        if isinstance(m, FSDPModule) and {p.dtype for p in m.parameters()} == {torch.float32}:
+            try:  # diagnostic only: was this unit going to reshard after forward?
+                group = m._get_fsdp_state()._fsdp_param_group
+                n_resharding += int(group is not None and group.post_forward_mesh_info is not None)
+            except Exception:
+                pass
+            m.set_reshard_after_forward(False, recurse=False)
+            n += 1
+    if int(os.environ.get("RANK", "0")) == 0:
+        print(f"[train] fp32 FSDP units: {n} total, {n_resharding} were set to reshard after forward")
+    return n
+
+
+def main():
+    cfg = parse_args_and_load_config(os.path.join(os.path.dirname(os.path.abspath(__file__)), "pretrain.yaml"))
+    resume = will_resume(cfg)
+    recipe = TrainFinetuneRecipeForNextTokenPrediction(cfg)
+    recipe.setup()
+    rank = int(os.environ.get("RANK", "0"))
+    if rank == 0:
+        print(f"[train] nemo_automodel from {os.path.dirname(nemo_automodel.__file__)}")
+    n_units = sum(keep_fp32_units_unsharded(part) for part in recipe.model_parts)
+    if rank == 0:
+        print(f"[train] fp32 FSDP units kept unsharded after forward: {n_units}")
+    if resume:
+        if rank == 0:
+            print("[train] resuming from checkpoint: keeping hash tables / mHC mixers from the checkpoint")
+    else:
+        seed = int(cfg.get("rng.seed", 0) or 0)
+        std = float(getattr(recipe.model_parts[0].config, "initializer_range", 0.02))
+        for part in recipe.model_parts:
+            n_tab = fill_hash_tables(part, seed=seed)
+            n_hc = init_hyper_connections(part, std=std)
+            if rank == 0:
+                print(f"[train] fresh start: filled {n_tab} hash table(s), initialised {n_hc} mHC tensors; "
+                      f"table balance {check_hash_tables(part)}")
+    recipe.run_train_validation_loop()
+
+
+if __name__ == "__main__":
+    main()