1. 模型优化器到底在解决什么问题第一次接触 Model-Optimizer 这个概念是在一个推荐系统的排序模型上。当时线上推理延迟卡在 120ms 下不去GPU 利用率却只有 30% 出头显存倒是先爆了。排查了一圈发现问题不在模型结构也不在数据管道而是优化器状态把显存吃掉了将近一半——Adam 的动量项和方差项各存一份参数量乘以 2 再乘以 4 字节一个 1.2B 参数的模型光优化器状态就接近 10GB。这就是 Model-Optimizer 这类工具要解决的核心矛盾训练时优化器带来的显存开销和计算效率问题。Model-Optimizer 本质上是一套围绕模型训练与推理全流程的优化工具集它做的事情可以拆成三条线。第一条线是显存优化通过优化器状态分片、梯度累积、混合精度等手段把单卡放不下的模型拆到多卡上或者让同样的卡能跑更大的 batch。第二条线是计算效率优化包括算子融合、通信压缩、计算图重写让每一步迭代的实际耗时降下来。第三条线是推理侧优化把训练好的模型做量化、剪枝、蒸馏让它在生产环境里跑得更快更省。适合谁来参考这份内容如果你正在做以下任意一件事这篇东西应该对你有用模型参数量超过单卡显存、训练吞吐上不去、推理延迟不达标、想用有限硬件跑更大模型。不需要你是分布式训练的专家但至少得能看懂 PyTorch 的DataParallel和DistributedDataParallel的区别知道torch.cuda.amp是干什么的。下面我会从设计思路、核心细节、实操过程到问题排查把 Model-Optimizer 这套东西拆开讲清楚。2. 整体设计思路与方案选型拆解2.1 为什么优化器状态是显存第一杀手要理解 Model-Optimizer 的设计逻辑得先算清楚一笔账。假设你有一个参数量为 P 的模型用 Adam 优化器训练精度是 FP32。模型本身占 4P 字节梯度占 4P 字节优化器状态里动量占 4P、方差占 4P加起来就是 16P 字节。如果 P 是 10 亿那就是 16GB还没算激活值和临时缓冲区。用混合精度训练的话模型权重会存一份 FP32 的 master copy 和一份 FP16 的计算副本显存占用进一步上升。Model-Optimizer 的第一个设计决策就是把优化器状态从每卡一份变成多卡分片。这个思路最早来自 ZeROZero Redundancy Optimizer系列工作核心逻辑是既然数据并行时每张卡上的梯度最终要 all-reduce 成一样的那优化器状态就没必要每张卡都存一份。把优化器状态按参数切分到不同卡上每张卡只维护自己那一片的优化器状态更新完再把参数广播回去。这样显存占用就从 16P 降到了 16P/NN 是卡数代价是增加了一次 all-gather 通信。注意分片策略不是银弹。当卡数很多时通信开销会线性增长需要根据实际带宽和模型大小做权衡。通常 8 卡以内收益最明显超过 32 卡就要考虑通信压缩了。2.2 混合精度训练的三个关键取舍Model-Optimizer 在精度上的设计也值得细说。混合精度训练不是简单地把 FP32 换成 FP16 就完事里面有三个关键取舍。第一个取舍是哪些算子用 FP16哪些保留 FP32。矩阵乘法和卷积用 FP16 收益最大因为这两个算子是计算密集型的FP16 的吞吐量通常是 FP32 的 2 到 8 倍。但像 softmax、layer norm、loss 计算这些涉及归约和指数运算的算子FP16 的动态范围不够容易溢出必须保留 FP32。第二个取舍是损失缩放loss scaling策略。FP16 能表示的最小正数大约是 6e-8梯度在反向传播过程中会越来越小很容易下溢成 0。解决办法是在计算 loss 时乘一个缩放因子反向传播得到的梯度也相应放大更新参数前再除回去。静态缩放是固定一个值动态缩放是根据梯度是否出现 inf/nan 自动调整。Model-Optimizer 默认用动态缩放初始值设 2^16连续 2000 步没有溢出就翻倍出现溢出就减半并跳过这一步。第三个取舍是master weight 的维护。FP16 的精度只有 10 位尾数参数更新时如果直接累加到 FP16 权重上小梯度会被舍入误差吃掉。所以需要维护一份 FP32 的 master weight每次更新在 FP32 上做更新完再转成 FP16 给前向计算用。这份 master weight 会额外占 4P 字节显存但这是保证收敛性必须付出的代价。2.3 通信优化的两条路线分布式训练里通信往往是瓶颈。Model-Optimizer 在通信优化上有两条路线梯度压缩和通信重叠。梯度压缩的思路是减少每次 all-reduce 传输的数据量。常见做法有量化和稀疏化。量化是把 FP32 梯度压成 INT8 或 INT16传输量直接降 2 到 4 倍代价是引入量化误差。稀疏化是只传梯度中绝对值最大的那部分比如只传 top 1% 的元素其余置零不传。实测下来稀疏化在梯度本身就很稀疏的场景比如 embedding 层效果很好但在稠密梯度上压缩率有限。通信重叠的思路是让通信和计算并行。反向传播是逐层计算的第 L 层的梯度算完就可以开始 all-reduce不用等所有层都算完。Model-Optimizer 会把梯度按层分组算完一组就触发一次异步 all-reduce这样通信时间就被计算时间掩盖掉了。这个优化在带宽受限的环境下收益特别明显我实测过一个 8 卡 A100 的集群开启通信重叠后训练吞吐提升了将近 40%。3. 核心细节解析与实操要点3.1 优化器状态分片的参数配置配置优化器状态分片时有几个参数需要重点关注。以 DeepSpeed 的 ZeRO 为例Stage 1 只分片优化器状态Stage 2 额外分片梯度Stage 3 连模型参数也分片。选择哪个 Stage 取决于你的瓶颈在哪里。如果显存瓶颈主要在优化器状态Stage 1 就够了通信开销最小。如果梯度也占了很多显存上 Stage 2。如果模型本身单卡都放不下必须用 Stage 3。但 Stage 3 的通信开销最大因为每次前向传播都要 all-gather 参数反向传播又要 reduce-scatter 梯度。# DeepSpeed ZeRO Stage 2 配置示例 { zero_optimization: { stage: 2, allgather_partitions: true, allgather_bucket_size: 5e8, overlap_comm: true, reduce_scatter: true, reduce_bucket_size: 5e8, contiguous_gradients: true }, fp16: { enabled: true, loss_scale: 0, loss_scale_window: 1000, initial_scale_power: 16, hysteresis: 2, min_loss_scale: 1 } }allgather_bucket_size和reduce_bucket_size这两个参数控制通信桶的大小。桶太小通信次数多启动开销大桶太大显存占用高而且通信和计算的重叠效果差。5e8 是 5 亿个元素按 FP16 算就是 1GB这个值在 8 卡 A100 上实测比较均衡。overlap_comm打开通信重叠contiguous_gradients把梯度放在连续内存里减少内存碎片。实操心得initial_scale_power设 16 意味着初始 loss scale 是 65536。如果你的模型 loss 本身很大比如上千这个值可能偏小会导致频繁溢出。可以先跑几百步观察溢出频率如果超过 5% 就调大这个值。3.2 梯度累积与 micro-batch 的配合当显存不够跑大 batch 时梯度累积是常用手段。原理很简单把一个大 batch 拆成 K 个小 batch每个小 batch 前向反向算梯度但不清零累加 K 次后再更新参数。这样等效 batch size 就是 micro-batch size 乘以 K 再乘以数据并行度。但梯度累积有个坑BatchNorm 层的行为会变。BatchNorm 在训练时用当前 batch 的均值和方差做归一化如果 micro-batch 太小统计量估计不准训练会不稳定。解决办法是用 SyncBatchNorm在多个卡和多个 micro-batch 之间同步统计量或者干脆把 BatchNorm 换成 GroupNorm 或 LayerNorm。另一个坑是学习率需要重新调。等效 batch size 变大后梯度的方差变小理论上可以用更大的学习率。经验法则是学习率随等效 batch size 线性缩放但超过某个阈值后要改成平方根缩放。我一般会先用小 batch 跑一个 baseline确定最佳学习率然后按线性缩放规则放大再微调。# 梯度累积 混合精度的标准写法 scaler torch.cuda.amp.GradScaler() accumulation_steps 4 for i, (inputs, labels) in enumerate(dataloader): with torch.cuda.amp.autocast(): outputs model(inputs) loss criterion(outputs, labels) loss loss / accumulation_steps scaler.scale(loss).backward() if (i 1) % accumulation_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad()注意 loss 要除以accumulation_steps这样累加后的梯度才和真实大 batch 的梯度在同一个量级。scaler.step和scaler.update只在累积满的时候调用zero_grad也是。3.3 算子融合的实际收益算子融合是 Model-Optimizer 里容易被忽视但收益很实在的一块。深度学习模型里有很多小算子比如add、relu、dropout每个算子单独启动一个 CUDA kernelkernel launch 的开销可能比计算本身还大。算子融合就是把这些小算子合并成一个 kernel减少 launch 次数和内存读写。最典型的例子是bias relu融合。原本需要两个 kernel第一个算x b写回显存第二个读出来算relu再写回。融合后一个 kernel 搞定省了一次显存读写。在 Transformer 模型里layer_norm dropout add这种组合很常见融合后能省 20% 到 30% 的显存带宽。PyTorch 2.0 之后引入了torch.compile底层用 TorchInductor 做算子融合基本是开箱即用。但要注意不是所有模型都能顺利编译动态控制流多的模型可能会 fallback 回 eager 模式。我一般会先用torch.compile跑一遍看编译日志里有多少算子被融合了如果融合率低于 50%就得手动检查哪里出了问题。4. 实操过程与核心环节实现4.1 环境准备与依赖安装先把环境搭起来。我用的组合是 PyTorch 2.1 CUDA 12.1 DeepSpeed 0.12这个组合在 A100 和 H100 上都比较稳。安装顺序很重要先装 PyTorch再装 DeepSpeed最后装其他依赖。# 创建虚拟环境 conda create -n model-opt python3.10 -y conda activate model-opt # 安装 PyTorch根据 CUDA 版本选择 pip install torch2.1.0 torchvision0.16.0 torchaudio2.1.0 --index-url https://download.pytorch.org/whl/cu121 # 安装 DeepSpeed pip install deepspeed0.12.0 # 安装其他依赖 pip install transformers4.36.0 accelerate0.25.0 datasets2.16.0装完先验证一下 DeepSpeed 能不能正确识别环境ds_report这个命令会输出 DeepSpeed 的编译选项、CUDA 版本、支持的优化器等信息。重点看Compatible NCCL和Compatible Apex是不是 True如果 NCCL 不兼容分布式训练会直接报错。注意DeepSpeed 安装时会编译一些 CUDA 扩展如果编译失败可能是 CUDA toolkit 版本和 PyTorch 的 CUDA 版本不一致。用nvcc --version和python -c import torch; print(torch.version.cuda)对比一下不一致就重装。4.2 训练脚本的改造步骤假设你有一个单卡训练脚本要改造成支持 Model-Optimizer 的分布式训练需要动这几个地方。第一步初始化分布式环境。用deepspeed.initialize替代原来的模型和优化器创建过程。import deepspeed # 原来的写法 # model MyModel() # optimizer torch.optim.Adam(model.parameters(), lr1e-4) # 改造后 model_engine, optimizer, _, _ deepspeed.initialize( argsargs, modelmodel, model_parametersmodel.parameters(), config_paramsds_config )deepspeed.initialize会返回一个DeepSpeedEngine它包装了模型、优化器和学习率调度器。后续的前向反向都用这个 engine。第二步改造训练循环。把loss.backward()和optimizer.step()换成 engine 的对应方法。for step, batch in enumerate(dataloader): inputs batch[input_ids].to(model_engine.device) labels batch[labels].to(model_engine.device) outputs model_engine(inputs, labelslabels) loss outputs.loss model_engine.backward(loss) model_engine.step()model_engine.backward内部处理了混合精度的 loss scalingmodel_engine.step处理了梯度累积和优化器状态分片。你不需要再手动调scaler.scale和scaler.step。第三步配置启动脚本。用deepspeed命令替代python启动指定卡数和配置文件。deepspeed --num_gpus8 train.py \ --deepspeed ds_config.json \ --model_name_or_path bert-base-uncased \ --per_device_train_batch_size 16 \ --gradient_accumulation_steps 4 \ --learning_rate 2e-5 \ --num_train_epochs 3--num_gpus8指定用 8 张卡--deepspeed指定配置文件路径。如果只想用部分卡可以用--include localhost:0,1,2,3指定。4.3 关键参数的计算与选择batch size 的配置需要算一笔账。假设你有 8 张卡每张卡 micro-batch size 是 16梯度累积 4 步那全局 batch size 就是 8 × 16 × 4 512。这个值要和你的学习率匹配。学习率的经验公式是lr base_lr × sqrt(global_batch_size / base_batch_size)。如果 base_batch_size 是 32 时 base_lr 是 1e-5那 global_batch_size 是 512 时lr 大约是 1e-5 × sqrt(512/32) 4e-5。但这只是起点实际还要看 loss 曲线微调。warmup steps 也很关键。分布式训练初期梯度噪声大直接上大学习率容易发散。一般设总步数的 5% 到 10% 做 warmup学习率从 0 线性增加到目标值。如果总步数是 10000warmup 设 500 到 1000 步。{ train_batch_size: 512, train_micro_batch_size_per_gpu: 16, gradient_accumulation_steps: 4, optimizer: { type: AdamW, params: { lr: 4e-5, betas: [0.9, 0.999], eps: 1e-8, weight_decay: 0.01 } }, scheduler: { type: WarmupDecayLR, params: { warmup_min_lr: 0, warmup_max_lr: 4e-5, warmup_num_steps: 1000, total_num_steps: 10000 } } }train_batch_size是全局 batch sizeDeepSpeed 会根据卡数和gradient_accumulation_steps自动算train_micro_batch_size_per_gpu如果算出来和你设的不一致会报错。4.4 推理侧优化的落地方法训练完的模型要上生产推理优化是另一套东西。Model-Optimizer 在推理侧主要做三件事量化、剪枝、图优化。量化是把 FP32 权重压成 INT8 或 INT4。INT8 量化基本不掉精度推理速度能提升 2 到 4 倍。INT4 量化压缩率更高但精度损失明显适合对精度不敏感的场景。PyTorch 自带的torch.quantization可以做动态量化但性能一般。生产环境更常用 TensorRT 或 ONNX Runtime 做量化。# 动态量化示例 import torch.quantization model_fp32 MyModel() model_fp32.eval() # 动态量化只量化 Linear 和 LSTM model_int8 torch.quantization.quantize_dynamic( model_fp32, {torch.nn.Linear}, dtypetorch.qint8 ) # 保存量化模型 torch.save(model_int8.state_dict(), model_int8.pth)剪枝是去掉模型中不重要的权重。结构化剪枝直接去掉整个通道或注意力头非结构化剪枝把单个权重置零。结构化剪枝对推理加速更友好因为硬件对稀疏矩阵的支持有限。剪枝后要 fine-tune 几个 epoch 恢复精度。图优化是把计算图里的冗余算子去掉比如恒等映射、常量折叠、算子融合。ONNX Runtime 和 TensorRT 都会自动做这些优化。导出 ONNX 模型时注意 opset 版本太低不支持某些算子太高可能不被推理引擎支持。一般用 opset 13 到 15 比较稳。5. 常见问题与排查技巧实录5.1 训练不收敛的排查路径分布式训练不收敛是最让人头疼的问题。我一般按这个顺序排查。先看 loss 曲线。如果 loss 从一开始就震荡或者直接变 nan大概率是学习率太大或者 loss scaling 有问题。把学习率降 10 倍再跑如果 loss 正常下降就是学习率的问题。如果还是 nan检查 loss scaling 的初始值调大initial_scale_power。如果 loss 前期正常跑着跑着突然变 nan通常是梯度爆炸。加梯度裁剪max_grad_norm设 1.0 到 5.0 之间。DeepSpeed 的配置里加{ gradient_clipping: 1.0 }如果 loss 下降但比单卡慢很多检查数据并行是否真的生效了。用nvidia-smi看每张卡的 GPU 利用率如果有的卡利用率很低可能是数据分配不均或者通信瓶颈。还有一个隐蔽的坑随机种子没对齐。分布式训练时每张卡的随机种子应该不同但数据加载的 shuffle 种子要一致否则不同卡可能读到重复数据。用torch.manual_seed(seed rank)设置模型初始化的种子用seed设置数据 shuffle 的种子。5.2 显存溢出的定位方法显存溢出OOM的报错信息通常只告诉你哪一步爆了不告诉你为什么爆。定位方法是用torch.cuda.memory_summary()打印显存分配详情。print(torch.cuda.memory_summary(device0, abbreviatedFalse))输出会列出当前显存分配、峰值显存、缓存分配器状态。重点看Allocated memory和Reserved memory的差值如果 Reserved 远大于 Allocated说明有内存碎片可以调PYTORCH_CUDA_ALLOC_CONF环境变量。export PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:128这个参数控制缓存分配器的最大分片大小设小一点可以减少碎片但会增加分配次数。128MB 是个比较均衡的值。如果显存是在某个特定层爆的比如 attention 层那可能是序列长度太长。Transformer 的 attention 显存占用是 O(n²)序列长度翻倍显存翻四倍。解决办法是用 FlashAttention 或者梯度检查点gradient checkpointing。梯度检查点用时间换空间前向传播时不存中间激活值反向传播时重新算一遍显存能省 50% 到 70%但训练速度会慢 20% 到 30%。5.3 通信超时的处理分布式训练跑着跑着卡住然后报 NCCL timeout这是通信问题。常见原因有三个。第一个是某张卡挂了。用nvidia-smi检查所有卡的状态如果有卡掉线重启训练。如果频繁掉线可能是硬件问题检查电源和散热。第二个是网络带宽不够。NCCL 默认用 ring all-reduce如果卡之间的网络带宽不对称会拖慢整个通信。可以用NCCL_DEBUGINFO打印通信日志看实际用的什么算法和带宽。export NCCL_DEBUGINFO export NCCL_IB_DISABLE0 # 如果有 InfiniBand 就打开 export NCCL_SOCKET_IFNAMEeth0 # 指定网卡第三个是通信桶大小设置不合理。桶太小通信次数多容易超时桶太大单次通信时间长也容易超时。调allgather_bucket_size和reduce_bucket_size从 5e8 开始每次减半或翻倍试。5.4 常见问题速查表问题现象可能原因排查方法解决方案loss 变 nan学习率过大、loss scaling 溢出降学习率、看溢出日志调小 lr、调大 initial_scale_power显存 OOMbatch 太大、激活值未释放memory_summary 看分配减 batch、开梯度检查点训练速度慢通信瓶颈、数据加载慢nvidia-smi 看利用率开通信重叠、增加 dataloader workersNCCL timeout网络问题、桶大小不合理NCCL_DEBUGINFO调桶大小、检查网卡精度下降量化损失、剪枝过度对比 FP32 输出减少量化位宽、降低剪枝率多卡结果不一致随机种子未对齐检查 seed 设置模型种子加 rank数据种子固定实操心得遇到问题先别急着改代码把NCCL_DEBUGINFO和TORCH_DISTRIBUTED_DEBUGDETAIL打开让日志告诉你发生了什么。我踩过的坑里有一半是因为没看日志瞎猜浪费了大半天时间。6. 几个容易被忽视的优化细节6.1 dataloader 的 num_workers 设置num_workers设多少合适经验值是 CPU 核数除以 GPU 数再乘以 2。比如 8 卡机器有 64 核那num_workers设 16。设太小GPU 等数据设太大CPU 上下文切换开销大。可以用htop看 CPU 利用率如果某个核跑满而其他核闲着说明num_workers不够。另外pin_memoryTrue一定要开它把数据放在锁页内存里GPU 可以直接通过 DMA 读取省一次内存拷贝。persistent_workersTrue也建议开避免每个 epoch 重新创建 worker 进程。6.2 梯度检查点的正确用法梯度检查点不是所有层都值得开。它用时间换空间开了之后训练速度会降。一般只在显存瓶颈的层开比如 Transformer 的每一层。PyTorch 的torch.utils.checkpoint可以精确控制哪些层用检查点。from torch.utils.checkpoint import checkpoint class TransformerLayer(nn.Module): def forward(self, x): if self.training and self.use_checkpoint: return checkpoint(self._forward, x) return self._forward(x)注意checkpoint要求输入张量的requires_gradTrue否则不会保存计算图反向传播会报错。另外 dropout 层在检查点里会有随机性问题因为前向算了两次dropout mask 不一致。解决办法是在检查点外面套一个固定的随机种子或者把 dropout 移到检查点外面。6.3 学习率调度的 warmup 策略warmup 不是越长越好。warmup 太长前期学习率太小浪费计算warmup 太短前期梯度噪声大容易发散。我一般用线性 warmup步数设总步数的 5%。如果模型特别大比如 10B 以上可以设到 10%。warmup 之后用 cosine decay 比 step decay 更平滑最终学习率降到峰值的 10% 左右。如果训练步数不多比如几千步可以用 constant learning rate 加 warmup避免后期学习率太小导致欠拟合。from transformers import get_cosine_schedule_with_warmup scheduler get_cosine_schedule_with_warmup( optimizer, num_warmup_steps1000, num_training_steps10000 )这个 scheduler 每个 step 调一次scheduler.step()注意是在optimizer.step()之后调。6.4 模型保存与恢复的注意事项分布式训练的模型保存和单卡不一样。DeepSpeed 的model_engine.save_checkpoint会保存分片的优化器状态和模型参数恢复时用model_engine.load_checkpoint。如果只想保存模型权重给推理用用model_engine.module.state_dict()拿到完整的 state dict。# 保存完整模型权重 if model_engine.local_rank 0: state_dict model_engine.module.state_dict() torch.save(state_dict, model_weights.pth) # 保存 DeepSpeed checkpoint含优化器状态 model_engine.save_checkpoint(checkpoints, tagstep_10000)注意save_checkpoint要在所有卡上调用但只有 rank 0 会真正写文件。load_checkpoint也是所有卡都调它会自动加载对应分片。注意DeepSpeed 的 checkpoint 和 PyTorch 的 checkpoint 格式不兼容。如果要从 DeepSpeed checkpoint 转成普通 PyTorch 权重用zero_to_fp32.py脚本转换这个脚本在 DeepSpeed 的安装目录里。7. 从训练到推理的完整链路打通7.1 模型导出与格式转换训练完的模型要上推理引擎第一步是导出。PyTorch 模型导出成 ONNX 或 TorchScript再转成 TensorRT 或 OpenVINO。# 导出 ONNX dummy_input torch.randn(1, 128, 768).cuda() torch.onnx.export( model, dummy_input, model.onnx, opset_version14, input_names[input_ids, attention_mask], output_names[logits], dynamic_axes{ input_ids: {0: batch, 1: sequence}, attention_mask: {0: batch, 1: sequence}, logits: {0: batch, 1: sequence} } )dynamic_axes指定哪些维度是动态的batch 和 sequence 通常都是动态的。opset 版本选 14 是因为它支持 FlashAttention 相关的算子而且 TensorRT 8.x 兼容性最好。导出后验证一下 ONNX 模型的正确性用onnxruntime跑一遍和 PyTorch 的输出对比误差在 1e-4 以内算正常。7.2 推理引擎的选型对比推理引擎优势劣势适用场景ONNX Runtime跨平台、易用性能一般快速验证、CPU 推理TensorRT性能最强只支持 NVIDIA GPU生产环境 GPU 推理OpenVINOIntel CPU 优化好GPU 支持弱Intel CPU 推理TorchScript与 PyTorch 无缝优化有限PyTorch 生态内生产环境如果用的是 NVIDIA GPU首选 TensorRT。它的 INT8 量化和 kernel 自动调优能把推理延迟压到最低。但 TensorRT 的模型转换比较麻烦需要写转换脚本而且不同版本的 TensorRT 兼容性差升级要重新转换。ONNX Runtime 胜在通用CPU 和 GPU 都能跑而且支持多种量化方式。如果推理服务要部署在多种硬件上ONNX Runtime 更省心。7.3 推理服务的性能调优推理服务的性能不只是模型本身还有服务框架。几个关键点。批处理batching把多个请求攒成一个 batch 一起推理GPU 利用率更高。但 batch 太大会增加延迟需要根据 SLA 权衡。Triton Inference Server 支持动态批处理可以设最大 batch size 和最大等待时间。并发concurrency多个请求同时进来时用多个 CUDA stream 并行推理。但 GPU 的计算资源有限并发太高反而会互相抢占。一般设 2 到 4 个 stream 比较合适。内存池memory pool推理时频繁分配释放显存会产生碎片用内存池预分配一块大显存重复使用。TensorRT 和 ONNX Runtime 都内置了内存池但需要配置初始大小。# ONNX Runtime 内存池配置 sess_options ort.SessionOptions() sess_options.enable_cpu_mem_arena True sess_options.enable_mem_pattern True sess_options.execution_mode ort.ExecutionMode.ORT_SEQUENTIALenable_mem_pattern会记录内存分配模式重复推理时直接复用减少分配开销。7.4 端到端性能对比我在一个 BERT-base 模型上做了完整的对比测试硬件是单卡 A100 40GB输入序列长度 128batch size 32。配置推理延迟 (ms)显存占用 (MB)精度 (F1)PyTorch FP324512000.912PyTorch FP16226800.912ONNX Runtime FP16186200.911TensorRT FP16125500.911TensorRT INT873800.908INT8 量化后延迟降到 7ms精度只掉了 0.4 个点这个收益非常划算。但 INT8 量化需要校准数据集校准集的质量直接影响量化精度。我一般从训练集里抽 500 到 1000 条样本做校准覆盖各种长度的输入。训练侧的优化效果也很明显。同一个模型单卡训练显存 24GB8 卡 ZeRO Stage 2 后每卡显存降到 6GB训练吞吐从 120 samples/s 提升到 850 samples/s。这个提升主要来自显存释放后能跑更大的 batch以及通信重叠减少了等待时间。踩过几次坑之后我的体会是Model-Optimizer 这套东西不是配置越多越好每个优化都有代价。优化器状态分片增加通信混合精度增加调参复杂度量化损失精度。关键是找到瓶颈在哪里然后针对性地开对应的优化。盲目把所有开关都打开往往适得其反。