Merge branch 'main' into shanea/tokenizer-package-data

Author: 2015aroras

CHANGELOG.md

@@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added caching to disk of HF datasets used in downstream evals
 - Added FLOPs logging
 - Added configs for OLMo tiny set of models
+- Added configuration field `optimizer.record_update_metrics`, which defaults to `False`, but when set to True will trigger AdamW to collect the step size norm and absolute max for each parameter.
 - Added `olmo_data`, a package holding data files like tokenizers.
 - Added ability to load tokenizers from `olmo_data` package data.
 
@@ -24,6 +25,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added original legacy unsharding implementation back, as the default. The new
 shared memory implementation can be used by passing `use_legacy_shared_mem_impl` to `unshard.py`.
 - Refactor weight initialization. IMPORTANT: this does not maintain backwards-compatibility with older configs; the jobs will still run, but may produce different outputs.
+- Changed the behavior of the Lion optimizer to only record the update cosine similarity when `optimizer.record_update_metrics` is `True` in order to be consistent with the API.
 
 ### Fixed
 

olmo/config.py

@@ -489,6 +489,12 @@ class OptimizerConfig(BaseConfig):
     If not set, defaults to the wandb `log_interval`.
     """
 
+    record_update_metrics: bool = False
+    """
+    Whether to record detailed metrics about the optimizer's parameter updates, like the norm and max
+    of the update with AdamW.
+    """
+
     def __post_init__(self):
         self.betas = tuple(self.betas)  # type: ignore[assignment]
 

olmo/optim.py

@@ -35,6 +35,11 @@
 
 
 class Optimizer(OptimizerBase):
+    def __init__(self, *args, record_update_metrics: bool = False, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._record_update_metrics = record_update_metrics
+        self._collecting_metrics = False
+
     def _clean_param_name(self, name: str) -> str:
         return name.replace("_fsdp_wrapped_module.", "")
 
@@ -50,6 +55,7 @@ def clip_grads_and_collect_metrics(
         Clips gradients for every group that has the field `max_grad_norm`.
         At the same time collect metrics for each parameter and its gradient.
         """
+        self._collecting_metrics = collect_param_metrics
         device = get_default_device() if device is None else device
 
         # NOTE (epwalsh): during distributed training we're making an assumption that the order of
@@ -365,12 +371,13 @@ def __init__(
         lr: float = 1e-4,
         betas: Tuple[float, float] = (0.9, 0.99),
         weight_decay: float = 0.0,
+        record_update_metrics: bool = False,
         device: Optional[torch.device] = None,
     ):
         assert lr > 0.0
         assert all([0.0 <= beta <= 1.0 for beta in betas])
         defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
-        super().__init__(params, defaults)
+        super().__init__(params, defaults, record_update_metrics=record_update_metrics)
         for group in self.param_groups:
             group["initial_lr"] = group["lr"]
         self._update_total_dot_prod: Optional[torch.Tensor] = None
@@ -391,6 +398,10 @@ def get_post_step_metrics(
         if update_total_dot_prod is None or update_total_norm is None or signed_update_total_norm is None:
             return {}
 
+        self._update_total_dot_prod = None
+        self._update_total_norm = None
+        self._signed_update_total_norm = None
+
         if is_distributed() and isinstance(module, FullyShardedDataParallel):
             # Reduce total dot prod and norms across all ranks.
             update_total_norm = update_total_norm**2.0
@@ -419,9 +430,13 @@ def step(self, closure=None) -> None:
             with torch.enable_grad():
                 closure()
 
-        update_total_dot_prod = torch.tensor(0.0, dtype=torch.float32)
-        update_norms = []
-        signed_update_norms = []
+        update_total_dot_prod: Optional[torch.Tensor] = None
+        update_norms: Optional[List[torch.Tensor]] = None
+        signed_update_norms: Optional[List[torch.Tensor]] = None
+        if self._collecting_metrics and self._record_update_metrics:
+            update_total_dot_prod = torch.tensor(0.0, dtype=torch.float32)
+            update_norms = []
+            signed_update_norms = []
 
         for group in self.param_groups:
             for p in group["params"]:
@@ -452,31 +467,169 @@ def step(self, closure=None) -> None:
 
                 # Track dot product and norms of update vs signed update in order to calculate
                 # their cosine similarity.
-                update_total_dot_prod = update_total_dot_prod.to(update.device)
-                update_total_dot_prod += torch.tensordot(update, signed_update, dims=len(update.shape))
-                update_norms.append(torch.linalg.vector_norm(update, 2.0, dtype=torch.float32))
-                signed_update_norms.append(torch.linalg.vector_norm(signed_update, 2.0, dtype=torch.float32))
+                if (
+                    update_total_dot_prod is not None
+                    and update_norms is not None
+                    and signed_update_norms is not None
+                ):
+                    update_total_dot_prod = update_total_dot_prod.to(update.device)
+                    update_total_dot_prod += torch.tensordot(update, signed_update, dims=len(update.shape))
+                    update_norms.append(torch.linalg.vector_norm(update, 2.0, dtype=torch.float32))
+                    signed_update_norms.append(torch.linalg.vector_norm(signed_update, 2.0, dtype=torch.float32))
 
         # Compute cosine similarity between update and signed update.
-        self._update_total_dot_prod = update_total_dot_prod.to(
-            get_default_device() if self._device is None else self._device
-        )
-        self._update_total_norm = torch.linalg.vector_norm(
-            torch.stack(update_norms),
-            2.0,
-            dtype=torch.float32,
-        ).to(get_default_device() if self._device is None else self._device)
-        self._signed_update_total_norm = torch.linalg.vector_norm(
-            torch.stack(signed_update_norms),
-            2.0,
-            dtype=torch.float32,
-        ).to(get_default_device() if self._device is None else self._device)
+        if update_total_dot_prod is not None and update_norms is not None and signed_update_norms is not None:
+            device = get_default_device() if self._device is None else self._device
+            self._update_total_dot_prod = update_total_dot_prod.to(device)
+            self._update_total_norm = torch.linalg.vector_norm(
+                torch.stack(update_norms),
+                2.0,
+                dtype=torch.float32,
+            ).to(device)
+            self._signed_update_total_norm = torch.linalg.vector_norm(
+                torch.stack(signed_update_norms),
+                2.0,
+                dtype=torch.float32,
+            ).to(device)
 
 
 class AdamW(torch.optim.AdamW, Optimizer):
+    def __init__(self, *args, record_update_metrics: bool = False, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Need to set these here just like in our base `Optimizer` class since our `Optimizer.__init__`
+        # won't be called.
+        self._record_update_metrics = record_update_metrics
+        self._collecting_metrics = False
+
+        self._step_size_param_names: Optional[List[str]] = None
+        self._step_size_norms: Optional[List[torch.Tensor]] = None
+        self._step_size_maxs: Optional[List[torch.Tensor]] = None
+
+    @torch.no_grad()
+    def step(self, closure=None) -> None:
+        if not (self._record_update_metrics and self._collecting_metrics):
+            return super().step(closure=closure)
+
+        device = get_default_device()
+        param_names = []
+        step_size_norms = []
+        step_size_maxs = []
+        for group in self.param_groups:
+            beta1, beta2 = group["betas"]
+            lr = group["lr"]
+            weight_decay = group["weight_decay"]
+            eps = group["eps"]
+            amsgrad = group["amsgrad"]
+            for name, param in zip(group["param_names"], group["params"]):
+                name = self._clean_param_name(name)
+                param_names.append(name)
+                grad = param.grad
+                if grad is None:
+                    step_size_norms.append(torch.tensor([0.0], device=device))
+                    step_size_maxs.append(torch.tensor([0.0], device=device))
+                    continue
+
+                state = self.state[param]
+                # init state if needed
+                if len(state) == 0:
+                    state["step"] = (
+                        torch.zeros((), dtype=torch.float32, device=param.device)
+                        if group["capturable"] or group["fused"]
+                        else torch.tensor(0.0, dtype=torch.float32)
+                    )
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(param, memory_format=torch.preserve_format)
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(param, 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(param, memory_format=torch.preserve_format)
+
+                exp_avg = state["exp_avg"]
+                exp_avg_sq = state["exp_avg_sq"]
+                step_t = state["step"]
+
+                # Update step.
+                step_t += 1
+
+                # Perform step weight decay.
+                param.mul_(1 - lr * weight_decay)
+
+                # Decay the first and second moment running average coefficient.
+                exp_avg.lerp_(grad, 1 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+                step = step_t.item()
+
+                bias_correction1 = 1 - beta1**step
+                bias_correction2 = 1 - beta2**step
+
+                step_size = lr / bias_correction1
+
+                bias_correction2_sqrt = sqrt(bias_correction2)
+
+                if amsgrad:
+                    max_exp_avg_sq = state["max_exp_avg_sq"]
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.maximum(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = (max_exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)
+                else:
+                    denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)
+
+                update = -step_size * torch.div(exp_avg, denom)
+                param.add_(update)
+                step_size_norms.append(torch.linalg.vector_norm(update, 2.0, dtype=torch.float32).unsqueeze(0))
+                step_size_maxs.append(update.abs().max().unsqueeze(0))
+
+        self._step_size_param_names = param_names
+        self._step_size_norms = step_size_norms
+        self._step_size_maxs = step_size_maxs
+
     def get_state_for_param(self, param: nn.Parameter) -> Dict[str, Optional[torch.Tensor]]:
         return {key: self.state[param].get(key) for key in ("exp_avg", "exp_avg_sq")}  # type: ignore
 
+    def get_post_step_metrics(
+        self, module: nn.Module, process_group: Optional[dist.ProcessGroup] = None
+    ) -> Dict[str, torch.Tensor]:
+        if not (self._record_update_metrics and self._collecting_metrics):
+            return {}
+        else:
+            device = get_default_device()
+            dst_rank = 0
+            if process_group is not None:
+                dst_rank = dist.get_global_rank(process_group, 0)
+            param_names = self._step_size_param_names
+            step_size_norms = self._step_size_norms
+            step_size_maxs = self._step_size_maxs
+            assert param_names is not None
+            assert step_size_norms is not None
+            assert step_size_maxs is not None
+
+            # Reduce metrics if needed.
+            if is_distributed() and isinstance(module, FullyShardedDataParallel):
+                # Reduce norms.
+                all_norms = torch.cat(step_size_norms).to(device) ** 2.0
+                dist.reduce(all_norms, dst_rank, op=dist.ReduceOp.SUM, group=process_group)
+                step_size_norms = (all_norms ** (0.5)).squeeze(0).split(1)
+
+                # Reduce maxs.
+                all_maxs = torch.cat(step_size_maxs).to(device)
+                dist.reduce(all_maxs, dst_rank, op=dist.ReduceOp.MAX, group=process_group)
+                step_size_maxs = all_maxs.split(1)
+
+            metrics = {}
+            for param_name, step_size_norm, step_size_max in zip(param_names, step_size_norms, step_size_maxs):  # type: ignore[arg-type]
+                metrics[f"step/{param_name}.norm"] = step_size_norm.squeeze(0)
+                metrics[f"step/{param_name}.max"] = step_size_max.squeeze(0)
+
+            self._step_size_param_names = None
+            self._step_size_norms = None
+            self._step_size_maxs = None
+            return metrics
+
 
 @dataclass
 class Scheduler(metaclass=ABCMeta):
@@ -745,13 +898,15 @@ def build_optimizer(cfg: TrainConfig, model: nn.Module) -> Optimizer:
             lr=cfg.optimizer.learning_rate,
             betas=cfg.optimizer.betas,
             weight_decay=cfg.optimizer.weight_decay,
+            record_update_metrics=cfg.optimizer.record_update_metrics,
         )
     elif cfg.optimizer.name == OptimizerType.adamw:
         return AdamW(
             param_groups,
             lr=cfg.optimizer.learning_rate,
             betas=cfg.optimizer.betas,
             weight_decay=cfg.optimizer.weight_decay,
+            record_update_metrics=cfg.optimizer.record_update_metrics,
             eps=cfg.optimizer.eps,
         )
     else:

olmo/train.py

@@ -1245,6 +1245,7 @@ def on_trace_ready(p):
                     log.info("Training epoch complete")
                     self.epoch = epoch + 1
                     self.global_train_examples_seen_this_epoch = 0
+                    self.dataset.start_index = 0
                     if self.epoch < self.max_epochs:
                         self.dataset.reshuffle()
                     continue