File size: 3,484 Bytes
1961af5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """Calibrate from original BF16 and measure actual native W4A4 + BF16 correction."""
import os,json,time,math
from pathlib import Path
import torch
from safetensors import safe_open
from safetensors.torch import load_file
from flashinfer import nvfp4_quantize,mm_fp4
from flashinfer.gemm import nvfp4_quantize_smooth,mm_nvfp4_svdquant
BASE=Path(os.environ['IMAGE21_BASE'])/'transformer'
OLD=Path(os.environ['IMAGE21_WORK_DIR'])
INDEX=json.loads((BASE/'diffusion_pytorch_model.safetensors.index.json').read_text())['weight_map']
def weight(name):
with safe_open(BASE/INDEX[name+'.weight'],framework='pt',device='cuda') as f:return f.get_tensor(name+'.weight')
def split(x):
g=torch.arange(len(x)//4,device=x.device)
return x[g*4+g%4].contiguous(),x[g*4+(g+1)%4].contiguous()
def error(y,ref):return ((y.float()-ref.float()).square().mean()/ref.float().square().mean()).sqrt().item()
def prepare(w,x,smoothing,rank,method="svd",mse=False):
amax=x.float().abs().amax(0).clamp_min(1e-5);wmax=w.float().abs().amax(0).clamp_min(1e-5)
smooth=torch.ones_like(amax) if smoothing==0 else (amax**smoothing/wmax**(1-smoothing)).clamp(.01,100)
smooth/=smooth.log().mean().exp()
pre=(1/smooth).to(torch.bfloat16)
ws=w.float()/pre.float()[None,:]
if rank:
torch.manual_seed(1234)
if method=='svd':
u,s,v=torch.svd_lowrank(ws,q=rank+32,niter=3)
up=u[:,:rank]*s[:rank];down=v[:,:rank]
elif method=='output':
u,s,v=torch.svd_lowrank(ws@(x.float()*pre.float()).T,q=rank+32,niter=3)
up=u[:,:rank];down=ws.T@up
else:
u,s,v=torch.svd_lowrank(x.float()*pre.float(),q=rank+32,niter=3)
down=v[:,:rank];up=ws@down
residual=ws-up@down.T
else:residual=ws
gx=(2688/((x.float()*pre.float()).abs().max()*1.2)).reshape(1)
gw=(2688/residual.abs().max()).reshape(1)
if mse:
from weight_quant import quantize_weight
h=(x.float()*pre.float()).square().mean(0)
h=h+.01*h.mean()
wq,wsc=quantize_weight(residual,gw,h)
else:wq,wsc=nvfp4_quantize(residual.to(torch.bfloat16),gw,backend='cute-dsl')
alpha=(1/(gx*gw)).reshape(1)
result={'weight':wq.view(torch.uint8),'sf':wsc.view(torch.uint8).reshape(-1),'pre':pre,'gx':gx,'alpha':alpha}
if rank:result.update(down=(pre.float()[:,None]*down).to(torch.bfloat16).contiguous(),up=(up/alpha).to(torch.bfloat16).contiguous())
return result
def run(x,p):
from dynamic_scale import scale
gx,alpha,up=scale(x,p['pre'],p['gx'],p['alpha'],p['up']) if 'up' in p else (p['gx'],p['alpha'],None)
q,sf=nvfp4_quantize_smooth(x,p['pre'],gx,backend='cute-dsl')
if 'down' in p:
d=x@p['down']
return mm_nvfp4_svdquant(q,p['weight'],sf,p['sf'],alpha,d,up,backend='cute-dsl')
return mm_fp4(q,p['weight'].T,sf,p['sf'],p['alpha'],backend='b12x')
@torch.inference_mode()
def main():
rows=[]
for block,leaf in [(0,'attn.to_q'),(0,'img_mlp.out'),(8,'attn.to_v'),(16,'img_mlp.proj'),(16,'img_mlp.gate_layer'),(24,'attn.to_out.0'),(31,'attn.to_k')]:
name=f'transformer_blocks.{block}.{leaf}';w=weight(name)
x=load_file(str(OLD/'calibration'/f'{name}.safetensors'))['activations'].cuda();fit,check=split(x)
rf=fit@w.T;rc=check@w.T
for smooth in [0.,.5,1.]:
for rank in [0,128]:
start=time.time();p=prepare(w,fit,smooth,rank);y=run(fit,p);yc=run(check,p)
row={'name':name,'smooth':smooth,'rank':rank,'fit_nrmse':error(y,rf),'check_nrmse':error(yc,rc),'seconds':time.time()-start};rows.append(row);print(json.dumps(row),flush=True)
Path('probe.json').write_text(json.dumps(rows,indent=2))
if __name__=='__main__':main()
|