diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1babcbb --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +__pycache__ +nvfp4gguf/__pycache__ diff --git a/__pycache__/gguf_info.cpython-313.pyc b/__pycache__/gguf_info.cpython-313.pyc deleted file mode 100644 index 1bdc4c0..0000000 Binary files a/__pycache__/gguf_info.cpython-313.pyc and /dev/null differ diff --git a/__pycache__/gguf_to_nvfp4_gguf.cpython-313.pyc b/__pycache__/gguf_to_nvfp4_gguf.cpython-313.pyc deleted file mode 100644 index 5837d9a..0000000 Binary files a/__pycache__/gguf_to_nvfp4_gguf.cpython-313.pyc and /dev/null differ diff --git a/__pycache__/hf_info.cpython-313.pyc b/__pycache__/hf_info.cpython-313.pyc deleted file mode 100644 index 30e5c86..0000000 Binary files a/__pycache__/hf_info.cpython-313.pyc and /dev/null differ diff --git a/__pycache__/hf_to_nvfp4_gguf.cpython-313.pyc b/__pycache__/hf_to_nvfp4_gguf.cpython-313.pyc deleted file mode 100644 index 1914030..0000000 Binary files a/__pycache__/hf_to_nvfp4_gguf.cpython-313.pyc and /dev/null differ diff --git a/__pycache__/hf_to_nvfp4_gguf.cpython-314.pyc b/__pycache__/hf_to_nvfp4_gguf.cpython-314.pyc deleted file mode 100644 index e06aad3..0000000 Binary files a/__pycache__/hf_to_nvfp4_gguf.cpython-314.pyc and /dev/null differ diff --git a/__pycache__/hf_to_nvfp4_safetensors.cpython-313.pyc b/__pycache__/hf_to_nvfp4_safetensors.cpython-313.pyc deleted file mode 100644 index 55285c1..0000000 Binary files a/__pycache__/hf_to_nvfp4_safetensors.cpython-313.pyc and /dev/null differ diff --git a/__pycache__/ollama_nvfp4.cpython-313.pyc b/__pycache__/ollama_nvfp4.cpython-313.pyc deleted file mode 100644 index 0d6eac5..0000000 Binary files a/__pycache__/ollama_nvfp4.cpython-313.pyc and /dev/null differ diff --git a/__pycache__/ollama_nvfp4.cpython-314.pyc b/__pycache__/ollama_nvfp4.cpython-314.pyc deleted file mode 100644 index b859ef7..0000000 Binary files a/__pycache__/ollama_nvfp4.cpython-314.pyc and /dev/null differ diff --git a/__pycache__/safetensors_to_nvfp4_gguf.cpython-313.pyc b/__pycache__/safetensors_to_nvfp4_gguf.cpython-313.pyc deleted file mode 100644 index 86d4f19..0000000 Binary files a/__pycache__/safetensors_to_nvfp4_gguf.cpython-313.pyc and /dev/null differ diff --git a/__pycache__/safetensors_to_nvfp4_gguf.cpython-314.pyc b/__pycache__/safetensors_to_nvfp4_gguf.cpython-314.pyc deleted file mode 100644 index 82fac26..0000000 Binary files a/__pycache__/safetensors_to_nvfp4_gguf.cpython-314.pyc and /dev/null differ diff --git a/nvfp4gguf/__init__.py b/nvfp4gguf/__init__.py index a6cae77..6e6e038 100644 --- a/nvfp4gguf/__init__.py +++ b/nvfp4gguf/__init__.py @@ -9,6 +9,10 @@ from .quant import ( DEFAULT_EARLY_STOP_STALL, DEFAULT_SCALE_OPTIMIZATION, DEFAULT_SCALE_REFINEMENT_ROUNDS, + FP4_LEVELS, + FP4_LEVELS_SORTED, + FP4_LEVEL_MAP, + FP4_TABLE, QUANT_VERSION, ShapeView, TensorQuantizer, @@ -25,16 +29,21 @@ from .quant import ( must_bake_nvfp4_scale, mtp_layer_tensor, bake_packed_nvfp4, + optimize_block_scales, + optimize_block_scales_torch, pack_nvfp4, preserve_scale_tensor, quant_profile, quant_settings, quantize_and_pack, + refine_codes_for_scale, + refine_codes_for_scale_torch, report_reconstruction, require_quant_settings, resolve_optimizer, scale_bake_notice, ue4m3_bits, + ue4m3_bits_torch, ) __all__ = [ @@ -46,6 +55,10 @@ __all__ = [ "DEFAULT_EARLY_STOP_STALL", "DEFAULT_SCALE_OPTIMIZATION", "DEFAULT_SCALE_REFINEMENT_ROUNDS", + "FP4_LEVELS", + "FP4_LEVELS_SORTED", + "FP4_LEVEL_MAP", + "FP4_TABLE", "QUANT_VERSION", "ShapeView", "TensorQuantizer", @@ -62,14 +75,19 @@ __all__ = [ "must_bake_nvfp4_scale", "mtp_layer_tensor", "bake_packed_nvfp4", + "optimize_block_scales", + "optimize_block_scales_torch", "pack_nvfp4", "preserve_scale_tensor", "quant_profile", "quant_settings", "quantize_and_pack", + "refine_codes_for_scale", + "refine_codes_for_scale_torch", "report_reconstruction", "require_quant_settings", "resolve_optimizer", "scale_bake_notice", "ue4m3_bits", + "ue4m3_bits_torch", ] diff --git a/nvfp4gguf/__pycache__/__init__.cpython-313.pyc b/nvfp4gguf/__pycache__/__init__.cpython-313.pyc deleted file mode 100644 index 0eaec0e..0000000 Binary files a/nvfp4gguf/__pycache__/__init__.cpython-313.pyc and /dev/null differ diff --git a/nvfp4gguf/__pycache__/bundle.cpython-313.pyc b/nvfp4gguf/__pycache__/bundle.cpython-313.pyc deleted file mode 100644 index 7aca80c..0000000 Binary files a/nvfp4gguf/__pycache__/bundle.cpython-313.pyc and /dev/null differ diff --git a/nvfp4gguf/__pycache__/quant.cpython-313.pyc b/nvfp4gguf/__pycache__/quant.cpython-313.pyc deleted file mode 100644 index e65e60f..0000000 Binary files a/nvfp4gguf/__pycache__/quant.cpython-313.pyc and /dev/null differ diff --git a/nvfp4gguf/quant.py b/nvfp4gguf/quant.py index 1e048f2..c172c8f 100644 --- a/nvfp4gguf/quant.py +++ b/nvfp4gguf/quant.py @@ -25,10 +25,26 @@ import torch log = logging.getLogger("nvfp4gguf") -QUANT_VERSION = 2 +QUANT_VERSION = 3 QK_NVFP4 = 64 NVFP4_BLOCK_BYTES = 36 F8_E4M3_MAX = 448.0 +# The 15 non-zero representable values in FP4 E2M1 (symmetric, unsigned half + zero). +# Codes 0-7 encode [0, 0.5, 1, 1.5, 2, 3, 4, 6]; codes 8-15 are negated. +_FP4_E2M1_VALUES = np.array( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=np.float32, +) +FP4_TABLE = _FP4_E2M1_VALUES.astype(np.float32) + +FP4_LEVELS = FP4_TABLE +FP4_LEVELS_ORDER = np.argsort(FP4_LEVELS) +FP4_LEVELS_SORTED = FP4_LEVELS[FP4_LEVELS_ORDER] +FP4_LEVEL_MAP = FP4_LEVELS_ORDER.astype(np.uint8) + +# Number of candidate scales to try in outlier-aware block scale selection. +_SCALE_CANDIDATES = 8 CTQ_NUM_ITER = 4000 # Joint trains the block scales inside the same loop. Iterative refreshes them # only on the original optimizer, so joint is the default that Prodigy, AdamW, @@ -255,8 +271,12 @@ def pack_nvfp4(codes: np.ndarray, scale_bits: np.ndarray) -> np.ndarray: code. ``scale_bits`` is ``(..., K / 16)``. A super-block is 4 scale bytes followed by 32 code bytes, matching ``block_nvfp4``. """ - codes = np.ascontiguousarray(codes, dtype=np.uint8) - scale_bits = np.ascontiguousarray(scale_bits, dtype=np.uint8) & np.uint8(0x7F) + codes = np.asarray(codes, dtype=np.uint8) + if not codes.flags.c_contiguous: + codes = np.ascontiguousarray(codes) + scale_bits = np.asarray(scale_bits, dtype=np.uint8) & np.uint8(0x7F) + if not scale_bits.flags.c_contiguous: + scale_bits = np.ascontiguousarray(scale_bits) if codes.shape[:-1] != scale_bits.shape[:-1]: raise ValueError(f"code shape {codes.shape} does not match scales {scale_bits.shape}") length = codes.shape[-1] @@ -318,6 +338,366 @@ def ue4m3_bits(scale_float: np.ndarray) -> np.ndarray: 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, + block_scale: torch.Tensor, + per_tensor: float, + bake: bool, + device: str | torch.device | None = None, +) -> torch.Tensor: + """Re-select FP4 codes against the actual UE4M3-rounded runtime scale directly on GPU/Torch.""" + 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_per_row = cols // 16 + + if bake: + product = bs_t * float(per_tensor) + runtime_scale = _ue4m3_roundtrip_torch(product) + else: + runtime_scale = _ue4m3_roundtrip_torch(bs_t) * float(per_tensor) + + 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) + + 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) + + old_recon = fp4_table[c_blocks.long()] * rs_3d + new_recon = fp4_table[best_codes.long()] * rs_3d + old_err = (w_blocks - old_recon).square() + new_err = (w_blocks - new_recon).square() + improved = new_err < old_err + refined = torch.where(improved, best_codes, c_blocks) + + n_changed = int(improved.sum().item()) + if n_changed > 0: + old_mse = float(old_err.mean().item()) + new_mixed_err = torch.where(improved, new_err, old_err) + new_mse = float(new_mixed_err.mean().item()) + log.info( + " code refinement: %d/%d codes changed, block MSE %.3e -> %.3e", + n_changed, int(refined.numel()), old_mse, new_mse, + ) + + return refined.reshape(rows, cols) + + +def optimize_block_scales_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]: + """Optimize block scales directly on GPU/Torch tensors.""" + 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 + w_blocks = w_t.reshape(rows, n_blocks, 16) + 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] + + 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 + if bake: + all_scales_raw = candidate_scales_raw * float(per_tensor) + all_rt = _ue4m3_roundtrip_torch(all_scales_raw) + else: + all_rt = _ue4m3_roundtrip_torch(candidate_scales_raw) * float(per_tensor) + + # 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) + else: + orig_rt = _ue4m3_roundtrip_torch(bs_t) * float(per_tensor) + + orig_recon = fp4_table[c_blocks.long()] * orig_rt.unsqueeze(-1) + orig_mse = torch.mean((w_blocks - orig_recon).square(), 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 + + 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) + + recon = fp4_table[opt_codes.long()] * rt_scale_3d + mse = torch.mean((w_blocks - recon).square(), dim=-1) # (R, B) + + 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_mse = torch.where(accept, mse, best_mse) + + n_changed = int((best_scale != bs_t).sum().item()) + if n_changed > 0: + orig_total_mse = float(orig_mse.mean().item()) + new_total_mse = float(best_mse.mean().item()) + 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 + + +def refine_codes_for_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, +) -> 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). + """ + 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) + + +def optimize_block_scales( + 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]: + """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. + """ + 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 + + 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) + + candidate_indices = list(range(15, 15 - _SCALE_CANDIDATES, -1)) + candidate_indices = [max(i, 0) for i in candidate_indices] + + 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) + + # 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 + + def _cuda_python_candidates() -> list[Path]: found: list[Path] = [] env_root = Path.home() / "anaconda3" / "envs" @@ -516,17 +896,18 @@ class TensorQuantizer: cached = torch.randn(samples, cols, generator=generator, dtype=torch.float32) self.calib_cache[cols] = cached - error = torch.from_numpy(np.ascontiguousarray(original - reconstructed, dtype=np.float32)) - error = error.to(self.device) - accumulator = torch.zeros(rows, dtype=torch.float32) + cached_device = cached.to(self.device) + error = torch.from_numpy(np.ascontiguousarray(original - reconstructed, dtype=np.float32)).to(self.device) + accumulator = torch.zeros(rows, dtype=torch.float32, device=self.device) batch = 256 for start in range(0, samples, batch): - batch_x = cached[start:start + batch].to(self.device) - accumulator += (batch_x @ error.T).sum(dim=0).float().cpu() - del batch_x - del error + batch_x = cached_device[start:start + batch] + accumulator += (batch_x @ error.T).sum(dim=0) + del error, cached_device + if str(self.device).startswith("cuda"): + torch.cuda.empty_cache() gc.collect() - return (accumulator / samples).numpy().astype(np.float32, copy=False) + return (accumulator / samples).cpu().numpy().astype(np.float32, copy=False) # Feed-forward matmuls in these graphs pass a null scale, so a sidecar @@ -883,11 +1264,15 @@ def quantize_and_pack( """Quantize one matrix, or a stack of expert matrices, into ggml NVFP4 bytes.""" 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) 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) 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]))