在这系列文章中,我们介绍了 HugeCTR,这是一个面向行业的推荐系统训练框架,针对具有模型并行嵌入和数据并行密集网络的大规模 CTR 模型进行了优化。
其中借鉴了HugeCTR源码阅读 这篇大作,特此感谢。
本系列其他文章如下:
[源码解析] NVIDIA HugeCTR,GPU 版本参数服务器 --(1)
[源码解析] NVIDIA HugeCTR,GPU版本参数服务器--- (2)
[源码解析] NVIDIA HugeCTR,GPU版本参数服务器---(3)
[源码解析] NVIDIA HugeCTR,GPU版本参数服务器--- (4)
[源码解析] NVIDIA HugeCTR,GPU版本参数服务器--- (5) 嵌入式hash表
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;};在 DistributedSlotSparseEmbeddingHash 之中,嵌入表中的一些插槽被分配给多个GPU,称为分布式插槽。例如,slot-0 被分配到GPU-0/GPU-1上,slot-1 被分配到GPU-0/GPU-1上。嵌入表被封装在哈希表中,或者说哈希表是嵌入表的前置条件。哈希表一些相关成员变量如下:
DistributedSlotSparseEmbeddingHash 类实现了嵌入层的训练过程所需的所有操作,包括前向传播和后向传播。前向传播对应于API forward()。反向传播分为两个阶段的API:backward()和update_params()。该类还提供将哈希表(包括哈希表键hash_table_key、hash_table_value_index和hash_table_value)从主机文件上载到GPU(load_parameters 方法)的操作,以及将哈希表从GPU下载到主机文件(dump_parameters方法)的操作。
我们先自行想想看如何实现这个嵌入层,这样会让我们更好的理清楚思路。
DistributedSlotSparseEmbeddingHash 的定义如下,主要变量/概念为:
CSR相关,可以结合CSR定义来印证。
输入/输出数据:
Hash相关:
中间数据:
反向传播:
这里有两点说明:
我们再从源码之中找出部分注释给大家看看几个变量之间的关系,其查找逻辑是从上到下。

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>;因为DistributedSlotSparseEmbeddingHash 用到了 using NvHashTable = HashTable<TypeHashKey, size_t>,所以我们先看看 HashTable。这部分对应的是上面总图第一步,就是如何从 hash table 之中拿到低维嵌入表的 index。在后文之中,我们用 HashTable/哈希表来指定 DistributedSlotSparseEmbeddingHash 内部使用的真正的哈希表。
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_;};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()) {}};为了更好的分析,在看 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。
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;};我们先看看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);}插入操作我们就看看之前的 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,容易和其他数据结构弄混。具体逻辑是:
所以,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)=3m_hf(667)=6m_hf(588)=53,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);}我们接下来看看如何构建 DistributedSlotSparseEmbeddingHash,代码之中需要留意的是:
具体就是分配内存,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;}我们要看看几个关键变量的内存配置。
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);}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。
前面提到了 DistributedSlotSparseEmbeddingHash 如下成员变量会保存一些嵌入表信息。
EmbeddingData<TypeHashKey, TypeEmbeddingComp> embedding_data_;我们来挖掘一下。
EmbeddingData 定义如下,这里有两套成员变量,Tensors2 和 SparseTensors。
train_value_tensors_,train_row_offsets_tensors_,train_nnz_array_ 都是Tensor2,是普通张量,而 train_keys_ 是 SparseTensors,可以一个变量就搞定前面所有概念。所以,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_;}这里有两套构建函数,可能维护者在从旧接口切换到新接口。结合前后文,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_ 可能(可能是因为有两种不同的构造方式,目前只是讨论其中一种)都会指向用户输入训练数据。

目前,我们只设置了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; }我们仔细看看 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 大部分成员变量都是可以被引用来修改的。
具体配置就是在 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之中的相应数据。

DistributedSlotSparseEmbeddingHash 内部也存在一些优化器。
std::vector<EmbeddingOptimizer<TypeHashKey, TypeEmbeddingComp>> embedding_optimizers_;我们接下来分析一下。
优化器定义如下:
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);};其内部主要是通过 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// weightstemplate <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表 基本概念介绍完成,我们接下来介绍前向传播,敬请期待。
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
无中生有:论推荐算法中的Embedding思想
tf.nn.embedding_lookup函数原理
求通俗讲解下tensorflow的embedding_lookup接口的意思?
【技术干货】聊聊在大厂推荐场景中embedding都是怎么做的
ctr预估算法对于序列特征embedding可否做拼接,输入MLP?与pooling
推荐系统中的深度匹配模型
土法炮制:Embedding 层是如何实现的?
不等距双杆模型_搜索中的深度匹配模型(下)
深度特征 快牛策略关于高低层特征融合
[深度学习] DeepFM 介绍与Pytorch代码解释
deepFM in pytorch
推荐算法之7——DeepFM模型
DeepFM 参数理解(二)
推荐系统遇上深度学习(三)--DeepFM模型理论和实践
[深度学习] DeepFM 介绍与Pytorch代码解释
https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/operations.html
带你认识大模型训练关键算法:分布式训练Allreduce算法