CodeVetter · Repo Unpack · Data flow
deepseek-ai/DeepSeek-V3
Data moves along three pipelines: (1) offline checkpoint conversion — HF safetensors are renamed via the convert.py mapping table, column/row-sharded across mp ranks, and experts are partitioned contiguously; (2) startup — JSON config -> ModelArgs -> Transformer construction -> safetensors load_model into each rank's shard; (3) inference — token IDs flow through vocab-parallel embedding, N Blocks of RMSNorm->MLA->residual and RMSNorm->(MLP|MoE)->residual, a final norm + head producing logits that are all-gathered across ranks, then sampled and appended back into the token buffer. FP8 support threads through as blockwise scale tensors attached to Linear weights.
- Config data flows: --config JSON -> json.load -> ModelArgs(**dict); every architecture knob (dim, n_layers, n_routed_experts, dtype, scale_fmt, rope params) is a ModelArgs field with defaults matching the 16B model.
- Weight-name conversion is driven by the convert.py `mapping` dict, which renames HF keys (e.g. self_attn.q_a_proj -> attn.wq_a, mlp.gate_proj -> ffn.w1, weight_scale_inv -> scale, e_score_correction_bias -> bias) and records the shard axis (0 or 1) per tensor type.
- Expert sharding in convert.py assigns expert i to rank floor(i / n_local_experts); non-expert tensors with a shard axis are narrowed into mp contiguous slices, asserted divisible by mp.
- At startup each rank loads only its shard file model{rank}-mp{world_size}.safetensors after a 2-token warmup generate call that exercises the model before weights are loaded.
- Token flow: input IDs -> ParallelEmbedding (each rank embeds only its vocab slice, masks others to zero, then all_reduce) -> h [bsz, seqlen, dim].
- MLA forward computes q via wq or wq_a->q_norm->wq_b, splits q into q_nope/q_pe, applies RoPE to q_pe and to the k_pe slice of wkv_a output, then either writes k/v caches (naive) or caches normalized kv latent + pe (absorb).
- Attention scores are softmaxed in float32 then cast back to the input dtype; the causal mask is an additive -inf upper-triangular tensor built only when seqlen > 1.
- MoE forward flattens tokens to [tokens, dim], routes via Gate (scores -> optional group masking -> top-k indices -> gather original scores -> normalize -> route_scale), then loops over local experts scattering weighted outputs, adds shared_experts output, and all_reduces across ranks.
- Final logits are produced by ColumnParallelLinear head and reassembled across ranks with dist.all_gather + concat along the vocab dimension.
- The decode loop writes prompt tokens into a -1-padded buffer, iterates cur_pos from min(prompt_lens) to total_len calling model.forward on the slice [prev_pos:cur_pos], and forces prompt positions via prompt_mask so prefilled tokens override sampled ones.
- FP8 quantization data flow: act_quant produces (float8_e4m3fn tensor, fp32 scales) per 128-element block; fp8_gemm multiplies scaled blocks accumulating in fp32; weight scales live as a .scale Parameter on each Linear whose weight element_size() == 1.
- fp8_cast_bf16 reads model.safetensors.index.json weight_map to locate tensors across shard files, keeps a 2-file LRU cache of loaded safetensors on cuda, dequantizes each FP8 weight with weight_dequant, and strips *_scale_inv entries from the output index.
- EOS handling: finished flags update only for generated positions (not prompt), generation breaks early when all sequences finish, and each completion is truncated at its first eos_id.