Skip to content

Commit f372505

Browse files
EddyLXJfacebook-github-bot
EddyLXJ
authored and
facebook-github-bot
committed
Remove dry run for feature score eviction
Summary: Previously we are using dry run to scan all weight blocks in backend and assign each feature score into eviction bucket and calculate the eviction threshold. This diff is removing dry run process and put the process assign eviction bucket into update and eviction block. This can save half time about total eviction duration. Reviewed By: emlin Differential Revision: D80425794
1 parent 84d2a72 commit f372505

File tree

5 files changed

+316
-470
lines changed

5 files changed

+316
-470
lines changed

‎fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache.h

Lines changed: 87 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -384,6 +384,19 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
384384
local_write_allocate_total_duration +=
385385
facebook::WallClockUtil::NowInUsecFast() -
386386
before_alloc_ts;
387+
if (feature_evict_config_.has_value() &&
388+
feature_evict_config_.value()->trigger_mode_ !=
389+
EvictTriggerMode::DISABLED &&
390+
feature_evict_) {
391+
auto* feature_score_evict = dynamic_cast<
392+
FeatureScoreBasedEvict<weight_type>*>(
393+
feature_evict_.get());
394+
if (feature_score_evict) {
395+
feature_score_evict
396+
->update_feature_score_statistics(
397+
block, 0, shard_id, true);
398+
}
399+
}
387400
}
388401
if (feature_evict_config_.has_value() &&
389402
feature_evict_config_.value()->trigger_mode_ !=
@@ -705,16 +718,21 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
705718
auto it = wlmap->find(id);
706719
if (it != wlmap->end()) {
707720
block = it->second;
721+
feature_score_evict
722+
->update_feature_score_statistics(
723+
block, engege_rate, shard_id, false);
708724
} else {
709725
// Key doesn't exist, allocate new block and
710726
// insert.
711727
block = pool->template allocate_t<weight_type>();
712728
FixedBlockPool::set_key(block, id);
729+
FixedBlockPool::set_feature_score_rate(
730+
block, engege_rate);
713731
wlmap->insert({id, block});
732+
feature_score_evict
733+
->update_feature_score_statistics(
734+
block, 0, shard_id, true);
714735
}
715-
716-
feature_score_evict->update_feature_score_statistics(
717-
block, engege_rate);
718736
updated_id_count++;
719737
}
720738
}
@@ -1489,61 +1507,75 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
14891507
const auto indexes = iter->second;
14901508
auto f =
14911509
folly::via(executor_.get())
1492-
.thenValue(
1493-
[this, shard_id, indexes, &indices, &weights_with_metaheader](
1494-
folly::Unit) {
1495-
FBGEMM_DISPATCH_INTEGRAL_TYPES(
1496-
indices.scalar_type(),
1497-
"dram_kv_set_with_metaheader",
1498-
[this,
1499-
shard_id,
1500-
indexes,
1501-
&indices,
1502-
&weights_with_metaheader] {
1503-
using index_t = scalar_t;
1504-
CHECK(indices.is_contiguous());
1505-
CHECK(weights_with_metaheader.is_contiguous());
1506-
CHECK_EQ(
1507-
indices.size(0), weights_with_metaheader.size(0));
1508-
{
1509-
auto wlmap = kv_store_.by(shard_id).wlock();
1510-
auto* pool = kv_store_.pool_by(shard_id);
1511-
int64_t stride = weights_with_metaheader.size(1);
1512-
auto indices_data_ptr = indices.data_ptr<index_t>();
1513-
auto weights_data_ptr =
1514-
weights_with_metaheader.data_ptr<weight_type>();
1515-
for (auto index_iter = indexes.begin();
1516-
index_iter != indexes.end();
1517-
index_iter++) {
1518-
const auto& id_index = *index_iter;
1519-
auto id = int64_t(indices_data_ptr[id_index]);
1520-
// Defensive programming
1521-
// used is false shouldn't occur under normal
1522-
// circumstances
1523-
FixedBlockPool::set_used(
1524-
weights_data_ptr + id_index * stride, true);
1525-
1526-
// use mempool
1527-
weight_type* block = nullptr;
1528-
// First check if the key already exists
1529-
auto it = wlmap->find(id);
1530-
if (it != wlmap->end()) {
1531-
block = it->second;
1532-
} else {
1533-
// Key doesn't exist, allocate new block and
1534-
// insert.
1535-
block =
1536-
pool->template allocate_t<weight_type>();
1537-
wlmap->insert({id, block});
1510+
.thenValue([this,
1511+
shard_id,
1512+
indexes,
1513+
&indices,
1514+
&weights_with_metaheader](folly::Unit) {
1515+
FBGEMM_DISPATCH_INTEGRAL_TYPES(
1516+
indices.scalar_type(),
1517+
"dram_kv_set_with_metaheader",
1518+
[this,
1519+
shard_id,
1520+
indexes,
1521+
&indices,
1522+
&weights_with_metaheader] {
1523+
using index_t = scalar_t;
1524+
CHECK(indices.is_contiguous());
1525+
CHECK(weights_with_metaheader.is_contiguous());
1526+
CHECK_EQ(
1527+
indices.size(0), weights_with_metaheader.size(0));
1528+
{
1529+
auto wlmap = kv_store_.by(shard_id).wlock();
1530+
auto* pool = kv_store_.pool_by(shard_id);
1531+
int64_t stride = weights_with_metaheader.size(1);
1532+
auto indices_data_ptr = indices.data_ptr<index_t>();
1533+
auto weights_data_ptr =
1534+
weights_with_metaheader.data_ptr<weight_type>();
1535+
for (auto index_iter = indexes.begin();
1536+
index_iter != indexes.end();
1537+
index_iter++) {
1538+
const auto& id_index = *index_iter;
1539+
auto id = int64_t(indices_data_ptr[id_index]);
1540+
// Defensive programming
1541+
// used is false shouldn't occur under normal
1542+
// circumstances
1543+
FixedBlockPool::set_used(
1544+
weights_data_ptr + id_index * stride, true);
1545+
1546+
// use mempool
1547+
weight_type* block = nullptr;
1548+
// First check if the key already exists
1549+
auto it = wlmap->find(id);
1550+
if (it != wlmap->end()) {
1551+
block = it->second;
1552+
} else {
1553+
// Key doesn't exist, allocate new block and
1554+
// insert.
1555+
block = pool->template allocate_t<weight_type>();
1556+
wlmap->insert({id, block});
1557+
if (feature_evict_config_.has_value() &&
1558+
feature_evict_config_.value()->trigger_mode_ !=
1559+
EvictTriggerMode::DISABLED &&
1560+
feature_evict_) {
1561+
auto* feature_score_evict = dynamic_cast<
1562+
FeatureScoreBasedEvict<weight_type>*>(
1563+
feature_evict_.get());
1564+
if (feature_score_evict) {
1565+
feature_score_evict
1566+
->update_feature_score_statistics(
1567+
block, 0, shard_id, true);
15381568
}
1539-
std::copy(
1540-
weights_data_ptr + id_index * stride,
1541-
weights_data_ptr + (id_index + 1) * stride,
1542-
block);
15431569
}
15441570
}
1545-
});
1546-
});
1571+
std::copy(
1572+
weights_data_ptr + id_index * stride,
1573+
weights_data_ptr + (id_index + 1) * stride,
1574+
block);
1575+
}
1576+
}
1577+
});
1578+
});
15471579
futures.push_back(std::move(f));
15481580
}
15491581
return folly::collect(futures);

‎fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache_wrapper.h

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -148,8 +148,7 @@ class DramKVEmbeddingCacheWrapper : public torch::jit::CustomClassHolder {
148148
at::Tensor processed_counts,
149149
at::Tensor eviction_threshold_with_dry_run,
150150
at::Tensor full_duration_ms,
151-
at::Tensor exec_duration_ms,
152-
at::Tensor dry_run_exec_duration_ms) {
151+
at::Tensor exec_duration_ms) {
153152
auto metrics = impl_->get_feature_evict_metric();
154153
if (metrics.has_value()) {
155154
evicted_counts.copy_(
@@ -164,8 +163,6 @@ class DramKVEmbeddingCacheWrapper : public torch::jit::CustomClassHolder {
164163
metrics.value().full_duration_ms); // full duration (Long)
165164
exec_duration_ms.copy_(
166165
metrics.value().exec_duration_ms); // exec duration (Long)
167-
dry_run_exec_duration_ms.copy_(
168-
metrics.value().dry_run_exec_duration_ms); // dry run exec duration
169166
}
170167
}
171168

0 commit comments

Comments
 (0)