HugeCTR源碼簡單走讀

前言

這段時間除了開發(fā)算子之外，還在做一些推薦系統(tǒng)相關(guān)的工作，這期間主要看的是HugeCTR的代碼，其性能優(yōu)異，系統(tǒng)不復雜，代碼結(jié)構(gòu)較扁平，整體還是比較清晰。在這段時間看源碼的過程中也算是對HugeCTR有一點了解，這篇博客主要梳理下HugeCTR代碼的結(jié)構(gòu)，以及他在MLPERF中做的一些優(yōu)化。

倉庫地址：NVIDIA-Merlin/HugeCTR

MLPERF博客：

Boosting NVIDIA MLPerf Training v1.1 Performance with Full Stack Optimization

MLPerf v1.0 Training Benchmarks: Insights into a Record-Setting NVIDIA Performance

整體結(jié)構(gòu)

倉庫主體結(jié)構(gòu)如下，一些不重要的目錄就省去了

- HugeCTR 主要的源碼目錄
- gpu_cache GPU 帶緩存機制的Embedding實現(xiàn)
- onnx_converter onnx模型轉(zhuǎn)換器
- samples 模型示例目錄，包含如DLRM, DeepFM等常用模型
- sparse_operation_kit 稀疏操作庫，可以作為一款插件搭配在其他框架上使用，如Tensorflow

我們主要還是關(guān)注 HugeCTR 這個主目錄，里面分別存放了頭文件(include)和對應的實現(xiàn)(src)，我們基于頭文件目錄來看下結(jié)構(gòu)：

HugeCTR
 | - include
     | - collectives 通信相關(guān)部分
     | - cpu CPU版本實現(xiàn)全部集中放在這一目錄
     | - data_readers 數(shù)據(jù)加載器部分
     | - embedding_training_cache Embedding訓練緩存機制
     | - embedding 不同類型Embedding實現(xiàn)部分，如混合Embedding等
     | - hashtable GPU哈希表
     | - embedding 不同類型Embedding實現(xiàn)部分，如混合Embedding等
     | - hps 系統(tǒng)組件實現(xiàn)部分，包含內(nèi)存池，各種數(shù)據(jù)庫后端等
     | - inference 推理實現(xiàn)
     | - layers 不同神經(jīng)網(wǎng)絡(luò)層實現(xiàn)
     | - optimizers 不同優(yōu)化器實現(xiàn)
     | - prims 我理解是提供cuda基礎(chǔ)操作部分，如reduce，矩陣求外積等基礎(chǔ)操作
     | - pybind cpp接口導出到python部分
     | - regularizer L1/2 regularizer實現(xiàn)
     | - resource_managers 系統(tǒng)資源管理
     | - shuffle （不知道這部分是干啥的，有了解的朋友也可以幫忙補充下）

直接硬翻源碼我覺得是有點難的，我的方法是從一個模型開始，看其分別涉及到了哪些代碼，下面我們就以官方的DLRM示例來看源碼，并針對我比較熟悉的算子實現(xiàn)展開。

Python這部分接口HugeCTR走的是Keras風格，習慣PyTorch的朋友可能一時還不太習慣

CreateSolver

第一部分是構(gòu)建了一個Solver，看起來是一個全局模型主體配置的東西：

solver?=?hugectr.CreateSolver(max_eval_batches?=?51,
??????????????????????????????batchsize_eval?=?1769472,
??????????????????????????????batchsize?=?55296,
??????????????????????????????vvgpu?=?[[0,1,2,3,4,5,6,7]],
??????????????????????????????repeat_dataset?=?True,
??????????????????????????????lr?=?24.0,
??????????????????????????????...
??????????????????????????????is_dlrm?=?True)

對應我們找到其pybind文件solver_wrapper.hpp，在cpp層創(chuàng)建了一個Solver對象，并將參數(shù)set進去：

??std::unique_ptr?solver(new?Solver());
??solver->model_name?=?model_name;
??solver->seed?=?seed;
??solver->lr_policy?=?lr_policy;
??solver->lr?=?lr;
??...

DataReaderParams

這部分是配置數(shù)據(jù)讀取的一些參數(shù)

reader?=?hugectr.DataReaderParams(data_reader_type?=?hugectr.DataReaderType_t.RawAsync,
??????????????????????????????????source?=?["./train_data.bin"],
??????????????????????????????????eval_source?=?"./test_data.bin"
??????????????????????????????????...

其作為一個結(jié)構(gòu)體，具體在 DataReaderParams::DataReaderParams，并在后續(xù)的model中傳入并調(diào)用

Model::Model(const?Solver&?solver,?const?DataReaderParams&?reader_params,...)

CreateOptimizer

顯然這部分是創(chuàng)建整個模型使用的優(yōu)化器

optimizer?=?hugectr.CreateOptimizer(optimizer_type?=?hugectr.Optimizer_t.SGD,
????????????????????????????????????update_type?=?hugectr.Update_t.Local,
????????????????????????????????????atomic_update?=?True)

在optimizer_wrapper.hpp里面，對不同優(yōu)化器的參數(shù)進行配置：

??OptHyperParams?opt_hyper_params;
??//?給Adam優(yōu)化器的beta1?beta2配置
??opt_hyper_params.adam.beta1?=?beta1;
??...
??//?設(shè)置Adagrad初始累加值
??opt_hyper_params.adagrad.initial_accu_value?=?initial_accu_value;
??...
??opt_hyper_params.momentum.factor?=?momentum_factor;
??opt_hyper_params.nesterov.mu?=?momentum_factor;
??opt_hyper_params.sgd.atomic_update?=?atomic_update;

實例化Model

做完前置工作后，這里正式實例化一個Model主體，傳入solver，reader，optimizer這三者

model?=?hugectr.Model(solver,?reader,?optimizer)

在其頭文件中model.hpp描述了Model對象的一些成員，這里面包含的挺雜的，諸如：

GpuLearningRateSchedulers?gpu_lr_sches_;?lr_sch_;?//?學習率調(diào)度器
std::vector<std::shared_ptrfloat>>>?train_weight_buff_list_;?//?用于訓練時float32類型權(quán)重分配顯存
std::vector<std::shared_ptr>>?train_weight_buff_half_list_;?//?用于開啟混合精度時，對half類型權(quán)重分配顯存
std::vector<std::string>?data_input_info_;?//?輸入data信息，這里是字符串形式
std::map<std::string,?std::vector<size_t>>?tensor_shape_info_;?//?tensor形狀信息
std::vector<std::pair<std::string,?std::string>>
??????input_output_info_;?//?每一層輸入，輸出信息
std::vector<std::string>?layer_info_;?//?layer信息，有一個枚舉類型維護string到實際網(wǎng)絡(luò)層的映射關(guān)系
...?

就不一一列舉了。

在Model::Model里面，根據(jù)傳進來的信息做一系列初始化：

//?使用多少gpu訓練
for?(size_t?i?=?0;?i?get_local_gpu_count();?i++)?{
????train_weight_buff_list_.emplace_back(blobs_buff_list_[i]->create_block<float>());
????...
????auto?id?=?resource_manager_->get_local_gpu(i)->get_local_id();
????//?如果開啟混合精度，那么就給half版本的wgrad設(shè)置以及allreduce通信操作，并給float版本的wgrad設(shè)置float類型分配器
????if?(solver_.use_mixed_precision)?{
??????wgrad_buff_half_list_.emplace_back(
??????????(solver_.grouped_all_reduce)
????????????????std::dynamic_pointer_cast>(exchange_wgrad_)
????????????????????->get_network_wgrad_buffs()[id]
??????????????:?std::dynamic_pointer_cast>(exchange_wgrad_)
????????????????????->get_network_wgrad_buffs()[id]);
??????wgrad_buff_list_.emplace_back(blobs_buff_list_[i]->create_block<float>());
????}?else?{
??????//?不開啟混合精度，就給float版本的wgrad設(shè)置對應的allreduce通信操作
??????wgrad_buff_list_.emplace_back(...);
??????wgrad_buff_half_list_.emplace_back(...);??//?placeholder
????}
??}
??
??...

??//?initialize?optimizer
??init_optimizer(opt_params_,?solver_,?opt_params_py);
??init_learning_rate_scheduler(lr_sch_,?solver_,?gpu_lr_sches_,?resource_manager_);

構(gòu)建Input，SparseEmbedding

model.add(hugectr.Input(label_dim?=?1,?label_name?=?"label",
????????????????????????dense_dim?=?13,?dense_name?=?"dense",
????????????????????????...))
model.add(hugectr.SparseEmbedding(embedding_type?=?hugectr.Embedding_t.HybridSparseEmbedding,?
????????????????????????????workspace_size_per_gpu_in_mb?=?15000,
????????????????????????????...)))

這里調(diào)用model的add方法添加輸入層和Embedding層，我們先看下add方法，在對應的pybind綁定中model_wrapper.hpp對應四種重載：

//?1.?輸入層
.def("add",?pybind11::overload_cast(&HugeCTR::Model::add),?pybind11::arg("input"))
//?2.?Embedding層
.def("add",?pybind11::overload_cast(&HugeCTR::Model::add),
?????pybind11::arg("sparse_embedding"))
//?3.?全連接層
.def("add",?pybind11::overload_cast(&HugeCTR::Model::add),
?????pybind11::arg("dense_layer"))
//?4.?多層全連接層
.def("add",?pybind11::overload_cast(&HugeCTR::Model::add),
?????pybind11::arg("group_dense_layer"))

Input和SparseEmbedding重載實現(xiàn)對應在model.cpp內(nèi)，這里就不再展開了。

增加全連接層

DLRM主體計算部分是由兩部分bottom_mlp和top_mlp組成的，因此模型后續(xù)代碼就是在構(gòu)建這部分：

model.add(hugectr.GroupDenseLayer(group_layer_type?=?hugectr.GroupLayer_t.GroupFusedInnerProduct,
????????????????????????????bottom_name_list?=?["dense"],
????????????????????????????top_name_list?=?["fc1",?"fc2",?"fc3"],
????????????????????????????num_outputs?=?[512,?256,?128],
????????????????????????????last_act_type?=?hugectr.Activation_t.Relu))???????????????????
model.add(hugectr.DenseLayer(bottom_names?=?["fc3","sparse_embedding1"],
????????????????????????????top_names?=?["interaction1",?"interaction1_grad"]...))
model.add(hugectr.GroupDenseLayer(...))
model.add(hugectr.DenseLayer(...))

• bottom_name_list和top_name_list來表示輸入Tensor列表，輸出Tensor列表，這樣后續(xù)層可以根據(jù)這個名字來實現(xiàn)網(wǎng)絡(luò)層相連。
• num_outputs表示的是全連接層輸出維度大小
• last_act_type表示最后一層全連接層的激活層類型，這里是ReLU

下面我們以全連接層來看一下一個網(wǎng)絡(luò)層的具體實現(xiàn)，首先看一下基類Layer的頭文件layer.hpp

??...
??//?前向計算邏輯
??virtual?void?fprop(bool?is_train)?=?0;
??//?反向計算邏輯
??virtual?void?bprop()?=?0;
??virtual?std::string?get_no_trained_params_in_string()?{?return?std::string();?}
??void?init_params(const?curandGenerator_t&?generator);

??Layer(const?std::shared_ptr&?gpu_resource,
????????std::vector?initializer_types?=?std::vector())
??????:?gpu_resource_(gpu_resource),?initializer_types_(initializer_types)?{}
??Layer(const?Layer&)?=?delete;
??Layer&?operator=(const?Layer&)?=?delete;
??virtual?~Layer()?=?default;
??//?參數(shù)初始化
??virtual?void?initialize()?{}
??//?算法搜索，比如全連接層，HugeCTR會通過試跑選擇一個最快的算法
??virtual?void?search_algorithm()?{}

我們以fused_relu_bias_fully_connected_layer.cu為具體例子：

構(gòu)造函數(shù)

在構(gòu)造函數(shù)中，根據(jù)tensor的形狀信息得到m,n,k，并推算出kernel，bias tensor的形狀等：

??size_t?m?=?bottom_tensor_dim[0];
??size_t?n?=?top_tensor_dim[1];
??size_t?k?=?bottom_tensor_dim[1];

??std::vector<size_t>?kernel_dim?=?{k,?n};
??std::vector<size_t>?bias_dim?=?{1,?n};
??std::vector<size_t>?identity_dim?=?{1,?m};

接著對這些tensor分配顯存：

??{
????Tensor2<float>?tensor;
????master_weights_buff->reserve(kernel_dim,?&tensor);
????weights_.push_back(tensor);
??}
??{
????Tensor2<float>?tensor;
????master_weights_buff->reserve(bias_dim,?&tensor);
????weights_.push_back(tensor);
??}
??...

initialize

這個Fuse矩陣乘是通過cublasLt實現(xiàn)的，因此在這個函數(shù)里做了一些cublasLt所需要的初始化，如矩陣信息，設(shè)置計算類型，設(shè)置epilogue(指定cublasLt的fuse模式)

??HCTR_LIB_THROW(cublasLtMatmulDescSetAttribute(cublas_op_desc_,?CUBLASLT_MATMUL_DESC_TRANSA,
????????????????????????????????????????????????&trans,?sizeof(trans)));
??HCTR_LIB_THROW(cublasLtMatmulDescSetAttribute(cublas_op_desc_,?CUBLASLT_MATMUL_DESC_TRANSB,
????????????????????????????????????????????????&trans,?sizeof(trans)));
??...
??cublasLtEpilogue_t?epi?=?CUBLASLT_EPILOGUE_RELU_AUX_BIAS;?//?設(shè)置epilogue
??...
??//?創(chuàng)建kernel,?bias,?output的矩陣維度，數(shù)據(jù)類型
??HCTR_LIB_THROW(cublasLtMatrixLayoutCreate(&cublas_kernel_desc_,?CUDA_R_16F,?n,?k,?n));
??HCTR_LIB_THROW(cublasLtMatrixLayoutCreate(&cublas_bottom_desc_,?CUDA_R_16F,?k,?m,?k));
??HCTR_LIB_THROW(cublasLtMatrixLayoutCreate(&cublas_top_desc_,?CUDA_R_16F,?n,?m,?n));
??...

initialize_dgrad() 和 initialize_wgrad() 則是給矩陣乘后向?qū)?次矩陣乘分別做上述類似初始化，這里不再贅述

fprop

調(diào)用cublasLtMatmul進行前向計算：

??const?__half*?kernel?=?weights_half_[0].get_ptr();
??const?__half*?bias?=?weights_half_[1].get_ptr();
??const?__half*?bottom?=?get_bottom_tensor_fprop(is_train).get_ptr();
??...
??HCTR_LIB_THROW(cublasLtMatmul(...));

bprop也是類似的邏輯

search_algorithm

HugeCTR設(shè)定了一個最大算法個數(shù)，通過cublasLtMatmulAlgoGetHeuristic啟發(fā)式搜索算法接口獲取當前全連接層可用的算法，并使用cudaEvent進行計時，選取表現(xiàn)最好的算法：

??//?獲取當前可用算法
??cublasLtMatmulHeuristicResult_t?heuristic_result[max_algo_count]?=?{0};
??int?algo_count?=?0;
??HCTR_LIB_THROW(cublasLtMatmulAlgoGetHeuristic(...));
??
??//?遍歷所有算法
??for?(int?algoIdx?=?0;?algoIdx?????cublasStatus_t?status?=?CUBLAS_STATUS_SUCCESS;
????const?float?alpha?=?1.0f;
????const?float?beta?=?0.0f;
????HCTR_LIB_THROW(cudaEventRecord(start,?get_gpu().get_stream()));
????//?執(zhí)行repeat_num次矩陣乘
????for?(size_t?i?=?0;?i???????status?=?cublasLtMatmul(...);
????}
????//?記錄時間
????HCTR_LIB_THROW(cudaEventRecord(stop,?get_gpu().get_stream()));
????HCTR_LIB_THROW(cudaEventSynchronize(stop));
????HCTR_LIB_THROW(cudaEventElapsedTime(&time,?start,?stop));
????//?Avg?Time(ms)?for?this?alorithm?for?fprop?GEMM
????time?=?time?/?repeat_num;
????...
????//?更新最佳時間
????if?(time???????shortestTime?=?time;
??????//?把當前最佳的算法拷貝到falgo_k，在計算過程中使用falgo_k
??????memcpy(&falgo_k_,?&heuristic_result[algoIdx].algo,?sizeof(falgo_k_));
????}
??}

這也是為什么在nsys前半部分，能看到一堆密密麻麻的矩陣乘

模型搭建完以后，后續(xù)調(diào)用compile, fit執(zhí)行訓練，這部分也可以在model.cpp看到對應實現(xiàn)，這里就不展開了，下面我們講下HugeCTR在MLPERF提及到的一些優(yōu)化

MLPERF1.0

Hybrid Embedding

頻繁的Embedding交換是模型訓練過程中的一個重要瓶頸，對此HugeCTR實現(xiàn)了HybridEmbedding。輸入中會存在重復id，因此一開始會剔除掉重復的數(shù)據(jù)，對應反向傳播也做相應處理。此外它還針對數(shù)據(jù)做了統(tǒng)計，根據(jù)頻率分為高頻Embedding，低頻Embedding。高頻Embedding以數(shù)據(jù)并行實現(xiàn)，這樣能夠在一個batch內(nèi)刪掉更多重復的數(shù)據(jù)，減少Embedding交換，而低頻Embedding以模型并行實現(xiàn)。

通信優(yōu)化

筆者不太了解通信方面的知識

all2all和allreduce耗時在拓展模型過程中是很重要的一環(huán)，對于比較小的message，多節(jié)點的all2all吞吐量受限于IB的消息速率的限制，為此將All2All分為節(jié)點內(nèi)All2All，節(jié)點間All2All。并將高頻Embedding和MLP的AllReduce放在一個AllReduce操作內(nèi)完成，以減少延遲。

數(shù)據(jù)讀取優(yōu)化

采用Linux的異步方式讀取，以達到IO峰值

重疊MLP和Embedding

在DLRM中，bottom_mlp部分和Embedding之間不存在依賴，因此做了如下的流水線重疊：

• BottomMLP前向過程和Embedding前向進行重疊
• 高頻Embedding在更新local權(quán)重時和AllReduce重疊
• MLP權(quán)重更新和節(jié)點內(nèi)All2All重疊

cublasLt算子融合

cublasLt可以通過epilogue來選擇不同算子融合方式，比如 matmul+bias, matmul+bias+relu，以及對應的后向矩陣乘等

CUDA Graph

為了減少kernel launch開銷，將模型的所有操作都包到一個 CUDA Graph 內(nèi)。

關(guān)于CUDA Graph可以參考 https://zhuanlan.zhihu.com/p/467466998

MLPERF1.1

Hybrid Embedding索引預計算

在之前的Hybrid Embedding中需要計算索引來決定在哪兒讀取對應的Embedding，而索引計算只依賴于輸入數(shù)據(jù)，這些數(shù)據(jù)可以在提前幾個iter時候預取好(Prefetch)，預先計算好Index，以隱藏延遲

通信與計算之間更好的重疊

這里就不過多闡述了，這里的圖描述的十分詳細

異步梯度計算

在矩陣乘中，其反向?qū)?個矩陣乘，而這兩個矩陣乘接受相同的輸入dy，分別輸出weight，input的梯度。因此這兩個計算可以重疊起來，具體在代碼中FusedReluBiasFullyConnectedLayer::bprop()

//?dgrad
HCTR_LIB_THROW(cublasLtMatmul(...,?get_gpu().get_stream()));

//?bgrad+wgrad
HCTR_LIB_THROW(cublasLtMatmul(...,?get_gpu().get_comp_overlap_stream()));

在nsys中，是這樣：

Better Fusion

其中一個是CublasLt提供了更多的fuse方式，另外一個是在混合精度情況下，將fp32權(quán)重cast成fp16的部分，放到Optimizer更新時候做，這樣就避免單獨啟動Cast Kernel，在SGD優(yōu)化器代碼中可以看到對應的操作sgd_optimizer.cu

template?<typename?T>
__device__?inline?void?sgd_update_device(int?len,?float*?weight,?__half*?weight_half,
?????????????????????????????????????????const?T*?wgrad,?float?lr,?float?scaler)?{
??...
??weight[i]?-=?lr?*?gi;
??weight_half[i]?=?(__half)weight[i];??
}