DeepSpeed 介绍

主要参考：

1. https://www.deepspeed.ai/getting-started/
2. Tutorials · MPI Tutorial

代码中使用 deepspeed

deepspeed.init_distributed() 

# 如果期间不需要使用到分布式环境，可以跳过上面的操作

model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args,
                                           model=model,    # 任意 torch.nn.module 类型模型
                                           model_parameters=params)


#load checkpoint
_, client_sd = model_engine.load_checkpoint(args.load_dir, args.ckpt_id)
step = client_sd['step']

#advance data loader to ckpt step
dataloader_to_step(data_loader, step + 1)

for step, batch in enumerate(data_loader):

    #forward() method
    loss = model_engine(batch)

    #runs backpropagation
    model_engine.backward(loss)

    #weight update
    model_engine.step()

    #save checkpoint
    if step % args.save_interval:
        client_sd['step'] = step  # 用户自定义需要存储的中间结果
        ckpt_id = loss.item()
        model_engine.save_checkpoint(args.save_dir, ckpt_id, client_sd = client_sd)

注意事项：

1. 如果原始代码使用了 torch.distributed.init_process_group(...) 来初始化分布式环境，需要去除或者替换为 deepspeed.init_distributed()
2. 在保存模型时，不光是 rank=0 的机器需要执行保存指令，所有的 worker 都需要执行保存指令来保存它们对应的 master weights and scheduler+optimizer states。如果只在 rank=0 的机器上调用，对应函数会挂起等待尝试与其他节点同步。

deepspeed 的配置

# ds_config.json
{
  "train_batch_size": 8,
  "gradient_accumulation_steps": 1,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.00015
    }
  },
  "fp16": {
    "enabled": true
  },
  "zero_optimization": true
}


# hostfile
# 如果没有显示指定，将默认查找 /job/hostfile；如果还没找到将默认使用本机GPU
worker-1 slots=4
worker-2 slots=4


deepspeed --hostfile=myhostfile <client_entry.py> <client args> \
  --deepspeed --deepspeed_config ds_config.json


# use only two nodes
deepspeed --num_nodes=2 \
	<client_entry.py> <client args> \
	--deepspeed --deepspeed_config ds_config.json


# except GPU 0 on node worker-2 and GPUs 0 and 1 on worker-3
deepspeed --exclude="worker-2:0@worker-3:0,1" \
	<client_entry.py> <client args> \
	--deepspeed --deepspeed_config ds_config.json


# use only GPUs 0 and 1 on worker-2
deepspeed --include="worker-2:0,1" \
	<client_entry.py> <client args> \
	--deepspeed --deepspeed_config ds_config.json

# .deepspeed_env 将文件中指定的环境变量同步到所有的worker上
# 在以下路径下依次查找：
#     1、当前代码执行路径
#     2、~/
#     3、环境变量 $DS_ENV_FILE 

NCCL_IB_DISABLE=1
NCCL_SOCKET_IFNAME=eth0

专业术语

NCCL：Nvidia Collective multi-GPU Communication Library的简称，它是一个实现多GPU的collective communication通信（all-gather, reduce, broadcast）库，Nvidia做了很多优化，以在PCIe、Nvlink、InfiniBand上实现较高的通信速度。

MPI：多个原语的消息传递接口，这一接口主要被用于多进程间的通信。

OpenMPI：针对 MPI 定义的抽象接口的一个具体实现。

Megatron-LM：一个基于 PyTorch 的分布式训练框架，用来训练基于Transformer的大型语言模型。Megatron-LM 综合应用了数据并行（Data Parallelism），模型并行（张量并行 Tensor Parallelism、流水线并行 Pipeline Parallelism）来复现 GPT-3。