mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +08:00
8a5209ee65
- examples/neuron: tiny random V4.1 checkpoint in the real HF layout, pure-torch replacements for the official reference kernels, golden generation, a teacher-forced DeepInfer runner (LLM offline API) and a comparison tool. CPU mode: 25/27 greedy tokens match, worst logits cos 0.9985 / rel 0.056 against a reference self-noise floor of 0.032-0.060. - model_loader: check parameters against the platform device, not is_cuda. - deep_gemm paged MQA logits: do not dereference block-table slots past the context. - bind_cpu: skip in Neuron CPU mode (1 core per NeuronCore share would serialize host compute). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
38 lines
2.1 KiB
Markdown
38 lines
2.1 KiB
Markdown
# DeepInfer on AWS Trainium — end-to-end check with a tiny checkpoint
|
|
|
|
Design and conventions of the port: [deepinfer/neuron_ops/NEURON_PORT.md](../../deepinfer/neuron_ops/NEURON_PORT.md).
|
|
|
|
The ground truth is DeepSeek's official reference implementation (`inference/` in the
|
|
`deepseek-ai/DeepSeek-V4.1-Flash-1017` HF repo) running on CPU, with its TileLang kernels replaced by
|
|
the pure-torch versions in `ref_kernels_torch.py`.
|
|
|
|
| Script | Purpose |
|
|
|---|---|
|
|
| `make_tiny_checkpoint.py` | Random tiny V4.1 checkpoint in the real HF layout (FP8 e4m3 + ue8m0 32x32 scales, FP4 e2m1 experts with 1x32 scales, engram tables) |
|
|
| `check_tiny_checkpoint.py` | DeepInfer loader coverage / shape / dtype checks on the meta device |
|
|
| `ref_kernels_torch.py`, `test_ref_kernels_torch.py` | Torch replacements for the reference `kernel.py` |
|
|
| `reference_golden.py` | Reference prefill + greedy decode; saves logits, tokens and per-layer hidden states |
|
|
| `run_deepinfer_tiny.py` | Same prompts through the DeepInfer engine (`LLM` offline API), teacher-forced |
|
|
| `compare_golden.py` | Candidate vs golden: argmax match, logits rel / cos / KL, per-layer hidden error |
|
|
|
|
Tiny config: 10 layers + 3 nextn, hidden 1024, 8 heads x 512 (MLA, rope 64), 16 experts / top-4,
|
|
indexer 4 x 128 with top-64, engram on layers 1 and 4, real tokenizer (vocab 129280).
|
|
|
|
```bash
|
|
export DEEPINFER_ACCELERATOR=neuron DEEPINFER_NEURON_DEVICE=cpu # or =neuron for the NeuronCore
|
|
python examples/neuron/run_deepinfer_tiny.py --threads 32 --hidden --compare
|
|
```
|
|
|
|
## Results so far (CPU mode, basic mode: eager, no DSpark/MTP, EP=DP=CP=1)
|
|
|
|
Teacher-forced, prefill + 8 decode steps per prompt; "floor" is the reference compared with itself
|
|
under a different CPU thread count (FP8 activation quantization amplifies sub-ulp differences).
|
|
|
|
| prompt (tokens) | argmax match | max logits rel (floor) | min cos (floor) |
|
|
|---|---|---|---|
|
|
| short (11) | 9/9 | 0.055 (0.032) | 0.9985 (0.9995) |
|
|
| medium (172, wraps the 128 SWA window) | 8/9 | 0.049 (0.039) | 0.9988 (0.9993) |
|
|
| long (715, candidate-block filter active) | 8/9 | 0.056 (0.060) | 0.9985 (0.9982) |
|
|
|
|
Per-step / per-layer numbers: `golden/deepinfer_cpu_vs_golden.json`.
|