ColossalAI NVMe Offload 使用指南将优化器状态卸载到磁盘突破 GPU 显存墙【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI本文讲解 ColossalAI 中基于 CPUAdam / HybridAdam 的 NVMe offload 特性当 Adam 优化器状态撑爆 GPU乃至 CPU内存时如何通过异步流水线把部分优化器状态卸载到 NVMe 磁盘从而用 I/O 换容量突破显存墙。读完本文你将掌握nvme_offload_fraction/nvme_offload_dir的配置语义、底层异步读写预取/回写机制以及如何在纯 CPU 训练与 GeminiZeRO chunk 内存管理两条路径下落地实践。GPU 显存墙为什么优化器状态会成为瓶颈训练深度学习模型时显存不仅要容纳模型本身还要容纳梯度与优化器状态。以最常用的 Adam 优化器为例它需要为每个参数维护一阶矩exp_avgmomentum与二阶矩exp_avg_sqvariance。当这两个状态以 fp32 保存时一个拥有N个参数的模型其优化器状态总计约为8N字节每个状态每个参数占 4 字节。cpu_adam.py 中num_fp32_shards_per_param 4的注释也印证了这一点——每个参数在 fp32 视角下对应 4 份数据Param weight, grad, momentum and variance。官方文档以 1.16 亿参数的 GPT2-Small 为例其优化器状态约需 0.928 GB0.116 B × 8 bytes而这一数字会随参数规模线性膨胀。当模型达到十亿规模时文档给出的结论是优化器状态至少需要 32 GB 量级的内存。显存容量由此构成训练规模的墙业界称之为GPU 显存墙GPU Memory Wall。NVMe offload 的思路正是针对这道墙将优化器状态中暂时用不到的部分异步写入磁盘待真正更新该参数时再读回内存计算计算完成后再次写回磁盘。这样内存/显存只需驻留正在使用的一小部分优化器状态其容量占用可以从 O(N) 降低到一个近似常数级的滑动窗口。该设计借鉴了两篇代表性论文的思路ZeRO-OffloadDemocratizing Billion-Scale Model Training将优化器状态与梯度卸载到 CPU 内存/计算ZeRO-InfinityBreaking the GPU Memory Wall for Extreme Scale Deep Learning进一步把卸载目标扩展到 NVMe 等异构存储。设计概览TensorNVMe 与读取—计算—卸载三段流水线ColossalAI 并未自行实现底层磁盘 I/O而是依赖同属 HPCAI Tech 生态的异步 Tensor I/O 库TensorNVMe可通过pip install tensornvme安装。文档特别提醒该库与各种磁盘HDD、SATA SSD 和 NVMe SSD兼容。由于 HDD 或 SATA SSD 的 I/O 带宽较低建议仅在 NVMe 磁盘上使用此库。换言之功能层面它对磁盘类型透明但性能收益的前提是存储介质本身具备足够高的随机 I/O 带宽。在更新某个优化器状态如exp_avg、exp_avg_sq时可以把一次优化过程切分为三个连续动作读取read把当前要更新的参数对应状态从磁盘异步读入内存计算compute在内存中完成 Adam 的动量/方差更新与参数更新offload写回把更新后的状态再次异步写回磁盘。ColossalAI 采用流水线方式编排这三步从而让下一个参数的磁盘读取与当前参数的计算相互重叠把 I/O 延迟藏在计算背后。源码拆解NVMeOptimizer基类如何实现异步流水线NVMe offload 的通用逻辑收敛在基类NVMeOptimizernvme_optimizer.py中CPUAdam与HybridAdam都继承自它。理解了这个类就理解了 offload 的全部机制。初始化与 I/O 后端选择构造时若nvme_offload_fraction 0.0会从tensornvme导入DiskOffloader并按顺序执行# nvme_optimizer.py 中的关键逻辑节选 assert 0.0 nvme_offload_fraction 1.0 self.offload_dir offload_dir or tempfile.mkdtemp() backend uring if uring in get_backends() else aio self.offloader DiskOffloader(self.offload_dir, 8, backendbackend)值得注意的源码细节nvme_offload_fraction必须在[0.0, 1.0]区间内越界会直接断言失败nvme_offload_dir传None时自动使用tempfile.mkdtemp()生成的随机临时目录进程退出后__del__会尝试清理该目录I/O 后端在io_uring与 Linux AIO 之间自动选择优先uringDiskOffloader(self.offload_dir, 8, ...)的第 2 个参数 8 表示并发 I/O 队列深度若未安装tensornvme却开启了 fraction 0会抛出ModuleNotFoundError(Please install tensornvme to use NVMeOptimizer)。offload 对象的选择只有 CPU 上的状态会被卸载nvme_offload_fraction描述的是占全部参数元素的比例而不是均匀地分散到每个参数上。在首个 step 时_pre_step会计算可卸载的总元素数self.total_numel self._get_numel() self.can_offload_numel math.floor(self.total_numel * self.nvme_offload_fraction)随后在_post_state_init中按参数遍历顺序先到先得地累计配额if (self.offloader is not None and param.device.type cpu and numel self.offloaded_numel self.can_offload_numel): self.is_on_nvme[param] True self.offloaded_numel numel也就是说只有param.device.type cpu的参数其优化器状态才可能进入 NVMe。这从源码层面印证了文档中那句醒目提示——⚠ 它只会卸载在 CPU 上的优化器状态。这意味着它只会影响 CPU 训练或者使用卸载的 Zero/Gemini。GPU 上参数的状态不在考虑范围之内因此该功能对纯 GPU 显存驻留的普通训练并无意义必须配合 CPU 训练或把参数/状态放置在 CPU 的 Zero/Gemini 场景。异步流水线的四个钩子NVMeOptimizer把流水线拆成_pre_step → _pre_update → _post_update → _post_step四个生命周期钩子子类CPUAdam / HybridAdam只需在自己的step()中于合适时机调用钩子时机行为_pre_step一个 step 开始遍历参数组之前计算can_offload_numel建立待预取参数表prefetch_params对第一个参数的状态发起async_read_pre_update更新某个 CPU 参数之前sync_read_events()等待该参数状态读回内存随后对下一个参数发起async_read预取_post_update更新某个 CPU 参数之后sync_write_events()等待写事件排空若该参数is_on_nvme对其状态发起async_write_post_step一个 step 结束时offloader.synchronize()等待全部异步写完成清空预取表可见流水线只保留一个正在计算 一个正在预取的滑动窗口前一个参数的读与当前参数的计算重叠写回则异步排队从而最大化计算与 I/O 的重叠度。checkpoint 相关的限制同样在NVMeOptimizer源码中能看到state_dict()与load_state_dict()在启用 offloadoffloader is not None时直接raise NotImplementedError代码注释说明状态本身体积可能超过内存容量直接序列化可能导致 OOM。这意味着启用 NVMe offload 的训练流程暂不支持常规的优化器 checkpoint 存取规划断点续训方案时需要格外留意。落地的两类优化器CPUAdam 与 HybridAdamNVMe offload 只在两类优化器上开放二者都以init.py 中导出的CPUAdam、HybridAdam形式对外提供。CPUAdamCPU SIMD 加速的 AdamCPUAdamcpu_adam.py继承NVMeOptimizer其 CPU 侧更新由CPUAdamLoader().load()加载的融合内核完成文档/类注释表明该内核使用 SIMD 指令AVX2 / AVX512加速且需要安装时以BUILD_EXT1方式构建扩展。一个容易踩的坑CPU Adam 融合内核目前不支持 bf16 梯度因此当p.grad.dtype is torch.bfloat16时会自动回退到类内的torch_adam_update基于原生 PyTorch 算子实现的逐参数更新功能等价但性能低于融合内核。HybridAdamCPU 走 CPUAdam、GPU 走 FusedAdam 的混合体HybridAdamhybrid_adam.py继承CPUAdam其 docstring 把设计目标说得很清楚它是 CPUAdam 与 FusedAdam 的混合——位于 CPU 的参数用 CPUAdam 语义更新位于 GPU 的参数走FusedOptimizerLoader加载的multi_tensor_adam融合内核。在step()内部按target_device.type分发cpu/npu校验状态与参数同设备后走 CPU 更新路径同样内嵌_pre_update/_post_update即 NVMe offload 的生效位置bf16 梯度或 npu 时回退torch_adam_updatecuda把梯度、参数、一阶矩、二阶矩收集进列表一次multi_tensor_applier批量更新天然适配 Gemini 将不同参数放在不同设备的场景。从源码可见HybridAdam的构造参数集合参数默认值说明lr1e-3学习率betas(0.9, 0.999)Adam 一阶/二阶矩系数eps1e-8数值稳定项weight_decay0权重衰减adamw_modeTrueTrue 表示解耦权重衰减AdamWFalse 表示 L2 正则bias_correctionTrue是否做偏差修正nvme_offload_fraction0.0卸载到 NVMe 的优化器状态比例取值[0.0, 1.0]nvme_offload_dirNone保存 offload 文件的目录为None时使用随机临时目录与并行方法的兼容性文档明确声明它与 ColossalAI 中的所有并行方法兼容。兼容性的关键恰恰在于上文那则 ⚠ 提示offload 只作用于位于 CPU 上的优化器状态。因此要真正吃到 NVMe offload 的收益训练配置必须满足其一纯CPU 训练参数与梯度都在 CPU带 offload 的 ZeRO / Gemini——即 chunk 化后由内存策略把参数与优化器状态放上 CPU 的训练方式GeminiPlugin 的 placement policy 设为cpu/auto/const。快速开始安装与最小用法先安装TensorNVMe以及其依赖packagingpip install packaging pip install tensornvme接着用官方支持的两类优化器之一构造带 NVMe offload 的优化器from colossalai.nn.optimizer import CPUAdam, HybridAdam # 以 HybridAdam 为例 optimizer HybridAdam(model.parameters(), lr1e-3, nvme_offload_fraction1.0, nvme_offload_dir./)参数语义nvme_offload_fraction要卸载到 NVMe 的优化器状态比例0.0 表示不启用1.0 表示在满足条件的前提下全量卸载nvme_offload_dir保存 NVMe offload 文件的目录为None时使用随机临时目录进程退出即清理。实战示例一纯 CPU 训练 GPT2下面用完整的端到端例子观察 offload 效果。示例依赖transformers与psutil先安装pip install psutil transformers导入必要模块import os import time from typing import Dict, Optional import psutil import torch import torch.nn as nn from transformers.models.gpt2.configuration_gpt2 import GPT2Config from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel import colossalai from colossalai.nn.optimizer import HybridAdam from colossalai.utils.model.colo_init_context import ColoInitContext from colossalai.booster import Booster from colossalai.booster.plugin import GeminiPlugin定义损失函数对logits做 shift 后计算交叉熵class GPTLMLoss(nn.Module): def __init__(self): super().__init__() self.loss_fn nn.CrossEntropyLoss() def forward(self, logits, labels): shift_logits logits[..., :-1, :].contiguous() shift_labels labels[..., 1:].contiguous() # Flatten the tokens return self.loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))再定义三个工具函数生成随机数据、统计模型参数量、统计当前进程 RSS 内存def get_data(batch_size: int, seq_len: int, vocab_size: int, device: Optional[str] None) - Dict[str, torch.Tensor]: device torch.cuda.current_device() if device is None else device input_ids torch.randint(vocab_size, (batch_size, seq_len), devicedevice) attn_mask torch.ones_like(input_ids) return dict(input_idsinput_ids, attention_maskattn_mask) def get_model_numel(model: nn.Module) - int: return sum(p.numel() for p in model.parameters()) def get_mem_usage() - int: proc psutil.Process(os.getpid()) return proc.memory_info().rss训练主函数接受nvme_offload_fraction其余保持一致便于对照def train_cpu(nvme_offload_fraction: float 0.0): config GPT2Config() model GPT2LMHeadModel(config) criterion GPTLMLoss() optimizer HybridAdam(model.parameters(), nvme_offload_fractionnvme_offload_fraction) print(fModel numel: {get_model_numel(model) / 1024**3:.3f} B) start time.time() for step in range(3): data get_data(4, 128, config.vocab_size, devicecpu) outputs model(**data) loss criterion(outputs.logits, data[input_ids]) loss.backward() optimizer.step() optimizer.zero_grad() print(f[{step}] loss: {loss.item():.3f}) print(fTime: {time.time() - start:.3f} s) print(fMem usage: {get_mem_usage() / 1024**2:.3f} MB)不使用 NVMe 卸载train_cpu(0.0)可能得到如下输出Model numel: 0.116 B [0] loss: 10.953 [1] loss: 10.974 [2] loss: 10.965 Time: 7.739 s Mem usage: 5966.445 MB使用全量 NVMe 卸载train_cpu(1.0)可能得到Model numel: 0.116 B [0] loss: 10.951 [1] loss: 10.994 [2] loss: 10.984 Time: 8.527 s Mem usage: 4968.016 MB结果解读对于约 1.16 亿参数的 GPT2-Small其优化器状态8N 字节约为 0.928 GB。NVMe 卸载让进程 RSS 从约 5966 MB 降到约 4968 MB节省约 998 MB与理论上限基本吻合代价是 3 步训练耗时从 7.739 s 增加到 8.527 s——这正是文档预期中用时间I/O换空间内存的体现。实战示例二GeminiPlugin 下把优化器状态卸载到 NVMe第二个场景结合 ColossalAI 的Gemini基于 Chunk 内存管理的 ZeRO前置教程见 zero_with_chunk。先用ColoInitContext让模型在 GPU 上惰性初始化再用GeminiPlugin接管训练。文档强调placement policy 应设置为auto、cpu或const——因为只有把参数/状态放到 CPU 的配置才谈得上对 CPU 上的优化器状态做 NVMe offload。def train_gemini_cpu(nvme_offload_fraction: float 0.0): colossalai.launch_from_torch() config GPT2Config() with ColoInitContext(devicetorch.cuda.current_device()): model GPT2LMHeadModel(config) criterion GPTLMLoss() optimizer HybridAdam(model.parameters(), nvme_offload_fractionnvme_offload_fraction) print(fModel numel: {get_model_numel(model) / 1024**3:.3f} B) plugin GeminiPlugin( strict_ddp_modeTrue, devicetorch.cuda.current_device(), placement_policycpu, pin_memoryTrue, hidden_dimconfig.n_embd, initial_scale2**5 ) booster Booster(plugin) model, optimizer, criterion, *_ booster.boost(model, optimizer, criterion) start time.time() for step in range(3): data get_data(4, 128, config.vocab_size) outputs model(**data) loss criterion(outputs.logits, data[input_ids]) booster.backward(loss, optimizer) optimizer.step() optimizer.zero_grad() print(f[{step}] loss: {loss.item():.3f}) print(fTime: {time.time() - start:.3f} s) print(fMem usage: {get_mem_usage() / 1024**2:.3f} MB)说明原文档此处写作model, optimizer, criterion, _* booster.boost(...)_*并非合法 Python 语法若需直接运行应写作model, optimizer, criterion, *_ booster.boost(...)将 boost 返回的其余组件收集进列表即可。Booster 与插件的更完整用法可参见 GeminiPlugin 源码。不使用 NVMe 卸载train_gemini_cpu(0.0)可能得到Model numel: 0.116 B searching chunk configuration is completed in 0.27 s. used number: 118.68 MB, wasted number: 0.75 MB total wasted percentage is 0.63% [0] loss: 10.953 [1] loss: 10.938 [2] loss: 10.969 Time: 2.997 s Mem usage: 5592.227 MB使用全量 NVMe 卸载train_gemini_cpu(1.0)可能得到Model numel: 0.116 B searching chunk configuration is completed in 0.27 s. used number: 118.68 MB, wasted number: 0.75 MB total wasted percentage is 0.63% [0] loss: 10.953 [1] loss: 10.938 [2] loss: 10.969 Time: 3.691 s Mem usage: 5298.344 MB这里只观察到约 294 MB 的内存下降而并非像纯 CPU 场景那样接近 1 GB。文档给出的解释与实验设计直接相关示例中开启了 GeminiPlugin 的pin_memoryTrue它会为 chunk 分配**锁页内存pinned memory**以加速 host↔device 拷贝而这部分内存同样计入进程 RSS抵消了一部分 offload 收益。文档同时指出如果关闭pin_memory仍可观察到约 900 MB 的内存占用下降。因此实际部署时需要在训练速度pin_memory 带来的拷贝加速与内存占用之间做权衡。正确性验证仓库中的单元测试仓库在 tests/test_optimizer/test_nvme.py 中提供了 offload 正确性的对照测试值得作为二次开发的参考parameterize(nvme_offload_fraction, [0.0, 0.5, 1.0]) parameterize(nvme_offload_dir, [./offload, None]) parameterize(adam_cls, [CPUAdam, HybridAdam]) def test_nvme_adam(nvme_offload_fraction, nvme_offload_dir, adam_cls): ...测试要点参数化覆盖nvme_offload_fraction ∈ {0.0, 0.5, 1.0}、nvme_offload_dir ∈ {./offload, None}以及CPUAdam/HybridAdam两种优化器构造部分参数在 CUDA、部分参数在 CPU的混合模型move_some_params_to_cuda再以随机梯度连跑 3 步与同参数的torch.optim.Adam逐参数比对torch.allclose容差atol1e-3验证 offload 前后数值行为不漂移注意该用例在文件头有pytest.mark.skip(reasonskip because of something wrong with CI)即当前 CI 中默认跳过此测试。另外本文档在 Docusaurus 体系中对应一个可运行的 doc-test原文档末尾的注释torchrun --standalone --nproc_per_node1 nvme_offload.py意味着整套示例本可按单机单卡方式直接跑通复现。关键结论与使用红线最后把使用 NVMe offload 时最容易踩的坑与前提收敛如下介质功能上兼容 HDD / SATA SSD / NVMe SSD但低带宽磁盘会让流水线退化为 I/O 瓶颈建议只在 NVMe 上开启生效范围只卸载CPU 上的优化器状态。纯 GPU 显存训练不会从中获益需配合 CPU 训练或带 CPU offload 的 ZeRO / Gemini如 GeminiPlugin 的placement_policycpu/auto/const比例语义nvme_offload_fraction是可卸载元素占全体参数元素的比例按参数遍历顺序累计分配取值范围[0.0, 1.0]目录语义nvme_offload_dirNone时使用随机临时目录进程退出自动清理自定义目录请保证路径可写且有足够磁盘空间性能权衡offload 减少内存占用但引入 I/O 时间pin_memory可加速拷贝但本身占用 RSS需按场景取舍checkpoint 限制启用 offload 时优化器的state_dict()/load_state_dict()会因潜在 OOM 抛出NotImplementedError断点续训需另行设计状态保存方案数值正确性可借助tests/test_optimizer/test_nvme.py的对照测试验证自己的接入没有改变优化轨迹。当模型规模持续增大、显存与内存都告急时把八字节每参数的 Adam 状态交给 NVMe是 ColossalAI 在 ZeRO-Offload / ZeRO-Infinity 路线上给出的一个工程化、即插即用的答案。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
ColossalAI NVMe Offload 使用指南：将优化器状态卸载到磁盘，突破 GPU 显存墙
ColossalAI NVMe Offload 使用指南将优化器状态卸载到磁盘突破 GPU 显存墙【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI本文讲解 ColossalAI 中基于 CPUAdam / HybridAdam 的 NVMe offload 特性当 Adam 优化器状态撑爆 GPU乃至 CPU内存时如何通过异步流水线把部分优化器状态卸载到 NVMe 磁盘从而用 I/O 换容量突破显存墙。读完本文你将掌握nvme_offload_fraction/nvme_offload_dir的配置语义、底层异步读写预取/回写机制以及如何在纯 CPU 训练与 GeminiZeRO chunk 内存管理两条路径下落地实践。GPU 显存墙为什么优化器状态会成为瓶颈训练深度学习模型时显存不仅要容纳模型本身还要容纳梯度与优化器状态。以最常用的 Adam 优化器为例它需要为每个参数维护一阶矩exp_avgmomentum与二阶矩exp_avg_sqvariance。当这两个状态以 fp32 保存时一个拥有N个参数的模型其优化器状态总计约为8N字节每个状态每个参数占 4 字节。cpu_adam.py 中num_fp32_shards_per_param 4的注释也印证了这一点——每个参数在 fp32 视角下对应 4 份数据Param weight, grad, momentum and variance。官方文档以 1.16 亿参数的 GPT2-Small 为例其优化器状态约需 0.928 GB0.116 B × 8 bytes而这一数字会随参数规模线性膨胀。当模型达到十亿规模时文档给出的结论是优化器状态至少需要 32 GB 量级的内存。显存容量由此构成训练规模的墙业界称之为GPU 显存墙GPU Memory Wall。NVMe offload 的思路正是针对这道墙将优化器状态中暂时用不到的部分异步写入磁盘待真正更新该参数时再读回内存计算计算完成后再次写回磁盘。这样内存/显存只需驻留正在使用的一小部分优化器状态其容量占用可以从 O(N) 降低到一个近似常数级的滑动窗口。该设计借鉴了两篇代表性论文的思路ZeRO-OffloadDemocratizing Billion-Scale Model Training将优化器状态与梯度卸载到 CPU 内存/计算ZeRO-InfinityBreaking the GPU Memory Wall for Extreme Scale Deep Learning进一步把卸载目标扩展到 NVMe 等异构存储。设计概览TensorNVMe 与读取—计算—卸载三段流水线ColossalAI 并未自行实现底层磁盘 I/O而是依赖同属 HPCAI Tech 生态的异步 Tensor I/O 库TensorNVMe可通过pip install tensornvme安装。文档特别提醒该库与各种磁盘HDD、SATA SSD 和 NVMe SSD兼容。由于 HDD 或 SATA SSD 的 I/O 带宽较低建议仅在 NVMe 磁盘上使用此库。换言之功能层面它对磁盘类型透明但性能收益的前提是存储介质本身具备足够高的随机 I/O 带宽。在更新某个优化器状态如exp_avg、exp_avg_sq时可以把一次优化过程切分为三个连续动作读取read把当前要更新的参数对应状态从磁盘异步读入内存计算compute在内存中完成 Adam 的动量/方差更新与参数更新offload写回把更新后的状态再次异步写回磁盘。ColossalAI 采用流水线方式编排这三步从而让下一个参数的磁盘读取与当前参数的计算相互重叠把 I/O 延迟藏在计算背后。源码拆解NVMeOptimizer基类如何实现异步流水线NVMe offload 的通用逻辑收敛在基类NVMeOptimizernvme_optimizer.py中CPUAdam与HybridAdam都继承自它。理解了这个类就理解了 offload 的全部机制。初始化与 I/O 后端选择构造时若nvme_offload_fraction 0.0会从tensornvme导入DiskOffloader并按顺序执行# nvme_optimizer.py 中的关键逻辑节选 assert 0.0 nvme_offload_fraction 1.0 self.offload_dir offload_dir or tempfile.mkdtemp() backend uring if uring in get_backends() else aio self.offloader DiskOffloader(self.offload_dir, 8, backendbackend)值得注意的源码细节nvme_offload_fraction必须在[0.0, 1.0]区间内越界会直接断言失败nvme_offload_dir传None时自动使用tempfile.mkdtemp()生成的随机临时目录进程退出后__del__会尝试清理该目录I/O 后端在io_uring与 Linux AIO 之间自动选择优先uringDiskOffloader(self.offload_dir, 8, ...)的第 2 个参数 8 表示并发 I/O 队列深度若未安装tensornvme却开启了 fraction 0会抛出ModuleNotFoundError(Please install tensornvme to use NVMeOptimizer)。offload 对象的选择只有 CPU 上的状态会被卸载nvme_offload_fraction描述的是占全部参数元素的比例而不是均匀地分散到每个参数上。在首个 step 时_pre_step会计算可卸载的总元素数self.total_numel self._get_numel() self.can_offload_numel math.floor(self.total_numel * self.nvme_offload_fraction)随后在_post_state_init中按参数遍历顺序先到先得地累计配额if (self.offloader is not None and param.device.type cpu and numel self.offloaded_numel self.can_offload_numel): self.is_on_nvme[param] True self.offloaded_numel numel也就是说只有param.device.type cpu的参数其优化器状态才可能进入 NVMe。这从源码层面印证了文档中那句醒目提示——⚠ 它只会卸载在 CPU 上的优化器状态。这意味着它只会影响 CPU 训练或者使用卸载的 Zero/Gemini。GPU 上参数的状态不在考虑范围之内因此该功能对纯 GPU 显存驻留的普通训练并无意义必须配合 CPU 训练或把参数/状态放置在 CPU 的 Zero/Gemini 场景。异步流水线的四个钩子NVMeOptimizer把流水线拆成_pre_step → _pre_update → _post_update → _post_step四个生命周期钩子子类CPUAdam / HybridAdam只需在自己的step()中于合适时机调用钩子时机行为_pre_step一个 step 开始遍历参数组之前计算can_offload_numel建立待预取参数表prefetch_params对第一个参数的状态发起async_read_pre_update更新某个 CPU 参数之前sync_read_events()等待该参数状态读回内存随后对下一个参数发起async_read预取_post_update更新某个 CPU 参数之后sync_write_events()等待写事件排空若该参数is_on_nvme对其状态发起async_write_post_step一个 step 结束时offloader.synchronize()等待全部异步写完成清空预取表可见流水线只保留一个正在计算 一个正在预取的滑动窗口前一个参数的读与当前参数的计算重叠写回则异步排队从而最大化计算与 I/O 的重叠度。checkpoint 相关的限制同样在NVMeOptimizer源码中能看到state_dict()与load_state_dict()在启用 offloadoffloader is not None时直接raise NotImplementedError代码注释说明状态本身体积可能超过内存容量直接序列化可能导致 OOM。这意味着启用 NVMe offload 的训练流程暂不支持常规的优化器 checkpoint 存取规划断点续训方案时需要格外留意。落地的两类优化器CPUAdam 与 HybridAdamNVMe offload 只在两类优化器上开放二者都以init.py 中导出的CPUAdam、HybridAdam形式对外提供。CPUAdamCPU SIMD 加速的 AdamCPUAdamcpu_adam.py继承NVMeOptimizer其 CPU 侧更新由CPUAdamLoader().load()加载的融合内核完成文档/类注释表明该内核使用 SIMD 指令AVX2 / AVX512加速且需要安装时以BUILD_EXT1方式构建扩展。一个容易踩的坑CPU Adam 融合内核目前不支持 bf16 梯度因此当p.grad.dtype is torch.bfloat16时会自动回退到类内的torch_adam_update基于原生 PyTorch 算子实现的逐参数更新功能等价但性能低于融合内核。HybridAdamCPU 走 CPUAdam、GPU 走 FusedAdam 的混合体HybridAdamhybrid_adam.py继承CPUAdam其 docstring 把设计目标说得很清楚它是 CPUAdam 与 FusedAdam 的混合——位于 CPU 的参数用 CPUAdam 语义更新位于 GPU 的参数走FusedOptimizerLoader加载的multi_tensor_adam融合内核。在step()内部按target_device.type分发cpu/npu校验状态与参数同设备后走 CPU 更新路径同样内嵌_pre_update/_post_update即 NVMe offload 的生效位置bf16 梯度或 npu 时回退torch_adam_updatecuda把梯度、参数、一阶矩、二阶矩收集进列表一次multi_tensor_applier批量更新天然适配 Gemini 将不同参数放在不同设备的场景。从源码可见HybridAdam的构造参数集合参数默认值说明lr1e-3学习率betas(0.9, 0.999)Adam 一阶/二阶矩系数eps1e-8数值稳定项weight_decay0权重衰减adamw_modeTrueTrue 表示解耦权重衰减AdamWFalse 表示 L2 正则bias_correctionTrue是否做偏差修正nvme_offload_fraction0.0卸载到 NVMe 的优化器状态比例取值[0.0, 1.0]nvme_offload_dirNone保存 offload 文件的目录为None时使用随机临时目录与并行方法的兼容性文档明确声明它与 ColossalAI 中的所有并行方法兼容。兼容性的关键恰恰在于上文那则 ⚠ 提示offload 只作用于位于 CPU 上的优化器状态。因此要真正吃到 NVMe offload 的收益训练配置必须满足其一纯CPU 训练参数与梯度都在 CPU带 offload 的 ZeRO / Gemini——即 chunk 化后由内存策略把参数与优化器状态放上 CPU 的训练方式GeminiPlugin 的 placement policy 设为cpu/auto/const。快速开始安装与最小用法先安装TensorNVMe以及其依赖packagingpip install packaging pip install tensornvme接着用官方支持的两类优化器之一构造带 NVMe offload 的优化器from colossalai.nn.optimizer import CPUAdam, HybridAdam # 以 HybridAdam 为例 optimizer HybridAdam(model.parameters(), lr1e-3, nvme_offload_fraction1.0, nvme_offload_dir./)参数语义nvme_offload_fraction要卸载到 NVMe 的优化器状态比例0.0 表示不启用1.0 表示在满足条件的前提下全量卸载nvme_offload_dir保存 NVMe offload 文件的目录为None时使用随机临时目录进程退出即清理。实战示例一纯 CPU 训练 GPT2下面用完整的端到端例子观察 offload 效果。示例依赖transformers与psutil先安装pip install psutil transformers导入必要模块import os import time from typing import Dict, Optional import psutil import torch import torch.nn as nn from transformers.models.gpt2.configuration_gpt2 import GPT2Config from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel import colossalai from colossalai.nn.optimizer import HybridAdam from colossalai.utils.model.colo_init_context import ColoInitContext from colossalai.booster import Booster from colossalai.booster.plugin import GeminiPlugin定义损失函数对logits做 shift 后计算交叉熵class GPTLMLoss(nn.Module): def __init__(self): super().__init__() self.loss_fn nn.CrossEntropyLoss() def forward(self, logits, labels): shift_logits logits[..., :-1, :].contiguous() shift_labels labels[..., 1:].contiguous() # Flatten the tokens return self.loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))再定义三个工具函数生成随机数据、统计模型参数量、统计当前进程 RSS 内存def get_data(batch_size: int, seq_len: int, vocab_size: int, device: Optional[str] None) - Dict[str, torch.Tensor]: device torch.cuda.current_device() if device is None else device input_ids torch.randint(vocab_size, (batch_size, seq_len), devicedevice) attn_mask torch.ones_like(input_ids) return dict(input_idsinput_ids, attention_maskattn_mask) def get_model_numel(model: nn.Module) - int: return sum(p.numel() for p in model.parameters()) def get_mem_usage() - int: proc psutil.Process(os.getpid()) return proc.memory_info().rss训练主函数接受nvme_offload_fraction其余保持一致便于对照def train_cpu(nvme_offload_fraction: float 0.0): config GPT2Config() model GPT2LMHeadModel(config) criterion GPTLMLoss() optimizer HybridAdam(model.parameters(), nvme_offload_fractionnvme_offload_fraction) print(fModel numel: {get_model_numel(model) / 1024**3:.3f} B) start time.time() for step in range(3): data get_data(4, 128, config.vocab_size, devicecpu) outputs model(**data) loss criterion(outputs.logits, data[input_ids]) loss.backward() optimizer.step() optimizer.zero_grad() print(f[{step}] loss: {loss.item():.3f}) print(fTime: {time.time() - start:.3f} s) print(fMem usage: {get_mem_usage() / 1024**2:.3f} MB)不使用 NVMe 卸载train_cpu(0.0)可能得到如下输出Model numel: 0.116 B [0] loss: 10.953 [1] loss: 10.974 [2] loss: 10.965 Time: 7.739 s Mem usage: 5966.445 MB使用全量 NVMe 卸载train_cpu(1.0)可能得到Model numel: 0.116 B [0] loss: 10.951 [1] loss: 10.994 [2] loss: 10.984 Time: 8.527 s Mem usage: 4968.016 MB结果解读对于约 1.16 亿参数的 GPT2-Small其优化器状态8N 字节约为 0.928 GB。NVMe 卸载让进程 RSS 从约 5966 MB 降到约 4968 MB节省约 998 MB与理论上限基本吻合代价是 3 步训练耗时从 7.739 s 增加到 8.527 s——这正是文档预期中用时间I/O换空间内存的体现。实战示例二GeminiPlugin 下把优化器状态卸载到 NVMe第二个场景结合 ColossalAI 的Gemini基于 Chunk 内存管理的 ZeRO前置教程见 zero_with_chunk。先用ColoInitContext让模型在 GPU 上惰性初始化再用GeminiPlugin接管训练。文档强调placement policy 应设置为auto、cpu或const——因为只有把参数/状态放到 CPU 的配置才谈得上对 CPU 上的优化器状态做 NVMe offload。def train_gemini_cpu(nvme_offload_fraction: float 0.0): colossalai.launch_from_torch() config GPT2Config() with ColoInitContext(devicetorch.cuda.current_device()): model GPT2LMHeadModel(config) criterion GPTLMLoss() optimizer HybridAdam(model.parameters(), nvme_offload_fractionnvme_offload_fraction) print(fModel numel: {get_model_numel(model) / 1024**3:.3f} B) plugin GeminiPlugin( strict_ddp_modeTrue, devicetorch.cuda.current_device(), placement_policycpu, pin_memoryTrue, hidden_dimconfig.n_embd, initial_scale2**5 ) booster Booster(plugin) model, optimizer, criterion, *_ booster.boost(model, optimizer, criterion) start time.time() for step in range(3): data get_data(4, 128, config.vocab_size) outputs model(**data) loss criterion(outputs.logits, data[input_ids]) booster.backward(loss, optimizer) optimizer.step() optimizer.zero_grad() print(f[{step}] loss: {loss.item():.3f}) print(fTime: {time.time() - start:.3f} s) print(fMem usage: {get_mem_usage() / 1024**2:.3f} MB)说明原文档此处写作model, optimizer, criterion, _* booster.boost(...)_*并非合法 Python 语法若需直接运行应写作model, optimizer, criterion, *_ booster.boost(...)将 boost 返回的其余组件收集进列表即可。Booster 与插件的更完整用法可参见 GeminiPlugin 源码。不使用 NVMe 卸载train_gemini_cpu(0.0)可能得到Model numel: 0.116 B searching chunk configuration is completed in 0.27 s. used number: 118.68 MB, wasted number: 0.75 MB total wasted percentage is 0.63% [0] loss: 10.953 [1] loss: 10.938 [2] loss: 10.969 Time: 2.997 s Mem usage: 5592.227 MB使用全量 NVMe 卸载train_gemini_cpu(1.0)可能得到Model numel: 0.116 B searching chunk configuration is completed in 0.27 s. used number: 118.68 MB, wasted number: 0.75 MB total wasted percentage is 0.63% [0] loss: 10.953 [1] loss: 10.938 [2] loss: 10.969 Time: 3.691 s Mem usage: 5298.344 MB这里只观察到约 294 MB 的内存下降而并非像纯 CPU 场景那样接近 1 GB。文档给出的解释与实验设计直接相关示例中开启了 GeminiPlugin 的pin_memoryTrue它会为 chunk 分配**锁页内存pinned memory**以加速 host↔device 拷贝而这部分内存同样计入进程 RSS抵消了一部分 offload 收益。文档同时指出如果关闭pin_memory仍可观察到约 900 MB 的内存占用下降。因此实际部署时需要在训练速度pin_memory 带来的拷贝加速与内存占用之间做权衡。正确性验证仓库中的单元测试仓库在 tests/test_optimizer/test_nvme.py 中提供了 offload 正确性的对照测试值得作为二次开发的参考parameterize(nvme_offload_fraction, [0.0, 0.5, 1.0]) parameterize(nvme_offload_dir, [./offload, None]) parameterize(adam_cls, [CPUAdam, HybridAdam]) def test_nvme_adam(nvme_offload_fraction, nvme_offload_dir, adam_cls): ...测试要点参数化覆盖nvme_offload_fraction ∈ {0.0, 0.5, 1.0}、nvme_offload_dir ∈ {./offload, None}以及CPUAdam/HybridAdam两种优化器构造部分参数在 CUDA、部分参数在 CPU的混合模型move_some_params_to_cuda再以随机梯度连跑 3 步与同参数的torch.optim.Adam逐参数比对torch.allclose容差atol1e-3验证 offload 前后数值行为不漂移注意该用例在文件头有pytest.mark.skip(reasonskip because of something wrong with CI)即当前 CI 中默认跳过此测试。另外本文档在 Docusaurus 体系中对应一个可运行的 doc-test原文档末尾的注释torchrun --standalone --nproc_per_node1 nvme_offload.py意味着整套示例本可按单机单卡方式直接跑通复现。关键结论与使用红线最后把使用 NVMe offload 时最容易踩的坑与前提收敛如下介质功能上兼容 HDD / SATA SSD / NVMe SSD但低带宽磁盘会让流水线退化为 I/O 瓶颈建议只在 NVMe 上开启生效范围只卸载CPU 上的优化器状态。纯 GPU 显存训练不会从中获益需配合 CPU 训练或带 CPU offload 的 ZeRO / Gemini如 GeminiPlugin 的placement_policycpu/auto/const比例语义nvme_offload_fraction是可卸载元素占全体参数元素的比例按参数遍历顺序累计分配取值范围[0.0, 1.0]目录语义nvme_offload_dirNone时使用随机临时目录进程退出自动清理自定义目录请保证路径可写且有足够磁盘空间性能权衡offload 减少内存占用但引入 I/O 时间pin_memory可加速拷贝但本身占用 RSS需按场景取舍checkpoint 限制启用 offload 时优化器的state_dict()/load_state_dict()会因潜在 OOM 抛出NotImplementedError断点续训需另行设计状态保存方案数值正确性可借助tests/test_optimizer/test_nvme.py的对照测试验证自己的接入没有改变优化轨迹。当模型规模持续增大、显存与内存都告急时把八字节每参数的 Adam 状态交给 NVMe是 ColossalAI 在 ZeRO-Offload / ZeRO-Infinity 路线上给出的一个工程化、即插即用的答案。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
