File size: 7,104 Bytes
851b203 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | # 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。
|