def pre_reduce_metrics(self, optimizer_metrics: dict[str, torch.Tensor]):
        """Preprocess metrics to reduce across ranks correctly."""
        # Only L2 norm metric keys are present, can skip sorting at this stage
        for metric in optimizer_metrics:
            # L2 norms need to be squared, before they are reduced via summation
            optimizer_metrics[metric] = optimizer_metrics[metric]**2
        return optimizer_metrics
Preprocess metrics to reduce across ranks correctly.
llmfoundry/optim/adaptive_lion.py