zefie/nvfp4_gguf_quant
1
0 Fork 0
Code Actions Activity

some edits

Browse Source
This commit is contained in:
zefie committed 2026-11-01 18:06:11 -05:00
1 parent 8c66990853
commit 0c49394911
16 files changed
+416 -11
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

nvfp4gguf/quant.py
+396 -11
View File
@@ -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]))