Replace LR access with wrapper (#2832)

Summary:
Pull Request resolved: #2832

X-link: pytorch/FBGEMM#3849

X-link: facebookresearch/FBGEMM#937

Currently, `learning_rate` has been accessed directly through `optimizer_args.learning_rate`. Hence, any changes to `learning_rate` will be affected.

This diff adds a wrapper function for accessing `learning_rate`.

**Usage**
```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(....)
lr = emb_op.get_learning_rate()
```

We plan to remove `learning_rate` from `optimizer_args` to avoid recompilation in PT2.
- PT2 adds a guard on the float inputs, and if the value is changed, it will be recompiled.
- Re-compilation is expensive
- Especially during warm up stage in e2e training, learning rate changes gradually for each iteration. If there is 10k warm up step, it recompiles 10k times.
- Hence, we cannot keep `learning_rate` as float

More context in D65511904.

Reviewed By: sryap, TroyGarden

Differential Revision: D71444136

fbshipit-source-id: e175eecac4c554c9e436b8dff8c9a7eb28b6578e

Author: spcyppt

torchrec/distributed/batched_embedding_kernel.py

@@ -190,7 +190,7 @@ def __init__(
         state: Dict[Any, Any] = {}
         param_group: Dict[str, Any] = {
             "params": [],
-            "lr": emb_module.optimizer_args.learning_rate,
+            "lr": emb_module.get_learning_rate(),
         }
 
         params: Dict[str, Union[torch.Tensor, ShardedTensor]] = {}
@@ -383,7 +383,7 @@ def get_optimizer_pointwise_shard_metadata_and_global_metadata(
         state: Dict[Any, Any] = {}
         param_group: Dict[str, Any] = {
             "params": [],
-            "lr": emb_module.optimizer_args.learning_rate,
+            "lr": emb_module.get_learning_rate(),
         }
 
         params: Dict[str, Union[torch.Tensor, ShardedTensor]] = {}

torchrec/modules/fused_embedding_modules.py

@@ -68,7 +68,7 @@ def __init__(  # noqa C901
         state: Dict[Any, Any] = {}
         param_group: Dict[str, Any] = {
             "params": [],
-            "lr": emb_module.optimizer_args.learning_rate,
+            "lr": emb_module.get_learning_rate(),
         }
 
         params: Dict[str, torch.Tensor] = {}