Unifying TBE API using List (Frontend) - reland

Summary:
Re-land of D68055168 with fixes.
------
Previous landing of D68055168 causes S498612 with error
```
RuntimeError: refined_slots[i]->isSubtypeOf(*attributes_[i].getType()) INTERNAL ASSERT FAILED at "fbcode/caffe2/aten/src/ATen/core/class_type.cpp":415, please report a bug to PyTorch.
```
Assertion failure point: https://fburl.com/code/q1xg0w8m 

**Root cause:**
This is a bug in PyTorch in which JIT type cannot infer the type correctly. This error occurs in the combination of the following:
- Only in JIT script 
- Only fails for module attributes

**Explanation**:
1. Module is scripted
e.g., [`module_factory/model_materializer_full_sync.py`](https://www.internalfb.com/code/fbsource/[ed9d26f3fda9f40054948e7c492bbbbf5dab987e]/fbcode/caffe2/torch/fb/module_factory/model_materializer_full_sync.py?lines=1445-1446%2C1622-1623), [module_factory/sync_sgd/train_module.py](https://www.internalfb.com/code/fbsource/[ed9d26f3fda9f40054948e7c492bbbbf5dab987e]/fbcode/caffe2/torch/fb/module_factory/sync_sgd/train_module.py?lines=1615)
2. JIT refines class type of `type __torch__.fbgemm_gpu.split_table_batched_embeddings_ops_training.SplitTableBatchedEmbeddingBagsCodegen` where module attributes are added to `refined_slots`
3. For each attribute, it asserts if the `refined_slot` attribute is a JIT subtype of the attribute ([here](https://www.internalfb.com/code/fbsource/[b5505c1f8e2a7945af8860fd29f89e220ffec919]/fbcode/caffe2/aten/src/ATen/core/class_type.cpp?lines=415) and [here](https://www.internalfb.com/code/fbsource/[d693ef764b7d8eb2e4205dd85a61a33d6fd7c977]/fbcode/caffe2/aten/src/ATen/core/jit_type_base.h?lines=396-409))
4. `SplitTableBatchedEmbeddingBagsCodegen` module has an attribute called `optimizer_args`. To avoid re-compilation in PT2, we need to change `learning_rate` to a Tensor (D65511904). We hence change the type of `optimizer_args` from `OptimizerArgs` (with `learning_rate` as float) to `OptimizerArgsPT2` (with `learning_rate_tensor`) as defined [here](https://www.internalfb.com/code/fbsource/[cec693b485d5fb800984943f4ce2bc6a1ca1c52a]/fbcode/deeplearning/fbgemm/fbgemm_gpu/codegen/training/python/lookup_args.template?lines=111).
5. JIT subtype sees the attribute as type `OptimizerArgsPT2` but `refined_slots` assumes the type as `Tuple` (which is the actual schema of `OptimizerArgsPT2`) and concludes that they are not the same.
 {F1976090817} 
See [log of refined_slot debug ](https://www.internalfb.com/phabricator/paste/view/P1756112218?lines=730-726)
6. Since JIT type sees that they are of different types, the assertion is triggered.

**Note**:
The JIT type is unable to figure out only when *Tensor* is added to the Class. 
- Previously prior to landing D68055168, `optimizer_args` was of type `OptimizerArgs` defined here. The class contains variables of types float, int and bool. JIT Type sees that they are both of type  `OptimizerArgs` and does not fail.
 {F1976090829} 
See [log of refined_slot debug](https://www.internalfb.com/intern/everpaste/?handle=GO2d1Bz063DvSSAEADUn74T4WyZPbsIXAAAB&phabricator_paste_number=1756112738)
- Adding any tensors to the class `OptimizerArgs` would cause errors as above. 
- Adding floats, ints or bools to the class `OptimizerArgs`is fine.
-------
__**Solution**__
We discussed several workarounds but the best solution is to 
- keep the `optimizer_args` as `OptimizerArgs` the same, i.e., `learning_rate` remains float. 
- create a `learning_rate_tensor` in invokers before passing it to `lookup_function`.

Moreover, as we just realize many users access `learning_rate` directly through `optimizer_args.learning_rate`, changes to `optimizer_args` may cause backward compatibility for their cases. 
This is also the best way to maintain backward compatibility. 

Note that, `learning_rate_tensor` is always created on CPU, so there should be no host-device synchronization.

Differential Revision: D71010630

Author: spcyppt

fbgemm_gpu/codegen/genscript/generate_backward_split.py

@@ -447,7 +447,7 @@ def generate() -> None:
                     ssd_optimizers.append(optim)
 
             BackwardSplitGenerator.generate_backward_split(
-                ssd_tensors=ssd_tensors, **optimizer
+                ssd_tensors=ssd_tensors, aux_args=aux_args, **optimizer
             )
         BackwardSplitGenerator.generate_rocm_backward_split()
 

fbgemm_gpu/codegen/training/backward/embedding_backward_split_template.cu

@@ -601,6 +601,7 @@ Tensor {{ embedding_cuda_op }}(
 
     {%- if "learning_rate" in args.split_kernel_arg_names %}
     // convert `learning rate` to float since `learning rate` is float in kernels
+    TORCH_CHECK(learning_rate_tensor.is_cpu(), "learning_rate_tensor tensor needs to be on CPU. Ensure learning_rate_tensor is on CPU or contact FBGEMM team if you get this error.")
     const float learning_rate = learning_rate_tensor.item<float>();
     {%- endif %}
 

fbgemm_gpu/codegen/training/python/lookup_args.template

@@ -50,7 +50,8 @@ class CommonArgs(NamedTuple):
     ssd_tensors: Dict[str, torch.Tensor]
     {%- endif %}
 
-
+# Do not add a parameter of Type tensor here. It will cause JIT script error due to a bug in PyTorch.
+# See more detail in D71010630.
 class OptimizerArgs(NamedTuple):
     stochastic_rounding: bool
     gradient_clipping: bool
@@ -108,36 +109,6 @@ class CommonArgsPT2(NamedTuple):
     ssd_tensors: Dict[str, torch.Tensor]
     {%- endif %}
 
-class OptimizerArgsPT2(NamedTuple):
-    """
-    Optimizer arguments for PT2 interface
-    """
-    stochastic_rounding: bool
-    gradient_clipping: bool
-    max_gradient: float
-    max_norm: float
-    learning_rate_tensor: torch.Tensor
-    eps: float
-    beta1: float
-    beta2: float
-    weight_decay: float
-    weight_decay_mode: int
-    eta: float
-    momentum: float
-    counter_halflife: int
-    adjustment_iter: int
-    adjustment_ub: float
-    learning_rate_mode: int
-    grad_sum_decay: int
-    tail_id_threshold: float
-    is_tail_id_thresh_ratio: int
-    total_hash_size: int  # Required for OptimType.NONE
-    weight_norm_coefficient: float
-    lower_bound: float
-    regularization_mode: int
-    use_rowwise_bias_correction: bool # Used for OptimType.ADAM
-
-
 class Momentum(NamedTuple):
     dev: torch.Tensor
     host: torch.Tensor