Sharded是一項新技術，它可以幫助您節省超過60%的記憶體，并将模型放大兩倍。
深度學習模型已被證明可以通過增加資料和參數來改善。即使使用175B參數的Open AI最新GPT-3模型，随着參數數量的增加，我們仍未看到模型達到平穩狀态。
對于某些領域，例如NLP，最主要的模型是需要大量GPU記憶體的Transformer。對于真實模型，它們隻是不适合存儲在記憶體中。微軟的Zero論文介紹了稱為Sharded的最新技術，他們開發了一種技術，可以使我們接近1萬億個參數。
在本文中，我将給出sharded工作原理，并向您展示如何利用PyTorch 在幾分鐘内用将使用相同記憶體訓練模型參數提升一倍。
由于Facebook AI Research的FairScale團隊與PyTorch Lightning團隊之間的合作，PyTorch中的此功能現已可用。
本文大綱
- 本文适用于誰？
- 如何在PyTorch中使用Sharded
- Sharded後的工作原理
- Sharded與模型并行
本文适用于誰？
本文适用于使用PyTorch訓練模型的任何人。Sharded适用于任何模型，無論它是哪種類型的模型，無論是NLP，視覺SIMCL，Swav，Resnets還是語音。
以下是這些模型類型在Sharded時可以看到的性能提升的快速快照。
SwAV是計算機視覺中自我監督學習的最新方法。
DeepSpeech2是最先進的語音方法。
圖像GPT是最先進的視覺方法。
Transformer 是NLP的最新方法。
如何在PyTorch中使用Sharded
對于那些沒有足夠的時間來了解Sharded工作原理的人，我将在前面解釋如何在您的PyTorch代碼中使用Sharded。但是，我鼓勵您通讀本文結尾，以了解Sharded的工作原理。
Sharded意味着可以與多個GPU一起使用以獲得所有好處。但是，在多個GPU上進行訓練會比較複雜，并且會造成巨大的痛苦。
使用Sharded為代碼添加代碼的最簡單方法是将模型轉換為PyTorch Lightning（這隻是一個簡單的重構）。
完成此操作後，在8個GPU上啟用Sharded就像更改一個标志一樣簡單，因為無需更改代碼。
如果您的模型來自另一個深度學習庫，那麼它仍然可以與Lightning（NVIDIA Nemo，fast.ai，huggingface transformers）一起使用。您需要做的就是将該模型導入LightningModule并運作訓練。
from argparse import ArgumentParser
import torch
import torch.nn as nn
import pytorch_lightning as pl
from pytorch_lightning.metrics.functional import accuracy
from transformers import BertModel
class LitBertClassifier(pl.LightningModule):
def __init__(self, n_classes, pretrained_model_name='bert-base-uncased'):
super().__init__()
self.save_hyperparameters()
self.bert = BertModel.from_pretrained(pretrained_model_name)
self.drop = nn.Dropout(p=0.3)
self.out = nn.Linear(self.bert.config.hidden_size, n_classes)
self.loss_fn = nn.CrossEntropyLoss()
def forward(self, input_ids, attention_mask):
outputs = self.bert(
input_ids=input_ids,
attention_mask=attention_mask,
return_dict=False
)
pooled_output = outputs[1]
output = self.drop(pooled_output)
return self.out(output)
def training_step(self, batch, batch_idx):
loss, acc = self._shared_step(batch, batch_idx)
self.log("acc", acc)
return loss
def validation_step(self, batch, batch_idx):
_, acc = self._shared_step(batch, batch_idx)
self.log("val_acc", acc)
def _shared_step(self, batch, batch_idx):
input_ids = batch["input_ids"]
attention_mask = batch["attention_mask"]
targets = batch["targets"]
outputs = self.forward(
input_ids=input_ids,
attention_mask=attention_mask
)
_, preds = torch.max(outputs, dim=1)
loss = self.loss_fn(outputs, targets)
acc = accuracy(preds, targets)
return loss, acc
def configure_optimizers(self):
return torch.optim.AdamW(self.parameters(), lr=2e-5)
if __name__ == '__main__':
# TODO: add your own dataset
train_dataloader = ...
val_dataloader = ...
bert = LitBertClassifier()
trainer = pl.Trainer(gpus=8, plugins='ddp_sharded')
trainer.fit(bert, train_dataloader) 複制
Sharded的工作原理
在許多GPU上進行有效訓練有幾種方法。在一種方法（DP）中，每批都配置設定給多個GPU。這是DP的說明，其中批處理的每個部分都轉到不同的GPU，并且模型多次複制到每個GPU。
但是，這種方法很糟糕，因為模型權重是在裝置之間轉移的。此外，第一個GPU維護所有優化器狀态。例如，Adam 優化器會保留模型權重的完整副本。
在另一種方法（分布式資料并行，DDP）中，每個GPU訓練資料的子集，并且梯度在GPU之間同步。此方法還可以在許多機器（節點）上使用。在此示例中，每個GPU擷取資料的子集，并在每個GPU上完全相同地初始化模型權重。然後，在向後傳遞之後，将同步所有梯度并進行更新。
但是，該方法仍然存在一個問題，即每個GPU必須維護所有優化器狀态的副本（大約是模型參數數量的2-3倍）以及所有向前和向後激活。
Sharded消除了這些備援。除了僅針對部分完整參數計算所有開銷（梯度，優化器狀态等）外，它的功能與DDP相同，是以，我們消除了在所有GPU上存儲相同的梯度和優化器狀态的備援。
是以，每個GPU僅存儲激活，優化器參數和梯度計算的子集。
使用分布式模式
通過使用這些優化方法中的任何一種，可以通過多種方法來壓縮分布式訓練中的最大效率。
好消息是，所有這些模式都可在PyTorch Lightning中使用，而零代碼更改則可用。您可以嘗試其中的任何一種，并根據需要根據您的特定模型進行調整。
最後論文在這裡：https://arxiv.org/abs/1910.02054
作者：William Falcon
原文位址：https://towardsdatascience.com/sharded-a-new-technique-to-double-the-size-of-pytorch-models-3af057466dba
deephub翻譯組譯