def pre_reduce_metrics(self, optimizer_metrics: dict[str, torch.Tensor]):
        """Preprocess metrics to reduce across ranks correctly."""
        # Sort L2 norms first so they are squared before other metrics, which depend on squared values
        metrics = optimizer_metrics.keys()
        metrics = sorted(
            metrics,
            key=lambda metric: 0 if 'l2_norm' in metric else 1,
        )
        for metric in metrics:
            if 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'):
                _, vectors, layer = tuple(metric.split('/'))

                A, B = tuple(vectors.split('_'))

                # L2 norm would've been squared in previous branch
                A_rank_subset_norm = math.sqrt(
                    optimizer_metrics[f'l2_norm/{A}/{layer}'],
                )
                B_rank_subset_norm = math.sqrt(
                    optimizer_metrics[f'l2_norm/{B}/{layer}'],
                )

                optimizer_metrics[metric
                                 ] *= A_rank_subset_norm * B_rank_subset_norm

        return optimizer_metrics
Preprocess metrics to reduce across ranks correctly.
llmfoundry/optim/adaptive_lion.py