change config name

torchtitan/checkpoint.py

@@ -245,7 +245,7 @@ def __init__(
                         model_parts,
                         optimizers,
                     )
-                    if not job_config.training.enable_optimizer_in_backward
+                    if not job_config.optimizer.backward
                     else OptimizerInBackwardWrapper(
                         model_parts,
                         optimizers,

torchtitan/config_manager.py

@@ -187,6 +187,14 @@ def __init__(self):
             action="store_true",
             help="Whether the fused implementation(CUDA only) is used.",
         )
+        self.parser.add_argument(
+            "--optimizer.backward",
+            type=bool,
+            default=False,
+            help="""
+            Whether to apply optimizer in the backward. Caution, optimizer_in_backward
+            is not compatible with gradients clipping.""",
+        )
 
         # training configs
         self.parser.add_argument(
@@ -270,14 +278,6 @@ def __init__(self):
             action="store_true",
             help="Whether to apply loss parallel when sequence parallel is enabled",
         )
-        self.parser.add_argument(
-            "--training.enable_optimizer_in_backward",
-            type=bool,
-            default=False,
-            help="""
-            Whether to apply optimizer in the backward. Caution, optimizer_in_backward
-            cannot compile with gradients clipping.""",
-        )
         self.parser.add_argument(
             "--experimental.enable_async_tensor_parallel",
             default=False,

torchtitan/optimizer.py

@@ -16,7 +16,7 @@ def build_optimizers(model_parts, job_config: JobConfig):
     """Wrap one optimizer per model part in an OptimizersContainer which provides a single
     step() and zero_grad() method for all the child optimizers.
     """
-    optim_in_bwd = job_config.training.enable_optimizer_in_backward
+    optim_in_bwd = job_config.optimizer.backward
 
     def _build_optimizer(model):
         name = job_config.optimizer.name
@@ -135,7 +135,7 @@ def linear_warmup_linear_decay(
 
 
 def build_lr_schedulers(optimizers, job_config: JobConfig):
-    optim_in_bwd = job_config.training.enable_optimizer_in_backward
+    optim_in_bwd = job_config.optimizer.backward
 
     def _build_lr_scheduler(optimizer):
         """Build a linear warmup and linear decay scheduler"""