Uno: 离散扩散加速 LLM 推理，无损兼高速 · 干货攻略

  • 链接: https://x.com/_akhaliq/status/2097454804272566527
  • 分类: x-tips
  • 来源: X @_akhaliq
  • 作者: Jay
  • 更新: 2026-09-13
  • 仓库: ifm-ai/uno

这是什么

Uno（发音「乌诺」）是 IFM（Institute of Foundation Models）团队开源的离散扩散增强型 LLM 推理框架，核心思想是：不让扩散模型取代自回归模型，而是让扩散模型辅助自回归模型并行生成多个 token，从而在保持输出无损的前提下显著提升吞吐。

传统 LLM 推理慢在「串行」——每个 token 必须等前一个 token 生成完才能生成。投机解码（Speculative Decoding）通过一个小模型「打草稿」再让大模型验证来加速，但需要训两个模型，且草稿质量直接影响效果。Uno 换了一条路：在同一个 LLM 里加入一套轻量扩散权重，专门负责「并行预测多个 token」，然后用一套叫 Ψ-Spec（Psi-Spec） 的采样器从 AR 分布中无损地采样。

项目主页：https://s-sahoo.github.io/uno/ arXiv 论文：https://arxiv.org/abs/2609.04010 开源仓库：https://github.com/ifm-ai/uno（MIT License）

已开源的 Uno 模型包括： - Uno 8B（基于 Qwen3-8B 训练）：s-sahoo/uno-qwen3-8B - Uno 7B：IFM/K2-Horizon-7B-Uno - Uno 1B：IFM/K2-Horizon-0.9B-Uno

为什么值得关注

@_akhaliq 在 X 上推荐了这篇工作，核心卖点是三个「不同于」：

  1. 不同于投机解码：不需要单独训练一个草稿模型，直接复用同一个 LLM 的扩散权重，额外显存开销极小（相比 EAGLE-3 等方案，引入了最少的额外参数）。
  2. 不同于纯扩散 LLM（如 Mercury 2、DiffusionGemma）：Uno 加速的是已有的 AR 分布，不会因为引入扩散过程而降低输出质量，是「无损」加速。
  3. 效果上：Uno 8B 在 agentic tool use、coding、long-context reasoning 三个维度的 benchmark 上，同时超越了 26B 的 DiffusionGemma 和闭源的 Mercury 2。

原帖还提到「~4.6x 吞吐对比 Mercury 2」。需注意：arXiv 论文和官方项目页的表述是「up to 3× speedups over the base AR model」以及「up to 2.5× per-request throughput speedup」。4.6× 这个数字在官方文档中未找到直接对应，可能来自特定 batch size 或特定 benchmark 配置下的对比，建议以论文数据为准。

核验过程

  • ✅ GitHub README（ifm-ai/uno）：确认了双权重架构（AR weights + diffusion weights）、Ψ-Spec 采样器、支持 FlashAttention-2/3、安装步骤、模型 checkpoint 列表。
  • ✅ arXiv 论文摘要（2609.04010）：确认 up to 3× base AR speedup、8B Uno benchmark 超越 Mercury 2 和 DiffusionGemma、无需单独 draft model、diffusion weights 通过 Diffusion Distillation 阶段训练（开销极小）。
  • ✅ Uno 官方项目页（s-sahoo.github.io/uno）：确认了架构图、关键数字（3× / 2.5×）、benchmark 覆盖范围（agentic / coding / long-context）。
  • ⚠️ 交叉验证：「~4.6× 吞吐对比 Mercury 2」——官方文档未找到此数字，论文明确数字为 up to 3× over base AR model。4.6× 可能为特定配置下的结果，攻略正文以官方 3× 为准，原帖 4.6× 说法标注为「原帖主张，未在官方文档中核验到对应数字」。

上手步骤

环境安装

# 创建 Python 3.10 环境
conda create -n nano-vllm-uno python=3.10 pip -y
conda activate nano-vllm-uno

# 升级 pip 并安装 PyTorch
python -m pip install --upgrade pip
python -m pip install torch==2.11.0 \
  --index-url https://download.pytorch.org/whl/cu128

# 安装 FlashAttention-2（线性解码必需）
python -m pip install \
  'https://github.com/lesj0610/flash-attention/releases/download/v2.8.3-cu12-torch2.11/flash_attn-2.8.3%2Bcu12torch2.11cxx11abiTRUE-cp310-cp310-linux_x86_64.whl'

# 安装 Uno（包含 eval 和 train 依赖）
python -m pip install -e '.[eval,train]'

# Tree 验证额外需要 FlashAttention-3
python -m pip install ninja==1.13.0
git clone --depth 1 --branch v2.8.3 \
  https://github.com/Dao-AILab/flash-attention.git
cd flash-attention/hopper
MAX_JOBS=16 python -m pip install --no-build-isolation .

下载模型

公开 checkpoint 无需 Hugging Face token：

# Uno Qwen3 8B（推荐，8B 参数）
# Base: s-sahoo/uno-qwen3-8B
# Diffuson weights (adapter): s-sahoo/uno-qwen3-8B/tree/main/adapter

# Uno 7B
# Base: IFM/K2-Horizon-7B
# Diffusion weights: IFM/K2-Horizon-7B-Uno

# Uno 1B
# Base: IFM/K2-Horizon-0.9B
# Diffusion weights: IFM/K2-Horizon-0.9B-Uno

运行推理

仓库使用 nano_vllm_uno/ 目录下的 Nano-vLLM 推理引擎，通过 generation.py 统一处理模型加载、LoRA 挂载、prompt 格式化和 TPS 统计：

# 使用 Uno Qwen3 8B 的推理脚本（线性采样，高系统吞吐）
bash examples/uno_qwen3_8B/run_linear.sh

# 使用 Tree 采样（高 per-request 吞吐）
bash examples/uno_qwen3_8B/run_tree.sh

或直接调用共享推理入口：

python -m inference \
  --model s-sahoo/uno-qwen3-8B \
  --sampler linear \    # 或 tree
  --prompt "What is the capital of France?"

推理脚本输出的指标包括 TPF（Tokens Per Forward）和 TPS（Tokens Per Second），可直接用于性能评估。

训练自己的 Uno 扩散权重

Uno 支持对已有 AR 模型添加扩散权重进行增强。以 Uno Qwen3 8B 为例：

# 1. 准备 OpenThoughts-3 数据集
python -m training.prepare_openthoughts \
  --output /path/to/openthoughts-uno-4095 \
  --cache-dir /path/to/hf-cache \
  --num-proc 32

# 2. 单卡训练示例（global batch size 128 通过梯度累积实现）
python -m training.train \
  --dataset-path /path/to/openthoughts-uno-4095 \
  --output-dir /path/to/uno-training \
  --curriculum training/configs/uno_3epoch_curriculum.yaml \
  --deepspeed training/configs/deepspeed_zero2.json \
  --per-device-batch-size 8 \
  --gradient-accumulation-steps 16 \
  --learning-rate 1e-5 \
  --warmup-steps 562 \
  --lora-target all \
  --lora-rank 128 \
  --lora-alpha 2048 \
  --ce-alpha 0 \
  --kl-beta 0 \
  --tv-gamma 1

关键参数：LoRA rank 128 / LoRA alpha 2048，使用 conditional-LoRA 机制挂载扩散权重到原模型。

坑与适用边界

  1. FlashAttention-3 依赖：Tree 采样模式需要 FA3，线性采样只需 FA2。没有 FA3 的硬件上不要尝试 tree 模式。
  2. 扩散权重 ≠ 完整模型：Uno 不是一个独立模型，是一套在原 AR 模型上「嫁接」的扩散权重。老型号 AR 模型加上 Uno 权重包才能生效。
  3. 3× speedup 是对比 base AR 模型的上限：具体能跑多快取决于 batch size、序列长度和硬件配置。官方 figure 显示 2.5× per-request throughput 是在特定配置下。
  4. 4.6× 数字来源不明：原帖提到 ~4.6× 吞吐对比 Mercury 2，论文正文和官方项目页均表述为 up to 3× over base AR model。建议在生产对比前先跑自己的 benchmark。
  5. Benchmark 超越不等于所有任务都好：Uno 在 agentic tool use、coding、long-context reasoning 上表现好，未覆盖纯闲聊、创意写作等场景。

一句话结论

Uno 通过在同一 LLM 里嵌入轻量扩散权重 + Ψ-Spec 无损采样器，实现无需草稿模型的无损推理加速，实测 up to 3× 吞吐提升，8B 模型在多个 benchmark 上超越 26B DiffusionGemma 和闭源 Mercury 2——对想给开源 LLM 提速且不愿牺牲输出质量的用户，值得关注。