智能計算系統

第五章編程框架原理中國科學院計算技術研究所李威副研究員liwei2017@2編程框架設計計算圖構建計算圖執行*深度學習編譯*分布式訓練本章小結提綱學習智能計算系統，了解編程框架的原理大有脾益編寫與框架底層更為契合、性能更優的代碼定制化擴展編程框架，為新算法、新設備提供支持智能計算系統中編程框架的四大模塊必備：計算圖構建模塊和計算圖執行模塊追求更高性能：深度學習編譯模塊和分布式訓練模塊3概述4編程框架設計計算圖構建計算圖執行*深度學習編譯*分布式訓練本章小結提綱1、設計原則5簡潔性（Simplicity）框架提供一套抽象機制，用戶僅需關心算法本身和部署策略易用性（Usability）高效性（Performance）簡潔抽象輸入數據模型輸出結果addbbaccuda()GPUCPU1、設計原則6簡潔性（Simplicity）易用性（Usability）熟悉的開發范式：如PyTorch始于Python，忠于Python直觀且用戶友好的接口：如PyTorch提供了命令式的動態圖編程方法高效性（Performance）1、設計原則7簡潔性（Simplicity）易用性（Usability）高效性（Performance）如采用靜態圖編程方式，可以生成完整的計算圖并進行全局優化，從而盡量提高用戶應用程序的運行效率支持深度學習編譯技術，多層級表示優化，充分利用用戶硬件的計算能力支持多機多卡條件的分布式訓練，從而高效支持大規模深度學習任務2、整體架構8四大模塊計算圖構建模塊：完成從輸入的用戶程序到編程框架內部原始計算圖的轉換過程，編程框架的入口模塊分布式訓練模塊：應對更大規模的神經網絡，將訓練、推理任務從一臺設備擴展到多臺設備深度學習編譯模塊：對計算圖分別進行圖層級和算子層級的編譯優化，從而提升單設備上的執行效率計算圖執行模塊：將優化后的計算圖中的張量和操作映射到指定設備上進行具體執行，并給出編程框架的輸出結果計算圖構建計算圖執行分布式訓練深度學習編譯dataMatmulweightoutput原始計算圖優化后的計算圖data_0OptMatmul_0weight_0output_0設備0data=torch.tensor(…)weight=torch.tensor(…)output=torch.matmul(data,weight)用戶程序編程框架輸出

output…拆分后的計算圖Matmul_0weight_0output_0設備0data_0Matmul_1weight_1output_1設備1data_1通信2、整體架構10編程框架設計計算圖構建計算圖執行*深度學習編譯*分布式訓練本章小結提綱正向圖與反向圖構建11計算圖由兩個基本元素構成：張量（Tensor）和張量操作（Operation）。計算圖是有向圖，有向邊指明了張量的流動方向1、正向傳播12輸入張量經過搭建的神經網絡層層計算傳遞，并最終獲得計算結果的過程構建形式動態圖：在執行函數時，按照函數順序逐條語句地生成節點，立即計算并返回結果；易調試但性能優化空間有限靜態圖：在執行計算之前構建好所有圖上的節點，在圖運行時才計算整個計算圖并返回最終結果；不易調試但性能好動態圖計算圖在函數運行過程中逐步構建的（On-the-fly）立即（eager）模式：每次調用語句就立刻執行計算PyTorch中的動態圖實現：每次執行，都會重新被構建W_h=torch.randn(20,20,requires_grad=True)W_x=torch.randn(20,10,requires_grad=True)x=torch.randn(1,10)prev_h=torch.randn(1,20)h2h=torch.matmul(W_h,prev_h.t())i2h=torch.matmul(W_x,x.t())next_h=h2h+12hnext_h=next_h.tanh()loss=next_h.sum()13

prev_h

xmatmulAddTanhnext_hlossh2hi2hmatmul靜態圖整個網絡的結構會在開始計算前就建立完成計算圖框架執行時接收整個計算圖而不是單一語句TensorFlow1.x中的靜態圖使用若干基本控制流算子（Switch、Merge、Enter、Exit和NexIteration）的不同組合來實現各種復雜控制流場景PyTorch2.0中的靜態圖PyTorch2.0中采取了圖捕獲（TorchDynamo）的技術將用戶的動態圖轉化為靜態圖142、反向傳播15正向計算得到的結果和目標結果存在損失函數值，對其求導得到梯度，并使用該梯度更新參數wz*

MultBackward反向正向y1y2

計算導數的方法16自動微分是一種計算導數的方法常見的求導方式手動求導：用鏈式法則求解出梯度公式，然后根據公式編寫代碼、代入數值計算得到梯度結果數值求導：直接代入數值近似求解符號求導：直接對代數表達式求解，最后才代入問題數字，出現表達式膨脹問題自動求導：用戶只需描述前向計算的過程，由編程框架自動推導反向計算圖，先建立表達式，再代入數值計算17手動用鏈式法則求解出梯度公式，代入數值，得到最終梯度值缺點：對于大規模的深度學習算法，手動用鏈式法則進行梯度計算并轉換成計算機程序非常困難需要手動編寫梯度求解代碼，且模型變化，算法也需要修改addmatmulbwx

addmatmulbwx

前向傳播

反向傳播手動求解法18數值求導法

19符號求導法

表達式膨脹示例20自動求導法對基本算子應用符號求導法代入數值，保留中間結果應用于整個函數數值求導法：直接代入數值近似求解符號求導法：直接對代數表達式求解，最后才代入問題數字自動求導法：介于數值求導和符號求導的方法21

22exp+*+1

前向計算

23exp+*+1

+

反向計算

24求導方式對比方法對圖的遍歷次數精度備注手動求解法NA高實現復雜數值求導法nI+1低計算量大，速度慢符號求導法NA高表達式膨脹自動求導法nO+1高對輸入維度較大的情況優勢明顯

AutoGrad是PyTorch的自動微分引擎，用戶只需要一行代碼tensor.backward()，即可調用其自動計算梯度并反向傳播AutoGrad模塊的backward函數實現1）正向圖解析2）構建反向計算圖的節點3）進行反向梯度傳播PyTorch中的自動求導AutoGrad模塊的backward函數實現1）正向圖解析PyTorch中的自動求導//創建根節點和梯度的列表，并預分配num_tensors大小的空間std::vector<Edge>roots;//反向傳播根節點集合variable_listgrads;//反向傳播的梯度集合

for(inti=0;i<num_tensors;i++){

//獲取正向圖的輸出張量

at::Tensortensor=py::handle(PyTuple_GET_ITEM(tensors,i)).cast<at::Tensor>();

//獲取從梯度函數指向輸出結果的邊

autogradient_edge=torch::autograd::impl::gradient_edge(tensor);

roots.push_back(std::move(gradient_edge));//將獲取的邊加入到根節點集合中

at::Tensorgrad_tensor=py::handle(PyTuple_GET_ITEM(grad_tensors,

i)).cast<at::Tensor>();

autograd_var=torch::autograd::make_variable(grad_tensor);

grads.push_back(grad_var);//將梯度變量加入到梯度集合中}AutoGrad模塊的backward函數實現2）構建反向計算圖的節點PyTorch中的自動求導std::vector<Edge>output_edges;//反向計算圖中所有邊的集合if(inputs!=nullptr){

for(inti=0;i<num_inputs;++i){//初始化列表

...

constautooutput_nr=tensor.output_nr();

autograd_fn=tensor.grad_fn();

if(!grad_fn){//沒梯度函數，則標記是葉子節點

output_edges.emplace_back(std::make_shared<Identity>(),0);

}else{//有梯度函數，創建梯度函數指向該節點的邊（構造反向計算圖）

output_edges.emplace_back(grad_fn,output_nr);

}

}}AutoGrad模塊的backward函數實現3）進行反向梯度傳播PyTorch中的自動求導//反向計算圖已經構建完成，可進行執行//roots中包含了反向傳播根節點//grads中包含了反向傳播產生的梯度，output_edges中是構建的反向計算圖的邊variable_listoutputs;{pybind11::gil_scoped_releaseno_gil;auto&engine=python::PythonEngine::get_python_engine();//進入引擎執行

outputs=engine.execute(roots,grads,keep_graph,create_graph,accumulate_grad,output_edges);}29編程框架設計計算圖構建計算圖執行*深度學習編譯*分布式訓練本章小結提綱計算圖執行將計算圖中的張量和操作（本節又稱算子）映射到給定設備上具體執行設備管理張量實現算子執行獲取算子執行序列實現算子：前端定義、后端實現、前后端綁定查找并調用算子301、設備管理設備是編程框架中計算圖執行時的硬件實體，每個設備都具體負責計算子圖中的張量存放和算子運算常見設備包括通用處理器（如CPU）和領域專用處理器（如GPU和DLP等）添加對領域專用處理器的設備管理支持（三個模塊）設備操作執行流管理事件管理31PyTorch中的設備類型32PyTorch中的設備被直接按照類型分類，例如CPU,CUDA,DLP等通過索引表示特定設備，設備索引唯一，在有多個特定類型的設備時標識特定的計算設備設備管理33在pytorch/c10/core/impl/DeviceGuardImplInterface.h中定義了抽象的設備管理類DeviceGuardImplInterface設備操作初始化設備運行環境、獲取設備句柄和關閉并釋放設備等執行流管理：設備上抽象出來的管理計算任務的軟件概念在異構編程模型下，完成設備上任務執行的下發和同步操作執行流創建、執行流同步和執行流銷毀等事件管理表示設備上任務運行的狀態和進展事件創建、事件記錄和事件銷毀等基本操作設備管理342、張量實現邏輯視圖：形狀、布局、步長、偏移量、數據類型和設備等。是框架使用者能直接控制和表達的基本屬性物理視圖：設備上的物理地址空間大小、指針、數據類型等。對框架使用者不可見35張量數據結構A的邏輯視圖是一個形狀為[2,2]的張量，物理視圖是物理地址空間中從0x10位置開始連續存儲的一塊數據邏輯視圖通過偏移量和步長來確定物理視圖中物理地址空間的尋址空間一個物理視圖可以對應多個邏輯視圖：切片的結果不是新的物理視圖，而是原本物理視圖下的一個新的邏輯視圖36

b)行切片a)列切片物理視圖邏輯視圖物理視圖邏輯視圖大?。?步長：2偏移量：無大?。?步長：1偏移量：21234A[1,:]1234A[:,0]AA0x100x140x180x1c0x200x100x140x180x1c0x20PyTorch中的張量抽象PyTorch中存在與張量對應的類Tensor持有一個指向底層TensorImpl對象的指針37TensorAllocatorStorageTensorImplTensorBaseCPUGPU…DLPATenC10PyTorch中的張量抽象通過張量（Tensor）抽象類和存儲（Storage）抽象類來分別表示張量數據結構中的邏輯視圖和物理視圖TensorImpl類：張量抽象的實現，包含了維度信息，步長信息，數據類型，設備，布局等邏輯視角的張量信息StorageImpl類：張量的存儲實現，包含了內存指針、數據總數等物理視角的張量信息，調用結構體Allocator進行張量數據空間的分配38Tensorimpl*impl_Storagestorage_TypeMeta

data_type_Device*device_opt_SizeAndStrides

size_and_stridesStorageImpl*storage_impl_Allocator*allocator_StorageImplStorageTensorImplTensorBase即時分配器內存池分配器…張量內存分配從邏輯視圖到物理視圖的轉換需要完成對張量的內存分配，即對張量進行內存管理根據設備的類型不同，張量管理的方式不同即時分配---CPU內存池分配---GPU39張量內存分配--即時分配每當需要分配張量的內存時，就立即從系統中申請一塊合適大小的內存空間代碼核心部分：malloc()和free()函數40張量內存分配--內存池分配預先分配一塊固定大小的內存池，然后在需要時從內存池中分配內存自我維護：內存塊的拆分和合并優點：節約設備內存使用，減少設備內存碎片化41張量初始化在CPU上創建一個空張量選擇分配器

創建StorageImpl類

創建TensoImpl類423、算子執行計算圖的執行過程=每個算子獨立執行的過程計算圖

執行序列（確保正確的數據流和依賴關系）針對每個算子進行算子實現：前端定義、后端實現和前后端綁定分派執行：查找適合給定輸入的算子實現，并調用相應的實現來執行具體的計算任務43執行序列分析計算圖節點之間的依賴關系

執行序列拓撲排序算法（可有多種可行的結果）44f(x1

,

x2)x1x2

matmul

cos

add

substract

exp

cos

matmul

add

exp

substract算子實現正向傳播實現和反向傳播實現分離用戶接口（前端）和具體實現（后端）分離算子實現流程前端定義：在編程框架中配置算子信息，包含算子的輸入、輸出以及相關的接口定義，最后生成前端接口（如PythonAPI）后端實現：使用C++或其他高級的編程語言，編寫算子的底層實現代碼，完成算子的計算邏輯部分實現前后端綁定：編程框架將前端定義的算子與后端的具體實現進行綁定45native_function模式PyTorch用于管理整個算子實現模塊在使用該模式進行算子實現時，需要修改配置文件native_functions.yaml以添加算子配置信息native函數格式（位于native_functions.yaml）func字段：定義了算子名稱和輸入輸出的參數類型variants字段：表示需要自動生成的高級方法dispatch字段：表示該算子所支持的后端類型和對應的實現函數46-func:func_name(ArgTypearg0[=default],ArgTypearg1[=default],...)->Returnvariants:function,methoddispatch:CPU:func_cpuCUDA:func_cudaPReLU算子的native函數47-func:prelu(Tensorself,Tensorweight)->Tensorvariants:function,methodautogen:prelu.out-func:_prelu_kernel(Tensorself,Tensorweight)->Tensordispatch:CPU,CUDA:_prelu_kernelQuantizedCPU:_prelu_kernel_quantized_cpuMkldnnCPU:mkldnn_preluMPS:prelu_mps-func:_prelu_kernel_backward(Tensorgrad_output,Tensorself,Tensorweight)->(Tensor,Tensor)dispatch:CPU,CUDA:_prelu_kernel_backwardMkldnnCPU:mkldnn_prelu_backwardMPS:prelu_backward_mpsPReLU算子實現、PReLU正向傳播函數實現和PReLU反向傳播函數實現前端定義前端實現代碼綁定正向傳播和反向傳播函數參數在配置文件中添加算子正向傳播函數和反向傳播函數的對應關系48-name:_prelu_kernel(Tensorself,Tensorweight)->Tensorself,weight:"grad.defined()?_prelu_kernel_backward(grad,self,weight):std::tuple<Tensor,Tensor>()"result:at::where(self_p>=0,self_t,weight_p*self_t+weight_t*self_p)后端實現表層實現：不同設備之間的抽象函數接口_prelu_kernel()和_prelu_kernel_backward()iter提供了統一的計算抽象，其封裝了前向計算的輸入input、權重weight，以及反向計算的梯度grad調用stub函數進行具體的實現49后端實現底層實現：具體到某個設備上的實際代碼實現正向prelu_kernel()反向prelu_backward_kernel()

這兩個函數都利用了SIMD指令實現向量優化底層實現中的prelu_kernel和表層實現中的prelu_stub會在前后端綁定中完成對應50前后端綁定同一個算子可能會有多個后端實現的代碼多種后端&多種輸入，根據不同情況調用相應的后端實現PyTorch使用分派機制來管理前后端對應關系，由Dispatcher管理分派表分派表的表項記錄著算子到具體的后端實現對應關系，縱軸表示PyTorch所支持的算子，橫軸表示支持的分派鍵（與后端相關的標識符）51CPUGPUDLP…addmulprelucpu_prelu…TORCH_LIBRARY_IMPL(aten,CPU,m){m.impl(“prelu”,cpu_prelu);}分派執行獲得算子執行序列

實現對應算子

對算子分派執行分派執行：在運行時根據輸入張量的類型和設備類型查找并調用合適的算子實現方法Dispatcher計算分派鍵，并由此找到對應的內核函數算子：Dispatcher的調度對象，代表了具體的計算任務分派鍵：根據輸入張量和其他信息計算，可簡單地理解為與硬件平臺相關聯的標識符內核函數：特定硬件平臺上實現算子功能的具體代碼5253編程框架設計計算圖構建計算圖執行*深度學習編譯*分布式訓練本章小結提綱1、為什么需要深度學習編譯編程框架中早期優化方式存在的問題框架維護成本高：對于新硬件和新算子，都需要程序員手動進行算子開發，開發數量呈平方級增長性能受限：性能受限于程序員人工優化算子的能力，且沒有充分探索計算圖的優化空間54在深度學習編程框架中引入深度學習編譯機制減少人工開發工作量：可針對不同硬件平臺進行代碼生成便于性能優化：（圖層級）對完整的計算圖進行靜態分析和全局優化；（算子層級）利用自動調優技術優化算子，最大限度提升硬件利用率什么是深度學習編譯器55接收以計算圖形式表示的深度學習任務，并在指定硬件平臺上生成高性能代碼多個層級中間表示&多個層級優化圖層級優化：子圖替換、常量折疊、公共子表達式刪除、布局優化以及算子融合等算子層級優化：自動調優，基于搜索的方法和基于多面體模型的方法常見深度學習框架中所采用的編譯技術和深度學習編譯器TVM、TC（TensorComprehensions）、XLA、MLIR…什么是深度學習編譯器56原始計算圖3x31x11x13x31x1圖層級

編譯優化算子層級

編譯優化硬件平臺CPUGPUDLPforiinrange(len(A)):forjinrange(len(B[0]):forkinrange(len(B)):C[i][j]+=A[i][k]*B[k][j]for(intx=j;x<j+BLOCK;x+=4){__m256dc=_mm256_loadu_pd(&C[y*N+x]);for(intz=k;z<k+BLOCK;++z){__m256da=_mm256_broadcast_sd(&A[y*N+z]);__m256db=_mm256_loadu_pd(&B[z*N+x]);…深度學習編譯器跟編程框架的關系57深度學習編程框架自行適配廠商提供的計算庫或者手寫算子來支持不同硬件，這帶來了極高的框架維護成本深度學習編譯器提供跨平臺統一的抽象和優化，較為靈活的適配不同的上層編程框架和底層硬件平臺經過圖層級優化和算子層級優化后，自動生成在目標硬件平臺上的高性能算子深度學習編譯器跟編程框架的關系58…硬件平臺CPUGPUDLPTF-CPUTF-GPUTF-DLPTorch-CPUTorch-GPUTorch-DLP原始計算圖圖優化算子優化框架對每個硬件平臺的計算庫適配編譯器提供跨平臺統一的抽象和優化CPUGPUDLP2、圖層級編譯優化不關心特定算子的具體執行過程，而關心數據在圖中的流動過程圖優化方法子圖替換、常量折疊、公共子表達式刪除、布局優化、算子融合…59圖優化方法子圖替換：將原計算圖中的節點（計算操作）替換為功能等價但運算邏輯更優的形式TensorFlow中人為設定的替換規則60圖優化方法常量折疊：如16*16*224結果為定值，則計算其值后帶入此定值公共子表達式消除Tensor0MatmulTensor1DivideSubtractTensor2Tensor3MatmulTensor0MatmulTensor1DivideSubtractTensor261圖優化方法代數化簡：將代價高的計算換為等價代價低的運算若：表達式中出現乘0，此表達式的結果直接為0布局優化：輸入布局影響執行性能使用TensorCore計算相同輸入數據，采用NHWC格式的性能普遍優于NCHW格式62圖優化方法算子融合：縱向函數調用有開銷，外設的函數調用（KernelLaunch）開銷巨大將多個小算子融合為一個大算子進行執行常見：FMA（FusedMultiply-Add）TensorTensorMultiplyAddTensorTensorTensorFMATensor63圖優化方法算子融合：橫向多個小的矩陣乘可以合并為一個大的矩陣乘TensorMatmulTensorMatmulMatmulTensorTensorTensorTensorFusedTensorTensor643、算子層級編譯優化接收圖層級編譯優化后的計算圖節點作為輸入，將其下降到算子層級中間表示上，最終生成目標硬件后端上的代碼算子層級中間表示：抽象建模一個計算及其在設備上的具體執行流程算子調度：針對目標硬件后端上的計算特性和訪存特性進行優化自動調優：自動確定最優的調度配置65算子層級中間表示計算與調度分離計算表示涵蓋了算子的計算定義信息，但不包括具體的實現信息。包括對張量和對計算本身的描述。計算加上調度的表示確定了算子的實現，可以使用嵌套循環程序對其進行表示。66A=placeholder((N,L),name=“A”,dtype=“float”)B=placeholder((L,M),name=“B”,dtype=“float”)k=reduce_axis((0,L),name=“k”)C=compute((N,M),lambdai,j:sum(A[i,k]*B[k,j]),axis=k),name=“C”)for(i:int32,0,1024)“parallel”:

for(j:int32,0,1024):C[((i*1024)+j)]=0f32for(k:int32,0,1024):letcse_var_2:int32=(i*1024)letcse_var_1:int32=(cse_var_2+j)C[cse_var_1]=(C[cse_var_1]+(A[(cse_var_2+k)]*B[((k*1024)+j)]))算子調度通過循環變換來匹配目標平臺的體系結構特性（包括計算特性和訪存特性）算子調度算子的具體實現通常表現為嵌套循環程序循環分塊（tiling）優化、循環向量化（Vectorize）等優點：提升緩存命中率和（在CPU平臺上）使用向量化加速67for(i:int32,0,1024):for(j:int32,0,1024):C[...]=0f32for(k:int32,0,1024):C[...]=C[...]+A[...]*B[...]算子調度一個完整的調度是由多個調度原語構成的常見的算子調度原語通過不同調度原語和調度參數的組合，編譯器可以構建一個包含海量不同程序重寫的優化空間68自動調優通過搜索的方式確定合適的調度配置空間探索：一個點代表一種配置性能測量：測試某配置下的程序性能代價模型：對性能進行評估，并選擇配置69空間探索代價模型性能測量計算表示優化程序性能預測程序性能數據調度配置自動調優自動調優的核心是空間搜索空間和搜索算法的設計常見的空間搜索方式基于手工模板的搜索基于序列構建的搜索層次化的搜索方法70基于手工模板的搜索方法依賴于給定的調度模板該模板包括手工設計的原語序列該序列通常只有調度參數沒有確定空間較為簡單，可用多種搜索算法隨機搜索、網格搜索、遺傳算法…手工模板的設計需要領域專家進行71參數搜索固定人工模板

fori.0in

range():

forj.0in

range():

fork.0in

range():

fori.1in

range():

forj.1in

range():C[...]+=A[…]*B[…]

fori.1in

range():

forj.1in

range():D[...]=max(C[...],0.0)???????基于序列構建的搜索方法逐條循環語句地構建優化程序編譯器選擇合適的調度原語以及調度參數使用代價模型進行性能評估搜索算法存在一定的限制隨機搜索、集束搜索、蒙特卡洛樹搜索…耗時且低效缺少程序優化的先驗知識對空間進行有效剪枝72fori.0inrange(512):提前終止的集束搜索未完成的程序

fori.0in

range(512):

forj.0in

range(512):

D[...]=max(C[...],0.0)如何建立下一條語句?候選1候選2候選3候選4剪枝剪枝保留保留層次化構建的搜索方法從粗到細的粒度構建優化程序粗粒度：程序所要采用的循環結構細粒度：具體的調度參數搜索算法存在一定的限制隨機搜索、遺傳算法…優缺點更大的搜索空間在粗粒度的結構選擇策略中引入優化的先驗知識需要領域專家設計73更好的程序高層次結構生成低層次細節采樣完整的程序

fori.0in

range(64):

forj.0in

range(64):

fork.0in

range(512):

fori.1in

range(512):

forj.1in

range(8):D[...]+=…

for…

for…

for…

for…

for…

for…

for…for…

for…

for…for…

for…

?????微調算子層級編譯優化的實現74模型輸入：將調度該算子的任務分發到搜索空間生成器搜索空間生成：為算子產生調度序列，及該調度序列所需的參數取值范圍搜索空間探索：通過特定搜索策略選取高性能的算子性能測量算子候選配置性能測量靜態分析任務抽取搜索空間輸入搜索空間生成調度配置tilingparallelvectorizeunroll調度規則代價模型優化程序搜索空間探索搜索中止4、常見深度學習編譯器介紹TVMTensorComprehensionsXLAMLIRTorchDynamo和TorchInductor…75TVM（TensorVirtualMachine）76TVM結構示意圖：兩個層級的編程抽象TensorFlowPyTorch編程框架計算圖優化計算圖算子融合布局變換算子優化后端…計算描述調度優化OpenCLLLVMCUDA…TVMTVM的核心思想：計算與調度分離計算：定義元素之間的運算關系調度：規劃具體計算執行的運算順序、數量C=pute((M,N),lambdam,n:

te.sum(A[m,k]*B[k,n],axis=k),

name="C")s=te.create_schedule(C.op)mo,no,mi,ni=s[C].tile(C.op.axis[0],C.op.axis[1],bn,bn)(kaxis,)=s[C].op.reduce_axisko,ki=s[C].split(kaxis,factor=kfactor)計算調度77兩個層級的編程抽象圖層級中間表示算子層級中間表示TVMTVM的自動調優實現C=pute((M,N),lambdam,n:

te.sum(A[m,k]*B[k,n],axis=k),

name="C")s=te.create_schedule(C.op)mo,no,mi,ni=s[C].tile(C.op.axis[0],C.op.axis[1],bn,bn)(kaxis,)=s[C].op.reduce_axisko,ki=s[C].split(kaxis,factor=kfactor)計算調度78候選程序集合產生自用戶定義的調度策略搜索空間數據庫CPUGPUDLP…調度器代價模型查詢更新記錄評估TensorComprehensions自動生成高性能代碼TensorComprehensions（TC）基于多面體模型（PolyhedralModel）的即時（JIT）編譯器將DAG轉換為具有內存管理和同步的CUDA內核函數多面體模型技術使用一個結構化的方式來捕獲和表示循環代碼的結構和語義，并可以在這個表示的基礎上應用各種優化和變換手段，在保持代碼語義不變的基礎上提高性能79XLA（AcceleratedLinearAlgebra）XLA由谷歌開發的，并作為TensorFlow的一部分提供對計算圖進行優化和編譯核心：HLO

IR（提供細粒度的算子抽象），組合成任意的算子提供多種與硬件架構無關的優化：公共子表達式消除、算子融合等BatchNorm算子，通過一系列的Broadcast、Reduce和Element_wise操作組成使用LLVMIR表示CPU和GPU后端80XLA（AcceleratedLinearAlgebra）81TensorFlowPyTorchJAXStableHLO（OpenXLA輸入）硬件無關優化CPU硬件相關優化&目標硬件代碼生成…GPUTPUOpenXLA項目源于將編譯器相關技術從TensorFlow獨立編程框架的輸入

StableHLO

硬件無關優化

硬件相關優化

代碼生成神經網絡、學習框架、硬件平臺數量蓬勃發展為減少重復工作量，MLIR（多層級中間表示）的概念被提出MobileNetDiffusionGANCPUGPUDLPTPUFPGADSP………N種網絡M個框架P個硬件平臺最差情況下，需要做N*M*P的適配MLIR提供了一套基礎設施為開發編譯器提供便利82MLIR（Multi-LevelIntermediateRepresentation）MLIR提供一套基礎設施為開發編譯器提供便利使用混合的中間表示，解決當下的軟件碎片化問題設計方言“Dialect”機制，便于接入各式語言和中間表示若存在不相容的特性，則使用Dialect機制擴充83方言A接口屬性類型操作屬性類型約束接口特征方言B接口屬性類型操作屬性類型約束接口特征方言轉換PyTorch編譯技術迭代PyTorch中編譯技術的發展歷程84版本號更新內容/使用Codegen而非C++模板，利用StructureKernels對算子生成進行描述；使用Dispatcher分派1.0引入JIT編譯器，使用jit.trace或jit.script1.8引入Torch.fx，使用中間表示FXIR實現Python到Python的代碼轉換1.12引入nvFuser，一個針對nVidia平臺的編譯器2.0引入pile()，結合多種編譯技術優化體驗0.1Jit.traceTorchscriptFXTracingnvFuserTorchdynamoTorchInductor1.01.81.122.0PyTorch2.0易用且高性能：pile()PyTorch2.0引入的新編譯相關技術85技術用途TorchDynamo圖捕獲工具，捕獲執行的Python幀并將其給Inductor編譯，用戶無感TorchInductor深度學習編譯器后端，將Dynamo捕獲的結果編譯到多個后端AOTAutograd重載了Pytorch的Autograd，使用trace技術生成反向的tracePrimTorch將PyTorch的2000多個算子抽象縮減為了約250個算子的集合TritonTorchInductor中使用的GPU后端，使用Python編寫算子并自動調優TorchDynamo&TorchInductor核心技術的組織結構86C++/OpenMP其他TorchDynamoAOTAutogradTritonTorchInductor代碼

生成

后端模型輸入Inductor循環層級中間表示ATen/PrimsTorch中間表示：精簡算子集合，2000+

->250Torch中間表示：使用Python語法的中間表示，算子調用…捕獲執行的Python幀，傳遞至TorchInductor編譯為反向傳播提供靜態圖支持將中間表示進行翻譯，代碼傳遞至后端模塊常見深度學習編譯器對比87編譯器名稱主要維護團體中間表示典型特點TVMApache圖層級：Relay/Relax中間表示

算子層級：Tensor中間表示基于算子調度的自動調優TensorComprehensions官方不再維護算子層級：基于Halide的中間表示基于多面體模型的自動優化XLAGoogle圖層級：HLO中間表示細粒度的算子抽象、算子融合MLIRLLVMDeveloperGroup基于方言機制的多層級中間表示編譯器基礎設施88編程框架設計計算圖構建計算圖執行*深度學習編譯*分布式訓練本章小結提綱891、為什么需要分布式訓練大模型及其相關應用蓬勃發展參數數量的增加帶來了模型的表達能力和擬合能力提高龐大的訓練數據使得模型能夠學習到更全面的知識和對數據分布的理解

許多從前難以實現的任務變得可行更大的參數量和更多的訓練數據導致訓練過程中的算力墻和存儲墻90更大的參數量大模型通常擁有更長更深的網絡結構網絡中有更多層次的神經元結構

更大的參數量從GPT-1到GPT-3，模型的參數量從1.2億增長到了1750億常見大模型參數量（截至2023年7月）BERT

340MGPT-1

117MGPT-2

1.5BGPT-3

175BLaMDA137BPaLM540BGLM-130B130BChinchilla70BMegatron-11B

11BGPT-41800BMT-NLG530B91更多的訓練數據大模型的訓練都需要海量的數據數據的多樣性可以幫助模型更好地理解和捕捉數據中的模式從GPT-1到GPT-3，訓練模型使用的數據集大小從5GB增長到了753GB；GPT-3訓練使用的數據集包含維基百科、書本、會議期刊、源代碼等；數據集大小/GBGPT-1GPT-2GPT-3GPT-J/GPT-NeoX-20BMegatron-11BMT-NLGGopher維基百科\\11611612書籍5\2111851182100學術雜志0\101244\77\Reddit鏈接\4050633863\Commoncrawl數據集\\5702271079833450其它\\\167\1274823總計54075382516113741038592模型很大&數據很多在訓練過程中產生了計算墻和存儲墻計算墻：GPT-3模型的訓練若使用8張V100顯卡，訓練預計耗時36年，而使用1024張A100可以將訓練耗時減少到1個月存儲墻：千億級別大模型的存儲(包括參數、訓練中間值等)需要2TB存儲空間，而單卡的顯存目前最大只有80GB單個計算設備的資源有限，無法存儲整個模型的參數或者計算全部的數據集，并且提升單個設備性能成本遠遠高于使用多個設備分布式訓練技術：拆分任務并由多個設備共同協作完成計算拆分訓練數據拆分模型（計算圖）932、分布式訓練基礎分布式架構和分布式同步策略是分布式訓練的基礎分布式架構：組織和管理分布式訓練任務的方式，以最大程度地利用計算資源和提高訓練效率分布式同步策略：在分布式環境中，為了保證計算節點之間的一致性和正確性，對不同計算節點之間的操作進行協調和同步的策略分布式架構94常使用以下兩種架構實現分布式訓練：參數服務器(ParameterServers)中心化的ParameterServers架構由李沐于2014年提出，“中心化”是指將模型參數進行中心化管理，以此實現模型參數的同步集合通信(CollectiveCommunication)去中心化的CollectiveCommunication架構中，每個訓練節點都有當前全局最新參數，節點間的參數同步通常采用多次設備之間的點對點通信完成的參數服務器95參數服務器將所有節點分成中心節點(Server節點)和計算節點(Worker節點)兩類中心節點用于存儲參數和梯度更新計算節點用于完成中心節點下發的實際計算任務，僅與中心節點通信以更新和檢索共享參數中心節點m中心節點1中心節點2計算節點1計算節點2計算節點3計算節點n……參數服務器96參數服務器架構的計算和存儲分離優點靈活：通過改變中心節點數量適應不同的負載和數據規模高效地參數共享：由中心節點統一管理模型參數缺點單點故障：單個中心節點故障會影響整個系統數據一致性的問題：多個計算節點可能同時讀取和更新模型參數網絡通信開銷：受通信帶寬的限制，中心節點成為系統的瓶頸集合通信97集合通信是指一個進程組的所有進程都參與全局通信操作集合通信中沒有中心節點（也被稱為去中心化的架構）每個計算節點都有當前全局最新參數節點間的參數同步通常采用多次設備之間的點對點通信完成的對芯片的算力和芯片之間的網絡通信要求較高計算節點1計算節點2計算節點4計算節點3集合通信原語集合通信98集合通信的基礎操作：發送（send）、接收（receive）、復制（copy）、組內進程障礙同步（barrier)以及節點間進程同步（signal+wait）基礎操作組合后可以得到集合通信中常用的的通信原語通信原語一對多通信原語：Broadcast、Scatter多對一通信原語：Gather、Reduce多對多通信原語：All-to-All、All-Gather、All-Reduce、Reduce-Scatter集合通信原語99一對多廣播（Broadcast）：將一個進程的數據廣播到所有進程，常用于分享模型參數BroadcastADLP0DLP1DLP2DLP3AAAADLP0DLP1DLP2DLP3集合通信原語100一對多散射（Scatter）：將一個進程中的數據按索引散射到多個進程，常用于更新權重ScatterDLP0DLP1DLP2DLP3DLP0DLP1DLP2DLP3ABCDABCD集合通信原語101多對一收集（Gather）：從多個進程收集數據到一個進程，常用于收集梯度DLP0DLP1DLP2DLP3GatherABCDDLP0DLP1DLP2DLP3ABCDAll-GatherABCDABCDABCDABCDABCDDLP0DLP1DLP2DLP3DLP0DLP1DLP2DLP3集合通信原語102多對多收集（All-Gather）：從多個進程收集數據，并廣播到所有進程，常用于數據同步集合通信原語103多對一歸約（Reduce）：從多個進程收集數據，并按某種運算（如求和運算）歸約到一個進程，常用于梯度累加ReduceABCDA

+

B

+

C

+

DDLP0DLP1DLP2DLP3DLP0DLP1DLP2DLP3All-ReduceABCDA

+

B

+

C

+

DA

+

B

+

C

+

DA

+

B

+

C

+

DA

+

B

+

C

+

DDLP0DLP1DLP2DLP3DLP0DLP1DLP2DLP3集合通信原語104多對多歸約（All-Reduce）：從多個進程收集數據，并按某運算歸約，再廣播到所有進程，常用于數據同步和梯度累加集合通信原語105多對多歸約散射（Reduce-Scatter）：從多個進程收集數據，并按某種運算歸約到一個進程，將該進程中的數據按索引散射到對應進程上，常用于更新權重Reduce-Scattersum{Ai}B1C1D1sum{Bi}C1D1B1sum{Ci}D1B1C1sum{Di}DLP0DLP1DLP2DLP3DLP0DLP1DLP2DLP3A1B1C1D1A2B2C2D2A3B3C3D3A4B4C4D4集合通信原語106多對多交換（All-to-All）：將每個進程中的數據按索引發射到其他進程，每個進程接收數據后以發送進程號為索引存儲到對應的數據塊中，常用于數據同步和信息傳遞A1B1C1D1A2B2C2D2A3B3C3D3A4B4C4D4All-to-AllA1A2A3A4D4D3D2D1B1B2B3B4C1C2C3C4DLP0DLP1DLP2DLP3DLP0DLP1DLP2DLP3分布式同步策略107設備之間的通信可以采用不同的同步策略同步通信異步通信選擇合適的分布式同步策略對于保證分布式系統的正確性、性能和可擴展性至關重要同步通信108采用同步通信作為分布式同步策略需要等待全部計算節點完成本輪計算后才進行通信時序性和順序性使用同步障確保計算節點之間的數據一致性可能會導致較大的延遲和通信開銷同步通信109計算輪次設備1設備2計算輪次空泡計算輪次設備3通信計算輪次計算輪次空泡空泡計算輪次同步障同步障同步障的存在確保全部設備完成通信后才可開始下一輪計算異步通信110采用異步通信作為分布式同步策略每個計算節點可以隨時和其他設備進行通信更加靈活提高整個分布式訓練系統的計算利用率但不能保證數據的一致性異步通信111設備1計算輪次計算輪次計算輪次計算輪次設備2計算輪次計算輪次計算輪次設備3計算輪次計算輪次計算輪次計算輪次通信每個設備可以隨時處理自己收到的信息，不會因為同步障而帶來互相等待的開銷3、分布式訓練方法概述112根據分布式計算中的分區情況，可以劃分出不同的分布式計算方法：數據并行：對輸入數據進行分區模型并行：對模型參數進行分區混合并行：同時對輸入數據和模型參數進行分區分布式計算步驟單卡節點訓練

多節點分布式訓練113輸入計算節點輸入1輸入2輸入3計算節點1計算節點2計算節點3輸出1輸出2輸出3輸入分區并行計算合并輸出輸出分布式計算一般包含三個步驟（1）將輸入進行分區（2）將每個分區發給不同的計算節點，實現并行計算（3）合并每個計算節點的輸出，得到和單節點等價的計算結果114數據并行(DataParallelism)往往用于解決單節點算力不足的問題。其中，每個設備共享完整的模型副本，輸入數據會被分發給這些設備，減少單個設備的負載。數據并行示意圖數據并行模型并行115模型并行(ModelParallelism)往往用于解決單節點內存不足的問題。一般將模型并行分為算子內并行和算子間并行。算子內并行大型算子計算所需內存超過單設備內存容量，對單個算子進行切分按行切分和按列切分算子間并行模型的總內存需求超過單設備的內存容量，在算子間進行切分算子內并行116模型并行示意圖——算子內并行反向：算子2的數據被廣播給設備1和設備2，兩設備根據本地的參數分區完成局部的反向計算正向：輸入被廣播給設備1和設備2，計算結果合并后傳給下游算子2模型中單個算子本身計算所需的內存已經超過單設備的內存容量，就需要對這些大型算子進行切分。對單個算子的切分一般可以分為按行切分和按列切分按行切分和按列切分117XMKNAKMY=XAN原始矩陣乘法XMKA1A2N/2N/2[Y1,Y2]=[XA1

,XA2]=XA=Y按列切分

參數矩陣AKA=[A1,A2]N/2N/2Y1=

XA1Y2=

XA2MK/2K/2

X=[X1,X2]

按行切分

參數矩陣ANA2A1Y1=X1A1NY2=X2A2NMX1K/2K/2MX2

算子間并行118模型中單個算子參數量較少，但整個模型的總內存需求超過單設備的內存容量，就需要在算子間進行切分。模型并行示意圖——算子間并行模型并行空泡(ModelParallelismBubble)現象：算子間并行中，下游設備需要等