[原始碼解析] NVIDIA HugeCTR,GPU版本引數伺服器---(3)

羅西的思考發表於2022-02-17

[原始碼解析] NVIDIA HugeCTR,GPU版本引數伺服器---(3)

0x00 摘要

在本系列中,我們介紹了 HugeCTR,這是一個面向行業的推薦系統訓練框架,針對具有模型並行嵌入和資料並行密集網路的大規模 CTR 模型進行了優化。

本文主要介紹HugeCTR所依賴的輸入資料和一些基礎資料結構。其中借鑑了HugeCTR原始碼閱讀 這篇大作,特此感謝。因為 HugeCTR 實際上是一個具體而微的深度學習系統,所以它也實現了眾多基礎功能,值得想研究深度學習框架的朋友仔細研讀。

本系列其他文章如下:

[原始碼解析] NVIDIA HugeCTR,GPU 版本引數伺服器 --(1)

[原始碼解析] NVIDIA HugeCTR,GPU版本引數伺服器--- (2)

0x01 回顧

我們首先回歸一下前文內容,流水線邏輯關係如下:

img

訓練流程如下:

img

基於前文知識,我們接下來看看如何處理資料。

0x02 資料集

HugeCTR 目前支援三種資料集格式,即NormRawParquet,具體格式參見如下:

Fig. 1: (a) Norm (b) Raw (c) Parquet Dataset Formats

2.1 Norm

為了最大化資料載入效能並最小化儲存,Norm 資料集格式由一組二進位制資料檔案和一個 ASCII 格式的檔案列表組成。模型檔案應指定訓練和測試(評估)集的檔名,樣本中的元素(鍵)最大數目和標籤維度,具體如圖 1(a)所示。

2.1.1 資料檔案

一個資料檔案是一個讀取執行緒的最小讀取粒度,因此每個檔案列表中至少需要10個檔案才能達到最佳效能。資料檔案由header和實際表格(tabular )資料組成。

Header定義:

typedef  struct DataSetHeader_ {
   long  long error_check;       //0: 沒有錯誤檢查;1:check_num 
  long  long number_of_records; //此資料檔案中的樣本數
  long  long label_dim;          //標籤的維度
  long  long density_dim;        //密集特徵的維度
  long  long slot_num;           //每個嵌入的 slot_num 
  long  long reserved[ 3 ];      //保留以備將來使用
資料集頭;

資料定義(每個樣本):

typedef struct Data_ {
  int length;                   //此示例中的位元組數(可選:僅在 check_sum 模式下)
  float label[label_dim];
  float dense[dense_dim];
  Slot slots[slot_num];
  char checkbits;               //此樣本的校驗位(可選:僅在checksum模式下)
} Data;

typedef struct Slot_ {
  int nnz;
  unsigned int*  keys; //可在配置檔案的 `solver` 物件中使用 `"input_key_type"` 更改為 `long long`
} Slot;

資料欄位(field)通常有很多樣本。每個樣本以格式化為整數的標籤開始,然後是nnz(非零數)和使用 long long(或無符號整數)格式的輸入key,如圖 1(a)所示。

categorical 的輸入key分佈到插槽(slot)中,不允許重疊。例如:slot[0] = {0,10,32,45}, slot[1] = {1,2,5,67}。如果有任何重疊,它將導致未定義的行為。例如,給定slot[0] = {0,10,32,45}, slot[1] = {1,10,5,67},查詢10鍵的表將產生不同的結果,結果根據插槽分配給 GPU 的方式。

2.1.2 檔案列表

檔案列表的第一行應該是資料集中資料檔案的數量,然後是這些檔案的路徑,具體如下所示:

$ cat simple_sparse_embedding_file_list.txt
10
./simple_sparse_embedding/simple_sparse_embedding0.data
./simple_sparse_embedding/simple_sparse_embedding1.data
./simple_sparse_embedding/simple_sparse_embedding2.data
./simple_sparse_embedding/simple_sparse_embedding3.data
./simple_sparse_embedding/simple_sparse_embedding4.data
./simple_sparse_embedding/simple_sparse_embedding5.data
./simple_sparse_embedding/simple_sparse_embedding6.data
./simple_sparse_embedding/simple_sparse_embedding7.data
./simple_sparse_embedding/simple_sparse_embedding8.data
./simple_sparse_embedding/simple_sparse_embedding9.data

使用例子如下:

reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Norm,
                                  source = ["./wdl_norm/file_list.txt"],
                                  eval_source = "./wdl_norm/file_list_test.txt",
                                  check_type = hugectr.Check_t.Sum)

2.2 Raw

Raw 資料集格式與 Norm 資料集格式的不同之處在於訓練資料出現在一個二進位制檔案中,並且使用 int32。圖 1 (b) 顯示了原始資料集樣本的結構。

注意:此格式僅接受獨熱資料。

Raw資料集格式只能與嵌入型別 LocalizedSlotSparseEmbeddingOneHot 一起使用。

例子:

reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Raw,
                                  source = ["./wdl_raw/train_data.bin"],
                                  eval_source = "./wdl_raw/validation_data.bin",
                                  check_type = hugectr.Check_t.Sum)

2.3 Parquet

Parquet 是一種面向列的、開源的資料格式。它可用於 Apache Hadoop 生態系統中的任何專案。為了減小檔案大小,它支援壓縮和編碼。圖 1 (c) 顯示了一個示例 Parquet 資料集。有關其他資訊,請參閱parquet 文件

請注意以下事項:

  • Parquet 資料載入器當前不支援巢狀列型別。
  • 不允許列中有任何缺失值。
  • 與 Norm 資料集格式一樣,標籤和密集特徵列應使用浮點格式。
  • Slot 特徵列應使用 Int64 格式。
  • Parquet 檔案中的資料列可以按任何順序排列。
  • 要從每個 parquet 檔案中的所有行和每個標籤、密集(數字)和槽(分類)特徵的列索引對映中獲取所需資訊,需要一個單獨的_metadata.json檔案。

例子 _metadata.json:

{
"file_stats": [{"file_name": "file1.parquet", "num_rows": 6528076}, {"file_name": "file2.parquet", "num_rows": 6528076}],
"cats": [{"col_name": "C11", "index": 24}, {"col_name": "C24", "index": 37}, {"col_name": "C17", "index": 30}, {"col_name": "C7", "index": 20}, {"col_name": "C6", "index": 19}],
"conts": [{"col_name": "I5", "index": 5}, {"col_name": "I13", "index": 13}, {"col_name": "I2", "index": 2}, {"col_name": "I10", "index": 10}],
"labels": [{"col_name": "label", "index": 0}]
}

使用如下:

reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Parquet,
                                  source = ["./criteo_data/train/_file_list.txt"],
                                  eval_source = "./criteo_data/val/_file_list.txt",
                                  check_type = hugectr.Check_t.Non,
                                  slot_size_array = [278899, 355877, 203750, 18573, 14082, 7020, 18966, 4, 6382, 1246, 49, 185920, 71354, 67346, 11, 2166, 7340, 60, 4, 934, 15, 204208, 141572, 199066, 60940, 9115, 72, 34])

我們提供了通過一個選項slot_size_array,可以為每個插槽新增偏移量。slot_size_array是一個長度等於槽數的陣列。為了避免新增offset後出現key重複,我們需要保證第i個slot的key範圍在0到slot_size_array[i]之間。我們將以這種方式進行偏移:對於第 i 個槽鍵,我們將其新增偏移量 slot_size_array[0] + slot_size_array[1] + ... + slot_size_array[i - 1]。在上面提到的配置片段中,對於第 0 個插槽,將新增偏移量 0。對於第一個插槽,將新增偏移量 278899。對於第三個插槽,將新增偏移量 634776。

0x03 CSR 格式

嵌入層是基於CSR格式基礎之上搭建的,所以我們首先看看CSR格式。

3.1 什麼是CSR

稀疏矩陣指的是矩陣中的元素大部分是0的矩陣,實際上現實問題中大多數的大規模矩陣都是稀疏矩陣,因此就出現了很多專門針對稀疏矩陣的高效儲存格式,Compressed Sparse Row(CSR)就是其中之一。

這是最簡單的一種格式,每一個元素需要用一個三元組來表示,分別是(行號,列號,數值),對應上圖右邊的一列。這種方式簡單,但是記錄單資訊多(行列),每個三元組自己可以定位,因此空間不是最優。

CSR需要三類資料來表達:數值,列號,行偏移。它不是用三元組來表示一個元素,而是一個整體編碼方式。

  • 數值:一個元素。
  • 列號 :元素的列號,
  • 行偏移:某一行的第一個元素在values裡面的起始偏移位置。

上圖中,第一行元素1是0偏移,第二行元素2是2偏移,第三行元素5是4偏移,第4行元素6是7偏移。最後會在行偏移之後加上矩陣總的元素個數,本例子中是9。

3.2 HugeCTR 之中的CSR

我們從中找出一個例子看看。因為只是用來儲存slot裡的sparse key,所以沒有列號,因為一個slot裡的sparse key可以直接順序儲存。

* For example data:
*   4,5,1,2
*   3,5,1
*   3,2
* Will be convert to the form of:
* row offset: 0,4,7,9
* value: 4,5,1,2,3,5,1,3,2

我們再從原始碼之中找一些資訊 samples/ncf/preprocess-20m.py。

def write_hugeCTR_data(huge_ctr_data, filename='huge_ctr_data.dat'):
    print("Writing %d samples"%huge_ctr_data.shape[0])
    with open(filename, 'wb') as f:
        #write header
        f.write(ll(0)) # 0: no error check; 1: check_num
        f.write(ll(huge_ctr_data.shape[0])) # the number of samples in this data file
        f.write(ll(1)) # dimension of label
        f.write(ll(1)) # dimension of dense feature
        f.write(ll(2)) # long long slot_num
        for _ in range(3): f.write(ll(0)) # reserved for future use

        for i in tqdm.tqdm(range(huge_ctr_data.shape[0])):
            f.write(c_float(huge_ctr_data[i,2])) # float label[label_dim];
            f.write(c_float(0)) # dummy dense feature
            f.write(c_int(1)) # slot 1 nnz: user ID
            f.write(c_uint(huge_ctr_data[i,0]))
            f.write(c_int(1)) # slot 2 nnz: item ID
            f.write(c_uint(huge_ctr_data[i,1]))

3.3 操作類

3.3.1 定義

這裡只給出成員變數,具體可以和上面csr格式進行印證。

class CSR {
 private:
  const size_t num_rows_;       /**< num rows. */
  const size_t max_value_size_; /**< number of element of value the CSR matrix will have for
                                num_rows rows. */

  Tensor2<T> row_offset_tensor_;
  Tensor2<T> value_tensor_; /**< a unified buffer for row offset and value. */
  T* row_offset_ptr_;       /**< just offset on the buffer, note that the length of it is
                             * slot*batchsize+1.
                             */
  T* value_ptr_;            /**< pointer of value buffer. */

  size_t size_of_row_offset_; /**< num of rows in this CSR buffer */
  size_t size_of_value_;      /**< num of values in this CSR buffer */

  size_t check_point_row_;   /**< check point of size_of_row_offset_. */
  size_t check_point_value_; /**< check point of size_of_value__. */
}

3.3.2 建構函式

建構函式之中,會在GPU之上進行分配記憶體。

/**
 * Ctor
 * @param num_rows num of rows is expected
 * @param max_value_size max size of value buffer.
 */
CSR(size_t num_rows, size_t max_value_size)
    : num_rows_(num_rows),
      max_value_size_(max_value_size),
      size_of_row_offset_(0),
      size_of_value_(0) {
  std::shared_ptr<GeneralBuffer2<CudaHostAllocator>> buff =
      GeneralBuffer2<CudaHostAllocator>::create();
  buff->reserve({num_rows + 1}, &row_offset_tensor_);
  buff->reserve({max_value_size}, &value_tensor_);
  buff->allocate();

  row_offset_ptr_ = row_offset_tensor_.get_ptr();
  value_ptr_ = value_tensor_.get_ptr();
}

3.3.3 生成新行

new_row 之中會生成新行,並且把目前value總數設定到row_offset之中。

/**
 * Insert a new row to CSR
 * Whenever you want to add a new row, you need to call this.
 * When you have pushed back all the values, you need to call this method
 * again.
 */
inline void new_row() {  // call before push_back values in this line
  if (size_of_row_offset_ > num_rows_) CK_THROW_(Error_t::OutOfBound, "CSR out of bound");
  row_offset_ptr_[size_of_row_offset_] = static_cast<T>(size_of_value_);
  size_of_row_offset_++;
}

3.3.4 插入資料

這裡會插入資料,並且增加value總數。

/**
 * Push back a value to this object.
 * @param value the value to be pushed back.
 */
inline void push_back(const T& value) {
  if (size_of_value_ >= max_value_size_)
    CK_THROW_(Error_t::OutOfBound, "CSR out of bound " + std::to_string(max_value_size_) +
                                       "offset" + std::to_string(size_of_value_));
  value_ptr_[size_of_value_] = value;
  size_of_value_++;
}

0x04 基礎資料結構

因為 HugeCTR 實際上是一個具體而微的深度學習系統,所以它也實現了眾多基礎功能,為了更好的進行分析,我們需要首先介紹一些基礎資料結構。以下只給出各個類的成員變數和必要函式。

4.1 張量

首先就是最基礎的張量概念。

4.1.1 TensorBuffer2

TensorBuffer2 是張量底層的資料,也許聯絡到 PyTorch 的 data 或者 storage 可以更好的理解。

class TensorBuffer2 {
 public:
  virtual ~TensorBuffer2() {}
  virtual bool allocated() const = 0;
  virtual void *get_ptr() = 0;
};

4.1.2 Tensor2

這就對應了TF或者PyTorch的張量。

template <typename T>
class Tensor2 {
  std::vector<size_t> dimensions_;
  size_t num_elements_;
  std::shared_ptr<TensorBuffer2> buffer_;
}

成員函式我們選介紹兩個如下:

static Tensor2 stretch_from(const TensorBag2 &bag) {
  return Tensor2(bag.dimensions_, bag.buffer_);
}

TensorBag2 shrink() const {
  return TensorBag2(dimensions_, buffer_, TensorScalarTypeFunc<T>::get_type());
}

具體如下:

4.1.3 Tensors2

Tensors2 就是 Tensor2 的一個vector。

template <typename T> using Tensors2 = std::vector<Tensor2<T>>;

4.1.4 TensorBag2

PyTorch 之中也有一些Bag字尾名字的類,比如 nn.Embedding和nn.EmbeddingBag。當構建袋子模型時,做一個Embedding跟隨Sum或是Mean常見的。對於可變長度序列,nn.EmbeddingBag 來提供了更加高效和更快速的處理方式,特別是對於可變長度序列。

在 HugeCTR,TensorBag2 可以認為是把 Tensor 放在袋子裡統一處理的類。

class TensorBag2 {
  template <typename T>
  friend class Tensor2;

  std::vector<size_t> dimensions_;
  std::shared_ptr<TensorBuffer2> buffer_;
  TensorScalarType scalar_type_;
};

using TensorBags2 = std::vector<TensorBag2>;

關於 Tensor 和 Bag 的聯絡,可以參見下面的函式。

template <typename T>
Tensors2<T> bags_to_tensors(const std::vector<TensorBag2> &bags) {
  Tensors2<T> tensors;
  for (const auto &bag : bags) {
    tensors.push_back(Tensor2<T>::stretch_from(bag));
  }
  return tensors;
}

template <typename T>
std::vector<TensorBag2> tensors_to_bags(const Tensors2<T> &tensors) {
  std::vector<TensorBag2> bags;
  for (const auto &tensor : tensors) {
    bags.push_back(tensor.shrink());
  }
  return bags;
}

4.1.5 SparseTensor

SparseTensor 是 Sparse 型別的張量,這是3.2 版本加入的,目的是為了統一處理CSR格式,或者說是統一處理稀疏矩陣,可以有效儲存和處理大多數元素為零的張量。後續在讀取資料到GPU時候會有分析。我們對比一下 CSR 格式,就可以看出來其內部機制就對應了CSR 的 rowoffset 和 value。其具體定義如下:

template <typename T>
class SparseTensor {
  std::vector<size_t> dimensions_;
  std::shared_ptr<TensorBuffer2> value_buffer_;
  std::shared_ptr<TensorBuffer2> rowoffset_buffer_;
  std::shared_ptr<size_t> nnz_;  // maybe size_t for FixedLengthSparseTensor
  size_t rowoffset_count_;
};

示意圖如下:

我們從中找出一個例子看看。因為只是用來儲存slot裡的sparse key,所以沒有列號,因為一個slot裡的sparse key可以直接順序儲存。

* For example data:
*   4,5,1,2
*   3,5,1
*   3,2
* Will be convert to the form of:
* row offset: 0,4,7,9
* value: 4,5,1,2,3,5,1,3,2

對應下圖:

成員函式介紹如下:

static SparseTensor stretch_from(const SparseTensorBag &bag) {
  return SparseTensor(bag.dimensions_, bag.value_buffer_, bag.rowoffset_buffer_, bag.nnz_,
                      bag.rowoffset_count_);
}

SparseTensorBag shrink() const {
  return SparseTensorBag(dimensions_, value_buffer_, rowoffset_buffer_, nnz_, rowoffset_count_,
                         TensorScalarTypeFunc<T>::get_type());
}
PyTorch

PyTorch 有 sparse_coo_tensor 可以實現類似的功能。PyTorch 支援不同layout的張量,大家可以從 torch/csrc/utils/tensor_layouts.cpp 找到,比如 at::Layout::Strided,at::Layout::Sparse,at::Layout::SparseCsr,at::Layout::Mkldnn 等等,這些對應了不同的記憶體佈局模式。

使用稀疏張量時候,提供一對 dense tensors:一個value張量,一個二維indice張量,也有其他輔助引數。

>>> i = [[1, 1]]
>>> v =  [3, 4]
>>> s=torch.sparse_coo_tensor(i, v, (3,))
>>> s
tensor(indices=tensor([[1, 1]]),
       values=tensor(  [3, 4]),
       size=(3,), nnz=2, layout=torch.sparse_coo)
TensorFlow

TensorFlow 也有 SparseTensor 型別來表示多維稀疏資料。一個 SparseTensor 使用三個稠密張量來表示:

  • indices 表示稀疏張量的非零元素座標。
  • values 則對應每個非零元素的值。
  • shape 表示本稀疏張量轉換為稠密形式後的形狀。

比如下面程式碼:

indices = tf.constant([[0, 0], [1, 1], [2,2]], dtype=tf.int64)
values = tf.constant([1, 2, 3], dtype=tf.float32)
shape = tf.constant([3, 3], dtype=tf.int64)
sparse = tf.SparseTensor(indices=indices,
                            values=values,
                            dense_shape=shape)
dense = tf.sparse_tensor_to_dense(sparse, default_value=0)
with tf.Session() as session:
    sparse, dense = session.run([sparse, dense])
    print('Sparse is :\n', sparse)
    print('Dense is :\n', dense)

列印出來如下:

Sparse is :
 SparseTensorValue(indices=array([[0, 0],
       [1, 1],
       [2, 2]]), values=array([1., 2., 3.], dtype=float32), dense_shape=array([3, 3]))
Dense is :
 [[1. 0. 0.]
 [0. 2. 0.]
 [0. 0. 3.]]

4.1.6 SparseTensorBag

這個類似 TensorBag 的功能,具體如下:

class SparseTensorBag {
  template <typename T>
  friend class SparseTensor;

  std::vector<size_t> dimensions_;
  std::shared_ptr<TensorBuffer2> value_buffer_;
  std::shared_ptr<TensorBuffer2> rowoffset_buffer_;
  std::shared_ptr<size_t> nnz_;
  size_t rowoffset_count_;
  TensorScalarType scalar_type_;

  SparseTensorBag(const std::vector<size_t> &dimensions,
                  const std::shared_ptr<TensorBuffer2> &value_buffer,
                  const std::shared_ptr<TensorBuffer2> &rowoffset_buffer,
                  const std::shared_ptr<size_t> &nnz, const size_t rowoffset_count,
                  TensorScalarType scalar_type)
      : dimensions_(dimensions),
        value_buffer_(value_buffer),
        rowoffset_buffer_(rowoffset_buffer),
        nnz_(nnz),
        rowoffset_count_(rowoffset_count),
        scalar_type_(scalar_type) {}

 public:
  SparseTensorBag() : scalar_type_(TensorScalarType::None) {}
  const std::vector<size_t> &get_dimensions() const { return dimensions_; }
};

4.1.7 向量類

以下是兩個向量類,用來方便使用者使用。

using TensorBags2 = std::vector<TensorBag2>;

template <typename T>
using SparseTensors = std::vector<SparseTensor<T>>;

4.2 記憶體

我們接下來看看一些記憶體相關類。

4.2.1 Allocator

首先看看如何為tensor等變數分配記憶體。

4.2.1.1 HostAllocator

HostAllocator 作用是在host之上管理記憶體。

class HostAllocator {
 public:
  void *allocate(size_t size) const { return malloc(size); }
  void deallocate(void *ptr) const { free(ptr); }
};

後面幾個實現都是呼叫了CUDA函式來進行記憶體分配,比如 cudaHostAlloc,有興趣讀者可以深入學習。

4.2.1.2 CudaHostAllocator

呼叫CUDA方法在主機上分配記憶體

class CudaHostAllocator {
 public:
  void *allocate(size_t size) const {
    void *ptr;
    CK_CUDA_THROW_(cudaHostAlloc(&ptr, size, cudaHostAllocDefault));
    return ptr;
  }
  void deallocate(void *ptr) const { CK_CUDA_THROW_(cudaFreeHost(ptr)); }
};
4.2.1.3 CudaManagedAllocator

cudaMallocManaged 分配旨在供主機或裝置程式碼使用的記憶體,算是一種統一分配記憶體的方法。

class CudaManagedAllocator {
 public:
  void *allocate(size_t size) const {
    void *ptr;
    CK_CUDA_THROW_(cudaMallocManaged(&ptr, size));
    return ptr;
  }
  void deallocate(void *ptr) const { CK_CUDA_THROW_(cudaFree(ptr)); }
};
4.2.1.4 CudaAllocator

該類是在裝置上分配記憶體。

class CudaAllocator {
 public:
  void *allocate(size_t size) const {
    void *ptr;
    CK_CUDA_THROW_(cudaMalloc(&ptr, size));
    return ptr;
  }
  void deallocate(void *ptr) const { CK_CUDA_THROW_(cudaFree(ptr)); }
};

4.2.2 GeneralBuffer2

分析完如何分配記憶體,我們接下來看看如何封裝記憶體,具體通過 GeneralBuffer2 完成的。GeneralBuffer2 可以認為是一個對大段記憶體的統一封裝,具體在其上可以有若干Tensor。

4.2.2.1 定義

這裡都忽略了成員函式,內部類也忽略了成員函式。

  • allocator :具體記憶體分配器,也區分在GPU分配還是CPU分配。
  • ptr_ :指向分配的記憶體;
  • total_size_in_bytes_ :記憶體大小;
  • reserved_buffers_ :前期預留buffer,後續會統一分配;

具體內部類為:

  • BufferInternal 是介面。
  • TensorBufferImpl 是 Tensor2 對應的buffer實現。
  • BufferBlockImpl 則是在構建網路時候會用到。

具體程式碼如下:

template <typename Allocator>
class GeneralBuffer2 : public std::enable_shared_from_this<GeneralBuffer2<Allocator>> {
  
  class BufferInternal {
   public:
    virtual ~BufferInternal() {}
    virtual size_t get_size_in_bytes() const = 0;
    virtual void initialize(const std::shared_ptr<GeneralBuffer2> &buffer, size_t offset) = 0;
  };

  class TensorBufferImpl : public TensorBuffer2, public BufferInternal {
    size_t size_in_bytes_;
    std::shared_ptr<GeneralBuffer2> buffer_;
    size_t offset_;
  };

  template <typename T>
  class BufferBlockImpl : public BufferBlock2<T>, public BufferInternal {
    size_t total_num_elements_;
    std::shared_ptr<TensorBufferImpl> buffer_impl_;
    Tensor2<T> tensor_;
    bool finalized_;
    std::vector<std::shared_ptr<BufferInternal>> reserved_buffers_;
  };

  Allocator allocator_;
  void *ptr_;
  size_t total_size_in_bytes_;
  std::vector<std::shared_ptr<BufferInternal>> reserved_buffers_;
}
4.2.2.2 TensorBufferImpl

就是指向了一個 GeneralBuffer2,然後設定了自己的offset和大小。

void initialize(const std::shared_ptr<GeneralBuffer2> &buffer, size_t offset) {
  buffer_ = buffer;
  offset_ = offset;
}
4.2.2.2 BufferBlockImpl 關鍵函式

BufferBlockImpl 和 TensorBufferImpl 可以來比較一下。

其中,BufferBlock2 是 BufferBlockImpl 的介面類。

template <typename T>
class BufferBlock2 {
 public:
  virtual ~BufferBlock2() {}
  virtual void reserve(const std::vector<size_t> &dimensions, Tensor2<T> *tensor) = 0;
  virtual Tensor2<T> &as_tensor() = 0;
};

BufferBlockImpl 是一組連續的 Tensor,某些特定的實現需要連續的記憶體,比如權重。

std::shared_ptr<BufferBlock2<float>> train_weight_buff = blobs_buff->create_block<float>();

// 省略其他程式碼......

network->train_weight_tensor_ = train_weight_buff->as_tensor();

BufferBlockImpl 多了一個reserve方法,用來預留記憶體空間,在此空間之上生成內部tensor。

void reserve(const std::vector<size_t> &dimensions, Tensor2<T> *tensor) override {
  if (finalized_) {
    throw std::runtime_error(ErrorBase + "Buffer block is finalized.");
  }
  size_t num_elements = get_num_elements_from_dimensions(dimensions);
  size_t size_in_bytes = num_elements * TensorScalarSizeFunc<T>::get_element_size();

  std::shared_ptr<TensorBufferImpl> buffer_impl =
      std::make_shared<TensorBufferImpl>(size_in_bytes);
  reserved_buffers_.push_back(buffer_impl);

  *tensor = Tensor2<T>(dimensions, buffer_impl);

  total_num_elements_ += num_elements;
}

initialize 會對內部進行配置

void initialize(const std::shared_ptr<GeneralBuffer2> &buffer, size_t offset) {
  size_t local_offset = 0;
  for (const std::shared_ptr<BufferInternal> &buffer_impl : reserved_buffers_) {
    buffer_impl->initialize(buffer, offset + local_offset);
    local_offset += buffer_impl->get_size_in_bytes();
  }
  reserved_buffers_.clear();

  if (!finalized_) {
    buffer_impl_ = std::make_shared<TensorBufferImpl>(
        total_num_elements_ * TensorScalarSizeFunc<T>::get_element_size());
    tensor_ = Tensor2<T>({total_num_elements_}, buffer_impl_);
    finalized_ = true;
  }
  buffer_impl_->initialize(buffer, offset);
}
4.2.2.3 GeneralBuffer2 關鍵函式

reserve 方法會把某一個張量對應的記憶體需求用 TensorBufferImpl 的形式記錄在reserved_buffers_之中,然後生成這個張量,而且就是用TensorBufferImpl 生成。

template <typename T>
void reserve(const std::vector<size_t> &dimensions, Tensor2<T> *tensor) {
  if (allocated()) {
    throw std::runtime_error(ErrorBase + "General buffer is finalized.");
  }

  size_t size_in_bytes =
      get_num_elements_from_dimensions(dimensions) * TensorScalarSizeFunc<T>::get_element_size();

  std::shared_ptr<TensorBufferImpl> buffer_impl =
      std::make_shared<TensorBufferImpl>(size_in_bytes);
  reserved_buffers_.push_back(buffer_impl);

  *tensor = Tensor2<T>(dimensions, buffer_impl);
}

create_block 會針對BufferBlock2進行建立。

template <typename T>
std::shared_ptr<BufferBlock2<T>> create_block() {
  if (allocated()) {
    throw std::runtime_error(ErrorBase + "General buffer is finalized.");
  }
  std::shared_ptr<BufferBlockImpl<T>> block_impl = std::make_shared<BufferBlockImpl<T>>();
  reserved_buffers_.push_back(block_impl);
  return block_impl;
}

allocate 會遍歷註冊的 BufferInternal,累積其總大小,最後呼叫 allocator_ 進行分配記憶體。

void allocate() {
  if (ptr_ != nullptr) {
    throw std::runtime_error(ErrorBase + "Memory has already been allocated.");
  }

  size_t offset = 0;
  for (const std::shared_ptr<BufferInternal> &buffer : reserved_buffers_) {
    // 對 BufferInternal(比如TensorBufferImpl)內部進行配置
    buffer->initialize(this->shared_from_this(), offset);
    size_t size_in_bytes = buffer->get_size_in_bytes();
    if (size_in_bytes % 32 != 0) {
      size_in_bytes += (32 - size_in_bytes % 32);
    }
    offset += size_in_bytes;
  }
  reserved_buffers_.clear();
  total_size_in_bytes_ = offset;

  if (total_size_in_bytes_ != 0) {
    ptr_ = allocator_.allocate(total_size_in_bytes_);
  }
}
4.2.4 小結

至此,Tensor的邏輯擴充一下:

  • TensorBufferImpl 的 buffer 是GeneralBuffer2;
  • GeneralBuffer2 的 ptr 是由CudaAllocator在GPU之中分配的;GeneralBuffer2 可以認為是一個對大段記憶體的統一封裝,在其上可以有若干Tensor。這些Tensor先reserve記憶體,然後統一分配。
  • TensorBufferImpl 的 offset_ 就指向了 GeneralBuffer2 的 ptr 之中具體的某一個記憶體偏移;
  • BufferBlockImpl 用來實現一個連續的Tensor記憶體。

如果還有另外一個 Tensor2,則其 TensorBufferImpl.offset 會指向 GPU記憶體的另外一個offset,比如下面有兩個張量,Tensor 1 和 Tensor 2。

0xFF 參考

https://developer.nvidia.com/blog/introducing-merlin-hugectr-training-framework-dedicated-to-recommender-systems/

https://developer.nvidia.com/blog/announcing-nvidia-merlin-application-framework-for-deep-recommender-systems/

https://developer.nvidia.com/blog/accelerating-recommender-systems-training-with-nvidia-merlin-open-beta/

HugeCTR原始碼閱讀

embedding層如何反向傳播

https://web.eecs.umich.edu/~justincj/teaching/eecs442/notes/linear-backprop.html

稀疏矩陣儲存格式總結+儲存效率對比:COO,CSR,DIA,ELL,HYB

相關文章