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. 可复现入口

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：

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。