| """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) |
| 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() |
| |
| 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) |
| |
| err = (qf - x) * sc |
| delta = (altf - qf) * sc |
| rv = r[:, None] |
| col = (rv * err).sum(0) |
| gain_unit = rv * delta |
| flipped = torch.zeros_like(valid) |
| added = (err + delta) ** 2 - err ** 2 |
| tol = 1e-4 * col.abs().max().clamp_min(1e-12) |
| for _ in range(ITERS): |
| |
| 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) |
| 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): |
| 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]) |
|
|
