Zero deepspeed是啥
时间: 2025-03-03 21:20:28 浏览: 31
### Zero Deepspeed 概览
Zero 是由微软开发的一种用于优化大规模模型训练的技术,旨在通过减少内存占用来提高分布式训练效率[^3]。该技术允许更高效的资源利用,在不牺牲性能的情况下支持更大规模的神经网络。
#### Stage 0: 基础优化
Stage 0 主要提供了一些基础性的改进措施,比如参数广播重叠通信与计算等基本操作上的并行化处理方式。这使得即使是在单节点多GPU环境下也能获得一定的加速效果。
#### Stage 1: 参数分片 (Parameter Sharding)
在这一阶段,除了继承自Stage 0 的特性外,还引入了参数矩阵切分机制——即将权重张量按照列方向拆分成多个子块,并分配给不同的设备负责更新;同时梯度也会相应地被分割开来累积求和后再同步回全局副本位置上完成一次完整的反向传播过程。这种方法可以有效降低每台机器所需承载的工作负载以及显存消耗水平。
#### Stage 2: 数据并行中的全精度激活检查点保存
此级别进一步扩展了前两个版本的功能集合,增加了对浮点数格式下中间状态变量存储空间的有效管理策略。具体而言就是仅保留必要的部分而非全部历史记录来进行后续迭代运算时所需的依赖关系重建工作,从而达到节省更多RAM的目的。此外,它也能够兼容更高阶别的混合精度模式以获取额外的速度增益优势。
#### Stage 3: 完整的 ZeRO-Offload 技术栈
这是目前最为先进的实现形式之一,不仅涵盖了上述提到过的所有优点而且还加入了CPU卸载功能,即把一部分不太紧急的任务交给后台处理器去执行进而释放出更多的即时可用算力供前台程序调用。这样一来便可以在保持良好吞吐率的同时显著削减整体功耗开销,非常适合应用于超大型集群环境中开展长时间稳定的科研项目研究活动。
```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import deepspeed
model_name_or_path = "bigscience/bloomz"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
training_args = {
'local_rank': -1,
'zero_optimization': {
'stage': 3,
'offload_optimizer': {'device': 'cpu'},
'allgather_partitions': True,
'allgather_bucket_size': 2e8,
'reduce_scatter': True,
'reduce_bucket_size': 2e8,
'overlap_comm': True,
'contiguous_gradients': True,
'sub_group_size': 1e9,
'elastic_checkpoint': True
}
}
# Initialize model with DeepSpeed
model_engine, optimizer, _, _ = deepspeed.initialize(
args=training_args,
model=model,
model_parameters=[p for p in model.parameters() if p.requires_grad],
config_params="./ds_config.json"
)
```
阅读全文
相关推荐


















