torchrec
torchrec copied to clipboard
add optimizer key in EmbeddingFusedOptimizer and KeyValueEmbeddingFusedOptimizer.
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