DeepSeek-V4.1-Flash

DeepSeek-V4.1 Flash uses a 40-layer causal encoder/decoder Mixture-of-Experts (MoE) backbone with 384 routed experts, one shared expert, and six selected experts per token. Its CSA2 attention shares compressed KV and index state between Full, Reuse, and Reindex layers. Single-pass mHC carries four residual streams, and two trainable Engram tables at layers 1 and 14 provide tokenizer-dependent N-gram memory. The approximately 552B backbone is accompanied by 196B Engram parameters.

Model Reference

Model Architecture

Fine-Tuning

The recipe trains all 40 text-backbone layers and both Engram tables. Dense parameters use generic FSDP2; routed experts use the shared NeMo AutoModel MoE implementation with native HybridEP dispatch and combine. Engram reuses the Qwen3.8 Flash Next row-owner embedding, distributing each table across the same world ranks without a full-table all-gather. Its gradients return to the row owners and receive the owner divisor once before clipping.

The configuration requests 100 updates with validation every 50 updates, activation checkpointing, and reshard-after-forward. It uses no gradient accumulation or additional expert FSDP shard axis, and requests a one-hour allocation. W&B logging is disabled by default; set wandb.enable: true to enable online logging. TE FusedAdam retains FP32 master-weight precision through int16 remainders for BF16 parameters.

Place all 64 ranks in one verified NVLink domain with static rank ordering. The model and recipe use attn: tilelang, linear: torch, rms_norm: torch_fp32, experts: torch_mm, and dispatcher: hybridep. Set model.backend.rms_norm: te to use Transformer Engine RMSNorm for the text backbone, attention projections, compressor, and indexer. This preserves parameter names and dtypes; the default retains eager FP32 normalization. TileLang and TileKernels are the optional dependencies also used by DeepSeek V4. HybridEP provides both token dispatch and expert-output combine. Router correction biases remain fixed during fine-tuning through gate_bias_update_factor: 0.

FSDP mixed precision must use param_dtype: bfloat16, reduce_dtype: float32, output_dtype: null, and cast_forward_inputs: false. Set lm_head_precision: float32. These settings preserve FP32 mHC coefficients across module boundaries; the language-model head returns FP32 logits.

See the Slurm Launcher Guide and Installation Guide for setup. Compilation caches and communication prewarming must match the backend, topology, and local input shapes.

Context Parallelism

Unpacked, right-padded text supports context parallelism with contiguous sequence shards. Set distributed.cp_size and retain the same global batch size when comparing CP sizes. The model’s sharder pads the global sequence so each rank owns complete compression groups and preserves global position IDs. Shift labels before sharding and use the recipe’s external loss; direct labels= loss is unavailable when CP is active.

Queries and hidden activations remain local. Window KV, newly compressed KV, and index keys are gathered across the CP group; reused compressed KV retains its gradient history across layers. Engram gathers integer token history for N-gram hashing, then looks up only the local tokens through the existing row owners. CP does not gather full Engram tables. TP, PP, packed documents, and image inputs remain outside this CP path.

The Tulu3 CP recipe uses 4,096-token sequences, TE RMSNorm, EP64, and 64 Engram owners. At global batch 64 and local batch 1, CP1 uses one microbatch per update and CP8 uses eight. Both runs must use the same prepared Arrow data, tokenizer, initial checkpoint, optimizer, and seed. The provided chat template follows the released chat-mode message format and masks loss to assistant content.

The Tulu3 CP recipes explicitly set clip_grad_norm.backend: triton. This uses the shared fused FP64 gradient-norm reduction from main and avoids gradient-sized norm temporaries for contiguous CUDA owner shards. Set clip_grad_norm.backend: te to select the shared Transformer Engine reducer.

H100 80 GB Systems

The GB200 recipes place 748B trainable parameters on 64 GPUs, about 117 GB per GPU for BF16 weights, BF16 gradients, BF16 moments, and int16 master remainders. On 80 GB H100 parts the same recipes need twice the ranks: with ep_size: 64 on 128 or more GPUs the mesh derives an additional ep_shard FSDP axis over the routed experts, which halves the per-rank expert state. The variants below were validated on 8-GPU H100 nodes with InfiniBand, 100 updates each, using the NGC nemo-automodel:26.08 container plus this repository at e2c47c5b (torch 2.13, NCCL 2.30.5).

The Tulu3 variants load allenai/tulu-3-sft-mixture from the Hugging Face Hub through ChatDataset with tulu3_chat_template.jinja: padded 4,096-token rows for CP8 and 32,768-token THD packs, built at startup, for the packed recipe. For offline runs, prepare_tulu3_data.py writes the same padded and prepacked layouts to disk.

Three points apply to every H100 variant:

• Set NCCL_PXN_DISABLE=1, or NCCL_NET_PLUGIN=none NCCL_NET=IB, on multi-node InfiniBand systems. With the container’s default NCCL network plugin and PXN both active, the Engram row-owner all-to-all received byte-shifted token IDs from the first node’s ranks and the recipe stopped in _validate_received_ids during the first forward, with NCCL 2.30.5 and 2.30.7 alike, and with both HPC-X 2.50 plugins the image ships (the default Spectrum-X spcx and IBext). Either switch alone avoids it: disabling PXN with the plugin kept, or bypassing the plugin with PXN kept, both completed 100 updates with matching losses. See issue #3959.
• Gradient accumulation needs defer_fsdp_grad_sync: false. The default keeps unsharded parameters and gradients of every FSDP unit resident across microbatches, about 33 GB here regardless of dp, which fits GB200 but not 80 GB parts. With per-microbatch gradient sync and reshard_after_forward: true, torch 2.13 hits an activation-checkpoint recompute mismatch on the expert all-gather copy-out (issue #3958). moe.reshard_after_forward: false avoids it at about 8.5 GB per rank and becomes unnecessary once PR #3961 is merged.
• Point the JIT caches at node-local storage (HYBRID_EP_CACHE_DIR, TILELANG_CACHE_DIR, TRITON_CACHE_DIR). HybridEP resolves its cache from the passwd home directory rather than $HOME, and compiling into a parallel file system can stretch the first step past the NCCL timeout.

Checkpoint saving is enabled in the packed recipe, as on GB200. One torch_save checkpoint of the full model and optimizer is about 8.2 TB and took 2.5 minutes to write from 512 ranks; reaching max_steps triggers a final save regardless of ckpt_every_steps. Resuming with checkpoint.restore_from: LATEST was verified to continue the loss curve.

Loading and Inputs

The recipes load revision df42c109f1defefcbfcedbe7d905718a12266e40 from the Hugging Face Hub. To use a shared local snapshot of that revision instead, with the same directory for parity and training, set DS41_CHECKPOINT to it and pass both overrides to the recipe launcher:

--model.config.pretrained_model_name_or_path="$DS41_CHECKPOINT" \
--model.config.name_or_path="$DS41_CHECKPOINT"

Text architecture settings belong to config.text_config; vision settings belong to config.vision_config. Enabled Engram layers require the checkpoint’s fast tokenizer during construction. Configuration-owned loading uses the resolved checkpoint revision, or callers can supply the tokenizer explicitly. Materialization restores the tokenizer-derived map and deterministic hash buffers. engram_layer_ids selects the enabled tables; disabling a table changes the model function.

The checkpoint adapter decodes released FP8 dense weights, FP4 experts, and FP8 Engram tables into floating-point parameters. Retain the original config’s quantization metadata for automatic base-weight dequantization; the recipe also sets dequantize_base_checkpoint: true. Floating-point checkpoint export preserves released parameter names, required FP32 weights, and logical Engram rows while removing the original quantization metadata. Quantized re-export is not supported.

Recipe initialization uses the NeMo AutoModel wrapper’s internal Checkpointer and DCP path to load the base checkpoint into the distributed model.

Text inputs use unpacked input_ids [batch, sequence], right padding, and zero-based positions. The optional labels [batch, sequence] API computes shifted causal cross-entropy with -100 ignored. Attention KV and index quantize and dequantize boundaries remain part of every forward and use straight-through gradients; there is no quantization toggle.

The released vision tower is enabled by the default model configuration and disabled by the HellaSwag recipe. DeepseekV41Processor supports ordinary text and image conversations and save/reload. Image inputs require unpacked input_ids [batch, sequence], pixel_values [all_patches, 3, patch_size, patch_size], image_grid_hws [images, 2], and vision_token_types [batch, sequence]. The image encoder and projector reuse DeepSeek V4 modules; image delimiters and visual router bias are retained. Specialized tool and reasoning formatting should be rendered with the official encoder before processing.

Numerical Validation

At PR commit d5c8b5e45176bd6de54c139b311937516be4e0d2, a real-weight comparison covered all 40 layers, 4,096 positions, and all 129,280 vocabulary entries against the pinned checkpoint’s unchanged official inference/ implementation. Official TP1 was compared with native EP4 and 4 Engram owners using torch_mm, native HybridEP dispatch and combine, and TileLang. Dense and expert reference weights were decoded for BF16 compute, with the original cache/index quantization retained; this was not a default FP8/FP4 GEMM comparison.

The native execution streamed one production block at a time while carrying its own residual, mHC, and CSA2 state. Embedding, final norm, and head weights were replicated. Rank 0 supplied the single valid sample; other ranks used masked inputs and participated in communication. The input repeated one 145-token passage to fill 4,096 positions, so these metrics describe that sample rather than broad prompt coverage.

The fixed gates still failed: mean KL exceeded 0.001 and cosine was below 0.999. P95 KL and top-1 agreement passed their 0.002 and 98% thresholds. This is a measured baseline, not a strict parity pass.

An earlier PR commit, ae380ea74b8417a000c1639ebd64991aafeb42b8, completed three resident forward and backward iterations using the now-removed custom FP32 combine. It used 4 GB200 GPUs, the released first four layers and full layer-1 Engram table, EP4, activation checkpointing, and resharding. All local trainable gradients were present and finite, including nonzero Engram and expert gradients. This check performed no clipping or optimizer updates and does not establish training behavior for native HybridEP combine.

The historical streamed forward used a diagnostic direct-loading entry point that has since been removed. These records identify the tested commits; they do not claim a new GPU run after removing that loader or the custom combine code, or full-resident 40-layer PR training.

Original Implementation Training Result

The original implementation at 0d5919a5a506dd35dc4f55769410feb98ae854b8 completed 100 HellaSwag optimizer updates on 16 nodes x 4 GB200 with EP64, 64 Engram owners, the now-removed custom FP32 combine, BF16 optimizer moments, and checkpoint saving disabled. Training loss changed from 2.9959 at update 1 to 1.8177 at update 100; validation losses were 2.0687 and 2.0173 after updates 50 and 100. All logged loss and gradient-norm values were finite, the job exited successfully, and W&B received every training step and both validations. This full-resident training result belongs to the original commit and does not establish training behavior for native HybridEP combine, checkpoint restoration, or full-model convergence.

Current Scope

THD packing, TP, PP, incremental generation caches, bounded decoder replay, and million-token training are unsupported. Image inputs also require unpacked batches. Indexers remain frozen because hard top-k selection has no language-model gradient; indexer distillation is outside this implementation. Native DSpark draft training is available through the DSpark training guide, with the target frozen and DSpark gradients isolated from the target backbone. Fine-tuning does not reproduce the report’s Engram pretraining optimizer. Bitwise training continuation is not established for the TileLang backward kernels.