使用DeepSpeed进行模型并行优化时的踩坑记录

在使用DeepSpeed进行模型并行优化时，踩坑是家常便饭。最近在部署一个7B参数模型时，遇到了RuntimeError: Expected all tensors to be on the same device的错误。

问题出现在设置--model-parallel-size为4时，训练过程报错。通过排查发现，模型的某些层未正确分配到对应设备上。解决方法是显式指定每个参数的设备位置：

# 在初始化模型前添加
import torch.distributed as dist
from deepspeed.runtime.engine import DeepSpeedEngine

device = torch.device(f'cuda:{dist.get_rank()}')
for name, param in model.named_parameters():
    if 'embedding' in name:
        param.data = param.data.to(device)

此外，为避免梯度同步问题，建议将--zero-stage 1改为--zero-stage 2，并设置合适的--gradient-accumulation-steps。实际测试中发现，当batch_size=8时，设置--gradient-accumulation-steps=4可有效缓解显存溢出。

在多机训练中还遇到通信超时问题，通过将--communication-queue-depth 16参数调高后问题解决。建议在大规模训练前先进行小规模验证，避免长时间等待。

经验总结：模型并行优化需要对设备分配、通信机制有深入理解，否则容易出现难以复现的bug。