LAB 01 · 训练预算、适配与分块归一化

import math, json
N = 8_000_000_000
D = 160_000_000_000
C = 6*N*D
print(f'基线 C≈{C:.3e} FLOPs（名义近似）')
for n in (4e9, 8e9, 16e9, 32e9):
    print(f'N={n/1e9:4.0f}B, 固定预算 D={C/(6*n)/1e9:6.1f}B tokens')
assert C == 7.68e21

P=8_000_000_000
cards=4
parts={'weight':P*2/1e9,'gradient':P*2/1e9,'master':P*4/1e9,'moments':P*8/1e9}
def state_ledger(k):
    w,g,o=parts['weight'],parts['gradient'],parts['master']+parts['moments']
    return {'DDP':w+g+o,'ZeRO-1':w+g+o/k,'ZeRO-2':w+(g+o)/k,'ZeRO-3':(w+g+o)/k}
print('持久状态 GB:',parts)
print('四卡，每卡持久状态 GB:',state_ledger(cards))
print('单卡时所有分片方案回到同一账本:',state_ledger(1))
assert list(state_ledger(4).values()) == [128,56,44,32]
assert len(set(state_ledger(1).values())) == 1

width=4096
rank=8
full=width*width
trainable=2*width*rank
print('完整 W 参数:',full,'LoRA 参数:',trainable,'比例:',trainable/full)
assert full==16777216 and trainable==65536
W=[[1.0,2.0],[0.0,1.0]]
A=[[1.0,-1.0]]
B=[[0.5],[1.0]]
x=[2.0,1.0]
s=1.0
W_before=json.dumps(W)
def mv(M,x): return [sum(a*b for a,b in zip(row,x)) for row in M]
def lora_output(W,A,B,x,s=1.0):
    base=mv(W,x); delta=mv(B,mv(A,x))
    return {'Wx':base,'sBAx':[s*v for v in delta],'sum':[u+s*v for u,v in zip(base,delta)]}
print('更新前:',lora_output(W,A,B,x,s))
B[0][0]+=0.1
print('只更新 B 后:',lora_output(W,A,B,x,s))
assert json.dumps(W)==W_before
for r in (8,64,2048,4096):
    print('rank',r,'适配参数/原矩阵参数=',2*r/width)

def quantize(values,bits=4):
    qmax=2**(bits-1)-1
    scale=max(abs(x) for x in values)/qmax if any(values) else 1.0
    ints=[max(-qmax,min(qmax,round(x/scale))) for x in values]
    restored=[q*scale for q in ints]
    err=sum((a-b)**2 for a,b in zip(values,restored))/len(values)
    return {'scale':scale,'integer':ints,'dequantized':restored,'MSE':err}
regular=[0.1,-0.2,0.5,1.0]
outlier=regular+[12.0]
for label,values in [('基线',regular),('只增加离群权重',outlier)]:
    print(label,quantize(values,4))
print('同一离群数组改为8bit:',quantize(outlier,8))
print('全零边界:',quantize([0,0,0]))
assert quantize(outlier,4)['MSE'] > quantize(regular,4)['MSE']

values=[10.,20.,30.]
def block_example(weights):
    block1=sum(a*b for a,b in zip(weights[:2],values[:2]))/sum(weights[:2])
    block2=values[2]
    exact=sum(a*b for a,b in zip(weights,values))/sum(weights)
    return {'exact':exact,'wrong_sum':block1+block2,'unweighted_mean':(block1+block2)/2,'denominators':[sum(weights[:2]),weights[2]]}
print('[2,1,3]:',block_example([2,1,3]))
print('[2,1,4]:',block_example([2,1,4]))
assert math.isclose(block_example([2,1,3])['exact'],130/6)
assert not math.isclose(block_example([2,1,4])['exact'],block_example([2,1,4])['unweighted_mean'])

def online_softmax_value(blocks):
    m=-math.inf; denom=0.; numer=0.; trace=[]
    for block in blocks:
        valid=[(z,v) for z,v in block if z != -math.inf]
        if not valid:
            trace.append({'empty':True,'m':m,'denom':denom,'numer':numer}); continue
        new_m=max(m,max(z for z,v in valid))
        scale=0. if m == -math.inf else math.exp(m-new_m)
        denom=denom*scale+sum(math.exp(z-new_m) for z,v in valid)
        numer=numer*scale+sum(math.exp(z-new_m)*v for z,v in valid)
        m=new_m
        trace.append({'m':m,'denom':denom,'numer':numer,'old_scale':scale})
    return (numer/denom if denom else None),trace
for shift in (0.,1000.):
    out,trace=online_softmax_value([[(math.log(2)+shift,10),(shift,20)],[],[(math.log(4)+shift,30)]])
    print('logit整体平移',shift,'输出',out,'trace',trace)
    assert math.isclose(out,160/7,rel_tol=1e-12)
assert online_softmax_value([[],[(-math.inf,10)]])[0] is None
print('全 mask 行：无有效结果（None），由调用方明确处理。')

def edit_distance(a,b):
    row=list(range(len(b)+1))
    for i,ca in enumerate(a,1):
        new=[i]
        for j,cb in enumerate(b,1): new.append(min(new[-1]+1,row[j]+1,row[j-1]+(ca!=cb)))
        row=new
    return row[-1]
for output in ('ABCDEF','ABCDEG','XXXXXX'):
    print({'人工输出':output,'目标':'ABCDEF','exact_match':output=='ABCDEF','edit_distance':edit_distance(output,'ABCDEF')})