code/flashsim/flash_adamw_sim.py

"""CPU simulation of FlashAdamW (Algorithm 4 of arXiv:2602.23349v2).

Independent reimplementation for mechanism-level tests on machines without
CUDA (the authors' library is Triton/CUDA-only). Semantics follow Alg. 4:
optimizer states live as (int8/uint8 q, fp16 scales); master weights live as
(bf16 theta', int8/int16 rho); each step dequantizes, applies the standard
AdamW update, then re-quantizes / re-splits.

`companded=False` switches BOTH moment tensors to plain linear group-32 absmax
quantization (no softsign, no sqrt) - the ablation of paper Figs. 4-5.
`quantize_states=False` keeps fp32 states (weight-splitting-only ablation).
`master_weight_bits` in {24, 32, None} as in the authors' API.
"""

from __future__ import annotations

import math

import torch

from . import quantization as Q
from . import splitting as S


class FlashAdamWSim(torch.optim.Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=0.0, companded=True, quantize_states=True,
                 master_weight_bits=24, simulate_split_on_fp32=False):
        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
        super().__init__(params, defaults)
        assert master_weight_bits in (24, 32, None)
        self.companded = companded
        self.quantize_states = quantize_states
        self.master_weight_bits = master_weight_bits
        self.rho_dtype = torch.int8 if master_weight_bits == 24 else torch.int16
        # CPU-experiment mode: params stay fp32 for compute speed/stability,
        # but after each update we replace them with the value they would have
        # under bf16+rho splitting (roundtrip error injection). This isolates
        # FlashOptim's precision semantics without bf16 CPU kernels.
        self.simulate_split_on_fp32 = simulate_split_on_fp32

    def _q_state(self, x: torch.Tensor, kind: str):
        if not self.quantize_states:
            return x.clone()
        if kind == "m":
            return Q.quantize(x, signed=True, sqrt=False, softsign=self.companded)
        # variance: companded -> sqrt+uint8; linear -> plain uint8 absmax
        return Q.quantize(x, signed=False, sqrt=self.companded, softsign=False)

    def _dq_state(self, st, kind: str, shape):
        if not self.quantize_states:
            return st.clone()
        q, s = st
        if kind == "m":
            return Q.dequantize(q, s, signed=True, sqrt=False, softsign=self.companded).reshape(shape)
        return Q.dequantize(q, s, signed=False, sqrt=self.companded, softsign=False).reshape(shape)

    @torch.no_grad()
    def step(self):
        for group in self.param_groups:
            beta1, beta2 = group["betas"]
            lr, eps, wd = group["lr"], group["eps"], group["weight_decay"]
            for p in group["params"]:
                if p.grad is None:
                    continue
                state = self.state[p]
                use_split = self.master_weight_bits is not None and p.dtype != torch.float32
                if len(state) == 0:
                    state["step"] = 0
                    z = torch.zeros(p.numel(), dtype=torch.float32)
                    state["m"] = self._q_state(z, "m")
                    state["v"] = self._q_state(z, "v")
                    if use_split:
                        # Alg. 4 line 4: split initial theta (params already low-precision)
                        _, state["rho"] = S.compress(p.detach().to(torch.float32),
                                                     p.dtype, self.rho_dtype)
                state["step"] += 1
                t = state["step"]
                g = p.grad.detach().to(torch.float32).reshape(-1)

                # Prologue: reconstruct
                m = self._dq_state(state["m"], "m", g.shape)
                v = self._dq_state(state["v"], "v", g.shape)
                if use_split:
                    theta = S.decompress(p.detach().reshape(-1), state["rho"].reshape(-1))
                else:
                    theta = p.detach().to(torch.float32).reshape(-1)

                # Standard AdamW update (decoupled weight decay, torch convention)
                if wd:
                    theta = theta * (1 - lr * wd)
                m = beta1 * m + (1 - beta1) * g
                v = beta2 * v + (1 - beta2) * g * g
                mhat = m / (1 - beta1 ** t)
                vhat = v / (1 - beta2 ** t)
                theta = theta - lr * mhat / (vhat.sqrt() + eps)

                # Epilogue: re-quantize / re-split
                state["m"] = self._q_state(m, "m")
                state["v"] = self._q_state(v, "v")
                if use_split:
                    theta_lp, rho = S.compress(theta, p.dtype, self.rho_dtype)
                    p.copy_(theta_lp.reshape(p.shape))
                    state["rho"] = rho
                else:
                    p.copy_(theta.reshape(p.shape).to(p.dtype))


def bytes_per_param(optimizer_kind: str, gradient_release: bool = False) -> dict:
    """Analytic bytes/param accounting (paper Table 1), incl. scale overhead."""
    scale = 2.0 / Q.GROUP  # fp16 scale per group of 32
    if optimizer_kind == "adamw_ref":
        parts = {"master": 4, "grad": 4, "m": 4, "v": 4}
    elif optimizer_kind == "flash_adamw":
        parts = {"master": 2, "ecc": 1, "grad": 0 if gradient_release else 2,
                 "m": 1 + scale, "v": 1 + scale}
        if gradient_release:
            parts["master"] = 2  # theta' always materialized
    elif optimizer_kind == "sgd_ref":
        parts = {"master": 4, "grad": 4, "m": 4}
    elif optimizer_kind == "flash_sgd":
        parts = {"master": 2, "ecc": 1, "grad": 0 if gradient_release else 2,
                 "m": 1 + scale}
    else:
        raise ValueError(optimizer_kind)
    parts["total"] = round(sum(parts.values()), 4)
    return parts