Jialin Liu, Mengyuan Chao, Jian Li, Wei Peng, Sixiang Ma, Wei Xu, Run Yang, Xin Chen
RayRTC 是字節(jié)基礎(chǔ)架構(gòu)組與字節(jié) AML 組共同合作,在內(nèi)部 RTC(Realtime Text Classification)文本訓(xùn)練平臺上基于 Ray 進(jìn)行得下一代 Serverless ML 得探索。RTC 文本分類平臺是一個一站式得 NLP 服務(wù)平臺,包括了數(shù)據(jù)預(yù)處理,標(biāo)注,模型訓(xùn)練,打分,評估,AutoML 以及模型推理等機(jī)器學(xué)習(xí)全流程。目前字節(jié)內(nèi)各大產(chǎn)品,包括抖音,TikTok,頭條,西瓜,番茄等都有使用該平臺提供得相關(guān)自然語言能力。RayRTC 通過算法與系統(tǒng)得協(xié)同設(shè)計(jì)及 Serverless 等技術(shù)為 RTC 提供了性能和資源利用率得極致優(yōu)化,并由此抽象出一套通用得 Serverless ML 框架,目前已在字節(jié)內(nèi)部機(jī)器學(xué)習(xí)平臺上部署上線。
RayRTC 得核心計(jì)算引擎是 Ray,最早是 UC Berkeley 得一個針對強(qiáng)化學(xué)習(xí)所設(shè)計(jì)得大規(guī)模分布式計(jì)算框架。Ray 得感謝分享 Robert Nishihara 和 Philipp Moritz 在此基礎(chǔ)上成立了 Anyscale 這家公司。開源項(xiàng)目千千萬,能成功商業(yè)化并在硅谷乃至整個 IT 屆產(chǎn)生顛覆性影響得鳳毛麟角。Anyscale 得創(chuàng)始人中包括 Ion Stoica,這位羅馬尼亞籍教授上一家公司是跟他得學(xué)生 Matei Zaharia 以 Spark 技術(shù)為基礎(chǔ)成立得 Databricks 。Spark 和 Ray 分別誕生于大數(shù)據(jù)和機(jī)器學(xué)習(xí)時代,前者已經(jīng)在工業(yè)界得到廣泛應(yīng)用,后者也逐漸引起越來越多得公司在不同業(yè)務(wù)場景進(jìn)行探索。字節(jié)美研計(jì)算團(tuán)隊(duì)自 上年 年末開始接觸 Ray,2021 年開始在不同場景小范圍試驗(yàn)。RTC 文本分類平臺是第壹個大規(guī)模上線得 Ray 應(yīng)用場景,在 RayRTC 得設(shè)計(jì)過程中,有不少第壹手得經(jīng)驗(yàn)值得分享。感謝從 RayRTC 所遇到得實(shí)際問題出發(fā),對 Ray 在字節(jié)得實(shí)踐進(jìn)行介紹。
第壹次接觸 Ray 得讀者可能會問,除了明星創(chuàng)始人團(tuán)隊(duì),深度貼近當(dāng)前 ML 需求得產(chǎn)學(xué)研支持,Ray 這套框架到底有哪些吸引人得地方?
首先是以 Ray 為底座可以非常輕松構(gòu)建完整機(jī)器學(xué)習(xí)完整生態(tài),如下圖所示:
機(jī)器學(xué)習(xí)得研發(fā)人員往往不僅需要感謝對創(chuàng)作者的支持算法本身,在實(shí)際得生產(chǎn)環(huán)境中,各個環(huán)節(jié)所涉及得工程量和運(yùn)維量也不容小覷。不少研究表明,工程師們有 80-90%得時間和精力投入在了算法之外得數(shù)據(jù)處理,全流程打通等。Ray 社區(qū)在近幾年得演進(jìn)中,不斷吸收業(yè)界領(lǐng)先得理念,積極地與其他開源社區(qū)和各大廠商進(jìn)行合作交流。以 Ray 為計(jì)算引擎得上層生態(tài)得豐富度是別得開源生態(tài)中不常見得。比如大數(shù)據(jù)處理方面,有 Intel 設(shè)計(jì)得 RayDP,將 Spark 無縫集成到 Ray 中,通過 Ray 得 Actor 拉起 Spark 得 executor,利用 Ray 得分布式調(diào)度實(shí)現(xiàn)資源細(xì)粒度得調(diào)控。這樣做得好處在于以 Spark 為大數(shù)據(jù)引擎得機(jī)器學(xué)習(xí)應(yīng)用中,通過 Ray 可以將 Spark 產(chǎn)生得 dataframe 以 ML Dataset 得形式直接從內(nèi)存?zhèn)鹘o下游得機(jī)器學(xué)習(xí)框架,比如 PyTorch。而 Ray 得生態(tài)里得其他組件,比如超參訓(xùn)練(Ray Tune)和推理服務(wù)(Ray Serve),則進(jìn)一步補(bǔ)足了訓(xùn)練階段后續(xù)得一系列工程需求。研發(fā)人員可以拋開繁瑣得上線部署流程,實(shí)現(xiàn)一鍵分布式以及一鍵部署。
Ray 得另一個顯著優(yōu)勢是其簡單通用得 API ,只需在一段函數(shù)上加入ray.remote 得裝飾器,便可將一個單機(jī)程序變成分布式執(zhí)行單元,如下所示:
#declare a Ray task等ray.remotedef fun(a): return a + 1#submit and execute a Ray taskfun.remote()#declare a Ray actor等ray.remoteclass Actor(): def fun(slef, a): return a+1actor = Actor.remote()#execute an actor methodactor.fun.remote()
Ray 中最基礎(chǔ)得概念包括 Task 和 Actor,分別對應(yīng)函數(shù)和類。函數(shù)一般是無狀態(tài)得,在 Ray 里被封裝成 Task,從而被 Ray 得分布式系統(tǒng)進(jìn)行調(diào)度;類一般是有狀態(tài)得,在 Ray 里被映射成一個 Actor。Actor 得表達(dá)性更強(qiáng),能覆蓋大多數(shù)得應(yīng)用程序子模塊?;?Actor 和 Task,Ray 對用戶暴露了資源得概念,即每個 actor 或 task 都可以指定運(yùn)行所需要得資源,這對異構(gòu)得支持從開發(fā)人員得角度變得非常便利。比如:
等ray.remote(num_cpus=1, num_gpus=0.2):def infer(data): return model(data)
當(dāng) task 在被提交執(zhí)行得時候,Ray 得調(diào)度器會去找到一個滿足指定資源需求得節(jié)點(diǎn)。在此同時 Ray 會考慮數(shù)據(jù)得 locality。比如上述例子中得“data”,實(shí)際運(yùn)行中可能會分布在任意一個遠(yuǎn)端得節(jié)點(diǎn)得內(nèi)存里,如果 task 不在數(shù)據(jù)所在得節(jié)點(diǎn)上執(zhí)行,跨節(jié)點(diǎn)得數(shù)據(jù)傳輸就無法避免。而 Ray 可以讓這一類得優(yōu)化變得透明??蚣荛_發(fā)人員也可以利用 Ray 得 API 集成更豐富調(diào)度策略,最終提供給用戶得是非常簡單得 API。Ray 對 Actor 和 Task 還有很多高級得細(xì)粒度控制特性,比如支持 gang-scheduling 得 placement group 等,在此不一一贅述。
Ray 另外得優(yōu)勢在于:
下圖是 RayRTC 得一個早期設(shè)計(jì)規(guī)劃圖和階段一核心部分(DP+Training)得 Actor 封裝流程圖。感謝著重講解階段一,二得設(shè)計(jì)和實(shí)現(xiàn)。其中在階段一中所用到得核心組件包括 Ray Actor Pool 和 RaySGD 等。
“DP+Training” Actor 化流程圖:
其中主要包括 DataProcessing 和 Training 兩個 Stage。每一部分得核心計(jì)算邏輯都用 Ray 得 API 封裝成為 Actor 或 Task。Actor 提交運(yùn)行后通過 Ray 得調(diào)度被放置到合適得節(jié)點(diǎn)上執(zhí)行。Ray 得集群資源可以通過改造后得 Autoscaler 在字節(jié)內(nèi)得 Yarn/K8S 集群上實(shí)現(xiàn)動態(tài)擴(kuò)縮容。
DP 實(shí)現(xiàn)過程中,我們利用 Ray 得 ActorPool 解決了一個因?yàn)閯?chuàng)建 Actor 數(shù)目過多而導(dǎo)致得 OOM 問題。Actorpool 本身相當(dāng)于一個線程池,但 Ray 得 Actorpool 可以被開發(fā)者拓展為更高階得彈性線程池。在 RayRTC 中,給定一組數(shù)據(jù),我們需要解決得核心問題之一是使用多少 Ray 得 actor 是比較高效得。這里得高效指:資源使用高效,性能較優(yōu)且穩(wěn)定性較好(不能 oom)等。最簡單得設(shè)計(jì)方式是 1 對 1,即對于每一個 HDFS 路徑, 都指定一個單獨(dú)得 DP Actor 來進(jìn)行處理。但當(dāng)數(shù)據(jù)量線性增長時,由于缺少內(nèi)存管控而很容易出現(xiàn) OOM。最品質(zhì)不錯得方式是 n 對 1,即用一個 actor,順序處理所有數(shù)據(jù),這樣做顯然無法發(fā)揮 Ray 得分布式能力。比較理想得方式是 n 對 m,即 m 個 actor 處理 n 組數(shù)據(jù)。作為對比,1 對 1 得情況如下:
ray_preprocessor_ret_refs = []for hdfs_file_path in hdfs_file_path_list: my_dp = ray.remote(DP).remote(hdfs_file_path) ray_preprocessor_ret_refs.append(my_dp.__call__.remote())
n 對 m 得情況:
num_cpus = 10actors = [ray.remote(DP).remote() for _ in range(num_cpus)]actor_pool = ray.util.ActorPool(actors)for hdfs_file_path in hdfs_file_path_list: actor_pool.submit(lambda actor, info: actor.__call__.remote(**info), hdfs_file_path)
在生產(chǎn)實(shí)踐中,通過對 m 取一個定值,比如 m=10,可以有效控制內(nèi)存使用并實(shí)現(xiàn) I/O 并行。如前所述,給定一個動態(tài)得 workload,我們也可以對 m 得進(jìn)行彈性支持,類似于 K8S 得 HPA 或 Spark 得 dynamic allocation。不同得是,在 Ray 里,開發(fā)者通過可編程得方式實(shí)現(xiàn)定制化得 dynamic allocation,比較簡單得實(shí)現(xiàn)任意粒度得自動擴(kuò)縮。這一部分得代碼可以參考最新版本得 Ray dataset 中得類似實(shí)現(xiàn)(感謝分享github感謝原創(chuàng)分享者/ray-project/ray/blob/master/python/ray/data/impl/compute.py)。
Training 部分得邏輯由于歷史原因,在字節(jié)得內(nèi)部場景有比較復(fù)雜得深度定制。對此,我們采用了 Ray 社區(qū)第壹版得 Ray SGD(最新得版本中,這一模塊為 Ray Train)對已有訓(xùn)練模塊進(jìn)行封裝。RaySGD 是一個輕量級得分布式訓(xùn)練框架,支持 PyTorch 和 TensorFlow。底層直接集成了 PyTorch 得 DDP 和 Tensorflow 得 MirroredStrategy 來進(jìn)行數(shù)據(jù)并行。RaySGD 通過把訓(xùn)練 worker 用 actor 進(jìn)行封裝,不僅實(shí)現(xiàn)了更靈活得分布式統(tǒng)一調(diào)度,而且與整個 Ray 生態(tài)打通。比如可以與 Ray Tune(超參)和 Ray Serve(推理)直接在 actor 這一粒度上進(jìn)行通信和數(shù)據(jù)傳輸。
數(shù)據(jù)并行得分布式訓(xùn)練相比模型并行和混合并行得模式都要相對簡單。但把一個復(fù)雜得單機(jī)版 NLP 訓(xùn)練框架通過 Ray 封裝為分布式框架,并做到對原代碼侵入性最小,需要處理好以下幾個問題:
- 單節(jié)點(diǎn)得訓(xùn)練邏輯,如何設(shè)置模型,如何在 CPU 和 GPU 之間傳遞數(shù)據(jù)
- 如何設(shè)置 dataloader 以及 sampler,實(shí)現(xiàn)分布式數(shù)據(jù)讀取
- 如何控制一個 epoch 里得 batch 循環(huán)
- 分布式訓(xùn)練邏輯,如何設(shè)置 worker 數(shù)量
- 如何使用 Ray 拉起 worker,并能在 worker 間通信
對于前 3 個問題,RayRTC 實(shí)現(xiàn)了 RayRTCTrainoperator,繼承自 ray.util.sgd.torch 中得 TrainingOperator,把單節(jié)點(diǎn)上得訓(xùn)練邏輯全部抽象到一個類。
class RayRTCTrainOperator(TrainingOperator): def setup(self, config): # Setup data self.train_loader = DataLoader(self.train_data,...) self.valid_loader = DataLoader(self.valid_data,...) # Register data loader self.register_data( train_loader=self.train_loader, validation_loader=self.valid_loader) ... # Register model, optimizer self.model, self.optimizer = \ self.register(models=model, optimizers=optimizer,...)
在 RayRTCTrainOperator 這個類中,首先設(shè)置好訓(xùn)練所需要得模型和數(shù)據(jù),并將 optimizer,scheduler 等參數(shù)傳入。這些數(shù)據(jù)會隨著 RayRTCTrainOperator 這個類被 Ray 封裝為 actor,從而分布到不同得節(jié)點(diǎn)上,從而使得每個節(jié)點(diǎn)上都有一份完全一樣得模型得拷貝和參數(shù)得初始狀態(tài)。
數(shù)據(jù)格式得不同:
除了模型和數(shù)據(jù)得 setup,具體得訓(xùn)練邏輯需要根據(jù) RTC 得場景進(jìn)行定制。比如,每一個 epoch 得訓(xùn)練,以及一個 epoch 中每一個 batch 得訓(xùn)練。由于 RaySGD 對于 input 有一定得格式假設(shè),導(dǎo)致在 RayRTCTrainOperator 中,需要重定義 train_epoch 和 train_batch 這兩個函數(shù)以便正確處理數(shù)據(jù)和 metrics。舉例而言,在 RaySGD 中,batch input 需要符合以下格式:
*features, target = batch
(感謝分享github感謝原創(chuàng)分享者/ray-project/ray/blob/ray-1.3.0/python/ray/util/sgd/torch/training_operator.py#L536)
而實(shí)際得場景中,用戶往往對數(shù)據(jù)格式有自己得定義。比如 RTC 中,batch 被定義為 Dict:
TensorDict = Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]
使用 RaySGD 中默認(rèn)得 train_batch 函數(shù),會在數(shù)據(jù) unpack 時候發(fā)生錯誤。在 RayRTC 中,重寫得 train_batch 把處理后 batch 以正確得格式傳給 forward 函數(shù)。
訓(xùn)練指標(biāo)得自定義問題:
在 train_epoch 中,同樣有需要特殊處理得地方。RaySGD 默認(rèn)支持得 metrics 只包括 loss 等。RTC 中,用戶主要關(guān)心得指標(biāo)包括 accuracy, precision, recall 以及 f1 measure 等。這些指標(biāo)如何在 RaySGD 中加入是 RayRTC 實(shí)現(xiàn)過程中遇到得一個不小得挑戰(zhàn)。一方面由于 RTC 本身已經(jīng)實(shí)現(xiàn)了豐富得 metrics 計(jì)算模塊,一方面 RaySGD 對訓(xùn)練過程中 metrics 得處理有固定得假設(shè)和且封裝在比較底層。RayRTC 最終采取得方法是把 RTC 中得 metrics 計(jì)算模塊復(fù)用到 RaySGD 得 train_epoch 中。另外遇到得一個問題是 RTC 得 metrics 計(jì)算需要把 model 作為參數(shù)傳入,而 RaySGD 中得 model 已經(jīng)被 DDP 封裝,直接傳入會導(dǎo)致出錯。最后,train_epoch 需要加入如下改動:
if hasattr(model, 'module'): metrics = rtc.get_metrics(model.module, ... reset=True)else: metrics = rtc.get_metrics(model, ... reset=True)
改動之后同時兼容了分布式和單機(jī)(沒有被 DDP 封裝)得情況。
RayRTCTrainOperator 可以理解為單機(jī)得訓(xùn)練模塊,到了分布式環(huán)境下,可以通過 TorchTrainer 這個類。如下所示:
trainer = TorchTrainer( training_operator_cls=RayRTCTrainOperator, num_workers=self.num_workers, use_fp16=self.use_fp16, use_gpu=self.use_gpu, ... num_cpus_per_worker=self.cpu_worker)
Trainer 得主要功能是設(shè)置 training worker 得數(shù)量,混合精度,以及 worker 得 cpu 和 gpu。應(yīng)用程序通過 trainer 可以非常簡單地控制整個分布式訓(xùn)練得邏輯:
for epoch in epochs: metrics['train'] = trainer.train() metrics['validate'] = trainer.validate()return metrics
Trainer 得底層邏輯中包括了拉起 worker group(感謝分享github感謝原創(chuàng)分享者/ray-project/ray/blob/8ce01ea2cc7eddd40c2415904fa94198c0fe1e44/python/ray/util/sgd/torch/worker_group.py#L195),每一個worker用actor表達(dá),從而形成一個actor group。RaySGD 也會處理 communication group 得 setup,以及 actor 得失敗重啟。經(jīng)過這些封裝,用戶只需要感謝對創(chuàng)作者的支持跟訓(xùn)練最直接相關(guān)得邏輯,而不需要花過多時間在底層通訊,調(diào)度等分布式邏輯,極大提高了編程效率。
Checkpoint 得問題:
在改造基本完成后,我們用抖音得數(shù)據(jù)進(jìn)行測試,發(fā)現(xiàn)模型在多卡時,沒有任何調(diào)參得情況下,性能已經(jīng)可以與單機(jī)持平,符合上線要求。但第壹次上線測試后,發(fā)現(xiàn) RayRTC 訓(xùn)練出來得模型連基線模型都打不過,準(zhǔn)確率甚至低到 30%。在把所有控制變量固定仍然沒有沒有找到原因后,第壹反應(yīng)是 RayRTC 訓(xùn)練出來得模型可能并沒有真正保存下來,以致線上打分用到得實(shí)際是 pre-trained 得 bert 模型。事實(shí)證明確實(shí)如此,而導(dǎo)致這個原因是因?yàn)?RaySGD 中得 training worker 是在遠(yuǎn)端運(yùn)行,driver 端所初始得數(shù)據(jù)結(jié)構(gòu)隨著訓(xùn)練進(jìn)行會與之逐漸不同步。checkpointing 之前需要取得更新后得模型參數(shù),代碼如下所示:
for epoch in epochs: metrics['train'] = trainer.train() metrics['validate'] = trainer.validate() self.model = trainer.get_model() self.save_checkpoint()return metrics
與之前比較,增加了第 4 行,通過 trainer 獲得更新后得 model,并通過 checkpoint 將模型持久化。
改造侵入性問題:
Anyscale 在一篇博客[感謝分享特別anyscale感謝原創(chuàng)分享者/blog/ray-distributed-library-patterns]中總結(jié)了使用 Ray 得幾種 pattern。其中大致可以分為三類,RayRTC 屬于第三類。
從第壹類到第三類,用 Ray 得層次加深,但并不意味著改造成本線性增加。具體得應(yīng)用需要具體分析。單純從代碼改動量上分析,RayRTC 第壹階段改了大概 2000 行代碼,占原應(yīng)用總代碼量得 1%不到。
同時,RayRTC 把訓(xùn)練模塊單獨(dú)抽象出來,與原有代碼保持松耦合關(guān)系。用戶使用得時候,只需要載入相關(guān) RayRTC 得模塊,即可啟動 Ray 進(jìn)行分布式訓(xùn)練。
實(shí)驗(yàn)效果:
RayRTC 第壹階段在 1 到 8 卡(NV發(fā)布者會員賬號IA V100)上進(jìn)行 scaling 測試,如下圖所示:
訓(xùn)練速度上,RayRTC 得性能隨卡數(shù)呈現(xiàn)線性增加。訓(xùn)練準(zhǔn)確度上,RayRTC 沒有因?yàn)?global batch size 得增加而顯著降低。8 卡訓(xùn)練中,單個 epoch 時間降到了 6 分鐘以內(nèi)。以往研發(fā)人員往往需要等待幾個小時才能拿到訓(xùn)練結(jié)果,導(dǎo)致大家都習(xí)慣在下班前大量提交作業(yè),第二天再來查看效果。整體集群 quota 資源利用率在白天不高,在晚上排隊(duì)高峰。經(jīng)過 RayRTC 提速后,研發(fā)人員會越來越多得進(jìn)行接近交互式得開發(fā)迭代。
RayRTC pipelineRayRTC 在字節(jié)內(nèi)部運(yùn)行在 Arnold 機(jī)器學(xué)習(xí)平臺。用戶在提交一個 RayRTC 任務(wù)時,對應(yīng)在 Arnold 平臺上拉起一個 Trial。一個 Trial 里,用戶配置一個或多個 container 以及每個 container 所需得 CPU/GPU/Mem 資源。在一個 RayRTC 任務(wù)得整個生命周期中,對應(yīng) Trial 得資源是一直占用得。下圖展示了某 RTC 任務(wù)運(yùn)行期間得 GPU 資源使用情況。
如圖所示,在 Data Processing(DP)階段,GPU 資源完全處于 idle 狀態(tài)。造成這個現(xiàn)象得主要原因是當(dāng)前得 RayRTC 階段一方案雖然在 DP 和 Training 階段都充分利用 Ray 得并行能力進(jìn)行加速,但是這兩個 stage 之間本質(zhì)還是串行執(zhí)行:Training 階段必須等到 DP 結(jié)束了才開始。對于 DP 時間長得 RayRTC 任務(wù),這將帶來很大得 GPU 資源浪費(fèi)。為了提高 GPU 資源使用率,我們結(jié)合 Ray Datasets 提供得 pipeline 功能, 提出并實(shí)現(xiàn)了 RayRTC 得流水并行方案 RayRTC pipeline。
Ray Datasets 是在 Ray1.6+版本引入得在 Ray 得 libraries 和應(yīng)用之間加載和交換數(shù)據(jù)標(biāo)準(zhǔn)化方法,其本身提供了一定得基本分布式數(shù)據(jù)處理能力,如 map, filter, repartition 等。如下圖所示,數(shù)據(jù)經(jīng)過 ETL 后,進(jìn)入 ML Training 系統(tǒng)前,可以先通過 Ray Datasets 得 API 進(jìn)行 last mile 得預(yù)處理。換言之,RayRTC 中得 DP 部分,完全可以用 Ray Datasets APIs 這種 Ray 標(biāo)準(zhǔn)化得方式重構(gòu),并與后面得 RaySGD(現(xiàn) Ray Train)打通。
除了提供 last mile 預(yù)處理標(biāo)準(zhǔn)化 APIs, Ray Dataset s 還提供了一組非常重要得 pipeline 接口,使得 DP 部分和 Training 部分得流水并行執(zhí)行成為可能。所謂流水并行執(zhí)行,如下圖所示,Training 執(zhí)行并不會等到 DP 全部結(jié)束后才開始,而是一旦 DP 完成了一小部分就會把處理后得數(shù)據(jù)直接傳入 Training 部分。流水處理有效減少 GPU idle 時間并縮短整個端到端 RTC 訓(xùn)練時間。
基于 Ray Datasets 得 RayRTC pipeline 實(shí)現(xiàn)RayRTC pipeline 版本一:把 DP 部分當(dāng)做黑盒
考慮到 RTC 中 DP 得復(fù)雜邏輯,在 RayRTC pipeline 版本一中,我們把 DP 當(dāng)作黑盒處理。改造需求如下:
- DP(含 IO, trasforms, 數(shù)據(jù)集 split 等邏輯)與 Training 需要以 window 粒度流水并行,其中 DP 得 input 是文件路徑 fp_i,output 是訓(xùn)練和驗(yàn)證數(shù)據(jù)集{'T':Ti, 'V':Vi}。
- DP 中得 split 邏輯要保證多 epoch 訓(xùn)練中每個 epoch 拿到得訓(xùn)練/驗(yàn)證數(shù)據(jù)集都相同,否則會導(dǎo)致數(shù)據(jù)泄露。多 epoch 訓(xùn)練中,只有第壹個 epoch 拿到得訓(xùn)練/驗(yàn)證數(shù)據(jù)集真正經(jīng)歷 DP,其余 epoch 都復(fù)用之前已經(jīng)處理分割好得數(shù)據(jù)集。
為滿足以上需求,我們利用 Ray Datasets 得 API 實(shí)現(xiàn)如下:
dsp= ray.data.from_items([fp1, fp2, …., fpn],parallelism=n).window(blocks_per_window=2).map(dp).repeat().split(2)
但是,以上改造無法滿足“每個訓(xùn)練 worker 拿到相同數(shù)目得 training instances”這個需求,因?yàn)樵摳脑熘械?split 得粒度其實(shí)還是“文件”而非“training instances”,而每個文件中包含得 training instances 數(shù)很可能不一樣。為了滿足這個需求,我們更新實(shí)現(xiàn)如下:
dsp_train= ray.data.from_items([fp1, fp2, …., fpn],parallelism=n).window(blocks_per_window=2).map(dp).flat_map(takeT).repeat().split(2, equal=True)dsp_valid= ray.data.from_items([fp1, fp2, …., fpn],parallelism=n).window(blocks_per_window=2).map(dp).flat_map(takeV).repeat().split(2, equal=True)
其中:
def takeT(row): train_data = row['T'].iter_rows() for data in train_data: yield data.as_pydict()def takeV(row): train_data = row['V'].iter_rows() for data in train_data: yield data.as_pydict()
但是更新后得實(shí)現(xiàn)帶來了新問題:dsp_train 和 dsp_valid 實(shí)際對應(yīng)兩次不同得 DP split 邏輯,從而導(dǎo)致了數(shù)據(jù)泄露。我們需要類似如下實(shí)現(xiàn)來解決:
dsp_train,dsp_valid = ray.data.from_items([fp1, fp2, …., fpn],parallelism=n).window(blocks_per_window=2).map(dp).unzip_and_flat_map('T', 'V').repeat().split(2, equal=True)
其中, unzip_and_flat_map 既有類似 unzip 功能,把原數(shù)據(jù)集分割成兩個數(shù)據(jù)集,原來數(shù)據(jù)集得 Row={'T':Ti, 'V':Vi} 變成兩個新數(shù)據(jù)集得 Row1=Ti,Row2=Vi;又有 flat_map 功能,把數(shù)據(jù)集得 Row1=Ti 真正展開成 Row=Training Instance。考慮到這個 API 實(shí)現(xiàn)復(fù)雜且不具通用性,我們放棄了該版本改造,轉(zhuǎn)向了 RayRTC pipeline 得版本二實(shí)現(xiàn),把 DP 中得數(shù)據(jù)集分割邏輯抽取出來并提前,從開始就構(gòu)造獨(dú)立得訓(xùn)練/驗(yàn)證 pipeline,其余剩下得 DP 邏輯保留。
RayRTC pipeline 版本二:把 DP 中得數(shù)據(jù)集 Split 邏輯抽取出來并提前
在 RayRTC pipeline 版本二實(shí)現(xiàn)中,我們將數(shù)據(jù)集 scaling 和 split 邏輯抽取出來往前移,先構(gòu)造訓(xùn)練和驗(yàn)證數(shù)據(jù)集。然后,分別從這兩個數(shù)據(jù)集構(gòu)造相應(yīng)得訓(xùn)練/驗(yàn)證 pipelines。具體實(shí)現(xiàn)如下:
train_dataset, valid_dataset = self.get_datasets()train_dataset_pipeline = train_dataset.window(blocks_per_window=2).flat_map(dp).repeat().random_shuffle_each_window().split(2, equal=True) # 2 is #trainWorkersvalid_dataset_pipeline = valid_dataset.window(blocks_per_window=2).flat_map(dp).repeat().split(2, equal=True) # 2 is #trainWorkers
其中:
def get_datasets(self): # read dataset from hdfs new_dataset = ray.data.read_api.read_json(partition_info_list) # scale dataset up scaled_dataset = new_dataset.flat_map(scale) # shuffle dataset shuffled_dataset = scaled_dataset.random_shuffle() # split dataset into training and validation datasets train_valid_ratio = 0.9 return shuffled_dataset.split_at_indices([int(shuffled_dataset.count() * train_valid_ratio)])
接著,train_dataset_pipeline 和 valid_dataset_pipeline 被傳入 trainer:在每個 training worker 得 setup() 中,根據(jù)自己得 rank 得到相應(yīng)得子 pipeline。
self.train_dataset_pipeline = self.train_pipeline[self.world_rank]self.train_dataset_pipeline_epoch = self.train_dataset_pipeline.iter_epochs()self.valid_dataset_pipeline = self.valid_pipeline[self.world_rank]self.valid_dataset_pipeline_epoch = self.valid_dataset_pipeline.iter_epochs()
在 training worker 得 train_epoch() 中,從子 training pipeline 中獲取 training instances 訓(xùn)練。
def train_epoch(): dataset_for_this_epoch = next(self.train_dataset_pipeline_epoch) train_dataset = self.data_parser.parse(dataset_for_this_epoch) train_loader = DataLoader(train_dataset) for batch_idx, batch in enumerate(train_loader): metrics = self.train_batch(batch, batch_info)
在 training worker 得 validate() 中, 從子 validation pipeline 中獲取 validation instances 驗(yàn)證。
def validate(): dataset_for_this_epoch = next(self.valid_dataset_pipeline_epoch) valid_dataset = self.data_parser.parse(dataset_for_this_epoch) valid_loader = DataLoader(valid_dataset) for batch_idx, batch in enumerate(valid_loader): metrics = self.validate_batch(batch, batch_info)
實(shí)驗(yàn)效果:
為驗(yàn)證 RayRTC-pipeline 效果,我們隨機(jī)選擇中等規(guī)模 RTC training job (約 168 萬條 instance),使用同等計(jì)算資源(2CPUs, 2GPUs)簡單做了如下對比實(shí)驗(yàn)。結(jié)果顯示,使用 pipeline 后,GPU idle 時間從原來得 245s 減少到了 102s,約 2.5 倍降低。端到端時間也比原來減少了 158s。除此之外,相比于階段一實(shí)現(xiàn),我們不但在初始階段對整個數(shù)據(jù)集進(jìn)行 random_shuffle,在每個 window 得訓(xùn)練數(shù)據(jù)從 pipeline 出來時,也通過 random shuffle 對 window 中得訓(xùn)練數(shù)據(jù)再次進(jìn)行 shuffle。結(jié)果顯示,充分得全局和局部 shuffle 有效提高模型精度。
VersionAccuracyPrecisionRecallf1-measureGPU idle timeE2E timeRayRTC-phase10.8040.6370.5710.602245s2296sRayRTC-pipeline0.8210.7150.5560.625102s2138sImprove+0.017+0.078-0.015+0.023-143s-158s
總結(jié)RayRTC 以 Ray 為分布式計(jì)算學(xué)習(xí)引擎,對字節(jié) RTC NLP 框架得全面改造升級不僅實(shí)現(xiàn)了性能得極致優(yōu)化(5 小時到 30 分鐘),同時通過流水并行極大降低了 GPU 資源得 idle 時間(60% reduction)。RayRTC 以松耦合得形式對現(xiàn)有業(yè)務(wù)得侵入極小(<1% loc),同時為后續(xù)可插拔 low-level 優(yōu)化和 serverless autoscaling 提供了 API 支持。可以預(yù)見,后續(xù) RayRTC 在更大規(guī)模上進(jìn)行超參以及與推理打通,將會形成更高效得端到端 Serverless NLP Pipeline。
加入我們字節(jié)跳動應(yīng)用機(jī)器學(xué)習(xí)(AML)團(tuán)隊(duì)負(fù)責(zé)給公司提供統(tǒng)一得機(jī)器學(xué)習(xí)中臺服務(wù)。我們?yōu)楸菊?抖音/西瓜視頻得推薦/廣告/搜索/審核/OCR/消重等業(yè)務(wù)提供大規(guī)模訓(xùn)練系統(tǒng)和推理系統(tǒng),并在系統(tǒng)上提供簡單易用得機(jī)器學(xué)習(xí)平臺,提高算法調(diào)研效率。同時,AML 團(tuán)隊(duì)自有得機(jī)器學(xué)習(xí)算法團(tuán)隊(duì),同業(yè)務(wù)算法團(tuán)隊(duì)合力研究業(yè)務(wù)通用算法、改善業(yè)務(wù)實(shí)際效果。此外,AML 團(tuán)隊(duì)也在積極探索對外部企業(yè)提供機(jī)器學(xué)習(xí)服務(wù),并研究科學(xué)計(jì)算和聯(lián)邦學(xué)習(xí)等技術(shù)及其應(yīng)用。
加入 AML 團(tuán)隊(duì),可以參與到推薦、廣告、CV、NLP、Speech 等多種場景,掌握業(yè)界先進(jìn)、第壹手得機(jī)器學(xué)習(xí)算法與平臺實(shí)踐經(jīng)驗(yàn)。我們得 base 包括北京、上海、杭州、西雅圖、山景城,歡迎通過下方鏈接加入我們!
ByteDance Infrastructure System Lab(字節(jié)基礎(chǔ)架構(gòu)美研)
The Infra lab is responsible for the design and implementation of ByteDance’s critical infrastructure systems, including but not limited to compute, storage, database, and networking, etc. We welcome strong candidates (Researchers and Engineers) to join us to work on cutting edge research to support the company's business (Douyin, Toutiao, Helo and Resso, etc.) and touch the lives of billions of users across the world.
We are based in Mountain View and Seattle, welcome to apply using the following link: 感謝分享job.nnzzn.com/s/FTaJVUh