首页
/ DGX Spark 统一内存(UMA)显存核算:用四项模型估算大模型训练能否装入 128GB 内存池

DGX Spark 统一内存(UMA)显存核算:用四项模型估算大模型训练能否装入 128GB 内存池

2026-09-09 15:20:34作者:农烁颖Land

本篇指南以 agents 仓库 dgx-spark-ops 插件中的 UMA Memory Accounting 工作表为核心,讲解如何在启动训练前估算"模型 + 微调方法 + batch/packing 组合"能否装入 NVIDIA DGX Spark 单机的 128GB 统一内存池:掌握权重、优化器状态、梯度、激活四项的核算公式与 dtype 字节表,学会用四个实测锚点校验估算结果,并理解它与 OOM 处置梯、热监控在插件内的协作关系,从而把"启动前按字节规划内存"从口号变成可复制的计算流程。

为什么 UMA 场景需要专门的显存核算

DGX Spark 的 GB10 芯片没有独立 GPU 显存——CPU 与 GPU 共享同一个 128GB 统一内存(UMA)池。这带来两个与离散 GPU 完全不同的行为,也正是核算工作表存在的前提(见 SKILL.md):

  • nvidia-smi / cudaMemGetInfo 会低估内存压力,甚至什么都不报。 它们只报告 CUDA 分配器可见的内存,页缓存与 mmap 页面同样消耗同一个池,机器完全可能显示"还有余量"却照样 OOM;在某些驱动/配置组合下,内存查询直接返回 [N/A], [N/A] 而非数值。因此在 Spark 上应以 free -g 作为预算基准,而不是 128GB 规格数:

    free -g | awk 'NR==2 {print "free:", $4, "GB"}'
    

    经验法则是:取 free 值,扣除数 GB 的操作系统/驱动开销,对剩余量做预算。

  • 模型加载是一个瞬时峰值而非稳态。 加载 safetensors 时会先 mmap 文件,再拷贝进 CUDA 张量——加载期间 mmap 页与 CUDA 拷贝同时在池上计费。一个训练时放得下的模型,如果余量是按加载后稳态而非这个"翻倍瞬态"规划的,仍然可能在加载阶段 OOM。

工作表 uma-accounting.md 的定位是规划数学而非保证(planning math, not a guarantee):它的输入是参数量、dtype、微调方法,输出是一个与已知锚点对比的内存估算;当估算被证明过于乐观时,按 SKILL.md 中的 OOM Ladder 处置,而不是按字节贴边规划——文档明确建议总是保留余量。文件头部还标注了维护纪律:"Last verified: 2026-07-13",当新的模型规模锚点在 Spark 上被验证、或官方 playbook 改变量化默认值时需要刷新,这提醒读者该工作表的结论带有硬件与工具链版本前提。

核算工作表:四项 + 一个可忽略项

总占用 ≈ 权重 + 优化器状态 + 梯度 + 激活,再加上 LoRA/QLoRA 适配器这一近似为零的项。逐项从参数量和 dtype 出发计算,再求和。

第 1 项:权重

params × bytes/param,按 dtype 取值:

dtype 每参数字节数
fp32 4
bf16 / fp16 2
int8 1
int4(QLoRA NF4) 0.5

一个 70B 模型在 bf16 下约 140GB——在其他任何组件加载之前就已经超出整个池;同一模型 4-bit(QLoRA)下约 35GB。这正是"70B 级模型之所以在 Spark 上可行,靠的是 QLoRA 而不是 bf16"的量化依据。

第 2 项:优化器状态

全量微调要为每个可训练参数携带优化器状态;LoRA/QLoRA 只为适配器参数携带,因此无论基座多大,这一项对它们都可忽略:

优化器 每参数字节数(仅可训练参数)
AdamW,fp32 状态 8(4B 动量 + 4B 方差)
AdamW 8-bit(bitsandbytes) ≈2(量化的动量 + 方差)

在 128GB 的机器上,adamw_8bit 是 Unsloth 的默认值是有原因的——对任何不是"仅 LoRA/QLoRA 适配器"的 run,fp32 变体大约会让这一项膨胀 4 倍。仓库中 lora-qlora-recipes 的超参数参考 给出了 optim="adamw_8bit" 的完整训练脚本示例,unsloth-trl-mapping.md 也确认该字符串在 Unsloth 与 TRL SFTConfig 间一一对应、无需翻译——与本文工作表直接兼容。

第 3 项:梯度

与计算精度同 dtype——通常是 bf16,即 2 字节/参数——并且和像优化器状态一样,只针对可训练参数。全量微调要为每个权重支付这份开销;LoRA/QLoRA 只为适配器支付,因为冻结的基座权重从不累积梯度。

第 4 项:激活

最难钉死在单一数字上的一项——它随 batch 大小、序列/packing 长度和架构(注意力变体、hidden 尺寸、层数)变化,而不只取决于参数量。工作表指出,比精确估算更重要的是两个抓手:

  • 梯度检查点(gradient checkpointing) 用重算换内存:相比不启用,期望在这一项上节省约 30%,代价是每个被检查点的段多一遍重算。这也是 Unsloth 将 use_gradient_checkpointing="unsloth" 设为默认的原因。
  • Packing/序列长度是比 batch 大小更直接的杠杆。 这也是 SKILL.md 的 OOM Ladder 把"缩短 packing 长度"放在"缩小 batch"之前的原因。

第 5 项:LoRA/QLoRA 适配器开销

r 的适配器在一个线性层上增加 r × (in + out) 个参数——Ar×inBout×r,两个矩阵合计 r·in + r·out。在实际使用的秩范围(RL 场景 1–32,规模化 SFT 最高约 256)下,这是基座模型规模的千分之几——除非使用异常高的秩,否则在工作表里直接取零。

四个实测锚点:估算之外的一致性校验

工作表第二部分是"已知在单台 Spark 上跑通的组合",用途不是替代公式,而是拿新计划去对照它做 sanity check:

模型规模 方法 实测总占用 备注
70B QLoRA ≈40GB 3 个 epoch 需 30–48 小时;"70B 靠 QLoRA 而非 bf16 能跑"的参考点
~120B 级 MoE 模型 NVFP4 原生 LoRA ≈68GB(据 Unsloth 官方 DGX Spark 教程,2025-12) 社区配方(nvfp4-lora-spark);实验性质,不能作为其他 100B+ MoE 模型的默认假设
27B LoRA pack ≤1024 时放得下 单台 Spark 上的 LoRA 上限——更大的 dense 模型需要多机 Spark 或更小的方法
9B 全量微调 宽裕地放得下 全量 FT 的上限——再往上,全量 FT 必须换 LoRA/QLoRA 或多机 Spark

对"上限"条目的正确读法是:它们是实践中装得下的最大规模,而非硬性架构边界——更小的 batch、更短的 packing 或更省的优化器有时能把某条上限再推一点点;反过来,同一模型类的更重配置也可能远低于上限就失败。超出这四个参考点的组合,应回到上面四项公式重新推导,并在估算被证明乐观时用 OOM Ladder 验证。

实操:70B QLoRA 的完整规划序列

结合工作表与 SKILL.md 的 Planning Sequence,一次启动前的核算按顺序进行:

  1. free -g,扣除 OS/驱动开销得到预算;
  2. uma-accounting.md 估算权重 + 优化器 + 梯度 + 激活;
  3. 与最接近的锚点(70B QLoRA、27B LoRA、9B 全量 FT)对比,而不是只信估算本身;
  4. 若估算贴近预算,先用更短的 packing 或更小的 batch 起步——比 run 中途撞 OOM Ladder 便宜。

SKILL.md 给出的 70B QLoRA sanity check 脚本,验证公式与 ≈40GB 锚点的一致性:

params = 70e9
weights_gb = params * 0.5 / 1e9      # NF4,第 1 项
adapter_gb = 0.5                     # 第 5 项,可忽略
total_gb = weights_gb + adapter_gb   # 再加激活
print(f"{total_gb:.0f}GB before activations")

仅权重一项就落在 ≈40GB 锚点附近——如果同一模型规模的计划估算远超这个数字,就是该回查 dtype 与方法的信号。

估算乐观了怎么办?不要跳级,按 SKILL.md 的 OOM Ladder 逐级执行,缩小 batch 永远不是第 1 步

  1. 刷缓冲缓存:上一轮 run 或大数据集读取留下的页缓存常常占掉"失踪"的数 GB 余量,零配置代价:

    sync; echo 3 > /proc/sys/vm/drop_caches
    

    需要 root;这是 run 之间的复位动作,不是训练中的常规步骤。背后完整诊断见 gotcha-checks.md 的 G3。

  2. 降 batch 或 packing 长度:flush 无效后才动 run 本身的行为;优先缩短 packing 长度,因为长上下文下它更直接地驱动激活占用。

  3. 降级方法:bf16 LoRA 优先于 QLoRA。 前两步仍 OOM 就降一档,方向是 bf16 LoRA 而不是反过来——QLoRA 的 bitsandbytes 反量化缓冲是瞬态的 CUDA 侧分配,可能先于等量的 bf16 LoRA run 触发 OOM,尽管 QLoRA 稳态占用更小。QLoRA OOM 不等于"模型放不下"。

三步走完仍放不下,才考虑更小的模型或多机 Spark。

插件内的落地:preflight 命令与 ops agent

这套核算在仓库里不是孤立的参考文档,而是 dgx-spark-ops 插件诊断链的一环:

  • spark-preflight.md 命令接收"计划中的工作负载"(如 QLoRA 8B, 3 epochs, 8k context),驱动 subagent 执行硬件身份确认、G1–G10 检查,然后**"用 spark-memory-thermal-ops 工作表计算该工作负载的内存余量"**,最后输出 env-report.json(含 headroom_gb 字段与 ready / ready-with-warnings / blocked 判定)。

  • dgx-spark-ops-engineer.md agent 在其第 3 步方法中要求:按工作表估算"权重 + 优化器 + 梯度 + 激活,外加模型加载瞬态峰值",与 free -g 余量(而非 nvidia-smi)对比,把任何贴着预算落地的计划标出来,并引用 Anchors 表中最接近的锚点而不是裸信估算——与本文的工作表用法完全一致;其行为准则中同样写明"nvidia-smi 的余量数字在统一内存下不可信,规划 run 前必须与 free -g 交叉核对"。

  • 多小时的 run 还会撞上持续功耗上限,配套的 thermal-sample.sh 以 30–60 秒间隔把 nvidia-smi 的温度/功耗采样写成 CSV,与训练日志按时间戳对齐:

    bash assets/thermal-sample.sh 30 thermal.log
    

    脚本自带提示:持续 ~100W 是平台功耗上限而非配置 bug(与 gotcha-checks.md G4 的诊断阈值一致),不应为了"修复"正常平台行为去重调 batch 或精度。

  • 启动即失败的场景(ABI 不匹配、flash-attn、playbook 过期)则由 spark-training-gotchas 覆盖,本文的记忆核算假设 job 已经启动。

适用边界

  • 工作表针对 GB10(aarch64、SM121、CUDA 13)单机 128GB UMA 的规划场景;多机 Spark 的跨机策略(DDP/FSDP、禁用 TP)属于 G10 范畴,不在本表适用面内。
  • 锚点是"实践上装得下"的经验值而非架构边界,且绑定验证时的工具链状态(文件标注 Last verified: 2026-07-13);换用新量化默认值或新模型规模时应重新验证。
  • 估算接近预算时,正确动作是启动前收缩 packing/batch,而不是启动后再爬 OOM Ladder;~120B 级 MoE 的 NVFP4 LoRA 锚点是实验性社区配方,不要外推到其他 100B+ MoE 模型。
登录后查看全文
热门项目推荐
相关项目推荐

项目优选

收起
kernelkernel
deepin linux kernel
C
33
18
docsdocs
暂无描述
Markdown
899
5.83 K
ops-transformerops-transformer
本项目是CANN提供的transformer类大模型算子库,实现网络在NPU上加速计算。
C++
1.14 K
2.76 K
pytorchpytorch
作为 Ascend for PyTorch 社区的核心组件,TorchNPU 是昇腾专为 PyTorch 打造的深度学习适配插件,使 PyTorch 框架能够直接调用昇腾 NPU,为开发者提供昇腾 AI 处理器的超强算力。
Python
860
1.35 K
ops-nnops-nn
本项目是CANN提供的神经网络类计算算子库,实现网络在NPU上加速计算。
C++
925
1.85 K
jiuwenswarmjiuwenswarm
JiuwenSwarm 是一款基于openJiuwen开发的智能AI Agent,它能够将大语言模型的强大能力,通过你日常使用的各类通讯应用,直接延伸至你的指尖。
Python
3.84 K
1.02 K
kernelkernel
openEuler内核是openEuler操作系统的核心,既是系统性能与稳定性的基石,也是连接处理器、设备与服务的桥梁。
C
533
601
ops-mathops-math
本项目是CANN提供的数学类基础计算算子库,实现网络在NPU上加速计算。
C++
1.37 K
1.46 K
AscendNPU-IRAscendNPU-IR
AscendNPU-IR是基于MLIR(Multi-Level Intermediate Representation)构建的,面向昇腾亲和算子编译时使用的中间表示,提供昇腾完备表达能力,通过编译优化提升昇腾AI处理器计算效率,支持通过生态框架使能昇腾AI处理器与深度调优
C++
548
395
cann-learning-hubcann-learning-hub
CANN 学习中心仓,支持在线互动运行、边学边练,提供教程、示例与优化方案,一站式助力昇腾开发者快速上手。
Jupyter Notebook
1.04 K
525