得物AI平臺-KubeAI推理訓練引擎設計和實踐
1、KubeAI介紹
KubeAI是得物AI平臺，是我們在容器化過程中，逐步收集和挖掘公司各業務域在AI模型研究和生產迭代過程中的需求，逐步建設而成的一個云原生AI平臺。KubeAI以模型為主線提供了從模型開發，到模型訓練，再到推理(模型)服務管理，以及模型版本持續迭代的整個生命周期內的解決方案。
在數據方面，KubeAI提供基于cvat的標注工具，與數據處理及模型訓練流程打通，助力線上模型快速迭代；提供任務/Pipeline編排功能，對接ODPS/NAS/CPFS/OSS數據源，為用戶提供一站式AI工作站。平臺自研推理引擎助力業務在提高模型服務性能的同時還能控制成本；自研訓練引擎提高了模型訓練任務吞吐量，縮短了模型的訓練時長，幫助模型開發者加速模型迭代。
此外，隨著AIGC的火熱發展，我們經過調研公司內部AI輔助生產相關需求，上線了AI制圖功能，為得物海報、營銷活動、設計師團隊等業務場景提供了基礎能力和通用AI制圖能力。
此前，我們通過一文讀懂得物云原生AI平臺-KubeAI的落地實踐過程一文，向大家介紹了KubeAI的建設和在業務中的落地過程。本文，我們將重點介紹下KubeAI平臺在推理、訓練和模型迭代過程中的核心引擎能力實踐經驗。
2、AI推理引擎設計實現
2.1 推理服務現狀及性能瓶頸分析
Python語言以其靈活輕盈的特點，以及其在神經網絡訓練與推理領域提供了豐富的庫支持，在模型研究和開發領域被廣泛使用，所以模型推理服務也主要以Python GPU推理為主。模型推理過程一般涉及預處理、模型推理、后處理過程，單體進程的方式下CPU前/后處理過程，與GPU推理過程需要串行，或者假并行的方式進行工作，大致流程如下圖所示：
上述架構的優勢是代碼寫起來比較通俗易懂，但在性能上有很大的弊端，所能承載的QPS比較低。通過在CV域的模型上進行壓測，我們發現推理QPS很難達到5，深入分析發現造成這一問題的原因如下：
（1）單線程模式下，CPU邏輯與GPU邏輯相互等待，GPU Kernel函數調度不足，導致GPU使用率不高，無法充分提升服務QPS。這種情況下只能開啟更多進程來提升QPS，但是更多進程會帶來更大的GPU顯存開銷。
（2）多線程模式下，由于Python的GIL鎖的原因，Python的多線程實際上是偽的多線程，并不是真正的并發執行，而是多個線程通過爭搶GIL鎖來執行，這種情況下GPU Kernel Launch線程不能得到充分的調度。此外，在Python推理服務中開啟多線程反而會導致GPU Kernel Launch線程頻繁被CPU的線程打斷，所以GPU算力也會一直“萎靡不振”，持續低下。
以上問題使得 如果推理服務想要支撐更多的流量，只能做橫向的增加服務實例數，伴隨著成本的上漲。
2.2 自研推理服務統一框架kubeai-inference-framework
針對以上問題，KubeAI的解決方案是把CPU邏輯與GPU邏輯分離在兩個不同的進程中：CPU進程主要負責圖片的前處理與后處理，GPU進程則主要負責執行CUDA Kernel 函數，即模型推理。
為了方便模型開發者更快速地接入我們的優化方案，我們基于Python開發了一個CPU與GPU進程分離的統一框架kubeai-inference-framework，舊有Flask或Kserve的服務，稍作修改即可接入推理引擎統一框架，新增服務按照框架實現指定function即可。推理服務統一框架構如下圖所示：
如前所述，推理服務統一框架的主要思路是把GPU邏輯與CPU邏輯分離到兩個進程，除此之外，還會拉起一個Proxy進程做路由轉發。
CPU進程
CPU進程主要負責推理服務中的CPU相關邏輯，包括前處理與后處理。前處理一般為圖片解碼，圖片轉換，后處理一般為推理結果判定等邏輯。CPU進程在前處理結束后，會調用GPU進程進行推理，然后繼續進行后處理相關邏輯。CPU進程與GPU進程通過共享內存或網絡進行通信，共享內存可以減少圖片的網絡傳輸。
GPU進程
GPU進程主要負責運行GPU推理相關的邏輯，它啟動的時候會加載很多模型到顯存，然后在收到CPU進程的推理請求后，直接觸發Kernel Lanuch調用模型進行推理。
kubeai-inference-framework框架中對模型開發者提供了一個Model類接口，他們不需要關心后面的調用邏輯，只需要填充其中的前處理，后處理的業務邏輯，就可以快速上線模型服務，自動拉起這些進程。
Proxy進程
Proxy進程是推理服務入口，對外提供調用接口，負責路由分發與健康檢查。當Proxy進程收到請求后，會輪詢調用CPU進程，分發請求給CPU進程進行處理。
自研的推理服務統一框架，把CPU邏輯（圖片解碼，圖片后處理等）與GPU邏輯（模型推理）分離到兩個不同的進程中后，有效解決了Python GIL鎖帶來的GPU Kernel Launch調度問題，提升了GPU利用率，提高了推理服務性能。針對線上的某個推理服務，使用我們的框架進行了CPU與GPU進程分離，壓測得出的數據如下表所示，可以看到QPS提升了近7倍。
推理服務框架類型 | QPS | 耗時(s) | GPU算力使用率(%) |
傳統多線程架構 | 4.5 | 1.05 | 2.0 |
自研推理服務框架（6個CPU進程+1個GPU進程） | 27.43 | 0.437 | 12.0 |
2.3 做的更好 — 引入TensorRT優化加速
在支持推理服務接入kubeai-inference-framework統一框架的過程中，我們繼續嘗試在模型本身做優化提升。經過調研和驗證，我們將現有pth格式模型通過轉成TensorRT格式，并開啟FP16，在推理階段取得了更好的QPS提升，最高可到10倍提升。
TensorRT是由英偉達公司推出的一款用于高性能深度學習模型推理的軟件開發工具包，可以把經過優化后的深度學習模型構建成推理服務部署在實際的生產環境中，并提供基于硬件級別的推理引擎性能優化。業內最常用的TensorRT優化流程，是把pytorch / tensorflow等模型先轉成onnx格式，然后再將onnx格式轉成TensorRT（trt）格式進行優化，如下圖所示：
TensorRT所做的工作主要在兩個時期，一個是網絡構建期，另外一個是模型運行期。
- 網絡構建期
- 模型解析與建立，加載onnx網絡模型。
- 計算圖優化，包括橫向算子融合，或縱向算子融合等。
- 節點消除，去除無用的節點。
- 多精度支持，支持FP32/FP16/int8等精度。
- 基于特定硬件的相關優化。
- 模型運行期
- 序列化，加載RensorRT模型文件。
- 提供運行時的環境，包括對象生命周期管理，內存顯存管理等
為了更好地幫助模型開發者使用TensorRT優化，KubeAI平臺提供了kubeai-trt-helper工具，用戶可以使用該工具把模型轉成TensorRT格式，如果在模型轉換的過程中出現精度丟失等問題，也可以使用該工具進行問題定位與解決。kubeai-trt-helper主要在兩個階段為用戶提供幫助：一個是問題定位，另一個階段是模型轉換。
問題定位
問題定位階段主要是為了解決模型轉TensorRT開啟FP16模式時出現的精度丟失問題。一般分類模型，對精度的要求不是極致的情況下，盡量開啟FP16，FP16模式下，NVIDIA對于FP16有專門的Tensor Cores可以進行矩陣運算，相比FP32來說吞吐量提升一倍以上。比如在轉TensorRT時，開啟FP16出現了精度丟失問題，kubeai-trt-helper工具在問題定位階段的大致工作流程如下：
第1步：設定模型轉換精度要求后，標記所有算子為輸出，然后對比所有算子的輸出精度。
第2步：找到最早的不符合精度要求的算子，對該算子進行如下幾種方式干預。
- 標記該算子為FP32。
- 標記其父類算子為FP32。
- 更改該算子的優化策略。
循環通過以上2個步驟，最終找到符合目標精度要求的模型參數。這些參數比如：需要額外開啟FP32的那些算子等。相關參數會輸出到配置文件中，如下：
配置項 | 釋義 |
FP32_LAYERS_FOR_FP16 | 開啟FP16模式下，哪些算子需要額外開啟FP32 |
TRT_EXCLUDE_TACTIC | TensorRT算子需要忽略的tactic策略（tactic可參考TensorRT相關資料） |
atol | 相對誤差 |
rtol | 絕對誤差 |
check-error-stat | 誤差的計算方法包括：mean, median, max |
模型轉換
模型轉換階段則直接使用上面問題定位階段得到的參數，調用TensorRT相關接口與工具進行轉換。此外，我們在模型轉換階段，針對TensorRT原有參數與API過于復雜的問題也做了一些封裝，提供了更為簡潔的接口，比如工具可以自動解析onnx，判斷模型的輸入與輸出shape，不需要用戶再提供相關shape信息等。
2.4 落地實踐成果
在實際應用中，我們幫助算法域的模型開發同學，能夠對一個推理基于自研推理服務統一框架進行實現的同時，也開啟TensorRT優化，這樣往往可以得到QPS兩次優化的疊加效果。
2.4.1 分類模型，CPU與GPU分離，TensorRT優化，并開啟FP16，得到10倍QPS提升
線上某個基于Resnet的分類模型，對精度損失可以接受誤差在0.001（誤差定義：median,atol,rtol）范圍內。因此我們對該推理服務進行了3項性能優化：
- 使用kubeai-inference-framework統一框架，對CPU進程和GPU進程進行分離改造。
- 對模型轉ONNX后，轉TensorRT。
- 開啟FP16模式，并使用自研工具定位到中間出現精度損失的算子，把這些算子標記為FP32。
經過以上優化，最終得到了10倍QPS的提升（與原來Pytorch直接推理比較），服務成本大幅削減。
2.4.2 檢測模型，CPU與GPU分離，TensorRT模型優化，QPS提升4-5倍左右。
線上某個基于Yolo的檢查模型，由于對精度要求比較高，所以不能開啟FP16，我們直接在FP32的模式下進行了TensorRT優化，并使用kubeai-inference-framework統一框架對GPU進程與CPU進程分離，最終得到QPS 4-5倍的提升。
2.4.3 模型推理進程多實例化，充分利用GPU算力資源
在實際的場景中，往往GPU的算力是充足的，而GPU顯存是不夠的。經過TensorRT優化后，模型運行時需要的顯存大小一般會降低到原來的1/3到1/2。所以為了充分利用GPU算力，kubeai-inference-framework統一框架進一步優化，支持可以把GPU進程在一個容器內復制多份，這種架構即保證了CPU可以提供充足的請求給GPU，也保證了GPU算力充分利用。
線上某個模型，經過TensorRT優化后，顯存由原來的2.4G降低到只需要1.2G。在保持推理服務配置5G顯存不變的情況下，我們將GPU進程為復制4份，充分利用了5G顯存，使得服務吞吐達到了原來的4倍。
3、AI訓練引擎優化實踐
3.1 PyTorch框架概況
PyTorch是近年來較為火爆的深度學習框架，幾乎占據了CV（Computer Vision，計算機視覺）、NLP（Natural Language Processing，自然語言處理）領域各業務方向，算法同學基本都在使用PyTorch框架來進行模型訓練。下圖是基于PyTorch框架進行模型訓練時的代碼基本流程：
第1步：從pytorch dataloader中將本step訓練過程中需要的數據拉出來。
第2步：將獲取到的數據，例如：樣本圖片、樣本標簽的tensor等數據，復制到GPU顯存里。
第3步：開始正式的模型訓練：前向計算、計算損失、計算梯度、 更新參數。
整個訓練過程的耗時，也主要分布在上面3個步驟。通常第2步不會是瓶頸，因為大部分訓練樣本圖片都是被resize變小之后才從內存拷貝到到GPU顯存上的。但由于模型的差異性、訓練數據的差異性，經常是第1、2步會在訓練過程中出現性能瓶頸，導致訓練耗時長，GPU利用率低下，影響模型迭代效率。
3.2 Dataloader瓶頸分析及優化
3.2.1 PyTorch Dataset/Dataloader分析
PyTorch訓練讀取數據部分主要是通過Dataset、Dataloader的方式完成的，其中Dataset為用戶自定義讀取數據的類（繼承自 torch.utils.data.Dataset），而Dataloader是PyTorch實現的在訓練過程中對Dataset的調度器。
torch.utils.data.dataloader
torch.utils.data.Dataset
train_loader = torch.utils.data.DataLoader( MyDataset,
batch_size=16,
num_workers=4,
shuffle=True,
drop_last=True,
pin_memory=False)
val_loader = torch.utils.data.DataLoader(MyDataset,
batch_size=batch_size,
num_workers=4,
shuffle=False)參數解釋如下：
- dataset(Dataset)：傳入的自定義Dataset（數據讀取的具體步驟）。
- batch_size(int, optional)：每個batch有多少個樣本，每個iter可以從dataloader中取出多少數據。
- shuffle(bool, optional)：在每個epoch開始的時候，對數據進行重新排序，可以使每個epoch讀取數據的組合和順序不同。
- num_workers (int, optional)：這個參數決定dataloader啟動幾個后臺進程來做數據拉取。0意味著所有的數據都會被load進主進程，默認為0。
- collate_fn (callable, optional)：將一個list的sample組成一個mini-batch的函數，一般CV場景是concat函數。
- pin_memory (bool, optional)：如果設置為True，那么data loader將會在返回batch之前，將tensors拷貝到CUDA中的固定內存（CUDA pinned memory）中， 這個參數某些場景下有妙用。
- drop_last (bool, optional)：該參數是對最后的未完成的batch來說的，比如batch_size設置為64，而一個epoch只有100個樣本，如果設置為True，那么訓練的時候后面的36個就被扔掉了，否則會繼續正常執行，只是最后的batch_size會小一點。默認設置為False。
上述參數中，比較重要的是num_workers，Dataloader在構造的時候，會啟動num_workers個worker進程，然后主進程會向worker進程分發讀取任務，worker進程讀到數據之后，再把數據放到隊列中供主進程取用。多進程模式使用的是torch.multiprocessing接口，可以實現worker進程與主進程之間共享內存，而且共享內存中可以存放tensor，這樣進程中如果返回tensor，可以通過共享內存的方式直接將結果返回給主進程，減少多進程間的通訊開銷。
當num_workers 為0 的時候： get_data()流程與train_model()過程是串行，效率非常低下，如下圖所示：
當num_workers 大于0開啟多進程讀取數據， 并且讀取一個batch數據的時間小于一個step訓練的時間時效率最高，GPU算力被充分利用，如下圖所示：
當num_workers 大于0開啟多進程度數據， 但是讀取一個batch數據的時間大于一個step訓練的時間時，會出現GPU訓練過程等待數據拉取，就會出現GPU算力空閑，訓練耗時增加，如下圖所示：
由此可見Dateset中的__getitem__函數非常重要，詳細分析它的源碼實現后我們發現，該函數的耗時主要包含2段時間：
- load_image_time：從磁盤或者遠程盤上讀取數據的耗時。
- transform_image_time：將圖片或文本數據進行預處理的耗時。
3.2.2 解問題 — 設置合理的參數很重要
通過上一小節的分析，訓練時相關參數的選擇至關重要?？偨Y如下：
- batch_size：根據數據量，以及期望訓練時長，用戶合理自定義設置
- 訓練環境（KubeAI Notebook/任務/流水線節點）的CPU配置：建議CPU配置為 GPU卡數*(單GPU卡配置的CPU核數)。
- num_workers：參數最小設置為 訓練環境的CPU配置-1，比如：任務配置為12C時，建議該參數設置為11 。另外，num_workers數值可以適當調大，因為dataset iter中有部分時間是在網絡或者磁盤IO， 這部分不消耗CPU；但是也不能設置太大，因為數據預處理部分是CPU密集型任務，并行進程過多，會造成CPU爭搶從而降低預處理效率。
優化案例一
線上一個基于MMDetection框架（其底層也是調用PyTorch框架）的CV模型訓練任務，在做參數調整之前，單個step耗時不穩定，平均在1.12s左右，其中拉取數據時長在0.3s左右：
mmengine - INFO - Epoch(train) [2][3050/6005] time: 1.1128 data_time: 0.1032
mmengine - INFO - Epoch(train) [2][3100/6005] time: 1.0193 data_time: 0.0055
mmengine - INFO - Epoch(train) [2][3150/6005] time: 1.0928 data_time: 0.3230
mmengine - INFO - Epoch(train) [2][3200/6005] time: 0.9927 data_time: 0.2304
mmengine - INFO - Epoch(train) [2][3250/6005] time: 1.3224 data_time: 0.5135
mmengine - INFO - Epoch(train) [2][3300/6005] time: 1.1044 data_time: 0.3427
mmengine - INFO - Epoch(train) [2][3350/6005] time: 1.0574 data_time: 0.2842調整參數之后，單個step耗時穩定，平均在0.78 s左右，其中拉取數據耗時0.004s，基本可以忽略。
mmengine - INFO - Epoch(train) [1][100/5592] time: 0.8508 data_time: 0.0049
mmengine - INFO - Epoch(train) [1][150/5592] time: 0.7743 data_time: 0.0043
mmengine - INFO - Epoch(train) [1][200/5592] time: 0.7736 data_time: 0.0044
mmengine - INFO - Epoch(train) [1][250/5592] time: 0.7880 data_time: 0.0044
mmengine - INFO - Epoch(train) [1][300/5592] time: 0.7761 data_time: 0.0045該模型訓練任務，通過上述優化調整，數據拉取時間縮短為0，單個step的耗時從原來的1.12s降到0.78s，整體訓練時間減少30%（從2天縮短到33小時），效果顯著。
優化案例二
線上某個多模態模型（輸入包含圖片和文字）訓練任務，使用2卡V100訓練，參數調整如下：
batch_size=32
CPU = 12 ---> 調整為 24
num-workers = 4 ---> 調整為 11調整后訓練300 step總消耗時405s，整體訓練時間減少45%左右（從10天縮短到5天左右）。
優化案例三
線上某YoloX模型訓練任務，使用單卡A100訓練，參數調整如下：
batch size ：48
num-workers = 4 ---> 調整為 16調整后整體訓練時長減少80%左右（從10天19小時，縮短至1天16小時）。
3.2.3 數據拉取IO瓶頸分析
當前，KubeAI平臺為訓練場景提供3種存儲介質：
- 本地盤：空間小，讀寫性能最好，單盤500G~3T空間可用。
- NAS網絡存儲：空間大，讀寫性能較差，成本適中。
- CPFS并行文件系統存儲：空間大，讀寫性能好，成本高。
對于小數據集，可以先將數據一次性拉取到本地盤，然后每個epoch從本地盤來讀數據，這樣避免了每一個epoch重復的從遠程NAS來拉取數據，相當于整個訓練只需要從遠程NAS拉取一次數據。對于大數據集，有2種解決方案：
- 將大數據集提前進行resize，存儲比較小的圖片來進行訓練，這樣避免了每個epoch都需要resize，而且resize之后，圖片變小，讀取更快。
- 將數據集放入并行文件系統CPFS存儲上，提高訓練吞吐。實驗表明CPFS 在圖片場景下是NAS盤讀性能的3～6倍。
3.3 TrainingModel優化
數據部分優化后，訓練過程中的主要時間開銷就在GPU訓練部分了。目前業內有一些比較成熟的方法可以參考，我們總結如下。
3.3.1 混合精度訓練（AMP）
PyTorch混合精度訓練在PyTorch官網有詳細介紹，以及開啟混合精度訓練的方法，可以閱讀這里獲取實現方法。當前許多CV訓練框架已經支持AMP訓練，比如：
- MMCV框架中AMP參數就是開啟混合精度訓練的選項。
- Pytorch Vision中也有相關參數來開啟AMP訓練。
需要說明的是，混合精度訓練過程中并不是將所有模型參數都轉為FP16來計算，只有部分做轉換?；旌暇戎阅芗铀儆柧氝^程，是因為大部分英偉達GPU機型在FP16這種數據格式的浮點算力比FP32要快一倍；此外，混合精度訓練顯存占用會更小。
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
with autocast(device_type='cuda', dtype=torch.float16):
output = model(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
# Unscales the gradients of optimizer's assigned params in-place
scaler.unscale_(optimizer)
# Since the gradients of optimizer's assigned params are unscaled, clips as usual:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
# optimizer's gradients are already unscaled, so scaler.step does not unscale them,
# although it still skips optimizer.step() if the gradients contain infs or NaNs.
scaler.step(optimizer)
# Updates the scale for next iteration.
scaler.update()3.3.2 單機多卡數據并行訓練
Pytorch原生支持多卡數據并行訓練，詳細開啟多卡訓練的方式參考官方文檔。多卡訓練過程中每一張卡的backword計算會多加一次多卡之間集合通訊all-reduce操作，用來計算多張卡上的梯度的平均值。
3.4 自研訓練引擎框架kubeai-training-framework
通過前面的分析我們可以看到，雖然PyTorch框架本身已經做的很好了，訓練方式、參數支持豐富，但在實際的模型研究和生產過程中，由于模型的差異性、訓練數據的差異性，以及模型開發者的經驗差異性，PyTorch框架本身的優勢不一定能夠發揮出來。
基于前述分析和實踐，KubeAI平臺開發了訓練引擎框架kubeai-training-framework，幫助模型開發者更好地匹配訓練腳本參數，快速接入使用合適的訓練方式。kubeai-training-framework中包含PyTorch Dataloader優化、GPU TrainModel（AMP）提速以及各種功能函數等。以Dataloader為例，用戶可通過以下方式使用：
import torch
from kubeai_training_framework.dataloader import Dataloader
def train(train_loader, model, criterion, optimizer, epoch):
train_dataset = .......
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True)
model.train()
my_train_loader = Dataloader(train_loader)
input, target = my_train_loader.next()
while input is not None:
## model train 代碼 input, target
..........
input, target = my_train_loader.next()4、AI Pipeline引擎助力AI業務快速迭代
通常模型的開發可以歸納為如下圖所示的過程：
可以看到，在需求場景確定、第一個模型版本上線之后，模型是需要反復迭代的，以期望取得更好的業務效果。KubeAI平臺在迭代建設的過程中，逐步上線了Notebook、模型管理、訓練任務管理、推理服務管理等一個個相對獨立的功能模塊。隨著業務需求的不斷變化，模型迭代效率直接影響了業務的上線效率，KubeAI平臺建設了AI Pipeline能力，重點解決AI場景的周期性迭代類需求，提高生產效率。
AI Pipeline是在ArgoWorkflow基礎上做了二次開發，以滿足模型迭代、推理任務管理、數據處理等對定時需求、任務啟動觸發方式、通用模板任務、指定節點啟動等需求。AI Pipeline上線之前，一個迭代任務可能會被配置為多個分散的任務，維護工作量大，調試周期長。如下圖是做一個類似任務需要單獨配置的任務情況：
AI Pipeline可以將整個工作流設計成如下圖所示：
Pipeline編排的方式，減少了模型開發者浪費在重復工作上的時間，可以將更多的時間投入到模型研究上。同時，通過合理編排任務，可以對有限的資源進行充分地利用。
5、展望
KubeAI平臺從得物AI業務場景的實際需求出發，以三大核心引擎為建設目標，著力解決AI模型研發過程中的訓練、推理性能問題，以及模型版本迭代過程中的效率問題。
在推理服務性能上，我們會以kubeai-inference-framework為起點，繼續在模型量化、算子優化、圖優化等方面進行深入探索。在模型訓練方面，我們會繼續在圖像數據預處理、Tensorflow GPU訓練框架支持、NLP模型訓練支持上發力，以kubeai-training-framework訓練引擎框架為接口，為模型開發者提供更高效、性能更高的訓練框架。此外，AI Pipeline引擎上，我們會支持更豐富的預置模型，以滿足通用數據處理任務、推理任務等需求。
