文 | Sherry

---

在2019年PyTorch開發者大會上，Facebook，Google和Salesforce Research聯合宣布啟動PyTorch-TPU項目。項目的目標是在保持PyTorch的靈活性的同時讓社區盡可能容易地利用云TPU提供的高性能計算。團隊創建了PyTorch/XLA這個repo，它可以讓使PyTorch連接到云TPU并使用TPU內核，同時Colab也支持在云TPU上使用PyTorch/XLA。現在，PyTorch/XLA 已在Google Cloud上達到通用標準（GA），并為我們提供了許多方便易用的接口。趕緊來看看PyTorch/XLA都有哪些功能吧！

GA版本有哪些新功能？

借助PyTorch/XLA GA，云TPU正式支持PyTorch 1.6。其他值得注意的新功能包括：

• 支持層內模型并行性：現在可以在reduce的時候在多組tensor上使用多種運算。添加了更多的通信原語從而可以實現有趣的應用，例如把嵌入詞向量詞分布到多個TPU內核上；

• 額外的XLA運算：隨著PyTorch/XLA在越來越廣泛的新模型中被使用，用戶要求將PyTorch運算映射到XLA。對此，從beta（1.5）版本開始，我們已經為Replication_pad1d，replication_pad2d，max_unpool2d，max_unpool3d等運算加入了低配XLA；

• 在Colab/Kaggle上更好的體驗：現在，不再需要在Colab/Kaggle上運行env-setup.py腳本，就可以開始訓練倆。

• 對深度學習VM映像的支持：Google Cloud 平臺提供了一組深度學習虛擬機（DLVM）映像，配置好了各種常用深度學習框架（包括PyTorch）和所需的一切。PyTorch/XLA 1.6現在已預安裝在DLVM中，并針對云TPU進行了優化。官方的PyTorch 1.6也預裝在相同的Conda環境中。

詳細使用方法參照用戶指南^[1]。

GA版本支持哪些模型？

PyTorch/XLA已被用于在云TPU上訓練眾多深度學習模型。包括：

• 圖像分類任務（ImageNet）：Torchvision的ResNet-50；

• 翻譯任務（WMT-18 en-de）：Fairseq Transformer；

• 通用的語言模型的預訓練和微調（GLUE等）：HuggingFace（BERT，DistilBERT，RoBERTa，XLNet等）和Fairseq RoBERTa；

• 深度學習推薦模型（1.6版的新增功能）：DLRM。

在大多數情況下，在云TPU上訓練這些模型幾乎不需要更改代碼。可以參考官方教程^[2]獲得上述模型的詳細教程，也可以參考PyTorch / XLA GitHub^[3]代碼庫找到訓練的其他模型架構的示例。

PyTorch/XLA是如何工作的？

PyTorch/XLA使用“lazy tensor”進行抽象。使用lazy tensor時，運算的evaluation會被推遲到這個運算被訪問之前。此時，運算被描述為中間表示圖，當我們需要運算結果時（訪問），就通過XLA編譯這些中間表示圖，并發送到TPU內核以執行。這種XLA編譯對CPU和GPU都支持。其他技術詳細信息，可以參考GitHub^[4]。

升級到云TPU要改哪些代碼？

首先需要創建一個配有PyTorch / XLA映像的Google Compute Engine虛擬機和一個云TPU實例。創建好虛擬機和云TPU實例后，需要配置conda環境并設置XRT_TPU_CONFIG環境變量指向云TPU實例：

export?XRT_TPU_CONFIG="tpu_worker;0;<TPU_IP_ADDRESS>:8470"

現在就可以開始在云TPU上訓練模型了！

實際代碼中需要改動的地方包括：

• 調取相關代碼包import torch_xla

• 設置訪問XLA設備抽象的方法以及并行數據加載器。

• 使用xm.optimizer_step(optimizer)在后臺進行reduce。

可以參考英文博客原文^[5]中的代碼實例看看具體的區別。

在Cloud TPU Pod上訓練

PyTorch/XLA支持將剛剛在單個云TPU上執行的訓練擴展到整個云TPU Pod或者任意Pod切片,只需要使用xla_dist包裝就可以了：

?python?-m?torch_xla.distributed.xla_dist?\--tpu=$TPU_NAME?\--conda-env=torch-xla-1.6?\--env?ANY_ENV_VAR=VALUE?\--?\python?/path/to/your/code.py?--train_arg1?\--train_arg2?...

官方提供的Colab notebook^[6]寫得非常詳細，現在就開始探索吧！

更多詳情請參考英文博客原文^[7]以及Google官方博文^[8]

后臺回復關鍵詞【入群】

加入賣萌屋NLP/IR/Rec與求職討論群

有頂會審稿人、大廠研究員、知乎大V和妹紙

等你來撩哦~

參考文獻

[1]用戶指南: https://docs.google.com/document/d/1RzNgpTK4wESImhIwgMbknqn9xVNUFtCTCknNfONlMlo/edit

[2]官方教程: https://cloud.google.com/tpu/docs/tutorials/

[3]PyTorch / XLA GitHub: https://github.com/pytorch/xla

[4]GitHub: https://github.com/pytorch/xla/blob/master/API_GUIDE.md#xla-tensor-deep-dive

[5]英文博客原文: https://medium.com/pytorch/pytorch-xla-is-now-generally-available-on-google-cloud-tpus-f9267f437832

[6]Colab notebook: https://github.com/pytorch/xla/tree/master/contrib/colab

[7]英文博客原文: https://medium.com/pytorch/pytorch-xla-is-now-generally-available-on-google-cloud-tpus-f9267f437832

[8]Google官方博文: https://cloud.google.com/blog/products/ai-machine-learning/pytorch-is-now-ga-on-google-cloud-tpu

總結

以上是生活随笔為你收集整理的Google Cloud TPUs支持Pytorch框架啦！的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。