cosmos-framework/cosmos_framework/callbacks/sigma_loss_analysis.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

from dataclasses import dataclass
from typing import Optional

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.distributed as dist
import wandb

from cosmos_framework.model._base import ImaginaireModel
from cosmos_framework.utils import distributed, misc
from cosmos_framework.utils.callback import Callback
from cosmos_framework.utils.easy_io import easy_io


def _get_quantile_bins(n=10) -> np.ndarray:
    """Get predefined bins based on logarithmically spaced values"""
    points = torch.linspace(0, 1, n + 1)
    return points.numpy()


@dataclass
class _SigmaLossCache:
    """A fixed-size queue for caching sigma and loss tensors.

    Stores sigma/loss pairs on CPU.
    When the total number of elements exceeds queue_size, the oldest entries
    are automatically removed to maintain the size limit.

    Args:
        queue_size: Maximum number of elements to store in the cache.
    """

    def __init__(self, queue_size: int = 2000):
        self.queue_size = queue_size
        self.reset()

    def reset(self):
        self.sigma_list: list[torch.Tensor] = []
        self.loss_list: list[torch.Tensor] = []
        self._total_elements: int = 0

    def add(self, sigma: torch.Tensor, loss: torch.Tensor):
        # Convert to bf16 and store on CPU
        sigma_cpu = sigma.detach().cpu().to(torch.bfloat16)
        loss_cpu = loss.detach().cpu().to(torch.bfloat16)

        self.sigma_list.append(sigma_cpu)
        self.loss_list.append(loss_cpu)
        self._total_elements += sigma_cpu.numel()

        # Remove oldest elements if queue exceeds max size
        while self._total_elements > self.queue_size and len(self.sigma_list) > 1:
            removed_sigma = self.sigma_list.pop(0)
            self.loss_list.pop(0)
            self._total_elements -= removed_sigma.numel()

    def get_arrays(self) -> tuple[torch.Tensor, torch.Tensor]:
        if not self.sigma_list:
            return torch.tensor([], dtype=torch.bfloat16), torch.tensor([], dtype=torch.bfloat16)

        sigma_arr = torch.cat(self.sigma_list, dim=0)  # [N_total]  (concatenated across cached batches)
        loss_arr = torch.cat(self.loss_list, dim=0)  # [N_total]

        return sigma_arr, loss_arr


class SigmaLossAnalysis(Callback):
    """Analyze the relationship between sigma (noise level) and flow matching loss.

    This callback tracks per-instance flow matching losses at different sigma values
    during training. It maintains separate caches for image and video batches,
    periodically aggregates statistics across all distributed ranks, and logs
    the results to wandb.

    The analysis helps understand how well the model learns to denoise at different
    noise levels, which is useful for diagnosing training dynamics in flow matching
    models.

    Args:
        every_n: Log statistics every N iterations.
        every_n_viz: Create visualization plots every N iterations (must be multiple of every_n).
        save_s3: If True, save raw data to S3 for offline analysis.
    """

    def __init__(
        self,
        every_n: int = 1,
        every_n_viz: int = 1,
        save_s3: bool = False,
    ) -> None:
        super().__init__()
        self.save_s3 = save_s3
        self.every_n = every_n
        assert every_n_viz % every_n == 0, "every_n_viz must be a multiple of every_n in sigma_loss_analysis callback"
        self.every_n_viz = every_n_viz
        self.name = self.__class__.__name__

        self.image_cache = _SigmaLossCache(queue_size=2000)
        self.video_cache = _SigmaLossCache(queue_size=2000)

    def _create_analysis_plots(
        self,
        sigma_arr: torch.Tensor,
        loss_arr: torch.Tensor,  # [N]  # [N]
    ) -> Optional[wandb.Image]:
        if len(sigma_arr) == 0:
            return None

        # Convert to numpy for plotting
        sigma_np = sigma_arr.cpu().float().numpy()
        loss_np = loss_arr.cpu().float().numpy()

        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

        # Get predefined bins based on logarithmically spaced values
        sigma_bins = _get_quantile_bins(10)

        # y_tick_min, y_tick_max = 0, 1.0
        y_tick_min, y_tick_max = 0, 1.0
        # 2D histogram with exponential sigma bins and fixed [0,1] loss range
        loss_bins = np.linspace(y_tick_min, y_tick_max, 20)

        counts, xedges, yedges = np.histogram2d(sigma_np, loss_np, bins=(sigma_bins, loss_bins))
        if counts.max() < 0.1:
            return None

        # Plot heatmap with exponential scale colormap
        im = ax1.imshow(
            counts.T,
            origin="lower",
            aspect="auto",
            extent=[sigma_bins[0], sigma_bins[-1], y_tick_min, y_tick_max],
            norm=matplotlib.colors.LogNorm(vmin=1, vmax=counts.max()),
        )
        plt.colorbar(im, ax=ax1)

        # Set fixed loss ticks from 0 to 1
        yticks = np.linspace(y_tick_min, y_tick_max, 6)
        ax1.set_yticks(yticks)
        ax1.set_yticklabels([f"{y:.1f}" for y in yticks])

        ax1.set_xlabel("Sigma")
        ax1.set_ylabel("Loss")
        title = "Sigma vs Loss Distribution"
        ax1.set_title(title)

        # Sigma histogram with loss statistics per bin
        hist_counts, _ = np.histogram(sigma_np, bins=sigma_bins)
        bin_indices = np.digitize(sigma_np, sigma_bins) - 1

        # Calculate statistics per bin
        n_bins = len(sigma_bins) - 1
        means = np.zeros(n_bins)
        stds = np.zeros(n_bins)
        for i in range(n_bins):
            bin_mask = bin_indices == i
            if bin_mask.any():
                means[i] = loss_np[bin_mask].mean()
                stds[i] = loss_np[bin_mask].std()
            else:
                means[i] = np.nan
                stds[i] = np.nan

        # Plot histogram
        bin_centers = (sigma_bins[:-1] + sigma_bins[1:]) / 2
        ax2.bar(bin_centers, hist_counts, width=np.diff(sigma_bins), alpha=0.3, align="center")

        # Plot loss statistics on twin axis
        ax2_twin = ax2.twinx()
        valid_mask = ~np.isnan(means)
        ax2_twin.errorbar(
            bin_centers[valid_mask], means[valid_mask], yerr=stds[valid_mask], color="red", fmt="o-", alpha=0.5
        )

        ax2.set_xlabel("Sigma (Log Scale)")
        ax2.set_ylabel("Count")
        ax2_twin.set_ylabel("Loss (mean ± std)")
        title = "Sigma Distribution with Loss Statistics"
        ax2.set_title(title)

        # Add grid for better readability
        ax1.grid(True, alpha=0.3)
        ax2.grid(True, alpha=0.3)

        # Create log-scale labels
        sigma_labels = [f"{val:.1e}" for val in sigma_bins]
        ax1.set_xticks(sigma_bins[1:-1])  # Skip boundary bins
        ax1.set_xticklabels(sigma_labels[1:-1], rotation=45)
        ax1.set_xscale("linear")
        ax2.set_xticks(sigma_bins[1:-1])
        ax2.set_xticklabels(sigma_labels[1:-1], rotation=45)
        ax2.set_xscale("linear")

        plt.tight_layout()
        fig_img = wandb.Image(fig)
        plt.close(fig)

        return fig_img

    def _process_stats(self, sigma: torch.Tensor, loss: torch.Tensor) -> dict:
        """Calculate summary statistics for sigma and loss distributions.

        Args:
            sigma: Tensor of sigma (noise level) values.
            loss: Tensor of corresponding loss values.

        Returns:
            Dictionary containing:
                - sigma_log_mean: Mean of log(sigma). Log-space is used since sigma spans
                    multiple orders of magnitude, a standard practice on flow matching / EDM models.
                - sigma_log_std: Standard deviation of log(sigma).
                - loss_mean: Average loss across all samples.
                - loss_std: Standard deviation of loss, measuring spread.
                - loss_min: Minimum loss value observed.
                - loss_max: Maximum loss value observed.
                - loss_median: Median (50th percentile) loss, robust to outliers.
                - loss_q1: First quartile (25th percentile) of loss.
                - loss_q3: Third quartile (75th percentile) of loss.
        """
        return {
            "sigma_log_mean": float(sigma.log().mean()),
            "sigma_log_std": float(sigma.log().std()),
            "loss_mean": float(loss.mean()),
            "loss_std": float(loss.std()),
            "loss_min": float(loss.min()),
            "loss_max": float(loss.max()),
            "loss_median": float(loss.median()),
            "loss_q1": float(torch.quantile(loss.float(), 0.25)),
            "loss_q3": float(torch.quantile(loss.float(), 0.75)),
        }

    def _gather_and_save(self, cache: _SigmaLossCache, iteration: int, prefix: str, log_viz: bool = True) -> dict:
        info = {}

        # Gather data from all ranks
        local_sigma, local_loss = cache.get_arrays()
        world_size = dist.get_world_size()

        if world_size > 1:
            # Gather sizes first
            local_size = torch.tensor([len(local_sigma)], dtype=torch.long, device="cuda")  # [1]
            sizes = [torch.zeros_like(local_size) for _ in range(world_size)]
            dist.all_gather(sizes, local_size)
            sizes = [s.item() for s in sizes]

            # Gather data
            max_size = max(sizes)
            if max_size > 0:
                # Move to GPU for gathering
                padded_sigma = torch.zeros(max_size, dtype=torch.bfloat16, device="cuda")  # [max_size]
                padded_loss = torch.zeros(max_size, dtype=torch.bfloat16, device="cuda")  # [max_size]

                if len(local_sigma) > 0:
                    padded_sigma[: len(local_sigma)] = local_sigma.cuda()
                    padded_loss[: len(local_loss)] = local_loss.cuda()

                all_sigma = [torch.zeros_like(padded_sigma) for _ in range(world_size)]
                all_loss = [torch.zeros_like(padded_loss) for _ in range(world_size)]

                dist.all_gather(all_sigma, padded_sigma)
                dist.all_gather(all_loss, padded_loss)

                if distributed.is_rank0():
                    # Combine data from all ranks
                    valid_sigma = []
                    valid_loss = []
                    for sigma, loss, size in zip(all_sigma, all_loss, sizes):
                        if size > 0:
                            valid_sigma.append(sigma[:size])
                            valid_loss.append(loss[:size])

                    if valid_sigma:
                        sigma_arr = torch.cat(valid_sigma)  # [N_total]  (across all ranks)
                        loss_arr = torch.cat(valid_loss)  # [N_total]

                        # Overall statistics
                        info[f"{prefix}/total_samples"] = sigma_arr.shape[0]

                        # Calculate statistics
                        stats = self._process_stats(sigma_arr, loss_arr)
                        info.update({f"{prefix}/{k}": v for k, v in stats.items()})

                        # Create visualization
                        if log_viz:
                            fig_img = self._create_analysis_plots(sigma_arr, loss_arr)
                            print(fig_img)
                            if fig_img is not None:
                                info[f"{prefix}/distribution_plot"] = fig_img

                        if self.save_s3:
                            save_data = {
                                "sigma": sigma_arr.cpu(),
                                "loss": loss_arr.cpu(),
                                "stats": {k: v for k, v in info.items() if not isinstance(v, wandb.Image)},
                            }
                            easy_io.dump(
                                save_data,
                                f"s3://rundir/{self.name}/{prefix}_Iter{iteration:09d}.pkl",
                            )

        cache.reset()
        return info

    def on_training_step_end(
        self,
        model: ImaginaireModel,
        data_batch: dict[str, torch.Tensor],
        output_batch: dict[str, torch.Tensor],
        loss: torch.Tensor,
        iteration: int = 0,
    ):
        sigma = output_batch["sigma"]
        fm_loss_vision_per_instance = output_batch["flow_matching_loss_vision_per_instance"]

        # sigma is [B] (base), [B,1] (TF), or [B,T_max] (DF); reduce to [B] for logging
        assert sigma.ndim <= 2, f"Sigma should be [B] or [B,T_max], got shape {sigma.shape}"
        if sigma.ndim == 2:
            sigma = sigma.mean(dim=-1)  # [B]  (reduced from [B,T_max] or [B,1])

        if model.is_image_batch(data_batch):
            self.image_cache.add(sigma, fm_loss_vision_per_instance)
        else:
            self.video_cache.add(sigma, fm_loss_vision_per_instance)

        if iteration % self.every_n == 0:
            info = {}

            with misc.timer("sigma_loss_analysis"):
                log_viz = iteration % self.every_n_viz == 0
                # Process image data
                if len(self.image_cache.sigma_list) > 0:
                    info.update(self._gather_and_save(self.image_cache, iteration, "sigma_loss_image", log_viz=log_viz))

                # Process video data
                if len(self.video_cache.sigma_list) > 0:
                    info.update(self._gather_and_save(self.video_cache, iteration, "sigma_loss_video", log_viz=log_viz))

                if distributed.is_rank0() and info and wandb.run:
                    wandb.log(info, step=iteration)