sgl-project--sglang/python/sglang/benchmark/dspark_sts_fit.py

from __future__ import annotations

import argparse
import glob
import logging
import math
from pathlib import Path
from typing import Optional

import torch

from sglang.srt.speculative.dspark_components.dspark_sts import (
    DSparkStsCalibration,
)

logger = logging.getLogger(__name__)

_EPS_PROB = 1e-8


def default_temperature_grid() -> torch.Tensor:
    return torch.logspace(math.log10(0.1), math.log10(10.0), steps=41)


def expected_calibration_error(
    *,
    probs: torch.Tensor,
    targets: torch.Tensor,
    num_bins: int,
) -> float:
    probs = probs.reshape(-1).to(torch.float64).clamp(_EPS_PROB, 1.0 - _EPS_PROB)
    targets = targets.reshape(-1).to(torch.float64)
    total = probs.numel()
    if total == 0:
        return float("nan")
    bin_index = (probs * num_bins).long().clamp_(0, num_bins - 1)
    count = torch.zeros(num_bins, dtype=torch.float64)
    pred_sum = torch.zeros(num_bins, dtype=torch.float64)
    target_sum = torch.zeros(num_bins, dtype=torch.float64)
    count.scatter_add_(0, bin_index, torch.ones_like(probs))
    pred_sum.scatter_add_(0, bin_index, probs)
    target_sum.scatter_add_(0, bin_index, targets)
    denom = count.clamp_min(1.0)
    bin_error = (pred_sum / denom - target_sum / denom).abs()
    return float((bin_error * count).sum().item() / total)


def fit_sts_temperatures(
    *,
    logits: torch.Tensor,
    prefix_mask: torch.Tensor,
    grid: torch.Tensor,
    num_bins: int = 15,
) -> dict[str, list[float]]:
    logits = logits.to(torch.float64)
    prefix_mask = prefix_mask.to(torch.float64)
    num_samples, gamma = logits.shape
    if num_samples == 0:
        raise ValueError("fit_sts_temperatures requires at least one sample.")
    grid_values = grid.to(torch.float64).tolist()

    temperatures: list[float] = []
    ece_before: list[float] = []
    ece_after: list[float] = []

    survival_at_one = torch.ones(num_samples, dtype=torch.float64)
    survival_fitted = torch.ones(num_samples, dtype=torch.float64)
    for position in range(gamma):
        position_logits = logits[:, position]
        position_target = prefix_mask[:, position]

        survival_at_one = survival_at_one * torch.sigmoid(position_logits)
        ece_before.append(
            expected_calibration_error(
                probs=survival_at_one,
                targets=position_target,
                num_bins=num_bins,
            )
        )

        best_temperature = grid_values[0]
        best_survival = survival_fitted * torch.sigmoid(
            position_logits / best_temperature
        )
        best_ece = expected_calibration_error(
            probs=best_survival, targets=position_target, num_bins=num_bins
        )
        for temperature in grid_values[1:]:
            candidate_survival = survival_fitted * torch.sigmoid(
                position_logits / temperature
            )
            candidate_ece = expected_calibration_error(
                probs=candidate_survival,
                targets=position_target,
                num_bins=num_bins,
            )
            if candidate_ece < best_ece:
                best_ece = candidate_ece
                best_temperature = temperature
                best_survival = candidate_survival

        temperatures.append(float(best_temperature))
        ece_after.append(float(best_ece))
        survival_fitted = best_survival

    return {
        "temperatures": temperatures,
        "ece_before": ece_before,
        "ece_after": ece_after,
    }


def load_collected_shards(*, data_glob: str) -> tuple[torch.Tensor, torch.Tensor]:
    shard_paths = sorted(glob.glob(data_glob))
    if not shard_paths:
        raise ValueError(f"No STS data shards matched {data_glob!r}.")

    logits_shards: list[torch.Tensor] = []
    prefix_mask_shards: list[torch.Tensor] = []
    shard_gamma: Optional[int] = None
    for shard_path in shard_paths:
        shard = torch.load(shard_path, map_location="cpu")
        shard_logits = shard["logits"]
        shard_prefix_mask = shard["prefix_mask"]
        if shard_logits.shape != shard_prefix_mask.shape:
            raise ValueError(
                f"Shard {shard_path!r} logits / prefix_mask shape mismatch: "
                f"{tuple(shard_logits.shape)} vs {tuple(shard_prefix_mask.shape)}."
            )
        if shard_gamma is None:
            shard_gamma = int(shard_logits.shape[1])
        elif int(shard_logits.shape[1]) != shard_gamma:
            raise ValueError(
                f"Shard {shard_path!r} gamma {int(shard_logits.shape[1])} disagrees "
                f"with earlier shards' gamma {shard_gamma}."
            )
        logits_shards.append(shard_logits)
        prefix_mask_shards.append(shard_prefix_mask)

    return torch.cat(logits_shards, dim=0), torch.cat(prefix_mask_shards, dim=0)


def fit(
    *,
    data_glob: str,
    out: Path,
    num_bins: int = 15,
    gamma: Optional[int] = None,
) -> None:
    logits, prefix_mask = load_collected_shards(data_glob=data_glob)
    resolved_gamma = int(logits.shape[1])
    if gamma is not None and gamma != resolved_gamma:
        raise ValueError(
            f"Collected shards have gamma={resolved_gamma} but --gamma={gamma}."
        )
    num_samples = int(logits.shape[0])

    result = fit_sts_temperatures(
        logits=logits,
        prefix_mask=prefix_mask,
        grid=default_temperature_grid(),
        num_bins=num_bins,
    )
    calibration = DSparkStsCalibration(
        temperatures=result["temperatures"],
        dataset=data_glob,
        num_samples=num_samples,
        ece_before=result["ece_before"],
        ece_after=result["ece_after"],
    )
    out.write_text(calibration.to_json(), encoding="utf-8")

    print(
        f"Fit STS temperatures over {num_samples} samples (gamma={resolved_gamma}) "
        f"-> {out}"
    )
    print("pos  temperature  ece_before  ece_after")
    for position in range(resolved_gamma):
        print(
            f"{position:>3}  {result['temperatures'][position]:>11.4f}  "
            f"{result['ece_before'][position]:>10.4f}  "
            f"{result['ece_after'][position]:>9.4f}"
        )


def main() -> None:
    logging.basicConfig(level=logging.INFO)
    parser = argparse.ArgumentParser(
        description="Fit DSpark Sequential Temperature Scaling (STS) calibration "
        "temperatures from collected confidence shards."
    )
    parser.add_argument(
        "--data-glob",
        required=True,
        help="Glob of collected .pt shards, each a dict with [n, gamma] "
        "'logits' and 'prefix_mask' tensors.",
    )
    parser.add_argument(
        "--out",
        required=True,
        type=Path,
        help="Output STS calibration JSON path.",
    )
    parser.add_argument(
        "--num-bins",
        type=int,
        default=15,
        help="Number of equal-width ECE bins.",
    )
    parser.add_argument(
        "--gamma",
        type=int,
        default=None,
        help="Optional gamma override to validate the shards against.",
    )
    args = parser.parse_args()

    fit(
        data_glob=args.data_glob,
        out=args.out,
        num_bins=args.num_bins,
        gamma=args.gamma,
    )


if __name__ == "__main__":
    main()