Make betas and weight_decay Adam(W) hyperparameters configurable (#1282)

This change exposes the `betas` and `weight_decay` hyperparameters for
Adam(W) to the user by adding them to the `Optimizer` dataclass.

The default behavior will be unchanged since I used the previously
hard-coded values as the default.

Author: runame

torchtitan/components/optimizer.py

@@ -267,7 +267,10 @@ def build_optimizers(
         )
     name = job_config.optimizer.name
     lr = job_config.optimizer.lr
+    beta1 = job_config.optimizer.beta1
+    beta2 = job_config.optimizer.beta2
     eps = job_config.optimizer.eps
+    weight_decay = job_config.optimizer.weight_decay
 
     optim_implementation = job_config.optimizer.implementation
     assert optim_implementation in ["fused", "foreach", "for-loop"]
@@ -277,9 +280,9 @@ def build_optimizers(
 
     optimizer_kwargs = {
         "lr": lr,
+        "betas": (beta1, beta2),
         "eps": eps,
-        "betas": (0.9, 0.95),
-        "weight_decay": 0.1,
+        "weight_decay": weight_decay,
         "fused": fused,
         "foreach": foreach,
     }

torchtitan/config_manager.py

@@ -126,9 +126,16 @@ class Optimizer:
     lr: float = 8e-4
     """Learning rate to use"""
 
+    beta1: float = 0.9
+    beta2: float = 0.95
+    """Exponential moving average hyperparameters to use"""
+
     eps: float = 1e-8
     """Epsilon value to use"""
 
+    weight_decay: float = 0.1
+    """Weight decay to use"""
+
     implementation: Literal["for-loop", "foreach", "fused"] = "fused"
     """
     Specify which optimizer implementation to use: