尧图网页设计 
Colossal-AI 基于 Chunk 的零冗余优化器ZeRO Gemini原理与实战指南【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI本文聚焦 Colossal-AI 中 ZeRO 数据并行与基于 Chunk 的内存管理这一对核心技术组合先厘清 ZeRO 通过切分优化器状态/梯度/参数三种模型状态来消除数据并行冗余的原理再深入讲解 Chunk连续内存块机制如何缓解小张量通信导致的带宽浪费、临时 buffer 引起的内存碎片以及 Gradient Checkpoint 场景下的重复通信问题最后给出从底层GeminiDDP到高层Booster GeminiPlugin的完整可运行示例。读完本文你将掌握 ZeRO Chunk 的存储布局、状态机、chunk 大小自动搜索机制以及用LazyInitContext初始化超大模型、使用optimizer.backward(loss)驱动混合精度训练的完整实践方法。ZeRO将三种模型状态分片而非复制经典的 Data ParallelDP训练中每个进程都要保留一份完整的模型副本在数据并行组内产生了大量的内存冗余。ZeROZero Redundancy Optimizer的核心思想是用分片partition取代复制把以下三种模型状态在数据并行进程间切分存放对应论文ZeRO: Memory Optimizations Toward Training Trillion Parameter Models提出的三阶段优化切分优化器状态Shard Optimizer States以 Adam 为例优化器状态包含 fp32 权重副本、一阶动量momentum与二阶动量variance。ZeRO 将这些状态按进程划分每个进程只负责更新自己分区内的参数。切分梯度Shard Gradient在数据并行组内完成梯度归约all-reduce之后梯度张量同样被切分使每个进程只持有与自身优化器状态分区相对应的那部分梯度。需要指出的是Colossal-AI 会把梯度转换为 fp32 格式参与参数更新见 gemini_ddp.py 的梯度处理逻辑。切分参数Shard Parameter16-bit 模型参数在数据并行进程间分片存放使用前通过通信临时聚合gather。此外Colossal-AI 在其之上叠加了Gemini——一个面向参数、梯度与优化器状态的动态异构内存空间管理器。ZeRO 负责并行切分Gemini 负责决定每个 Chunk 何时驻留 GPU、何时被驱逐到 CPU二者协作形成完整的训练方案。采用分片之后各节点内存天然负载均衡这是张量级 ZeRO 的优势但它也带来两个不可忽视的缺陷相关分析见 zero_with_chunk.md内存碎片通信过程中需要临时分配一个 buffer通信结束后再释放反复分配/释放会加剧内存碎片。带宽利用率低以单个张量为粒度做通信消息长度往往很短。一般规律是传输消息越长带宽利用率越高大量小张量的 NVLINK/PCI-E 传输无法打满带宽而且每次通信与内存移动都伴随一次 kernel launch 开销。Chunk 机制把张量粒度升级为连续内存块粒度为解决上述问题Colossal-AI 从 v0.1.8 起引入Chunk 机制本仓库的实现位于 colossalai/zero/gemini/chunk按照参数初始化顺序把一组连续的参数塞进一个 Chunk。Chunk 本质上是一段等大小的连续内存空间每个 Chunk 大小相同具体值由搜索算法决定见下文。通信、内存移动都以 Chunk 为粒度进行一次操作把 Chunk 内全部参数作为一个大 Tensor搬运/传输从而提高 PCI-E、NVLINK 与 GPU-GPU 之间网络带宽的利用率显著减少通信次数减少 kernel launch 次数从根本上避免小 buffer 反复分配释放带来的内存碎片。文档中特别剖析了一个 v0.1.8 之前的典型低效场景当一个参数在连续多个算子中被多次使用时会发生重复的通信操作。这种情况在配合Gradient Checkpoint梯度检查点时非常常见——反向传播阶段需要重算前向。以 GPT 为例Checkpoint 通常作用在每个 GPT Block 上而每个 GPT Block 包含一个 Self-Attention 层和一个 MLP 层反向传播时先依次重算 Self-Attention 层与 MLP 层的前向再依次计算 MLP 层与 Self-Attention 层的反向。同一个 Block 的参数于是被反复访问。若按张量粒度管理每次访问都要重新 gather/通信而按 Chunk 粒度管理后整个 Chunk 一次 gather 便可在多次前向/反向计算中复用这正是 Chunk 机制在分片参数 梯度检查点组合下收益格外明显的原因。在 Chunk 之上Colossal-AI 还提供了轻量级 Chunk 尺寸搜索机制根据模型的参数大小分布自动寻找内存碎片最小的 chunk 大小。搜索逻辑见 search_chunk_configuration按数据并行度dp world size把参数分组若某组参数总量小于min_chunk_size直接令其keep_gatheredTrue让这些小参数常驻 GPU不参与分片对剩余参数先用均值 3 倍标准差过滤掉异常超大参数再以隐藏维度作为搜索步长在search_range内逐一枚举候选 chunk 大小选择装箱浪费字节数最小的值最终 chunk 大小还会向上取整使其能整除各分组尺寸的最小公倍数保证各进程分片均衡。这一搜索由 ChunkManager/init_chunk_manager 在未显式提供chunk_config_dict时自动触发用户提供的hidden_dim隐藏维度作为最佳搜索间隔search_range_m控制搜索范围默认 32 MiBmin_chunk_size_m设定最小分片块默认 32 MiB。仓库中的 Chunk 体系Chunk、ChunkManager 与张量状态机为了让 Chunk 机制可感知、可调度仓库把它拆成了三层清晰的抽象源码目录colossalai/zero/gemini/chunkChunkchunk.py一段连续内存空间。每个被收纳的张量记录其在该 Chunk 内的offset/end偏移即TensorInfoChunk 持有 GPU 上的全量副本gathered或 GPU/CPU 上的分片shard并提供shard_move分片搬移、reduce块内梯度归约、access/release聚合与释放等原子操作。通过can_release/can_reduce判定该 Chunk 当前是否处于可释放、可归约状态。TensorState状态机Chunk 内每个张量都遵循一条严格的状态流转见 chunk.py#L13-L30COMPUTE正在被算子使用→HOLD_AFTER_BWD反向结束后持有→READY_FOR_REDUCE等待归约→HOLD归约完成、持有中非法跳转会直接报错从而保证内存回收的时序安全。ChunkManagermanager.pyChunk 的总调度器。它以group_type区分fp16_param与fp32_param两组 Chunkfp32 为主权重 master weights通过register_tensor把参数按初始化顺序逐块装箱、装不下时close旧块并新建 Chunk对外提供access_chunkCPU→GPU 聚合、release_chunkGPU 分片回收、move_chunk跨设备搬移、reduce_chunk块级梯度归约、init_grad_chunk为梯度分配独立 Chunk等接口并持续记账total_mem与accessed_mem供 Gemini 的内存策略做驱逐决策。在 ZeRO 分片语义下fp16_param的 Chunk 负责保存 16-bit 工作参数fp32_param的 Chunk 保存 fp32 优化器主权重开启reuse_fp16_chunkmaster_weightsTrue时后梯度的存储可以复用 fp16 参数块的空间进一步省内存。梯度归约同样是块粒度的grad_handle把整块梯度reduce_scatter/all-reduce后在进程间均分参见 gemini_ddp.py 中 grad_handle这正对应文档所述多个小 Tensor 通信合并为一次大 Tensor 通信。使用方式文档同时提供了底层GeminiDDP手动封装与高层BoosterAPI两种接入方式官方推荐优先使用后者更友好、自动处理 checkpoint、dataloader 等。本文均以仓库实际实现为准展开。LazyInitContext超大模型初始化省显存GeminiDDP是 Colossal-AI 的torch.nn.Module包装器它同时利用 ZeRO-DP 做并行切分、利用 Gemini 做内存管理。Gemini 支持LazyInitContext——在多 GPU 上初始化超大模型时先不真正分配参数内存等到GeminiDDP封装时再按分片需求落地从而省下集中初始化带来的瞬时显存峰值实现在 colossalai/lazy/lazy_init.py。文档给出一个经验法则设模型有N十亿参数、单卡显存为MGB当4N M时建议使用LazyInitContext否则它是可选的from colossalai.lazy import LazyInitContext with LazyInitContext(default_devicetorch.device(cuda)): model gpt2_medium(checkpointTrue)底层用法GeminiDDP GeminiAdamOptimizer如果不使用Booster可手动完成 ZeRO Chunk 封装from colossalai.zero import GeminiDDP, GeminiAdamOptimizer model GeminiDDP(model, hidden_dimhidden_dim, min_chunk_size_mmin_chunk_size_m) optimizer GeminiAdamOptimizer(model, lr1e-3, initial_scale2**5)其中hidden_dimDNN 的隐藏维度用于加速 chunk 尺寸搜索不确定时可省略默认取 1024。min_chunk_size_m浮点数表示最小 chunk 大小 ÷ 2^20。例如min_chunk_size_m2.5时最小 chunk 为2.5 * 2^20字节。若某组参数总量仍小于该值则所有参数会被压缩进一个小的 chunk并keep_gathered常驻 GPU见上文搜索逻辑。训练循环注意一个关键约定不要使用loss.backward()标准写法是optimizer.backward(loss)optimizer.zero_grad() outputs model(input_ids, attn_mask) loss criterion(outputs, input_ids) optimizer.backward(loss) optimizer.step()GeminiAdamOptimizer实现在 gemini_optimizer.py是面向分片参数的 ZeRO 优化器它内部维护 fp32 主权重与动态 loss scaler并在backward阶段驱动 Chunk 级的梯度聚合与归约。完整的GeminiDDP参数清单gemini_ddp.py#L56-L104还包括placement_policystatic/auto、shard_param_frac、offload_optim_frac、offload_param_frac、search_range_m、pin_memory、mixed_precisiontorch.float16/torch.bfloat16、enable_async_reduce、scatter_after_inference等。GeminiPlugin静态/自动两种放置策略GeminiPlugin把上述底层能力收敛为一个 Booster 插件完整参数语义见 gemini_plugin.py#L369-L440。其中几个高频参数如下参数默认值语义placement_policystaticstatic静态放置 /auto自动放置。auto需要先 warmup 迭代收集内存画像MemStatsNPU 后端仅支持staticshard_param_frac1.0被切分的参数占比仅 static。为1.0等价 ZeRO-3为0.0等价 ZeRO-2offload_optim_frac0.0优化器状态卸载到 CPU 的比例仅 static。shard_param_frac1.0且此值为0.0时等价旧版cuda放置offload_param_frac0.0参数卸载到 CPU 的比例仅 static。三者均取 1.0 时等价旧版cpu放置推荐先调shard_param_frac再调offload_optim_fracwarmup_non_model_data_ratio0.8warmup 阶段非模型数据内存占比预期仅 autosteady_cuda_cap_ratio0.9稳态阶段允许模型数据占用的 CUDA 容量比例仅 autoprecisionfp16支持fp16与bf16search_range_m32chunk 搜索范围 ÷ 2^20hidden_dimNone隐藏维度提供可加速 chunk 搜索缺省用 1024 作搜索间隔min_chunk_size_m32最小 chunk 大小 ÷ 2^20initial_scale2**16动态 loss scaler 的初始 scale2**5亦可见文档示例max_norm0.0clip_grad_norm的范数上限。使用 ZeRO DDP 时不要自行做梯度裁剪由 ZeRO 优化器统一处理0 表示不裁剪master_weightsTrue是否在优化器中保留 fp32 主权重pin_memoryFalseCPU 上的 Chunk 是否使用 pin-memory开启enable_async_reduce时建议设为True静态策略StaticPlacementPolicy见 placement_policy.py#L47-L84在初始化阶段一次性决定哪些 Chunk 保持聚合、哪些 Chunk 留在 GPU 分片、哪些卸载到 CPU自动策略则依据运行时内存追踪器给出的访存画像在 warmup 后把warmup_non_model_data_ratio、steady_cuda_cap_ratio转化为对每个 Chunk 的驱逐release/offload决策需要看更长远的收益时也可以结合max_prefetch做异步预取。端到端示例用 Booster 训练 GPTGemini ZeRO DP下面示例与文档及仓库 examples/language/gpt/gemini/train_gpt_demo.py 保持一致模型直接复用HuggingFace Transformers的GPT2LMHeadModel无需自行定义或修改模型结构。运行前需安装transformers为了演示简洁输入使用随机数据。1) 定义 GPT 模型与 lossimport torch import torch.nn as nn from transformers import GPT2LMHeadModel, GPT2Config class GPTLMModel(nn.Module): def __init__(self, hidden_size768, num_layers12, num_attention_heads12, max_seq_len1024, vocab_size50257, checkpointFalse): super().__init__() self.checkpoint checkpoint self.model GPT2LMHeadModel( GPT2Config(n_embdhidden_size, n_layernum_layers, n_headnum_attention_heads, n_positionsmax_seq_len, n_ctxmax_seq_len, vocab_sizevocab_size)) if checkpoint: self.model.gradient_checkpointing_enable() def forward(self, input_ids, attention_mask): return self.model(input_idsinput_ids, attention_maskattention_mask, use_cachenot self.checkpoint)[0] def gpt2_medium(checkpointFalse): # GPT2 Mediumhidden1024, 24 层, 16 个注意力头 return GPTLMModel(hidden_size1024, num_layers24, num_attention_heads16, checkpointcheckpoint) 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() return self.loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))2) 随机数据生成与 main 训练循环from colossalai.booster import Booster from colossalai.booster.plugin import GeminiPlugin from colossalai.lazy import LazyInitContext from colossalai.nn.optimizer import HybridAdam def get_data(batch_size, seq_len, vocab_size): input_ids torch.randint(0, vocab_size, (batch_size, seq_len), devicetorch.cuda.current_device()) attention_mask torch.ones_like(input_ids) return input_ids, attention_mask def main(): BATCH_SIZE 8 SEQ_LEN 1024 VOCAB_SIZE 50257 NUM_STEPS 10 colossalai.launch_from_torch() # 从 torchrun 环境读取分布式配置 criterion GPTLMLoss() model gpt2_medium(checkpointTrue) optimizer HybridAdam(model.parameters(), lr0.001) # 初始化阶段就进入 Chunk 分片管理比手动 GeminiDDP 更省内存路径 with LazyInitContext(default_devicetorch.device(cuda)): # 若改用普通初始化可省略该上下文示例见 train_gpt_demo.py model gpt2_medium(checkpointTrue) # Gemini ZeRO DP plugin GeminiPlugin(max_norm1.0, initial_scale2**5) booster Booster(pluginplugin) model, optimizer, criterion, _, _ booster.boost(model, optimizer, criterion) torch.cuda.synchronize() model.train() for n in range(NUM_STEPS): input_ids, attn_mask get_data(BATCH_SIZE, SEQ_LEN, VOCAB_SIZE) optimizer.zero_grad() outputs model(input_ids, attn_mask) loss criterion(outputs, input_ids) booster.backward(loss, optimizer) # 等价于优化器内部的 backward optimizer.step() torch.cuda.synchronize()示例关键点解读HybridAdam从 colossalai/nn/optimizer 导入是 Colossal-AI 的高性能混合 Adam 实现能感知分片参数并配合动态 loss scalerGeminiPlugin(initial_scale2**5)中的 scale 会被注入到优化器内部。booster.boost返回(model, optimizer, criterion, train_dataloader, lr_scheduler)五元组model 会被自动包装成内部使用GeminiDDP的形式。文档主例程为演示旧 API 语义保留了一次空model创建实际应用只需保留LazyInitContext内的那一次构造完整可直接运行的脚本参考 train_gpt_demo.py 与配套启动脚本 run_gemini.sh。单机单卡运行方式对应文档 doc-test 命令torchrun --standalone --nproc_per_node1 zero_with_chunk.py多卡/多机则通过 Colossal-AI 的分布式启动方式torchrun/hostfile配置进程数即可每个进程对应数据并行组中的一个分片。重要注意事项使用该方案时请留意以下三点均为文档明确强调或源码可验证不要调用loss.backward()必须经由optimizer.backward(loss)低层 API或booster.backward(loss, optimizer)Booster 路径否则无法触发 Chunk 级的梯度归约与 fp32 转换甚至可能因状态机未正确流转而报错。不要与 Booster 的 Gradient Accumulation 混用文档明确提示若使用 Gemini 模块请勿使用 gradient accumulation with booster 中描述的梯度累积功能Gemini 有独立的enable_gradient_accumulation开关二者语义不同。建议先通读前置文档本机制建立在 Booster 训练范式之上入门请先阅读 Train with booster想进一步理解 Gemini 内存管理静态/自动策略、内存画像可阅读 Meet Gemini。关于 Chunk 机制的更多底层细节内存记账、Chunk 开合、梯度块复用可继续阅读源码 chunk/manager.py 与 chunk/chunk.py并在测试目录中检索对应单测验证其行为。本文描述的设计思想可对照 ZeRO/ZeRO-Offload/ZeRO-Infinity/DeepSpeed 以及 PatrickStar首次提出基于 Chunk 的内存管理等公开论文进一步理解背景实际行为以当前仓库实现为准。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Colossal-AI 基于 Chunk 的零冗余优化器（ZeRO + Gemini）原理与实战指南
Colossal-AI 基于 Chunk 的零冗余优化器ZeRO Gemini原理与实战指南【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI本文聚焦 Colossal-AI 中 ZeRO 数据并行与基于 Chunk 的内存管理这一对核心技术组合先厘清 ZeRO 通过切分优化器状态/梯度/参数三种模型状态来消除数据并行冗余的原理再深入讲解 Chunk连续内存块机制如何缓解小张量通信导致的带宽浪费、临时 buffer 引起的内存碎片以及 Gradient Checkpoint 场景下的重复通信问题最后给出从底层GeminiDDP到高层Booster GeminiPlugin的完整可运行示例。读完本文你将掌握 ZeRO Chunk 的存储布局、状态机、chunk 大小自动搜索机制以及用LazyInitContext初始化超大模型、使用optimizer.backward(loss)驱动混合精度训练的完整实践方法。ZeRO将三种模型状态分片而非复制经典的 Data ParallelDP训练中每个进程都要保留一份完整的模型副本在数据并行组内产生了大量的内存冗余。ZeROZero Redundancy Optimizer的核心思想是用分片partition取代复制把以下三种模型状态在数据并行进程间切分存放对应论文ZeRO: Memory Optimizations Toward Training Trillion Parameter Models提出的三阶段优化切分优化器状态Shard Optimizer States以 Adam 为例优化器状态包含 fp32 权重副本、一阶动量momentum与二阶动量variance。ZeRO 将这些状态按进程划分每个进程只负责更新自己分区内的参数。切分梯度Shard Gradient在数据并行组内完成梯度归约all-reduce之后梯度张量同样被切分使每个进程只持有与自身优化器状态分区相对应的那部分梯度。需要指出的是Colossal-AI 会把梯度转换为 fp32 格式参与参数更新见 gemini_ddp.py 的梯度处理逻辑。切分参数Shard Parameter16-bit 模型参数在数据并行进程间分片存放使用前通过通信临时聚合gather。此外Colossal-AI 在其之上叠加了Gemini——一个面向参数、梯度与优化器状态的动态异构内存空间管理器。ZeRO 负责并行切分Gemini 负责决定每个 Chunk 何时驻留 GPU、何时被驱逐到 CPU二者协作形成完整的训练方案。采用分片之后各节点内存天然负载均衡这是张量级 ZeRO 的优势但它也带来两个不可忽视的缺陷相关分析见 zero_with_chunk.md内存碎片通信过程中需要临时分配一个 buffer通信结束后再释放反复分配/释放会加剧内存碎片。带宽利用率低以单个张量为粒度做通信消息长度往往很短。一般规律是传输消息越长带宽利用率越高大量小张量的 NVLINK/PCI-E 传输无法打满带宽而且每次通信与内存移动都伴随一次 kernel launch 开销。Chunk 机制把张量粒度升级为连续内存块粒度为解决上述问题Colossal-AI 从 v0.1.8 起引入Chunk 机制本仓库的实现位于 colossalai/zero/gemini/chunk按照参数初始化顺序把一组连续的参数塞进一个 Chunk。Chunk 本质上是一段等大小的连续内存空间每个 Chunk 大小相同具体值由搜索算法决定见下文。通信、内存移动都以 Chunk 为粒度进行一次操作把 Chunk 内全部参数作为一个大 Tensor搬运/传输从而提高 PCI-E、NVLINK 与 GPU-GPU 之间网络带宽的利用率显著减少通信次数减少 kernel launch 次数从根本上避免小 buffer 反复分配释放带来的内存碎片。文档中特别剖析了一个 v0.1.8 之前的典型低效场景当一个参数在连续多个算子中被多次使用时会发生重复的通信操作。这种情况在配合Gradient Checkpoint梯度检查点时非常常见——反向传播阶段需要重算前向。以 GPT 为例Checkpoint 通常作用在每个 GPT Block 上而每个 GPT Block 包含一个 Self-Attention 层和一个 MLP 层反向传播时先依次重算 Self-Attention 层与 MLP 层的前向再依次计算 MLP 层与 Self-Attention 层的反向。同一个 Block 的参数于是被反复访问。若按张量粒度管理每次访问都要重新 gather/通信而按 Chunk 粒度管理后整个 Chunk 一次 gather 便可在多次前向/反向计算中复用这正是 Chunk 机制在分片参数 梯度检查点组合下收益格外明显的原因。在 Chunk 之上Colossal-AI 还提供了轻量级 Chunk 尺寸搜索机制根据模型的参数大小分布自动寻找内存碎片最小的 chunk 大小。搜索逻辑见 search_chunk_configuration按数据并行度dp world size把参数分组若某组参数总量小于min_chunk_size直接令其keep_gatheredTrue让这些小参数常驻 GPU不参与分片对剩余参数先用均值 3 倍标准差过滤掉异常超大参数再以隐藏维度作为搜索步长在search_range内逐一枚举候选 chunk 大小选择装箱浪费字节数最小的值最终 chunk 大小还会向上取整使其能整除各分组尺寸的最小公倍数保证各进程分片均衡。这一搜索由 ChunkManager/init_chunk_manager 在未显式提供chunk_config_dict时自动触发用户提供的hidden_dim隐藏维度作为最佳搜索间隔search_range_m控制搜索范围默认 32 MiBmin_chunk_size_m设定最小分片块默认 32 MiB。仓库中的 Chunk 体系Chunk、ChunkManager 与张量状态机为了让 Chunk 机制可感知、可调度仓库把它拆成了三层清晰的抽象源码目录colossalai/zero/gemini/chunkChunkchunk.py一段连续内存空间。每个被收纳的张量记录其在该 Chunk 内的offset/end偏移即TensorInfoChunk 持有 GPU 上的全量副本gathered或 GPU/CPU 上的分片shard并提供shard_move分片搬移、reduce块内梯度归约、access/release聚合与释放等原子操作。通过can_release/can_reduce判定该 Chunk 当前是否处于可释放、可归约状态。TensorState状态机Chunk 内每个张量都遵循一条严格的状态流转见 chunk.py#L13-L30COMPUTE正在被算子使用→HOLD_AFTER_BWD反向结束后持有→READY_FOR_REDUCE等待归约→HOLD归约完成、持有中非法跳转会直接报错从而保证内存回收的时序安全。ChunkManagermanager.pyChunk 的总调度器。它以group_type区分fp16_param与fp32_param两组 Chunkfp32 为主权重 master weights通过register_tensor把参数按初始化顺序逐块装箱、装不下时close旧块并新建 Chunk对外提供access_chunkCPU→GPU 聚合、release_chunkGPU 分片回收、move_chunk跨设备搬移、reduce_chunk块级梯度归约、init_grad_chunk为梯度分配独立 Chunk等接口并持续记账total_mem与accessed_mem供 Gemini 的内存策略做驱逐决策。在 ZeRO 分片语义下fp16_param的 Chunk 负责保存 16-bit 工作参数fp32_param的 Chunk 保存 fp32 优化器主权重开启reuse_fp16_chunkmaster_weightsTrue时后梯度的存储可以复用 fp16 参数块的空间进一步省内存。梯度归约同样是块粒度的grad_handle把整块梯度reduce_scatter/all-reduce后在进程间均分参见 gemini_ddp.py 中 grad_handle这正对应文档所述多个小 Tensor 通信合并为一次大 Tensor 通信。使用方式文档同时提供了底层GeminiDDP手动封装与高层BoosterAPI两种接入方式官方推荐优先使用后者更友好、自动处理 checkpoint、dataloader 等。本文均以仓库实际实现为准展开。LazyInitContext超大模型初始化省显存GeminiDDP是 Colossal-AI 的torch.nn.Module包装器它同时利用 ZeRO-DP 做并行切分、利用 Gemini 做内存管理。Gemini 支持LazyInitContext——在多 GPU 上初始化超大模型时先不真正分配参数内存等到GeminiDDP封装时再按分片需求落地从而省下集中初始化带来的瞬时显存峰值实现在 colossalai/lazy/lazy_init.py。文档给出一个经验法则设模型有N十亿参数、单卡显存为MGB当4N M时建议使用LazyInitContext否则它是可选的from colossalai.lazy import LazyInitContext with LazyInitContext(default_devicetorch.device(cuda)): model gpt2_medium(checkpointTrue)底层用法GeminiDDP GeminiAdamOptimizer如果不使用Booster可手动完成 ZeRO Chunk 封装from colossalai.zero import GeminiDDP, GeminiAdamOptimizer model GeminiDDP(model, hidden_dimhidden_dim, min_chunk_size_mmin_chunk_size_m) optimizer GeminiAdamOptimizer(model, lr1e-3, initial_scale2**5)其中hidden_dimDNN 的隐藏维度用于加速 chunk 尺寸搜索不确定时可省略默认取 1024。min_chunk_size_m浮点数表示最小 chunk 大小 ÷ 2^20。例如min_chunk_size_m2.5时最小 chunk 为2.5 * 2^20字节。若某组参数总量仍小于该值则所有参数会被压缩进一个小的 chunk并keep_gathered常驻 GPU见上文搜索逻辑。训练循环注意一个关键约定不要使用loss.backward()标准写法是optimizer.backward(loss)optimizer.zero_grad() outputs model(input_ids, attn_mask) loss criterion(outputs, input_ids) optimizer.backward(loss) optimizer.step()GeminiAdamOptimizer实现在 gemini_optimizer.py是面向分片参数的 ZeRO 优化器它内部维护 fp32 主权重与动态 loss scaler并在backward阶段驱动 Chunk 级的梯度聚合与归约。完整的GeminiDDP参数清单gemini_ddp.py#L56-L104还包括placement_policystatic/auto、shard_param_frac、offload_optim_frac、offload_param_frac、search_range_m、pin_memory、mixed_precisiontorch.float16/torch.bfloat16、enable_async_reduce、scatter_after_inference等。GeminiPlugin静态/自动两种放置策略GeminiPlugin把上述底层能力收敛为一个 Booster 插件完整参数语义见 gemini_plugin.py#L369-L440。其中几个高频参数如下参数默认值语义placement_policystaticstatic静态放置 /auto自动放置。auto需要先 warmup 迭代收集内存画像MemStatsNPU 后端仅支持staticshard_param_frac1.0被切分的参数占比仅 static。为1.0等价 ZeRO-3为0.0等价 ZeRO-2offload_optim_frac0.0优化器状态卸载到 CPU 的比例仅 static。shard_param_frac1.0且此值为0.0时等价旧版cuda放置offload_param_frac0.0参数卸载到 CPU 的比例仅 static。三者均取 1.0 时等价旧版cpu放置推荐先调shard_param_frac再调offload_optim_fracwarmup_non_model_data_ratio0.8warmup 阶段非模型数据内存占比预期仅 autosteady_cuda_cap_ratio0.9稳态阶段允许模型数据占用的 CUDA 容量比例仅 autoprecisionfp16支持fp16与bf16search_range_m32chunk 搜索范围 ÷ 2^20hidden_dimNone隐藏维度提供可加速 chunk 搜索缺省用 1024 作搜索间隔min_chunk_size_m32最小 chunk 大小 ÷ 2^20initial_scale2**16动态 loss scaler 的初始 scale2**5亦可见文档示例max_norm0.0clip_grad_norm的范数上限。使用 ZeRO DDP 时不要自行做梯度裁剪由 ZeRO 优化器统一处理0 表示不裁剪master_weightsTrue是否在优化器中保留 fp32 主权重pin_memoryFalseCPU 上的 Chunk 是否使用 pin-memory开启enable_async_reduce时建议设为True静态策略StaticPlacementPolicy见 placement_policy.py#L47-L84在初始化阶段一次性决定哪些 Chunk 保持聚合、哪些 Chunk 留在 GPU 分片、哪些卸载到 CPU自动策略则依据运行时内存追踪器给出的访存画像在 warmup 后把warmup_non_model_data_ratio、steady_cuda_cap_ratio转化为对每个 Chunk 的驱逐release/offload决策需要看更长远的收益时也可以结合max_prefetch做异步预取。端到端示例用 Booster 训练 GPTGemini ZeRO DP下面示例与文档及仓库 examples/language/gpt/gemini/train_gpt_demo.py 保持一致模型直接复用HuggingFace Transformers的GPT2LMHeadModel无需自行定义或修改模型结构。运行前需安装transformers为了演示简洁输入使用随机数据。1) 定义 GPT 模型与 lossimport torch import torch.nn as nn from transformers import GPT2LMHeadModel, GPT2Config class GPTLMModel(nn.Module): def __init__(self, hidden_size768, num_layers12, num_attention_heads12, max_seq_len1024, vocab_size50257, checkpointFalse): super().__init__() self.checkpoint checkpoint self.model GPT2LMHeadModel( GPT2Config(n_embdhidden_size, n_layernum_layers, n_headnum_attention_heads, n_positionsmax_seq_len, n_ctxmax_seq_len, vocab_sizevocab_size)) if checkpoint: self.model.gradient_checkpointing_enable() def forward(self, input_ids, attention_mask): return self.model(input_idsinput_ids, attention_maskattention_mask, use_cachenot self.checkpoint)[0] def gpt2_medium(checkpointFalse): # GPT2 Mediumhidden1024, 24 层, 16 个注意力头 return GPTLMModel(hidden_size1024, num_layers24, num_attention_heads16, checkpointcheckpoint) 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() return self.loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))2) 随机数据生成与 main 训练循环from colossalai.booster import Booster from colossalai.booster.plugin import GeminiPlugin from colossalai.lazy import LazyInitContext from colossalai.nn.optimizer import HybridAdam def get_data(batch_size, seq_len, vocab_size): input_ids torch.randint(0, vocab_size, (batch_size, seq_len), devicetorch.cuda.current_device()) attention_mask torch.ones_like(input_ids) return input_ids, attention_mask def main(): BATCH_SIZE 8 SEQ_LEN 1024 VOCAB_SIZE 50257 NUM_STEPS 10 colossalai.launch_from_torch() # 从 torchrun 环境读取分布式配置 criterion GPTLMLoss() model gpt2_medium(checkpointTrue) optimizer HybridAdam(model.parameters(), lr0.001) # 初始化阶段就进入 Chunk 分片管理比手动 GeminiDDP 更省内存路径 with LazyInitContext(default_devicetorch.device(cuda)): # 若改用普通初始化可省略该上下文示例见 train_gpt_demo.py model gpt2_medium(checkpointTrue) # Gemini ZeRO DP plugin GeminiPlugin(max_norm1.0, initial_scale2**5) booster Booster(pluginplugin) model, optimizer, criterion, _, _ booster.boost(model, optimizer, criterion) torch.cuda.synchronize() model.train() for n in range(NUM_STEPS): input_ids, attn_mask get_data(BATCH_SIZE, SEQ_LEN, VOCAB_SIZE) optimizer.zero_grad() outputs model(input_ids, attn_mask) loss criterion(outputs, input_ids) booster.backward(loss, optimizer) # 等价于优化器内部的 backward optimizer.step() torch.cuda.synchronize()示例关键点解读HybridAdam从 colossalai/nn/optimizer 导入是 Colossal-AI 的高性能混合 Adam 实现能感知分片参数并配合动态 loss scalerGeminiPlugin(initial_scale2**5)中的 scale 会被注入到优化器内部。booster.boost返回(model, optimizer, criterion, train_dataloader, lr_scheduler)五元组model 会被自动包装成内部使用GeminiDDP的形式。文档主例程为演示旧 API 语义保留了一次空model创建实际应用只需保留LazyInitContext内的那一次构造完整可直接运行的脚本参考 train_gpt_demo.py 与配套启动脚本 run_gemini.sh。单机单卡运行方式对应文档 doc-test 命令torchrun --standalone --nproc_per_node1 zero_with_chunk.py多卡/多机则通过 Colossal-AI 的分布式启动方式torchrun/hostfile配置进程数即可每个进程对应数据并行组中的一个分片。重要注意事项使用该方案时请留意以下三点均为文档明确强调或源码可验证不要调用loss.backward()必须经由optimizer.backward(loss)低层 API或booster.backward(loss, optimizer)Booster 路径否则无法触发 Chunk 级的梯度归约与 fp32 转换甚至可能因状态机未正确流转而报错。不要与 Booster 的 Gradient Accumulation 混用文档明确提示若使用 Gemini 模块请勿使用 gradient accumulation with booster 中描述的梯度累积功能Gemini 有独立的enable_gradient_accumulation开关二者语义不同。建议先通读前置文档本机制建立在 Booster 训练范式之上入门请先阅读 Train with booster想进一步理解 Gemini 内存管理静态/自动策略、内存画像可阅读 Meet Gemini。关于 Chunk 机制的更多底层细节内存记账、Chunk 开合、梯度块复用可继续阅读源码 chunk/manager.py 与 chunk/chunk.py并在测试目录中检索对应单测验证其行为。本文描述的设计思想可对照 ZeRO/ZeRO-Offload/ZeRO-Infinity/DeepSpeed 以及 PatrickStar首次提出基于 Chunk 的内存管理等公开论文进一步理解背景实际行为以当前仓库实现为准。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考