"""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 # 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])