[FEA] dynamic embedding training cache

[shijieliu]

Is your feature request related to a problem? Please describe.
A clear and concise description of what the problem is. Ex. I wish I could use recsys-examples to do [...]

Describe the solution you'd like
A clear and concise description of what you want to happen.

Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.

Additional context
Add any other context, code examples, or references to existing implementations about the feature request here.

---

By submitting this issue, you agree to follow our code of conduct and our contributing guidelines.

Roadmaps

Milestone 0: Verify cache performance

• verify cache performance on pow-law input

Milestone 1: Refactor with cache & storage

• refactor with cache & storage
• unit tests

Milestone 2: Replace old version with Cache&Storage version

• optimizer test
• incremental dump
• integration testing(cover existed and external PS test)
• dump & load test
• prefetch pipeline in example

Milestone 3: Optimization Cache&Storage

• benchmark it and list possible optimization chance
• align with old version when no caching

[shijieliu]

@jiashuy update this issue

[jiashuy]

After implementation, we found two issues with training cache:

1. backward is too slow: solved by optimizing kernels and replacing with cub calls in HKV.
2. evict times > missed items when capacity:batch=16M:1M.
   It may be caused by evicting key_a which already existed in the table, and also existed in current batch.
   using workflow update->lock->insert_and_evict(missed)->unlock to replace insert_and_evict will solve it in theory(code).

Encounter new hang out and python abort issue in HKV's insert_and_evict_kernel after adjusting the input(only insert those missed), and I will solve it.

[jiashuy]

Hang out problem is fixed in HKV

Benchmark result of no cache version(adagrad optimizer) on H100 PCIe:

Benchmark result of cache version(adagrad optimizer) on H100 PCIe:

Cache version is a little slow, because it introduced more operation.
I'm trying to optimize it. For example, we can classify the keys into found and missed in forward, and reuse it in backward.

[shijieliu]

1. support for dump.
2. discussion about integrating into inference pipeline
3. perf optimization

[jiashuy]

Previous:

We update (embedding + optimizer states) in a temp buffer, then insert into HKV.
In this case, we fetch the embedding from cpu to gpu in forward, but have to load optimizer states from cpu in backward'update.
And the temp buffer brings one more value movement.

Now

Compare with torchrec's UVM_CACHING(STBE, static embedding table)
We insert missing keys into HKV in forward, so backward's update will happend on gpu only, and change the value in the table inplace.
We collect more cache metrics to help to analysis.

Next:

Sort host addresses before load embedding and optimizer states from cpu(this part is slow)
Seperate embedding initialization and host embedding loading(two standalone kernel).
Enlarge cpu table's capacity in full scale.

Potential：
Overlap loading cache missing values from cpu and storing cache evicted values

[jiashuy]

In 25.08, will release new interface based on Cache&Storage as as BatchedDynamicEmbeddingTablesV2.
In 25.09, will replace BatchedDynamicEmbeddingTables with BatchedDynamicEmbeddingTablesV2.