From f1e394f1c9d3cf2c1b30cb98701ccd20a0f29117 Mon Sep 17 00:00:00 2001 From: zefie Date: Sun, 1 Nov 2026 18:30:02 -0500 Subject: [PATCH] improve quantization --- gguf_to_nvfp4_gguf.py | 1 + hf_to_nvfp4_gguf.py | 1 + hf_to_nvfp4_safetensors.py | 1 + nvfp4gguf/__init__.py | 12 + nvfp4gguf/quant.py | 571 +++++++++++++++++++++---------------- 5 files changed, 335 insertions(+), 251 deletions(-) diff --git a/gguf_to_nvfp4_gguf.py b/gguf_to_nvfp4_gguf.py index 1c61087..49a928f 100644 --- a/gguf_to_nvfp4_gguf.py +++ b/gguf_to_nvfp4_gguf.py @@ -584,6 +584,7 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument("--num-iter", type=int, default=CTQ_NUM_ITER, help="Learned-rounding iterations per matrix (convert_to_quant default: 4000)") parser.add_argument("--lr", type=float, default=1.0, help="Initial learning rate") parser.add_argument("--lr-min", type=float, default=1e-8, help="Lowest learning rate. The flat-loss check starts here (default: 1e-8).") + parser.add_argument("--lr-patience", type=int, default=5, help="Patience for learning rate decay (default: 5)") parser.add_argument("--lr-schedule", default="plateau", choices=("plateau", "exponential", "adaptive")) parser.add_argument( "--scale-optimization", default=DEFAULT_SCALE_OPTIMIZATION, choices=("fixed", "iterative", "joint"), diff --git a/hf_to_nvfp4_gguf.py b/hf_to_nvfp4_gguf.py index e2c0011..804fd72 100644 --- a/hf_to_nvfp4_gguf.py +++ b/hf_to_nvfp4_gguf.py @@ -1051,6 +1051,7 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: help="Learned-rounding optimizer. auto uses prodigy when it is installed, otherwise AdamW") parser.add_argument("--num-iter", type=int, default=CTQ_NUM_ITER, help="Learned-rounding iterations per matrix (convert_to_quant default: 4000)") parser.add_argument("--lr", type=float, default=1.0, help="Initial learning rate") + parser.add_argument("--lr-patience", type=int, default=5, help="Patience for learning rate decay (default: 5)") parser.add_argument("--lr-min", type=float, default=1e-8, help="Lowest learning rate. The flat-loss check starts here (default: 1e-8).") parser.add_argument("--lr-schedule", default="plateau", choices=("plateau", "exponential", "adaptive")) parser.add_argument( diff --git a/hf_to_nvfp4_safetensors.py b/hf_to_nvfp4_safetensors.py index 3e72042..5042fb7 100644 --- a/hf_to_nvfp4_safetensors.py +++ b/hf_to_nvfp4_safetensors.py @@ -291,6 +291,7 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: help="Learned-rounding iterations per matrix (convert_to_quant default: 4000)") parser.add_argument("--lr", type=float, default=1.0, help="Initial learning rate") parser.add_argument("--lr-min", type=float, default=1e-8, help="Lowest learning rate. The flat-loss check starts here (default: 1e-8).") + parser.add_argument("--lr-patience", type=int, default=5, help="Patience for learning rate decay (default: 5)") parser.add_argument("--lr-schedule", default="plateau", choices=("plateau", "exponential", "adaptive")) parser.add_argument( "--scale-optimization", default=DEFAULT_SCALE_OPTIMIZATION, choices=("fixed", "iterative", "joint"), diff --git a/nvfp4gguf/__init__.py b/nvfp4gguf/__init__.py index 6e6e038..1acba59 100644 --- a/nvfp4gguf/__init__.py +++ b/nvfp4gguf/__init__.py @@ -44,6 +44,12 @@ from .quant import ( scale_bake_notice, ue4m3_bits, ue4m3_bits_torch, + coordinate_descent_refine, + coordinate_descent_refine_torch, + sweep_per_tensor_scale, + sweep_per_tensor_scale_torch, + get_ue4m3_lut, + get_fp4_torch_constants, ) __all__ = [ @@ -84,6 +90,12 @@ __all__ = [ "quantize_and_pack", "refine_codes_for_scale", "refine_codes_for_scale_torch", + "coordinate_descent_refine", + "coordinate_descent_refine_torch", + "sweep_per_tensor_scale", + "sweep_per_tensor_scale_torch", + "get_ue4m3_lut", + "get_fp4_torch_constants", "report_reconstruction", "require_quant_settings", "resolve_optimizer", diff --git a/nvfp4gguf/quant.py b/nvfp4gguf/quant.py index c172c8f..99cd0d8 100644 --- a/nvfp4gguf/quant.py +++ b/nvfp4gguf/quant.py @@ -54,8 +54,9 @@ DEFAULT_SCALE_REFINEMENT_ROUNDS = 2 # Patience 1 decays the step on every miss and the lr-floor stop then ends # the search while the loss is still the best. Prodigy adapts its own step, # so a slower plateau and no floor abort let the budget actually run. -DEFAULT_LR_PATIENCE = 32 +DEFAULT_LR_PATIENCE = 5 DEFAULT_EARLY_STOP_LR = 0.0 +DEFAULT_LR_FLOOR = 1e-8 # A wobble of 2.000e-7 to 2.002e-7 is not progress, but only once the step # size has floored. Before that, any new best has to count or the search # stops while the loss is still around 1e-2. @@ -83,7 +84,7 @@ _TIGHT_ATTENTION = ( ) DEFAULT_LLAMA_CPP = Path.home() / "llama.cpp" -DEFAULT_CONVERT_TO_QUANT = Path("/mnt/d/dev/convert_to_quant") +DEFAULT_CONVERT_TO_QUANT = Path.home() / "convert_to_quant" def _ensure_import_path(path: Path, label: str) -> None: @@ -183,10 +184,99 @@ def resolved_lr_min(args) -> float: """Plateau floor from ``--lr-min``. ``--lr`` is only the starting rate.""" lr_min = getattr(args, "lr_min", None) if lr_min is None: - return 1e-8 + return DEFAULT_LR_FLOOR return float(lr_min) +def ue4m3_bits(scale_float: np.ndarray) -> np.ndarray: + """Encode positive scales as UE4M3 using torch's E4M3FN bit patterns. + + ggml's own encoder saturates every value from 256 to 448 down to 448. + Torch's float8 bits decode back to the right magnitude with ggml's decoder, + including the exponent-15 range that NVFP4 block scales actually use. + """ + array = np.ascontiguousarray(np.clip(scale_float, 0.0, F8_E4M3_MAX), dtype=np.float32) + bits = torch.from_numpy(array).to(torch.float8_e4m3fn).view(torch.uint8).numpy() + return np.ascontiguousarray(bits & np.uint8(0x7F)) + + +def _ue4m3_roundtrip(scale_float: np.ndarray) -> np.ndarray: + """Encode to UE4M3 and decode back, returning the actual runtime scale values.""" + bits = ue4m3_bits(scale_float) + return torch.from_numpy(bits).view(torch.float8_e4m3fn).float().numpy() + + +def ue4m3_bits_torch(scale_float: torch.Tensor) -> torch.Tensor: + """Encode positive scales as UE4M3 using torch's float8_e4m3fn directly on tensor/device.""" + clipped = torch.clamp(scale_float, 0.0, F8_E4M3_MAX).to(torch.float32) + bits = clipped.to(torch.float8_e4m3fn).view(torch.uint8) + return bits & 0x7F + + +def _ue4m3_roundtrip_torch(scale_float: torch.Tensor) -> torch.Tensor: + """Encode to UE4M3 and decode back directly in Torch without CPU/NumPy roundtrips.""" + bits = ue4m3_bits_torch(scale_float) + return bits.view(torch.float8_e4m3fn).to(torch.float32) + + +_UE4M3_LUT: torch.Tensor | None = None + + +def get_ue4m3_lut(device: str | torch.device | None = None) -> torch.Tensor: + """Cached 128-element lookup table for all positive UE4M3 float values.""" + global _UE4M3_LUT + if _UE4M3_LUT is None: + bits = torch.arange(128, dtype=torch.uint8) + vals = bits.view(torch.float8_e4m3fn).to(torch.float32) + vals[127] = F8_E4M3_MAX # 0x7F is NaN in float8_e4m3fn; saturate to 448.0 + _UE4M3_LUT = vals + if device is not None: + target_device = torch.device(device) if isinstance(device, str) else device + if _UE4M3_LUT.device != target_device: + return _UE4M3_LUT.to(device=target_device) + return _UE4M3_LUT + + +_FP4_TORCH_CACHE: dict[torch.device, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {} + + +def get_fp4_torch_constants(device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pre-loaded FP4 levels, argmin map, and reconstruction table on the target device.""" + target_device = torch.device(device) if isinstance(device, str) else device + if target_device not in _FP4_TORCH_CACHE: + levels_sorted = torch.from_numpy(FP4_LEVELS_SORTED).to(device=target_device, dtype=torch.float32) + level_map = torch.from_numpy(FP4_LEVEL_MAP).to(device=target_device, dtype=torch.long) + table = torch.from_numpy(FP4_TABLE).to(device=target_device, dtype=torch.float32) + _FP4_TORCH_CACHE[target_device] = (levels_sorted, level_map, table) + return _FP4_TORCH_CACHE[target_device] + + +def snap_to_fp4_codes_torch( + scaled: torch.Tensor, + fp4_levels_sorted: torch.Tensor, + fp4_level_map: torch.Tensor, +) -> torch.Tensor: + """Vectorized O(N) snap of pre-scaled weights to nearest FP4 nibble codes.""" + idx = torch.searchsorted(fp4_levels_sorted, scaled) + idx = idx.clamp(1, 15) + left = fp4_levels_sorted[idx - 1] + right = fp4_levels_sorted[idx] + choose_right = (scaled - right).abs() < (scaled - left).abs() + nearest = torch.where(choose_right, idx, idx - 1) + return fp4_level_map[nearest].to(torch.uint8) + + +def snap_to_fp4_torch( + scaled: torch.Tensor, + fp4_levels_sorted: torch.Tensor, + fp4_level_map: torch.Tensor, + fp4_table: torch.Tensor, +) -> torch.Tensor: + """Vectorized O(N) snap of pre-scaled weights to reconstructed FP4 float values.""" + codes = snap_to_fp4_codes_torch(scaled, fp4_levels_sorted, fp4_level_map) + return fp4_table[codes.long()] + + class GgufLearnedNVFP4: """LearnedNVFP4Converter plus the total scale the codes were optimized against.""" @@ -326,37 +416,6 @@ def bake_packed_nvfp4(raw: np.ndarray, per_tensor: float) -> np.ndarray: return np.ascontiguousarray(baked.reshape(raw.shape)) -def ue4m3_bits(scale_float: np.ndarray) -> np.ndarray: - """Encode positive scales as UE4M3 using torch's E4M3FN bit patterns. - - ggml's own encoder saturates every value from 256 to 448 down to 448. - Torch's float8 bits decode back to the right magnitude with ggml's decoder, - including the exponent-15 range that NVFP4 block scales actually use. - """ - array = np.ascontiguousarray(np.clip(scale_float, 0.0, F8_E4M3_MAX), dtype=np.float32) - bits = torch.from_numpy(array).to(torch.float8_e4m3fn).view(torch.uint8).numpy() - return np.ascontiguousarray(bits & np.uint8(0x7F)) - - -def _ue4m3_roundtrip(scale_float: np.ndarray) -> np.ndarray: - """Encode to UE4M3 and decode back, returning the actual runtime scale values.""" - bits = ue4m3_bits(scale_float) - return torch.from_numpy(bits).view(torch.float8_e4m3fn).float().numpy() - - -def ue4m3_bits_torch(scale_float: torch.Tensor) -> torch.Tensor: - """Encode positive scales as UE4M3 using torch's float8_e4m3fn directly on tensor/device.""" - clipped = torch.clamp(scale_float, 0.0, F8_E4M3_MAX).to(torch.float32) - bits = clipped.to(torch.float8_e4m3fn).view(torch.uint8) - return bits & 0x7F - - -def _ue4m3_roundtrip_torch(scale_float: torch.Tensor) -> torch.Tensor: - """Encode to UE4M3 and decode back directly in Torch without CPU/NumPy roundtrips.""" - bits = ue4m3_bits_torch(scale_float) - return bits.view(torch.float8_e4m3fn).to(torch.float32) - - def refine_codes_for_scale_torch( weight: torch.Tensor, codes: torch.Tensor, @@ -370,24 +429,27 @@ def refine_codes_for_scale_torch( w_t = weight.to(device=target_device, dtype=torch.float32) c_t = codes.to(device=target_device, dtype=torch.uint8) bs_t = block_scale.to(device=target_device, dtype=torch.float32) - fp4_table = torch.from_numpy(FP4_TABLE).to(device=target_device, dtype=torch.float32) + fp4_levels_sorted, fp4_level_map, fp4_table = get_fp4_torch_constants(target_device) + lut = get_ue4m3_lut(device=target_device) rows, cols = w_t.shape n_blocks_per_row = cols // 16 + pts = max(float(per_tensor), 1e-30) if bake: - product = bs_t * float(per_tensor) - runtime_scale = _ue4m3_roundtrip_torch(product) + scale_bits = ue4m3_bits_torch(bs_t * pts) + runtime_scale = lut[scale_bits.long()] else: - runtime_scale = _ue4m3_roundtrip_torch(bs_t) * float(per_tensor) + scale_bits = ue4m3_bits_torch(bs_t) + runtime_scale = lut[scale_bits.long()] * pts w_blocks = w_t.reshape(rows, n_blocks_per_row, 16) c_blocks = c_t.reshape(rows, n_blocks_per_row, 16) - rs_3d = runtime_scale.unsqueeze(-1) # (rows, n_blocks, 1) + rs_3d = runtime_scale.unsqueeze(-1) - grid = fp4_table.view(1, 1, 1, 16) * rs_3d.unsqueeze(-1) # (rows, n_blocks, 1, 16) - dist = (w_blocks.unsqueeze(-1) - grid).abs() # (rows, n_blocks, 16, 16) - best_codes = dist.argmin(dim=-1).to(torch.uint8) # (rows, n_blocks, 16) + safe_rs = rs_3d.clamp_min(1e-30) + scaled = w_blocks / safe_rs + best_codes = snap_to_fp4_codes_torch(scaled, fp4_levels_sorted, fp4_level_map) old_recon = fp4_table[c_blocks.long()] * rs_3d new_recon = fp4_table[best_codes.long()] * rs_3d @@ -417,12 +479,11 @@ def optimize_block_scales_torch( bake: bool, device: str | torch.device | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """Optimize block scales directly on GPU/Torch tensors.""" + """Optimize block scales using windowed discrete UE4M3 local search and outlier mitigation (Optimizations 3 & 5).""" target_device = device if device is not None else weight.device w_t = weight.to(device=target_device, dtype=torch.float32) c_t = codes.to(device=target_device, dtype=torch.uint8) bs_t = block_scale.to(device=target_device, dtype=torch.float32) - fp4_table = torch.from_numpy(FP4_TABLE).to(device=target_device, dtype=torch.float32) rows, cols = w_t.shape n_blocks = cols // 16 @@ -430,55 +491,77 @@ def optimize_block_scales_torch( c_blocks = c_t.reshape(rows, n_blocks, 16) abs_blocks = w_blocks.abs() - candidate_indices = list(range(15, 15 - _SCALE_CANDIDATES, -1)) - candidate_indices = [max(i, 0) for i in candidate_indices] + fp4_levels_sorted, fp4_level_map, fp4_table = get_fp4_torch_constants(target_device) + lut = get_ue4m3_lut(device=target_device) - scale_factor = 6.0 * max(float(per_tensor), 1e-30) - sorted_abs, _ = torch.sort(abs_blocks, dim=-1) # (R, B, 16) - candidate_scales_raw = torch.stack( - [sorted_abs[:, :, idx] / scale_factor for idx in candidate_indices], - dim=0, - ) - candidate_scales_raw = torch.clamp(candidate_scales_raw, min=1e-30) - - # Batched UE4M3 roundtrip across all candidates at once + pts = max(float(per_tensor), 1e-30) if bake: - all_scales_raw = candidate_scales_raw * float(per_tensor) - all_rt = _ue4m3_roundtrip_torch(all_scales_raw) + base_bits = ue4m3_bits_torch(bs_t * pts) else: - all_rt = _ue4m3_roundtrip_torch(candidate_scales_raw) * float(per_tensor) + base_bits = ue4m3_bits_torch(bs_t) + + sorted_abs, _ = torch.sort(abs_blocks, dim=-1) + scale_denom = 6.0 * pts + + # Percentile-based candidate bit patterns + b_max = ue4m3_bits_torch(sorted_abs[:, :, 15] / 6.0 if bake else sorted_abs[:, :, 15] / scale_denom) + b_sub = ue4m3_bits_torch(sorted_abs[:, :, 14] / 6.0 if bake else sorted_abs[:, :, 14] / scale_denom) + b_3rd = ue4m3_bits_torch(sorted_abs[:, :, 13] / 6.0 if bake else sorted_abs[:, :, 13] / scale_denom) + + # Outlier-dampened candidates (mitigates isolated spikes stretching scale) + b_damp1 = ue4m3_bits_torch((0.85 * sorted_abs[:, :, 15]) / 6.0 if bake else (0.85 * sorted_abs[:, :, 15]) / scale_denom) + b_damp2 = ue4m3_bits_torch((0.75 * sorted_abs[:, :, 15]) / 6.0 if bake else (0.75 * sorted_abs[:, :, 15]) / scale_denom) + + candidate_bit_list: list[torch.Tensor] = [] + + # Window around current scale (+/- 4 discrete steps in UE4M3 grid) + 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)) + + # 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) + candidate_bit_list.append(torch.clamp(b_damp1, 1, 126)) # Baseline MSE from original (codes, scale) so we never regress if bake: - orig_product = bs_t * float(per_tensor) - orig_rt = _ue4m3_roundtrip_torch(orig_product) + orig_rt = lut[base_bits.long()] else: - orig_rt = _ue4m3_roundtrip_torch(bs_t) * float(per_tensor) + orig_rt = lut[base_bits.long()] * pts orig_recon = fp4_table[c_blocks.long()] * orig_rt.unsqueeze(-1) - orig_mse = torch.mean((w_blocks - orig_recon).square(), dim=-1) + orig_mse = (w_blocks - orig_recon).square().mean(dim=-1) best_codes = c_blocks.clone() best_scale = bs_t.clone() best_mse = orig_mse.clone() - for ci in range(len(candidate_indices)): - cs_raw = candidate_scales_raw[ci] - rt_scale = all_rt[ci] - valid = rt_scale > 0 + for cand_b in candidate_bit_list: + cand_scale_val = lut[cand_b.long()] + if bake: + cand_rt = cand_scale_val + cand_raw = cand_scale_val / pts + else: + cand_rt = cand_scale_val * pts + cand_raw = cand_scale_val - rt_scale_3d = rt_scale.unsqueeze(-1) # (R, B, 1) - grid = fp4_table.view(1, 1, 1, 16) * rt_scale_3d.unsqueeze(-1) # (R, B, 1, 16) - dist = (w_blocks.unsqueeze(-1) - grid).abs() - opt_codes = dist.argmin(dim=-1).to(torch.uint8) # (R, B, 16) + valid = cand_rt > 0 + rs_3d = cand_rt.unsqueeze(-1) - recon = fp4_table[opt_codes.long()] * rt_scale_3d - mse = torch.mean((w_blocks - recon).square(), dim=-1) # (R, B) + safe_rs = 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()] * rs_3d + mse = (w_blocks - recon).square().mean(dim=-1) accept = (mse < best_mse) & valid if torch.any(accept): best_codes = torch.where(accept.unsqueeze(-1), opt_codes, best_codes) - best_scale = torch.where(accept, cs_raw, best_scale) + best_scale = torch.where(accept, cand_raw, best_scale) best_mse = torch.where(accept, mse, best_mse) n_changed = int((best_scale != bs_t).sum().item()) @@ -493,6 +576,103 @@ def optimize_block_scales_torch( return best_codes.reshape(rows, cols), best_scale +def coordinate_descent_refine_torch( + weight: torch.Tensor, + codes: torch.Tensor, + block_scale: torch.Tensor, + per_tensor: float, + bake: bool, + rounds: int = DEFAULT_SCALE_REFINEMENT_ROUNDS, + device: str | torch.device | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Multi-round coordinate descent alternating block scale optimization and code refinement (Optimization 2).""" + target_device = device if device is not None else weight.device + w_t = weight.to(device=target_device, dtype=torch.float32) + c_t = codes.to(device=target_device, dtype=torch.uint8) + bs_t = block_scale.to(device=target_device, dtype=torch.float32) + + 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: + log.info(" coordinate descent converged at round %d/%d", r + 1, total_rounds) + break + if scale_diff > 0: + c_t = refine_codes_for_scale_torch(w_t, c_t, bs_t, per_tensor, bake, device=target_device) + + return c_t, bs_t + + +def sweep_per_tensor_scale_torch( + weight: torch.Tensor, + codes: torch.Tensor, + block_scale: torch.Tensor, + per_tensor: float, + bake: bool, + device: str | torch.device | None = None, +) -> tuple[torch.Tensor, torch.Tensor, float]: + """Sweep per-tensor scale over fine multipliers to align block scales with UE4M3 mantissa bins (Optimization 4).""" + if bake: + return codes, block_scale, per_tensor + + target_device = device if device is not None else weight.device + w_t = weight.to(device=target_device, dtype=torch.float32) + c_t = codes.to(device=target_device, dtype=torch.uint8) + bs_t = block_scale.to(device=target_device, dtype=torch.float32) + rows, cols = w_t.shape + n_blocks = cols // 16 + w_blocks = w_t.reshape(rows, n_blocks, 16) + c_blocks = c_t.reshape(rows, n_blocks, 16) + + fp4_levels_sorted, fp4_level_map, fp4_table = get_fp4_torch_constants(target_device) + lut = get_ue4m3_lut(device=target_device) + + orig_bits = ue4m3_bits_torch(bs_t) + orig_rt = lut[orig_bits.long()] * float(per_tensor) + orig_recon = fp4_table[c_blocks.long()] * orig_rt.unsqueeze(-1) + best_total_mse = float((w_blocks - orig_recon).square().mean().item()) + + best_pt = per_tensor + 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] + + for m in multipliers: + 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( + " per-tensor scale sweep: adjusted by %+.2f%% (MSE %.3e -> %.3e)", + (best_pt / per_tensor - 1.0) * 100.0, + orig_mse_val, + best_total_mse, + ) + + return best_codes, best_bs, best_pt + + def refine_codes_for_scale( weight: np.ndarray | torch.Tensor, codes: np.ndarray | torch.Tensor, @@ -501,79 +681,16 @@ def refine_codes_for_scale( bake: bool, device: str | torch.device | None = None, ) -> np.ndarray | torch.Tensor: - """Re-select FP4 codes against the actual UE4M3-rounded runtime scale. - - The learned-rounding optimizer picks codes against full-precision scales. - After UE4M3 encoding the scale changes slightly, so some codes are no - longer optimal. This pass does a per-element argmin over all 16 FP4 - values. - - Returns new codes (same shape and type as input codes). - """ + """Re-select FP4 codes against the actual UE4M3-rounded runtime scale.""" if torch.is_tensor(weight): return refine_codes_for_scale_torch(weight, codes, block_scale, per_tensor, bake, device=device) - use_gpu = False - if device is not None: - use_gpu = str(device).startswith("cuda") - elif torch.cuda.is_available(): - use_gpu = True - - if use_gpu: - dev = device if device is not None else "cuda" - w_t = torch.from_numpy(weight).to(dev) - c_t = torch.from_numpy(codes).to(dev) - bs_t = torch.from_numpy(block_scale).to(dev) - refined_t = refine_codes_for_scale_torch(w_t, c_t, bs_t, per_tensor, bake, device=dev) - return refined_t.cpu().numpy() - - rows, cols = weight.shape - n_blocks_per_row = cols // 16 - - # Compute the actual scale the runtime will use after UE4M3 encoding. - if bake: - product = block_scale * np.float32(per_tensor) - runtime_scale = _ue4m3_roundtrip(product) - else: - runtime_scale = _ue4m3_roundtrip(block_scale) * np.float32(per_tensor) - - # Avoid repeated reshape calls (Optimization 3) - w_blocks = weight.reshape(rows, n_blocks_per_row, 16) - c_blocks = codes.reshape(rows, n_blocks_per_row, 16) - rs_3d = runtime_scale[..., None] - - # Replace brute-force FP4 argmin with lookup math (Optimization 1) - safe_rs = np.maximum(rs_3d, np.float32(1e-30)) - scaled = w_blocks / safe_rs - idx = np.searchsorted(FP4_LEVELS_SORTED, scaled) - idx = np.clip(idx, 1, 15) - left = FP4_LEVELS_SORTED[idx - 1] - right = FP4_LEVELS_SORTED[idx] - choose_right = np.abs(scaled - right) < np.abs(scaled - left) - nearest = np.where(choose_right, idx, idx - 1) - best_codes = FP4_LEVEL_MAP[nearest] - - # Precomputed FP4_TABLE (Optimization 4) & Inplace ops (Optimization 6) - old_recon = FP4_TABLE[c_blocks] * rs_3d - new_recon = FP4_TABLE[best_codes] * rs_3d - old_err = w_blocks - old_recon - np.square(old_err, out=old_err) - new_err = w_blocks - new_recon - np.square(new_err, out=new_err) - improved = new_err < old_err - refined = np.where(improved, best_codes, c_blocks) - - n_changed = int(np.count_nonzero(improved)) - if n_changed > 0: - old_mse = float(np.mean(old_err)) - new_mixed_err = np.where(improved, new_err, old_err) - new_mse = float(np.mean(new_mixed_err)) - log.info( - " code refinement: %d/%d codes changed, block MSE %.3e -> %.3e", - n_changed, int(refined.size), old_mse, new_mse, - ) - - return refined.reshape(rows, cols) + dev = device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu") + w_t = torch.from_numpy(weight).to(dev) + c_t = torch.from_numpy(codes).to(dev) + bs_t = torch.from_numpy(block_scale).to(dev) + refined_t = refine_codes_for_scale_torch(w_t, c_t, bs_t, per_tensor, bake, device=dev) + return refined_t.cpu().numpy() def optimize_block_scales( @@ -584,118 +701,57 @@ def optimize_block_scales( bake: bool, device: str | torch.device | None = None, ) -> tuple[np.ndarray | torch.Tensor, np.ndarray | torch.Tensor]: - """Try multiple candidate scales per block and pick the lowest-MSE pair. - - For each block of 16 values, the initial scale is set by the block's - absolute max, but a single outlier can waste precision for the other 15. - This pass tries several candidate scales derived from percentiles of - the block's absolute values, picks the best FP4 codes for each, and - keeps whichever (scale, codes) pair gives the smallest reconstruction - error. - - Returns (new_codes, new_block_scale) with the same shapes and types as inputs. - """ + """Try multiple candidate scales per block and pick the lowest-MSE pair.""" if torch.is_tensor(weight): return optimize_block_scales_torch(weight, codes, block_scale, per_tensor, bake, device=device) - use_gpu = False - if device is not None: - use_gpu = str(device).startswith("cuda") - elif torch.cuda.is_available(): - use_gpu = True + dev = device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu") + w_t = torch.from_numpy(weight).to(dev) + c_t = torch.from_numpy(codes).to(dev) + bs_t = torch.from_numpy(block_scale).to(dev) + best_codes_t, best_scale_t = optimize_block_scales_torch(w_t, c_t, bs_t, per_tensor, bake, device=dev) + return best_codes_t.cpu().numpy(), best_scale_t.cpu().numpy() - if use_gpu: - dev = device if device is not None else "cuda" - w_t = torch.from_numpy(weight).to(dev) - c_t = torch.from_numpy(codes).to(dev) - bs_t = torch.from_numpy(block_scale).to(dev) - best_codes_t, best_scale_t = optimize_block_scales_torch(w_t, c_t, bs_t, per_tensor, bake, device=dev) - return best_codes_t.cpu().numpy(), best_scale_t.cpu().numpy() - # NumPy path with optimizations 1, 2, 3, 4, 6 - rows, cols = weight.shape - n_blocks = cols // 16 - w_blocks = weight.reshape(rows, n_blocks, 16) - c_blocks = codes.reshape(rows, n_blocks, 16) - abs_blocks = np.abs(w_blocks) +def coordinate_descent_refine( + weight: np.ndarray | torch.Tensor, + codes: np.ndarray | torch.Tensor, + block_scale: np.ndarray | torch.Tensor, + per_tensor: float, + bake: bool, + rounds: int = DEFAULT_SCALE_REFINEMENT_ROUNDS, + device: str | torch.device | None = None, +) -> tuple[np.ndarray | torch.Tensor, np.ndarray | torch.Tensor]: + """Multi-round coordinate descent alternating block scale optimization and code refinement (Optimization 2).""" + if torch.is_tensor(weight): + return coordinate_descent_refine_torch(weight, codes, block_scale, per_tensor, bake, rounds=rounds, device=device) - candidate_indices = list(range(15, 15 - _SCALE_CANDIDATES, -1)) - candidate_indices = [max(i, 0) for i in candidate_indices] + dev = device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu") + w_t = torch.from_numpy(weight).to(dev) + c_t = torch.from_numpy(codes).to(dev) + bs_t = torch.from_numpy(block_scale).to(dev) + best_c_t, best_bs_t = coordinate_descent_refine_torch(w_t, c_t, bs_t, per_tensor, bake, rounds=rounds, device=dev) + return best_c_t.cpu().numpy(), best_bs_t.cpu().numpy() - scale_factor = np.float32(6.0) * max(np.float32(per_tensor), np.float32(1e-30)) - sorted_abs = np.sort(abs_blocks, axis=-1) - candidate_scales_raw = np.stack( - [sorted_abs[:, :, idx] / scale_factor for idx in candidate_indices], - axis=0, - ) - candidate_scales_raw = np.maximum(candidate_scales_raw, np.float32(1e-30)) - # Optimization 2: Batched UE4M3 roundtrip across all candidates at once - if bake: - all_scales_raw = candidate_scales_raw * np.float32(per_tensor) - all_rt = _ue4m3_roundtrip(all_scales_raw.reshape(-1)).reshape(candidate_scales_raw.shape) - else: - all_rt = _ue4m3_roundtrip(candidate_scales_raw.reshape(-1)).reshape(candidate_scales_raw.shape) * np.float32(per_tensor) +def sweep_per_tensor_scale( + weight: np.ndarray | torch.Tensor, + codes: np.ndarray | torch.Tensor, + block_scale: np.ndarray | torch.Tensor, + per_tensor: float, + bake: bool, + device: str | torch.device | None = None, +) -> tuple[np.ndarray | torch.Tensor, np.ndarray | torch.Tensor, float]: + """Sweep per-tensor scale over fine multipliers to minimize global reconstruction error (Optimization 4).""" + if torch.is_tensor(weight): + return sweep_per_tensor_scale_torch(weight, codes, block_scale, per_tensor, bake, device=device) - # Baseline MSE from original (codes, scale) so we never regress - if bake: - orig_product = block_scale * np.float32(per_tensor) - orig_rt = _ue4m3_roundtrip(orig_product) - else: - orig_rt = _ue4m3_roundtrip(block_scale) * np.float32(per_tensor) - - orig_rt_3d = orig_rt[..., None] - orig_recon = FP4_TABLE[c_blocks] * orig_rt_3d - orig_err = w_blocks - orig_recon - np.square(orig_err, out=orig_err) - orig_mse = orig_err.mean(axis=-1) - - best_codes = c_blocks.copy() - best_scale = block_scale.copy() - best_mse = orig_mse.copy() - - for ci in range(len(candidate_indices)): - cs_raw = candidate_scales_raw[ci] - rt_scale = all_rt[ci] - valid = rt_scale > 0 - - # Optimization 3: Avoid repeated reshape calls - rt_scale_3d = rt_scale[..., None] - - # Optimization 1: Replace brute-force FP4 argmin with lookup math - safe_rs = np.maximum(rt_scale_3d, np.float32(1e-30)) - scaled = w_blocks / safe_rs - idx = np.searchsorted(FP4_LEVELS_SORTED, scaled) - idx = np.clip(idx, 1, 15) - left = FP4_LEVELS_SORTED[idx - 1] - right = FP4_LEVELS_SORTED[idx] - choose_right = np.abs(scaled - right) < np.abs(scaled - left) - nearest = np.where(choose_right, idx, idx - 1) - opt_codes = FP4_LEVEL_MAP[nearest] - - # Optimization 4 & 6: Precomputed FP4_TABLE and inplace ops - recon = FP4_TABLE[opt_codes] * rt_scale_3d - err = w_blocks - recon - np.square(err, out=err) - mse = err.mean(axis=-1) - - accept = (mse < best_mse) & valid - if np.any(accept): - accept3 = accept[:, :, np.newaxis] - best_codes = np.where(accept3, opt_codes, best_codes) - best_scale = np.where(accept, cs_raw, best_scale) - best_mse = np.where(accept, mse, best_mse) - - n_changed = int(np.count_nonzero(best_scale != block_scale)) - if n_changed > 0: - orig_total_mse = float(np.mean(orig_mse)) - new_total_mse = float(np.mean(best_mse)) - log.info( - " scale optimization: %d/%d blocks rescaled, MSE %.3e -> %.3e", - n_changed, int(rows * n_blocks), orig_total_mse, new_total_mse, - ) - - return best_codes.reshape(rows, cols), best_scale + dev = device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu") + w_t = torch.from_numpy(weight).to(dev) + c_t = torch.from_numpy(codes).to(dev) + bs_t = torch.from_numpy(block_scale).to(dev) + best_c_t, best_bs_t, best_pt = sweep_per_tensor_scale_torch(w_t, c_t, bs_t, per_tensor, bake, device=dev) + return best_c_t.cpu().numpy(), best_bs_t.cpu().numpy(), best_pt def _cuda_python_candidates() -> list[Path]: @@ -1262,17 +1318,26 @@ def quantize_and_pack( quantizer: TensorQuantizer, weight: np.ndarray, bake: bool, name: str = "", ) -> tuple[np.ndarray, np.ndarray | None]: """Quantize one matrix, or a stack of expert matrices, into ggml NVFP4 bytes.""" + rounds = getattr(quantizer.args, "scale_refinement_rounds", DEFAULT_SCALE_REFINEMENT_ROUNDS) if weight.ndim == 2: codes, block_scale, per_tensor, _dequant = quantizer.quantize_matrix(weight, name) - codes, block_scale = optimize_block_scales(weight, codes, block_scale, per_tensor, bake, device=quantizer.device) - codes = refine_codes_for_scale(weight, codes, block_scale, per_tensor, bake, device=quantizer.device) + codes, block_scale = coordinate_descent_refine( + weight, codes, block_scale, per_tensor, bake, rounds=rounds, device=quantizer.device, + ) + codes, block_scale, per_tensor = sweep_per_tensor_scale( + weight, codes, block_scale, per_tensor, bake, device=quantizer.device, + ) return _pack_matrix(quantizer, codes, block_scale, per_tensor, bake) raws = [] scales = [] for index in range(weight.shape[0]): codes, block_scale, per_tensor, _dequant = quantizer.quantize_matrix(weight[index], name) - codes, block_scale = optimize_block_scales(weight[index], codes, block_scale, per_tensor, bake, device=quantizer.device) - codes = refine_codes_for_scale(weight[index], codes, block_scale, per_tensor, bake, device=quantizer.device) + codes, block_scale = coordinate_descent_refine( + weight[index], codes, block_scale, per_tensor, bake, rounds=rounds, device=quantizer.device, + ) + codes, block_scale, per_tensor = sweep_per_tensor_scale( + weight[index], codes, block_scale, per_tensor, bake, device=quantizer.device, + ) raw, scale = _pack_matrix(quantizer, codes, block_scale, per_tensor, bake) raws.append(raw) scales.append(1.0 if scale is None else float(scale[0])) @@ -1372,9 +1437,13 @@ def add_quant_arguments(parser: argparse.ArgumentParser) -> None: help="If the loss never reaches --target-loss, stop after this many steps without a new best once it has otherwise flattened. 0 disables it (default: 64)", ) parser.add_argument( - "--lr-min", type=float, default=1e-8, + "--lr-min", type=float, default=DEFAULT_LR_FLOOR, help="Lowest learning rate. The flat-loss check starts here (default: 1e-8).", ) + parser.add_argument( + "--lr-patience", type=int, default=DEFAULT_LR_PATIENCE, + help="Patience for learning rate decay (default: 5)", + ) parser.add_argument("--use-speed", action="store_true", help="Enable Prodigy's use_speed") parser.add_argument("--calib-samples", type=int, default=3072) parser.add_argument("--seed", type=int, default=42)