From 9d2e71521eb3428b0cc5dd9d6645c3977cfb0c13 Mon Sep 17 00:00:00 2001 From: Jose Javier <26491792+josejg@users.noreply.github.com> Date: Mon, 21 Jul 2025 18:06:33 -0700 Subject: [PATCH 1/2] Log gradient clipping norm and frequency --- .../gradient_clipping/gradient_clipping.py | 43 ++++++++++++++++--- 1 file changed, 38 insertions(+), 5 deletions(-) diff --git a/composer/algorithms/gradient_clipping/gradient_clipping.py b/composer/algorithms/gradient_clipping/gradient_clipping.py index abd712b152..9554ff5ce1 100644 --- a/composer/algorithms/gradient_clipping/gradient_clipping.py +++ b/composer/algorithms/gradient_clipping/gradient_clipping.py @@ -25,7 +25,7 @@ def apply_gradient_clipping( clipping_type: str, clipping_threshold: float, fsdp_enabled: bool, -): +) -> Union[torch.Tensor, None]: """Clips all gradients in model based on specified clipping_type. Args: @@ -41,12 +41,16 @@ def apply_gradient_clipping( threshold by which if grad_norm / weight_norm is greater than this threshold then scale gradients by this threshold * (weight_norm / grad_norm) (for 'adaptive'). fsdp_enabled (bool): Bool of if the model is a FSDP model or not. + + Returns: + Union[torch.Tensor, None]: The total gradient norm before clipping for 'norm' clipping type, + None otherwise. """ if fsdp_enabled: for module in model.modules(): if isinstance(module, FullyShardedDataParallel) and module.check_is_root(): if clipping_type == 'norm': - module.clip_grad_norm_(max_norm=clipping_threshold) + return module.clip_grad_norm_(max_norm=clipping_threshold) elif clipping_type == 'value': module.clip_grad_norm_(max_norm=clipping_threshold, norm_type=float('inf')) else: @@ -56,11 +60,13 @@ def apply_gradient_clipping( if clipping_type == 'adaptive': _apply_agc(parameters, clipping_threshold=clipping_threshold) elif clipping_type == 'norm': - torch.nn.utils.clip_grad_norm_(parameters, max_norm=clipping_threshold) + return torch.nn.utils.clip_grad_norm_(parameters, max_norm=clipping_threshold) elif clipping_type == 'value': torch.nn.utils.clip_grad_value_(parameters, clip_value=clipping_threshold) else: raise ValueError(f"clipping_type must be 'adaptive', 'norm', or 'value' not {clipping_type} ") + + return None def _apply_agc( @@ -122,23 +128,50 @@ class GradientClipping(Algorithm): to (for 'value'), what values to clip the gradient norms to (for 'norm'), and threshold by which if grad_norm / weight_norm is greater than this threshold then scale gradients by this threshold * (weight_norm / grad_norm) (for 'adaptive'). + clipping_frequency_window (int, optional): Number of steps to use for calculating + the rolling average of clipping frequency. Only used for 'norm' clipping type. + Defaults to 100. """ - def __init__(self, clipping_type: str, clipping_threshold: float): + def __init__(self, clipping_type: str, clipping_threshold: float, clipping_frequency_window: int = 100): self.clipping_type = clipping_type self.clipping_threshold = clipping_threshold + self.clipping_frequency_window = clipping_frequency_window + self._clipping_history = [] def match(self, event: Event, state: State) -> bool: return event in [Event.INIT, Event.AFTER_TRAIN_BATCH] def apply(self, event: Event, state: State, logger: Logger) -> Optional[int]: if event == Event.AFTER_TRAIN_BATCH: - apply_gradient_clipping( + maybe_grad_norm = apply_gradient_clipping( model=state.model, clipping_type=self.clipping_type, clipping_threshold=self.clipping_threshold, fsdp_enabled=state.fsdp_config_version == 1, ) + + if self.clipping_type == 'norm': + if maybe_grad_norm is None: + raise RuntimeError("Expected gradient norm to be returned for 'norm' clipping type, but got None") + + grad_norm = maybe_grad_norm.item() + + # Log the gradient norm before clipping + logger.log_metrics({'gradient_norm/unclipped_magnitude': grad_norm}) + + # Log whether clipping was applied + clipping_applied = grad_norm > self.clipping_threshold + logger.log_metrics({'gradient_norm/clipped': float(clipping_applied)}) + + # Track clipping frequency + self._clipping_history.append(float(clipping_applied)) + # Keep only last N steps for frequency calculation + if len(self._clipping_history) > self.clipping_frequency_window: + self._clipping_history.pop(0) + + clipping_frequency = sum(self._clipping_history) / len(self._clipping_history) + logger.log_metrics({'gradient_norm/clipping_frequency': clipping_frequency}) def _get_clipped_gradient_coeff(weights: torch.Tensor, grad: torch.Tensor, clipping_threshold: float = 0.01): From 3a261808ebc01c27b4e903584796e19b66322b0b Mon Sep 17 00:00:00 2001 From: Jose Javier <26491792+josejg@users.noreply.github.com> Date: Mon, 28 Jul 2025 12:27:03 -0700 Subject: [PATCH 2/2] Auto-format files with yapf after pre-commit --- .../gradient_clipping/gradient_clipping.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/composer/algorithms/gradient_clipping/gradient_clipping.py b/composer/algorithms/gradient_clipping/gradient_clipping.py index 9554ff5ce1..73464d9c8e 100644 --- a/composer/algorithms/gradient_clipping/gradient_clipping.py +++ b/composer/algorithms/gradient_clipping/gradient_clipping.py @@ -41,7 +41,7 @@ def apply_gradient_clipping( threshold by which if grad_norm / weight_norm is greater than this threshold then scale gradients by this threshold * (weight_norm / grad_norm) (for 'adaptive'). fsdp_enabled (bool): Bool of if the model is a FSDP model or not. - + Returns: Union[torch.Tensor, None]: The total gradient norm before clipping for 'norm' clipping type, None otherwise. @@ -65,7 +65,7 @@ def apply_gradient_clipping( torch.nn.utils.clip_grad_value_(parameters, clip_value=clipping_threshold) else: raise ValueError(f"clipping_type must be 'adaptive', 'norm', or 'value' not {clipping_type} ") - + return None @@ -150,26 +150,26 @@ def apply(self, event: Event, state: State, logger: Logger) -> Optional[int]: clipping_threshold=self.clipping_threshold, fsdp_enabled=state.fsdp_config_version == 1, ) - + if self.clipping_type == 'norm': if maybe_grad_norm is None: raise RuntimeError("Expected gradient norm to be returned for 'norm' clipping type, but got None") - + grad_norm = maybe_grad_norm.item() - + # Log the gradient norm before clipping logger.log_metrics({'gradient_norm/unclipped_magnitude': grad_norm}) - + # Log whether clipping was applied clipping_applied = grad_norm > self.clipping_threshold logger.log_metrics({'gradient_norm/clipped': float(clipping_applied)}) - + # Track clipping frequency self._clipping_history.append(float(clipping_applied)) # Keep only last N steps for frequency calculation if len(self._clipping_history) > self.clipping_frequency_window: self._clipping_history.pop(0) - + clipping_frequency = sum(self._clipping_history) / len(self._clipping_history) logger.log_metrics({'gradient_norm/clipping_frequency': clipping_frequency})