rustwx_stormscope/optimize.py

#!/usr/bin/env python
"""StormScope inference optimizations, validated on RTX 5090 (sm_120).

apply_opts(model, ...) mutates an earth2studio StormScope model in place:
  - sampler: deterministic EDM (S_churn=0) + reduced num_steps  [validated near-lossless vs
    real obs: det-32 corr 0.9591 vs stoch-100 0.9610; det ODE error 0.088% @32 / 0.13% @24]
  - torch.compile(default) on each diffusion expert network        [-26%, 0.1% numerical]
  - optional channels_last (marginal once compiled)

Defaults chosen as the validated-safe operating point. Set num_steps=24 for ~5.8x (still skillful),
or 18 for more speed with a small additional skill cost.
"""
import torch

# DiT linears safe to quantize: attention proj + MLP. NOTE the MLP is DiTBlock.linear (a timm Mlp)
# whose linears are `.layers.0` / `.layers.2` (NOT fc1/fc2). Keep adaLN modulation/embed/output bf16.
FP8_TARGETS = ["attention.qkv", "attention.proj", "linear.layers.0", "linear.layers.2"]


def _mxfp8_quantize(net):
    """Cast net to bf16 (torchao MX needs bf16 weights) and MXFP8-quantize the DiT attn/MLP linears.
    Gives ~1.29x on the network forward under torch.compile (fused quant), ~0.3% vs bf16. RTX 5090."""
    import torch.nn as nn
    from torchao.quantization import quantize_
    from torchao.prototype.mx_formats import MXDynamicActivationMXWeightConfig
    net = net.to(torch.bfloat16)

    def filt(mod, fqn):
        return isinstance(mod, nn.Linear) and any(s in fqn for s in FP8_TARGETS)
    quantize_(net, MXDynamicActivationMXWeightConfig(), filter_fn=filt)
    return net


def apply_opts(model, num_steps=32, deterministic=True, compile_net=True,
               channels_last=False, compile_mode="default", fp8=False, mxfp8=False, verbose=True):
    if num_steps:
        model.sampler_args["num_steps"] = int(num_steps)
    if deterministic:
        model.sampler_args["S_churn"] = 0.0
    nets = model.stage_models
    n_fp8 = 0
    for i in range(len(nets)):
        net = nets[i]
        if mxfp8:
            net = _mxfp8_quantize(net)
        elif fp8:
            from .fp8_linear import swap_linears
            n_fp8 += swap_linears(net, FP8_TARGETS)
        if channels_last:
            net = net.to(memory_format=torch.channels_last)
        if compile_net:
            net = torch.compile(net, mode=compile_mode, dynamic=False)
        nets[i] = net
    if verbose:
        print(f"[ss_optimize] sampler_args={model.sampler_args} compile={compile_net}"
              f" mxfp8={mxfp8} channels_last={channels_last} fp8_linears={n_fp8}", flush=True)
    return model