diff --git a/nvfp4gguf/quant.py b/nvfp4gguf/quant.py index 99cd0d8..bcb28cd 100644 --- a/nvfp4gguf/quant.py +++ b/nvfp4gguf/quant.py @@ -518,13 +518,18 @@ def optimize_block_scales_torch( for delta in (-4, -3, -2, -1, 1, 2, 3, 4): candidate_bit_list.append(torch.clamp(base_bits.to(torch.int16) + delta, 1, 126).to(torch.uint8)) + # Max-derived candidates (the natural block scale from the largest element) + for delta in (-1, 0, 1): + candidate_bit_list.append(torch.clamp(b_max.to(torch.int16) + delta, 1, 126).to(torch.uint8)) + # Outlier-clipped scales (2nd and 3rd max) for delta in (-1, 0, 1): candidate_bit_list.append(torch.clamp(b_sub.to(torch.int16) + delta, 1, 126).to(torch.uint8)) candidate_bit_list.append(torch.clamp(b_3rd, 1, 126)) - # Outlier dampening (85% max) + # Outlier dampening (85% and 75% of max — clips isolated spikes) candidate_bit_list.append(torch.clamp(b_damp1, 1, 126)) + candidate_bit_list.append(torch.clamp(b_damp2, 1, 126)) # Baseline MSE from original (codes, scale) so we never regress if bake: @@ -591,10 +596,15 @@ def coordinate_descent_refine_torch( c_t = codes.to(device=target_device, dtype=torch.uint8) bs_t = block_scale.to(device=target_device, dtype=torch.float32) + # Initial code refinement against the post-UE4M3 runtime scale before + # entering the scale optimization loop. The learned rounding chose codes + # against the full-precision scale, but the runtime scale is quantized to + # UE4M3. Snapping codes first gives the scale optimizer a better baseline. + c_t = refine_codes_for_scale_torch(w_t, c_t, bs_t, per_tensor, bake, device=target_device) + total_rounds = max(int(rounds), 1) for r in range(total_rounds): prev_bs = bs_t - prev_c = c_t c_t, bs_t = optimize_block_scales_torch(w_t, c_t, bs_t, per_tensor, bake, device=target_device) scale_diff = int((bs_t != prev_bs).sum().item()) if r > 0 and scale_diff == 0: @@ -639,7 +649,7 @@ def sweep_per_tensor_scale_torch( best_bs = bs_t best_codes = c_t - multipliers = [0.95, 0.97, 0.98, 0.99, 1.01, 1.02, 1.03, 1.05] + multipliers = [0.93, 0.95, 0.97, 0.98, 0.99, 1.01, 1.02, 1.03, 1.05, 1.07] for m in multipliers: cand_pt = per_tensor * m @@ -661,6 +671,30 @@ def sweep_per_tensor_scale_torch( best_bs = cand_bs best_codes = opt_codes.reshape(rows, cols) + # Fine refinement around the best multiplier from the coarse pass + if best_pt != per_tensor: + coarse_m = best_pt / per_tensor + for fine_delta in (-0.005, -0.003, -0.001, 0.001, 0.003, 0.005): + m = coarse_m + fine_delta + cand_pt = per_tensor * m + cand_bs = bs_t / m + cand_bits = ue4m3_bits_torch(cand_bs) + cand_rt = lut[cand_bits.long()] * float(cand_pt) + cand_rs_3d = cand_rt.unsqueeze(-1) + + safe_rs = cand_rs_3d.clamp_min(1e-30) + scaled = w_blocks / safe_rs + opt_codes = snap_to_fp4_codes_torch(scaled, fp4_levels_sorted, fp4_level_map) + + recon = fp4_table[opt_codes.long()] * cand_rs_3d + cand_mse = float((w_blocks - recon).square().mean().item()) + + if cand_mse < best_total_mse: + best_total_mse = cand_mse + best_pt = cand_pt + best_bs = cand_bs + best_codes = opt_codes.reshape(rows, cols) + if best_pt != per_tensor: orig_mse_val = float((w_blocks - orig_recon).square().mean().item()) log.info(