# DMax-16B 官方推理复现

状态：官方单样本推理 smoke 已完成（2026-07-13）

## 1. 复现目标

先复现官方仓库 Quick Start 的单样本 soft parallel decoding，不进行训练：

- 模型：`Zigeng/DMax-16B`
- 固定 revision：`225df0ce2818396fd3f18996c8fbd59b0666e4fe`
- 权重：7 个 safetensors 分片，共 `32,529,274,876` bytes
- `gen_length=2048`
- `block_length=32`
- `steps=32`
- `threshold=0.0`
- dtype：`bfloat16`

官方 `generate_spd` 在每个 block 内最多执行 32 次 forward。每一步都会刷新已经得到的 token，并将其 embedding 按置信度与 mask embedding 混合，所以这不是普通的“一旦 unmask 就冻结”的离散解码。

## 2. 远端路径与作业

远端代码/日志根目录：`/home/l/liyj/shiying/hku_diffusion_dllm`

DMax 权重与 Hugging Face cache 已切到 scratch，避免占用 `shiying` quota：

- 模型路径：`/mnt/scratch/l/liyj/hku_diffusion_dllm/models/dmax-16b`
- HF cache：`/mnt/scratch/l/liyj/hku_diffusion_dllm/.cache/huggingface`
- scratch 容量检查：`/mnt/scratch` 约 `61T` 可用

此前下载 job `680433` 在 `shiying` 路径完成 4/7 个 shard 后超时，重试时触发 `Disk quota exceeded`。之后切到 `/mnt/scratch/l/liyj`，下载 job `680834` 已完成全量权重。

| 阶段 | Job ID | 资源 | 状态 |
| --- | ---: | --- | --- |
| custom class import + CUDA smoke | `680446` | TITAN RTX | completed |
| 固定 revision 权重下载 | `680834` | CPU normal | completed |
| 旧官方 smoke 候选 1 | `680835` | H100-96G | cancelled |
| 旧官方 smoke 候选 2 | `680836` | H200-141G | cancelled |
| 旧官方 smoke 候选 3 | `680837` | A100-80G | failed: project venv torch mmap/import |
| scratch runtime 修复 | `682107` | CPU normal | completed |
| 官方 smoke | `682109` | H100-96G | completed |
| 官方 smoke | `682112` | A100-80G | completed |

`680837` 失败原因是项目目录 `.venv-dflash-cu128` 在 GPU node 上 PyTorch 动态库导入失败。后续将最小 runtime 拷到 scratch，并补齐 `torchgen`、`torch/testing`、`pyyaml` metadata 后通过 import smoke。

## 3. 环境

复用远端 `.venv-dflash-cu128` 的最小 Transformers 推理环境：

| 组件 | 版本 |
| --- | --- |
| Python | 3.12.3 |
| PyTorch | 2.8.0+cu128 |
| Transformers | 4.57.1 |
| Accelerate | 1.14.0 |
| huggingface_hub | 0.36.2 |
| safetensors | 0.8.0 |

模型配置声明的原始 Transformers 版本是 `4.56.0`。当前环境只高一个 minor release，先用官方 remote code 做 smoke；若出现 API 兼容错误，再建立固定 `4.56.0` 的独立环境。

环境 smoke `680446` 已验证：CUDA tensor 运算成功，`LLaDA2MoeConfig` 和 `LLaDA2MoeModelLM` 均可通过 Hugging Face remote code 导入。该测试不实例化 16B 权重，因此只证明 API/CUDA 环境兼容，不替代完整 model-load smoke。

## 4. 指标定义

| 指标 | 本次计算方式 |
| --- | --- |
| NFE | 直接读取官方 `generate_spd` 返回值 |
| generated tokens | `generated_tokens.shape[-1]`，到首个 EOS 或 2048 上限 |
| TPF | `generated_tokens / NFE` |
| tokens/s | generated tokens / `generate_spd` CUDA 同步 wall time |
| TTFT proxy | 从调用 `generate_spd` 到第一次 model forward 完成；官方 API 没有逐 token callback，因此不能声称是真实流式 TTFT |
| 模型加载显存 | load 后 `memory_allocated` / `memory_reserved` |
| 峰值显存 | generation 区间 CUDA peak allocated / reserved |

这里单独标注 TTFT proxy，避免把完整生成时延或首个 block 完成时延误写成标准流式 TTFT。

## 5. 可复现入口

```bash
cd /home/l/liyj/shiying/hku_diffusion_dllm
sbatch experiments/dmax_reproduction/download_dmax.sbatch
sbatch experiments/dmax_reproduction/fix_dmax_runtime_env.sbatch
sbatch --gres=gpu:h100-96:1 experiments/dmax_reproduction/run_dmax_smoke.sbatch
```

脚本：

- `experiments/dmax_reproduction/download_dmax.py`
- `experiments/dmax_reproduction/download_dmax.sbatch`
- `experiments/dmax_reproduction/prepare_dmax_env.sbatch`
- `experiments/dmax_reproduction/run_dmax_smoke.py`
- `experiments/dmax_reproduction/run_dmax_smoke.sbatch`
- `experiments/dmax_reproduction/submit_dmax_scratch.sh`
- `experiments/dmax_reproduction/submit_dmax_after_env.sh`
- `experiments/dmax_reproduction/fix_dmax_runtime_env.sbatch`
- `experiments/dmax_reproduction/submit_dmax_smoke_one.sh`
- `experiments/dmax_reproduction/check_recent_dmax.sh`

## 6. 实验结果

官方 Quick Start 单样本 prompt：

```text
A robe takes 2 bolts of blue fiber and half that much white fiber. How many bolts in total does it take?
Let's think step by step
```

| Job | GPU | NFE | tokens | TPF | tokens/s | TTFT proxy | load time | load mem | peak mem | 状态 |
| ---: | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |
| 680837 | A100-80G | - | - | - | - | - | - | - | - | failed before torch import |
| 682109 | H100 NVL 96G | 24 | 131 | 5.46 | 16.55 | 2.29s | 123.25s | 30.28 GiB | 30.62 GiB | completed |
| 682112 | A100 80G PCIe | 18 | 131 | 7.28 | 15.45 | 1.83s | 75.89s | 30.28 GiB | 30.60 GiB | completed |

两个作业输出的答案一致，最终答案为 `3 bolts`。原始 stdout/stderr 保存在远端 `runs/dmax_reproduction/logs/`；结构化指标和答案已拉回本地：

- `runs/dmax_reproduction/dmax16b_official_g2048_b32_t0_job682109/metrics.json`
- `runs/dmax_reproduction/dmax16b_official_g2048_b32_t0_job682109/answer.txt`
- `runs/dmax_reproduction/dmax16b_official_g2048_b32_t0_job682112/metrics.json`
- `runs/dmax_reproduction/dmax16b_official_g2048_b32_t0_job682112/answer.txt`

注意：A100 和 H100 在同一 seed 下 NFE 不同（18 vs 24），说明官方 `generate_spd` 的 confidence/commit 路径存在硬件或数值细节敏感性。下一步做多样本评测时需要固定 GPU 型号并报告 seed 方差，不能把单样本 TPF 当成稳定论文指标。

## 7. 后续

进入 7 月 16 日计划：用论文数学阈值 `0.5` 和代码阈值 `0.65` 做小规模 accuracy/TPF，对比 `threshold=0.0`，再区分 soft revise 带来的并行度收益与真实 wall-clock 收益。

## 8. Threshold Sweep（7 月 16 日计划启动）

已完成 H100/A100 小样本 sweep，记录见 `runs/dmax_reproduction/DMAX_THRESHOLD_SWEEP_2026-07-13.md`。

H100 `682193`：

| threshold | task | samples | accuracy | mean NFE | mean TPF | mean latency |
| ---: | --- | ---: | ---: | ---: | ---: | ---: |
| 0.00 | math | 4 | 1.000 | 21.50 | 6.53 | 5.04s |
| 0.50 | math | 4 | 1.000 | 17.00 | 7.40 | 3.67s |
| 0.65 | math | 4 | 1.000 | 19.50 | 6.61 | 4.24s |
| 0.00 | code | 3 | 1.000 | 7.67 | 10.58 | 1.59s |
| 0.50 | code | 3 | 0.667 | 7.33 | 10.22 | 1.49s |
| 0.65 | code | 3 | 1.000 | 7.00 | 10.21 | 1.42s |

当前结论只适用于 smoke：math 上 `threshold=0.5` 在不掉 accuracy 的情况下最快；code 小样本不稳定，需要扩大样本并保存失败代码。

A100 `682194` 也完成，趋势相近：math 在 `threshold=0.5` 时从 5.90s 降到 5.13s，code 在 `threshold=0.5` 同样出现 1/3 失败，`0.65` 恢复到 3/3。