Add optimizer-owned gradient clipping extension point

tests/unit_tests/cpu/components/optim/test_optim.py

@@ -11,7 +11,7 @@
 import pytest
 import torch
 
-from torchtitan.components.optim import Optim
+from torchtitan.components.optim import GradientClipper, Optim
 
 
 @dataclass
@@ -25,6 +25,26 @@ class _OptimizersConfig:
     build: MagicMock
 
 
+class _TestGradientClipper(GradientClipper):
+    @dataclass(kw_only=True, slots=True)
+    class Config(GradientClipper.Config):
+        grad_norm: float = 2.0
+
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+        self.calls = []
+
+    def __call__(
+        self,
+        parameters: list[torch.nn.Parameter],
+        max_norm: float | int,
+        *,
+        parallelism_context,
+    ) -> torch.Tensor:
+        self.calls.append((parameters, max_norm, parallelism_context))
+        return torch.tensor(self.config.grad_norm)
+
+
 def test_optim_config_rejects_negative_max_norm() -> None:
     with pytest.raises(ValueError, match="max_norm"):
         Optim.Config(max_norm=-1)
@@ -48,6 +68,7 @@ def test_optim_builds_owned_components(enable_cuda_graph: bool) -> None:
         lr_scheduler=SimpleNamespace(build=MagicMock(return_value=lr_schedulers)),
         ema=SimpleNamespace(build=MagicMock(return_value=ema)),
         max_norm=1.0,
+        gradient_clipper=None,
         enable_cuda_graph=enable_cuda_graph,
     )
     wrapped_update = MagicMock()
@@ -104,6 +125,7 @@ def test_optim_cuda_graph_falls_back_for_unsupported_device(
         lr_scheduler=SimpleNamespace(build=MagicMock()),
         ema=None,
         max_norm=1.0,
+        gradient_clipper=None,
         enable_cuda_graph=True,
     )
 
@@ -128,6 +150,57 @@ def test_optim_cuda_graph_falls_back_for_unsupported_device(
     assert "Optim CUDA graph is disabled" in caplog.text
 
 
+def test_optim_builds_configured_grad_clip() -> None:
+    model = torch.nn.Linear(2, 2)
+    config = SimpleNamespace(
+        optimizer=_OptimizersConfig(
+            optimizers=[_OptimizerConfig()],
+            build=MagicMock(return_value=MagicMock()),
+        ),
+        lr_scheduler=SimpleNamespace(build=MagicMock()),
+        ema=None,
+        max_norm=1.0,
+        gradient_clipper=_TestGradientClipper.Config(grad_norm=3.0),
+        enable_cuda_graph=False,
+    )
+
+    optim = Optim(
+        config,
+        model_parts=[model],
+        parallelism_context=MagicMock(),
+        training_steps=10,
+        pp_has_last_stage=True,
+    )
+
+    assert isinstance(optim.gradient_clipper, _TestGradientClipper)
+    assert optim.gradient_clipper.config.grad_norm == 3.0
+
+
+def test_optim_update_uses_configured_grad_clip() -> None:
+    parameters = [torch.nn.Parameter(torch.tensor(1.0))]
+    parallelism_context = MagicMock()
+    parallelism_context.pp_enabled = False
+    parallelism_context.get_optional_mesh.return_value = None
+    gradient_clipper = _TestGradientClipper(_TestGradientClipper.Config())
+    optim = object.__new__(Optim)
+    optim.config = SimpleNamespace(max_norm=1.0)
+    optim.parallelism_context = parallelism_context
+    optim.pp_has_last_stage = True
+    optim.parameters = parameters
+    optim.optimizers = MagicMock()
+    optim.gradient_clipper = gradient_clipper
+
+    with patch(
+        "torchtitan.components.optim.optim.dist_utils.clip_grad_norm_"
+    ) as eager_grad_clip:
+        grad_norm = optim._update(torch.tensor(1.0))
+
+    eager_grad_clip.assert_not_called()
+    assert gradient_clipper.calls == [(parameters, 1.0, parallelism_context)]
+    torch.testing.assert_close(grad_norm, torch.tensor(2.0))
+    optim.optimizers.step.assert_called_once_with()
+
+
 @pytest.mark.parametrize("cuda_initialized", [False, True])
 def test_optim_step_keeps_eager_state_outside_update(cuda_initialized: bool) -> None:
     events = []

tests/unit_tests/cpu/components/test_invalid_loss.py

@@ -66,6 +66,7 @@ def _make_trainer(self, loss_value: float, should_log: bool) -> Trainer:
         optim.pp_has_last_stage = True
         optim.parameters = []
         optim.optimizers = MagicMock()
+        optim.gradient_clipper = None
         optim.lr_schedulers = MagicMock()
         optim.lr_schedulers.get_metrics.return_value = {}
         optim.ema = None

tests/unit_tests/cpu/test_trainer.py

@@ -779,6 +779,7 @@ def test_optim_update_clips_before_parameter_update() -> None:
             ),
             parameters=[],
             optimizers=optimizers,
+            gradient_clipper=None,
         ),
     )
 

torchtitan/components/optim/__init__.py

@@ -6,7 +6,7 @@
 
 from .ema import EMA
 from .lr_scheduler import LRSchedulersContainer
-from .optim import Optim
+from .optim import GradientClipper, Optim
 from .optimizer import Adam, AdamW, BaseOptimizer, DistMuon, OptimizersContainer
 
 __all__ = [
@@ -15,6 +15,7 @@
     "BaseOptimizer",
     "DistMuon",
     "EMA",
+    "GradientClipper",
     "LRSchedulersContainer",
     "OptimizersContainer",
     "Optim",

torchtitan/components/optim/optim.py

@@ -26,6 +26,30 @@
 logger = logging.getLogger(__name__)
 
 
+class GradientClipper(Configurable):
+    """Base class for optimizer-owned gradient clipping.
+
+    Implementations must define their own nested ``Config`` class.
+    """
+
+    @dataclass(kw_only=True, slots=True)
+    class Config(Configurable.Config):
+        pass
+
+    def __init__(self, config: Config) -> None:
+        self.config = config
+
+    def __call__(
+        self,
+        parameters: list[nn.Parameter],
+        max_norm: float | int,
+        *,
+        parallelism_context: ParallelismContext,
+    ) -> torch.Tensor:
+        """Clip gradients in place and return their total norm."""
+        raise NotImplementedError
+
+
 class Optim(Configurable):
     """Own the parameter update and its eager state."""
 
@@ -41,6 +65,8 @@ class Config(Configurable.Config):
         """Optional online EMA of model weights."""
         max_norm: float | int = 1.0
         """Maximum norm used for gradient clipping."""
+        gradient_clipper: GradientClipper.Config | None = None
+        """Optional optimizer-owned gradient clipping implementation."""
         enable_cuda_graph: bool = False
         """Capture gradient clipping and optimizer updates in one CUDA graph."""
 
@@ -51,6 +77,7 @@ def __post_init__(self) -> None:
     optimizers: OptimizersContainer
     lr_schedulers: LRSchedulersContainer
     ema: EMA | None
+    gradient_clipper: GradientClipper | None
 
     def __init__(
         self,
@@ -93,6 +120,11 @@ def __init__(
             if config.ema is not None
             else None
         )
+        self.gradient_clipper = (
+            config.gradient_clipper.build()
+            if config.gradient_clipper is not None
+            else None
+        )
 
         self._run_update = self._update
         if enable_cuda_graph:
@@ -120,14 +152,24 @@ def step(self, loss: torch.Tensor, *, current_step: int) -> torch.Tensor:
             self.ema.step(current_step)
         return grad_norm
 
-    def _update(self, loss: torch.Tensor) -> torch.Tensor:
-        """Clip gradients, check loss and gradient norm, then update parameters."""
-        grad_norm = dist_utils.clip_grad_norm_(
+    def _clip_grad_norm(self) -> torch.Tensor:
+        """Clip gradients and return their total norm."""
+        if self.gradient_clipper is None:
+            return dist_utils.clip_grad_norm_(
+                self.parameters,
+                self.config.max_norm,
+                foreach=True,
+                pp_mesh=self.parallelism_context.get_optional_mesh("pp"),
+            )
+        return self.gradient_clipper(
             self.parameters,
             self.config.max_norm,
-            foreach=True,
-            pp_mesh=self.parallelism_context.get_optional_mesh("pp"),
+            parallelism_context=self.parallelism_context,
         )
+
+    def _update(self, loss: torch.Tensor) -> torch.Tensor:
+        """Clip gradients, check loss and gradient norm, then update parameters."""
+        grad_norm = self._clip_grad_norm()
         loss_is_finite = torch.isfinite(loss).all().to(torch.int32)
         if not self.parallelism_context.pp_enabled or self.pp_has_last_stage:
             loss_mesh = self.parallelism_context.get_optional_mesh("loss")

torchtitan/experiments/torchft/optimizer.py

@@ -17,7 +17,6 @@
     get_flat_optim_state_dict,
     init_optim_state,
 )
-from torchtitan.distributed import utils as dist_utils
 
 if TYPE_CHECKING:
     from torchtitan.experiments.torchft.manager import TorchFTManager
@@ -126,13 +125,8 @@ def configure_fault_tolerance(self, ft_manager: "TorchFTManager") -> None:
     def _update(self, loss: torch.Tensor) -> torch.Tensor:
         if self._quorum_manager is None:
             return super()._update(loss)
-        # Duplicates Optim._update except for the finite check.
-        grad_norm = dist_utils.clip_grad_norm_(
-            self.parameters,
-            self.config.max_norm,
-            foreach=True,
-            pp_mesh=self.parallelism_context.get_optional_mesh("pp"),
-        )
+        # Matches Optim._update except for the finite check.
+        grad_norm = self._clip_grad_norm()
         loss_is_finite = torch.isfinite(loss).all().to(torch.int32)
         if not self.parallelism_context.pp_enabled or self.pp_has_last_stage:
             loss_mesh = self.parallelism_context.get_optional_mesh("loss")

torchtitan/experiments/torchft/tests/test_torchft_optimizer.py

@@ -164,6 +164,21 @@ def test_finite_step_does_not_report_error(self):
         optim._update(torch.tensor(1.0))
         manager.report_error.assert_not_called()
 
+    def test_finite_step_uses_configured_grad_clip(self):
+        optim, manager = self._optim()
+        gradient_clipper = Mock(return_value=torch.tensor(2.0))
+        optim.gradient_clipper = gradient_clipper
+
+        grad_norm = optim._update(torch.tensor(1.0))
+
+        gradient_clipper.assert_called_once_with(
+            optim.parameters,
+            optim.config.max_norm,
+            parallelism_context=optim.parallelism_context,
+        )
+        torch.testing.assert_close(grad_norm, torch.tensor(2.0))
+        manager.report_error.assert_not_called()
+
     def test_persistent_non_finite_step_raises(self):
         optim, manager = self._optim()
         limit = optim.MAX_CONSECUTIVE_NON_FINITE_STEPS