torchrec
torchrec copied to clipboard
Swap CPU TBE lookup kernel with SSD TBE lookup kernel for tables being offloaded to SSD
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