istributed]

每个device上都拷贝一份完整模型每个device分别处理1个batch的一部分(如batch_size64,

2个device,

每device处理32个样本)梯度反向传播时每个设备上的梯度求和(求和才是一个完整batch所有样本的loss)汇入中心设备/参数服务器默认gpu0对模型进行梯度优化。

模型并行数据拷贝per

https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html

DDP:

nn.parallel.DistributedDataParallel

instead

DataLoader(datasetRandomDataset(input_size,

data_size),batch_sizebatch_size,

shuffleTrue)

并行化默认使用所有device加载数据torch.nn.DataParallel(module,

device_idsNone,

有哪些device_idsgpus默认None是使用全部device

output_deviceNone

model.to(cuda:0)不同的是tensor的to(device)是在device上生成一个拷贝不改变原来cpu上的tensor而model是直接将原model转移到gpu上。

step2加载到device0

设置中心设备(参数服务器)用于反向传播时的梯度汇总一般指定cuda:0

将模型从cpu放在gpu

model.to(device)step3forward前向传播

模型forward时将data_loader加载的一个batch的数据进行切分送入不同device的模型进行计算再将结果合并输出。

for

loss.backward()分别在每个device上计算loss的梯度average_gradients(model)将梯度聚合到中心设备/参数服务器(cuda:0)上进行梯度优化

criterion(output,

对各个设备上的梯度进行求和average_gradients(model)#

使用原始设备模型进行梯度优化optimizer.step()2.

分布式数据并行DDP(nn.parallel.DistributedDataParallel)

multiple

始终在卡间维持着模型参数和优化器状态的同步一致性在整个训练过程中

Data

ring环其同步过程不需要等待所有的卡都计算完一轮梯度经过这个同步过程之后所有的卡的

models/optimizers

分别计算损失forward和梯度backward蓝线梯度的聚合到中心device/参数服务器上(gpu0)绿线模型/优化器参数的更新及广播broadcast

其实参数服务器可以是一个GPU0也可以是CPU也可以是所有GPU

DDP基本概念

表示包含所有进程的组(所有gpu的集合)。

每个进程通常对应一个

GPU

中的进程可以相互通信这使得使用分布式数据并行Distributed

Data

表示分布式训练环境中的总进程数/gpu数。

每个进程都会被分配一个唯一的标识符rank从

rank进程标识符

是分配给world中每个进程的唯一标识符用于标识每个进程在分布式训练中的角色。

local

rank是分配个单个node中每个进程的标识符world中可能有多个node。

node节点

可以理解为一个服务器代表着物理设备上的一个实体。

在多机分布式训练中每台机器被视为一个节点节点之间需要进行通信。

例如如果有2

个node/server每个

one聚合过程reduce所有GPU(model和optiminizer状态)汇聚到参数服务器

one

torchrun运行分布式train.py脚本nproc-per-node设置每个node服务器上的gpu个数(一般是1个服务器)ddp_gpus_torchrun.py脚本名称--max_epochs

--batch_size

如单卡情况下data_len1024batch_size32则一个gpu的step1024/3232

多卡情况下2个gpudata_len1024batch_size32则每个gpu的step(1024/32)/232/216

import

destroy_process_groupstep2ddp_setup函数

这个函数用于设置分布式训练的环境。

它调用了init_process_group函数来初始化进程组使用的通信backend后端是ncclNVIDIA

Collective

Library然后使用torch.cuda.set_device函数根据环境变量设置当前进程使用的GPU设备。

def

间通信init_process_group(backendnccl)torch.cuda.set_device(int(os.environ[LOCAL_RANK]))step3Trainer类

这个类定义了一个模型训练的封装器。

在初始化方法中它接收一个模型backend、一个训练数据加载器train_dataloader、一个优化器train_dataloader作为参数并将模型移动到GPU上然后使用DistributedDataParallel对模型进行包装以实现数据并行。

(model先放cuda再DDP封装)

_run_batch方法实现了一次批量的训练过程包括前向传播、计算损失、反向传播和更新参数。

_run_epoch方法用于遍历整个训练集进行训练self.train_dataloader.sampler.set_epoch(epoch)是用于设置数据加载器的epoch以保证每个GPU在每个epoch开始时加载的数据都是不同的。

train方法则用于控制训练的总体流程。

class

int(os.environ[LOCAL_RANK])self.model

model.to(self.gpu_id)self.train_dataloader

optimizerself.model

ys):self.optimizer.zero_grad()output

self.model(xs)loss

ys)loss.backward()self.optimizer.step()def

_run_epoch(self,

len(next(iter(self.train_dataloader))[0])print(f[GPU:

{self.gpu_id}]

{len(self.train_dataloader)})self.train_dataloader.sampler.set_epoch(epoch)for

xs,

ys.to(self.gpu_id)self._run_batch(xs,

ys)def

range(max_epoch):self._run_epoch(epoch)step4MyTrainDataset类

这个类定义了一个自定义的训练数据集。

在初始化方法中它接收一个大小参数并生成一组随机的数据样本。

__len__方法返回数据集的大小__getitem__方法用于获取指定索引处的数据样本。

class

这个函数是程序的主函数。

在函数内部首先调用了ddp_setup函数来设置分布式训练的环境。

然后创建了一个自定义的训练数据集和相应的数据加载器以及一个线性模型和一个优化器。

DistributedSampler是PyTorch提供的一个分布式采样器用于确保每个进程加载的数据都是不同的且顺序随机。

sampler对象被传入训练数据集的构造函数可以通过数据加载器如torch.utils.data.DataLoader的sampler参数指定。

在每个进程中DistributedSampler会根据进程ID和进程数量将整个训练数据集划分成多个部分并为每个进程提供其应加载的数据索引。

这样在分布式训练过程中每个进程只会加载自己负责的数据部分避免了数据重复加载。

接下来创建了一个Trainer对象并调用其train方法进行模型训练。

最后调用destroy_process_group函数销毁进程组。

def

MyTrainDataset(2048)train_dataloader

pin_memoryTrue,

之间)samplerDistributedSampler(train_dataset))model

torch.nn.Linear(20,

torch.optim.SGD(model.parameters(),

lr1e-3)trainer

train_dataloadertrain_dataloader)trainer.train(max_epochs)destroy_process_group()step6解析命令行参数并运行主函数

在这个步骤中首先使用argparse模块解析命令行参数包括最大训练周期数max_epochs和批量大小batch_size。

然后调用main函数并将解析后的参数传递给它进行模型训练。

__name__

argparse.ArgumentParser(descriptionsimple

distributed

job)parser.add_argument(--max_epochs,

typeint,

model)parser.add_argument(--batch_size,

default32,

torch.cuda.device_count()main(args.max_epochs,

args.batch_size)3.

模型并行单卡放不下一份模型将一份大模型不同的层切分到不同的卡上forward时串行执行

Huggingface实现

device_mapHuggingface支持自动实现模型并行

balanced,

如下如果有两种gpudevice_mapauto使模型的layers的parameter分别加载到两张gpu上(各一半)

from

LlamaForCausalLM.from_pretrained(decapoda-research/llama-7b-hf,load_in_8bitTrue,device_mapauto,

for

enumerate(model.named_parameters()):

print(f{i},

pytorch模拟模型并行原理分别用to(device)将不同的层加载到不同的gpu上forward时将parameter也加载到对应gpu。

import

self.net2(self.relu(x.to(cuda:1))return

x进行一个batch的train每个batch_size20样本5分类

model

Train的数据4部分组成model模型参数、backward的梯度gradient、optimizer优化器参数、forward的数据tensor

Deepspeed、ZeRO技术方案分发Partitioning(按gpu数量N等分数据)、卸载Offload(不用的数据放入CPU)、模型并行Pipeline(模型参数按层切分到不同gpu上)

step1deepspeed初始化

deepspeed.init_distributed()加载参数local_rank

def

argparse.ArgumentParser(descriptiondeepspeed

training

script.)parser.add_argument(--local_rank,

typeint,

deepspeed.add_config_arguments(parser)args

deepspeed.initialize()封装model和dataset相当于将模型和数据集交给deepspeed进行托管engine就是deepspeed封装后的model其他返回值同样都是deepspeed封装过的。

(其中optimizer和lr_scheduler

后面是用不到的)我们只需要模型engine和数据加载器training_dataloader。

还要传入一个deepspeed的配置文件deepspeed_config。

init

deepspeed.initialize(argsargs,modelmodel,model_parametersmodel.parameters(),training_datads,configdeepspeed_config,

load

engine.load_checkpoint(./data/checkpoints/MyClassifier/)step3训练

在使用DeepSpeed进行分布式训练时通常不需要手动调用optimizer.zero_grad()来清零梯度。

DeepSpeed会自动处理梯度累积和梯度清零的操作无需手动调用zero_grad()。

当使用DeepSpeed进行分布式训练时一般会在engine.backward(loss)之后调用engine.step()来执行梯度更新操作。

在engine.step()中DeepSpeed会执行优化器的step()方法来更新模型参数并在必要的时候自动清零梯度以便进行下一轮的反向传播。

engine.train()for

enumerate(training_dataloader):step

1data

不需要梯度清零optimizer.zero_grad()outputs

engine(data)

)engine.backward(loss)engine.step()单机节点node多卡gpu运行

deepspeed

$PWD/deepspeed_config.jsondeepspeed_config.json

{train_micro_batch_size_per_gpu:

1,optimizer:

1000}},activation_checkpointing:

true,contiguous_memory_optimization:

null,synchronize_checkpoint_boundary:

false,profile: