智能計算系統(tǒng)

第五章編程框架原理中國科學院計算技術研究所李威副研究員liwei2017@2編程框架設計計算圖構建計算圖執(zhí)行*深度學習編譯*分布式訓練本章小結提綱學習智能計算系統(tǒng)，了解編程框架的原理大有脾益編寫與框架底層更為契合、性能更優(yōu)的代碼定制化擴展編程框架，為新算法、新設備提供支持智能計算系統(tǒng)中編程框架的四大模塊必備：計算圖構建模塊和計算圖執(zhí)行模塊追求更高性能：深度學習編譯模塊和分布式訓練模塊3概述4編程框架設計計算圖構建計算圖執(zhí)行*深度學習編譯*分布式訓練本章小結提綱1、設計原則5簡潔性（Simplicity）框架提供一套抽象機制，用戶僅需關心算法本身和部署策略易用性（Usability）高效性（Performance）簡潔抽象輸入數(shù)據(jù)模型輸出結果addbbaccuda()GPUCPU1、設計原則6簡潔性（Simplicity）易用性（Usability）熟悉的開發(fā)范式：如PyTorch始于Python，忠于Python直觀且用戶友好的接口：如PyTorch提供了命令式的動態(tài)圖編程方法高效性（Performance）1、設計原則7簡潔性（Simplicity）易用性（Usability）高效性（Performance）如采用靜態(tài)圖編程方式，可以生成完整的計算圖并進行全局優(yōu)化，從而盡量提高用戶應用程序的運行效率支持深度學習編譯技術，多層級表示優(yōu)化，充分利用用戶硬件的計算能力支持多機多卡條件的分布式訓練，從而高效支持大規(guī)模深度學習任務2、整體架構8四大模塊計算圖構建模塊：完成從輸入的用戶程序到編程框架內(nèi)部原始計算圖的轉換過程，編程框架的入口模塊分布式訓練模塊：應對更大規(guī)模的神經(jīng)網(wǎng)絡，將訓練、推理任務從一臺設備擴展到多臺設備深度學習編譯模塊：對計算圖分別進行圖層級和算子層級的編譯優(yōu)化，從而提升單設備上的執(zhí)行效率計算圖執(zhí)行模塊：將優(yōu)化后的計算圖中的張量和操作映射到指定設備上進行具體執(zhí)行，并給出編程框架的輸出結果計算圖構建計算圖執(zhí)行分布式訓練深度學習編譯dataMatmulweightoutput原始計算圖優(yōu)化后的計算圖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編程框架設計計算圖構建計算圖執(zhí)行*深度學習編譯*分布式訓練本章小結提綱正向圖與反向圖構建11計算圖由兩個基本元素構成：張量（Tensor）和張量操作（Operation）。計算圖是有向圖，有向邊指明了張量的流動方向1、正向傳播12輸入張量經(jīng)過搭建的神經(jīng)網(wǎng)絡層層計算傳遞，并最終獲得計算結果的過程構建形式動態(tài)圖：在執(zhí)行函數(shù)時，按照函數(shù)順序逐條語句地生成節(jié)點，立即計算并返回結果；易調試但性能優(yōu)化空間有限靜態(tài)圖：在執(zhí)行計算之前構建好所有圖上的節(jié)點，在圖運行時才計算整個計算圖并返回最終結果；不易調試但性能好動態(tài)圖計算圖在函數(shù)運行過程中逐步構建的（On-the-fly）立即（eager）模式：每次調用語句就立刻執(zhí)行計算PyTorch中的動態(tài)圖實現(xiàn)：每次執(zhí)行，都會重新被構建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靜態(tài)圖整個網(wǎng)絡的結構會在開始計算前就建立完成計算圖框架執(zhí)行時接收整個計算圖而不是單一語句TensorFlow1.x中的靜態(tài)圖使用若干基本控制流算子（Switch、Merge、Enter、Exit和NexIteration）的不同組合來實現(xiàn)各種復雜控制流場景PyTorch2.0中的靜態(tài)圖PyTorch2.0中采取了圖捕獲（TorchDynamo）的技術將用戶的動態(tài)圖轉化為靜態(tài)圖142、反向傳播15正向計算得到的結果和目標結果存在損失函數(shù)值，對其求導得到梯度，并使用該梯度更新參數(shù)wz*

MultBackward反向正向y1y2

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

addmatmulbwx

前向傳播

反向傳播手動求解法18數(shù)值求導法

19符號求導法

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

22exp+*+1

前向計算

23exp+*+1

+

反向計算

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

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

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

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

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

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

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

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

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函數(shù)實現(xiàn)2）構建反向計算圖的節(jié)點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){//沒梯度函數(shù)，則標記是葉子節(jié)點

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

}else{//有梯度函數(shù)，創(chuàng)建梯度函數(shù)指向該節(jié)點的邊（構造反向計算圖）

output_edges.emplace_back(grad_fn,output_nr);

}

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

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

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

data_type_Device*device_opt_SizeAndStrides

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

創(chuàng)建StorageImpl類

創(chuàng)建TensoImpl類423、算子執(zhí)行計算圖的執(zhí)行過程=每個算子獨立執(zhí)行的過程計算圖

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

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

,

x2)x1x2

matmul

cos

add

substract

exp

cos

matmul

add

exp

substract算子實現(xiàn)正向傳播實現(xiàn)和反向傳播實現(xiàn)分離用戶接口（前端）和具體實現(xiàn)（后端）分離算子實現(xiàn)流程前端定義：在編程框架中配置算子信息，包含算子的輸入、輸出以及相關的接口定義，最后生成前端接口（如PythonAPI）后端實現(xiàn)：使用C++或其他高級的編程語言，編寫算子的底層實現(xiàn)代碼，完成算子的計算邏輯部分實現(xiàn)前后端綁定：編程框架將前端定義的算子與后端的具體實現(xiàn)進行綁定45native_function模式PyTorch用于管理整個算子實現(xiàn)模塊在使用該模式進行算子實現(xiàn)時，需要修改配置文件native_functions.yaml以添加算子配置信息native函數(shù)格式（位于native_functions.yaml）func字段：定義了算子名稱和輸入輸出的參數(shù)類型variants字段：表示需要自動生成的高級方法dispatch字段：表示該算子所支持的后端類型和對應的實現(xiàn)函數(shù)46-func:func_name(ArgTypearg0[=default],ArgTypearg1[=default],...)->Returnvariants:function,methoddispatch:CPU:func_cpuCUDA:func_cudaPReLU算子的native函數(shù)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算子實現(xiàn)、PReLU正向傳播函數(shù)實現(xiàn)和PReLU反向傳播函數(shù)實現(xiàn)前端定義前端實現(xiàn)代碼綁定正向傳播和反向傳播函數(shù)參數(shù)在配置文件中添加算子正向傳播函數(shù)和反向傳播函數(shù)的對應關系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)后端實現(xiàn)表層實現(xiàn)：不同設備之間的抽象函數(shù)接口_prelu_kernel()和_prelu_kernel_backward()iter提供了統(tǒng)一的計算抽象，其封裝了前向計算的輸入input、權重weight，以及反向計算的梯度grad調用stub函數(shù)進行具體的實現(xiàn)49后端實現(xiàn)底層實現(xiàn)：具體到某個設備上的實際代碼實現(xiàn)正向prelu_kernel()反向prelu_backward_kernel()

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

實現(xiàn)對應算子

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

編譯優(yōu)化算子層級

編譯優(yōu)化硬件平臺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深度學習編程框架自行適配廠商提供的計算庫或者手寫算子來支持不同硬件，這帶來了極高的框架維護成本深度學習編譯器提供跨平臺統(tǒng)一的抽象和優(yōu)化，較為靈活的適配不同的上層編程框架和底層硬件平臺經(jīng)過圖層級優(yōu)化和算子層級優(yōu)化后，自動生成在目標硬件平臺上的高性能算子深度學習編譯器跟編程框架的關系58…硬件平臺CPUGPUDLPTF-CPUTF-GPUTF-DLPTorch-CPUTorch-GPUTorch-DLP原始計算圖圖優(yōu)化算子優(yōu)化框架對每個硬件平臺的計算庫適配編譯器提供跨平臺統(tǒng)一的抽象和優(yōu)化CPUGPUDLP2、圖層級編譯優(yōu)化不關心特定算子的具體執(zhí)行過程，而關心數(shù)據(jù)在圖中的流動過程圖優(yōu)化方法子圖替換、常量折疊、公共子表達式刪除、布局優(yōu)化、算子融合…59圖優(yōu)化方法子圖替換：將原計算圖中的節(jié)點（計算操作）替換為功能等價但運算邏輯更優(yōu)的形式TensorFlow中人為設定的替換規(guī)則60圖優(yōu)化方法常量折疊：如16*16*224結果為定值，則計算其值后帶入此定值公共子表達式消除Tensor0MatmulTensor1DivideSubtractTensor2Tensor3MatmulTensor0MatmulTensor1DivideSubtractTensor261圖優(yōu)化方法代數(shù)化簡：將代價高的計算換為等價代價低的運算若：表達式中出現(xiàn)乘0，此表達式的結果直接為0布局優(yōu)化：輸入布局影響執(zhí)行性能使用TensorCore計算相同輸入數(shù)據(jù)，采用NHWC格式的性能普遍優(yōu)于NCHW格式62圖優(yōu)化方法算子融合：縱向函數(shù)調用有開銷，外設的函數(shù)調用（KernelLaunch）開銷巨大將多個小算子融合為一個大算子進行執(zhí)行常見：FMA（FusedMultiply-Add）TensorTensorMultiplyAddTensorTensorTensorFMATensor63圖優(yōu)化方法算子融合：橫向多個小的矩陣乘可以合并為一個大的矩陣乘TensorMatmulTensorMatmulMatmulTensorTensorTensorTensorFusedTensorTensor643、算子層級編譯優(yōu)化接收圖層級編譯優(yōu)化后的計算圖節(jié)點作為輸入，將其下降到算子層級中間表示上，最終生成目標硬件后端上的代碼算子層級中間表示：抽象建模一個計算及其在設備上的具體執(zhí)行流程算子調度：針對目標硬件后端上的計算特性和訪存特性進行優(yōu)化自動調優(yōu)：自動確定最優(yōu)的調度配置65算子層級中間表示計算與調度分離計算表示涵蓋了算子的計算定義信息，但不包括具體的實現(xiàn)信息。包括對張量和對計算本身的描述。計算加上調度的表示確定了算子的實現(xiàn)，可以使用嵌套循環(huán)程序對其進行表示。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)]))算子調度通過循環(huán)變換來匹配目標平臺的體系結構特性（包括計算特性和訪存特性）算子調度算子的具體實現(xiàn)通常表現(xiàn)為嵌套循環(huán)程序循環(huán)分塊（tiling）優(yōu)化、循環(huán)向量化（Vectorize）等優(yōu)點：提升緩存命中率和（在CPU平臺上）使用向量化加速67for(i:int32,0,1024):for(j:int32,0,1024):C[...]=0f32for(k:int32,0,1024):C[...]=C[...]+A[...]*B[...]算子調度一個完整的調度是由多個調度原語構成的常見的算子調度原語通過不同調度原語和調度參數(shù)的組合，編譯器可以構建一個包含海量不同程序重寫的優(yōu)化空間68自動調優(yōu)通過搜索的方式確定合適的調度配置空間探索：一個點代表一種配置性能測量：測試某配置下的程序性能代價模型：對性能進行評估，并選擇配置69空間探索代價模型性能測量計算表示優(yōu)化程序性能預測程序性能數(shù)據(jù)調度配置自動調優(yōu)自動調優(yōu)的核心是空間搜索空間和搜索算法的設計常見的空間搜索方式基于手工模板的搜索基于序列構建的搜索層次化的搜索方法70基于手工模板的搜索方法依賴于給定的調度模板該模板包括手工設計的原語序列該序列通常只有調度參數(shù)沒有確定空間較為簡單，可用多種搜索算法隨機搜索、網(wǎng)格搜索、遺傳算法…手工模板的設計需要領域專家進行71參數(shù)搜索固定人工模板

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)???????基于序列構建的搜索方法逐條循環(huán)語句地構建優(yōu)化程序編譯器選擇合適的調度原語以及調度參數(shù)使用代價模型進行性能評估搜索算法存在一定的限制隨機搜索、集束搜索、蒙特卡洛樹搜索…耗時且低效缺少程序優(yōu)化的先驗知識對空間進行有效剪枝72fori.0inrange(512):提前終止的集束搜索未完成的程序

fori.0in

range(512):

forj.0in

range(512):

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

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…

?????微調算子層級編譯優(yōu)化的實現(xiàn)74模型輸入：將調度該算子的任務分發(fā)到搜索空間生成器搜索空間生成：為算子產(chǎn)生調度序列，及該調度序列所需的參數(shù)取值范圍搜索空間探索：通過特定搜索策略選取高性能的算子性能測量算子候選配置性能測量靜態(tài)分析任務抽取搜索空間輸入搜索空間生成調度配置tilingparallelvectorizeunroll調度規(guī)則代價模型優(yōu)化程序搜索空間探索搜索中止4、常見深度學習編譯器介紹TVMTensorComprehensionsXLAMLIRTorchDynamo和TorchInductor…75TVM（TensorVirtualMachine）76TVM結構示意圖：兩個層級的編程抽象TensorFlowPyTorch編程框架計算圖優(yōu)化計算圖算子融合布局變換算子優(yōu)化后端…計算描述調度優(yōu)化OpenCLLLVMCUDA…TVMTVM的核心思想：計算與調度分離計算：定義元素之間的運算關系調度：規(guī)劃具體計算執(zhí)行的運算順序、數(shù)量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的自動調優(yōu)實現(xiàn)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候選程序集合產(chǎn)生自用戶定義的調度策略搜索空間數(shù)據(jù)庫CPUGPUDLP…調度器代價模型查詢更新記錄評估TensorComprehensions自動生成高性能代碼TensorComprehensions（TC）基于多面體模型（PolyhedralModel）的即時（JIT）編譯器將DAG轉換為具有內(nèi)存管理和同步的CUDA內(nèi)核函數(shù)多面體模型技術使用一個結構化的方式來捕獲和表示循環(huán)代碼的結構和語義，并可以在這個表示的基礎上應用各種優(yōu)化和變換手段，在保持代碼語義不變的基礎上提高性能79XLA（AcceleratedLinearAlgebra）XLA由谷歌開發(fā)的，并作為TensorFlow的一部分提供對計算圖進行優(yōu)化和編譯核心：HLO

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

StableHLO

硬件無關優(yōu)化

硬件相關優(yōu)化

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

生成

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

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

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

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

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

340MGPT-1

117MGPT-2

1.5BGPT-3

175BLaMDA137BPaLM540BGLM-130B130BChinchilla70BMegatron-11B

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

+

B

+

C

+

DDLP0DLP1DLP2DLP3DLP0DLP1DLP2DLP3All-ReduceABCDA

+

B

+

C

+

DA

+

B

+

C

+

DA

+

B

+

C

+

DA

+

B

+

C

+

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

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

,XA2]=XA=Y按列切分

參數(shù)矩陣AKA=[A1,A2]N/2N/2Y1=

XA1Y2=

XA2MK/2K/2

X=[X1,X2]

按行切分

參數(shù)矩陣ANA2A1Y1=X1A1NY2=X2A2NMX1K/2K/2MX2

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