Model-Optimizer实战：显存优化与推理加速全解析

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 的第一个设计决策就是把优化器状态从"每卡一份"变成"多卡分片"。这个思路最早来自 ZeRO（Zero Redundancy Optimizer）系列工作，核心逻辑是：既然数据并行时每张卡上的梯度最终要 all-reduce 成一样的，那优化器状态就没必要每张卡都存一份。把优化器状态按参数切分到不同卡上，每张卡只维护自己那一片的优化器状态，更新完再把参数广播回去。这样显存占用就从 16P 降到了 16P/N（N 是卡数），代价是增加了一次 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 kernel，kernel 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 python=3.10 -y conda activate model-opt # 安装 PyTorch（根据 CUDA 版本选择） pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121 # 安装 DeepSpeed pip install deepspeed==0.12.0 # 安装其他依赖 pip install transformers==4.36.0 accelerate==0.25.0 datasets==2.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(), lr=1e-4) # 改造后 model_engine, optimizer, _, _ = deepspeed.initialize( args=args, model=model, model_parameters=model.parameters(), config_params=ds_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, labels=labels) loss = outputs.loss model_engine.backward(loss) model_engine.step()

model_engine.backward内部处理了混合精度的 loss scaling，model_engine.step处理了梯度累积和优化器状态分片。你不需要再手动调scaler.scale和scaler.step。

第三步，配置启动脚本。用deepspeed命令替代python启动，指定卡数和配置文件。

deepspeed --num_gpus=8 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_gpus=8指定用 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 线性增加到目标值。如果总步数是 10000，warmup 设 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 size，DeepSpeed 会根据卡数和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}, dtype=torch.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(device=0, abbreviated=False))

输出会列出当前显存分配、峰值显存、缓存分配器状态。重点看Allocated memory和Reserved memory的差值，如果 Reserved 远大于 Allocated，说明有内存碎片，可以调PYTORCH_CUDA_ALLOC_CONF环境变量。

export PYTORCH_CUDA_ALLOC_CONF=max_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_DEBUG=INFO打印通信日志，看实际用的什么算法和带宽。

export NCCL_DEBUG=INFO export NCCL_IB_DISABLE=0 # 如果有 InfiniBand 就打开 export NCCL_SOCKET_IFNAME=eth0 # 指定网卡

第三个是通信桶大小设置不合理。桶太小，通信次数多，容易超时；桶太大，单次通信时间长，也容易超时。调allgather_bucket_size和reduce_bucket_size，从 5e8 开始，每次减半或翻倍试。

5.4 常见问题速查表

实操心得：遇到问题先别急着改代码，把NCCL_DEBUG=INFO和TORCH_DISTRIBUTED_DEBUG=DETAIL打开，让日志告诉你发生了什么。我踩过的坑里，有一半是因为没看日志瞎猜，浪费了大半天时间。

6. 几个容易被忽视的优化细节

6.1 dataloader 的 num_workers 设置

num_workers设多少合适？经验值是 CPU 核数除以 GPU 数，再乘以 2。比如 8 卡机器有 64 核，那num_workers设 16。设太小，GPU 等数据；设太大，CPU 上下文切换开销大。可以用htop看 CPU 利用率，如果某个核跑满而其他核闲着，说明num_workers不够。

另外pin_memory=True一定要开，它把数据放在锁页内存里，GPU 可以直接通过 DMA 读取，省一次内存拷贝。persistent_workers=True也建议开，避免每个 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_grad=True，否则不会保存计算图，反向传播会报错。另外 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_steps=1000, num_training_steps=10000 )

这个 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', tag='step_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_version=14, 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 推理引擎的选型对比

生产环境如果用的是 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_SEQUENTIAL

enable_mem_pattern会记录内存分配模式，重复推理时直接复用，减少分配开销。

7.4 端到端性能对比

我在一个 BERT-base 模型上做了完整的对比测试，硬件是单卡 A100 40GB，输入序列长度 128，batch size 32。

INT8 量化后延迟降到 7ms，精度只掉了 0.4 个点，这个收益非常划算。但 INT8 量化需要校准数据集，校准集的质量直接影响量化精度。我一般从训练集里抽 500 到 1000 条样本做校准，覆盖各种长度的输入。

训练侧的优化效果也很明显。同一个模型，单卡训练显存 24GB，8 卡 ZeRO Stage 2 后每卡显存降到 6GB，训练吞吐从 120 samples/s 提升到 850 samples/s。这个提升主要来自显存释放后能跑更大的 batch，以及通信重叠减少了等待时间。

踩过几次坑之后，我的体会是：Model-Optimizer 这套东西不是配置越多越好，每个优化都有代价。优化器状态分片增加通信，混合精度增加调参复杂度，量化损失精度。关键是找到瓶颈在哪里，然后针对性地开对应的优化。盲目把所有开关都打开，往往适得其反。