add layerwise lr

Author: bmosaicml

composer/callbacks/__init__.py

@@ -16,6 +16,7 @@
 from composer.callbacks.optimizer_monitor import OptimizerMonitor
 from composer.callbacks.speed_monitor import SpeedMonitor
 from composer.callbacks.threshold_stopper import ThresholdStopper
+from composer.callbacks.loss_spike_intervention import LossSpikeIntervention
 
 __all__ = [
     'OptimizerMonitor',
@@ -27,5 +28,6 @@
     'EarlyStopper',
     'ExportForInferenceCallback',
     'ThresholdStopper',
+    'LossSpikeIntervention',
     'ImageVisualizer',
 ]

composer/callbacks/loss_spike_intervention.py

@@ -0,0 +1,192 @@
+# Copyright 2022 MosaicML Composer authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Monitor gradients during training."""
+
+import torch
+
+from composer.core import Callback, State
+from composer.loggers import Logger
+from composer.utils import dist
+import collections
+
+__all__ = ['LossSpikeIntervention']
+
+
+
+class MetricSpikeDetector:
+    
+    def __init__(self,
+                 window_moving_average=25,
+                 increase_factor=5,
+                 increase_lookback=500,
+                 plateau_min_duration=100,
+                 end_spike_factor=1.10):
+          
+        self.window_moving_average=window_moving_average
+        self.increase_factor=increase_factor
+        self.plateau_min_duration=plateau_min_duration
+        self.increase_lookback = increase_lookback
+        self.fast_moving_average = collections.deque(maxlen=window_moving_average)
+        self.intermediate_data_queue = collections.deque(maxlen=increase_lookback-window_moving_average)
+        self.slow_moving_average = collections.deque(maxlen=increase_lookback)
+        self.end_spike_factor = end_spike_factor
+        self.in_spike = False
+        self.mva_before_spike = None
+        self.spike_batch_idx_start = None
+
+
+    
+    def insert_observation(self, obs, batch_idx):
+        if len(self.fast_moving_average) >= self.fast_moving_average.maxlen:
+            # move the oldest obs out of the fast moving average into the
+            # intermediate data queue
+            fast_obs = self.fast_moving_average.popleft()
+            
+            if len(self.intermediate_data_queue) >= self.intermediate_data_queue.maxlen:
+                # move data from intermediate quque to slow MCVA queue
+                intermediate_obs = self.intermediate_data_queue.popleft()
+                self.slow_moving_average.append(intermediate_obs)
+
+            self.intermediate_data_queue.append(fast_obs)
+        
+        self.fast_moving_average.append(obs)
+        
+        fast_mva = sum(self.fast_moving_average) / len(self.fast_moving_average)
+        if not self.in_spike:
+            if len(self.slow_moving_average) > self.window_moving_average:
+                if self.mva_before_spike is None:
+                    slow_mva = sum(self.slow_moving_average) / len(self.slow_moving_average)
+                else:
+                    slow_mva = self.mva_before_spike
+                                    
+                
+                if fast_mva >= self.increase_factor * slow_mva:
+                    self.in_spike = True
+                    self.mva_before_spike = slow_mva
+                    self.spike_batch_idx_start = batch_idx
+        else:
+            if batch_idx - self.spike_batch_idx_start > self.plateau_min_duration:
+                # kill the layer!
+                return True
+            else:
+                if fast_mva <= self.mva_before_spike * self.end_spike_factor:
+                    self.in_spike = False
+                    self.spike_batch_idx_start = None
+        
+        return False
+  
+  
+
+class LossSpikeIntervention(Callback):
+    
+    def __init__(self,
+            metric = 'l2_norm/moment',
+            window_moving_average=25,
+            increase_factor=5,
+            increase_lookback=500,
+            plateau_min_duration=100,
+            end_spike_factor=1.10,
+            lr_scale=0.0
+    ):
+        self.metric = metric
+        self.lr_scale = lr_scale
+        self.window_moving_average = window_moving_average
+        self.increase_factor = increase_factor
+        self.increase_lookback = increase_lookback
+        self.plateau_min_duration = plateau_min_duration
+        self.end_spike_factor = end_spike_factor
+
+        self.metric_spike_detectors = {}
+        self.frozen_layers = set()
+        self.all_layers = set()
+
+    def fit_start(self, state: State, logger: Logger) -> None:
+        for name, p in state.model.named_parameters():
+            if p.requires_grad:
+                self.all_layers.add(name)
+                full_metric_name = f"{self.metric}/{name}"
+                self.metric_spike_detectors[full_metric_name] = MetricSpikeDetector(
+                    self.window_moving_average,
+                    self.increase_factor,
+                    self.increase_lookback,
+                    self.plateau_min_duration,
+                    self.end_spike_factor,
+                )
+
+    def batch_end(self, state: State, logger: Logger):
+        norm = 0.0
+        optimizer_metrics = {}
+
+        for name, p in state.model.named_parameters():
+            if p.grad is not None and p.requires_grad:
+
+                metric_reporter = getattr(state.optimizers[0], 'report_per_parameter_metrics', None)
+                if callable(metric_reporter):
+                    optimizer_metrics = metric_reporter(p, name, optimizer_metrics)
+
+                if f'l2_norm/grad/{name}' not in optimizer_metrics:
+                    param_grad_norm = torch.linalg.vector_norm(p.grad)
+                    optimizer_metrics[f'l2_norm/grad/{name}'] = param_grad_norm
+
+        if state.fsdp_enabled and dist.get_world_size() > 0 :
+            pre_reduce_metrics = getattr(state.optimizers[0], 'pre_reduce_metrics', None)
+            if callable(pre_reduce_metrics):
+                optimizer_metrics = pre_reduce_metrics(optimizer_metrics)
+
+            dist_reduce_metrics = getattr(state.optimizers[0], 'dist_reduce_metrics', None)
+            if callable(dist_reduce_metrics):
+                optimizer_metrics = dist_reduce_metrics(optimizer_metrics)
+
+        for metric in optimizer_metrics:
+            if metric.startswith('l2_norm/grad'):
+                norm += optimizer_metrics[metric]**2
+
+        optimizer_metrics['l2_norm/grad/global'] = norm**0.5
+
+        for metric in optimizer_metrics:
+            if isinstance(optimizer_metrics[metric], torch.Tensor):
+                optimizer_metrics[metric] = optimizer_metrics[metric].item()
+
+        batch_idx = state.timestamp.batch.value
+        newly_failed_layers = self.detect_failed_layers(optimizer_metrics, batch_idx)
+
+        if len(newly_failed_layers) > 0:
+            self.freeze_layers(newly_failed_layers, state)
+            for optimizer in state.optimizers:
+                for group in optimizer.param_groups:
+                    group['lr'] *= self.lr_scale
+            
+            for scheduler in state.schedulers:
+                scheduler.base_lrs = [self.lr_scale * lr for lr in scheduler.base_lrs]
+            
+
+        optimizer_metrics['num_frozen_layers'] = len(self.frozen_layers)
+        logger.log_metrics(optimizer_metrics)
+
+        if len(self.all_layers) == 0:
+            state.stop_training()
+
+    
+    def freeze_layers(self, newly_failed_layers, state):
+        for layer in newly_failed_layers:
+            self.all_layers.remove(layer)
+            if layer not in self.frozen_layers:
+                self.frozen_layers.add(layer)
+
+        for name, p in state.model.named_parameters():
+            if name in self.frozen_layers:
+                p.requires_grad = False
+
+
+    def detect_failed_layers(self, optimizer_metrics, batch_idx):
+        newly_failed = []
+        for logger_name, value in optimizer_metrics.items():
+            if logger_name.startswith(self.metric):
+                layer_name = logger_name.split('/')[-1]
+                if layer_name in self.frozen_layers:
+                    continue
+                if self.metric_spike_detectors[logger_name].insert_observation(value, batch_idx):
+                    newly_failed.append(layer_name)
+
+        return newly_failed

composer/optim/decoupled_weight_decay.py

@@ -232,12 +232,22 @@ def __init__(self,
         super().__init__(params=params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad)
         for group in self.param_groups:
             group['initial_lr'] = group['lr']
+        self.layer_to_scale = None
+           
+    def get_scaling(self, param):
+        if self.layer_to_scale:
+            if param not in self.layer_to_scale:
+                raise Exception(f"Couldn't find param: {param} in layer to scale: {self.layer_to_scale}")
+            else:
+                return self.layer_to_scale[param]
+        else:
+            return 1.0
 
     @staticmethod
     def adamw(params: List[torch.Tensor], grads: List[torch.Tensor], exp_avgs: List[torch.Tensor],
               exp_avg_sqs: List[torch.Tensor], max_exp_avg_sqs: List[torch.Tensor], state_steps: List[int], *,
-              amsgrad: bool, beta1: float, beta2: float, lr: float, initial_lr: float, weight_decay: float,
-              eps: float) -> None:
+              amsgrad: bool, beta1: float, beta2: float, lr: float, initial_lr: float, weight_decay: float, eps: float,
+              layerwise_lrs) -> None:
         r"""Functional API that performs AdamW algorithm computation with decoupled weight decay.
 
         Args:
@@ -280,10 +290,26 @@ def adamw(params: List[torch.Tensor], grads: List[torch.Tensor], exp_avgs: List[
             else:
                 denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
 
-            step_size = lr / bias_correction1
+            step_size = lr * layerwise_lrs[i] / bias_correction1
 
             param.addcdiv_(exp_avg, denom, value=-step_size)
 
+    def reset_state(self):
+        for group in self.param_groups:
+            amsgrad = group['amsgrad']
+            for p in group['params']:
+                if not p.requires_grad:
+                    continue
+                state = self.state[p]
+                state['step'] = 0
+                # Exponential moving average of gradient values
+                state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                # Exponential moving average of squared gradient values
+                state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                if amsgrad:
+                    # Maintains max of all exp. moving avg. of sq. grad. values
+                    state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+
     @torch.no_grad()
     def step(self, closure=None):
         """Performs a single optimization step.
@@ -304,6 +330,7 @@ def step(self, closure=None):
             exp_avg_sqs = []
             max_exp_avg_sqs = []
             state_steps = []
+            layerwise_lrs = []
             amsgrad = group['amsgrad']
             beta1, beta2 = group['betas']
             eps = group['eps']
@@ -312,7 +339,7 @@ def step(self, closure=None):
             weight_decay = group['weight_decay']
 
             for p in group['params']:
-                if p.grad is None:
+                if p.grad is None or not p.requires_grad:
                     continue
                 params_with_grad.append(p)
                 if p.grad.is_sparse:
@@ -322,7 +349,7 @@ def step(self, closure=None):
                 state = self.state[p]
 
                 # State initialization
-                if len(state) == 0:
+                if 'step' not in state:
                     state['step'] = 0
                     # Exponential moving average of gradient values
                     state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
@@ -334,7 +361,7 @@ def step(self, closure=None):
 
                 exp_avgs.append(state['exp_avg'])
                 exp_avg_sqs.append(state['exp_avg_sq'])
-
+                layerwise_lrs.append(self.get_scaling(p))
                 if amsgrad:
                     max_exp_avg_sqs.append(state['max_exp_avg_sq'])
 
@@ -355,13 +382,16 @@ def step(self, closure=None):
                        lr=lr,
                        initial_lr=initial_lr,
                        weight_decay=weight_decay,
-                       eps=eps)
+                       eps=eps,
+                       layerwise_lrs=layerwise_lrs)
 
         return loss
 
     def dist_reduce_metrics(self, optimizer_metrics):
         for metric in optimizer_metrics:
-            if metric.startswith('l2_norm'):
+            if metric.startswith('layerwise_lr_scaling'):
+                continue
+            elif metric.startswith('l2_norm'):
                 reduced = optimizer_metrics[metric]
                 if dist.get_world_size() > 1:
                     dist.all_reduce(reduced, reduce_operation='SUM')
@@ -385,15 +415,17 @@ def dist_reduce_metrics(self, optimizer_metrics):
                 if dist.get_world_size() > 1:
                     dist.all_reduce(reduced, reduce_operation='SUM')
                 optimizer_metrics[metric] = reduced / dist.get_world_size()
-        
+
         return optimizer_metrics
 
     def pre_reduce_metrics(self, optimizer_metrics):
         # some of the metrics need to be modified before being reduced in order for the
         # reduction to work properly
 
         for metric in optimizer_metrics:
-            if metric.startswith('l2_norm'):
+            if metric.startswith('layerwise_lr_scaling'):
+                continue
+            elif metric.startswith('l2_norm'):
                 # l2 norms need to be squared, before they are reduced via summation
                 optimizer_metrics[metric] = optimizer_metrics[metric]**2
             elif metric.startswith('cosine'):
@@ -418,16 +450,19 @@ def report_per_parameter_metrics(self, param: torch.Tensor, name: str, optimizer
         beta1, beta2 = self.param_groups[0]['betas']
         if param in self.state:
             param_optim_state = self.state[param]
+            local_lr = lr * self.get_scaling(param)
             step = param_optim_state['step']
             bias_correction1 = 1 - beta1**step
             bias_correction2 = 1 - beta2**step
             denom = (param_optim_state['exp_avg_sq'].sqrt() / math.sqrt(bias_correction2)).add_(eps)
-            step_size = lr / bias_correction1
+            step_size = local_lr / bias_correction1
             step_tensor = step_size * param_optim_state['exp_avg'].div(denom)
-            decay_factor = (lr / initial_lr) if initial_lr else 1.0
+            decay_factor = (local_lr / initial_lr) if initial_lr else 1.0
             step_tensor.add_(param, alpha=-weight_decay * decay_factor)
             for metric in self.metric_functions:
                 optimizer_metrics[f'{metric}/{name}'] = self.metric_functions[metric](param, param_optim_state,
                                                                                       step_tensor)
 
+            optimizer_metrics[f'layerwise_lr_scaling/{name}'] = self.get_scaling(param)
+
         return optimizer_metrics