improve quantization

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)