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

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

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

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

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

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

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

目前所取得的進(jìn)展

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

技術(shù)經(jīng)驗(yàn)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

展望未來

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

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

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

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

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

特別致謝

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

(按字母順序)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)這個(gè)項(xiàng)目上提供的初始技術(shù)領(lǐng)導(dǎo)、Martin Wicke 在第一年對(duì) TFRT 的支持、Alex Zaks 在第二年對(duì) TFRT 的支持以及見證其有效地登陸 Google 的 ML 服務(wù)工作負(fù)載。

原文標(biāo)題：TFRT 進(jìn)展與更新

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

審核編輯：湯梓紅