新火種
2023-09-07 
HuggingFace發布PyTorch新庫：適用于多GPU、TPU、混合精度訓練
機器之心報道
作者：力元
多數 PyTorch 高級庫都支持分布式訓練和混合精度訓練，但是它們引入的抽象化往往需要用戶學習新的 API 來定制訓練循環。許多 PyTorch 用戶希望完全控制自己的訓練循環，但不想編寫和維護訓練所需的樣板代碼。Hugging Face 最近發布的新庫 Accelerate 解決了這個問題。
「Accelerate」提供了一個簡單的 API，將與多 GPU 、 TPU 、 fp16 相關的樣板代碼抽離了出來，保持其余代碼不變。PyTorch 用戶無須使用不便控制和調整的抽象類或編寫、維護樣板代碼，就可以直接上手多 GPU 或 TPU。
通過將如下 5 行代碼添加到原始的 PyTorch 訓練循環中，腳本即可在本地以及任何分布式設置上運行。
import torch import torch.nn.functional as F from datasets import load_dataset+ from accelerate import Accelerator+ accelerator = Accelerator()- device = 'cpu'+ device = accelerator.device model = torch.nn.Transformer().to(device) optim = torch.optim.Adam(model.parameters()) dataset = load_dataset('my_dataset') data = torch.utils.data.DataLoader(dataset, shuffle=True)+ model, optim, data = accelerator.prepare(model, optim, data) model.train() for epoch in range(10): for source, targets indata: source = source.to(device) targets = targets.to(device) optimizer.zero_grad() output = model(source) loss = F.cross_entropy(output, targets)+ accelerator.backward(loss)- loss.backward() optimizer.step()
Accelerate 甚至可以通過處理設備的放置（需要對代碼進行一些更改，但通常更安全）進一步簡化訓練循環，代碼如下所示：
import torch import torch.nn.functional as F from datasets import load_dataset+ from accelerate import Accelerator+ accelerator = Accelerator()- device = 'cpu'+ model = torch.nn.Transformer()- model = torch.nn.Transformer().to(device) optim = torch.optim.Adam(model.parameters()) dataset = load_dataset('my_dataset') data = torch.utils.data.DataLoader(dataset, shuffle=True)+ model, optim, data = accelerator.prepare(model, optim, data) model.train() for epoch in range(10): for source, targets indata:- source = source.to(device)- targets = targets.to(device) optimizer.zero_grad() output = model(source) loss = F.cross_entropy(output, targets)+ accelerator.backward(loss)- loss.backward() optimizer.step()
除了 API，Accelerate 還提供了一個 CLI 工具，方便啟動腳本之前快速配置和測試訓練環境，然后啟動腳本。指令如下所示：
accelerate config
acceleratelaunchmy_script.py--args_to_my_script
如果不想自己編寫訓練循環， PyTorch 之上有許多可以替代 Accelerate 的高級庫。
Accelerate 的運作原理
accelerator = Accelerator()
除了提供要使用的主要對象之外，此行還將從環境中分析分布式訓練運行的類型并執行必要的初始化。用戶可以通過將 cpu = True 或 fp16 = True 傳遞給此 init 來強制進行 CPU 訓練或混合精度訓練。這兩個選項都可以使用腳本的啟動器進行設置。
model, optim, data = accelerator.prepare(model, optim, data)
這是 API 的主體，將準備三種主要類型的對象：models (torch.nn.Module)、optimizers (torch.optim.Optimizer)、dataloaders (torch.data.dataloader.DataLoader)。
模型
模型的準備包括將其包裝在適當的容器（例如 DistributedDataParallel）中，然后將其放置在適當的設備上。與普通分布式訓練一樣，進行保存或訪問其特定的方法時，需要先通過 accelerator.unwrap_model（model）解開模型。
優化器
優化器也包裝于一個為使混合精度訓練運轉執行必要操作的特殊容器里。如果狀態 dict 是非空的或從檢查點加載的，它會準確地處理狀態 dict 的設備放置。
數據加載器
此庫并不依賴于 DistributedSampler，它實際上可以與傳遞到數據加載器的采樣器一起使用。數據加載器包裝于僅在采樣器中獲取與當前進程相關的索引并將批次放入設備的容器中。
為此，Accelerate 提供了一種實用程序功能來同步更多 RNGs。該功能將在分布式訓練期間運行的每個進程上同步隨機數生成器。默認情況下，它僅同步采樣器的生成器，因此在每個進程中數據擴充都將不同，但是隨機改組將是相同的。
accelerator.backward(loss)
此行代碼為向后傳遞添加了必要的步驟來提高混合精度，但對于其他集成則需要進行一些自定義。
Accelerate 支持的集成包括：
CPU單 GPU單一節點多 GPU多節點多 GPUTPU帶有本地 AMP 的 FP16（路線圖上的頂點）
相關推薦
- 免責聲明
- 本文所包含的觀點僅代表作者個人看法，不代表新火種的觀點。在新火種上獲取的所有信息均不應被視為投資建議。新火種對本文可能提及或鏈接的任何項目不表示認可。 交易和投資涉及高風險，讀者在采取與本文內容相關的任何行動之前，請務必進行充分的盡職調查。最終的決策應該基于您自己的獨立判斷。新火種不對因依賴本文觀點而產生的任何金錢損失負任何責任。
