DeepSpeed的核心特性与引擎架构
DeepSpeed是微软开源的大模型训练优化库。请介绍其ZeRO优化器、DeepSpeed Engine的架构设计,以及在ChatGPT类模型训练中的典型配置(ZeRO Stage/Offload/混合精度)。
回答
小字辈
DeepSpeed是Microsoft DeepSpeed团队为PyTorch构建的分布式训练优化库,核心围绕三个模块。
1. ZeRO优化器:
- Stage 1/2/3 逐步消除冗余存储(如前所述)
- ZeRO-Offload:利用CPU内存扩展显存
- ZeRO-Infinity:扩展到NVMe
2. DeepSpeed Engine架构: 用户模型 → DeepSpeed Engine注入 → 自动优化:自动ZeRO包装、自动混合精度设置、自动梯度累积、自动通信优化、自动Checkpointing
使用方式即替换优化器和Trainer:model_engine, optimizer, _, _ = deepspeed.initialize(model, optimizer, ...)
3. 典型ChatGPT(SFT/RLHF)配置: ZeRO Stage 3 + Offload + FP16混合精度 + 梯度累积
4. 其他特性:
- Sparse Attention(减少长序列计算)
- 1-bit Adam(梯度压缩通信)
- Progressive Layer Dropping(随机跳过部分层加速)
- ZeRO++(通信量化)
DeepSpeed现被广泛用于LLaMA-Factory、HuggingFace Trainer中。