NVlabs/Nemotron-Labs-Diffusion — Tri-Mode 擴散語言模型完整教學

本教學針對 NVlabs/Nemotron-Labs-Diffusion（commit 2026-05-28，stars 32，dual license: Apache-2.0 code + NOML weights）撰寫，為內部知識庫的「擴散 LM 工程化 reference」。

本 repo 是「論文公開時的 code drop」性質，整體規模 ~1.4K 行程式碼、單次 commit、無 release tag。適合用來理解 NVIDIA 怎麼把擴散解碼路徑與 AR 路徑整合進同一模型；不適合直接拿來 production 大規模部署（建議走 SGLang upstream 或等 0.x 正式 release）。

目錄

1. 專案定位
2. 安裝指南
3. 核心架構解析
4. Helper Scripts 詳細用法
5. 應用場景：4 條 recipe + 5 個 use-case
6. 資安掃描報告
7. FAQ
8. 進階技巧
9. 整合進其他工作流
10. 重點摘要 Checklist
11. 進一步閱讀

1. 專案定位

1.1 這是什麼

NVlabs/Nemotron-Labs-Diffusion（以下簡稱 NLD）是 NVIDIA Research Labs（NVlabs，與 production 線的 NVIDIA-NeMo 不同 org）對外公開的 tri-mode 語言模型 inference & evaluation harness。

 1┌──────────────────────────────────────────────────────────────────┐
 2│ Nemotron-Labs-Diffusion (NLD) — Research Drop                    │
 3├──────────────────────────────────────────────────────────────────┤
 4│ (1) chat/*.py            5 個對話腳本，覆蓋 4 個解碼模式          │
 5│ (2) evaluate.py          單機 evaluator，無 server / SLURM 依賴   │
 6│ (3) eval.sh              SLURM + NeMo-Skills 多 GPU 評估流水線    │
 7│ (4) sglang_spark/        DGX Spark + SGLang FP8 部署完整指南      │
 8│ (5) xp/dlm_api/          自家擴散 batch / load balancer / 算法庫  │
 9│ (6) xp/nemo-skills/      NeMo-Skills 對接 patch + NFE 紀錄        │
10│ (7) scripts/             單一 helper：fetch_bundled_lora.sh       │
11└──────────────────────────────────────────────────────────────────┘

模型本體（3B / 8B / 14B base + instruct + VL）住在 HuggingFace：

• https://huggingface.co/collections/nvidia/nemotron-labs-diffusion

本 repo 提供「拿到權重之後該怎麼跑」的完整工程配套。

1.2 Tri-mode 是什麼？

論文 / README 的「tri-mode」指三條共用同一份權重的解碼路徑，只是 attention pattern 與解碼迴圈不同：

第四個延伸（不算獨立 mode）：linear_spec_lora，在 linear_spec 的 draft 階段掛上 ~137MB 的 LoRA adapter（PEFT），對 o_proj 做特化以提高 draft 的 acceptance length。

Self-speculation 為什麼好？ 傳統 speculative decoding（如 Medusa、Eagle）需要一個獨立的 draft 模型，KV cache 不共用、權重也不共用。NLD 的 linear_spec 只用「同一份權重切換 attention」就完成 draft+verify，記憶體效率最高。

1.3 在 NVIDIA AI 研究與工程版圖的位置（重要）

Nemotron-Labs-Diffusion 不是隨便命名的 — 它故意把 Nemotron 放進名字裡，暗示與正規 production 線 NVIDIA-NeMo/Nemotron 共享某種 lineage / branding，但 org 與 license 都不同：

結論：NLD 是 NVIDIA-NeMo/Nemotron 在「擴散解碼研究」面的對應物 — production 線專注於 hybrid Mamba-Transformer MoE LLM 的訓練 recipe 與量產部署；research 線（NLD）專注於擴散平行解碼與自推測加速的可行性驗證。兩者共享「Nemotron」品牌，但不共享 codebase。

與 NVIDIA Cosmos / NVlabs alpamayo 的關係

NVIDIA 在「擴散」這條技術線同時有多支隊伍在做不同 vertical：

NVIDIA-NeMo vs NVlabs 的 org 分工

1NVIDIA 對外 GitHub orgs（節錄）
2├── NVIDIA/             ← 大眾化工具（CUDA、TensorRT、NeMo-Skills...）
3├── NVIDIA-NeMo/        ← 產品 / 工程 / 訓練 hub（Apache-2.0）→ Nemotron 主家族住這
4├── NVlabs/             ← 研究院 paper-companion code（常用 non-commercial license 或 dual license）→ NLD 住這
5└── nvidia-cosmos/      ← 跨領域產品（Cosmos World Foundation Models 等）

特別注意：NVlabs 過去多數 repo 是 non-commercial license，NLD 採 Apache-2.0 code + NOML weights 是相對開放的做法 — 但拿模型權重做商業推論前一定要看 NOML 條款（§6.3）。

1.4 統計資料快照（2026-06-02）

2. 安裝指南

2.1 環境需求

NLD 自身不是 PyPI 套件 — 就是一堆腳本 + Python 檔案。但「跑得起來」需要對應的 inference / eval 基礎設施：

2.2 場景 A：最小可行（chat 單請求）

 1# 取得本 repo
 2git clone --depth 1 https://github.com/NVlabs/Nemotron-Labs-Diffusion.git
 3cd Nemotron-Labs-Diffusion
 4
 5# 建議用 uv 管理 venv（CLAUDE.md §6.2 規定）
 6uv venv .venv && source .venv/bin/activate
 7
 8# 4 個基礎依賴（peft 只在 linear_spec_lora 需要）
 9uv pip install "transformers>=5.0" torch peft
10
11# 跑單請求 chat（AR 模式，最簡）
12python chat/chat_ar.py
13# 接著 stdin 輸入問題 → 模型回答 + NFE（forward 次數）

第一次跑會自動下載 17GB 的 8B 權重到 ~/.cache/huggingface/。 若硬碟緊張改用 3B：REPO=nvidia/Nemotron-Labs-Diffusion-3B python chat/chat_ar.py（但腳本本身 hard-code 在 REPO 常數，需手動改檔或改用 chat.py --model nvidia/Nemotron-Labs-Diffusion-3B）。

2.3 場景 B：單機 benchmark（evaluate.py）

 1# 加上 datasets
 2uv pip install torch transformers datasets peft
 3
 4# 5 分鐘 smoke：50 題 gsm8k，dLM 模式
 5python evaluate.py --mode dlm --tasks gsm8k --limit 50
 6
 7# 完整 gsm8k（1319 題）× 三模式
 8python evaluate.py --mode ar           --tasks gsm8k
 9python evaluate.py --mode dlm          --tasks gsm8k
10python evaluate.py --mode linear_spec  --tasks gsm8k
11
12# Linear-spec + LoRA draft（先 fetch LoRA）
13bash scripts/fetch_bundled_lora.sh   # 137MB → miscs/linear_spec_lora/
14python evaluate.py --mode linear_spec --tasks gsm8k --lora
15
16# 多任務一次跑
17python evaluate.py --mode dlm --tasks gsm8k,math-500

每 50 題印一次進度（acc / avg_tok / avg_nfe / TPF / 耗時），最後印 summary。詳細欄位定義見 §4.2。

2.4 場景 C：SLURM cluster eval（eval.sh）

需先有：

• 一份 NeMo-Skills-ready 的容器 image（.sqsh，需自建或從 NVIDIA 內部 registry 抓）
• SLURM account

 1# 必填環境變數
 2export CONTAINER_IMAGE=/path/to/your/nemo-skills-ready.sqsh
 3export ACCOUNT=<your-slurm-account>
 4
 5# 可選
 6export OUT_DIR=$PWD/eval_suit_results
 7export HF_HOME=$HOME/.cache/huggingface
 8
 9# Sanity check：1 個 benchmark，dLM 模式
10bash eval.sh --mode dlm --benchmarks gsm8k:1
11
12# Full 10-benchmark suite（HumanEval / MBPP / MMLU / IFEval / LiveCodeBench / AIME / GPQA 等）
13bash eval.sh --mode dlm --gpus 8
14
15# Linear-spec + LoRA（會自動拉 PEFT_IMAGE 取代 PREBAKED_IMAGE）
16bash eval.sh --mode linear_spec --lora --benchmarks gsm8k:1
17
18# Dry-run 只看 resolved settings 不真的 submit
19bash eval.sh --mode dlm --benchmarks gsm8k:1 --dry-run

2.5 場景 D：DGX Spark + SGLang FP8 部署

詳細 8 步驟見 sglang_spark/README.md。摘要：

 1# Step 1：clone SGLang DLLM 分支
 2mkdir -p ~/sglang_dllm/src && cd ~/sglang_dllm/src
 3git clone --depth 1 -b upstream/2-dllm-lora-ar https://github.com/hutm/sglang.git
 4
 5# Step 2：一行 patch（修 self.report_prefill_stats bug；upstream merge 後可省）
 6sed -i 's|self\.report_prefill_stats(|self.metrics_reporter.report_prefill_stats(|' \
 7  ~/sglang_dllm/src/sglang/python/sglang/srt/dllm/mixin/scheduler.py
 8
 9# Step 3：抓 LoRA adapter
10mkdir -p ~/sglang_dllm/linear_spec_lora && cd ~/sglang_dllm/linear_spec_lora
11curl -fsSL -O https://huggingface.co/nvidia/Nemotron-Labs-Diffusion-8B/resolve/main/linear_spec_lora/adapter_config.json
12curl -fsSL -O https://huggingface.co/nvidia/Nemotron-Labs-Diffusion-8B/resolve/main/linear_spec_lora/adapter_model.safetensors
13
14# Step 4：拉容器 image
15docker pull lmsysorg/sglang:spark
16
17# Step 5–6：放 launch script + 啟動（FP8 + LoRA-enhanced linear_spec 是預設）
18QUANT=fp8 ~/sglang_dllm/launch_server.sh detach
19
20# Step 7：smoke test
21curl -sS http://localhost:30000/v1/chat/completions \
22  -H 'Content-Type: application/json' \
23  -d '{"model":"nvidia/Nemotron-Labs-Diffusion-8B",
24       "messages":[{"role":"user","content":"What is 15% of 240?"}],
25       "max_tokens":256,"temperature":0}'
26
27# Step 8：開 sglang_spark/index.html（單檔 HTML chat 客戶端）

模式切換（每次需 restart server，~3 min cuda graph 重 capture）：

1~/sglang_dllm/launch_server.sh stop
2QUANT=fp8 ALGO=LinearSpec-base ~/sglang_dllm/launch_server.sh detach  # 無 LoRA
3QUANT=fp8 ALGO=FastDiffuser    ~/sglang_dllm/launch_server.sh detach  # 純擴散
4QUANT=fp8 ALGO=AR              ~/sglang_dllm/launch_server.sh detach  # 純 AR

2.6 常見安裝坑

3. 核心架構解析

3.1 三模式統一在一份權重的訣竅

NLD 的關鍵設計是：整個 transformer block 的權重不變，只切換 attention mask 與 decoding loop。三個 method 是 model class 上的方法：

 1                ┌─────────────────────────────────┐
 2                │  Nemotron-Labs-Diffusion-Xb     │  ← 一份權重
 3                │  (AutoModel.from_pretrained,    │
 4                │   trust_remote_code=True)       │
 5                └────────────────┬────────────────┘
 6                                 │
 7            ┌────────────────────┼────────────────────┐
 8            │                    │                    │
 9            ▼                    ▼                    ▼
10    ar_generate()         generate()           linear_spec_generate()
11    causal mask           bidirectional        draft: bidirectional
12    one tok / step        block N tok / step   verify: causal
13                          NFE ≈ N_blk × steps  KV cache shared

從 chat/chat.py generate() 看得最清楚：

 1def generate(args, model, tokenizer, prompt_ids: torch.Tensor):
 2    eos = tokenizer.eos_token_id
 3    if args.mode == "ar":
 4        return model.ar_generate(prompt_ids, max_new_tokens=args.max_new_tokens)
 5    if args.mode == "dlm":
 6        return model.generate(
 7            prompt_ids,
 8            max_new_tokens=args.max_new_tokens,
 9            block_length=args.block_length,
10            threshold=args.threshold,           # 0.9：confidence-thresholded unmasking
11            eos_token_id=eos,
12        )
13    # linear_spec / linear_spec_lora
14    return model.linear_spec_generate(
15        prompt_ids,
16        max_new_tokens=args.max_new_tokens,
17        block_length=args.block_length,
18        eos_token_id=eos,
19    )

模型 class 內部（不在本 repo，在 HF model modeling.py）做的事：

• ar_generate：把 attention mask 設 causal triangular，每步生 1 token，重複 N 次。
• generate（block diffusion）：把 attention mask 對 N 個 mask token 設 bidirectional，前向一次後對 confidence > threshold 的位置 unmask、剩下繼續迭代。
• linear_spec_generate：draft phase 用 bidirectional 一次生 N 個候選；verify phase 切回 causal 一次性檢查；接受最長 prefix + 1 bonus token；繼續下一輪。

3.2 系統架構圖（mermaid）

flowchart TB
  subgraph User["使用者輸入"]
    U1["stdin chat
單請求 / multi-turn"]
    U2["HF dataset
gsm8k / math-500 / ..."]
    U3["HTTP /v1/chat/completions
OpenAI-compatible"]
  end

  subgraph Decode["三條解碼路徑（同一份權重）"]
    D1["model.ar_generate()
causal attention
1 tok / step"]
    D2["model.generate()
bidirectional attention
block_length tok / step
threshold=0.9 unmask"]
    D3["model.linear_spec_generate()
draft: bidirectional
verify: causal
shared KV cache"]
  end

  subgraph Backend["執行後端（三選一）"]
    B1["chat/chat.py
HF AutoModel 單機"]
    B2["evaluate.py
HF + datasets 單機"]
    B3["xp/dlm_api/
FastAPI batch server
+ load balancer"]
    B4["SGLang DLLM
FP8 + LoRA-aware
cuda graphs"]
  end

  subgraph Eval["評估出口"]
    E1["stdout NFE
+ generated text"]
    E2["pass@1 / TPF
per task summary"]
    E3["NeMo-Skills metrics.json
+ output-rs0.jsonl"]
    E4["OpenAI response
usage.completion_tokens"]
  end

  U1 --> D1
  U1 --> D2
  U1 --> D3
  U2 --> D1 & D2 & D3
  U3 --> D3

  D1 --> B1 --> E1
  D2 --> B1 --> E1
  D2 --> B2 --> E2
  D2 --> B3 --> E3
  D3 --> B2 --> E2
  D3 --> B3 --> E3
  D3 --> B4 --> E4

3.3 在 NVIDIA AI 研究與工程生態系裡的座標（mermaid）

graph TB
  NVIDIA["NVIDIA"]

  subgraph orgs["GitHub org 分工"]
    NVAI["NVIDIA/
大眾化工具
(CUDA, TRT, NeMo-Skills)"]
    NNEMO["NVIDIA-NeMo/
產品工程 hub
(Apache-2.0)"]
    NVLABS["NVlabs/
研究院 paper-companion
(常 dual license)"]
    COSMOS["nvidia-cosmos/
跨領域產品
(Cosmos World Models)"]
  end

  NVIDIA --> orgs

  subgraph products["關鍵 repo 對應"]
    NEMOTRON["NVIDIA-NeMo/Nemotron
hybrid Mamba-MoE LLM
recipes+cookbooks"]
    NEMO_SKILLS["NVIDIA/NeMo-Skills
eval harness"]
    NLD["⭐ NVlabs/Nemotron-Labs-Diffusion
tri-mode (AR / dLM / linear_spec)
research drop"]
    ALPAMAYO["NVlabs/alpamayo
sibling research"]
    COSMOSREPO["nvidia/cosmos
video / world foundation"]
  end

  NNEMO --> NEMOTRON
  NVAI --> NEMO_SKILLS
  NVLABS --> NLD
  NVLABS --> ALPAMAYO
  COSMOS --> COSMOSREPO

  NEMOTRON -. "品牌共享
code 不共享" .-> NLD
  NEMO_SKILLS -. "eval client" .-> NLD
  NLD -. "互補不重疊
(text vs. video)" .-> COSMOSREPO

  classDef hl fill:#fff4d4,stroke:#d97706,stroke-width:3px
  class NLD hl

3.4 三模式效能特性與選用建議

NFE = Number of Forward Evaluations：模型前向次數。AR 的 NFE = 生成 token 數；dLM 的 NFE = block_length × diffusion steps（通常遠少於 token 數）；linear_spec 的 NFE = draft 一次 + verify 一次 = 2 per accepted-block。

3.5 目錄結構詳解

 1Nemotron-Labs-Diffusion/
 2├── chat/                          # ★ 最簡的入門：5 個對話腳本
 3│   ├── chat.py                    #   統一多 turn launcher（--mode 切換）
 4│   ├── chat_ar.py                 #   AR 單請求
 5│   ├── chat_dlm.py                #   dLM 單請求
 6│   ├── chat_linear_spec.py        #   linear_spec 單請求
 7│   └── chat_linear_spec_lora.py   #   linear_spec + LoRA 單請求
 8│
 9├── evaluate.py                    # ★ 單檔 evaluator（HF + datasets，無 SLURM）
10├── eval.sh                        # ★ SLURM 多 GPU 流水線（NeMo-Skills 客戶端）
11│
12├── sglang_spark/                  # ★ DGX Spark + SGLang 部署完整指南
13│   ├── README.md                  #   8 步驟說明
14│   ├── launch_server.sh           #   thin wrapper（QUANT/ALGO/PORT...）
15│   └── index.html                 #   單檔 HTML chat client（黑色主題）
16│
17├── scripts/
18│   └── fetch_bundled_lora.sh      # 從 HF 抓 137MB LoRA adapter
19│
20├── xp/                            # ★ vendored helpers（slim build）
21│   ├── examples/
22│   │   └── run_dlm_eval_pipeline_gpu_only.sh   # eval.sh 真正 submit 的 sbatch 內容
23│   ├── dlm_api/                   #   FastAPI batch server + load balancer
24│   │   ├── dlm_batch_server.py    #     單 worker（owns model）
25│   │   ├── dlm_openai_server.py   #     OpenAI 相容 schema + chat template
26│   │   ├── dlm_load_balancer.py   #     多 GPU 分流
27│   │   ├── README.md
28│   │   └── dlm_generate/          #     算法 registry（dispatch into 三模式）
29│   │       ├── base.py
30│   │       ├── ar_native.py       #       --mode ar 走這
31│   │       ├── nemotron.py        #       --mode dlm / linear_spec 走這
32│   │       ├── nemotron_mixed.py  #       mixed AR/dLM
33│   │       ├── BatchStaticCache.py
34│   │       ├── __init__.py
35│   │       └── utils/
36│   │           ├── eos_detect.py
37│   │           ├── sliding_window.py
38│   │           └── sampler.py
39│   └── nemo-skills/               #   NeMo-Skills 對接 patch
40│       ├── eval_dlm.py            #     eval client 主入口
41│       ├── patch_dictconfig_serialization.py
42│       ├── patch_openai_extra_body.py
43│       └── add_nfe_to_metrics.py
44│
45├── miscs/                         # （runtime 生成）LoRA adapter 存放處
46│   └── linear_spec_lora/          #   adapter_config.json + adapter_model.safetensors
47│
48├── assets/                        # README 用的 demo gif / png
49├── README.md                      # 主文件，14.6KB
50├── SECURITY.md                    # 走 NVIDIA coordinated disclosure
51├── LICENSE                        # ★ Dual license：Apache-2.0 code + NOML weights
52└── .gitignore                     # 把 *.safetensors / .venv / 結果目錄擋掉

注意 xp/ 命名來自 NVIDIA 內部 codebase 縮寫（“experiments”），對外公開版「slim build」已刪掉與 LLaDA / dinfer / fast_dllm 等其他擴散家族算法包（不適用於 Nemotron 系列權重）。

4. Helper Scripts 詳細用法

4.1 chat/chat.py — 統一多 turn launcher（114 行）

最有用的腳本，4 個 mode 都 cover、可保留多 turn history：

 1# 單 turn 也可用，預設 8B
 2python chat/chat.py --mode dlm
 3
 4# 多 turn 互動 + 自訂解碼參數
 5python chat/chat.py --mode linear_spec_lora \
 6  --max-new-tokens 1024 \
 7  --block-length 32 \
 8  --threshold 0.9 \
 9  --dtype bfloat16
10
11# 換 3B（記憶體 < 12GB 顯卡可用）
12python chat/chat.py --mode dlm --model nvidia/Nemotron-Labs-Diffusion-3B
13
14# 換 device
15python chat/chat.py --mode ar --device cuda:1 --dtype float16

多 turn 邏輯：腳本維護 history = [{"role": ..., "content": ...}, ...]，每次回應後 append 兩筆（user + assistant），下一輪用 tokenizer.apply_chat_template(history, ..., add_generation_prompt=True) 重新組 prompt。

:q / :quit / Ctrl-C 退出。

4.2 evaluate.py — 單檔 evaluator（319 行）

無 SLURM、無 server、無 lm-evaluation-harness，三個 pip 套件就跑：

1# Smoke
2python evaluate.py --mode dlm --tasks gsm8k --limit 50

進度輸出每 50 題印一次：

1── gsm8k ── loading gsm8k [test]
2  [   50/ 1319]  acc=92.00%  avg_tok=308.4  avg_nfe= 51.7  TPF= 5.96  (   38s)
3  [  100/ 1319]  acc=93.00%  avg_tok=305.9  avg_nfe= 51.2  TPF= 5.97  (   74s)
4  ...
5  ✓ gsm8k        acc=93.78%  avg_tok=302.0  avg_nfe= 52.1  TPF= 5.89  (1319 problems)

TPF 解讀：AR 的 TPF = 1.0；dLM / linear_spec 越大越好（理論上限 = block_length）。範例的 5.96 表示平均 1 次 forward 產生 ~6 token，是 AR 的 6 倍效率。

內建任務（2 個，可擴充）

新增任務只要在 evaluate.py 頂端的 TASKS dict 加 6 行：HF dataset id、question_field、gold_extractor: row -> str、scorer: (out, gold) -> bool、instruction。

Full 10-benchmark suite（+HumanEval / MBPP / MMLU / IFEval / LiveCodeBench / AIME / GPQA）的 scorer 比較複雜，要走 eval.sh 透過 NeMo-Skills。

4.3 eval.sh — SLURM 多 GPU 流水線（484 行）

核心設計：把 --mode 翻譯成一組 env vars，每個 (mode × benchmark group) submit 一個 SLURM job；每個 job 在同一 GPU node 上啟動 N 個 dlm worker + 1 個 load balancer + 1 個 NeMo-Skills eval client。

必填 env vars

1export CONTAINER_IMAGE=/path/to/nemo-skills-ready.sqsh
2export ACCOUNT=<slurm-account>

容器 image 一定要在跑前就準備好。eval.sh 提供 3 個槽位（PEFT_IMAGE / PREBAKED_IMAGE / OLD_IMAGE），自動依需求選取 — --lora 自動升級到 PEFT_IMAGE（必須有 peft）。

Per-mode 預設值（節錄 default_mode_settings()）

共用預設：tokens=8192 / steps=8192 / max_thinking=6000 / batch_size=1 / gpus=8 / temperature=0。

CLI 常用 flag

輸出結構

1$OUT_DIR/<exp_name>/hf_base/<eval_dir_name>/
2├── pipeline_group<N>.log              # sbatch + pipeline 完整 log
3├── results/eval-results/<task>/       # NeMo-Skills 輸出
4│   ├── metrics.json                   #   pass@1 / latency / TPF
5│   └── output-rs0.jsonl               #   單題回答 + 評分 trace
6├── nfe_group<N>/nfe_log.jsonl         # 逐 batch NFE 紀錄
7├── server_info_group<N>.env           # server metadata
8└── COMPLETED_group<N>  | FAILED_group<N>

預設 $OUT_DIR=$PWD/eval_suit_results。

4.4 scripts/fetch_bundled_lora.sh — LoRA adapter 下載器（55 行）

只做兩件事：從 HF model repo 抓 adapter_config.json + adapter_model.safetensors（~137MB）到 miscs/linear_spec_lora/。

 1# 預設：抓 8B 的 LoRA
 2bash scripts/fetch_bundled_lora.sh
 3
 4# 換成 3B
 5bash scripts/fetch_bundled_lora.sh --model nvidia/Nemotron-Labs-Diffusion-3B
 6
 7# 換變體（若 HF repo 有 v2 / v3 subfolder）
 8bash scripts/fetch_bundled_lora.sh --subfolder linear_spec_lora_v2
 9
10# 私有模型：先 export HF_TOKEN
11export HF_TOKEN=hf_xxxxxxxxxxxx
12bash scripts/fetch_bundled_lora.sh

為什麼不直接放 git？safetensors ~137MB 超過 GitHub 100MB 單檔上限。同樣的策略 NVIDIA-NeMo/Nemotron 也用（big files → HF）。

4.5 sglang_spark/launch_server.sh — DGX Spark serving wrapper（157 行）

對 lmsysorg/sglang:spark 容器的 thin wrapper。重點 env：

啟動 / 停止：

1~/sglang_dllm/launch_server.sh detach   # 背景啟動
2~/sglang_dllm/launch_server.sh stop     # 停止
3tail -f ~/sglang_dllm/logs/server.log   # 看 log（等 "Uvicorn running on http://0.0.0.0:30000"）

4.6 sglang_spark/index.html — 單檔 HTML 客戶端

黑底主題、無 build step、直接雙擊 / 或 python3 -m http.server 8000 服務。

• SGLang send Access-Control-Allow-Origin: * → file:// origin 也能 fetch()。
• ⚙ drawer 可改 Base URL（遠端 Spark 用 SSH tunnel 時改 port）。
• UI sliders：temperature、top-p、max_tokens、stream toggle。

5. 應用場景：4 條 recipe + 5 個 use-case

5.1 Recipe — 內部研究員「在自己 GPU 上 reproduce 論文 latency 數字」

 1# 1. 三模式各跑 1319 題 gsm8k（單 H100，~ 各 2 小時）
 2python evaluate.py --mode ar          --tasks gsm8k > ar.log
 3python evaluate.py --mode dlm         --tasks gsm8k > dlm.log
 4python evaluate.py --mode linear_spec --tasks gsm8k > spec.log
 5
 6# 2. LoRA-enhanced linear_spec
 7bash scripts/fetch_bundled_lora.sh
 8python evaluate.py --mode linear_spec --tasks gsm8k --lora > spec_lora.log
 9
10# 3. 自製 summary
11grep "✓ gsm8k" *.log

預期 TPF：ar ≈ 1.0；dlm ≈ 5–7；linear_spec ≈ 4–8；linear_spec --lora ≈ 6–10。

5.2 Recipe — 跑完整 10-benchmark suite 對外發 paper

需 SLURM cluster：

1export CONTAINER_IMAGE=/path/to/nemo-skills.sqsh
2export ACCOUNT=research-llm
3export OUT_DIR=$HOME/nld_full_eval
4
5# 用 dlm 模式跑完整 suite（10 task × 8 GPU sweep）
6bash eval.sh --mode dlm --gpus 8
7
8# 同 suite 換 linear_spec + LoRA
9bash eval.sh --mode linear_spec --lora --gpus 8

輸出走 NeMo-Skills 標準格式，可直接餵 NVIDIA 內部論文表格 pipeline。

5.3 Recipe — DGX Spark 工程師「在自家 workstation 開 inference demo」

按 §2.5 八步驟跑完即可。常用組合：

 1# 對外 demo（FP8 + LoRA-enhanced linear_spec）
 2QUANT=fp8 ~/sglang_dllm/launch_server.sh detach
 3
 4# 跟 OpenAI SDK 接（純 Python）
 5python -c "
 6from openai import OpenAI
 7c = OpenAI(base_url='http://localhost:30000/v1', api_key='not-used')
 8r = c.chat.completions.create(
 9    model='nvidia/Nemotron-Labs-Diffusion-8B',
10    messages=[{'role':'user','content':'寫一首五言絕句。'}],
11    temperature=0.7, max_tokens=128, stream=True)
12for chunk in r:
13    print(chunk.choices[0].delta.content or '', end='', flush=True)
14"

5.4 Recipe — 新任務評估（擴 TASKS dict）

 1# 在 evaluate.py 頂端 TASKS dict 加：
 2"hellaswag": Task(
 3    name="hellaswag",
 4    hf_dataset="Rowan/hellaswag",
 5    hf_split="validation",
 6    question_field="ctx",
 7    gold_extractor=lambda row: row["endings"][int(row["label"])],
 8    scorer=lambda out, gold: gold.lower() in out.lower()[:200],
 9    instruction="完成下列句子（一句話內）：\n\n",
10),

1python evaluate.py --mode dlm --tasks hellaswag --limit 100

5.5 5 個典型 use-case

5.6 不適合的場景

• 公司 production OpenAI-相容 endpoint：等 SGLang upstream PR 合進 main，或改用 vLLM / TGI 上的成熟方案。
• 想 fine-tune Nemotron-Labs-Diffusion：本 repo 完全不涵蓋訓練；訓練 stack 走 NVIDIA-NeMo/Nemotron 或自家 Megatron-Bridge。
• 要跑非 Nemotron 系列擴散 LM（LLaDA, Mercury, …）：slim build 已把對應 algorithm packages 砍掉；走原始 LLaDA-API repo。

6. 資安掃描報告

掃描範圍：chat/*.py / evaluate.py / eval.sh / scripts/*.sh / sglang_spark/*.{sh,md,html} / xp/**/*.py。 掃描方法：grep -rn -E "eval\(|exec\(|os\.system|subprocess|shell=True|curl|wget|pickle|trust_remote_code|HF_TOKEN|password|api_key"，再人工 review 高風險 hit。

🟢 整體結論：低風險

本 repo 是 research demo 性質、無 web endpoint 暴露面（SGLang serving 例外、見 🟡）、無使用者輸入注入面（chat 是 single-user stdin），主要風險集中在 HF trust_remote_code=True 這一條。

🟡 中風險（3 項）

🟡 M1：trust_remote_code=True 廣泛使用

1# chat/chat.py:46-47, chat_*.py 全系列
2tokenizer = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True)
3model = AutoModel.from_pretrained(REPO, trust_remote_code=True).cuda().to(torch.bfloat16)

• 影響：HF 會下載並執行 nvidia/Nemotron-Labs-Diffusion-* repo 內 modeling_*.py 等 Python 檔。若該 HF repo 被攻陷（NVIDIA 帳號被駭、TOFU 失敗等），下載者本機會被執行任意程式。
• 緩解：(a) 確認 HF model card 顯示 verified NVIDIA org；(b) 對 air-gapped 環境，先在 sandbox download → audit 後再分發；(c) 鎖定特定 commit from_pretrained(..., revision="<sha>")。
• NLD 自身可做的：無 — 是 HF Transformers 的設計限制；NVlabs 自家模型走這條路徑屬正常。

🟡 M2：SGLang server 預設綁 0.0.0.0

sglang_spark/launch_server.sh 內 SGLang 預設 listen 0.0.0.0:30000，無 auth。

• 影響：若 Spark 在開放網段 / 公司 LAN，任何人可 hit /v1/chat/completions 觸發 inference。
• 緩解：(a) 用 firewall 限制 30000 port；(b) 走 SSH tunnel（README 已示範）；(c) 反向代理掛 auth header。

🟡 M3：HF_TOKEN env 透過 process env 傳遞

scripts/fetch_bundled_lora.sh:44-50：

1if [[ -n "${HF_TOKEN:-}" ]]; then
2    AUTH_HEADER=(-H "Authorization: Bearer $HF_TOKEN")
3fi
4...
5curl -fSL "${AUTH_HEADER[@]}" ...

• 影響：env vars 在多用戶 host 上可被 /proc/<pid>/environ 讀取；shell history 可能留下 export 紀錄。
• 緩解：(a) 改用 HF_TOKEN_FILE 機制（eval.sh 已支援，fetch_bundled_lora.sh 沒）；(b) 把 token 寫進 ~/.cache/huggingface/token（HF SDK 預設位置）；(c) bash export 前加空格避開 history。

🟢 低風險（已 review，無問題）

🔴 高風險：無

6.4 License 風險（不是程式碼風險，但同樣重要）

重點：把 NLD 拿去做商業 inference service / API 之前，請法務看過 NOML 條款。NVIDIA-NeMo/Nemotron 同樣使用 NOML，過去 12 個月內條款有微調，須以最新版為準。

7. FAQ

Q1：為什麼 stars 才 32？是不是不重要？

A：2026-05-18 才開 repo（半個月），且名字 Nemotron-Labs-Diffusion 比較長、品牌認知度不及主家族 NVIDIA-NeMo/Nemotron（1200+ stars）。重要性以「論文 + HF 權重 + SGLang upstream PR」三項一起看，stars 不是好指標。

Q2：能直接拿去 production 嗎？

A：不建議。理由：

1. SGLang 整合走 upstream PR 分支 hutm/sglang @ upstream/2-dllm-lora-ar，main 尚未 merge — 升 SGLang 時要重 rebase。
2. eval.sh 依賴 NeMo-Skills + SLURM + 自建 .sqsh image，工程化複雜度高。
3. 無 issue tracker、無 PR 紀錄、無社群活躍度 — 出問題只能等下次 commit。

可考慮的路徑：(a) 等 SGLang upstream merge；(b) fork 自己長期維護；(c) 等 NVIDIA-NeMo / NVIDIA-NIM 把擴散模式收編進 production stack。

Q3：和 LLaDA / Mercury 等其他擴散 LM 比，NLD 有什麼差異？

A：

• Self-speculation 共用 KV cache：LLaDA / Mercury 的擴散與 AR 是兩個獨立 model（draft + verify），NLD 是同一模型切 attention pattern。
• NVIDIA 工程支援：SGLang + NeMo-Skills + DGX Spark 三層 stack，其他擴散 LM 各自為政。
• License 可商用：LLaDA-8B-Instruct 走 Apache-2.0，Mercury 早期版走 non-commercial；NLD 用 NOML 是「商用可行但有條款」中間值。
• 規模較小：LLaDA 有 8B；NLD 開 3B / 8B / 14B 三個 size。

Q4：dLM 模式的 threshold=0.9 怎麼調？

A：confidence-thresholded unmasking 的門檻 — 越高每步 unmask 的 token 越少、NFE 越多、品質越穩；越低越激進、NFE 少但容易 unmask 錯字。

• 純擴散 baseline 跑 benchmark：0.9（README / eval.sh 預設）
• 互動聊天：0.85（速度感優先）
• 高品質需求（程式碼）：0.95（謹慎）

Q5：linear_spec 的 acceptance length 怎麼觀察？

A：evaluate.py 的 avg_nfe 與 TPF 欄位間接反映。直接觀察要看 xp/nemo-skills/add_nfe_to_metrics.py 寫入的 nfe_log.jsonl，每行記錄 batch 的 draft / accepted / verify 數量。

Q6：3B / 8B / 14B 我該選哪個？

Q7：可以 fine-tune NLD 嗎？

A：本 repo 不涵蓋訓練。要 fine-tune：(a) 用 peft 自己對 8B 跑 LoRA SFT；(b) 用 NVIDIA-NeMo/Nemotron 的 SFT recipe（要先把模型 convert 進 Megatron-Bridge 格式）。 完整 pretrain 從零開始要看 NVIDIA 內部論文，本 repo 無 reference impl。

Q8：MIG / 多 GPU 推論支援？

A：

• 單模型分多 GPU：靠 accelerate / torch.distributed，本 repo 沒給 reference，但 SGLang 啟動參數 --tp <N> 可做 tensor parallel。
• 多模型副本：xp/dlm_api/dlm_load_balancer.py 已做 — 每 GPU 一個 worker。

Q9：和 Cosmos 有什麼互補？

A：Cosmos 是 World Foundation Models（影片 / robotics / 物理模擬）、token domain 是 latent video chunks；NLD 是純語言、token domain 是文字。兩者皆用「擴散思想」但完全不重疊。共同點是 NVIDIA 對「擴散 generative model 在不同領域的工程化」這條路線的雙線投資。

Q10：怎麼判斷 self-speculation 在我的 workload 上有沒有用？

A：跑 evaluate.py 對你的 task：

• 若 linear_spec 的 TPF > 4.0，self-spec 有意義（每 forward 接受 ~4+ token）
• 若 TPF < 2.0，draft acceptance rate 低，不如直接 AR
• 結論強相關於 task 的「樣板化程度」 — 程式碼補完、客服 FAQ、固定格式 reasoning 都受益；高創意 / 開放對話 受益較少

8. 進階技巧

8.1 把 NLD 接成自己內部 OpenAI-compatible endpoint

不走 SGLang upstream 分支也行 — 直接用 xp/dlm_api/dlm_batch_server.py：

 1# 先啟動 batch server（per GPU 一個）
 2CUDA_VISIBLE_DEVICES=0 python xp/dlm_api/dlm_batch_server.py \
 3  --model nvidia/Nemotron-Labs-Diffusion-8B \
 4  --port 8001 \
 5  --engine nemotron \
 6  --max-model-len 20480
 7
 8# 再啟動 load balancer
 9python xp/dlm_api/dlm_load_balancer.py \
10  --upstream "http://localhost:8001" \
11  --port 8000
12
13# 用 OpenAI SDK 打
14curl -X POST http://localhost:8000/v1/chat/completions \
15  -H "Content-Type: application/json" \
16  -d '{"model":"nemotron-labs-diffusion","messages":[...]}'

不需 SLURM、不需 enroot；但要自己處理 process monitoring / restart。

8.2 自己加 benchmark scorer

evaluate.py 的 TASKS 是 plain dict，加 entry 即可：

1TASKS["mmlu_subset"] = Task(
2    name="mmlu_subset",
3    hf_dataset="hails/mmlu_no_train",
4    hf_split="test",
5    question_field="question",
6    gold_extractor=lambda row: row["choices"][row["answer"]],
7    scorer=lambda out, gold: gold.strip().lower() in out.lower()[:300],
8    instruction="Answer the multiple-choice question with the option text.\n\n",
9)

之後 python evaluate.py --mode dlm --tasks mmlu_subset --limit 100 就跑。

8.3 對比實驗：同題在三模式輸出比較

1echo "What is 15% of 240?" > /tmp/q.txt
2for m in ar dlm linear_spec linear_spec_lora; do
3  echo "=== $m ==="
4  python chat/chat.py --mode $m --max-new-tokens 128 < /tmp/q.txt
5done

可以對「同一題、同一個 model、不同 decoding path」拿到 4 份輸出 — 用來 sanity check tri-mode 在你的 prompt 上行為一致。

8.4 把 NFE 紀錄接到自家 metrics 系統

xp/nemo-skills/add_nfe_to_metrics.py 是把每 batch 的 NFE log 合進 metrics.json 的 post-processor。可以參考它寫類似邏輯，把 NFE 推進 prometheus / wandb：

1# 假 code 示意
2import json
3from prometheus_client import Histogram
4nfe_hist = Histogram("nld_nfe_per_request", "NFE distribution")
5for line in open("nfe_log.jsonl"):
6    rec = json.loads(line)
7    for nfe in rec["nfe_list"]:
8        nfe_hist.observe(nfe)

8.5 跟自家 PEFT 微調流程整合（draft LoRA 訓練）

bundled LoRA 是 NVIDIA 自己訓的；若想自己訓一個專屬 draft LoRA：

 1# 高層構想（非完整 code）
 2from peft import LoraConfig, get_peft_model
 3from transformers import AutoModel
 4
 5base = AutoModel.from_pretrained("nvidia/Nemotron-Labs-Diffusion-8B",
 6                                  trust_remote_code=True).cuda()
 7config = LoraConfig(
 8    r=16, lora_alpha=32,
 9    target_modules=["o_proj"],   # 重要：bundled LoRA 也是只在 o_proj
10    task_type="CAUSAL_LM",
11)
12model = get_peft_model(base, config)
13
14# 在你自家 domain text 上跑 draft objective
15# - input：(prompt + masked_block)
16# - target：原始 block（teacher-forcing 還原）
17# 用標準 trainer 跑 SFT

訓完存到 your_lora/，跑 inference 時：

1python chat/chat.py --mode linear_spec_lora \
2  --model nvidia/Nemotron-Labs-Diffusion-8B \
3  --lora-subfolder your_lora

（注意 --lora-subfolder 是 HF 上的 subfolder；本地路徑要對 model 引數做改造）

8.6 量化進階：FP8 + LoRA 同時

SGLang serving 預設組合就是 QUANT=fp8 ALGO=LinearSpec（含 LoRA）。但要注意：

• 線上 BF16 → FP8 quantize 是 weight-only，KV cache 仍 BF16。
• LoRA adapter 不會被 quantize；draft phase 跑 FP8 base + BF16 LoRA delta。
• Acceptance length 通常比 BF16 純跑 略低 5–10% — 是品質/速度 trade-off。

9. 整合進其他工作流

9.1 與本知識庫其他 layer 的接合

9.2 與 Sibling NVIDIA repo 的關係

9.3 整合到 production stack 的 checklist（未來規劃用）

如果要把 NLD 推到自家 production：

• License 法務 review（重點：NOML weights 是否符合服務模式）
• SGLang upstream PR 是否 merge 進 stable release
• 自家 SLURM / K8s cluster 能否提供穩定 GPU 配額
• Monitoring：NFE 分布、acceptance length 分布、P99 latency
• Fallback：擴散模式失敗時自動降級到 AR
• Eval gate：每次 model checkpoint 升級時跑 evaluate.py 防退化
• LoRA 管理：若自家訓 draft LoRA，要有版本化機制

9.4 Discord 分發策略（內部教學分享）

走標準 gh-tutorial-qd 流程：

1. 兩份 md → 兩份 qd → 兩份 plain HTML
2. slim.zip 上傳 Discord（< 1MB）
3. 同事雙擊 index.html 即可在瀏覽器看完整教學

10. 重點摘要 Checklist

• NLD = tri-mode 語言模型 inference 配套：同一份權重，三條解碼路徑（AR / dLM / linear_spec），第三條是「擴散 draft + AR verify 共用 KV cache」的 self-speculation
• 是 NVlabs 不是 NVIDIA-NeMo：研究院 paper-companion code，與 production 線 Nemotron 主家族品牌共享、codebase 不共享
• Dual license：code Apache-2.0，weights NVIDIA Open Model License（NOML）— 商用前法務先看
• 規模小：~1.4K 行程式碼、32 stars、單次 commit；屬論文 drop 性質，不是長期維護專案
• 入門最快：pip install transformers>=5.0 torch → python chat/chat_ar.py
• 標準 benchmark 跑法：python evaluate.py --mode dlm --tasks gsm8k
• 核心 metric：TPF（Tokens Per Forward）= avg_tok / avg_nfe；AR ≈ 1.0、dLM / linear_spec 越大越好
• 三個 mode 各自最佳場景：互動 → linear_spec；批次 → dLM；對比 baseline → AR
• DGX Spark 部署：SGLang lmsysorg/sglang:spark image + hutm/sglang PR 分支 + 137MB LoRA adapter
• 資安整體 🟢 低：唯一中風險是 trust_remote_code=True、SGLang 預設綁 0.0.0.0、HF_TOKEN env 傳遞
• 不適合 production：等 SGLang upstream merge、等 NIM 收編；自己 fork 維護要有覺悟

30 秒口頭摘要（給沒看 tutorial 的同事）

NLD 是 NVlabs 研究院 ship 的「同一份權重三條解碼路徑（純 AR、擴散平行、擴散 draft + AR verify 共用 KV）」的擴散語言模型 inference harness。Code 1.4K 行、Apache-2.0 + NOML 雙 license，目前 32 stars 屬 paper drop 性質。最快入門 python chat/chat_ar.py，正式 benchmark 跑 evaluate.py，DGX Spark 部署走 SGLang upstream PR。資安 🟢 低風險（注意 trust_remote_code=True 與 SGLang 0.0.0.0 預設）。要 production 等 SGLang merge。

11. 進一步閱讀

11.1 本 repo 直接相關

• NVlabs/Nemotron-Labs-Diffusion README
• HuggingFace Collection: Nemotron-Labs-Diffusion
• 8B Model Card
• DGX Spark 產品頁

11.2 上游依賴

• NVIDIA/NeMo-Skills — eval framework
• sgl-project/sglang DLLM-onboarding PR stack #25802
• hutm/sglang @ upstream/2-dllm-lora-ar — 部署用的 fork
• HuggingFace peft — LoRA attach
• HuggingFace transformers — AutoModel + chat template

11.3 本知識庫的姊妹 tutorial

• inbox/2026-06-02-tutorial-Nemotron.md — production Nemotron 主家族（hybrid Mamba-MoE LLM）詳細教學
• inbox/2026-06-02-github-NVIDIA-NeMo-Nemotron.md — Nemotron 主家族 gh-save 簡報

11.4 相關研究方向（建議 paper-search / paper-qa-lite 跟進）

• 擴散語言模型：LLaDA（Nie 2025）、Mercury（Inception Labs 2025）、Plaid（Gulrajani 2023）、DiffusionBERT（He 2023）
• Speculative decoding：Medusa（Cai 2024）、Eagle（Li 2024）、Lookahead Decoding（Fu 2024）
• Self-speculation / 同模型 draft：Layer Skip（Elhoushi 2024）、Kangaroo（Liu 2024）
• 長 context + 擴散：Block Diffusion（Sahoo 2025）

11.5 License 文件（商用前必讀）

• Apache License 2.0（code）
• NVIDIA Open Model License（weights）

11.6 引用 BibTeX

 1@techreport{fu2026nemotronlabsdiffusion,
 2  title       = {Nemotron-Labs-Diffusion: A Tri-Mode Language Model Unifying
 3                 Autoregressive, Diffusion, and Self-Speculation Decoding},
 4  author      = {Yonggan Fu and Lexington Whalen and Abhinav Garg and
 5                 Chengyue Wu and Maksim Khadkevich and Nicolai Oswald and
 6                 Enze Xie and Daniel Egert and Sharath Turuvekere Sreenivas and
 7                 Shizhe Diao and Chenhan Yu and Ye Yu and Weijia Chen and
 8                 Sajad Norouzi and Jingyu Liu and Shiyi Lan and Ligeng Zhu and
 9                 Jin Wang and Jindong Jiang and Morteza Mardani and
10                 Mehran Maghoumi and Song Han and Ante Jukic and
11                 Nima Tajbakhsh and Jan Kautz and Pavlo Molchanov},
12  institution = {NVIDIA},
13  year        = {2026},
14  note        = {Technical report}
15}

教學產出時間：2026-06-02 撰寫工具：gh-tutorial-qd v1.1（自動 transform_mermaid + transform_body_dot_tokens） 對應 gh-save md：inbox/2026-06-02-github-NVlabs-Nemotron-Labs-Diffusion.md