reencode/mxfp8_reencode.py

"""MXFP8 re-encoding of the trial-117 bake's BF16 attention output projections (qwen50, 2026-10-21).

Format matches the NVIDIA base for these layers: E4M3 codes [N, K] plus one E8M0 scale byte per
32 input elements, exponent = ceil(log2(amax / 448)) (the base's rule for ~95% of blocks; the rest
differ only because the base derived amax before quantization).

Constrained rounding: every element takes one of its two nearest E4M3 neighbours on the unchanged
block scale (never more than one step from the BF16 value). Starting from round-to-nearest, a
greedy pass flips elements to cancel, per input column, the error along the layer's removed output
direction r (the left singular vector of bake - base, which is rank-1). This keeps r^T Q(W') close
to r^T W', so the abliteration survives the re-encoding.

Usage (in the serving image, CPU):
  python3 mxfp8_reencode.py report <layers...>        # fidelity for a few layers, no files written
  python3 mxfp8_reencode.py build /out                 # write rewritten shards + metadata to /out;
                                                        # unchanged shards are listed for host-side hard links
Mounts: /base (NVIDIA FP4-Local), /bake (trial-117 bake), /out (new checkpoint dir, empty).
"""
import json
import os
import re
import shutil
import sys
import time

import torch
from safetensors import safe_open
from safetensors.torch import save_file

BASE, BAKE = '/base', '/bake'
IB = json.load(open(f'{BASE}/model.safetensors.index.json'))['weight_map']
IK = json.load(open(f'{BAKE}/model.safetensors.index.json'))
WK = IK['weight_map']
TARGET = re.compile(r'^model\.language_model\.layers\.(\d+)\.(linear_attn\.out_proj|self_attn\.o_proj)$')
ITERS = 1024


def tensor(root, idx, name):
    with safe_open(f'{root}/{idx[name]}', 'pt', device='cpu') as f:
        return f.get_tensor(name)


def base_dequant(n):
    w = tensor(BASE, IB, n + '.weight').float()
    e = tensor(BASE, IB, n + '.weight_scale').to(torch.int32) - 127
    return (w.view(w.shape[0], -1, 32) * torch.pow(2.0, e.float())[..., None]).view(w.shape)


def removed_direction(n, wp):
    d = wp - base_dequant(n)
    u, s, _ = torch.linalg.svd(d, full_matrices=False)
    return u[:, 0], (s[1] / s[0]).item() if s.numel() > 1 else 0.0


def encode(wp, r):
    """Return (codes fp8 [N,K], scale uint8 [N,K/32], dequantized float [N,K]) with constrained rounding."""
    n_rows, k = wp.shape
    blocks = wp.view(n_rows, -1, 32)
    amax = blocks.abs().amax(-1, keepdim=True).clamp_min(2.0 ** -126)
    exp = torch.ceil(torch.log2(amax / 448.0)).clamp(-127, 127)
    scale = torch.pow(2.0, exp)
    x = (blocks / scale).view(n_rows, k)                       # scaled targets, |x| <= 448
    sc = scale.expand(-1, -1, 32).reshape(n_rows, k)
    q8 = x.to(torch.float8_e4m3fn)
    bits = q8.view(torch.uint8).to(torch.int16)
    qf = q8.float()
    # the other neighbour: one E4M3 step from q towards x (magnitude up or down)
    mag, sign = bits & 0x7F, bits & 0x80
    toward_up = (x.abs() > qf.abs()) | ((qf == 0) & (x != 0))
    alt_mag = torch.where(toward_up, (mag + 1).clamp(max=0x7E), (mag - 1).clamp(min=0))
    alt_sign = torch.where((qf == 0) & (x < 0), torch.full_like(sign, 0x80), sign)
    alt_bits = (alt_sign | alt_mag).to(torch.uint8)
    altf = alt_bits.view(torch.float8_e4m3fn).float()
    valid = (alt_bits != bits.to(torch.uint8)) & (alt_mag != mag)
    # error contributions along r, per column, in original units
    err = (qf - x) * sc                                        # current error per element
    delta = (altf - qf) * sc                                   # change if flipped
    rv = r[:, None]
    col = (rv * err).sum(0)                                    # r^T E per column
    gain_unit = rv * delta                                     # change in r^T E per flip
    flipped = torch.zeros_like(valid)
    added = (err + delta) ** 2 - err ** 2                      # extra squared error if flipped (>= 0)
    tol = 1e-4 * col.abs().max().clamp_min(1e-12)
    for _ in range(ITERS):
        # per column, the flip with the best |r^T E| reduction per unit of added squared error
        new_abs = (col[None, :] + gain_unit).abs()
        reduction = col.abs()[None, :] - new_abs
        ok = valid & ~flipped & (reduction > 0)
        score = torch.where(ok, reduction / added.clamp_min(1e-30), torch.full_like(reduction, -1.0))
        best, idx = score.max(0)
        take = (best > 0) & (col.abs() > tol)
        if not take.any():
            break
        cols = torch.nonzero(take).squeeze(1)
        rows = idx[cols]
        flipped[rows, cols] = True
        col[cols] += gain_unit[rows, cols]
    fidelity.last_flips = flipped.float().mean().item()
    codes = torch.where(flipped, alt_bits, bits.to(torch.uint8)).view(torch.float8_e4m3fn)
    deq = codes.float() * sc
    scale_u8 = (exp.squeeze(-1) + 127).to(torch.uint8)
    return codes, scale_u8, deq


def fidelity(n):
    wp = tensor(BAKE, WK, n + '.weight').float()
    r, ratio = removed_direction(n, wp)
    w0 = base_dequant(n)
    codes, scale, deq = encode(wp, r)
    e = deq - wp
    keep_base = (r @ w0).norm()
    return {'layer': n, 'rank1_ratio_s2_s1': round(ratio, 5),
            'removed_fraction_bake': round(1 - ((r @ wp).norm() / keep_base).item(), 5),
            'removed_fraction_reencoded': round(1 - ((r @ deq).norm() / keep_base).item(), 5),
            'dir_err_over_intervention': round(((r @ e).norm() / (r @ (wp - w0)).norm()).item(), 5),
            'rel_err': round((e.norm() / wp.norm()).item(), 5),
            'flipped_fraction': round(fidelity.last_flips, 6)}, codes, scale


def report(layers):
    for layer in layers:
        for kind in ('linear_attn.out_proj', 'self_attn.o_proj'):
            n = f'model.language_model.layers.{layer}.{kind}'
            if n + '.weight' in WK:
                t = time.time(); f, _, _ = fidelity(n); f['seconds'] = round(time.time() - t, 1); print(json.dumps(f), flush=True)


def build(out):
    assert os.path.isdir(out) and not os.listdir(out), 'output dir must exist and be empty'
    targets = sorted({k[:-len('.weight')] for k in WK if k.endswith('.weight') and TARGET.match(k[:-len('.weight')])})
    assert len(targets) == 48, len(targets)
    by_shard = {}
    for n in targets:
        by_shard.setdefault(WK[n + '.weight'], []).append(n)
    reports, new_map, links = [], dict(WK), []
    os.makedirs(f'{out}/tensors', exist_ok=True)
    for shard in sorted(set(WK.values())):
        src, dst = f'{BAKE}/{shard}', f'{out}/{shard}'
        if shard not in by_shard:
            links.append(shard)                                # unchanged: hard-linked on the host afterwards
            continue
        with safe_open(src, 'pt', device='cpu') as f:
            meta = f.metadata()
            tensors = {k: f.get_tensor(k) for k in f.keys()}
        for n in by_shard[shard]:
            rep, codes, scale = fidelity(n)
            tensors[n + '.weight'] = codes.contiguous()
            tensors[n + '.weight_scale'] = scale.contiguous()
            new_map[n + '.weight_scale'] = shard
            reports.append(rep); print(json.dumps(rep), flush=True)
        save_file(tensors, dst, metadata=meta)
        del tensors
    for name in os.listdir(BAKE):                                # small files: copy
        p = f'{BAKE}/{name}'
        if os.path.isfile(p) and name not in ('model.safetensors.index.json', 'config.json', 'SHA256SUMS', 'README.md'):
            shutil.copy2(p, f'{out}/{name}')
    idx = dict(IK); idx['weight_map'] = dict(sorted(new_map.items()))
    json.dump(idx, open(f'{out}/model.safetensors.index.json', 'w'), indent=2)
    cfg = json.load(open(f'{BAKE}/config.json')); base_cfg = json.load(open(f'{BASE}/config.json'))
    q, bq = cfg['quantization_config'], base_cfg['quantization_config']
    for n in targets:
        q['quantized_layers'][n] = bq['quantized_layers'][n]
    q['ignore'] = [i for i in q['ignore'] if i not in set(targets)]
    cfg['quantization_config'] = q
    cfg['qwen50_reencode'] = {'attention_output_projections': 'MXFP8 (base format), constrained rounding along each layer removed direction',
                              'source_revision': '57f88b00e153fedbc10843ba0cd6284d4c2b5c6c', 'layers': len(targets)}
    json.dump(cfg, open(f'{out}/config.json', 'w'), indent=2)
    json.dump(reports, open(f'{out}/reencode-report.json', 'w'), indent=1)
    json.dump(links, open(f'{out}/link-unchanged-shards.json', 'w'), indent=1)
    print(json.dumps({'event': 'BUILT', 'layers': len(reports), 'rewritten_shards': len(by_shard)}))


if __name__ == '__main__':
    torch.set_num_threads(min(8, os.cpu_count() or 8))
    if sys.argv[1] == 'report':
        report([int(x) for x in sys.argv[2:]])
    elif sys.argv[1] == 'build':
        build(sys.argv[2])