torchrec icon indicating copy to clipboard operation
torchrec copied to clipboard

add optimizer key in EmbeddingFusedOptimizer and KeyValueEmbeddingFusedOptimizer.

Open tristanhy opened this issue 1 year ago • 3 comments

Summary: Add field optimizer_key in Torchrec EmbeddingFusedOptimizer. During the initialization of embedding module BatchedFusedEmbeddingBag, pass the optimizer_key information from fused parameters when creating the EmbeddingFusedOptimizer.

In sparse arch, update the optimizer key for TBE when creating a combined optimizer for different TBE modules. The optimizer_key in the combined fused optimizer can be extracted for learning rate warm up processing.

Differential Revision: D61668809

tristanhy avatar Sep 30 '24 22:09 tristanhy