NVIDIA Star Elastic 实战：一个 Checkpoint 切出 30B/23B/12B 三种推理模型

💡 一句话总结：一个 checkpoint，三种模型——NVIDIA Star Elastic 让你在 30B、23B、12B 之间零成本切换，训练一次，部署三套。

问题：为什么需要弹性推理模型

部署推理模型时，团队经常面临一个两难选择：

NVIDIA 的 Star Elastic 方法给出了第四种选择：训练一个 30B 的模型，里面直接嵌套 23B 和 12B 两个子模型。三个变体共享权重、共存于同一个 checkpoint，提取时只做索引切片，不需要额外微调。

这篇论文已被 ICML 2026 接收，2026 年 5 月 7 日公开发布，模型可商用。

Star Elastic 的核心思路

传统做法是训练三个独立模型。Star Elastic 换了一个思路：

1. 训练一个 30B 的父模型（Nemotron Nano v3，混合 Mamba-Transformer-MoE 架构，3.6B 活跃参数）
2. 在训练过程中，用重要性估计标记每个组件的优先级——embedding channels、attention heads、Mamba SSM heads、MoE experts、FFN channels
3. 较小子模型直接复用父模型中重要性最高的权重子集——23B 复用最重要的权重（2.8B 活跃），12B 复用更小的子集（2.0B 活跃）

关键设计决策是宽度压缩优于深度压缩：

所以 Star Elastic 保持全部 52 层不变，只缩小每层的宽度：

训练使用两阶段课程：

• 短上下文阶段：序列长度 8192 tokens，均匀采样三个变体
• 扩展上下文阶段：序列长度 49152 tokens，非均匀采样（p(30B)=0.5, p(23B)=0.3, p(12B)=0.2），优先保障大模型质量

总训练量约 160B tokens，比从头预训练减少 360 倍。

环境准备

依赖安装

# Python 3.10+，CUDA 12.4+
pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu124
pip install transformers==4.48.0 accelerate==1.2.0 \
    huggingface_hub safetensors sentencepiece

⚠️ 注意：NVFP4 量化需要 NVIDIA Blackwell/Ada 架构 GPU 和对应的 TensorRT-LLM。BF16 和 FP8 版本在 Ampere 及以上架构均可运行。

下载模型

Star Elastic 在 HuggingFace 上提供三种精度的 checkpoint：

# BF16 全精度（约 58.9GB）
huggingface-cli download nvidia/NVIDIA-Nemotron-Labs-3-Elastic-30B-A3B-BF16 \
    --local-dir ./elastic-30b-bf16

# FP8 量化（显存减半）
huggingface-cli download nvidia/NVIDIA-Nemotron-Labs-3-Elastic-30B-A3B-FP8 \
    --local-dir ./elastic-30b-fp8

# NVFP4 量化（仅 18.7GB，RTX 5080 可跑）
huggingface-cli download nvidia/NVIDIA-Nemotron-Labs-3-Elastic-30B-A3B-NVFP4 \
    --local-dir ./elastic-30b-nvfp4

💡 存储对比：三个独立的 BF16 checkpoint（12B+23B+30B）需要 126.1GB，而一个弹性 checkpoint 只需 58.9GB——节省 53%。

加载父模型（30B）

加载完整 30B 模型和加载普通 HuggingFace 模型没有区别：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_path = "./elastic-30b-bf16"

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    device_map="auto",        # 自动分配到可用 GPU
    trust_remote_code=True,   # Nemotron Nano v3 使用自定义架构
)

# 基础推理测试
prompt = "Solve this step by step: What is 23 * 47?"
messages = [{"role": "user", "content": prompt}]
input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=2048,
        temperature=0.6,
        top_p=0.95,
        do_sample=True,
    )

response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
print(response)

提取子模型（23B / 12B）

Star Elastic 的核心价值在这一步——从同一个 checkpoint 中零样本提取不同规模的子模型。

方法一：使用模型配置提取

Nemotron Nano v3 的弹性配置文件中定义了三个变体的维度映射。提取子模型只需要指定目标规模：

from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
import torch
import json

model_path = "./elastic-30b-bf16"

# 加载弹性配置
config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)

# 查看可用的弹性变体
# config 中包含 elastic_config 字段，定义了每个变体的维度
print(f"可用变体: {list(config.elastic_config.keys())}")
# 输出: 可用变体: ['30b', '23b', '12b']

# 提取 23B 子模型
config_23b = AutoConfig.from_pretrained(
    model_path,
    trust_remote_code=True,
    elastic_variant="23b",  # 指定要提取的变体
)

model_23b = AutoModelForCausalLM.from_pretrained(
    model_path,
    config=config_23b,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)

print(f"23B 子模型参数量: {sum(p.numel() for p in model_23b.parameters()) / 1e9:.1f}B")

方法二：手动切片提取

如果你想了解底层机制，可以手动根据重要性排序做权重切片：

import torch
from safetensors.torch import load_file
import json
from pathlib import Path

model_path = Path("./elastic-30b-bf16")

# 加载重要性索引
# 弹性 checkpoint 包含每个组件的重要性排序
importance_path = model_path / "elastic_importance.json"
with open(importance_path) as f:
    importance = json.load(f)

# 以 attention heads 为例：
# 30B 使用 32 个 heads，23B 使用排名前 N 个
attn_head_ranking = importance["attention_heads"]  # [head_idx, ...]，按重要性降序

# 12B 变体只取重要性最高的子集
target_heads_12b = attn_head_ranking[:importance["variants"]["12b"]["num_heads"]]
print(f"12B 变体使用的 attention heads: {sorted(target_heads_12b)}")

# 对 MoE experts 做同样的操作
expert_ranking = importance["moe_experts"]
target_experts_12b = expert_ranking[:importance["variants"]["12b"]["num_experts"]]
print(f"12B 变体使用的 experts 数量: {len(target_experts_12b)}")

保存独立的子模型 checkpoint

提取后可以保存为独立的 checkpoint 方便部署：

# 保存 23B 子模型为独立 checkpoint
output_path = "./nemotron-23b-extracted"
model_23b.save_pretrained(output_path)
tokenizer.save_pretrained(output_path)

print(f"独立 23B checkpoint 大小: {sum(f.stat().st_size for f in Path(output_path).rglob('*.safetensors')) / 1e9:.1f} GB")

弹性预算推理：小模型思考 + 大模型作答

Star Elastic 引入了一种新的推理模式——弹性预算控制：让较小变体处理 CoT（思维链）推理阶段，然后切换到较大变体生成最终答案。

这个思路很直觉：思考过程对精度的要求没有最终答案那么高，用小模型”打草稿”就够了。

实现弹性预算推理

from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
import torch

model_path = "./elastic-30b-bf16"
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 加载 23B 用于 thinking
config_23b = AutoConfig.from_pretrained(
    model_path, trust_remote_code=True, elastic_variant="23b"
)
model_23b = AutoModelForCausalLM.from_pretrained(
    model_path, config=config_23b,
    torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True
)

# 加载 30B 用于 answering
model_30b = AutoModelForCausalLM.from_pretrained(
    model_path, torch_dtype=torch.bfloat16,
    device_map="auto", trust_remote_code=True
)


def elastic_budget_inference(question: str, max_thinking_tokens: int = 1024,
                              max_answer_tokens: int = 512) -> dict:
    """
    弹性预算推理：
    Phase 1 - 用 23B 模型生成思考过程（CoT）
    Phase 2 - 用 30B 模型基于思考结果生成最终答案
    """
    # Phase 1: 23B thinking
    thinking_prompt = f"Think step by step about this problem:\n{question}\n\nThinking:"
    messages = [{"role": "user", "content": thinking_prompt}]
    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(input_text, return_tensors="pt").to(model_23b.device)

    with torch.no_grad():
        thinking_output = model_23b.generate(
            **inputs,
            max_new_tokens=max_thinking_tokens,
            temperature=0.6,
            top_p=0.95,
            do_sample=True,
        )
    thinking_text = tokenizer.decode(
        thinking_output[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True
    )

    # Phase 2: 30B answering
    answer_prompt = (
        f"Based on the following reasoning, provide the final answer.\n\n"
        f"Question: {question}\n\n"
        f"Reasoning: {thinking_text}\n\n"
        f"Final Answer:"
    )
    messages = [{"role": "user", "content": answer_prompt}]
    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(input_text, return_tensors="pt").to(model_30b.device)

    with torch.no_grad():
        answer_output = model_30b.generate(
            **inputs,
            max_new_tokens=max_answer_tokens,
            temperature=0.6,
            top_p=0.95,
            do_sample=True,
        )
    answer_text = tokenizer.decode(
        answer_output[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True
    )

    return {
        "question": question,
        "thinking": thinking_text,
        "answer": answer_text,
    }


# 测试
result = elastic_budget_inference(
    "Find all positive integers n such that n^2 + 2n + 3 is a perfect square."
)
print(f"思考过程:\n{result['thinking'][:500]}...")
print(f"\n最终答案:\n{result['answer']}")

弹性预算的性能收益

论文中的实测数据：

23B-thinking + 30B-answering 是最佳配置：准确率提升 16%，延迟降低 1.9 倍。这看起来反直觉——用更小的模型思考反而更准？论文的解释是，较小模型的 CoT 更简洁、噪声更少，给 30B 的上下文更干净。

⚠️ 注意：目前 vLLM 尚不原生支持推理过程中动态切换模型变体。上面的实现是在应用层手动编排两次推理调用。如果用 TensorRT-LLM 部署，可以通过 Triton Inference Server 的 ensemble model 实现更高效的路由。

性能全面对比

推理基准

关键发现：Star Elastic 的 12B 子模型（78.54）已经接近 Qwen3-30B-A3B 的全尺寸模型（80.00），而它的活跃参数只有 2.0B。

显存与速度

训练成本对比

Star Elastic 比从头预训练节省 360 倍，比传统逐个模型压缩节省 7 倍。

部署建议

场景一：资源充足，追求最高精度

# 直接加载 30B BF16，需要 A100 80GB 或同级 GPU
# 适合离线批处理、研究评测
python inference.py --model ./elastic-30b-bf16 --variant 30b

场景二：单卡部署，平衡精度与速度

# FP8 量化，A100 40GB / RTX 4090 可跑
# 适合在线服务的主力配置
python inference.py --model ./elastic-30b-fp8 --variant 23b

场景三：消费级显卡，极致性价比

# NVFP4 量化 + 12B 变体，RTX 5080 16GB 可跑
# 7426 tokens/s，适合本地开发和边缘部署
python inference.py --model ./elastic-30b-nvfp4 --variant 12b

场景四：弹性预算控制（推荐生产配置）

# 思考用 23B，回答用 30B
# 准确率提升 16%，延迟降低 1.9x
config = {
    "thinking_variant": "23b",
    "answering_variant": "30b",
    "max_thinking_tokens": 2048,
    "max_answer_tokens": 512,
    "quantization": "fp8",  # 视显存选择
}

用 vLLM 部署 30B 推理服务

虽然 vLLM 不支持弹性预算路由，但部署单个变体做推理服务是完全可以的：

# 安装 vLLM（需要支持 Nemotron 架构的版本）
pip install vllm>=0.6.0

# 启动 30B BF16 推理服务
python -m vllm.entrypoints.openai.api_server \
    --model ./elastic-30b-bf16 \
    --trust-remote-code \
    --tensor-parallel-size 2 \
    --max-model-len 32768 \
    --port 8000

客户端调用：

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="unused")

response = client.chat.completions.create(
    model="./elastic-30b-bf16",
    messages=[
        {"role": "user", "content": "Prove that the sum of two odd numbers is always even."}
    ],
    max_tokens=2048,
    temperature=0.6,
)

print(response.choices[0].message.content)

如果显存不够跑 30B，先提取 23B 或 12B 子模型保存为独立 checkpoint，然后用 vLLM 部署提取后的模型即可。

和同类方案的对比

Star Elastic 的独特之处在于它同时支持 Transformer、Mamba 和 MoE 三种架构的弹性压缩，这是之前的工作没有覆盖的。

局限性与注意事项

1. 弹性预算路由尚未集成到主流推理框架：vLLM、TGI 都还不支持推理过程中的动态变体切换，需要在应用层自行编排
2. NVFP4 量化需要特定硬件：目前只有 Blackwell 和 Ada 架构 GPU 支持 NVFP4，Ampere（A100）不支持
3. 当前仅验证于 Nemotron Nano v3 架构：论文没有证明该方法能直接迁移到纯 Transformer 架构（如 Llama），但宽度压缩的思路是通用的
4. 子模型数量有限：目前只嵌入了两个子模型（23B 和 12B），更细粒度的弹性（比如每隔 5B 一个变体）还没有探索

总结

Star Elastic 解决了一个实际的工程问题：怎么用最小的训练和存储开销，获得多种规模的推理模型。核心收益：

• 一个 checkpoint 包含三个变体（30B/23B/12B），存储从 126.1GB 降到 58.9GB
• 训练成本极低：160B tokens，比从头训练减少 360 倍
• 子模型质量高：12B 变体在 AIME-2025 上接近 Qwen3-30B-A3B 全尺寸
• 弹性预算推理：23B 思考 + 30B 回答，准确率提升 16%，延迟降低 1.9 倍
• 消费级硬件可跑：12B NVFP4 在 RTX 5080 上达 7426 tokens/s

对于需要在不同硬件条件下部署同一个推理模型的团队来说，Star Elastic 提供了一个非常优雅的解决方案。

相关资源：

• 论文：arxiv.org/abs/2605.07182
• 模型（BF16）：nvidia/NVIDIA-Nemotron-Labs-3-Elastic-30B-A3B-BF16
• 模型（NVFP4）：nvidia/NVIDIA-Nemotron-Labs-3-Elastic-30B-A3B-NVFP4
• 许可证：Nvidia Open Model License（可商用）