尧图网站建设 
基于 Colossal-AI Auto-Offload 的 GPT-2 参数自动卸载训练实践指南【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI导读本文基于 examples/language/gpt/experiments/auto_offload 目录下的 README 及其配套训练脚本讲解如何利用 Colossal-AI 的 Auto-Offload 机制在显存预算受限例如 16 GB的条件下训练 GPT-2 规模及以上的模型。你将学会环境安装与一键启动方法理解四个关键运行参数模型类型、批大小、显存预算、卸载求解器的作用并掌握colossalai.auto_parallel.offload从图追踪、区域划分、策略求解到运行时图重写的完整实现链路从而具备将该能力迁移到自定义模型上的能力。一、背景为什么要做参数自动卸载训练大模型时模型参数、优化器状态与激活值三者共同决定了显存峰值。在单卡显存不足而又无法横向扩容时常见的缓解手段包括 ZeRO、Gemini 以及显存换入换出offload。其中参数卸载的核心矛盾是GPU 显存与 PCIe/主机内存之间搬运参数会引入额外通信开销卸载得越多单次迭代耗时越长。Colossal-AI 的 Auto-Offload 思路是不再手工规定每 N 层卸载一次而是把计算图切分为若干Region区域借助离线训练模拟器在内存节省与通信/计算开销之间做量化权衡自动搜索出一个满足给定显存预算的卸载与预取方案。GPT 实验目录正是这一能力的端到端演示。二、运行环境与依赖安装原文档要求的环境为 PyTorch 1.12 CUDA 11.3 Colossal-AI v0.2.0 transformers具体安装方式如下。1. 安装 PyTorch#conda conda install pytorch1.12.0 torchvision0.13.0 torchaudio0.12.0 cudatoolkit11.3 -c pytorch #pip pip install torch1.12.0cu113 torchvision0.13.0cu113 torchaudio0.12.0 --extra-index-url https://download.pytorch.org/whl/cu1132. 安装 Colossal-AI v0.2.0pip install colossalai0.2.0torch1.12cu11.3 -f https://release.colossalai.org版本匹配非常关键示例代码直接导入colossalai.auto_parallel.offload与colossalai.nn.optimizer等模块必须保证 Colossal-AI 与你的 PyTorch/CUDA 版本一一对应否则会出现二进制不兼容。若使用本仓库主线版本请同步参考 requirements/requirements.txt 与 setup.py 中声明的依赖约束。3. 安装 transformers 与其它依赖pip install transformers实验目录自带的 requirements.txt 还要求colossalai 0.1.12与torch 1.8.1可作为最低版本约束的参考。提示Auto-Offload 的求解器需要读取 GPU 实时功耗与 PCIe 带宽见后文solver.py中的 profiling因此建议安装pynvml。在 train_gpt_offload.py 与测试代码中都用pytest.mark.skipif(NOT_NVML, ...)在缺少 pynvml 时跳过相关流程。三、数据集使用随机生成数据为了聚焦验证卸载机制本身示例没有引入真实语料。数据生成逻辑位于 model_zoo.pydef gpt2_data_gen(devicecuda): input_ids torch.randint(0, vocab_size, (batch_size, seq_len), devicedevice) attention_mask torch.ones_like(input_ids, devicedevice) kwargs dict(input_idsinput_ids, attention_maskattention_mask) return kwargs其中vocab_size1024、seq_len8标签label在训练脚本中通过torch.randint(0, 128, (64, 8))随机生成。这意味着你可以完全不依赖下载与预处理开箱即用地验证卸载能否跑通 显存是否受控之后再替换为真实数据集即可。四、一键启动训练原文档给出的核心命令只有一行#Run the auto offload on GPT with default setting and a dummy dataset. bash run.shrun.sh 的内容如下export BATCH_SIZE${BATCH_SIZE:-64} export MODEL_TYPE${MODEL_TYPE:-gpt2_medium} export MEMORY_BUDGET${MEMORY_BUDGET:-16} export SOLVER_TYPE${SOLVER_TYPE:-asyn} mkdir -p offload_logs python train_gpt_offload.py --model_type${MODEL_TYPE} --memory_budget${MEMORY_BUDGET} --solver_type${SOLVER_TYPE} --batch_size${BATCH_SIZE} 21 | tee ./offload_logs/${MODEL_TYPE}_bs_${BATCH_SIZE}_st_${SOLVER_TYPE}.log它做了两件事一是通过环境变量提供四个默认参数二是把训练输出同时打印到终端并落盘到offload_logs/模型_bs_批大小_st_求解器.log方便事后对比实验。参数速查表命令行参数对应环境变量默认值含义与可选范围--model_typeMODEL_TYPEgpt2_medium模型规模见model_zoo.py可选gpt2_medium/gpt2_xl/gpt2_10b/gpt2_14b/gpt2_20b/gpt2_24b--batch_sizeBATCH_SIZE64每个 step 的批大小--memory_budgetMEMORY_BUDGET16显存预算单位为GB脚本内会乘上1024**3换算成字节--solver_typeSOLVER_TYPEasyn卸载策略求解器可选syn/asyn见第五节其中--model_type在 model_zoo.py 中映射到不同的 GPT-2 结构超参例如gpt2_mediumhidden1024、24 层、16 头gpt2_xlhidden1600、48 层、32 头gpt2_10bhidden4096、50 层、16 头gpt2_24bhidden8192、30 层、16 头。注意模型越大、层数越深求解器搜索空间与运行时长都会显著上升建议从小模型gpt2_medium起步。五、训练脚本逐段拆解train_gpt_offload.py 是全流程的入口其调用链如下。1. 参数解析与模型构建parser.add_argument(--model_type, typestr, defaultgpt2_medium) parser.add_argument(--batch_size, typeint, default64) parser.add_argument(--solver_type, typestr, defaultasyn) parser.add_argument(--memory_budget, typefloat, default16)随后通过get_gpt2_components(model_type, batch_size)拿到模型工厂和数据工厂并用parameter_size(model)打印模型初始参数量按 FP32 统计后除以 2因为后续会转成 FP16。2. 核心 APImemory_optimize整个演示的关键只有一行调用model memory_optimize(model, data_args, memory_budget, solver_type)在 mem_optimize.py 中memory_optimize(model, inps, memory_budget-1.0, solver_nameasyn)依次完成六件事精度与设备转换model.cpu().half()先把模型放到 CPU 并转为 FP16输入数据也相应做半精度包装显存统计因此减半符号追踪用ColoTracer对模型做 meta 级 trace得到torch.fx.Graph元信息传播MetaInfoProp在图上传播张量形状与内存占用区域划分与求解RegionManager(graph, solver_name, memory_budget)构建 Region 列表并打印act_peak_mem / max_param_mem / total_param_mem运行时图重写按求解结果调用runtime_syn_offload_apply_pass或runtime_asyn_offload_apply_pass把上传/卸载动作以算子形式插入计算图封装将重写后的GraphModule包进BaseOffloadModule返回得到带卸载能力的模型对象。3. 优化器与梯度流hybrid_optimizer HybridAdam(model.model.parameters(), lr1e-3) optim AMPOptimizer(hybrid_optimizer, model)HybridAdam位于 hybrid_adam.py是 Colossal-AI 面向混合并行/异构内存场景的 Adam 实现参数可以直接驻留 CPUAMPOptimizer见 amp_optimizer.py是自动混合精度 卸载的桥接层提供backward(loss)、zero_grad()、step()、clip_grad_norm()等接口内部完成 loss 缩放默认initial_scale2**16、梯度溢出检测、并把被卸载 region 的梯度回拷到对应 CPU 参数切片上。4. 训练循环与指标统计脚本固定跑 10 个 step为剔除启动抖动exec_time取最快 5 个 step 的平均exec_time sum(sorted(time_list)[:5]) / 5 runtime_peak_mem_alc torch.cuda.max_memory_allocated() / 1024**2 runtime_peak_mem_res torch.cuda.max_memory_reserved() / 1024**2最终输出形如| exec_time0.982 s | param_size347.014 MB | runtime_peak_mem_alc10882.117 MB| runtime_peak_mem_res11128.012 MB|其中param_size是 FP16 参数量runtime_peak_mem_alc与runtime_peak_mem_res分别对应torch.cuda.max_memory_allocated与max_memory_reserved是验证显存是否被压在 MEMORY_BUDGET 内的直接证据。六、从源码看 Auto-Offload 的底层原理原 README 只给出运行方式要理解自动二字为何成立需要进入colossalai/auto_parallel/offload/目录对应 Colossal-Auto 在旧版 API 中的 offload 能力。1. Region卸载的基本单元直接对算子做卸载代价过高因此代码先把线性化后的计算图打包成 Region见 region_manager.py 与 region.py。Region 具有以下关键能力move_param_to_cuda()/move_grad_to_cpu()/free_cuda_data()控制参数与梯度在 GPU/CPU 间迁移can_release()、has_inf_or_nan()供运行期判断该 region 能否释放显存split(cut_node_idx, cut_param_idx)RegionManager 的后处理会合并小 region、切分超大连续段进一步降低内存碎片与调度粒度_merge_small_regions、_search_block_size等即服务于该目的。2. 训练模拟器把调度算出来直接试跑每种卸载方案代价不可接受因此 Colossal-AI 用离散事件模拟来预估每种策略下的峰值显存与迭代耗时。training_simulator.py 提供SynTrainingSimulator同步卸载与AsynTrainingSimulator异步预取它们都依赖两类硬件画像计算能力_extract_computing_power()对应 util.py 中的NvDevicePower表A100 FP16≈78 TFLOPS 等传输带宽_profile_bandwidth()对h2d/d2h/p2p链路的实际带宽做 benchmark见benchmark_func。3. 求解器syn / asyn / 修复策略solver.py 通过SolverFactory注册了两个求解器solvers: Dict[str, Type[Solver]] {syn: SynGreedySolver, asyn: AsynGreedySolver}syn同步卸载SynGreedySolver采用贪心策略反复挑选单位额外通信代价换来内存节省最大的 region 进行同步卸载_call_solver_l2l还保留了一种逐层全量卸载的 l2l 变体思路。当预算用error_factor0.95折减后仍无法满足时会抛出带最少所需显存信息的异常asyn异步卸载 前向预取AsynGreedySolver除选取卸载对象外还要为被卸载 region 决定在哪个后续 region 处预取回 GPUsearch_window_size3的滑动窗口内枚举 host region见 solver.py目标是让 H2D/D2H 传输与前向/反向计算重叠。若预取导致峰值不降反升_repair_strategy()会把部分异步预取回退为同步上传。在 util.py 的GlobalRuntimeInfo中可以看到异步化的物理支撑系统维护独立的h2d_stream与d2h_stream并用 CUDA Eventfwd_prefetch_event_map/bwd_prefetch_event_map同步计算流与搬运流的顺序。4. 运行时图重写与执行求解完成只是得到哪些 region 卸载、何时预取的决策表最终要落到可执行的图决策下发RegionManager._post_process()结合模拟器结果回填每个 region 的need_offload/is_syn/fwd_prefetch_region/bwd_prefetch_region标记图插入runtime.py 中的runtime_syn_offload_apply_pass与runtime_asyn_offload_apply_pass把convert_fwd_upload_bwd_offload_to_action/convert_fwd_prefetch_bwd_offload_to_action等动作节点插入torch.fx.Graph前反向驱动封装模型 base_offload_module.py 在forward/backward的_pre_forward/_post_backward钩子中注册梯度 hook保证计算前参数在 GPU、计算后梯度回写 CPU。七、如何验证与横向对比仓库在 tests/test_auto_parallel/test_offload 中给出了两层测试test_solver.py对gpt2_与bert_两种模型、syn与asyn两种求解器分别跑求解流程并断言solver.best_ts.peak_mem memory_budget即求解器给出的方案峰值显存必须低于预算test_perf.py在同一份代码里并排跑 Geminiplacement_policycpu与 asyn offload 两种方案打印各自的exec_time与runtime_peak_mem_*方便直接对照吞吐与显存表现注意该用例默认被pytest.mark.skip跳过需要显式开启。结合 test_perf.py 的做法你可以在本演示中自行做横向实验# 同一显存预算下对比同步与异步卸载 bash run.sh # 默认 asyngpt2_medium16GB SOLVER_TYPEsyn bash run.sh MODEL_TYPEgpt2_xl MEMORY_BUDGET24 bash run.sh每组运行都会在offload_logs/下留下独立日志可直接比对solver_time、exec_time、runtime_peak_mem_alc三项指标。需要强调的是是否出现卸载、卸载多少个 region完全由求解器结合硬件画像与显存预算动态决定——当MEMORY_BUDGET设得足够大时求解结果可能趋向于几乎不卸载这也是该方案被称为自动的原因。八、常见问题与使用注意求解时间长求解器内部会多次运行训练模拟器做贪心搜索solver.py模型越大、候选 region 越多solver_time越长。它是离线成本训练前一次性支付报 cant find the offload strategy ...说明即使把所有可卸载 region 全部卸出峰值显存仍高于预算。需要调大MEMORY_BUDGET或选用更小的MODEL_TYPE需要 pynvml带宽/算力画像依赖 pynvml 读取真实设备信息solver.py顶部try: from pynvml import *缺失时相关测试会被NOT_NVML标记跳过半精度前提memory_optimize内部强制model.cpu().half()输入中的浮点张量也需转换为 FP16脚本用wrap_fntree_map完成模型请勿残留 FP32 必须算子或敏感 BatchNorm 等结构适配新模型只要模型结构可被ColoTracer追踪就能套用构建模型 → 生成 dummy 输入 →memory_optimize→AMPOptimizer这一通用模板无需手工规划卸载点。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 基于 Colossal-AI Auto-Offload 的 GPT-2 参数自动卸载训练实践指南
基于 Colossal-AI Auto-Offload 的 GPT-2 参数自动卸载训练实践指南【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI导读本文基于 examples/language/gpt/experiments/auto_offload 目录下的 README 及其配套训练脚本讲解如何利用 Colossal-AI 的 Auto-Offload 机制在显存预算受限例如 16 GB的条件下训练 GPT-2 规模及以上的模型。你将学会环境安装与一键启动方法理解四个关键运行参数模型类型、批大小、显存预算、卸载求解器的作用并掌握colossalai.auto_parallel.offload从图追踪、区域划分、策略求解到运行时图重写的完整实现链路从而具备将该能力迁移到自定义模型上的能力。一、背景为什么要做参数自动卸载训练大模型时模型参数、优化器状态与激活值三者共同决定了显存峰值。在单卡显存不足而又无法横向扩容时常见的缓解手段包括 ZeRO、Gemini 以及显存换入换出offload。其中参数卸载的核心矛盾是GPU 显存与 PCIe/主机内存之间搬运参数会引入额外通信开销卸载得越多单次迭代耗时越长。Colossal-AI 的 Auto-Offload 思路是不再手工规定每 N 层卸载一次而是把计算图切分为若干Region区域借助离线训练模拟器在内存节省与通信/计算开销之间做量化权衡自动搜索出一个满足给定显存预算的卸载与预取方案。GPT 实验目录正是这一能力的端到端演示。二、运行环境与依赖安装原文档要求的环境为 PyTorch 1.12 CUDA 11.3 Colossal-AI v0.2.0 transformers具体安装方式如下。1. 安装 PyTorch#conda conda install pytorch1.12.0 torchvision0.13.0 torchaudio0.12.0 cudatoolkit11.3 -c pytorch #pip pip install torch1.12.0cu113 torchvision0.13.0cu113 torchaudio0.12.0 --extra-index-url https://download.pytorch.org/whl/cu1132. 安装 Colossal-AI v0.2.0pip install colossalai0.2.0torch1.12cu11.3 -f https://release.colossalai.org版本匹配非常关键示例代码直接导入colossalai.auto_parallel.offload与colossalai.nn.optimizer等模块必须保证 Colossal-AI 与你的 PyTorch/CUDA 版本一一对应否则会出现二进制不兼容。若使用本仓库主线版本请同步参考 requirements/requirements.txt 与 setup.py 中声明的依赖约束。3. 安装 transformers 与其它依赖pip install transformers实验目录自带的 requirements.txt 还要求colossalai 0.1.12与torch 1.8.1可作为最低版本约束的参考。提示Auto-Offload 的求解器需要读取 GPU 实时功耗与 PCIe 带宽见后文solver.py中的 profiling因此建议安装pynvml。在 train_gpt_offload.py 与测试代码中都用pytest.mark.skipif(NOT_NVML, ...)在缺少 pynvml 时跳过相关流程。三、数据集使用随机生成数据为了聚焦验证卸载机制本身示例没有引入真实语料。数据生成逻辑位于 model_zoo.pydef gpt2_data_gen(devicecuda): input_ids torch.randint(0, vocab_size, (batch_size, seq_len), devicedevice) attention_mask torch.ones_like(input_ids, devicedevice) kwargs dict(input_idsinput_ids, attention_maskattention_mask) return kwargs其中vocab_size1024、seq_len8标签label在训练脚本中通过torch.randint(0, 128, (64, 8))随机生成。这意味着你可以完全不依赖下载与预处理开箱即用地验证卸载能否跑通 显存是否受控之后再替换为真实数据集即可。四、一键启动训练原文档给出的核心命令只有一行#Run the auto offload on GPT with default setting and a dummy dataset. bash run.shrun.sh 的内容如下export BATCH_SIZE${BATCH_SIZE:-64} export MODEL_TYPE${MODEL_TYPE:-gpt2_medium} export MEMORY_BUDGET${MEMORY_BUDGET:-16} export SOLVER_TYPE${SOLVER_TYPE:-asyn} mkdir -p offload_logs python train_gpt_offload.py --model_type${MODEL_TYPE} --memory_budget${MEMORY_BUDGET} --solver_type${SOLVER_TYPE} --batch_size${BATCH_SIZE} 21 | tee ./offload_logs/${MODEL_TYPE}_bs_${BATCH_SIZE}_st_${SOLVER_TYPE}.log它做了两件事一是通过环境变量提供四个默认参数二是把训练输出同时打印到终端并落盘到offload_logs/模型_bs_批大小_st_求解器.log方便事后对比实验。参数速查表命令行参数对应环境变量默认值含义与可选范围--model_typeMODEL_TYPEgpt2_medium模型规模见model_zoo.py可选gpt2_medium/gpt2_xl/gpt2_10b/gpt2_14b/gpt2_20b/gpt2_24b--batch_sizeBATCH_SIZE64每个 step 的批大小--memory_budgetMEMORY_BUDGET16显存预算单位为GB脚本内会乘上1024**3换算成字节--solver_typeSOLVER_TYPEasyn卸载策略求解器可选syn/asyn见第五节其中--model_type在 model_zoo.py 中映射到不同的 GPT-2 结构超参例如gpt2_mediumhidden1024、24 层、16 头gpt2_xlhidden1600、48 层、32 头gpt2_10bhidden4096、50 层、16 头gpt2_24bhidden8192、30 层、16 头。注意模型越大、层数越深求解器搜索空间与运行时长都会显著上升建议从小模型gpt2_medium起步。五、训练脚本逐段拆解train_gpt_offload.py 是全流程的入口其调用链如下。1. 参数解析与模型构建parser.add_argument(--model_type, typestr, defaultgpt2_medium) parser.add_argument(--batch_size, typeint, default64) parser.add_argument(--solver_type, typestr, defaultasyn) parser.add_argument(--memory_budget, typefloat, default16)随后通过get_gpt2_components(model_type, batch_size)拿到模型工厂和数据工厂并用parameter_size(model)打印模型初始参数量按 FP32 统计后除以 2因为后续会转成 FP16。2. 核心 APImemory_optimize整个演示的关键只有一行调用model memory_optimize(model, data_args, memory_budget, solver_type)在 mem_optimize.py 中memory_optimize(model, inps, memory_budget-1.0, solver_nameasyn)依次完成六件事精度与设备转换model.cpu().half()先把模型放到 CPU 并转为 FP16输入数据也相应做半精度包装显存统计因此减半符号追踪用ColoTracer对模型做 meta 级 trace得到torch.fx.Graph元信息传播MetaInfoProp在图上传播张量形状与内存占用区域划分与求解RegionManager(graph, solver_name, memory_budget)构建 Region 列表并打印act_peak_mem / max_param_mem / total_param_mem运行时图重写按求解结果调用runtime_syn_offload_apply_pass或runtime_asyn_offload_apply_pass把上传/卸载动作以算子形式插入计算图封装将重写后的GraphModule包进BaseOffloadModule返回得到带卸载能力的模型对象。3. 优化器与梯度流hybrid_optimizer HybridAdam(model.model.parameters(), lr1e-3) optim AMPOptimizer(hybrid_optimizer, model)HybridAdam位于 hybrid_adam.py是 Colossal-AI 面向混合并行/异构内存场景的 Adam 实现参数可以直接驻留 CPUAMPOptimizer见 amp_optimizer.py是自动混合精度 卸载的桥接层提供backward(loss)、zero_grad()、step()、clip_grad_norm()等接口内部完成 loss 缩放默认initial_scale2**16、梯度溢出检测、并把被卸载 region 的梯度回拷到对应 CPU 参数切片上。4. 训练循环与指标统计脚本固定跑 10 个 step为剔除启动抖动exec_time取最快 5 个 step 的平均exec_time sum(sorted(time_list)[:5]) / 5 runtime_peak_mem_alc torch.cuda.max_memory_allocated() / 1024**2 runtime_peak_mem_res torch.cuda.max_memory_reserved() / 1024**2最终输出形如| exec_time0.982 s | param_size347.014 MB | runtime_peak_mem_alc10882.117 MB| runtime_peak_mem_res11128.012 MB|其中param_size是 FP16 参数量runtime_peak_mem_alc与runtime_peak_mem_res分别对应torch.cuda.max_memory_allocated与max_memory_reserved是验证显存是否被压在 MEMORY_BUDGET 内的直接证据。六、从源码看 Auto-Offload 的底层原理原 README 只给出运行方式要理解自动二字为何成立需要进入colossalai/auto_parallel/offload/目录对应 Colossal-Auto 在旧版 API 中的 offload 能力。1. Region卸载的基本单元直接对算子做卸载代价过高因此代码先把线性化后的计算图打包成 Region见 region_manager.py 与 region.py。Region 具有以下关键能力move_param_to_cuda()/move_grad_to_cpu()/free_cuda_data()控制参数与梯度在 GPU/CPU 间迁移can_release()、has_inf_or_nan()供运行期判断该 region 能否释放显存split(cut_node_idx, cut_param_idx)RegionManager 的后处理会合并小 region、切分超大连续段进一步降低内存碎片与调度粒度_merge_small_regions、_search_block_size等即服务于该目的。2. 训练模拟器把调度算出来直接试跑每种卸载方案代价不可接受因此 Colossal-AI 用离散事件模拟来预估每种策略下的峰值显存与迭代耗时。training_simulator.py 提供SynTrainingSimulator同步卸载与AsynTrainingSimulator异步预取它们都依赖两类硬件画像计算能力_extract_computing_power()对应 util.py 中的NvDevicePower表A100 FP16≈78 TFLOPS 等传输带宽_profile_bandwidth()对h2d/d2h/p2p链路的实际带宽做 benchmark见benchmark_func。3. 求解器syn / asyn / 修复策略solver.py 通过SolverFactory注册了两个求解器solvers: Dict[str, Type[Solver]] {syn: SynGreedySolver, asyn: AsynGreedySolver}syn同步卸载SynGreedySolver采用贪心策略反复挑选单位额外通信代价换来内存节省最大的 region 进行同步卸载_call_solver_l2l还保留了一种逐层全量卸载的 l2l 变体思路。当预算用error_factor0.95折减后仍无法满足时会抛出带最少所需显存信息的异常asyn异步卸载 前向预取AsynGreedySolver除选取卸载对象外还要为被卸载 region 决定在哪个后续 region 处预取回 GPUsearch_window_size3的滑动窗口内枚举 host region见 solver.py目标是让 H2D/D2H 传输与前向/反向计算重叠。若预取导致峰值不降反升_repair_strategy()会把部分异步预取回退为同步上传。在 util.py 的GlobalRuntimeInfo中可以看到异步化的物理支撑系统维护独立的h2d_stream与d2h_stream并用 CUDA Eventfwd_prefetch_event_map/bwd_prefetch_event_map同步计算流与搬运流的顺序。4. 运行时图重写与执行求解完成只是得到哪些 region 卸载、何时预取的决策表最终要落到可执行的图决策下发RegionManager._post_process()结合模拟器结果回填每个 region 的need_offload/is_syn/fwd_prefetch_region/bwd_prefetch_region标记图插入runtime.py 中的runtime_syn_offload_apply_pass与runtime_asyn_offload_apply_pass把convert_fwd_upload_bwd_offload_to_action/convert_fwd_prefetch_bwd_offload_to_action等动作节点插入torch.fx.Graph前反向驱动封装模型 base_offload_module.py 在forward/backward的_pre_forward/_post_backward钩子中注册梯度 hook保证计算前参数在 GPU、计算后梯度回写 CPU。七、如何验证与横向对比仓库在 tests/test_auto_parallel/test_offload 中给出了两层测试test_solver.py对gpt2_与bert_两种模型、syn与asyn两种求解器分别跑求解流程并断言solver.best_ts.peak_mem memory_budget即求解器给出的方案峰值显存必须低于预算test_perf.py在同一份代码里并排跑 Geminiplacement_policycpu与 asyn offload 两种方案打印各自的exec_time与runtime_peak_mem_*方便直接对照吞吐与显存表现注意该用例默认被pytest.mark.skip跳过需要显式开启。结合 test_perf.py 的做法你可以在本演示中自行做横向实验# 同一显存预算下对比同步与异步卸载 bash run.sh # 默认 asyngpt2_medium16GB SOLVER_TYPEsyn bash run.sh MODEL_TYPEgpt2_xl MEMORY_BUDGET24 bash run.sh每组运行都会在offload_logs/下留下独立日志可直接比对solver_time、exec_time、runtime_peak_mem_alc三项指标。需要强调的是是否出现卸载、卸载多少个 region完全由求解器结合硬件画像与显存预算动态决定——当MEMORY_BUDGET设得足够大时求解结果可能趋向于几乎不卸载这也是该方案被称为自动的原因。八、常见问题与使用注意求解时间长求解器内部会多次运行训练模拟器做贪心搜索solver.py模型越大、候选 region 越多solver_time越长。它是离线成本训练前一次性支付报 cant find the offload strategy ...说明即使把所有可卸载 region 全部卸出峰值显存仍高于预算。需要调大MEMORY_BUDGET或选用更小的MODEL_TYPE需要 pynvml带宽/算力画像依赖 pynvml 读取真实设备信息solver.py顶部try: from pynvml import *缺失时相关测试会被NOT_NVML标记跳过半精度前提memory_optimize内部强制model.cpu().half()输入中的浮点张量也需转换为 FP16脚本用wrap_fntree_map完成模型请勿残留 FP32 必须算子或敏感 BatchNorm 等结构适配新模型只要模型结构可被ColoTracer追踪就能套用构建模型 → 生成 dummy 输入 →memory_optimize→AMPOptimizer这一通用模板无需手工规划卸载点。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考