deeprec
/
deeprec
mirror of https://github.com/alibaba/deeprec.git
0
0
Fork 0
Code Issues Projects Releases Wiki Activity

[Embedding] Fix HBM_DRAM Restore core when id_num > cachesize 

Signed-off-by: RobertLou <2874395462@qq.com>
This commit is contained in:
RobertLou 2024-04-22 16:51:14 +08:00
parent 04413cf0ee
commit 4e7c352d9e
1 changed files with 11 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
tensorflow/core/framework/embedding/hbm_dram_storage.h
View File

@ -49,6 +49,7 @@ class HbmDramStorage : public MultiTierStorage<K, V> {
~HbmDramStorage() override {
MultiTierStorage<K, V>::DeleteFromEvictionManager();
//delete restore_cache_;
delete hbm_;
delete dram_;
delete dram_feat_desc_;
@ -227,7 +228,7 @@ class HbmDramStorage : public MultiTierStorage<K, V> {
}
void BatchEviction() override {
constexpr int EvictionSize = 10000;
constexpr int EvictionSize = 5000;
K evic_ids[EvictionSize];
if (!MultiTierStorage<K, V>::ready_eviction_) {
return;
@ -287,16 +288,18 @@ class HbmDramStorage : public MultiTierStorage<K, V> {
partition_id, partition_num,
is_incr, reset_version, reader);
restore_cache_.reset(CacheFactory::Create<K>(CacheStrategy::LFU, "ads"));
restorer.RestoreCkpt(emb_config, device);
int64 num_of_hbm_ids =
std::min(MultiTierStorage<K, V>::cache_capacity_,
(int64)MultiTierStorage<K, V>::cache_->size());
(int64)restore_cache_->size());
if (num_of_hbm_ids > 0) {
K* hbm_ids = new K[num_of_hbm_ids];
int64* hbm_freqs = new int64[num_of_hbm_ids];
int64* hbm_versions = nullptr;
MultiTierStorage<K, V>::cache_->get_cached_ids(hbm_ids, num_of_hbm_ids,
restore_cache_->get_cached_ids(hbm_ids, num_of_hbm_ids,
hbm_versions, hbm_freqs);
ImportToHbm(hbm_ids, num_of_hbm_ids, value_len, emb_config.emb_index);
MultiTierStorage<K, V>::cache_thread_pool_->Schedule(
@ -329,10 +332,10 @@ class HbmDramStorage : public MultiTierStorage<K, V> {
Status s = filter->Restore(key_num, bucket_num, partition_id,
partition_num, value_len, is_filter,
true/*to_dram*/, is_incr, restore_buff);
MultiTierStorage<K, V>::cache_->update((K*)restore_buff.key_buffer, key_num,
(int64*)restore_buff.version_buffer,
(int64*)restore_buff.freq_buffer);
restore_cache_->update((K*)restore_buff.key_buffer, key_num,
(int64*)restore_buff.version_buffer,
(int64*)restore_buff.freq_buffer);
return s;
}
@ -574,6 +577,7 @@ class HbmDramStorage : public MultiTierStorage<K, V> {
DramStorage<K, V>* dram_ = nullptr;
FeatureDescriptor<V>* hbm_feat_desc_ = nullptr;
FeatureDescriptor<V>* dram_feat_desc_ = nullptr;
std::unique_ptr<BatchCache<K>> restore_cache_ = nullptr;
Allocator* gpu_alloc_;
const int copyback_flag_offset_bits_ = 60;
};