训练的内存需求 本页的大部分内容将集中在减少训练神经网络所需的内存的技术上。因此，为了更好地理解和欣赏这些技术，我将首先分解训练过程中内存消耗的所有方式。
根据Sohoni等人的说法，我将使用以下术语来描述训练期间内存消耗的来源（Sohoni等人只描述了三个来源，但我将梯度内存和优化器内存分开，因为DeepSpeed中也做了同样的区分，稍后将展示）：
优化器内存 ：存储优化器所需的任何额外状态所需的内存。 激活内存 ：存储在前向传递期间计算的中间值所需的内存。 测量四种内存消耗源 让我们现在看看如何测量这四种来源消耗的内存量，从一个简单的全连接神经网络开始：
model = torch.nn.Sequential( torch.nn.Linear( 512 , 1024 ), torch.nn.Linear( 1024 , 1024 ), torch.nn.Linear( 1024 , 1024 ), torch.nn.Linear( 1024 , 512 ) ) 要获得总的 模型内存 ，我们计算模型中的参数数量，并乘以每个参数用来表示的字节数（在这种情况下是4个字节，因为每个权重是一个32位浮点数）：
def get_model_memory (model: torch.nn.Module) : ''' 返回给定模型的内存使用情况 ''' total_memory = 0 for param in model.parameters(): total_memory += param.numel() * param.element_size() return total_memory print( '模型内存：{:,} 字节' .format(get_model_memory(model))) 模型内存：12,597,248 字节 梯度内存 与模型内存相同，因为在后向传递期间我们需要为每个权重存储梯度。
消耗的总 优化器内存 取决于训练期间使用的优化器类型。下表显示了PyTorch中一些流行的优化器所需的优化器内存，其中M_model代表模型内存[2]：
SGD SGD Adam Adam优化器需要我们为每个梯度维护第一和第二矩（每个4字节）。
最后，要计算总的 激活内存 ，我们需要计算在前向传递期间计算的隐藏层输出。有几种不同的方法可以做到这一点；这里，我们将递归地注入PyTorch前向Hook到每个子模块中以捕获该子模块的输出大小。
首先，我们定义了Hook以及一个简单的辅助类 ActivationCounter ，用于跟踪模型隐藏层输出的总大小：
class ActivationCounter : def __init__ (self) : self.activation_bytes = 0 def add_activations (self, tensor) : self.activation_bytes += tensor.numel() * tensor.element_size() def activation_counter_hook (counter: ActivationCounter) : def hook (self, input, output) : counter.add_activations(output.data) return hook 接下来，我们定义一个函数，它将递归地将Hook注入到模型的每个子模块中：
def register_hooks_recursive (model, counter: ActivationCounter) : for module in model.children(): module.register_forward_hook(activation_counter_hook(counter)) register_hooks_recursive(module, counter) activation_counter = ActivationCounter() register_hooks_recursive(model, activation_counter) 现在要找到激活内存消耗的总量，我们执行模型的一个前向传递，然后打印 activation_counter.activation_bytes 的值：
inputs = torch.randn( 4 , 512 ) outputs = model(inputs) # 因为Hook只捕获层输出，我们需要单独添加原始输入张量的大小 activation_counter.add_activations(inputs) print( '激活内存：{:,} 字节' .format( activation_counter.activation_bytes )) 激活内存：65,536 字节 计算密集神经网络模型内存和激活内存的完整脚本可以在GitHub仓库中找到：estimate_nn_memory.py
Transformer模型的内存需求 由于Rajbhandari, Samyam等人。（在他们的论文中介绍了DeepSpeed实现的许多训练优化）将实验集中在Transformer模型上，我想更好地了解这类模型的内存需求。
使用上述相同的方法，我计算了一个基本Transformer模型（详见附录A）的模型和激活内存，作为验证我估计Transformer模型内存的公式，以及Korthikanti, Vijay Anand等人提供的Transformer激活内存公式（都显示在下面）：
给定以下条件：
n：层数 h：隐藏层大小 m：注意力头数 b：批量大小 l：序列长度
分布式训练 DeepSpeed提供的优化主要适用于分布式训练环境（例如，当多个GPU用于训练单个模型时）。有三种常见的多GPU训练并行化范式： 数据并行 、 模型并行 和 流水线并行 。每个都在以下部分中简要解释。
数据并行 数据并行（DP）描述了一个分布式训练过程，整个模型在多个设备（例如GPU）上复制，每个设备并行地在不同的输入数据批次上执行训练。有多种不同的方法可以实现数据并行训练，但这里我只描述PyTorch DistributedDataParallel 模块使用的实现[4]。
每个设备接收不同的输入数据批次，并执行前向传递以计算该批次的损失。
当在每个设备上计算梯度时，这些梯度与所有其他设备交换。然后使用梯度的平均值更新每个设备上的模型权重，确保在下一个训练步骤开始时，所有设备都有相同的一组模型权重。
AllReduce 设备之间的梯度交换是使用称为AllReduce的集体通信算法执行的。AllReduce算法在跨多个设备分布的数据上执行归约操作。
下面的可视化展示了在3个设备上执行AllReduce操作以计算分布在3个设备上的值向量的总和：
PyTorch通过其 torch.distributed 模块提供了对包括AllReduce在内的几种集体通信算法的支持。下面的脚本演示了在3个GPU上进行AllReduce操作：
import os import torch import torch.distributed as dist import torch.multiprocessing as mp def create_process_group (rank, world_size) : os.environ[ 'MASTER_ADDR' ] = 'localhost' os.environ[ 'MASTER_PORT' ] = '12355' dist.init_process_group( backend= 'nccl' , world_size=world_size, rank=rank ) def all_reduce_example (rank, world_size) : create_process_group(rank, world_size) # 在每个设备上创建不同的张量 if rank == 0 : tensor = torch.tensor([ 1 , 2 , 3 ]).to(rank) elif rank == 1 : tensor = torch.tensor([ 10 , 20 , 30 ]).to(rank) elif rank == 2 : tensor = torch.tensor([ 4 , 5 , 6 ]).to(rank) print( 'Before AllReduce: Rank ' , rank, ' has data ' , tensor) dist.all_reduce(tensor, op=dist.ReduceOp.SUM) print( 'After AllReduce: Rank ' , rank, ' has data ' , tensor) if __name__ == '__main__' : device_count = 3 mp.spawn(all_reduce_example, args=(device_count,), nprocs=device_count) Before AllReduce: Rank 0 has data tensor([1, 2, 3], device= 'cuda:0' ) Before AllReduce: Rank 1 has data tensor([10, 20, 30], device= 'cuda:1' ) Before AllReduce: Rank 2 has data tensor([4, 5, 6], device= 'cuda:2' ) After AllReduce: Rank 2 has data tensor([15, 27, 39], device= 'cuda:2' ) After AllReduce: Rank 0 has data tensor([15, 27, 39], device= 'cuda:0' ) After AllReduce: Rank 1 has data tensor([15, 27, 39], device= 'cuda:1' ) 关于数据并行的一些说明 重要的是，数据并行不会减少训练所需的总内存，因为每个设备必须有足够的内存来执行整个模型的前向和后向传递。
因此，数据并行的主要好处是它加快了训练过程，因为每个设备只需要在总训练数据的一小部分上进行训练。然而，加速并不随着设备数量的增加而线性扩展，因为设备之间交换梯度引入了通信开销。例如，Li, Shen等人。2020年观察到，当在256个GPU上训练BERT模型时，扩展因子大约是128。
下图显示了使用1-8个GPU进行数据并行训练，在1000步训练中BERT模型所花费的时间：
模型并行 模型并行（MP） 描述了一个分布式训练过程，模型被分割在多个设备上，以至于每个设备只包含模型权重的一部分。前向传递在每个设备上顺序执行，一个设备的输出成为下一个设备的输入。
模型并行在训练非常大的模型时至关重要，这些模型无法适应单个设备。
模型并行训练的一个主要缺点是，前向和后向传递的顺序性质导致每个设备上的空闲时间。例如，在上图中，GPU 0在完成输入批次的前向传递后，它在GPU 1和GPU 2执行它们的前向和后向传递时处于空闲状态。
通过使用PyTorch前向和后向Hook，我们可以测量每个GPU在训练期间的空闲时间：
def idle_time_hook(self, device, forward=True, entering=True): '' “创建一个PyTorch Hook，记录设备的空闲时间。 '' ' def hook(*args, **kwargs): current_timestamp = time.time() last_timestamp = self.previous_timestamp.get(device, None) message = ' {} {} pass on device {} '.format( ' Entering ' if entering else ' Finished ', ' forward ' if forward else ' backward ', device ) if entering and last_timestamp is not None: idle_time_ms = (current_timestamp - last_timestamp) * 1000 self.device_idle_time[device] = ( self.device_idle_time[device][0] + idle_time_ms, self.device_idle_time[device][1] + 1 ) message += f' . Idle time: {idle_time_ms:.2f}ms ' self.previous_timestamp[device] = current_timestamp self.log(message) return hook 2023-06-04 00:20:36.786470 - Entering forward pass on device 0. 2023-06-04 00:20:36.789441 - Finished forward pass on device 0 2023-06-04 00:20:36.790807 - Entering forward pass on device 1. 2023-06-04 00:20:36.793664 - Finished forward pass on device 1 2023-06-04 00:20:36.794982 - Entering forward pass on device 2. 2023-06-04 00:20:36.796839 - Finished forward pass on device 2 2023-06-04 00:20:36.798351 - Entering forward pass on device 3. 2023-06-04 00:20:36.799569 - Finished forward pass on device 3 2023-06-04 00:20:37.911536 - Entering backward pass on device 3. Idle time: 1111.96ms 2023-06-04 00:20:37.913496 - Finished backward pass on device 3 2023-06-04 00:20:37.914812 - Entering backward pass on device 2. Idle time: 1117.97ms 2023-06-04 00:20:37.916677 - Finished backward pass on device 2 2023-06-04 00:20:37.918056 - Entering backward pass on device 1. Idle time: 1124.39ms 2023-06-04 00:20:37.920743 - Finished backward pass on device 1 2023-06-04 00:20:37.922119 - Entering backward pass on device 0. Idle time: 1132.67ms 2023-06-04 00:20:37.923938 - Finished backward pass on device 0 下面的时间线图示出了在三个GPU上执行模型并行训练时每个GPU的前向和后向执行。每个时间线中的灰色空间代表该GPU空闲的时间。
流水线并行 流水线并行（PP） 是模型并行的一种变体，它通过将每个输入数据批次分割成多个较小的“微批次”来减少设备空闲时间[6]。模型参数仅在每个微批次被整个模型处理后更新，这意味着每个设备可以在其他设备仍在处理前一个微批次时开始处理下一个微批次。
支持流水线并行的功能已经内置在PyTorch中，通过 torch.distributed.pipeline.sync.Pipe 类。然而，这个类有两个主要限制，即（1）它只在模型被实现为 torch.nn.Sequential 模块时工作，以及（2）它要求每个模块的输入和输出是单个张量或张量元组[7]。
由于这些限制，我不得不修改HuggingFace Transformers库中的BERT模型实现以支持流水线并行。修改后的模型可以在这里找到：distributed-training-and-deepspeed/model/bert_mp.py
使用这个定制的BERT实现，我们可以通过首先将模型转换为 torch.nn.Sequential 模块，然后将其包装在 torch.distributed.pipeline.sync.Pipe 对象中来启用流水线并行：
def to_pipeline (self, chunks) : '' “将模型转换为流水线并行。 ''' rpc.init_rpc( name='worker', rank=0, world_size=1, rpc_backend_options=rpc.TensorPipeRpcBackendOptions( init_method='file://{}'.format(tempfile.NamedTemporaryFile().name) ) ) sequential = torch.nn.Sequential( self.embeddings, *self.encoders, self.head ) return Pipe(sequential, chunks=chunks) 比较模型并行和流水线并行 为了比较模型并行和流水线并行的性能，我使用我的定制BERT实现执行了一些训练运行。下图显示了在2、4、6和8个GPU上使用模型并行和流水线并行执行250个训练步骤所花费的时间。对于这些测试，批量大小设置为16，对于流水线并行，微批次的数量设置为4。
DeepSpeed 有了对分布式训练技术的基本情况的了解，我们现在可以开始看看DeepSpeed。
DeepSpeed是微软创建的一个开源深度学习优化工具库，涵盖了模型训练、推理和模型压缩领域。这里，我将只关注与训练相关的优化，并将推理和压缩的主题留给未来的探索。
如前所述，我们可以通过在多个设备上复制整个模型（数据并行）和/或分割模型并将其不同部分存储在不同设备上（模型并行/流水线并行）来进行分布式训练。通常，数据并行在计算效率上比模型并行更高；然而，在模型太大而无法适应单个设备的可用内存的情况下，模型并行是必要的。
DeepSpeed提供的培训优化的核心是Zero Redundancy Optimizer（ZeRO），这是一套减少分布式模型训练所需内存的技术。
ZeRO 数据并行在设备之间引入了显著的内存冗余。例如，考虑在8个GPU上训练一个10亿参数模型。存储模型参数本身所需的内存（假设是32位浮点数）约为3.7GB。当使用数据并行时，每个设备必须存储整个模型的副本，这意味着在所有设备上总共使用了约29.8GB的内存。换句话说，26.1GB的设备内存被冗余的模型参数占据。激活和优化器内存也在设备之间类似地冗余。
ZeRO-DP通过以下方式解决这种冗余：
当使用 deepspeed 库时，我们通过指定ZeRO Stage 来选择应用哪些内存优化。可用的 Stage 包括：
在接下来的部分中，我们将详细探讨每个 Stage 。所有示例脚本已在装有8个16GB NVIDIA V100 GPU的单台机器上进行了测试。
Stage 0 - 入门 DeepSpeed库像任何其他Python包一样安装：
当使用DeepSpeed执行数据并行训练时，我们不需要像使用PyTorch的 DistributedDataParallel 那样设置进程组或显式地启动多个进程。相反，我们通过调用 deepspeed.initialize 将模型包装在 DeepSpeedEngine 中，它在内部处理所有的分布式训练逻辑。
对于这些示例，我们将使用HuggingFace的预训练的、560M参数变体的BLOOM模型[9]。
# deepspeed_stage_0.py # # DeepSpeed自动设置LOCAL_RANK环境变量 # 为当前设备的索引。 rank = int(os.getenv( 'LOCAL_RANK' , '0' )) model = BloomForCausalLM .from_pretrained( 'bigscience/bloom-560m' ) deepspeed_config = { 'train_micro_batch_size_per_gpu' : 1 } model_engine, _, _, _ = deepspeed.initialize( model=model, model_parameters=model.parameters(), config=deepspeed_config ) print( f'Device {rank} - ZeRO Stage: {model_engine.zero_optimization_stage()} ' ) 要启动分布式训练作业，我们使用 deepspeed 命令行实用程序，它与 deepspeed Python包一起安装：
deepspeed deepspeed_stage_0.py Device 1 - ZeRO Stage: 0 Device 4 - ZeRO Stage: 0 Device 2 - ZeRO Stage: 0 Device 3 - ZeRO Stage: 0 Device 6 - ZeRO Stage: 0 Device 7 - ZeRO Stage: 0 Device 0 - ZeRO Stage: 0 Device 5 - ZeRO Stage: 0 在这里我们看到模型复制在所有四个设备上。要允许DeepSpeed创建和管理训练期间使用的优化器，我们向 deepspeed_config 对象添加适当的配置选项：
deepspeed_config = { 'train_micro_batch_size_per_gpu' : 1 , 'optimizer' : { 'type' : 'Adam' , 'params' : { 'lr' : 5e-5 } }, } model_engine, optimizer, _, _ = deepspeed.initialize( model=model, model_parameters=model.parameters(), config=deepspeed_config ) # 打印优化器状态信息 optimizer_state = optimizer.param_groups[ 0 ] print( f'Device {rank} - Optimizer: lr= {optimizer_state[ 'lr' ]} ; ' f'betas= {optimizer_state[ 'betas' ]} ; eps= {optimizer_state[ 'eps' ]} ; ' f'parameter count= {sum([torch.numel(p) for p in optimizer_state[ 'params' ]]):,} ' ) Device 1 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592 Device 4 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592 Device 2 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592 Device 3 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592 Device 6 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592 Device 7 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592 Device 0 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592 Device 5 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=559,214,592 训练 DeepSpeedEngine 与使用数据并行训练PyTorch模型非常相似。唯一的区别是 backward() 和 step() 方法直接在 DeepSpeedEngine 对象上调用，而不是在损失或优化器上：
# DeepSpeed自动设置WORLD_SIZE环境变量 # 为参与训练作业的设备数量。 world_size = int(os.getenv( 'WORLD_SIZE' , '1' )) tokenizer = AutoTokenizer.from_pretrained( 'bigscience/bloom-560m' ) collator = DataCollatorForLanguageModeling(tokenizer, mlm= False ) # 辅助函数加载wikitext数据集 # 实现可以在这里找到： # https://github.com/gnovack/distributed-training-and-deepspeed/blob/main/util.py train_dataset = load_wikitext(tokenizer, collator).select(range( 64 )) train_dataloader = DataLoader( train_dataset, batch_size= 1 , shuffle= False , sampler=DistributedSampler(train_dataset, num_replicas=world_size) ) for batch in train_dataloader: device = torch.device( 'cuda' , rank) input_ids = batch[ 'input_ids' ].to(device) labels = batch[ 'labels' ].to(device) outputs = model_engine(input_ids, labels=labels) model_engine.backward(outputs.loss) model_engine.step() Stage 1 - 优化器状态分割 ZeRO Stage-1将优化器状态（例如，当使用Adam时，一阶和二阶动量）在所有设备上进行分割，使得每个设备只包含这部分状态的一部分。
下图展示了使用传统数据并行（顶部）和ZeRO Stage-1（底部）训练BLOOM模型时每个设备的内存分配情况：
基于之前讨论的Transformer模型内存公式，我们可以估计存储一个560M参数模型的模型参数所需的内存约为2.24GB。当以32位模式使用Adam时，优化器状态将占用每个模型参数8字节的内存，因此优化器将需要4.48GB的内存。
假设我们以批量大小1和输入序列长度512来训练模型，我们可以使用上述公式估计激活所需的内存为1.75GB。
因此，使用传统数据并行训练BLOOM 560M模型所需的总内存大致为：
deepspeed 库包括几个实用工具函数，我们可以使用它们来分析训练期间的内存使用情况。例如，要测量当前和峰值内存使用情况，我们在训练循环内调用 memory_status() 函数：
from deepspeed.runtime.utils import memory_status for batch in train_dataloader: input_ids = batch[ 'input_ids' ].to(device) labels = batch[ 'labels' ].to(device) outputs = model_engine(input_ids, labels=labels) model_engine.backward(outputs.loss) model_engine.step() if rank == 0 : memory_status( 'Memory stats after training step' ) 如果我们现在使用 deepspeed CLI执行训练脚本，我们将在每个步骤后看到以下行打印出来：
RANK=0 MEMSTATS Memory stats after training step device=cuda:0 current alloc=6.8847GB (delta=0.0000GB max=10.8823GB) current cache=13.4453GB (delta=0.0000GB max=13.4453GB) 如所示，GPU的峰值内存使用量为10.88GB，非常接近我们估计的10.71GB。接下来，让我们看看使用ZeRO Stage-1可以节省多少内存。要启用ZeRO Stage-1，我们更新 deepspeed_config 对象以包含 zero_optimization 配置选项：
deepspeed_config = { 'train_micro_batch_size_per_gpu' : 1 , 'optimizer' : { 'type' : 'Adam' , 'params' : { 'lr' : 5e-5 } }, 'zero_optimization' : { 'stage' : 1 , } } model_engine, optimizer, _, _ = deepspeed.initialize( model=model, model_parameters=model.parameters(), config=deepspeed_config ) 现在执行训练脚本时，我们会注意到一些差异。首先，在每个设备上打印优化器参数数量时，显示的参数数量现在是69,901,824，或者是模型中总参数数量的18分之1：
Device 0 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824 Device 6 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824 Device 7 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824 Device 4 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824 Device 5 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824 Device 3 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824 Device 2 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824 Device 1 - Optimizer: lr=5e-05; betas=(0.9, 0.999); eps=1e-08; parameter count=69,901,824 其次，当打印每个训练步骤后的内存使用情况时，GPU的峰值内存使用量现在是7.48GB，比使用传统数据并行时测量的10.88GB少了3.4GB：
RANK=0 MEMSTATS Memory stats after training step: device=cuda:0 current alloc=3.4843GB (delta=0.0000GB max=7.4816GB) current cache=12.7695GB (delta=0.0000GB max=12.7695GB) 那么这额外的3.4GB来自哪里呢？回想一下，优化器状态占用了大约4.48GB。以前，当使用传统数据并行时，每个设备都持有所有4.48GB的优化器状态。现在，有了ZeRO Stage-1，每个设备只持有优化器状态的18分之1，即0.56GB。
理论上，这意味着使用ZeRO Stage-1时的峰值GPU内存应该是：
2.24 + 2.24 + 0.56 + 1.75 = 6.79GB
我们测量的7.48GB而不是6.79GB可能归因于用于通信更新模型权重的中间缓冲区的使用，但需要进一步调查以确认这一点。
Stage 2 - 梯度分割 ZeRO Stage-2通过在设备间分割优化器状态和梯度，更进一步。下图展示了使用ZeRO Stage-2训练时每个设备的内存分配情况：
Stage-2同样可以通过在 deepseed_config 中设置 stage 选项来启用：
deepspeed_config = { 'train_micro_batch_size_per_gpu' : 1 , 'optimizer' : { 'type' : 'Adam' , 'params' : { 'lr' : 5e-5 } }, 'zero_optimization' : { 'stage' : 2 } } 如果我们启用ZeRO Stage-2并执行训练脚本，我们将发现GPU的峰值内存使用量略有增加，达到8.22GB，相比之下，使用 Stage-1训练时为7.48GB：
RANK=0 MEMSTATS Memory stats after training step: device=cuda:0 current alloc=3.4843GB (delta=0.0000GB max=8.2187GB) current cache=12.7715GB (delta=0.0000GB max=12.7715GB) 为了理解GPU峰值内存意外增加的原因，我们需要了解一下ZeRO Stage-2的实现。在后向传递过程中，梯度被平均化并通过一系列Reduce操作放置在一个设备上。
虽然我们可以在计算每个梯度后执行Reduce操作，但通过分块减少梯度可以实现更好的性能。
这些块的大小由 deepspeed_config 中的 reduce_bucket_size 选项控制，默认值为5×10^8个元素，或者当每个元素是一个32位浮点数时为2GB。虽然这个默认值适用于更大的模型，但对于我们的560M模型来说太大了，它的梯度总共只占用了大约2.24GB。
如果我们将 reduce_bucket_size 减少到5×10^6个元素，我们可以将峰值内存使用量降低到6.37GB，与 Stage-1相比额外减少了1.11GB：
RANK=0 MEMSTATS Memory stats after training step: device=cuda:0 current alloc=3.4843GB (delta=0.0000GB max=6.3750GB) current cache=10.9277GB (delta=0.0000GB max=10.9277GB) 值得考虑的是6.37GB的峰值内存是否符合我们的预期。鉴于每个设备持有优化器状态的18分之1和梯度的18分之1，我们应该预期峰值内存在：
2.24 + 2.24/8 + 4.48/8 + 1.75 = 4.83GB
这个理论估计假设后向传递中每个操作产生的梯度小于桶大小。这篇文章探讨了在多个GPU上分布式训练神经网络的技术，并检验了微软的DeepSpeed库提供的多种分布式训练优化。
阶段3 - 参数分割 ZeRO第三阶段在第二阶段的基础上，进一步分割模型参数以及优化器状态和梯度。当我们使用ZeRO第三阶段训练BLOOM 560M模型时，我们会注意到每个训练步骤后分配的内存从第二阶段的3.48GB减少到1.94GB。然而，GPU的峰值增加到7.48GB。
RANK=0 MEMSTATS Memory stats after training step: device=cuda:0 current alloc=1.9410GB (delta=0.0000GB max=7.4082GB) current cache=10.7129GB (delta=0.0000GB max=10.7129GB) 目前，GPU峰值内存使用量增加的原因尚不清楚，但我在DeepSpeed GitHub仓库中有一个开放的问题，以更好地理解这种行为：DeepSpeed/issues/3734
通信开销 虽然ZeRO可以减少训练所需的总内存量，但它确实引入了额外的通信开销，因为优化器状态、梯度和模型参数必须在设备之间频繁交换。
DeepSpeed允许我们通过其通信日志设置来测量通信操作所花费的时间。通信日志通过在 deepspeed_config 中添加 comms_logger 部分来启用：
deepspeed_config = { 'train_micro_batch_size_per_gpu' : 1 , 'optimizer' : { 'type' : 'Adam' , 'params' : { 'lr' : 5e-5 } }, 'zero_optimization' : { 'stage' : 0 }, 'comms_logger' : { 'enabled' : True , 'verbose' : False , 'prof_all' : True , 'debug' : False } } 现在，执行训练脚本时，每个通信操作（AllReduce、AllGather、Reduce Scatter等）将被捕获和测量。要查看这些测量结果，我们调用 deepspeed.comm.log_summary() 函数：
## 训练循环后 import deepspeed.comm as dist dist.log_summary() Comm. Op Message Size Count Total Latency(ms) Avg Latency(ms) tput_avg (Gbps) busbw_avg (Gbps) broadcast 4.0 KB 221 20.58 0.09 0.36 0.36 8.0 KB 1 0.17 0.17 0.38 0.38 16.0 KB 24 2.55 0.10 1.25 1.25 113.27 KB 1 0.09 0.09 10.70 10.70 2.0 MB 1 0.44 0.44 38.01 38.01 4.0 MB 97 71.00 0.72 46.41 46.41 16.0 MB 48 155.34 3.23 41.67 41.67 113.27 MB 1 178.26 178.26 5.33 5.33 all_reduce 1.24 GB 50 22385.67 445.17 48.11 36.08 log_summary_barrier 0B 1 69.95 69.95 0.00 0.00 利用这些日志，我测量了ZeRO每个阶段引入的通信开销。结果如下所示：
多节点训练 到目前为止，我们所看到的所有示例都在单个节点上展示了多个GPU的分布式训练。DeepSpeed也可以应用于多节点训练。要执行多节点训练DeepSpeed，我们可以使用与之前相同的训练脚本，但需要一些额外的设置，以允许多个节点相互通信。本节概述了在多个AWS EC2实例上执行多节点训练的步骤。
假设我们有三个能够相互通信的所有端口的EC2实例。使用AWS CDK创建这些实例的示例包含在GitHub仓库中：distributed-training-and-deepspeed/aws/multi_node_training_stack.py。
为简化此示例，我将使用 worker-1 、 worker-2 和 worker-3 代替三个GPU实例的实际IP地址，并使用 localhost 来指代我的本地机器。如果你想使用工作器名称进行操作，你可以将它们添加到你的 ~/.ssh/config 文件中，如下所示，将HostNames替换为你的实例的实际IP地址，将IdentityFile替换为你的AWS密钥对的路径：
Host worker-1 HostName ec2-34-222-121-132.us-west-2.compute.amazonaws.com IdentityFile /home/me/keypair.pem User ubuntu Host worker-2 HostName ec2-32-27-113-7.us-west-2.compute.amazonaws.com IdentityFile /home/me/keypair.pem User ubuntu Host worker-3 HostName ec2-52-88-92-190.us-west-2.compute.amazonaws.com IdentityFile /home/me/keypair.pem User ubuntu 要运行多节点训练作业，每个节点都需要能够通过无密码SSH连接到每个其他节点。为此，我们将创建一个新的SSH密钥对并将其分发到每个节点：
# 在localhost上 # 创建一个新的SSH密钥对 mkdir -p multi-node-training-keys ssh-keygen -t rsa -N '' -f ./multi-node-training-keys/id_rsa # 将SSH密钥复制到worker-1 scp ./multi-node-training-keys/id_rsa.pub worker-1:/home/ubuntu/.ssh/id_rsa.pub scp ./multi-node-training-keys/id_rsa worker-1:/home/ubuntu/.ssh/id_rsa # 在工作器节点上将SSH密钥添加到authorized_keys ssh worker-1 'cat /home/ubuntu/.ssh/id_rsa.pub >> /home/ubuntu/.ssh/authorized_keys' # 为worker-2和worker-3重复上述三个命令 接下来，我们需要在每个工作器上安装DeepSpeed及其依赖项。这个页面的GitHub仓库包括一个 requirements.txt 文件和一个 worker-prereqs.sh 脚本，可以用来安装所有先决条件。在每个节点上，我们将克隆GitHub仓库并安装先决条件：
# 在localhost上 # 在worker-1上安装依赖项 ssh worker-1 'rm -rf distributed-training-and-deepspeed && git clone https://github.com/gnovack/distributed-training-and-deepspeed.git && pip install -r distributed-training-and-deepspeed/requirements.txt && distributed-training-and-deepspeed/scripts/worker-prereqs.sh' # 为worker-2和worker-3重复上述命令 最后，虽然不是必需的，但在每个工作器上配置SSH主机名别名是很好的。我们可以通过在每个工作器上写入 /home/ubuntu/.ssh/config 文件来实现：
# 在localhost上；假设你已经在你的~/.ssh/config文件中添加了SSH主机名别名 worker_ip_1=$(ssh -G worker-1 | awk '$1 == 'hostname' { print $2 }' ) worker_ip_2=$(ssh -G worker-2 | awk '$1 == 'hostname' { print $2 }' ) worker_ip_3=$(ssh -G worker-3 | awk '$1 == 'hostname' { print $2 }' ) ssh worker-1 'cat > /home/ubuntu/.ssh/config' << EOF Host worker-1 HostName $worker_ip_1 StrictHostKeyChecking no Host worker-2 HostName $worker_ip_2 StrictHostKeyChecking no Host worker-3 HostName $worker_ip_3 StrictHostKeyChecking no EOF # 为worker-2和worker-3重复上述命令 通过设置 StrictHostKeyChecking no ，我们绕过了手动接受每个工作器的SSH主机密钥的需要。这简化了这个例子，但不建议在生产环境中这样做。
一个执行所有这些命令的完整脚本适用于所有三个工作器节点，在distributed-training-and-deepspeed/scripts/generate-keys.sh中可用。
现在每个节点都可以通过无密码SSH连接到每个其他节点，下一步是创建一个 hostfile ，列出每个节点的IP地址，并指定每个节点上可用的GPU数量。假设我们的 ~/.ssh/config 中有每个工作器节点的别名，我们可以使用这些别名而不是IP地址在我们的hostfile中：
worker-1 slots=1 worker-2 slots=1 worker-3 slots=1 随着hostfile的创建，我们准备 运行一个多节点训练作业。我们可以从任何一个工作器节点启动作业。对于这个例子，我们将从worker-1开始作业，通过将hostfile从localhost复制到worker-1，然后使用 --hostfile 参数运行训练脚本：
# 在localhost上 # 将hostfile复制到worker-1 scp hostfile worker-1:/home/ubuntu/distributed-training-and-deepspeed/hostfile # 从SSH配置文件中获取worker-1的IP地址 MASTER_ADDR=$(ssh -G worker-1 | awk '$1 == 'hostname' { print $2 }' ) # ssh到第一个工作器节点并启动训练 ssh worker-1 'cd distributed-training-and-deepspeed && PATH='/home/ubuntu/.local/bin:$PATH' deepspeed --master_addr=$MASTER_ADDR --hostfile=./hostfile zero_dp_training.py --stage=2 --model_name=facebook/opt-125m' 总结 这次对分布式训练和DeepSpeed的探索让我学到了很多关于分布式模型训练的细微差别和挑战，以及零冗余优化器如何在多个GPU上训练时减少内存需求。DeepSpeed中还有许多其他优化没有在本文中涵盖，这为我未来的探索留下了几个话题：
卸载：DeepSpeed提供了将优化器状态和模型参数卸载到CPU内存和磁盘的能力，这可以进一步减少内存消耗，以训练时间为代价。
专家混合（MoE）：MoE模型是一类使用稀疏激活层（即在每次前向传递中选择一部分权重使用的层）来增加模型参数总数而不增加计算复杂性的模型。DeepSpeed库包括稀疏激活的PyTorch模型层，可以用来实现MoE模型。
渐进式层丢弃（PLD）：PLD是一种通过允许在训练过程中的不同点切换特定的Transformer层来加速Transformer模型训练的技术。
参考文献 References [1] Sohoni, Nimit S., et al. “Low-memory neural network training: A technical report.” arXiv preprint arXiv:1904.10631 (2019).
[2] https://huggingface.co/docs/transformers/v4.20.1/en/perf_train_gpu_one#anatomy-of-models-memory
[3] Korthikanti, Vijay Anand, et al. “Reducing activation recomputation in large transformer models.” Proceedings of Machine Learning and Systems 5 (2023).
[4] Li, Shen, et al. “Pytorch distributed: Experiences on accelerating data parallel training.” arXiv preprint arXiv:2006.15704 (2020).
[5] Sergeev, Alexander, and Mike Del Balso. “Horovod: fast and easy distributed deep learning in TensorFlow.” arXiv preprint arXiv:1802.05799 (2018).
[6] Huang, Yanping, et al. “Gpipe: Efficient training of giant neural networks using pipeline parallelism.” Advances in neural information processing systems 32 (2019).
[7] https://huggingface.co/docs/transformers/v4.15.0/parallelism#naive-model-parallel-vertical-and-pipeline-parallel
[8] Rajbhandari, Samyam, et al. “Zero: Memory optimizations toward training trillion parameter models.” SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020.
[9] Scao, Teven Le, et al. “Bloom: A 176b-parameter open-access multilingual language model.” arXiv preprint arXiv:2211.05100 (2022).
附录 附录A：计算Transformer内存需求 不幸的是，上述用于计算激活内存的方法（即注入前向钩子）与HuggingFace Transformers库中的Transformer模型并不兼容，因为这些模型中的许多操作不是torch.nn.Module实例，因此不会被钩子捕获。
为了克服这一点，我编写了一个基本的TransformerBlock类，基于HuggingFace的OPTDecoderLayer实现，但我使用了torch.nn.Module实例来执行所有操作，以确保每个中间输出都能被activation_counter_hook钩子捕获。
class MatMul(nn.Module): '' ' PyTorch Module wrapper for torch.bmm. ' '' def __init__(self): super().__init__() def forward(self, x, y): return torch.bmm(x, y) class TransformerBlock(nn.Module): def __init__(self, config): super().__init__() self.config = config self.hidden_size = config.hidden_size self.num_attention_heads = config.num_attention_heads self.head_dim = config.hidden_size // config.num_attention_heads self.ffn_dim = config.ffn_dim self.scaling = self.head_dim**-0.5 self.pre_attention_layer_norm = nn.LayerNorm(self.hidden_size) self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=config.enable_bias) self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=config.enable_bias) self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=config.enable_bias) self.out_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=config.enable_bias) self.compute_attention_weights = MatMul() self.attention_weights_softmax = nn.Softmax(dim=-1) self.attention_weights_dropout = nn.Dropout(config.dropout) self.compute_attentions = MatMul() self.attention_output_dropout = nn.Dropout(config.dropout) self.final_layer_norm = nn.LayerNorm(self.hidden_size) self.ffn_1 = nn.Linear(self.hidden_size, self.ffn_dim, bias=config.enable_bias) self.ffn_2 = nn.Linear(self.ffn_dim, self.hidden_size, bias=config.enable_bias) self.activation = nn.GELU() self.final_dropout = nn.Dropout(config.dropout) def forward(self, hidden_states): 这个类的完整实现可以在distributed-training-and-deepspeed/model/transformer.py找到。计算Transformer模型内存需求的脚本可以在这里找到：distributed-training-and-deepspeed/estimate_transformer_memory.py。
参考：https://tinkerd.net/blog/machine-learning/distributed-training/
