设置DeepSpeed Stage3后，Lightning模型未层并行且报CUDA OOM错误

我想用PyTorch Lightning搭建基于DeepSpeed的示例，已设置deepspeed_stage_3策略，但模型层未被并行化。为测试，我添加了18000个全连接层，期望这些层能分布到6个GPU上，却触发了CUDA显存不足错误：

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate
2.00 MiB (GPU 3; 15.00 GiB total capacity; 14.00 GiB already allocated; 5.25 MiB free; 14.00 GiB reserved in total by PyTorch) If
reserved memory is >> allocated memory try setting max_split_size_mb
to avoid fragmentation. See documentation for Memory Management and
PYTORCH_CUDA_ALLOC_CONF

推测所有层都被部署到了单个GPU。核心代码如下：

模型定义代码 #

class TelModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        embed_dim = 512
        component_list = [
                nn.Linear(512, embed_dim)
        #] + [nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True) for _ in range(n_layers)] + [
        ] + [nn.Linear(embed_dim, 512) for _ in range(n_layers)] + [
                nn.Linear(embed_dim, 512)
        ]
        self.net = torch.nn.Sequential(*component_list)

DeepSpeed初始化代码 #

tel_model = TelModel()
train_ds = RandomDataset(100)
train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE)
trainer = L.Trainer(accelerator="gpu", devices=6, strategy="deepspeed_stage_3", precision=32)
trainer.fit(tel_model, train_loader)

运行命令 #

deepspeed lightning-deepspeed-tel.py

1. 明确DeepSpeed Stage3的并行逻辑 #

DeepSpeed Stage3是ZeRO优化+流水线并行的组合，默认不会自动将nn.Sequential中的层拆分到多个GPU。要实现层级分布，需要手动划分模型阶段，让流水线并行机制识别并分配。

2. 手动拆分Sequential模型为多阶段 #

将大量全连接层按GPU数量拆分到不同的子模块中，让DeepSpeed能将每个子模块分配到对应GPU：

class TelModel(L.LightningModule):
    def __init__(self, num_gpus=6):
        super().__init__()
        embed_dim = 512
        n_layers = 18000
        layers_per_gpu = n_layers // num_gpus
        
        # 按GPU数量拆分模型为多个阶段
        self.stages = torch.nn.ModuleList()
        # 初始输入层
        self.stages.append(nn.Linear(512, embed_dim))
        
        # 拆分全连接层到各个阶段
        for gpu_idx in range(num_gpus):
            start = gpu_idx * layers_per_gpu
            # 处理最后一个GPU的剩余层
            end = start + layers_per_gpu if gpu_idx < num_gpus - 1 else n_layers
            stage_layers = [nn.Linear(embed_dim, 512) for _ in range(start, end)]
            self.stages.append(torch.nn.Sequential(*stage_layers))
        
        # 最终输出层
        self.stages.append(nn.Linear(embed_dim, 512))

    def forward(self, x):
        for stage in self.stages:
            x = stage(x)
        return x

3. 配置流水线并行参数 #

在Trainer中显式指定流水线并行的阶段数，或者通过DeepSpeed配置文件更精细地控制：

方式1：直接在Trainer中设置

trainer = L.Trainer(
    accelerator="gpu", 
    devices=6, 
    strategy="deepspeed_stage_3", 
    precision=32,
    pipeline_parallel_size=6  # 与GPU数量匹配
)

方式2：使用DeepSpeed配置文件（推荐，更灵活）

from pytorch_lightning.strategies import DeepSpeedStrategy

ds_config = {
    "train_batch_size": BATCH_SIZE * 6,  # 全局batch size
    "train_micro_batch_size_per_gpu": BATCH_SIZE,  # 单GPU的micro batch
    "pipeline_parallel": {
        "enabled": True,
        "num_stages": 6  # 流水线阶段数，对应GPU数量
    },
    "zero_optimization": {
        "stage": 3,
        "offload_param": {
            "device": "cpu"  # 开启参数CPU卸载，进一步降低显存占用
        }
    }
}

trainer = L.Trainer(
    accelerator="gpu", 
    devices=6, 
    strategy=DeepSpeedStrategy(config=ds_config), 
    precision=32
)

4. 关键注意事项 #

• ZeRO Stage3主要优化显存使用，但不会自动做层拆分，流水线并行需要手动划分模型阶段。
• 确保pipeline_parallel_size与实际使用的GPU数量一致，或者根据需求调整阶段数。
• 开启ZeRO的参数卸载（offload_param）可以大幅降低单GPU的显存占用，尤其适合超大模型。

内容的提问来源于stack exchange，提问作者Romeo Kienzler