repro-bundle/v1/scripts/ste_sim_torch.py

"""Fast batched STE simulation in PyTorch.

Simulates one-pass SGD on the quantized linear-regression model
    y   = (1/sqrt(d)) x^T w* + xi
    yhat = (1/sqrt(d)) psi_w(w)^T psi_x(x)

with the straight-through estimator gradient
    g = (yhat - y) * psi_x(x) / sqrt(d) + (lambda / d) psi_w(w)
    w <- w - eta * g

All runs (seeds x configurations) are batched as a single (B, d) tensor.
Quantization uses torch.bucketize (single fused kernel) for speed on GPU.
"""
from __future__ import annotations

import math
import time
from dataclasses import dataclass

import numpy as np
import torch


@dataclass
class STEConfig:
    d: int
    eta: float
    lam: float
    rho: float = 1.0
    sigma2: float = 0.0
    n_steps: int = 1_000_000
    log_period: int = 100
    n_seeds: int = 5
    w_init_std: float = 1.0
    w_star_value: float = 1.0
    sigma2_x: float = 1.0
    kappa_x: float = 1.0
    device: str = "cuda"
    dtype: torch.dtype = torch.float32


def run_ste(qw_pack, qx_pack, cfg: STEConfig, seed_offset: int = 0):
    """Run batched STE simulation.

    qw_pack : (levels, theta, omega, Delta) for weight quantizer
    qx_pack : (levels, theta, omega, Delta) for input quantizer, or None for identity
    cfg     : STEConfig

    Returns (taus, metrics_dict, wall_seconds).
    """
    use_cuda = cfg.device == "cuda" and torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    B = cfg.n_seeds
    d = cfg.d

    qw_levels, qw_theta, qw_omega, qw_Delta = qw_pack
    qw_levels = torch.tensor(qw_levels, dtype=cfg.dtype, device=device)
    qw_theta = torch.tensor(qw_theta, dtype=cfg.dtype, device=device)

    use_id_x = qx_pack is None
    if not use_id_x:
        qx_levels, qx_theta, qx_omega, qx_Delta = qx_pack
        qx_levels = torch.tensor(qx_levels, dtype=cfg.dtype, device=device)
        qx_theta = torch.tensor(qx_theta, dtype=cfg.dtype, device=device)

    gen = torch.Generator(device=device).manual_seed(1234 + seed_offset)
    w = torch.randn(B, d, dtype=cfg.dtype, device=device, generator=gen) * cfg.w_init_std

    w_star = torch.full((d,), cfg.w_star_value, dtype=cfg.dtype, device=device)
    rho = cfg.rho

    n_logs = cfg.n_steps // cfg.log_period + 1
    m_log = np.empty((n_logs, B), dtype=np.float64)
    q_log = np.empty((n_logs, B), dtype=np.float64)
    m_psi_log = np.empty((n_logs, B), dtype=np.float64)
    q_psi_log = np.empty((n_logs, B), dtype=np.float64)
    eps_log = np.empty((n_logs, B), dtype=np.float64)

    inv_sqrt_d = 1.0 / math.sqrt(d)
    inv_d = 1.0 / d
    eta = cfg.eta
    lam = cfg.lam
    sigma2_x = cfg.sigma2_x
    kappa_x = cfg.kappa_x
    sigma2 = cfg.sigma2
    noise_std = math.sqrt(sigma2) if sigma2 > 0 else 0.0

    def log_state(idx, w):
        with torch.no_grad():
            psi_w = qw_levels[torch.bucketize(w, qw_theta)]  # weight always quantized
            m = (w * w_star).sum(dim=1) * inv_d
            q = (w * w).sum(dim=1) * inv_d
            m_psi = (psi_w * w_star).sum(dim=1) * inv_d
            q_psi = (psi_w * psi_w).sum(dim=1) * inv_d
            eps = sigma2 + rho + sigma2_x * q_psi - 2.0 * kappa_x * m_psi
        m_log[idx] = m.cpu().numpy()
        q_log[idx] = q.cpu().numpy()
        m_psi_log[idx] = m_psi.cpu().numpy()
        q_psi_log[idx] = q_psi.cpu().numpy()
        eps_log[idx] = eps.cpu().numpy()

    # initial log
    t0 = time.time()
    idx = 0
    log_state(idx, w)
    idx += 1
    torch.cuda.synchronize() if use_cuda else None
    t0 = time.time()

    log_period = cfg.log_period
    n_steps = cfg.n_steps
    # tight inner loop: do log_period steps between log calls
    for outer in range(n_steps // log_period):
        for _ in range(log_period):
            x = torch.randn(B, d, dtype=cfg.dtype, device=device, generator=gen)
            if use_id_x:
                psi_x = x
            else:
                psi_x = qx_levels[torch.bucketize(x, qx_theta)]
            psi_w = qw_levels[torch.bucketize(w, qw_theta)]
            y = (x * w_star).sum(dim=1) * inv_sqrt_d
            if noise_std > 0:
                y = y + torch.randn(B, dtype=cfg.dtype, device=device, generator=gen) * noise_std
            yhat = (psi_w * psi_x).sum(dim=1) * inv_sqrt_d
            r = yhat - y
            # STE gradient: g_i = r * psi_x_i / sqrt(d) + (lam/d) psi_w_i
            w = w - eta * (r.unsqueeze(1) * psi_x * inv_sqrt_d + lam * inv_d * psi_w)
        log_state(idx, w)
        idx += 1

    if use_cuda:
        torch.cuda.synchronize()
    wall = time.time() - t0

    taus = np.arange(n_logs) * (log_period / d)
    metrics = {
        "m": m_log[:idx], "q": q_log[:idx],
        "m_psi": m_psi_log[:idx], "q_psi": q_psi_log[:idx],
        "eps_g": eps_log[:idx],
    }
    return taus, metrics, wall