use Fused AdamW as default

tests/integration_tests.py

@@ -267,11 +267,12 @@ def build_test_list():
         OverrideDefinitions(
             [
                 [
-                    "--optimizer.name Adam --optimizer.fused",
-                    "--optimizer.name AdamW --optimizer.fused",
+                    "--optimizer.name AdamW --optimizer.implementation foreach",
                 ]
             ],
-            "Fused Optimizer Test",
+            "Foreach Optimizer Test",
+            "optimizer_foreach",
+            ngpu=2,
         ),
         OverrideDefinitions(
             [

torchtitan/components/optimizer.py

@@ -204,13 +204,19 @@ def build_optimizers(
         )
     name = job_config.optimizer.name
     lr = job_config.optimizer.lr
-    fused = job_config.optimizer.fused
+
+    optim_implementation = job_config.optimizer.implementation
+    assert optim_implementation in ["fused", "foreach", "for-loop"]
+
+    fused = optim_implementation == "fused"
+    foreach = optim_implementation == "foreach"
+
     optimizer_kwargs = {
         "lr": lr,
         "betas": (0.9, 0.95),
         "weight_decay": 0.1,
         "fused": fused,
-        "foreach": not fused,
+        "foreach": foreach,
     }
 
     return (

torchtitan/config_manager.py

@@ -218,9 +218,17 @@ def __init__(self):
             "--optimizer.lr", type=float, default=8e-4, help="Learning rate to use"
         )
         self.parser.add_argument(
-            "--optimizer.fused",
-            action="store_true",
-            help="Whether the fused implementation(CUDA only) is used.",
+            "--optimizer.implementation",
+            type=str,
+            default="fused",
+            choices=["for-loop", "foreach", "fused"],
+            help="""
+            Specify which optimizer implementation to use:
+            - 'fused': Use fused implementation (CUDA only) for best performance.
+            - 'foreach': Use some horizontal fusion of tensors for better performance.
+            - 'for-loop': Use the default implementation for the optimizer (slowest).
+            - more info: https://pytorch.org/docs/stable/optim.html
+            """,
         )
         self.parser.add_argument(
             "--optimizer.early_step_in_backward",