fix: correct Nvfp4Linear dimension inference from checkpoint weights
Weight shape (N_packed, K_packed) means: - out_features = N_packed (GEMM output dim in BF16) - in_features = K_packed * 2 (BF16 input dim, for activation buffer)
This commit is contained in:
@@ -131,12 +131,12 @@ def make_nvfp4_linear(in_features, out_features, device, all_w, pfx, proj_name):
|
||||
d = device
|
||||
weight, ws, ws2, isc = get_nvfp4_weight(all_w, pfx, proj_name)
|
||||
assert weight is not None, f"{pfx}.{proj_name}.weight not found"
|
||||
# Infer actual dimensions from checkpoint weight shape
|
||||
# weight is (O, I//2) uint8 where O = out_features//2
|
||||
actual_out = weight.shape[0] * 2
|
||||
actual_in = weight.shape[1] * 2
|
||||
if actual_out != out_features or actual_in != in_features:
|
||||
log.warning(f"{pfx}.{proj_name}: declared ({in_features},{out_features}) but weight is ({actual_in},{actual_out})")
|
||||
# Checkpoint weight is (N_packed, K_packed) uint8
|
||||
# NVFP4 GEMM output dim = N_packed BF16 elements
|
||||
# Activation buffer needs K_packed FP4 columns = in_features BF16
|
||||
# So: in_features = K_packed * 2, out_features = N_packed
|
||||
actual_out = weight.shape[0] # N_packed = GEMM output dimension
|
||||
actual_in = weight.shape[1] * 2 # K_packed * 2 = BF16 input dim (for buffer allocation)
|
||||
lin = Nvfp4Linear(actual_in, actual_out, max_num_tokens=8192, device=d)
|
||||
lin.fp4 = [weight.to(d)]; lin.sf = [ws.to(d)]
|
||||
lin.ws2 = [ws2.to(d) if ws2 is not None else None] # weight_scale_2
|
||||
|
||||
Reference in New Issue
Block a user