torchrec icon indicating copy to clipboard operation
torchrec copied to clipboard

Swap CPU TBE lookup kernel with SSD TBE lookup kernel for tables being offloaded to SSD

Open faran928 opened this issue 6 months ago • 1 comments

Summary: Once, we TW shard the SSD tables here: D74444699 along with remaining tables, we traverse the model def and swap cpu tbe look up kernel for these ssd tables with the ssd look up kernel introduced here: D74536076 during TGIF publish

If and when SSD embedding db based look up kernel is open sourced we can directly map the ssd tbe kernel during sharding without any swap

Differential Revision: D74537413

faran928 avatar May 13 '25 01:05 faran928