DSV3 Native PyTorch Reference

This project is a single-card native PyTorch DeepSeek-V3 / DSV3 training reference. It is built by incrementally improving the existing dsv3_native_reference codebase, not by rewriting the project from scratch.

The goal is correctness and alignment. This code is meant to compare forward, causal LM loss, backward gradients, and optimizer-step parameter updates against Megatron / Megatron-Core or another training framework.

What It Implements

• Native PyTorch DeepseekV3Config with official DeepSeek-V3 field names.
• RMSNorm, default RoPE, MLP, grouped top-k router, naive MoE, MLA attention, decoder layers, base model, and causal LM head.
• Causal LM next-token loss with ignore_index=-100.
• Toy deterministic random-token batches for sanity and alignment tests.
• Single-card AdamW training loop with gradient norm, optional clipping, gradient accumulation, checkpoint save/load, and resume.
• Tensor, state-dict, and activation recording helpers for precision comparison.

Dependencies

Allowed runtime dependencies are:

• torch
• torch.nn
• torch.nn.functional
• torch.optim
• torch.utils.data
• Python standard library

The optional oracle/test path also uses:

• numpy

The optional HuggingFace alignment script additionally uses:

• transformers with transformers.models.deepseek_v3

The model code does not depend on these HuggingFace engineering components:

• PreTrainedModel
• GenerationMixin
• Cache / DynamicCache
• HF output dataclasses
• HF attention registries
• HF masking utilities
• HF Trainer / generate / AutoModel / AutoTokenizer

The local dsv3_native/configuration_deepseek_v3.py is a self-contained config shim that preserves the official DeepSeek-V3 config field surface used by this reference.

Quick Sanity

cd dsv3_native_reference
python scripts/run_sanity.py

Expected output includes the device, loss before/after one optimizer step, gradient norms, logits shape, and number of cache layers.

Single-Card Training

python scripts/train_single_gpu.py --steps 3 --batch-size 2 --seq-len 16 --log-interval 1

Useful options:

• --device auto|cpu|cuda
• --seed
• --deterministic
• --grad-accum-steps
• --lr
• --weight-decay
• --grad-clip
• --save-interval
• --out-dir
• --resume

Each log line reports step, loss, gradient norm before clipping, gradient norm after clipping, tokens/s, elapsed time, and learning rate.

Checkpoints

Save/load helpers live in dsv3_native/checkpoint.py.

from dsv3_native import save_checkpoint, load_checkpoint

save_checkpoint(path, model, optimizer, step=step, extra={"note": "alignment"})
meta = load_checkpoint(path, model, optimizer, map_location="cpu", strict=True)

Roundtrip check:

python scripts/check_checkpoint_roundtrip.py

Compare Helpers

python scripts/compare_two_native_models.py

This constructs two identical native models, loads one state dict into the other, runs the same batch, and compares logits, loss, hidden state, and parameters.

ActivationRecorder can be attached to selected modules with substring patterns to capture intermediate tensors for layer-by-layer alignment.

NumPy Oracle

The NumPy oracle is a deliberately small and slow mathematical reference for checking the native PyTorch implementation itself. It uses the same PyTorch state_dict, converts tensors to NumPy float64, and compares forward, causal LM loss, finite-difference gradients, and AdamW parameter deltas.

python scripts/compare_numpy_native_step.py --case dense_micro
python scripts/compare_numpy_native_step.py --case moe_micro
python scripts/compare_numpy_native_step.py --case tiny_sampled

Cases:

• dense_micro: tiny dense-only DSV3, full-parameter finite-difference gradient.
• moe_micro: tiny MoE DSV3, grouped router / routed experts / shared expert.
• tiny_sampled: current tiny_dsv3_config() forward/loss plus sampled gradient checks for representative parameters.

The pytest suite also includes a representative forward/loss shape sweep:

• dense q-LoRA shape: B=2, S=5
• dense non-interleave RoPE shape: B=1, S=6
• grouped MoE shape with padding mask and labels=-100: B=2, S=5
• current tiny mixed dense+MoE shape: B=2, S=8

Full-parameter finite-difference gradient checks intentionally stay on micro configs. Larger shape cases focus on forward/loss/hidden-state coverage so the oracle remains practical to run in CI or on a remote validation box.

For larger sequence lengths, use the large-shape NumPy oracle script. It still uses NumPy as the correctness oracle, but switches from full-parameter finite difference to sampled finite-difference gradients:

python scripts/compare_numpy_large_shape.py --case dense_lora --seq-len 128
python scripts/compare_numpy_large_shape.py --case moe_grouped --seq-len 128 --pad-width 16
python scripts/compare_numpy_large_shape.py --case non_interleave --seq-len 512 --skip-grad
python scripts/compare_numpy_large_shape.py --case tiny_mixed --seq-len 1024 --skip-grad

By default, large-shape gradient checks use sampled finite differences. To force full-parameter finite differences on a larger shape, pass --grad-mode full. This runs two full NumPy forwards per trainable parameter element, so the script has a safety limit by default:

python scripts/compare_numpy_large_shape.py \
  --case dense_lora \
  --seq-len 128 \
  --grad-mode full \
  --allow-slow-full-grad

The script prints checked_param_elements, finite_difference_points, and numpy_forward_count before running the finite-difference loop. For full mode:

numpy_forward_count = 2 * checked_param_elements

The intended layering is:

• micro shapes: NumPy full-parameter finite-difference gradients and AdamW delta.
• medium/large shapes: NumPy forward/loss/hidden-state checks and sampled finite-difference gradients by default; full finite differences are available explicitly with --grad-mode full.
• external frameworks: compared against the native reference after NumPy has validated the covered native paths.

Passing these commands means the NumPy oracle and native PyTorch reference agree on the covered math. It does not mean HuggingFace or Megatron alignment has been covered; those are separate comparison stages.

HuggingFace Official Alignment

The HuggingFace alignment script is an optional compatibility checker. It imports the official transformers.models.deepseek_v3 implementation, builds a matching HF model from the native config fields, loads the native state_dict into HF, and compares parameters, logits, loss, hidden states, selected attention/MLP/MoE activations, and gradients.

python scripts/compare_hf_deepseek_v3.py --case dense_micro
python scripts/compare_hf_deepseek_v3.py --case moe_micro
python scripts/compare_hf_deepseek_v3.py --case tiny --batch-size 2 --seq-len 16

Useful options:

• --device cpu|cuda|auto
• --no-backward
• --no-activations
• --atol, --rtol
• --grad-atol, --grad-rtol

Passing this script means the installed HuggingFace DeepSeek-V3 implementation matches the native reference on the covered path. HuggingFace is not treated as the correctness oracle for this project; NumPy remains the oracle used to validate native-reference math.

Tests

pytest tests/test_sanity.py
pytest tests/test_numpy_reference.py
pytest tests/test_numpy_large_shape.py
pytest tests/test_hf_alignment.py

If pytest is not installed, use the environment's normal package management flow before running tests. tests/test_hf_alignment.py skips automatically when Transformers or transformers.models.deepseek_v3 is not installed.

Megatron Alignment Order

1. Disable TP/PP/DP and first align single-card behavior.
2. Use FP32 first; only move to bf16/fp16 after FP32 passes.
3. Disable dropout.
4. Fix seed, input_ids, labels, and attention_mask.
5. Compare logits.
6. Compare loss.
7. Compare every layer's hidden states.
8. Compare attention output, MLP output, and MoE output.
9. Compare gradients.
10. Compare parameter changes after optimizer.step().
11. If Megatron uses tensor-parallel shards, gather full tensors before comparing them to this native reference.

Not Supported In This Reference

• yarn / dynamic RoPE
• FlashAttention
• SDPA
• fused kernels
• mixed precision
• tensor, pipeline, data, expert, or sequence parallelism
• distributed optimizer
• HuggingFace from_pretrained, Trainer, generate, AutoModel, or AutoTokenizer