DeepSpeed

深度学习训练（与推理）优化库：让大规模分布式训练更省显存、更好扩展。原微软团队发起，现由 DeepSpeed（deepspeedai）社区维护。

它是做什么的

围绕 ZeRO、3D Parallelism、ZeRO-Infinity、Sequence Parallelism（如 Ulysses）、DeepSpeed-MoE 等系统能力，降低大模型训练的显存与通信成本。相对「只做模型并行、必须大改网络结构」的路线，ZeRO 教程强调：许多场景下改配置 JSON 即可把优化器状态、梯度乃至参数切分到多卡，从而在数据并行框架里塞进更大模型。官方称曾支撑当时领先规模的语言模型训练（如 Megatron-Turing NLG、BLOOM 等，见官网/README 列举）。可与 Hugging Face Transformers、Accelerate、PyTorch Lightning、MosaicML Composer 等集成；也支持多类加速器（NVIDIA、AMD ROCm、Intel XPU/Gaudi、华为 Ascend 等，见 Getting Started / accelerator 指南）。近期文档与博客还覆盖长序列训练、编译优化、新优化器对接等，以官网 Latest News 为准。

要点

• ZeRO：按阶段切分优化器状态（Stage 1）、梯度（Stage 2）、参数（Stage 3），放大数据并行可用显存；Stage 3 可配合 Infinity 卸到 CPU/NVMe
• 配置驱动：ds_config.json 开关 fp16、zero_optimization、offload 等；ZeRO 教程强调常改配置即可，不必大改模型代码
• 引擎 API：deepspeed.initialize 包装 nn.Module；训练环用可调用 forward、backward、step；checkpoint 用 save_checkpoint / load_checkpoint（各 rank 都要调用）
• 启动：deepspeed 入口做单机/多机；可用 hostfile、--include/--exclude，或 --no_ssh 模式
• 框架集成：Transformers 可通过 --deepspeed + config；Lightning / Accelerate 等有官方文档路径
• 和 Megatron 的关系：两者常被一起提到；DeepSpeed 偏 ZeRO/配置与引擎包装，Megatron 偏模型并行参考实现——可按团队栈二选一或组合，先看你现有训练入口

适合谁

• 单卡/少卡显存不够，要更大模型或更大有效 batch
• 已在用 Transformers Trainer / Accelerate / Lightning，想接 ZeRO
• 需要多机启动与分布式 checkpoint，而不是从零写通信

官方入口

• 官网
• Getting Started
• ZeRO 教程
• GitHub（deepspeedai/DeepSpeed）（旧址 microsoft/DeepSpeed 会 301 重定向）
• Transformers × DeepSpeed

建议怎么用

1. pip install deepspeed，按 Getting Started 跑通 deepspeed.initialize 最小循环：forward → backward → step，并用 deepspeed <script> --deepspeed_config ds_config.json 启动。
2. 若用 Hugging Face：在已有 Trainer / Accelerate 流程上加 config，先开 ZeRO Stage 1 或 2（对照 ZeRO 教程 样例 JSON），确认能训、能存盘，再考虑 Stage 3 与 offload_optimizer / offload_param。
3. 一次只加一类优化；显存仍不够再升级 stage 或 offload。需要取出权重时，注意 ZeRO-3 下 state_dict 可能只是占位，按文档使用 save_16bit_model 或 zero_to_fp32.py，不要假设和单卡 checkpoint 行为完全一样。