大約兩年前，我們發(fā)布了一項全新機器學習 (ML) 運行時工作：TFRT(TensorFlow 運行時的簡稱)。同時，我們提供了初始技術(shù)設(shè)計的深度教程，并將其代碼庫開源。

在 ML 生態(tài)系統(tǒng)趨勢的推動下——模型越來越大，ML 被部署到更多樣化的執(zhí)行環(huán)境，以及跟上持續(xù)研究和模型創(chuàng)新的需求——我們啟動了 TFRT，希望可以實現(xiàn)以下目標：

為 ML 模型提供更快、更便宜的執(zhí)行

實現(xiàn)更靈活的部署

提供更多模塊化的可擴展基礎(chǔ)架構(gòu)，以促進 ML 基礎(chǔ)架構(gòu)和建模方面的創(chuàng)新

本文，我們將分享截至目前所取得的進展、在過去兩年的開發(fā)過程中收獲的經(jīng)驗和教訓(xùn)，以及未來的規(guī)劃。

目前所取得的進展

過去兩年的開發(fā)主要集中于通過為用戶啟用 Google 最重要的內(nèi)部工作負載(比如 Ads 和搜索)，來實施和驗證我們的想法。迄今為止，我們已經(jīng)在 Google 內(nèi)部廣泛部署了 TFRT，用于各種訓(xùn)練和推斷工作負載，并取得了很好的結(jié)果。

技術(shù)經(jīng)驗

我們是如何實現(xiàn)上述目標的?除了原始設(shè)計中的經(jīng)驗，我們還收獲了一些有趣的技術(shù)經(jīng)驗：

首先，異步支持對于一些關(guān)鍵的工作負載(例如重疊計算和 I/O，以及驅(qū)動異構(gòu)設(shè)備)十分重要，而快速同步執(zhí)行對許多其他工作負載至關(guān)重要，包括小型“嵌入式”ML 模型。

我們花費了大量精力設(shè)計和改進 AsyncValue，這是 TFRT 中的一個關(guān)鍵低級抽象，它允許主機運行時異步驅(qū)動設(shè)備以及調(diào)用內(nèi)核。由于它能夠在主機和設(shè)備間重疊更多的計算和通信，從而可以使設(shè)備利用率提高。例如，通過將模型拆分為多個階段并使用 TFRT 來將下一階段的變量傳輸與當前階段的 TPU 計算重疊，我們能夠以高性能在一個 TPU 芯片上成功運行 80B 參數(shù)模型的批量推理。

另一方面，在應(yīng)用進程中(而不是通過 RPC/REST 調(diào)用)調(diào)用嵌入在應(yīng)用服務(wù)器里的小型 CPU 模型，對 Google 的一些來自用戶的商業(yè)工作負載(例如 Ads)仍然十分重要。對于這些模型，TFRT 的異步優(yōu)先內(nèi)部設(shè)計一開始導(dǎo)致了性能和資源衰退。我們與 Ads 團隊合作，通過使用同步解釋器擴展 TFRT 設(shè)計以及實驗性內(nèi)存規(guī)劃優(yōu)化，成功解決了該問題，避免了內(nèi)核執(zhí)行期間的堆分配。我們正在努力將此擴展產(chǎn)品化。

下圖展示了由此產(chǎn)生的 TFRT 設(shè)計對基準測試的影響，通過與在部署 TFRT 之前運行舊運行時的“當前 TF”對比。這個基準測試專注于執(zhí)行微型 CPU 模型，其中大量的小型 matmul 按順序執(zhí)行。值得注意的是，TFRT 中的優(yōu)化執(zhí)行 (265 ns) 正接近我們設(shè)置的最佳基線 (204 ns)，它通過手寫 C++ 代碼執(zhí)行，從而沒有任何 ML 運行時開銷。

其次，雖然更快的運行時執(zhí)行很重要，但是優(yōu)化輸入程序以降低執(zhí)行復(fù)雜度也很重要。

請注意，雖然應(yīng)盡可能在將 TF SavedModel 保存到磁盤時執(zhí)行基于編譯器的圖優(yōu)化，但也有一些重要的推理時編譯器優(yōu)化只能在推理上下文中執(zhí)行(例如，在訓(xùn)練變量保持不變時)。

當我們將 ML 模型加入 TFRT 時，我們可以在執(zhí)行之前深入檢查一些模型，并確定重寫和簡化程序的新方法。簡化的程序，以及計算圖程序中每個內(nèi)核的更快速執(zhí)行，在縮短執(zhí)行延遲時間和減少資源成本方面產(chǎn)生了很好的復(fù)合效果。

例如，在下面左側(cè)的計算圖程序中，我們能夠提升標量算子歸一化計算(例如，將浮點值除以其域的最大值)，這 18 個用于形成“concat”算子的輸入標量均相同，因而我們對串聯(lián) 1D 浮點張量啟用歸一化矢量執(zhí)行。

雖然也可以在模型訓(xùn)練時執(zhí)行此優(yōu)化，但用于生成訓(xùn)練模型的編譯器+運行時不包括此優(yōu)化。

此外，我們還發(fā)現(xiàn)，盡可能將計算從模型執(zhí)行時間提升到加載時間至關(guān)重要(例如，const 折疊)。

再者，基于成本的執(zhí)行不僅僅適用于 SQL 查詢。

我們?yōu)?TF 算子內(nèi)核開發(fā)了一個簡單的編譯時成本模型(類似于 SQL 查詢優(yōu)化器的成本模型)，并為 ML 模型執(zhí)行應(yīng)用了基于成本的優(yōu)化(參見流分析)，同時在一組線程池線程間實現(xiàn)了更好的內(nèi)核執(zhí)行負載平衡。相比之下，TF1 有一個基于運行時的成本模型，其中每個操作的運行時成本都被分析并用于指導(dǎo)該操作的調(diào)度。在 TFRT 中，我們將成本分析移至編譯時，從而消除了運行時成本。此外，我們的編譯器方法可以分析整個計算圖，從而產(chǎn)生在更全局范圍內(nèi)最佳的調(diào)度決策。

有關(guān)數(shù)據(jù)和 ML 基礎(chǔ)架構(gòu)之間的更多相似之處，請觀看此技術(shù)講座。

展望未來

雖然我們確實取得了一些不錯的進展，尤其是在我們的第一個目標，即追求更快更經(jīng)濟的執(zhí)行上，但我們承認，在實現(xiàn)更模塊化的設(shè)計和通過硬件集成實現(xiàn)更靈活的部署方面，仍道阻且長。

在模塊化方面，隨著 JAX 采用 TFRT 設(shè)備運行時(例如 CPU)等初始集成成功，我們將繼續(xù)探索 TFRT 可以如何支持 TensorFlow 之外的工作負載。我們希望部分 TFRT 組件未來也將有利于 PyTorch/XLA 工作負載的發(fā)展。

此外，我們成功集成了 CPU 和 TPU(下一步是集成到 Cloud TPU 中)，即 Google 用于 ML 計算的兩種最重要的設(shè)備類型，NVIDIA GPU 集成也正在進行中。

在訓(xùn)練工作負載方面，TFRT 已被用作 Google 大規(guī)模分布式訓(xùn)練框架的基本模塊，目前正在積極開發(fā)中。

展望未來，我們的組織一直在探索與 Pixel 的硬件 SOC 設(shè)備(例如 Google Tensor)的集成。此外，由于 TFRT 已被成功證明可用于 Google 的內(nèi)部工作負載，它也被集成到 GCP 的 Vertex AI 和 Waymo 等新場景。

特別致謝

TFRT 團隊非常享受致力于在這個新基礎(chǔ)架構(gòu)項目上的工作。它讓人感覺像是在引導(dǎo)一家新的初創(chuàng)公司。在此，我們想向在這非凡的 2 年旅程中為 TFRT 提供建議、作出貢獻和給予支持的每個人高聲致謝：

(按字母順序)Adi Agrawal、Andrew Bernard、Andrew Leaver、Andy Selle、Ayush Dubey、Bangda Zhou、Bramandia Ramadhana、Catherine Payne、Ce Zheng、Chiachen Chou、Chao Xie、Christina Sorokin、Chuanhao Zhuge、Dan Hurt、Dong Lin、Eugene Zhulenev、Ewa Matejska、Hadi Hashemi、Haoliang Zhang、HanBin Yoon、Haoyu Zhang、Hongmin Fan、Jacques Pienaar、Jeff Dean、Jeremy Lau、Jordan Soyke、Jing Dong、Juanli Shen、Kemal El Moujahid、Kuangyuan Chen、Mehdi Amini、Ning Niu、Peter Gavin、Phil Sun、Pulkit Bhuwalka、Qiao Zhang、Raziel Alvarez、Russell Power、Sanjoy Das、Shengqi Zhu、Smit Hinsu、Tatiana Shpeisman、Tianrun Li、Tim Davis、Tom Black、Victor Akabutu、Vilobh Meshram、Xiao Yu、Xiaodan Song、Yiming Zhang、YC Ling、Youlong Chen 和 Zhuoran Liu。

我們還要特別感謝 Chris Lattner 在引導(dǎo)這個項目上提供的初始技術(shù)領(lǐng)導(dǎo)、Martin Wicke 在第一年對 TFRT 的支持、Alex Zaks 在第二年對 TFRT 的支持以及見證其有效地登陸 Google 的 ML 服務(wù)工作負載。

原文標題：TFRT 進展與更新

文章出處：【微信公眾號：谷歌開發(fā)者】歡迎添加關(guān)注!文章轉(zhuǎn)載請注明出處。

審核編輯：湯梓紅