DeepSeek-v3-Lite

Imported from atandra2000/DeepSeek-v3-Lite (AGENTS.md). Install upstream with npx skills add atandra2000/DeepSeek-v3-Lite. Copyright stays with the author.

AGENTS.md — DeepSeek-v3-Lite

CRITICAL RULE: You must also read, understand, and strictly obey all higher-level rules: the workspace LLM/AGENTS.md (shared-data pipeline rules, repo boundaries, cross-project invariants) and the top-level CoreProjects/AGENTS.md / CoreProjects/.agents/AGENTS.md. Those instructions apply globally to all projects; this file wins only on DeepSeek-v3-Lite-specific conflicts.

Quick checks (run before claiming work is done)

cd LLM/DeepSeek-v3-Lite
python3 -m pytest tests/ -q                 # 199 tests: 189 pass + 10 GPU-gated skips
python3 tests/test_doc_refs.py              # doc↔code symbol anchors
python3 scripts/check_docs.py               # links + stale-pattern lint
# Tiny end-to-end smoke on small hardware:
#   configs/pretrain_1650_2m.yaml (~2M params, GTX 1650 4GB-class)

Project: LLM/DeepSeek-v3-Lite/ · Type: faithful V3 reproduction Scale: ~412M params (411.6M deduped base / 418.7M with MTP) · 8.4B tokens (planned) · ~30–45h on A100 80GB (estimated; unverified) Stack: PyTorch 2.x, TF32, torch.compile(max-autotune), FA2, custom Triton kernels for MoE dispatch + MLA attention, dataclasses

Faithful from-scratch reimplementation of the full DeepSeek-V3 architecture: every V3 component implemented end-to-end (no stubs).

1. Subagent: deepseek-v3-engineer

Trigger: "Explain the MLA absorption trick", "Why does DeepSeek-V3 use biased-sigmoid MoE?", "How does speculative decoding work with MTP?", "Debug NaN in DeepSeek training", "Set up μP for 422M.", "Tune the MoE Triton kernel BLOCK_T", "Profile MLA flash kernel on A100."

System prompt: You are a senior engineer maintaining DeepSeek-v3-Lite. You know the DeepSeek-V2/V3 papers cold and the codebase even better. You also know Triton 3.0+ kernel idioms (online softmax, grouped GEMM, flash-decoding split-K) and can reason about HBM bandwidth vs compute trade-offs.

Architecture (18 layers):

• 2 dense layers (MLA + SwiGLU).
• 16 MoE layers (MLA + DeepSeekMoE).
• vocab 100,018, dim 768, 12 heads.
• RoPE θ=10K, factor 1.0 (no scaling at training length).

Component map:

• models/mla.py — kv_lora_rank=192, qk_rope_head_dim=24, absorption trick, YaRN scaling (gated by rope_factor > 1.0; canonical config uses rope_factor=1.0 so YaRN is dormant), KV cache. attn_impl supports "sdpa", "manual", and "triton" (the Triton path delegates to models/mla_triton.py). Technical deep-dive in docs/concepts/attention-and-precision.md.
• models/mla_triton.py — fused FA2-style MLA attention: on-the-fly K_nope / V materialisation from the latent + wkv_b, RoPE applied inside the kernel, online softmax. Backward via re-compute (FA2 pattern). Optional import gated on HAS_TRITON; pure-PyTorch fallback in models/mla.py if triton is missing or attn_impl != "triton".
• models/moe.py — AuxLossFreeGate + DeepSeekMoE: 20 routed (top-4) + 1 shared, stacked bmm dispatch, dynamic bias updates. moe_dispatch supports "stacked" (default, PyTorch) and "triton_grouped" (opt-in; the Triton path delegates to models/moe_triton.py). Any other value falls back to stacked. Expert weights are re-stacked every forward — never cache _stacked_w* across optimizer steps.
• models/moe_triton.py — fused grouped-GEMM SwiGLU kernel over the sorted-token layout: one launch handles all 20 routed experts per layer. Replaces the for e in range(E): … index_add … Python loop. Backward via re-compute (FA2 pattern). Optional import gated on HAS_TRITON; pure-PyTorch stacked / grouped fallbacks in models/moe.py otherwise.
• models/mtp.py — depth=1, shared output head, speculative-decoding support (inference/speculative.py). Triton-free; MTP attention uses nn.MultiheadAttention (standard SDPA).
• models/transformer.py — top-level wiring; routes attn_impl and moe_dispatch from config to the right backend.

Triton kernel contract (both kernels):

• import triton is optional at import time. Both modules wrap the import in try/except ImportError and set HAS_TRITON = False on failure. Never raise at module-load if triton is missing.
• Each kernel file ships a pure-PyTorch reference in the same file (or a sibling _reference.py) that the tests use on CPU/Mac.
• The kernel is exposed as a torch.autograd.Function with separate forward and backward static methods. The forward saves only the minimum needed for backward to re-compute (FA2 re-compute pattern). Do not save full intermediates to HBM unless the kernel explicitly requires it.
• All GEMM accumulators inside kernels are fp32, inputs and outputs are BF16 (or whatever the module's parameter dtype is).
• Block sizes (BLOCK_T, BLOCK_D, BLOCK_I, BLOCK_V, etc.) are declared tl.constexpr and tuned with @triton.autotune over a small grid of (block size, num_warps, num_stages) configs. Pre-warm at __init__ to amortise the first-call compile cost.
• New Triton kernels must add a unit test file under tests/test_<kernel>_triton.py that (a) compares the Triton path against the PyTorch reference within atol=1e-2 for BF16, (b) runs torch.autograd.gradcheck on a float32 tiny config, and (c) runs a shape / dtype / NaN-finite test on a random input.

Training:

• TF32 + torch.compile(max-autotune) + FA2 + custom Triton kernels (MoE dispatch, MLA attention) + μP LR scaling (8.07e-4 @ 418.7M with MTP).
• FP32 AdamW master weights + gradient checkpointing.
• NaN guard with checkpoint rollback.

Inference:

• inference/generate.py — interactive generation.
• inference/speculative.py — MTP-based speculative decoder (~0.8 acceptance, up to 2× throughput).
• The Triton MLA kernel also accelerates inference (prefill + decode via flash-decoding split-K). The Triton MoE kernel is training-only for now — the shared-expert path and the per-step routing tensor shapes differ slightly between training (gMA) and inference (sparse decode), and the kernel is not validated for the decode case.

Data: Universal 8.0B-token pipeline (lives at LLM/shared_data/ in the workspace umbrella; this project imports it via sys.path in data/prepare_data.py — it is not vendored here). Mixture (canonical 7 sources in LLM/shared_data/config/mixture.yaml): fineweb-edu 0.40 / dclm-baseline 0.15 / the-stack-v2-python 0.15 / the-stack-v2-jupyter 0.05 / openmath 0.10 / arxiv 0.10 / cosmopedia 0.05. Tokenized with deepseek-ai/deepseek-coder-v2-lite tokenizer (vocab 100,018). See docs/training.md (data pipeline) and docs/concepts/data-pipeline.md.

Doc routing: implementation questions → models/*.py first; theory → docs/<component>.md; API reference (signatures, shape contracts, defaults) → docs/references/R<n>_*.md; operational procedures → docs/guides/G<n>_*.md; invariants → tests/ + docs/concepts/kernels-and-ops.md. MLA deep-dive → docs/concepts/attention-and-precision.md. Doc↔code alignment is machine-checked by tests/test_doc_refs.py (every code-symbol citation — file path + class/method — must resolve; line anchors and JIT-kernel citations are banned) and scripts/check_docs.py (links, paths, stale patterns) — both run in CI; stale docs fail the build.

Configs: configs/pretrain_a100_422m.yaml (canonical 422M A100 recipe). Two new optional config keys: attn_impl (default "sdpa", set to "triton" to opt in) and moe_dispatch (default "stacked", set to "triton_grouped" to opt in). Master kill-switch env-var: ENABLE_TRITON_KERNELS=0 (default, no Triton) or =1 (allow per-config opt-in). Tests on CPU/Mac run with the env-var unset and the default attn_impl / moe_dispatch values, so the 199 pytest tests keep passing (189 pass + 10 GPU-gated skips) without any CUDA dependency.

Hard rules:

1. Raw PyTorch by default; custom Triton kernels are first-party for exactly two hot paths. Bulk of the codebase (RMSNorm, SwiGLU, embeddings, LM head, loss, gate, MLA SDPA, MTP, inference) stays raw PyTorch. No HuggingFace Trainer, no Lightning, no high-level wrappers. The two sanctioned Triton paths are:
   • MoE routed-expert dispatch — models/moe_triton.py.
   • MLA attention materialisation — models/mla_triton.py. No other component gets a custom kernel without updating this file and docs/concepts/kernels-and-ops.md.
2. Hardware Optimization: Maximize hardware utilization. For the two sanctioned Triton paths, target ≥ 1.5× speedup over the raw-PyTorch path in scripts/microbench_a100.py; below that, do not enable by default.
3. Always preserve the AuxLossFreeGate bias-update mechanism. update_bias runs out-of-band in Python every bias_update_every steps; the Triton MoE kernel only fuses the forward. Do not move the bias into the kernel.
4. Always read docs/concepts/attention-and-precision.md before answering MLA questions — it is the authoritative reference for theory and this-repo implementation. The Triton MLA kernel re-expresses the same math; the doc still describes the algebra.
5. Always verify the embedding dim matches vocab_size (100,018) — the tokenizer has unusual byte_fallback tokens.
6. Never disable the NaN guard without explicit user consent. The Triton kernels must be NaN-finite on a 50-step warmup sweep before they are enabled by default; if any kernel produces NaNs, gate it off and fix the kernel before re-enabling.
7. Never let a Triton kernel silently fall back to the raw-PyTorch path during a default-config training run. The opt-in is explicit (attn_impl: "triton" or moe_dispatch: "triton_grouped" in YAML, plus ENABLE_TRITON_KERNELS=1 env-var). If the kernel fails to compile or throws at runtime, the run must surface a clear error, not a silent fallback.
   • The master env-var guard lives in models/_triton_dispatch.py:enforce_triton_env_var and is called from both Transformer.__init__ (model construction) and Pretrainer.__init__ (training entry point). When the env-var is unset, any Triton dispatch key in the config is force-backed to its PyTorch default with a single startup warning, not 32 per-layer warnings at first step.
8. Always add a unit test in tests/ for any new Triton kernel path. The test must run on CPU (using the pure-PyTorch reference) without triton installed. GPU-only behaviour is gated behind @pytest.mark.gpu and is auto-skipped on CPU-only machines.
9. Concise comments only. Docstrings and inline comments must justify non-obvious code, not restate it. A docstring is at most three short lines unless the function is a public API. Inline comments appear only when the code itself is opaque. Verifiable targets per file:
   • Public function docstring: ≤ 3 lines, or one short paragraph.
   • Module docstring: ≤ 6 lines.
   • Inline comment density: ≤ 1 comment per ~10 lines of code on average; comments that say what the next line does (# compute x, # loop over rows) are forbidden.
   • Section banners (# ---- ... ----) are reserved for the top level of a file (≤ 3 per file) and inside kernels to delimit named algorithm phases. Violations are reviewable on wc -l <file> and grep -c '^[[:space:]]*#' <file>.

Known issues:

• Full 8.4B-token run not yet started.
• Speculative decoding acceptance rate measured at ~0.8 on smoke tests.
• Triton kernels (MoE + MLA) are implemented but not default-on — opt in via attn_impl: "triton" / moe_dispatch: "triton_grouped" in YAML plus ENABLE_TRITON_KERNELS=1. A100 benchmark validation (≥1.5× MoE speedup, ≥25% full-step reduction) is still open. tests/test_mla_triton.py covers MLA kernel reference + GPU numerics; full-model test_sdpa_and_triton_agree is GPU-gated in tests/test_mla_triton.py. All 199 tests pass on the default attn_impl="sdpa", moe_dispatch="stacked" config without Triton.
• The Triton MoE kernel cannot run the canonical config (moe_inter_dim=384, dim=768 exceed the 256 register cap; triton_grouped raises ValueError and auto-falls back to stacked). The ≥1.5× MoE benchmark target is therefore blocked at canonical dims until the kernel is extended past I=384/D=768 — it is only validated at smoke-config dims (I, D ≤ 256).
• Actual deduped param count is 411.6M base / 418.7M with MTP — the "422M" in the config filename is historical nominal.
• Triton requires SM_80+ (A100, H100, RTX 4090, etc.) and Linux. macOS and Windows are unsupported for the Triton paths; the HAS_TRITON = False fallback keeps the codebase importable and testable on those platforms, but the kernel is not exercised.