DeepSpeed 是微软开源的一个深度学习优化库，专门用来训练超大规模的模型。它的核心能力不只是多卡并行，更在于通过一系列显存优化策略，让单个 GPU 也能跑起原本塞不下的模型。而这一切的入口，几乎都藏在一个 json 配置文件里。不管你用的是 DeepSpeed 的命令行启动方式，还是直接调用它的 API，最终都要面对那个密密麻麻的参数列表。这篇文章就从头梳理一下 DeepSpeed 配置文件中最重要的几个模块，让你不再只是复制粘贴，而是真的理解每一行在做什么。
基础训练参数：定下训练的基调
配置文件的第一个大块通常用来定义训练的超参数，这部分参数直接决定了训练的基本行为。最容易让人混淆的是 train_batch_size、train_micro_batch_size_per_gpu 和 gradient_accumulation_steps 三者之间的关系。DeepSpeed 里全局的 batch size 等于 micro batch size 乘以 gradient accumulation steps 再乘以 GPU 数量。也就是说，如果你有 4 张卡，每张卡上的 micro batch size 设为 2，gradient accumulation steps 设为 8，那么全局的 train_batch_size 就是 4 * 2 * 8 = 64。这个参数在配置里是必填的，DeepSpeed 会根据它来反向校验你的设置是否自洽，如果算出来不一致，启动的时候会直接报错。
另一个很容易被忽视的参数是 gradient_clipping，它控制梯度裁剪的阈值。大模型训练时梯度很容易爆炸，设一个合理的值（比如 1.0）能有效防止 loss 变成 NaN。此外，steps_per_print 决定了每隔多少步打印一次训练日志。不要小看这个参数，在单步时间很长的训练任务里，如果设得太小，日志会刷屏；设得太大，出问题又很难及时定位。一般根据总步数取一个适中的值，比如每 10 步或每 100 步输出一次。
还有一个隐藏的细节是 wall_clock_breakdown，它默认是 false。当你发现训练速度不太对劲，想分析时间都花在哪里了，把它设为 true，DeepSpeed 就会在日志里详细打印出前向传播、反向传播、参数更新等各个环节的耗时，是定位性能瓶颈的重要工具。
ZeRO 优化：显存节省的核心
ZeRO（Zero Redundancy Optimizer）是 DeepSpeed 的招牌特性，也是配置文件里权重最高的部分。ZeRO 分为三个阶段，通过 zero_optimization 块下的 stage 参数来控制。Stage 1 是对优化器状态进行分片，把原本每个 GPU 都完整保存的 Adam 动量、方差等状态切分到不同卡上，显存节省明显，而且几乎不影响计算速度。Stage 2 在 Stage 1 的基础上，进一步把梯度也分片，节省更多显存，但会引入少量的通信开销。Stage 3 则连模型参数本身都分片，显存节省最激进，但实现最复杂，通信量也最大。
对于绝大多数百亿参数以下的模型，Stage 2 是性价比最高的选择。你只需要在配置里把 stage 设成 2，DeepSpeed 就会自动处理剩余逻辑。不过 Stage 2 下有几个配套参数值得留意。比如 contiguous_gradients 设为 true，可以让梯度在通信前先拼成连续内存，提高 allreduce 效率。 reduce_bucket_size 和 allgather_bucket_size 控制通信时的桶大小，默认值通常够用，但如果你发现通信与计算交织得不好，可以试着调整这两个值，数值越大单次通信量越大，启动次数越少。
当模型大到 Stage 2 都撑不住时，才需要考虑 Stage 3。Stage 3 的配置额外多出几个重要参数。 stage3_prefetch_bucket_size 和 stage3_param_persistence_threshold 用来控制参数预取的粒度和哪些小参数可以常驻显存。 stage3_max_live_parameters 是一个强力的显存控制开关，它限定了在前向/反向过程中最多同时保持多少完整参数，值越小显存占用越低，但重计算和通信次数会增加。如果不清楚怎么调，可以先设成 1e9 这样的较大值，等跑通后再逐步收紧。还有一个容易导致训练开不起来的是 stage3_gather_16bit_weights_on_model_save，建议设为 true，这样保存 checkpoint 时会自动把分片的参数收集合并成完整的 fp16 权重，否则存下来的状态只有 DeepSpeed 自己能读，做推理或转换格式时非常麻烦。
混合精度与优化器配置
DeepSpeed 默认支持 FP16 混合精度训练，通过 fp16 块来启用。这里面最关键的是 enabled 设为 true，以及可以选配一个 loss_scale 策略。动态损失缩放 loss_scale_window 和 initial_scale_power 基本不需要手动改动，DeepSpeed 内置的自动调整机制已经很鲁棒。如果想尝试 BF16，只需把 bf16 块里的 enabled 设为 true，同时记得把 fp16 设为 false，二者不能同时启用。BF16 在 A100 等新架构上可以获得更好的数值稳定性，但需要硬件和 PyTorch 版本支持。
优化器部分通常在配置里直接指定，比如 "optimizer": { "type": "AdamW", "params": { "lr": 1e-4, "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 0.01 } }。DeepSpeed 内置了对 Adam、AdamW、LAMB 等优化器的原生支持，并且这些优化器在与 ZeRO 结合时，内部实现针对分片和通信做了优化。如果你需要更复杂的优化器调度，也可以在训练脚本里自己构建 optimizer，然后通过 DeepSpeed 的接口传进去，这时候 json 里的 optimizer 字段就可以留空，但需要确保代码与配置的 zero stage 逻辑不冲突。
学习率调度器与激活检查点
学习率调度器同样可以在配置文件里直接声明，例如 "scheduler": { "type": "WarmupLR", "params": { "warmup_min_ratio": 0.005, "warmup_num_steps": 2000 } }。DeepSpeed 支持 WarmupLR、WarmupDecayLR 以及任意 PyTorch 原生 scheduler。如果你的训练脚本里已经手动定义了 scheduler 并传入 DeepSpeed 初始化函数，那么 json 里的 scheduler 配置会被忽略，不会发生冲突。
激活检查点（Activation Checkpointing）是另一个显存大户的调节器。在配置中添加 "activation_checkpointing": { "partition_activations": true, "cpu_checkpointing": true, "number_checkpoints": null } 就可以启用。这里的 partition_activations 指是否把激活值分片到不同 GPU，cpu_checkpointing 则允许把激活值卸载到 CPU 内存，虽然速度会慢一些，但在显存极度紧张时能救命。 number_checkpoints 代表每层做多少个检查点，一般不需要设，让框架自动决定即可。用上这一功能，配合 ZeRO Stage 2 或 3，通常能把支持的最大模型参数量再提升一倍以上。
一个可运行的配置示例
下面给出一个典型的大模型训练配置，模型规模约 7B，使用 4 张 A100 40GB 显卡。这个配置启用了 ZeRO Stage 2、FP16 混合精度、Warmup 调度器和激活检查点，batch size 设置为全局 64。
{
"train_batch_size": 64,
"train_micro_batch_size_per_gpu": 2,
"gradient_accumulation_steps": 8,
"gradient_clipping": 1.0,
"steps_per_print": 10,
"fp16": {
"enabled": true
},
"zero_optimization": {
"stage": 2,
"contiguous_gradients": true,
"overlap_comm": true,
"reduce_bucket_size": 5e8,
"allgather_bucket_size": 5e8
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": 3e-5,
"betas": [0.9, 0.95],
"eps": 1e-8,
"weight_decay": 0.1
}
},
"scheduler": {
"type": "WarmupDecayLR",
"params": {
"warmup_min_ratio": 0.005,
"warmup_num_steps": 2000,
"total_num_steps": 50000
}
},
"activation_checkpointing": {
"partition_activations": true,
"cpu_checkpointing": false
},
"wall_clock_breakdown": false
}
实际使用时，把这个文件保存为 ds_config.json，然后在启动命令里指定 --deepspeed ./ds_config.json 即可。随着模型规模和硬件环境的变化，需要调整的主要是 zero stage、micro batch size 和 bucket 大小这几个关键参数。理解每个参数背后的含义之后，配置就不再是玄学，而是一套可以清晰推理的显存与速度平衡术。
DeepSpeed配置json配置文件ZeRO优化修改时间：2026-08-12 06:03:35