模子的多GPU并行训练，DDP （附代码）

DDP全称是DistributedDataParallel, 在torch.nn.parallel内里。
本日总结一下用DDP举行多GPU并行训练的方法，
内容来自build gpt2加上自己的增补。
本文有些长，必要耐烦阅读，渴望整理的这些对各人有所资助。
假如你有多块GPU，就可以充实利用它们。
DDP会创建多个process（进程，不是线程哦）, 每个process分配一个GPU，这些process同时运行，于是就到达了多个GPU同时训练的目标。
之前写的训练代码你就可以想象它们同时在被多个process同时运行，
但每个process处置惩罚的是差别部门的数据（数据固然要差别了，否则不成了重复训练了嘛）。
这时间你还必要注意一个标题就是梯度的盘算，好比如今有8个GPU，它们同时盘算了差别部门数据的梯度，
而你的训练目标是不是全部数据汇总的梯度，你就必要把8个GPU盘算的梯度加起来求匀称。
DDP变量

那么你该怎样区分哪个呆板上的哪个进程呢，有下面几个变量帮你区分：
WORLD_SIZE: 总进程数
LOCAL_RANK: 当进步程在本呆板上的局部排名,从0开始计数
RANK: 当进步程在全部进程中的全局排名,从0到 WORLD_SIZE-1
举个例子：
在2台呆板上运行4个进程时:
呆板1:

• 进程0: RANK=0, LOCAL_RANK=0
  
• 进程1: RANK=1, LOCAL_RANK=1
  

呆板2:

• 进程2: RANK=2, LOCAL_RANK=0
  
• 进程3: RANK=3, LOCAL_RANK=1
  

全部进程的 WORLD_SIZE=4
假如只有一个呆板，那local_rank和rank的值是雷同的。
启动train.py

可以手动设置分布式情况

1. import torch.distributed as dist
3. # 手动设置环境变量
4. os.environ['RANK'] = '0'
5. os.environ['WORLD_SIZE'] = '4'
6. os.environ['LOCAL_RANK'] = '0'
7. os.environ['MASTER_ADDR'] = 'localhost'
8. os.environ['MASTER_PORT'] = '29500'
10. # 初始化进程组
11. dist.init_process_group(backend='nccl')

复制代码

更方便的是用torchrun下令，它可以自动做如下变乱：

• 自动设置情况变量(RANK, LOCAL_RANK, WORLD_SIZE等)
  
• 为每个GPU创建一个进程
  
• 设置进程组通讯,初始化分布式情况
  
• 在每个进程中运行train.py
  

1. # 单机4卡
2. torchrun --standalone --nproc_per_node=4 train.py
3. torchrun --nproc_per_node=4 train.py #单机训练，--standalone参数会被自动设置为默认值
5. # 多机训练
6. torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr="192.168.1.1" --master_port=29500 train.py

复制代码

torchrun会根据GPU序号设置RANK。比方在单机4卡时：
GPU 0: RANK=0
GPU 1: RANK=1
GPU 2: RANK=2
GPU 3: RANK=3
在多机时，RANK会跨呆板累加。比方2机4卡：
呆板1: RANK=0,1,2,3
呆板2: RANK=4,5,6,7
那你说，我的train.py偶尔会用单个GPU，偶尔会用多个，我想让它具有通用性，不必要每次都修改代码。
可以下面如许做，反面都说的是单个呆板的情况。
GPU之间的通讯一样平常用’nccl’,
NCCL (NVIDIA Collective Communications Library) 是NVIDIA开发的GPU通讯库，专门优化了GPU之间的通讯服从。在PyTorch分布式训练中，它是GPU训练的默认后端选项。
初始化阶段

1. from torch.distributed import init_process_group
2.     #如果用torchrun启动，RANK的值在0~WORLD_SIZE-1,WORLD_SIZE为进程数，即GPU总数
3.     ddp = int(os.environ.get('RANK', -1)) != -1
4.     if ddp:
5.         assert torch.cuda.is_available() #DDP中CUDA是必须的
6.         init_process_group(backend='nccl')
7.         ddp_rank = int(os.environ['RANK'])
8.         ddp_local_rank = int(os.environ['LOCAL_RANK'])
9.         ddp_world_size = int(os.environ['WORLD_SIZE'])
10.         device = f'cuda:{ddp_local_rank}'
11.         torch.cuda.set_device(device)
12.         master_process = ddp_rank==0 #防止log重复输出,指定master process输出
13.     else:
14.         ddp_rank = 0
15.         ddp_local_rank = 0
16.         ddp_world_size = 1
17.         master_process = True
18.         device = 'cpu'
19.         if torch.cuda.is_available():
20.             device = 'cuda'
21.         elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
22.             device = 'mps'  #较新的MAC用
23.         print(f'using device: {device}')

复制代码

在代码内里加上这么一句话，就可以看到代码正在被哪个GPU实验：

1. print("I am GPU: ", ddp_rank)

复制代码

打印出来的是随机序次，并不能包管先后，且它们是并行发生的。
batch size凌驾GPU内存的情况

好比gpt2, paper中的batch size为0.5M个token, 一样平常取2的N次方会利于GPU加快盘算，
以是令total_batch_size=219=524288, 这么大的batch size一样平常GPU很难容纳，如安在有限的GPU内存上实现云云大的batch size?
可以每次处置惩罚小批量的batch, 分多次处置惩罚。
好比gpt的sequence length设为T=128，batch size B=32, 这是一次处置惩罚的batch量（设为你的GPU能蒙受的batch size），
统共要处置惩罚grad_accum_steps=total_batch_size // (B * T)次。
这是在分解使命，统共batch, 但是我的GPU容纳不下，于是分批处置惩罚，末了再合起来。

1. total_batch_size = 524288
2. B = 32 #小的batch size(GPU能承受的)
3. T = 128
4. assert total_batch_size % (B*T) == 0
5. grad_accum_steps = total_batch_size // (B*T)  #一共需要累加多少轮
6. print(f"total desired batch size: {total_batch_size}")
7. print(f"=> calculated grad accumulation steps: {grad_accum_steps}")

复制代码

于是到了训练的时间，我们要在step内里再分解为micro_step,
我们的目标是末了把每个iteration内里的全部micro_steps的梯度累加起来，才算一个大batch的梯度，累加完了之后再更新optimizer, 梯度归0。
用到了这个特性：
假如你不消optimizer.zero_grad()，梯度梯度是累加的，以是每个iteration开始时用一次, micro_step内里（小batch竣事时）不消。optimizer.step()是用梯度来更新变量，以是也是每个iteration竣事时用一次，micro_step内里不消。
不外loss不会自动累加，你必要设一个变量用来累加loss.
在累加loss的时间，你要注意一个标题，
举例来分析，假如如今batch size=B, 每个micro_step中的batch size=b. 一共累加B//b=N次。
一样平常loss内里会对这个batch内里的loss求匀称，会有mean的运算。
那么每一个micro_step内里的loss就是1/b * 小batch的loss和，
末了我们把全部micro_step的loss累加起来之后loss就是1/b * 大batch的loss.
发现标题了吗，原来大batch的loss应该是1/B * 大batch的loss, 如今变成了1/b, 这个求匀称的基数不对。
以是在全部micro_steps的loss累加时，要别的再除grad_accum_steps.

1.     for step in range(max_steps):
2.         t0 = time.time()
3.         optimizer.zero_grad()
4.         loss_accum = 0.0 #用于print loss
5.         #在这些步里面累加梯度(用loss.backward的累加梯度效果)，模拟一个0.5M的batch size的效果
6.         for micro_step in range(grad_accum_steps):
7.             x, y = train_loader.next_batch()
8.             x = x.to(device)
9.             y = y.to(device)
11.             #一般的loss里面会用mean来求这个batch里面loss的平均，所以每个batch会除以batch_size
12.             #那么大batch的loss加起来，也只相当于除以了小batch的batch_size,并不是大batch_size
13.             #所以在累加的loss基础上还要除以（大batch_size / 小batch_size), 也就是grad_accum_steps
14.             loss /= grad_accum_steps
15.             loss_accum += loss.detach() #detach():变成叶子节点，梯度累加，loss可没有累加，累加loss要自己算
16.             loss.backward()  # 更新梯度，会累积，所以每次optimizer要设置zero_grad,它不会自动清零
17.         #把梯度的norm限制在1，什么是norm,就是一个向量的L2 norm
18.         norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) #根据gpt3 paper
19.         optimizer.step() #更新参数

复制代码

如今用了DDP，就有多个GPU一起来处置惩罚这个batch, 每个GPU处置惩罚的grad_accum_steps变为：

1. grad_accum_steps=total_batch_size // (B * T * ddp_world_size)

复制代码

分别差别数据地区

这部门讨论怎样让每个process访问差别的数据部门，也就是说每个GPU处置惩罚差别的数据。
以是我们必要把数据分别，让差别的GPU访问差别的部门。
怎样辨认差别的GPU，这就用到了前面的RANK.
你就必要把ddp_rank和ddp_world_size传给dataloader.

1. class DataLoaderLite:
2.     def __init__(self, B, T, process_rank, num_processes):
3.         self.B = B
4.         self.T = T
5.         self.process_rank = process_rank  #ddp_rank
6.         self.num_processes = num_processes  #ddp_world_size
8.         with open('./data/input.txt','r') as f:
9.             text = f.read()
10.         enc = tiktoken.get_encoding('gpt2')
11.         tokens = enc.encode(text)
12.         self.tokens = torch.tensor(tokens)
13.         print(f"loaded {len(self.tokens)} tokens")
14.         print(f"1 epoch = {len(self.tokens) // (B*T)} batches")
16.                 #每个GPU读不同的数据
17.         self.current_position = self.B * self.T * self.process_rank
19.     def next_batch(self):
20.         B, T = self.B, self.T
21.         buf = self.tokens[self.current_position:self.current_position+B*T+1]
22.         x = buf[:-1].view(B,T)
23.         y = buf[1:].view(B,T)
24.         self.current_position += B*T * self.num_processes  #每个GPU处理的数据，有点类似CNN中的batch
25.                
26.                 #数据应该不会小到一个循环之内self.B * self.T * self.process_rank的时候就超出边界了
27.         if self.current_position + (B*T*self.num_processes+1) >= len(self.tokens):
28.             self.current_position = B * T * self.process_rank
29.         return x, y
31. train_loader = DataLoaderLite(B=B, T=T, process_rank=ddp_rank, num_processes=ddp_world_size)

复制代码

怎样处置惩罚model

由于一个代码被N个GPU同时实验，这里假设有8个GPU并行。
那么model就会被创建8次，要包管这8个model是划一的，就必要固定seed.
还必要把model装进DDP container.
这个可以参考DistributedDataParallel官方文档，device_ids这个参数指的是ddp_local_rank, 而不是ddp_rank.
DDP 如许包装 model 的紧张缘故原由是为了实现高效的分布式训练，
device_ids=[ddp_local_rank] 的作用：
将模子绑定到指定的 GPU 上
ddp_local_rank 通常是 0,1,2,3… 如许的局部 GPU 编号
确保每个进程只利用分配给它的谁人 GPU
DDP 包装后会实现：
自动梯度同步：差别 GPU 上的梯度会自动举行匀称
参数广播：确保全部 GPU 上的模子参数保持划一
通讯优化：利用高效的 ring-allreduce 算法举行梯度同步
包装了之后model就变成了DDP model.
而必要调用model内里的method时，必要的是raw model.

1. #这里model会被创建GPU个数次，比如8个GPU，就会创建8个model, torch.compile也会发生8次
2. #固定seed, 让每个model计算一致
3. torch.manual_seed(1337)
4. if torch.cuda.is_avaliable():
5.     torch.cuda.manual_seed(1337)
6. model = GPT(GPTConfig(vocab_size=50534))
7. model.to(device)
8. model = torch.compile(model)
9. if ddp:
10.         model = DDP(model, device_ids=[ddp_local_rank])
11. raw_model = model.module if ddp else model
13. #当需要调用model里面的method时，要用raw_model
14. optimizer = raw_model.configure_optimizers(weight_decay=0.1, ...)

复制代码

同步梯度盘算

forward过程每个GPU都是一样的，
backward涉及到每个GPU都在盘算梯度，以是在backward竣事的时间，管帐算全部GPU盘算的梯度的匀称，
同步这个匀称梯度到每个GPU上。实在是调用了all_reduce
现实上并不是竣事的时间才如许做，而是在backward正在举行的时间就在同步做这些变乱。
在loss.backward()时，就会发生通讯，并同步梯度信息到各GPU。
前面提到了batch size过大时，会分解成每次处置惩罚小的batch.
可以看到每个小batch盘算时都有loss.backward(), 导致多次通讯，同步梯度，影响服从，
我们如今不想这么做，想让一轮batch盘算竣事时再同步。

1.     for step in range(max_steps):
2.         t0 = time.time()
3.         optimizer.zero_grad()
4.         loss_accum = 0.0
5.       
6.         for micro_step in range(grad_accum_steps):
7.             x, y = train_loader.next_batch()
8.             x = x.to(device)
9.             y = y.to(device)
11.             loss /= grad_accum_steps
12.             loss_accum += loss.detach()
13.             loss.backward()
14.             
15.         norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
16.         optimizer.step() #更新参数

复制代码

一样平常是用no_sync(), 看上去比力复杂

而当你看了no_sync()的源码，你会发现起作用的是一个flag

以是直接用这个flag, 会比力简便，风险就是未来版本升级了这个flag大概就不存在了-_-
必要在loss.backward()前面设置是否必要通讯同步梯度。
别的，梯度已经求了全部GPU的匀称并同步，但是loss并没有求各GPU的匀称，
每个GPU只能瞥见它自己的loss.
这时必要调用all_reduce, 前面提到了同步梯度时也是all_reduce.
在完成了一整个大batch的盘算后（从for loop出来），调用all_reduce.
all_reduce做的变乱是对每个GPU上的loss_accum, 盘算全部GPU（rank）上的匀称值，然后同步到每个GPU。
当你print(loss_accum)时，全部GPU上的数值应该是划一的。

1. import torch.distributed as dist
3.     for step in range(max_steps):
4.         t0 = time.time()
5.         optimizer.zero_grad()
6.         loss_accum = 0.0
7.         
8.         for micro_step in range(grad_accum_steps):
9.             x, y = train_loader.next_batch()
10.             x = x.to(device)
11.             y = y.to(device)
13.             loss /= grad_accum_steps
14.             loss_accum += loss.detach()
15.             if ddp:
16.                     #到最后一轮时才需要同步梯度
17.                     model.require_backward_grad_sync = (micro_step == grad_accum_steps-1)
18.             loss.backward()
19.         #从for loop中出来后所有GPU里面的梯度都同步了  
20.         if ddp:
21.                 dist.all_reduce(loss_accum, op=dist.ReduceOp.AVG) #average
22.         norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
23.         optimizer.step()

复制代码

输出log

应该意识到如今每个GPU都在运行这个代码，以是你会输出重复的log,
设置一个master, 只输出一个log，
master可以在前面初始化的时间设置
竣事时要调用destroy_process_group

1. #master_process = ddp_rank==0 #前面初始化的部分已经设置
3. for step in range(max_steps):
4.    ...
5. token_processed = train_loader.B * train_loader.T * grad_accum_steps * ddp_world_size
6. token_per_sec = token_processed / dt
7. if master_process:
8.         print(f"step {step:4d} | loss: {loss.item():.6f} | norm: {norm:.4f} | dt: {dt*1000:.2f}ms, token_per_sec: {token_per_sec:.2f}ms")
9.        
10. if ddp:
11.         destroy_process_group()

复制代码

你会发现多GPU并行时的loss大概会和单GPU差别，这是由于data_loader中读取数据的方式不一样，
从前是一个一个的batch, 超出边界了就从0开始，如今把N个GPU一次读的数据看成一个page, page凌驾边界时重置，这个地方会导致batch数据大概会出现差别。
假如想让loss雷同，必要调治total batch size, 让单个GPU重置的边界和多个GPU雷同。
拓展一下，假如数据很大，存储在差别的地方，也可通过DDP实现训练，设置差别rank的GPU读取差别地方的数据。
不外要注意以下几点：
确保数据分片巨细相近,克制负载不平衡
根据现实存储位置和网络带宽调解数据分配战略
大概必要处置惩罚数据分片边界的标题