code/flashsim/splitting.py

"""Independent reimplementation of FlashOptim's ULP-based weight splitting.

Written from Algorithm 1 of arXiv:2602.23349v2 (FlashOptim) ONLY — no code was
copied from github.com/databricks/flashoptim. Used to independently check the
paper's mechanism (challenge Claims 1 and 6) and later cross-validated against
the authors' Triton kernels on a CUDA GPU.

Algorithm 1 (paper §3.1):
    C(theta):
        theta' = Downcast(theta)                     # round-to-nearest
        e      = theta - Float32(theta')
        l      = floor(log2(ULP(theta'))) - 1
        h      = floor(-l/2)                         # split scaling for stability
        e_norm = (e * 2^h) * 2^(-l-h)
        rho    = Int(Round(Clamp(e_norm, -1, 1) * N))   # N=127 int8, 32767 int16
    C^-1(theta', rho):
        l = floor(log2(ULP(theta'))) - 1
        h = floor(l/2)
        e = ((Float32(rho)/N) * 2^h) * 2^(l-h)
        return Float32(theta') + e
"""

from __future__ import annotations

import torch

_FORMAT = {
    torch.bfloat16: dict(mant=7, emin=-126, int_view=torch.int16, exp_mask=0x7F80, exp_shift=7),
    torch.float16: dict(mant=10, emin=-14, int_view=torch.int16, exp_mask=0x7C00, exp_shift=10),
}

N_FOR = {torch.int8: 127.0, torch.int16: 32767.0}
QMIN = {torch.int8: -127, torch.int16: -32767}


def ulp_log2(theta_lp: torch.Tensor) -> torch.Tensor:
    """log2(ULP(x)) for each element of a bf16/fp16 tensor, as int32.

    For normals with biased exponent field E: ULP = 2^(E - bias - mant).
    For subnormals/zero (E == 0):             ULP = 2^(emin - mant)  (min spacing).
    Infs/NaNs are not expected (finite training tensors).
    """
    fmt = _FORMAT[theta_lp.dtype]
    bits = theta_lp.view(fmt["int_view"]).to(torch.int32)
    e_field = (bits & fmt["exp_mask"]) >> fmt["exp_shift"]
    bias = -(fmt["emin"] - 1)  # bf16/fp32 bias 127, fp16 bias 15
    unbiased = torch.where(e_field > 0, e_field - bias, torch.full_like(e_field, fmt["emin"]))
    return unbiased - fmt["mant"]


def compress(theta: torch.Tensor, lp_dtype: torch.dtype = torch.bfloat16,
             rho_dtype: torch.dtype = torch.int8) -> tuple[torch.Tensor, torch.Tensor]:
    """Split fp32 theta into (theta' in lp_dtype, rho in rho_dtype). Alg. 1 C()."""
    assert theta.dtype == torch.float32
    n = N_FOR[rho_dtype]
    theta_lp = theta.to(lp_dtype)  # round-to-nearest-even downcast
    e = theta - theta_lp.to(torch.float32)  # exact (Sterbenz) for finite values
    l = ulp_log2(theta_lp) - 1
    h = torch.div(-l, 2, rounding_mode="floor")
    # e_norm = e * 2^h * 2^(-l-h), two steps to avoid overflow of 2^(-l)
    e_norm = torch.ldexp(torch.ldexp(e, h), -l - h)
    rho_f = torch.round(torch.clamp(e_norm, -1.0, 1.0) * n)
    rho = rho_f.to(rho_dtype)
    return theta_lp, rho


def decompress(theta_lp: torch.Tensor, rho: torch.Tensor) -> torch.Tensor:
    """Reconstruct fp32 theta_hat from (theta', rho). Alg. 1 C^-1()."""
    n = N_FOR[rho.dtype]
    l = ulp_log2(theta_lp) - 1
    h = torch.div(l, 2, rounding_mode="floor")
    e = torch.ldexp(torch.ldexp(rho.to(torch.float32) / n, h), l - h)
    return theta_lp.to(torch.float32) + e


def roundtrip(theta: torch.Tensor, lp_dtype: torch.dtype = torch.bfloat16,
              rho_dtype: torch.dtype = torch.int8) -> torch.Tensor:
    return decompress(*compress(theta, lp_dtype, rho_dtype))


def roundtrip_float_error(theta: torch.Tensor, lp_dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
    """Baseline from prior work (Zamirai et al.): store e = theta - theta' in the
    SAME low-precision float format (e.g. BF16+BF16). Paper Fig. 3 baseline."""
    theta_lp = theta.to(lp_dtype)
    e_lp = (theta - theta_lp.to(torch.float32)).to(lp_dtype)
    return theta_lp.to(torch.float32) + e_lp.to(torch.float32)


def roundtrip_no_correction(theta: torch.Tensor, lp_dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
    """Baseline: plain downcast, no error correction."""
    return theta.to(lp_dtype).to(torch.float32)