plan_vit: add the muP / scaling-study ViT as a torchtitan experiment

tests/unit_tests/test_optimizer_param_groups.py

@@ -263,6 +263,84 @@ def test_lr_override(self):
         # Default group
         self.assertEqual(groups[1]["lr"], 1e-3)
 
+    def test_base_lr_with_lr_mult(self):
+        """A base lr scales each group by its lr_mult."""
+        model = SimpleModel()
+        config = OptimizersContainer.Config(
+            implementation="for-loop",
+            lr=1e-3,
+            param_groups=[
+                ParamGroupConfig(
+                    pattern=r"embed_tokens\.",
+                    optimizer_name="AdamW",
+                    optimizer_kwargs={"weight_decay": 0.1},
+                    lr_mult=0.25,
+                ),
+                ParamGroupConfig(
+                    pattern=r".*",
+                    optimizer_name="AdamW",
+                    optimizer_kwargs={"weight_decay": 0.1},
+                    lr_mult=1.0,
+                ),
+            ],
+        )
+        opt = config.build(model_parts=[model]).optimizers[0]
+
+        self.assertAlmostEqual(opt.param_groups[0]["lr"], 1e-3 * 0.25)
+        self.assertAlmostEqual(opt.param_groups[1]["lr"], 1e-3 * 1.0)
+
+    def test_base_lr_default_lr_mult_unscaled(self):
+        """lr_mult defaults to 1.0 so a single base lr applies unscaled."""
+        model = SimpleModel()
+        config = OptimizersContainer.Config(
+            implementation="for-loop",
+            lr=5e-4,
+            param_groups=[
+                ParamGroupConfig(
+                    pattern=r".*",
+                    optimizer_name="AdamW",
+                    optimizer_kwargs={"weight_decay": 0.1},
+                ),
+            ],
+        )
+        opt = config.build(model_parts=[model]).optimizers[0]
+
+        self.assertAlmostEqual(opt.param_groups[0]["lr"], 5e-4)
+
+    def test_base_lr_rejects_group_lr(self):
+        """Setting lr in optimizer_kwargs while a base lr is set raises ValueError."""
+        model = SimpleModel()
+        config = OptimizersContainer.Config(
+            implementation="for-loop",
+            lr=1e-3,
+            param_groups=[
+                ParamGroupConfig(
+                    pattern=r".*",
+                    optimizer_name="AdamW",
+                    optimizer_kwargs={"lr": 1e-4, "weight_decay": 0.1},
+                ),
+            ],
+        )
+        with self.assertRaises(ValueError):
+            config.build(model_parts=[model])
+
+    def test_lr_mult_without_base_lr_raises(self):
+        """Setting lr_mult without a base lr raises ValueError."""
+        model = SimpleModel()
+        config = OptimizersContainer.Config(
+            implementation="for-loop",
+            param_groups=[
+                ParamGroupConfig(
+                    pattern=r".*",
+                    optimizer_name="AdamW",
+                    optimizer_kwargs={"lr": 1e-3, "weight_decay": 0.1},
+                    lr_mult=0.5,
+                ),
+            ],
+        )
+        with self.assertRaises(ValueError):
+            config.build(model_parts=[model])
+
     def test_first_match_wins(self):
         """When patterns overlap, the first match wins."""
         model = SimpleModel()

torchtitan/components/optimizer.py

@@ -65,7 +65,14 @@ class ParamGroupConfig:
 
     optimizer_kwargs: dict[str, Any] = field(default_factory=dict)
     """Keyword arguments passed to the optimizer constructor.
-    Must include all required kwargs (e.g. ``lr``). No implicit defaults."""
+    Must include all required kwargs (e.g. ``lr``) unless the container sets a
+    base ``lr`` (see ``OptimizersContainer.Config.lr``). No implicit defaults."""
+
+    lr_mult: float = 1.0
+    """Multiplier on the container's base ``lr`` for this group. Only applies when
+    ``OptimizersContainer.Config.lr`` is set; the group's learning rate is then
+    ``lr * lr_mult`` and ``lr`` must not also appear in ``optimizer_kwargs``.
+    Defaults to 1.0 (group uses the base lr unscaled)."""
 
 
 T = TypeVar("T", bound=Optimizer)
@@ -107,6 +114,14 @@ class Config(Configurable.Config):
         regex pattern and a self-contained optimizer setup.
         Patterns are checked in order; first match wins."""
 
+        lr: float | None = None
+        """Optional base learning rate shared by all param groups. When set, each
+        group's learning rate is ``lr * ParamGroupConfig.lr_mult`` and groups must
+        not set their own ``lr`` in ``optimizer_kwargs``. When ``None`` (default),
+        each group provides its own ``lr`` in ``optimizer_kwargs``. A single base
+        lr lets the native ``--optimizer.lr`` override (and lr sweeps) scale every
+        group at once while preserving per-group ratios (e.g. muP lr scaling)."""
+
         implementation: Literal[
             "for-loop", "foreach", "fused", "fused_opt_states_bf16"
         ] = "fused"
@@ -158,10 +173,13 @@ def _build_param_groups(
         model: nn.Module,
         param_group_configs: list[ParamGroupConfig],
         impl_kwargs: dict[str, Any],
+        base_lr: float | None = None,
     ) -> tuple[dict[str, list[dict[str, Any]]], dict[str, list[str]]]:
         """Build PyTorch param groups from model parameters, partitioned by optimizer.
 
         Each parameter is assigned to the first matching ParamGroupConfig pattern.
+        When ``base_lr`` is set, each group's learning rate is ``base_lr *
+        ParamGroupConfig.lr_mult`` and groups must not set their own ``lr``.
 
         Returns two dicts keyed by optimizer name and aligned by index: the param
         group dicts to pass to the optimizer constructor, and the regex pattern of
@@ -192,12 +210,27 @@ def _build_param_groups(
                     f"matched no parameters"
                 )
 
+            group_kwargs = {**impl_kwargs, **pg.optimizer_kwargs}
+            if base_lr is not None:
+                if "lr" in pg.optimizer_kwargs:
+                    raise ValueError(
+                        f"Optimizer param_groups pattern '{pg.pattern}' sets 'lr' in "
+                        f"optimizer_kwargs while the optimizer Config also sets a base "
+                        f"lr; use lr_mult to scale from the base lr instead"
+                    )
+                group_kwargs["lr"] = base_lr * pg.lr_mult
+            elif pg.lr_mult != 1.0:
+                raise ValueError(
+                    f"Optimizer param_groups pattern '{pg.pattern}' sets lr_mult "
+                    f"but the optimizer Config sets no base lr; lr_mult only "
+                    f"applies when the Config sets a base lr"
+                )
+
             groups[pg.optimizer_name].append(
                 {
                     "params": params,
                     "param_names": param_names,
-                    **impl_kwargs,
-                    **pg.optimizer_kwargs,
+                    **group_kwargs,
                 }
             )
             patterns[pg.optimizer_name].append(pg.pattern)
@@ -213,7 +246,7 @@ def __init__(self, config: Config, *, model_parts: list[nn.Module]) -> None:
 
         for part_idx, model in enumerate(self.model_parts):
             groups_by_opt_name, patterns_by_opt_name = self._build_param_groups(
-                model, param_group_configs, impl_kwargs
+                model, param_group_configs, impl_kwargs, base_lr=config.lr
             )
             for opt_name, opt_param_groups in groups_by_opt_name.items():
                 optimizer = self._resolve_optimizer_cls(opt_name)(opt_param_groups)