GPU如何训练大批量模型？方法在这里

　　扩展到极致

　　你可以在 GPU 上训练连一个样本都无法加载的模型吗?

　　如果你的架构没有太多跳过连接，这就是可能的!解决方案是使用梯度检查点(gradient-checkpointing)来节省计算资源。

　　基本思路是沿着模型将梯度在小组件中进行反向传播，以额外的前馈传递为代价，节约存储完整的反向传播图的内存。这个方法比较慢，因为我们需要添加额外的计算来减少内存要求，但在某些设置中挺有意思，比如在非常长的序列上训练 RNN 模型(示例参见 https://medium.com/huggingface/from-zero-to-research-an-introduction-to-meta-learning-8e16e677f78a)。

　　这里不再赘述，读者可以查看以下链接：

　　TensorFlow：https://github.com/openai/gradient-checkpointing

　　PyTorch 文档：https://pytorch.org/docs/stable/checkpoint.html

　　「节约内存」(Memory-poor)策略需要 O(1) 的内存(但是要求 O(n²) 的计算步)。

　　充分利用多 GPU 机器

　　现在我们具体来看如何在多 GPU 上训练模型。

　　在多 GPU 服务器上训练 PyTorch 模型的首选策略是使用 torch.nn.DataParallel。该容器可以在多个指定设备上分割输入，按照批维度(batch dimension)分割，从而实现模块应用的并行化。

　　DataParallel 非常容易使用，我们只需添加一行来封装模型：

　　parallel_model = torch.nn.DataParallel(model) # Encapsulate the model

　　predictions = parallel_model(inputs) # Forward pass on multi-GPUs

　　loss = loss_function(predictions, labels) # Compute loss function

　　loss.backward() # Backward pass

　　optimizer.step() # Optimizer step

　　predictions = parallel_model(inputs) # Forward pass with new parameters

　　但是，DataParallel 有一个问题：GPU 使用不均衡。

　　在一些设置下，GPU-1 会比其他 GPU 使用率高得多。

　　这个问题从何而来呢?下图很好地解释了 DataParallel 的行为：

　　使用 torch.nn.DataParallel 的前向和后向传播。

　　在前向传播的第四步(右上)，所有并行计算的结果都聚集在 GPU-1 上。这对很多分类问题来说是件好事，但如果你在大批量上训练语言模型时，这就会成为问题。

　　我们可以快速计算语言模型输出的大小：

　　语言模型输出中的元素数量。

　　假设我们的数据集有 4 万词汇，每一条序列有 250 个 token、每个 batch 中有 32 条序列，那么序列中的每一个元素需要 4 个字节的内存空间，模型的输出大概为 1.2GB。要储存相关的梯度张量，我们就需要把这个内存翻倍，因此我们的模型输出需要 2.4GB 的内存。

　　这是典型 10GB GPU 内存的主要部分，意味着相对于其它 GPU，GPU - 1 会被过度使用，从而限制了并行化的效果。

　　如果不调整模型和/或优化方案，我们就无法轻易减少输出中的元素数量。但我们可以确保内存负载在 GPU 中更均匀地分布。

　　多 GPU 机器上的均衡负载

　　解决办法是把每部分输出保留在其 GPU 上，而不是将它们聚集到 GPU-1 上。我们也需要分配损失标准计算，计算损失并进行反向传播。

　　幸而，张航开源了一个名为 PyTorch-Encoding 的 PyTorch 包，它包含了这些定制的并行化功能。

　　我提取并稍稍改动了这个模块，你可以从以下地址下载 gist(parallel.py)来纳入并调用你的代码。它主要包括两个模块：DataParallelModel 和 DataParallelCriterion，它们的用途如下：

　　下载地址：https://gist.github.com/thomwolf/7e2407fbd5945f07821adae3d9fd1312

　　from parallel import DataParallelModel, DataParallelCriterion

　　parallel_model = DataParallelModel(model) # Encapsulate the model

　　parallel_loss = DataParallelCriterion(loss_function) # Encapsulate the loss function

　　predictions = parallel_model(inputs) # Parallel forward pass

　　# "predictions" is a tuple of n_gpu tensors

　　loss = parallel_loss(predictions, labels) # Compute loss function in parallel

　　loss.backward() # Backward pass

　　optimizer.step() # Optimizer step

　　predictions = parallel_model(inputs) # Parallel forward pass with new parameters

　　DataParallelModel 和 torch.nn.DataParallel 的区别在于，前向传播的输出(predictions)没有聚集在 GPU-1 上，而是作为 n_gpu 张量的元组，每个张量分布在相应的 GPU 上。

　　DataParallelCriterion 容器封装了损失函数，并把 n_gpu 张量元组和目标标签张量作为输入。它在每个 GPU 上并行计算损失函数，像 DataParallel 分割模型输入一样分割目标标签张量。

　　下图说明了 DataParallelModel/DataParallelCriterion 的内部情况：

　　使用 DataParallelModel 和 DataParallelCriterion。

　　以下是你可能会遇到的两个特定案例的解决办法：

　　你的模型输出几个张量：你可能想分解它们：output_1, output_2 = zip(*predictions)

　　有时候你并不想使用并行损失函数：收集 CPU 上的所有张量：gathered_predictions = parallel.gather(predictions)