From aec3a811967236c1b03411ba54dac24e46d709db Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Wed, 14 Oct 2026 17:35:24 +0000 Subject: [PATCH] Add optimizer-owned gradient clipping extension point --- .../cpu/components/optim/test_optim.py | 75 ++++++++++++++++++- .../cpu/components/test_invalid_loss.py | 1 + tests/unit_tests/cpu/test_trainer.py | 1 + torchtitan/components/optim/__init__.py | 3 +- torchtitan/components/optim/optim.py | 52 +++++++++++-- torchtitan/experiments/torchft/optimizer.py | 10 +-- .../torchft/tests/test_torchft_optimizer.py | 15 ++++ 7 files changed, 142 insertions(+), 15 deletions(-) diff --git a/tests/unit_tests/cpu/components/optim/test_optim.py b/tests/unit_tests/cpu/components/optim/test_optim.py index e0fb25a4afb..ea7ad8a8b6e 100644 --- a/tests/unit_tests/cpu/components/optim/test_optim.py +++ b/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 = [] diff --git a/tests/unit_tests/cpu/components/test_invalid_loss.py b/tests/unit_tests/cpu/components/test_invalid_loss.py index 927bf2f3a9f..d574669faf0 100644 --- a/tests/unit_tests/cpu/components/test_invalid_loss.py +++ b/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 diff --git a/tests/unit_tests/cpu/test_trainer.py b/tests/unit_tests/cpu/test_trainer.py index c79f1b0798f..e2ee08f7304 100644 --- a/tests/unit_tests/cpu/test_trainer.py +++ b/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, ), ) diff --git a/torchtitan/components/optim/__init__.py b/torchtitan/components/optim/__init__.py index f9adbd9d470..a21a6d37828 100644 --- a/torchtitan/components/optim/__init__.py +++ b/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", diff --git a/torchtitan/components/optim/optim.py b/torchtitan/components/optim/optim.py index 247ac3cafa4..d0a72e561c7 100644 --- a/torchtitan/components/optim/optim.py +++ b/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") diff --git a/torchtitan/experiments/torchft/optimizer.py b/torchtitan/experiments/torchft/optimizer.py index 72c2a0e4adc..2b87442fd3f 100644 --- a/torchtitan/experiments/torchft/optimizer.py +++ b/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") diff --git a/torchtitan/experiments/torchft/tests/test_torchft_optimizer.py b/torchtitan/experiments/torchft/tests/test_torchft_optimizer.py index 7155657178d..48dcd2d8fa4 100644 --- a/torchtitan/experiments/torchft/tests/test_torchft_optimizer.py +++ b/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