[源碼解析] NVIDIA HugeCTR，GPU版本參數服務器--- (6) --- Distributed hash表
目錄
0x00 摘要
在這系列文章中，我們介紹了 HugeCTR，這是一個面向行業的推薦系統訓練框架，針對具有模型並行嵌入和數據並行密集網絡的大規模 CTR 模型進行了優化。
其中借鑑了HugeCTR源碼閱讀 這篇大作，特此感謝。
本系列其他文章如下：
[源碼解析] NVIDIA HugeCTR，GPU 版本參數服務器 --(1)
[源碼解析] NVIDIA HugeCTR，GPU版本參數服務器--- (2)
[源碼解析] NVIDIA HugeCTR，GPU版本參數服務器---(3)
[源碼解析] NVIDIA HugeCTR，GPU版本參數服務器--- (4)
[源碼解析] NVIDIA HugeCTR，GPU版本參數服務器--- (5) 嵌入式hash表
0x01 簡述
1.1 基類
DistributedSlotSparseEmbeddingHash類繼承自 IEmbedding，Embedding 是所有嵌入層的接口。
class IEmbedding {
public:
virtual ~IEmbedding() {}
virtual TrainState train(bool is_train, int i, TrainState state) { return TrainState(); }
// TODO: can we remove the default argument?
virtual void forward(bool is_train, int eval_batch = -1) = 0;
virtual void backward() = 0;
virtual void update_params() = 0;
virtual void init_params() = 0;
virtual void load_parameters(std::string sparse_model) = 0;
virtual void dump_parameters(std::string sparse_model) const = 0;
virtual void set_learning_rate(float lr) = 0;
// TODO: a workaround to enable GPU LR for HE only; need a better way
virtual GpuLearningRateSchedulers get_learning_rate_schedulers() const {
return GpuLearningRateSchedulers();
}
virtual size_t get_params_num() const = 0;
virtual size_t get_vocabulary_size() const = 0;
virtual size_t get_max_vocabulary_size() const = 0;
virtual Embedding_t get_embedding_type() const = 0;
virtual void load_parameters(BufferBag& buf_bag, size_t num) = 0;
virtual void dump_parameters(BufferBag& buf_bag, size_t* num) const = 0;
virtual void reset() = 0;
virtual void reset_optimizer() = 0;
virtual void dump_opt_states(std::ofstream& stream) = 0;
virtual void load_opt_states(std::ifstream& stream) = 0;
virtual const SparseEmbeddingHashParams& get_embedding_params() const = 0;
virtual std::vector<TensorBag2> get_train_output_tensors() const = 0;
virtual std::vector<TensorBag2> get_evaluate_output_tensors() const = 0;
virtual void check_overflow() const = 0;
virtual void get_forward_results_tf(const bool is_train, const bool on_gpu,
void* const forward_result) = 0;
virtual cudaError_t update_top_gradients(const bool on_gpu, const void* const top_gradients) = 0;
};
1.2 功能
在 DistributedSlotSparseEmbeddingHash 之中，嵌入表中的一些插槽被分配給多個GPU，稱爲分佈式插槽。例如，slot-0 被分配到GPU-0/GPU-1上，slot-1 被分配到GPU-0/GPU-1上。嵌入表被封裝在哈希表中，或者說哈希表是嵌入表的前置條件。哈希表一些相關成員變量如下：
- 哈希表中的鍵稱爲 hash_table_key。
- 哈希表中的值稱爲 hash_table_value_index，表示嵌入特徵在嵌入表中的行號（row number）。
- 嵌入特徵稱爲 hash_table_value，就是利用 hash_table_value_index（行號）在嵌入表之中找到的那一行。
DistributedSlotSparseEmbeddingHash 類實現了嵌入層的訓練過程所需的所有操作，包括前向傳播和後向傳播。前向傳播對應於API forward()。反向傳播分爲兩個階段的API：backward()和update_params()。該類還提供將哈希表（包括哈希表鍵hash_table_key、hash_table_value_index和hash_table_value）從主機文件上載到GPU（load_parameters 方法）的操作，以及將哈希表從GPU下載到主機文件（dump_parameters方法）的操作。
0x02 定義
2.1 思路
我們先自行想想看如何實現這個嵌入層，這樣會讓我們更好的理清楚思路。
- 高維矩陣 ：假設不考慮field的情況下，embedding矩陣大小是 A * B，A 是 one-hot 的長度，B是embedding size，假如 one-hot 長10000000，embedding size是64。Hash_key 是一個one-hot [0,0,..0,1,0,..,0]，其可以定位到 embedding矩陣的一行。假設 hash_key 的第367位置上是1，則就會找到embedding矩陣的367行，從 367 行得到一個64長度的dense vector。
- 數據特點 ：前面提到過，CTR的特點是高維，稀疏，這說明嵌入表10000000 行之中可能只有500行是有意義的數值，其餘爲空，
- **低維矩陣 ** ： HugeCTR 內部實際上不可能內部存放一個巨大矩陣，肯定是改用一個小型矩陣來存儲，比如1000 x 64 的小型矩陣。
- 轉換機制 ：所以需要有一個機制，把 367 這個高維嵌入表的 row index 映射到這個小型低維矩陣的 row index，通過一系列複雜的操作用時間來換取空間。這也就是 DistributedSlotSparseEmbeddingHash 的一系列成員變量所起到的作用。
2.2 代碼
DistributedSlotSparseEmbeddingHash 的定義如下，主要變量/概念爲：
CSR相關，可以結合CSR定義來印證。
- @param row_offset ：row_offset (CSR format of input sparse tensors)。
- @param hash_key ：value (CSR format of input sparse tensors)。
- @param nnz ：non-zero feature number per batch。
輸入/輸出數據：
- embedding_data_ ：這裏包括很多數據。
- 前面提到的 DataReader.output_ 就會被保存在這裏，就是 sparse input 信息。
- 這裏的 train_output_tensors_ 成員變量則是嵌入層最終的輸出，就是多個GPU之間互相作用之後得出的輸出。注意，train_output_tensors_ 在反向傳播時候居然還被用來作爲輸入梯度。
Hash相關：
- hash_tables_ ：這是一個哈希表vector，每一個元素都是一個hash_table（NvHashTable），本地每一個GPU對應這個vector之中的一個NvHashTable。目的是爲了把高維矩陣的row offset 轉換爲低維矩陣的 row offset。
- 在 hash_table 內部，邏輯上來看每一個元素可以認爲是 <key, value_index>（其實內部是個黑盒子，只是對外邏輯表示爲一個哈希表 <key, value_index>）；
- 哈希表中的鍵稱爲 hash_table_key，其格式是 CSR (CSR format of input sparse tensors)相關。
- 哈希表中的值稱爲 hash_table_value_index，表示 CSR 對應的嵌入特徵在嵌入表中的行號。
- hash_value_index_tensors_ ：embedding vector表的row index。就是低維矩陣的 row offset。
- 需要注意，其類型是 Tensors2，其類型是 std::vector<Tensor2
>，所以每一個GPU對應了該vector之中的一個元素。 - index 和 value 的行數相關。
- 內容是hash table value_index(row index of embedding)。
- 需要注意，其類型是 Tensors2，其類型是 std::vector<Tensor2
- hash_table_value_tensors_ ：embedding vector表的value。就是低維矩陣。
- 需要注意，其類型是 Tensors2，其類型是 std::vector<Tensor2
>，所以每一個GPU對應了該vector之中的一個元素。 - 其內容是embedding vector。
- 用hash_value_index_tensors_的結果在這裏查找一個 embedding vector。
- 需要注意，其類型是 Tensors2，其類型是 std::vector<Tensor2
中間數據：
- embedding_feature_tensors_ ： 嵌入層前向傳播的中間輸出，就是上面查找到的embedding vector的結果，但是沒有經過GPU之間的操作（ reduce-scatter等）；
- row_offset_allreduce_tensors_ ：allreduce之後的row_offset。
反向傳播：
- wgrad_tensors_ ：後向傳播的梯度，是backward之後產生的結果；
- embedding_optimizers_ : 嵌入層對應的優化器。
這裏有兩點說明：
- 爲了方便起見，hash_value_index_tensors_ 這樣雖然是一個向量的向量，我們後續都省略一步，當作向量來考慮。
- 需要對 hash_value_index_tensors_ 做進一步解釋：
- hash_value_index_tensors_ 起到了解耦合的作用，把低維矩陣和哈希表進行解耦合。因爲解耦合的原因，hash_value_index_tensors_ 並不應該知道 哈希表內部把高維矩陣的維度映射到了多大的低維矩陣，而 hash_value_index_tensors_ 大小也不應該隨之變化。
- 所以，hash_value_index_tensors_ 大小被固定爲：batch_size * nnz_per_slot，可以認爲就是CSR之中元素個數。因此 hash_value_index_tensors_ 實際上記錄了每個元素對應的低維矩陣offset 數值，hash_value_index_tensors_ 其事實上就是和CSR之中元素位置一一對應。
- 因此，最終嵌入表查找時候，是通過CSR row offset 來找到 CSR之中每個元素，也找到了hash_value_index_tensors_ 這個表的index，從而就能找到其低維矩陣offset。
我們再從源碼之中找出部分註釋給大家看看幾個變量之間的關係，其查找邏輯是從上到下。
DistributedSlotSparseEmbeddingHash 具體定義如下：
template <typename TypeHashKey, typename TypeEmbeddingComp>
class DistributedSlotSparseEmbeddingHash : public IEmbedding {
using NvHashTable = HashTable<TypeHashKey, size_t>;
private:
// 前面提到的 DataReader.output_ 就會被保存在這裏。就是sparse input信息
EmbeddingData<TypeHashKey, TypeEmbeddingComp> embedding_data_;
// 是 hash_value, hash_value_index的實際存儲位置
std::vector<DistributedFilterKeyStorage<TypeHashKey>> filter_keys_storage_;
std::vector<std::shared_ptr<NvHashTable>> hash_tables_; /**< Hash table. */
// define tensors
Tensors2<float> hash_table_value_tensors_; /**< Hash table value. */
Tensors2<size_t> hash_value_index_tensors_; /**< Hash table value index. The index is
corresponding to the line number of the value. */
Tensors2<TypeEmbeddingComp>
embedding_feature_tensors_; /**< the output tensor of the forward(). */
Tensors2<TypeEmbeddingComp> wgrad_tensors_; /**< the input tensor of the backward(). */
Tensors2<TypeHashKey>
row_offset_allreduce_tensors_; /**< The temp memory to store the row_offset after all_reduce
operation among multi-gpu in forward(). */
Tensors2<TypeEmbeddingComp> utest_forward_temp_tensors_;
size_t max_vocabulary_size_; /**< Max vocabulary size for each GPU. */
size_t max_vocabulary_size_per_gpu_; /**< Max vocabulary size for each GPU. */
SparseEmbeddingFunctors functors_;
std::vector<EmbeddingOptimizer<TypeHashKey, TypeEmbeddingComp>> embedding_optimizers_;
}
因爲定義是模版類，所以具體拓展爲如下：
template class DistributedSlotSparseEmbeddingHash<unsigned int, float>;
template class DistributedSlotSparseEmbeddingHash<long long, float>;
template class DistributedSlotSparseEmbeddingHash<unsigned int, __half>;
template class DistributedSlotSparseEmbeddingHash<long long, __half>;
0x03 HashTable
因爲DistributedSlotSparseEmbeddingHash 用到了 using NvHashTable = HashTable<TypeHashKey, size_t>，所以我們先看看 HashTable。這部分對應的是上面總圖第一步，就是如何從 hash table 之中拿到低維嵌入表的 index。在後文之中，我們用 HashTable/哈希表來指定 DistributedSlotSparseEmbeddingHash 內部使用的真正的哈希表。
3.1 定義
HashTable 之中，很重要的成員變量是container_。
/**
* The HashTable class is wrapped by cudf library for hash table operations on single GPU.
* In this class, we implement the GPU version of the common used operations of hash table,
* such as insert() / get() / set() / dump()...
*/
template <typename KeyType, typename ValType>
class HashTable {
const KeyType empty_key = std::numeric_limits<KeyType>::max();
private:
static const int BLOCK_SIZE_ =
256; /**< The block size of the CUDA kernels. The default value is 256. */
const float LOAD_FACTOR = 0.75f;
const size_t capacity_;
HashTableContainer<KeyType, ValType>* container_; /**< The object of the Table class which is
defined in the concurrent_unordered_map class. */
// Counter for value index
size_t* d_counter_; /**< The device counter for value index. */
size_t* d_container_size_;
};
3.2 HashTableContainer
container_ 的類型是HashTableContainer，其是 concurrent_unordered_map 的派生類，所以我們還是需要看看 concurrent_unordered_map。
template <typename KeyType, typename ValType>
class HashTableContainer
: public concurrent_unordered_map<KeyType, ValType, std::numeric_limits<KeyType>::max()> {
public:
HashTableContainer(size_t capacity)
: concurrent_unordered_map<KeyType, ValType, std::numeric_limits<KeyType>::max()>(
capacity, std::numeric_limits<ValType>::max()) {}
};
3.3 調用
爲了更好的分析，在看 concurrent_unordered_map 之前，我們需要看看如何調用HashTable。調用代碼是HugeCTR/src/embeddings/forward_per_gpu_functor.cu 之中的forward_per_gpu方法。這裏已經是 CUDA 代碼了。
emplate <typename TypeHashKey, typename TypeEmbeddingComp>
void SparseEmbeddingFunctors::forward_per_gpu(
size_t batch_size, size_t slot_num, size_t embedding_vec_size, int combiner, bool train,
const Tensor2<TypeHashKey> &row_offset, const Tensor2<TypeHashKey> &hash_key, size_t nnz,
HashTable<TypeHashKey, size_t> &hash_table, const Tensor2<float> &hash_table_value,
Tensor2<size_t> &hash_value_index, Tensor2<TypeEmbeddingComp> &embedding_feature,
cudaStream_t stream) {
try {
if (train) {
// 這裏會調用插入代碼
hash_table.get_insert(hash_key.get_ptr(), hash_value_index.get_ptr(), nnz, stream);
} else {
hash_table.get_mark(hash_key.get_ptr(), hash_value_index.get_ptr(), nnz, stream);
}
// do sum reduction
// 省略其他代碼
return;
}
可以看到，hash_key.get_ptr(), hash_value_index.get_ptr() 分別對應的是 _d_keys, _d_vals。
template <typename KeyType, typename ValType>
void NvHashTable<KeyType, ValType>::get_insert(const void *d_keys, void *d_vals, size_t len, cudaStream_t stream) {
const KeyType *_d_keys = reinterpret_cast<const KeyType*>(d_keys);
ValType *_d_vals = reinterpret_cast<ValType*>(d_vals);
return hashtable_.get_insert(_d_keys, _d_vals, len, stream);
}
然後調用到 get_insert。
template <typename KeyType, typename ValType>
void HashTable<KeyType, ValType>::get_insert(const KeyType* d_keys, ValType* d_vals, size_t len,
cudaStream_t stream) {
if (len == 0) {
return;
}
const int grid_size = (len - 1) / BLOCK_SIZE_ + 1;
get_insert_kernel<<<grid_size, BLOCK_SIZE_, 0, stream>>>(container_, d_keys, d_vals, len,
d_counter_);
}
template <typename Table>
__global__ void get_insert_kernel(Table* table, const typename Table::key_type* const keys,
typename Table::mapped_type* const vals, size_t len,
size_t* d_counter) {
ReplaceOp<typename Table::mapped_type> op;
const size_t i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < len) {
auto it = table->get_insert(keys[i], op, d_counter);
vals[i] = it->second;
}
}
所以最終調用到 concurrent_unordered_map 的 get_insert。
3.4 concurrent_unordered_map
concurrent_unordered_map 定義在 HugeCTR/include/hashtable/cudf/concurrent_unordered_map.cuh。
這是位於顯存中的map。從其註釋可知，其支持併發插入，但是不支持同時insert和probping。結合HugeCTR看，hugeCTR是同步訓練，pull操作只會調用 get，push操作只會調用insert，不存在同時insert和probping，所以滿足需求。
/**
* Does support concurrent insert, but not concurrent insert and probping.
*
* TODO:
* - add constructor that takes pointer to hash_table to avoid allocations
* - extend interface to accept streams
*/
template <typename Key, typename Element, Key unused_key, typename Hasher = default_hash<Key>,
typename Equality = equal_to<Key>,
typename Allocator = managed_allocator<thrust::pair<Key, Element>>,
bool count_collisions = false>
class concurrent_unordered_map : public managed {
public:
using size_type = size_t;
using hasher = Hasher;
using key_equal = Equality;
using allocator_type = Allocator;
using key_type = Key;
using value_type = thrust::pair<Key, Element>;
using mapped_type = Element;
using iterator = cycle_iterator_adapter<value_type*>;
using const_iterator = const cycle_iterator_adapter<value_type*>;
private:
const hasher m_hf;
const key_equal m_equal;
const mapped_type m_unused_element;
allocator_type m_allocator;
size_type m_hashtbl_size;
size_type m_hashtbl_capacity;
value_type* m_hashtbl_values; // 這個纔是hash數據結構位置
unsigned long long m_collisions;
};
3.4.1 get
我們先看看get操作，就是find方法。
// __forceinline__ 的意思是編譯爲內聯函數
// __host__ __device__ 表示是此函數同時爲主機和設備編譯
__forceinline__ __host__ __device__ const_iterator find(const key_type& k) const {
// 對key進行hash操作
size_type key_hash = m_hf(k);
// 進而得到table的相應index
size_type hash_tbl_idx = key_hash % m_hashtbl_size;
value_type* begin_ptr = 0;
size_type counter = 0;
while (0 == begin_ptr) {
value_type* tmp_ptr = m_hashtbl_values + hash_tbl_idx;
const key_type tmp_val = tmp_ptr->first;
// 找到key，跳出
if (m_equal(k, tmp_val)) {
begin_ptr = tmp_ptr;
break;
}
// key的位置是空，或者在table之內沒有找到
if (m_equal(unused_key, tmp_val) || counter > m_hashtbl_size) {
begin_ptr = m_hashtbl_values + m_hashtbl_size;
break;
}
hash_tbl_idx = (hash_tbl_idx + 1) % m_hashtbl_size;
++counter;
}
return const_iterator(m_hashtbl_values, m_hashtbl_values + m_hashtbl_size, begin_ptr);
}
3.4.2 insert
插入操作我們就看看之前的 get_insert。
hash_table.get_insert(hash_key.get_ptr(), hash_value_index.get_ptr(), nnz, stream);
就是以 csr 部分信息作爲 hash key，來獲得一個低維嵌入表之中的index，在 hash_value_index之中返回。我們首先看一個CSR示例。
* For example data:
* 3356
* 667
* 588
* Will be convert to the form of:
* row offset: 0,1,2,3
* value: 3356,667,588,3
我們就是使用 3356 作爲 hash_key，獲取 3356 對應的 hash_value_index，如果能找到就返回，找不到就插入一個構建的value，然後這個 value 會返回給 hash_value_index。
但是這裏有幾個繞的地方，因爲 HashTable內部也分桶，也有自己的key，hash_value，容易和其他數據結構弄混。具體邏輯是：
- 傳入一個數字 3356（CSR格式相關），還有一個 value_counter，就是目前 hash_value_index 的數值。
- 先 hash_value = m_hf(3356)。
- 用 current_index = hash_value % hashtbl_size 找到 m_hashtbl_values 之中的位置。
- 用 current_hash_bucket = &(hashtbl_values[current_index]) 這找到了一個bucket。
- key_type& existing_key = current_hash_bucket->first，這個纔是 hash table key
- volatile mapped_type& existing_value = current_hash_bucket->second，這個纔是我們最終需要的 table value。如果沒有，就遞增傳入的 value_counter。
所以，CSR 3356 是一個one-hot 的index，它對應了embeding表的一個index，但是因爲沒有那麼大的embedding，所以後面會構建一個小數據結構（低維矩陣） hash_value，傳入的 value_counter 就是這個 hash_value的index，value_counter 是遞增的，因爲 hash_value 的行號就是遞增的。
比如一共有1億個單詞，3356表示第3356個單詞。如果想表示 3356，667，588 這三個位置在這一億個單詞是有效的，最笨的辦法是弄個1億長度數組，把3356，667，588這三個位置設置爲 1，其他位置設置爲0，但是這樣太佔據空間且沒有意義。如果想省空間，就弄一個hash函數 m_hf，假如是選取最高位數爲 value，則得到：
m_hf(3356)=3
m_hf(667)=6
m_hf(588)=5
3，5，6 就是內部的 hash_value，叫做 hash_value（對應下面代碼），對應的內部存儲數組叫做 hashtbl_values。再梳理一下：3356是哈希表的key，3 是哈希表的value，但是因爲分桶了，所以在哈希表內部是放置在 hashtbl_values 之中。
hashtbl_values[3] = 1，hashtbl_values[6] = 2, hashtbl_values[5] =3
於是 1，2，3 就是我們外部想得到的 3356, 667, 588 對應的數據，就是低維矩陣的 row offset，對應下面代碼就是 existing_value。簡化版本的邏輯如下:
具體代碼如下：
// __forceinline__ 的意思是編譯爲內聯函數
// __host__ __device__ 表示是此函數同時爲主機和設備編譯
template <typename aggregation_type, typename counter_type, class comparison_type = key_equal,
typename hash_value_type = typename Hasher::result_type>
__forceinline__ __device__ iterator get_insert(const key_type& k, aggregation_type op,
counter_type* value_counter,
comparison_type keys_equal = key_equal(),
bool precomputed_hash = false,
hash_value_type precomputed_hash_value = 0) {
const size_type hashtbl_size = m_hashtbl_size;
value_type* hashtbl_values = m_hashtbl_values;
hash_value_type hash_value{0};
// If a precomputed hash value has been passed in, then use it to determine
// the write location of the new key
if (true == precomputed_hash) {
hash_value = precomputed_hash_value;
}
// Otherwise, compute the hash value from the new key
else {
hash_value = m_hf(k); // 3356作爲key，得到了一個hash_value
}
size_type current_index = hash_value % hashtbl_size; // 找到哪個位置
value_type* current_hash_bucket = &(hashtbl_values[current_index]); // 找到該位置的bucket
const key_type insert_key = k;
bool insert_success = false;
size_type counter = 0;
while (false == insert_success) {
// Situation %5: No slot: All slot in the hashtable is occupied by other key, both get and
// insert fail. Return empty iterator
// hash表已經滿了
if (counter++ >= hashtbl_size) {
return end();
}
key_type& existing_key = current_hash_bucket->first; // 這個纔是table key
volatile mapped_type& existing_value = current_hash_bucket->second; // 這個纔是table value
// 如果 existing_key == unused_key時，則當前哈希位置爲空，所以existing_key由atomicCAS更新爲insert_key。
// 如果 existing_key == insert_key時，這個位置已經被插入這個key了。
// 在任何一種情況下，都要執行existing_value和insert_value的atomic聚合，因爲哈希表是用聚合操作的標識值初始化的，所以在existing_value仍具有其初始值時，執行該操作是安全的
// Try and set the existing_key for the current hash bucket to insert_key
const key_type old_key = atomicCAS(&existing_key, unused_key, insert_key);
// If old_key == unused_key, the current hash bucket was empty
// and existing_key was updated to insert_key by the atomicCAS.
// If old_key == insert_key, this key has already been inserted.
// In either case, perform the atomic aggregation of existing_value and insert_value
// Because the hash table is initialized with the identity value of the aggregation
// operation, it is safe to perform the operation when the existing_value still
// has its initial value
// TODO: Use template specialization to make use of native atomic functions
// TODO: How to handle data types less than 32 bits?
// Situation #1: Empty slot: this key never exist in the table, ready to insert.
if (keys_equal(unused_key, old_key)) { // 如果沒有找到hash key
existing_value = (mapped_type)(atomicAdd(value_counter, 1)); // hash value 就遞增
break;
} // Situation #2+#3: Target slot: This slot is the slot for this key
else if (keys_equal(insert_key, old_key)) {
while (existing_value == m_unused_element) {
// Situation #2: This slot is inserting by another CUDA thread and the value is not yet
// ready, just wait
}
// Situation #3: This slot is already ready, get successfully and return (iterator of) the
// value
break;
}
// Situation 4: Wrong slot: This slot is occupied by other key, get fail, do nothing and
// linear probing to next slot.
// 此位置已經被其他key佔了，只能向後遍歷
current_index = (current_index + 1) % hashtbl_size;
current_hash_bucket = &(hashtbl_values[current_index]);
}
return iterator(m_hashtbl_values, m_hashtbl_values + hashtbl_size, current_hash_bucket);
}
0x04 構建
4.1 初始化
我們接下來看看如何構建 DistributedSlotSparseEmbeddingHash，代碼之中需要留意的是：
- hash_tables_ 之中，每一個元素對應一個GPU。
- train_keys 就是前面提到的 sparse_input，就是CSR format 相關的 row offset。
具體就是分配內存，hash_tables_的大小是本地GPU數目，即每個GPU對應一個hash表，用一個gpu卡上的最大sparse key 的個數來初始化hash table，這樣每個hash table能容納元素的最大數值就被固定住了。
template <typename TypeHashKey, typename TypeEmbeddingComp>
DistributedSlotSparseEmbeddingHash<TypeHashKey, TypeEmbeddingComp>::
DistributedSlotSparseEmbeddingHash(const SparseTensors<TypeHashKey> &train_keys,
const SparseTensors<TypeHashKey> &evaluate_keys,
const SparseEmbeddingHashParams &embedding_params,
const std::shared_ptr<ResourceManager> &resource_manager)
: embedding_data_(Embedding_t::DistributedSlotSparseEmbeddingHash, train_keys, evaluate_keys,
embedding_params, resource_manager) {
try {
// 得到一個gpu卡上最大sparse key個數
max_vocabulary_size_per_gpu_ = embedding_data_.embedding_params_.max_vocabulary_size_per_gpu;
max_vocabulary_size_ = max_vocabulary_size_per_gpu_ *
embedding_data_.get_resource_manager().get_global_gpu_count();
// 構建上下文
CudaDeviceContext context;
for (size_t id = 0; id < embedding_data_.get_resource_manager().get_local_gpu_count(); id++) {
context.set_device(embedding_data_.get_local_gpu(id).get_device_id());
// buf用來分配內存
// new GeneralBuffer objects
const std::shared_ptr<GeneralBuffer2<CudaAllocator>> &buf = embedding_data_.get_buffer(id);
embedding_optimizers_.emplace_back(max_vocabulary_size_per_gpu_,
embedding_data_.embedding_params_, buf);
{ // train_value_tensors_ 配置內存
Tensor2<TypeHashKey> tensor;
buf->reserve({embedding_data_.embedding_params_.get_batch_size(true),
embedding_data_.embedding_params_.max_feature_num},
&tensor);
embedding_data_.train_value_tensors_.push_back(tensor);
}
{ // evaluate_value_tensors_ 配置內存
Tensor2<TypeHashKey> tensor;
buf->reserve({embedding_data_.embedding_params_.get_batch_size(false),
embedding_data_.embedding_params_.max_feature_num},
&tensor);
embedding_data_.evaluate_value_tensors_.push_back(tensor);
}
{ // train_row_offsets_tensors_配置內存
Tensor2<TypeHashKey> tensor;
buf->reserve({embedding_data_.embedding_params_.get_batch_size(true) *
embedding_data_.embedding_params_.slot_num +
1},
&tensor);
embedding_data_.train_row_offsets_tensors_.push_back(tensor);
}
{ // evaluate_row_offsets_tensors_ 配置內存
Tensor2<TypeHashKey> tensor;
buf->reserve({embedding_data_.embedding_params_.get_batch_size(false) *
embedding_data_.embedding_params_.slot_num +
1},
&tensor);
embedding_data_.evaluate_row_offsets_tensors_.push_back(tensor);
}
{ embedding_data_.train_nnz_array_.push_back(std::make_shared<size_t>(0)); }
{ embedding_data_.evaluate_nnz_array_.push_back(std::make_shared<size_t>(0)); }
// new hash table value vectors
{ // hash_table_value_tensors_ 配置內存
Tensor2<float> tensor;
buf->reserve(
{max_vocabulary_size_per_gpu_, embedding_data_.embedding_params_.embedding_vec_size},
&tensor);
hash_table_value_tensors_.push_back(tensor);
}
// new hash table value_index that get() from HashTable
{ // hash_value_index_tensors_配置內存，注意，這裏配置的大小是 batch_size * max_feature_number
Tensor2<size_t> tensor;
buf->reserve({1, embedding_data_.embedding_params_.get_universal_batch_size() *
embedding_data_.embedding_params_.max_feature_num},
&tensor);
hash_value_index_tensors_.push_back(tensor);
}
// new embedding features reduced by hash table values(results of forward)
{ // embedding_feature_tensors_ 配置內存
Tensor2<TypeEmbeddingComp> tensor;
buf->reserve({embedding_data_.embedding_params_.get_universal_batch_size() *
embedding_data_.embedding_params_.slot_num,
embedding_data_.embedding_params_.embedding_vec_size},
&tensor);
embedding_feature_tensors_.push_back(tensor);
}
// new wgrad used by backward
{ // wgrad_tensors_ 配置內存
Tensor2<TypeEmbeddingComp> tensor;
buf->reserve({embedding_data_.embedding_params_.get_batch_size(true) *
embedding_data_.embedding_params_.slot_num,
embedding_data_.embedding_params_.embedding_vec_size},
&tensor);
wgrad_tensors_.push_back(tensor);
}
// new temp tensors used by update_params
{ // row_offset_allreduce_tensors_ 配置內存
Tensor2<TypeHashKey> tensor;
buf->reserve({1, embedding_data_.embedding_params_.get_universal_batch_size() *
embedding_data_.embedding_params_.slot_num +
1},
&tensor);
row_offset_allreduce_tensors_.push_back(tensor);
}
{ // utest_forward_temp_tensors_ 配置內存
Tensor2<TypeEmbeddingComp> tensor;
buf->reserve({embedding_data_.embedding_params_.get_universal_batch_size() *
embedding_data_.embedding_params_.slot_num,
embedding_data_.embedding_params_.embedding_vec_size},
&tensor);
utest_forward_temp_tensors_.push_back(tensor);
}
// temp storage for filter keys
{
size_t max_nnz = embedding_data_.embedding_params_.get_universal_batch_size() *
embedding_data_.embedding_params_.max_feature_num;
size_t rowoffset_count = embedding_data_.embedding_params_.slot_num *
embedding_data_.embedding_params_.get_universal_batch_size() +
1;
filter_keys_storage_.emplace_back(
buf, max_nnz, rowoffset_count, embedding_data_.get_local_gpu(id).get_global_id(),
embedding_data_.get_resource_manager().get_global_gpu_count());
}
// init GenenralBuffers to do real allocation
}
// hash_tables_的大小是本地GPU數目，即每個GPU對應一個hash表
hash_tables_.resize(embedding_data_.get_resource_manager().get_local_gpu_count());
#pragma omp parallel num_threads(embedding_data_.get_resource_manager().get_local_gpu_count())
{ // 並行分配內存
size_t id = omp_get_thread_num();
CudaDeviceContext context(embedding_data_.get_local_gpu(id).get_device_id());
// construct HashTable object: used to store hash table <key, value_index>
// 用一個gpu卡上的最大sparse key的個數來初始化hash table，這樣每個hash table能容納元素的最大數值就被固定住了。
hash_tables_[id].reset(new NvHashTable(max_vocabulary_size_per_gpu_));
embedding_data_.get_buffer(id)->allocate();
}
// 遍歷本地的GPU
for (size_t id = 0; id < embedding_data_.get_resource_manager().get_local_gpu_count(); id++) {
context.set_device(embedding_data_.get_local_gpu(id).get_device_id());
embedding_optimizers_[id].initialize(embedding_data_.get_local_gpu(id));
} // end of for(int id = 0; id < embedding_data_.get_local_gpu_count(); id++)
if (!embedding_data_.embedding_params_.slot_size_array.empty()) {
std::vector<TypeHashKey> embedding_offsets;
TypeHashKey slot_sizes_prefix_sum = 0;
for (size_t i = 0; i < embedding_data_.embedding_params_.slot_size_array.size(); i++) {
embedding_offsets.push_back(slot_sizes_prefix_sum);
slot_sizes_prefix_sum += embedding_data_.embedding_params_.slot_size_array[i];
}
for (size_t id = 0; id < embedding_data_.get_resource_manager().get_local_gpu_count(); ++id) {
CudaDeviceContext context(embedding_data_.get_local_gpu(id).get_device_id());
CK_CUDA_THROW_(
cudaMemcpy(embedding_data_.embedding_offsets_[id].get_ptr(), embedding_offsets.data(),
embedding_offsets.size() * sizeof(TypeHashKey), cudaMemcpyHostToDevice));
}
}
functors_.sync_all_gpus(embedding_data_.get_resource_manager());
} catch (const std::runtime_error &rt_err) {
std::cerr << rt_err.what() << std::endl;
throw;
}
return;
}
4.2 配置內存
我們要看看幾個關鍵變量的內存配置。
4.2.1 hash_table_value_tensors_
hash_table_value_tensors_ 的內存是 max_vocabulary_size_per_gpu_ * embedding_vec_size。
{ // hash_table_value_tensors_ 配置內存
Tensor2<float> tensor;
buf->reserve(
{max_vocabulary_size_per_gpu_, embedding_data_.embedding_params_.embedding_vec_size},
&tensor);
hash_table_value_tensors_.push_back(tensor);
}
而 max_vocabulary_size_per_gpu_計算如下：
max_vocabulary_size_per_gpu_ = embedding_data_.embedding_params_.max_vocabulary_size_per_gpu;
max_vocabulary_size_per_gpu 是在這裏做了配置。
SparseEmbedding::SparseEmbedding(Embedding_t embedding_type, size_t workspace_size_per_gpu_in_mb,
size_t embedding_vec_size, const std::string& combiner_str,
std::string sparse_embedding_name, std::string bottom_name,
std::vector<size_t>& slot_size_array,
std::shared_ptr<OptParamsPy>& embedding_opt_params,
const HybridEmbeddingParam& hybrid_embedding_param)
: embedding_type(embedding_type),
workspace_size_per_gpu_in_mb(workspace_size_per_gpu_in_mb),
embedding_vec_size(embedding_vec_size),
sparse_embedding_name(sparse_embedding_name),
bottom_name(bottom_name),
slot_size_array(slot_size_array),
embedding_opt_params(embedding_opt_params),
hybrid_embedding_param(hybrid_embedding_param) {
if (combiner_str == "sum") {
combiner = 0;
} else if (combiner_str == "mean") {
combiner = 1;
} else {
CK_THROW_(Error_t::WrongInput, "No such combiner type: " + combiner_str);
}
max_vocabulary_size_per_gpu =
(workspace_size_per_gpu_in_mb * 1024 * 1024) / (sizeof(float) * embedding_vec_size);
}
4.2.2 hash_value_index_tensors_
hash_value_index_tensors_ 大小爲 batch_size * max_feature_number。
// new hash table value_index that get() from HashTable
{ // hash_value_index_tensors_配置內存，注意，這裏配置的大小是 batch_size * max_feature_number
Tensor2<size_t> tensor;
buf->reserve({1, embedding_data_.embedding_params_.get_universal_batch_size() *
embedding_data_.embedding_params_.max_feature_num},
&tensor);
hash_value_index_tensors_.push_back(tensor);
}
max_feature_number 按照如下規則計算。
DataReaderSparseParam(const std::string& top_name_, const std::vector<int>& nnz_per_slot_,
bool is_fixed_length_, int slot_num_)
: top_name(top_name_),
nnz_per_slot(nnz_per_slot_),
is_fixed_length(is_fixed_length_),
slot_num(slot_num_),
type(DataReaderSparse_t::Distributed) {
max_feature_num = std::accumulate(nnz_per_slot.begin(), nnz_per_slot.end(), 0);
max_nnz = *std::max_element(nnz_per_slot.begin(), nnz_per_slot.end());
}
所以，hash_value_index_tensors_ 大小就是 batch_size * nnz_per_slot。
0x05 EmbeddingData
前面提到了 DistributedSlotSparseEmbeddingHash 如下成員變量會保存一些嵌入表信息。
EmbeddingData<TypeHashKey, TypeEmbeddingComp> embedding_data_;
我們來挖掘一下。
5.1 定義
EmbeddingData 定義如下，這裏有兩套成員變量，Tensors2 和 SparseTensors。
- Tensors2 如下：
- train_value_tensors_ 這個就會記錄sparse input，是CSR 的value。
- train_row_offsets_tensors_ 是CSR 的 row offset。
- train_nnz_array_ 是CSR 相關的nnz。
- train_output_tensors_ 這個是前向傳播的輸出。
- SparseTensors 如下：
- train_keys_ 會把 value，offset，nnz都整合在一起，這裏懷疑是在接口遷移，所以維護了兩套。爲何遷移？因爲
train_value_tensors_，train_row_offsets_tensors_，train_nnz_array_ 都是Tensor2，是普通張量，而 train_keys_ 是 SparseTensors，可以一個變量就搞定前面所有概念。 - evaluate_keys_ 是驗證集相關。
- train_keys_ 會把 value，offset，nnz都整合在一起，這裏懷疑是在接口遷移，所以維護了兩套。爲何遷移？因爲
所以，embedding_data_ 就是包攬了嵌入層的輸入和輸出。需要注意的是，這裏都是 Tensors2，可以認爲是 Tensor2 的列表，列表之中每一個Tensor2 對應了一個GPU。
template <typename TypeKey, typename TypeEmbeddingComp>
class EmbeddingData {
public:
const Embedding_t embedding_type_;
SparseEmbeddingHashParams embedding_params_; /**< Sparse embedding hash params. */
std::vector<std::shared_ptr<GeneralBuffer2<CudaAllocator>>>
bufs_; /**< The buffer for storing output tensors. */
Tensors2<TypeEmbeddingComp> train_output_tensors_; /**< The output tensors. */
Tensors2<TypeEmbeddingComp> evaluate_output_tensors_; /**< The output tensors. */
Tensors2<TypeKey> train_row_offsets_tensors_; /**< The row_offsets tensors of the input data. */
Tensors2<TypeKey> train_value_tensors_; /**< The value tensors of the input data. */
std::vector<std::shared_ptr<size_t>> train_nnz_array_;
Tensors2<TypeKey>
evaluate_row_offsets_tensors_; /**< The row_offsets tensors of the input data. */
Tensors2<TypeKey> evaluate_value_tensors_; /**< The value tensors of the input data. */
std::vector<std::shared_ptr<size_t>> evaluate_nnz_array_;
std::shared_ptr<ResourceManager> resource_manager_; /**< The GPU device resources. */
SparseTensors<TypeKey> train_keys_;
SparseTensors<TypeKey> evaluate_keys_;
Tensors2<TypeKey> embedding_offsets_;
}
5.2 構建
這裏有兩套構建函數，可能維護者在從舊接口切換到新接口。結合前後文，sparse_input 在 DistributedSlotSparseEmbeddingHash 構造函數之中是 train_keys 參數，在EmbeddingData 這裏就是train_value_tensors，所以，value_tensors 就是我們要關注的，從註釋可以知道，這是輸入數據的value tensors，指向了稀疏矩陣的 value vector。
/**
* The constructor of Embedding class.
* @param row_offsets_tensors the row_offsets tensors of the input data(refer to row offset vector
* in sparse matrix CSR format).
* @param value_tensors the value tensors of the input data(refer to value vector in sparse matrix
* CSR format).
* @param batchsize the batch size of the input data
* @param slot_num the number of slots of the hash table
* @param embedding_vec_size the dim size of the embedding feature vector.
* @param resource_manager the GPU device resource group
* @param scaler scaler factor for mixed precision
*/
EmbeddingData(const Tensors2<TypeKey>& train_row_offsets_tensors,
const Tensors2<TypeKey>& train_value_tensors,
const std::vector<std::shared_ptr<size_t>>& train_nnz_array,
const Tensors2<TypeKey>& evaluate_row_offsets_tensors,
const Tensors2<TypeKey>& evaluate_value_tensors,
const std::vector<std::shared_ptr<size_t>>& evaluate_nnz_array,
const Embedding_t embedding_type, const SparseEmbeddingHashParams& embedding_params,
const std::shared_ptr<ResourceManager>& resource_manager)
: embedding_type_(embedding_type),
embedding_params_(embedding_params),
train_row_offsets_tensors_(train_row_offsets_tensors),
train_value_tensors_(train_value_tensors),
train_nnz_array_(train_nnz_array),
evaluate_row_offsets_tensors_(evaluate_row_offsets_tensors),
evaluate_value_tensors_(evaluate_value_tensors),
evaluate_nnz_array_(evaluate_nnz_array),
resource_manager_(resource_manager) {
try {
// Error check
if (embedding_params.train_batch_size < 1 || embedding_params.evaluate_batch_size < 1 ||
embedding_params.slot_num < 1 || embedding_params.embedding_vec_size < 1) {
CK_THROW_(Error_t::WrongInput, "batchsize < 1 || slot_num < 1 || embedding_vec_size < 1");
}
if (embedding_params.embedding_vec_size > 1024) {
CK_THROW_(Error_t::WrongInput,
"the embedding_vec_size can not be more than 1024 in embedding layer");
}
size_t total_gpu_count = resource_manager_->get_global_gpu_count();
size_t local_gpu_count = resource_manager_->get_local_gpu_count();
if (train_row_offsets_tensors.size() != local_gpu_count ||
train_value_tensors.size() != local_gpu_count ||
evaluate_row_offsets_tensors.size() != local_gpu_count ||
evaluate_value_tensors.size() != local_gpu_count) {
CK_THROW_(
Error_t::WrongInput,
"either row_offsets_tensors.size() or value_tensors.size() isn't local_gpu_count_");
}
assert(bufs_.empty());
for (size_t i = 0; i < local_gpu_count; i++) {
std::shared_ptr<GeneralBuffer2<CudaAllocator>> buf =
GeneralBuffer2<CudaAllocator>::create();
bufs_.push_back(buf);
Tensor2<TypeEmbeddingComp> tensor;
buf->reserve({get_batch_size_per_gpu(true), embedding_params_.slot_num,
embedding_params_.embedding_vec_size},
&tensor);
train_output_tensors_.push_back(tensor);
buf->reserve({get_batch_size_per_gpu(false), embedding_params_.slot_num,
embedding_params_.embedding_vec_size},
&tensor);
evaluate_output_tensors_.push_back(tensor);
}
// value，offset，nnz又整合了進來
for (size_t i = 0; i < local_gpu_count; i++) {
train_keys_.emplace_back(train_value_tensors_[i], train_row_offsets_tensors_[i],
train_nnz_array_[i]);
evaluate_keys_.emplace_back(evaluate_value_tensors_[i], evaluate_row_offsets_tensors_[i],
evaluate_nnz_array_[i]);
}
} catch (const std::runtime_error& rt_err) {
std::cerr << rt_err.what() << std::endl;
throw;
}
return;
}
我們最終拓展如下，經過第 C 步之後，DistributedSlotSparseEmbeddingHash的成員變量 也指向了 GPU 內存，這裏依據構建函數的不同，train_output_tensors_，和 train_keys_ 可能（可能是因爲有兩種不同的構造方式，目前只是討論其中一種）都會指向用戶輸入訓練數據。
5.3 怎麼得到row_offset
5.3.1 問題
目前，我們只設置了EmbeddingData的train_keys/train_value_tensors_，但這是SparseTensor，其內部不僅僅有value，還有row_offset等專門針對稀疏矩陣的信息，所以這部分也要進行設置。
我們提前看看前向傳播，會發現其使用了類似 embedding_data_.get_row_offsets_tensors 進行運算。但是我們目前並沒有配置這樣的參數，只是配置了 train_keys。這個地方很繞，仔細看代碼，原來在前向傳播之中有使用 filter_keys_per_gpu 進行設置類似參數。
void forward(bool is_train, int eval_batch = -1) override {
// Read data from input_buffers_ -> look up -> write to output_tensors
#pragma omp parallel num_threads(embedding_data_.get_resource_manager().get_local_gpu_count())
{
size_t i = omp_get_thread_num();
CudaDeviceContext context(embedding_data_.get_local_gpu(i).get_device_id());
if (embedding_data_.embedding_params_.is_data_parallel) {
// 在這裏有操作
filter_keys_per_gpu(is_train, i, embedding_data_.get_local_gpu(i).get_global_id(),
embedding_data_.get_resource_manager().get_global_gpu_count());
}
// 部分前向操作
functors_.forward_per_gpu(embedding_data_.embedding_params_.get_batch_size(is_train),
embedding_data_.embedding_params_.slot_num,
embedding_data_.embedding_params_.embedding_vec_size, 0, is_train,
embedding_data_.get_row_offsets_tensors(is_train)[i],
embedding_data_.get_value_tensors(is_train)[i],
*embedding_data_.get_nnz_array(is_train)[i], *hash_tables_[i],
hash_table_value_tensors_[i], hash_value_index_tensors_[i],
embedding_feature_tensors_[i],
embedding_data_.get_local_gpu(i).get_stream());
}
// 省略後面代碼
// do reduce scatter
// scale for combiner=mean after reduction
// do average
}
return;
}
5.3.2 引用
我們仔細看看 EmbeddingData 的一些成員函數，發現他們都返回了引用。這就是關鍵，這些成員函數可以修改 EmbeddingData的內部成員變量，比如：get_row_offsets_tensors返回了一個引用。
Tensors2<TypeKey>& get_row_offsets_tensors(bool is_train) {
if (is_train) {
return train_row_offsets_tensors_;
} else {
return evaluate_row_offsets_tensors_;
}
}
類似的，比如get_output_tensors，get_input_keys，get_row_offsets_tensors，get_value_tensors，get_nnz_array 都返回引用，這說明 EmbeddingData 大部分成員變量都是可以被引用來修改的。
5.3.3 修改
具體配置就是在 filter_keys_per_gpu 這裏進行，就是利用 train_keys 進行配置其他成員變量，具體方法涉及到CUDA一些集合運算，有興趣的讀者可以自行研究。
template <typename TypeHashKey, typename TypeEmbeddingComp>
void DistributedSlotSparseEmbeddingHash<TypeHashKey, TypeEmbeddingComp>::filter_keys_per_gpu(
bool is_train, size_t id, size_t global_id, size_t global_num) {
const SparseTensor<TypeHashKey> &all_gather_key = embedding_data_.get_input_keys(is_train)[id];
// 這裏拿到了get_row_offsets_tensors
Tensor2<TypeHashKey> rowoffset_tensor = embedding_data_.get_row_offsets_tensors(is_train)[id];
Tensor2<TypeHashKey> value_tensor = embedding_data_.get_value_tensors(is_train)[id];
std::shared_ptr<size_t> nnz_ptr = embedding_data_.get_nnz_array(is_train)[id];
auto &filter_keys_storage = filter_keys_storage_[id];
auto &stream = embedding_data_.get_local_gpu(id).get_stream();
if (all_gather_key.get_dimensions().size() != 2) {
CK_THROW_(Error_t::WrongInput, "distributed embedding all gather key dimension != 2");
}
size_t batch_size = embedding_data_.embedding_params_.get_batch_size(is_train);
size_t slot_num = (all_gather_key.rowoffset_count() - 1) / batch_size;
size_t rowoffset_num = batch_size * slot_num + 1;
size_t rowoffset_num_without_zero = rowoffset_num - 1;
if (rowoffset_tensor.get_num_elements() != rowoffset_num) {
std::cout << rowoffset_tensor.get_num_elements() << " " << rowoffset_num << std::endl;
CK_THROW_(Error_t::WrongInput, "filter rowoffset size not match.");
}
// select value
{
distributed_embedding_kernels::HashOp<TypeHashKey> select_op{global_id, global_num};
size_t size_in_bytes = filter_keys_storage.temp_value_select_storage.get_size_in_bytes();
cub::DeviceSelect::If(filter_keys_storage.temp_value_select_storage.get_ptr(), size_in_bytes,
all_gather_key.get_value_ptr(), value_tensor.get_ptr(),
filter_keys_storage.value_select_num.get_ptr(), all_gather_key.nnz(),
select_op, stream);
}
// select rowoffset
{
cudaMemsetAsync(filter_keys_storage.rowoffset_select.get_ptr(), 0,
filter_keys_storage.rowoffset_select.get_size_in_bytes(), stream);
{
constexpr int block_size = 512;
int grid_size = (rowoffset_num_without_zero - 1) / block_size + 1;
distributed_embedding_kernels::select_rowoffset<<<grid_size, block_size, 0, stream>>>(
all_gather_key.get_rowoffset_ptr(), rowoffset_num_without_zero,
all_gather_key.get_value_ptr(), filter_keys_storage.rowoffset_select.get_ptr(), global_id,
global_num);
}
{
// 這裏會進行修改設置rowoffset_tensor
size_t size_in_bytes =
filter_keys_storage.temp_rowoffset_select_scan_storage.get_size_in_bytes();
cub::DeviceScan::InclusiveSum(
filter_keys_storage.temp_rowoffset_select_scan_storage.get_ptr(), size_in_bytes,
filter_keys_storage.rowoffset_select.get_ptr(), rowoffset_tensor.get_ptr(), rowoffset_num,
stream);
}
// select nnz
cudaMemcpyAsync(nnz_ptr.get(), filter_keys_storage.value_select_num.get_ptr(), sizeof(size_t),
cudaMemcpyDeviceToHost, stream);
}
}
於是，在進行具體前向操作之前，會把EmbeddingData內部都進行配置，分別指向GPU之中的相應數據。
0x06 優化器
DistributedSlotSparseEmbeddingHash 內部也存在一些優化器。
std::vector<EmbeddingOptimizer<TypeHashKey, TypeEmbeddingComp>> embedding_optimizers_;
我們接下來分析一下。
6.1 定義
優化器定義如下：
template <typename TypeHashKey, typename TypeEmbeddingComp>
class EmbeddingOptimizer {
Tensor2<void> temp_storage_encode_tensors_;
Tensor2<void> temp_storage_sort_tensors_; /**< The temp memory for the CUB lib sorting
API in update_params(). */
Tensor2<void> temp_storage_scan_tensors_; /**< The temp memory for the CUB lib scaning API
in update_params(). */
Tensor2<TypeHashKey> sample_id_tensors_; /**< The temp memory to store the sample ids of hash
table value in update_params(). */
Tensor2<TypeHashKey> sample_id_sort_tensors_; /**< The temp memory to store the sorted sample
ids of hash table value in update_params(). */
Tensor2<size_t> hash_value_index_sort_tensors_; /**< The temp memory to store the sorted hash
table value indexes in update_params(). */
Tensor2<size_t> hash_value_index_sort_unique_tensors_;
Tensor2<uint32_t> hash_value_index_count_tensors_;
Tensor2<uint32_t> new_hash_value_flag_tensors_;
Tensor2<uint32_t> hash_value_flag_sumed_tensors_;
Tensor2<uint32_t>
hash_value_index_count_offset_tensors_; /**< The temp memory to store the offset of each count
of hash table value indexes in update_params(). */
Tensor2<uint32_t> hash_value_index_count_counter_tensors_; /**< The temp memory to store the
counter of the count of hash table
value indexes in update_params(). */
SparseEmbeddingHashParams& param;
public:
OptimizerTensor<TypeEmbeddingComp> opt_tensors_;
EmbeddingOptimizer(size_t max_vocabulary_size_per_gpu_, SparseEmbeddingHashParams& param,
const std::shared_ptr<GeneralBuffer2<CudaAllocator>>& buf);
void initialize(const GPUResource& local_gpu);
void reset(GPUResource const& local_gpu) { initialize(local_gpu); }
void update(size_t batch_size, size_t slot_num, size_t embedding_vec_size,
size_t max_vocabulary_size_per_gpu, size_t nnz,
const Tensor2<TypeHashKey>& row_offset, Tensor2<size_t>& hash_value_index,
const Tensor2<TypeEmbeddingComp>& wgrad, Tensor2<float>& hash_table_value,
size_t sm_count, cudaStream_t stream);
};
6.2 更新
其內部主要是通過 opt_adagrad_kernel 進行更新。
template <typename TypeHashKey, typename TypeEmbeddingComp>
void EmbeddingOptimizer<TypeHashKey, TypeEmbeddingComp>::update(
size_t batch_size, size_t slot_num, size_t embedding_vec_size,
size_t max_vocabulary_size_per_gpu, size_t nnz, const Tensor2<TypeHashKey> &row_offset,
Tensor2<size_t> &hash_value_index, const Tensor2<TypeEmbeddingComp> &wgrad,
Tensor2<float> &hash_table_value, size_t sm_count, cudaStream_t stream) {
OptimizerTensor<TypeEmbeddingComp> &opt_tensor = opt_tensors_;
OptParams &opt_params = param.opt_params;
Tensor2<TypeHashKey> &sample_id = sample_id_tensors_;
Tensor2<TypeHashKey> &sample_id_sort = sample_id_sort_tensors_;
Tensor2<size_t> &hash_value_index_sort = hash_value_index_sort_tensors_;
Tensor2<uint32_t> &hash_value_index_count_offset = hash_value_index_count_offset_tensors_;
Tensor2<uint32_t> &new_hash_value_flag = new_hash_value_flag_tensors_;
Tensor2<uint32_t> &hash_value_flag_sumed = hash_value_flag_sumed_tensors_;
Tensor2<uint32_t> &hash_value_index_count_counter = hash_value_index_count_counter_tensors_;
Tensor2<void> &temp_storage_sort = temp_storage_sort_tensors_;
Tensor2<void> &temp_storage_scan = temp_storage_scan_tensors_;
size_t block_size, grid_size;
try {
// step1: expand sample IDs
block_size = 64;
grid_size = (batch_size * slot_num - 1) / block_size + 1;
sample_id_expand_kernel<<<grid_size, block_size, 0, stream>>>(
batch_size, slot_num, row_offset.get_ptr(), sample_id.get_ptr());
if (opt_params.optimizer == Optimizer_t::SGD &&
opt_params.hyperparams.sgd.atomic_update) { // for SGD, do atomic update
const size_t block_size = embedding_vec_size;
const size_t grid_size = min(max(1ul, nnz), sm_count * 32);
float lr_scale = opt_params.lr / opt_params.scaler;
opt_sgd_atomic_kernel<<<grid_size, block_size, 0, stream>>>(
nnz, embedding_vec_size, lr_scale, hash_value_index.get_ptr(), sample_id.get_ptr(),
wgrad.get_ptr(), hash_table_value.get_ptr());
} else {
// step3: sort by hash_value_index
int end_bit = static_cast<int>(log2(static_cast<float>(max_vocabulary_size_per_gpu))) + 1;
size_t temp_storage_sort_size = temp_storage_sort.get_size_in_bytes();
CK_CUDA_THROW_(cub::DeviceRadixSort::SortPairs(
temp_storage_sort.get_ptr(), temp_storage_sort_size, hash_value_index.get_ptr(),
hash_value_index_sort.get_ptr(), sample_id.get_ptr(), sample_id_sort.get_ptr(), nnz, 0,
end_bit, stream, false));
// step4: count the number for each unduplicated hash_value_index
CK_CUDA_THROW_(
cudaMemsetAsync(hash_value_index_count_counter.get_ptr(), 0, sizeof(uint32_t), stream));
constexpr size_t max_grid_size = 384;
block_size = 256;
grid_size = min(max_grid_size, (nnz - 1) / block_size + 1);
value_count_kernel_1<<<grid_size, block_size, 0, stream>>>(
nnz, hash_value_index_sort.get_ptr(), new_hash_value_flag.get_ptr());
// a pinned memroy
CK_CUDA_THROW_(cudaMemcpyAsync(&hash_hash_value_index_count_num,
hash_value_index_count_counter.get_ptr(), sizeof(uint32_t),
cudaMemcpyDeviceToHost, stream));
// step5: use optimizer method to compute deltaw and update the parameters
block_size = embedding_vec_size;
grid_size = max(1, hash_hash_value_index_count_num);
switch (opt_params.update_type) {
case Update_t::Global: {
switch (opt_params.optimizer) {
case Optimizer_t::Adam: {
}
case Optimizer_t::AdaGrad: {
opt_adagrad_kernel<<<grid_size, block_size, 0, stream>>>(
hash_hash_value_index_count_num, embedding_vec_size, opt_params.lr,
opt_params.hyperparams.adagrad, opt_tensor.opt_accm_tensors_.get_ptr(),
sample_id_sort.get_ptr(), hash_value_index_sort.get_ptr(),
hash_value_index_count_offset.get_ptr(), wgrad.get_ptr(),
hash_table_value.get_ptr(), opt_params.scaler);
break;
}
case Optimizer_t::MomentumSGD:
case Optimizer_t::Nesterov:
case Optimizer_t::SGD:
default:
CK_THROW_(Error_t::WrongInput, "Error: Invalid opitimizer type");
} // switch (optimizer)
break;
}
case Update_t::Local: {
switch (opt_params.optimizer) {
case Optimizer_t::Adam: {
}
case Optimizer_t::AdaGrad: {
opt_adagrad_kernel<<<grid_size, block_size, 0, stream>>>(
hash_hash_value_index_count_num, embedding_vec_size, opt_params.lr,
opt_params.hyperparams.adagrad, opt_tensor.opt_accm_tensors_.get_ptr(),
sample_id_sort.get_ptr(), hash_value_index_sort.get_ptr(),
hash_value_index_count_offset.get_ptr(), wgrad.get_ptr(),
hash_table_value.get_ptr(), opt_params.scaler);
break;
}
case Optimizer_t::MomentumSGD:
case Optimizer_t::Nesterov:
case Optimizer_t::SGD:
default:
CK_THROW_(Error_t::WrongInput, "Error: Invalid opitimizer type");
} // switch (optimizer)
break;
}
case Update_t::LazyGlobal: {
switch (opt_params.optimizer) {
case Optimizer_t::Adam: {
}
case Optimizer_t::AdaGrad:
case Optimizer_t::MomentumSGD:
case Optimizer_t::Nesterov:
case Optimizer_t::SGD: {
CK_THROW_(Error_t::WrongInput,
"Error: lazy global update is only implemented for Adam");
break;
}
default:
CK_THROW_(Error_t::WrongInput, "Error: Invalid opitimizer type");
}
break;
}
default:
CK_THROW_(Error_t::WrongInput, "Error: Invalid update type");
} // switch (update type)
}
#ifndef NDEBUG
cudaDeviceSynchronize();
CK_CUDA_THROW_(cudaGetLastError());
#endif
} catch (const std::runtime_error &rt_err) {
std::cerr << rt_err.what() << std::endl;
throw;
}
return;
}
其本質就是更新 hash_table_value，也就是嵌入層的權重。具體我們後文會結合反向傳播進行分析。
// Local update for the Adagrad optimizer: compute the gradients and update the accumulators and the
// weights
template <typename TypeKey, typename TypeEmbeddingComp>
__global__ void opt_adagrad_kernel(uint32_t hash_value_index_count_num, int embedding_vec_size,
float lr, const AdaGradParams adagrad,
TypeEmbeddingComp *accum_ptr, const TypeKey *sample_id,
const size_t *hash_value_index_sort,
const uint32_t *hash_value_index_count_offset,
const TypeEmbeddingComp *wgrad, float *hash_table_value,
float scaler) {
int bid = blockIdx.x;
int tid = threadIdx.x;
if (tid < embedding_vec_size && bid < hash_value_index_count_num) {
uint32_t offset = hash_value_index_count_offset[bid];
float gi = accumulate_gradients(embedding_vec_size, sample_id, hash_value_index_count_offset,
wgrad, scaler, offset, bid, tid);
size_t row_index = hash_value_index_sort[offset];
size_t feature_index = row_index * embedding_vec_size + tid;
float accum =
TypeConvertFunc<float, TypeEmbeddingComp>::convert(accum_ptr[feature_index]) + gi * gi;
accum_ptr[feature_index] = TypeConvertFunc<TypeEmbeddingComp, float>::convert(accum);
float weight_diff = -lr * gi / (sqrtf(accum) + adagrad.epsilon);
hash_table_value[feature_index] += weight_diff; // 更新權重
}
}
至此，Distributed hash表 基本概念介紹完成，我們接下來介紹前向傳播，敬請期待。
0xFF 參考
https://web.eecs.umich.edu/~justincj/teaching/eecs442/notes/linear-backprop.html
稀疏矩陣存儲格式總結+存儲效率對比:COO,CSR,DIA,ELL,HYB
求通俗講解下tensorflow的embedding_lookup接口的意思？
【技術乾貨】聊聊在大廠推薦場景中embedding都是怎麼做的
ctr預估算法對於序列特徵embedding可否做拼接，輸入MLP?與pooling
https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/operations.html