diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..c3faa42 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "cmake.sourceDirectory": "/home/zefie/dev/nvfp4_gguf_quant/llama.cpp" +} \ No newline at end of file diff --git a/hf_to_nvfp4_safetensors.py b/hf_to_nvfp4_safetensors.py index baa434a..932812e 100644 --- a/hf_to_nvfp4_safetensors.py +++ b/hf_to_nvfp4_safetensors.py @@ -21,12 +21,20 @@ import shutil import struct from pathlib import Path +import torch + from nvfp4gguf import ( CTQ_NUM_ITER, DEFAULT_CONVERT_TO_QUANT, + DEFAULT_EARLY_STOP_FLOOR_STALL, DEFAULT_EARLY_STOP_STALL, + DEFAULT_LR_FLOOR, + DEFAULT_LR_PATIENCE, + DEFAULT_LR_THRESHOLD, DEFAULT_SCALE_OPTIMIZATION, DEFAULT_SCALE_REFINEMENT_ROUNDS, + DEFAULT_TARGET_LOSS, + ensure_cuda_interpreter, load_learned_converter, quant_profile, resolve_optimizer, @@ -171,31 +179,62 @@ def merge_safetensors(paths: list[Path], dest: Path) -> None: def quantize_shard(convert_to_nvfp4, args, optimizer: str, source: Path, destination: Path) -> None: if destination.exists(): destination.unlink() - convert_to_nvfp4( - str(source), - str(destination), - exclude_layers=args.exclude_layers, - output_dtype=args.output_dtype, - simple=args.simple, - num_iter=args.num_iter, - heur=args.heur, - calib_samples=args.calib_samples, - seed=args.seed, - optimizer=optimizer, - lr=args.lr, - lr_schedule=args.lr_schedule, - top_p=args.top_p, - min_k=args.min_k, - max_k=args.max_k, - full_matrix=args.full_matrix, - early_stop_loss=args.early_stop_loss, - early_stop_stall=args.early_stop_stall, - scale_refinement_rounds=args.scale_refinement_rounds, - scale_optimization=args.scale_optimization, - low_memory=args.low_memory, - use_speed=args.use_speed, - verbose_output=args.verbose, - ) + + import convert_to_quant.formats.nvfp4_conversion as nvfp4_conversion + + orig_save_file = nvfp4_conversion.save_file + + def patched_save_file(tensors, file_path, metadata=None): + if getattr(args, "no_bias_correction", False): + from safetensors import safe_open + with safe_open(str(source), framework="pt", device="cpu") as src: + for k in list(tensors.keys()): + if k.endswith(".bias") and k in src.keys(): + tensors[k] = src.get_tensor(k) + return orig_save_file(tensors, file_path, metadata=metadata) + + nvfp4_conversion.save_file = patched_save_file + + orig_cuda_avail = torch.cuda.is_available + device_arg = getattr(args, "device", "auto") + if device_arg == "cpu": + torch.cuda.is_available = lambda: False + elif str(device_arg).startswith("cuda:") and torch.cuda.is_available(): + torch.cuda.set_device(device_arg) + + try: + convert_to_nvfp4( + str(source), + str(destination), + exclude_layers=args.exclude_layers, + output_dtype=args.output_dtype, + simple=args.simple, + num_iter=args.num_iter, + heur=args.heur, + calib_samples=args.calib_samples, + seed=args.seed, + optimizer=optimizer, + lr=args.lr, + lr_patience=args.lr_patience, + lr_min=args.lr_min, + lr_threshold=args.lr_threshold, + lr_schedule=args.lr_schedule, + top_p=args.top_p, + min_k=args.min_k, + max_k=args.max_k, + full_matrix=args.full_matrix, + early_stop_loss=args.early_stop_loss, + early_stop_stall=args.early_stop_stall, + scale_refinement_rounds=args.scale_refinement_rounds, + scale_optimization=args.scale_optimization, + low_memory=args.low_memory, + use_speed=args.use_speed, + verbose_output=args.verbose, + ) + finally: + nvfp4_conversion.save_file = orig_save_file + torch.cuda.is_available = orig_cuda_avail + if not destination.is_file() or destination.stat().st_size == 0: raise SystemExit(f"convert_to_quant did not write {destination}") @@ -225,11 +264,12 @@ def _apply_directory_profile(args, source: Path) -> None: args.exclude_layers = profile.skip print(f"{profile.arch} profile: leaving {profile.label} unquantized.") else: - print(f"{profile.arch} profile: --exclude-layers replaces the default skip of {profile.label}.") + print(f"{profile.arch} profile: --exclude/--exclude-layers replaces the default skip of {profile.label}.") return def convert_directory(args) -> None: + ensure_cuda_interpreter(args.device) source = Path(args.model).resolve() _apply_directory_profile(args, source) if not source.is_dir(): @@ -242,6 +282,7 @@ def convert_directory(args) -> None: shards = discover_shards(source) load_learned_converter(args.convert_to_quant) + import convert_to_quant.formats.nvfp4_conversion as nvfp4_conversion from convert_to_quant.formats.nvfp4_conversion import convert_to_nvfp4 from convert_to_quant.utils.logging import setup_logging @@ -249,9 +290,82 @@ def convert_directory(args) -> None: # and only the inner 4000-step bar is visible. setup_logging("VERBOSE" if args.verbose else "NORMAL") + base_converter = nvfp4_conversion.LearnedNVFP4Converter + + class EarlyStoppingLearnedNVFP4(base_converter): + target_loss = float(getattr(args, "target_loss", DEFAULT_TARGET_LOSS)) + early_stop_floor_stall = int(getattr(args, "early_stop_floor_stall", DEFAULT_EARLY_STOP_FLOOR_STALL)) + + def convert(self, *args_c, **kwargs_c): + self._last_lr = self.lr + self._floor_stall = 0 + self._best_seen = None + self._settled = False + return super().convert(*args_c, **kwargs_c) + + def _learning_rate_floored(self) -> bool: + return getattr(self, "_last_lr", self.lr) <= self.lr_min * 1.01 + + def _loss_has_settled(self, loss: float) -> bool: + target = float(getattr(self, "target_loss", DEFAULT_TARGET_LOSS) or 0) + return target > 0 and 0 < loss <= target + + def _check_improvement(self, current_loss: float, best_loss: float) -> bool: + settled = getattr(self, "_settled", False) + if self.lr_threshold > 0 and (settled or self._learning_rate_floored()): + return super()._check_improvement(current_loss, best_loss) + return current_loss < best_loss + + def _check_early_stop(self, current_loss: float, curr_lr: float, worse_loss_counter: int) -> bool: + self._last_lr = curr_lr + if 0 < current_loss < float("inf"): + best = getattr(self, "_best_seen", None) + self._best_seen = current_loss if best is None else min(best, current_loss) + best_seen = getattr(self, "_best_seen", None) + self._settled = self._loss_has_settled(best_seen if best_seen is not None else current_loss) + limit = int(getattr(self, "early_stop_floor_stall", 0) or 0) + if self._settled: + from convert_to_quant.utils.logging import info + info(f"\n - Loss reached {float(self.target_loss):g}. Stopping early.") + return True + if limit > 0 and self._learning_rate_floored() and worse_loss_counter >= limit: + from convert_to_quant.utils.logging import info + info(f"\n - Loss stalled at lr-min ({worse_loss_counter} steps). Stopping early.") + return True + return super()._check_early_stop(current_loss, curr_lr, worse_loss_counter) + + def _compute_svd_components(self, W_float32, verbose: bool = True): + if not self.full_matrix: + return super()._compute_svd_components(W_float32, verbose) + import math + rows, cols = W_float32.shape + max_rank = min(rows, cols) + k = min(self.max_k, max(self.min_k, int(math.floor(self.top_p * max_rank)))) + k = min(k, max_rank) + if verbose: + print(f" - Tensor shape: [{rows}, {cols}], Max rank: {max_rank}. Using k={k} components.") + print(" - Using torch.linalg.svd with full_matrices=False") + left, _, right = torch.linalg.svd(W_float32, full_matrices=False) + return left[:, :k], right[:k, :], k + + nvfp4_conversion.LearnedNVFP4Converter = EarlyStoppingLearnedNVFP4 + optimizer = resolve_optimizer(args.optimizer) + device_label = args.device + if device_arg := getattr(args, "device", "auto"): + if device_arg == "auto": + device_label = "cuda" if torch.cuda.is_available() else "cpu" + else: + device_label = device_arg + print(f"Quantizing {len(shards)} shard(s) from {source}") print(f"Output file: {output}") + print( + f"NVFP4 quantizer: {'simple round-to-grid' if args.simple else f'learned rounding ({optimizer})'} on {device_label}, " + f"{args.num_iter if not args.simple else 0} iterations, scale mode {args.scale_optimization}, " + f"lr patience {args.lr_patience}, stop at loss {args.target_loss:g}" + ) + if len(shards) == 1: quantize_shard(convert_to_nvfp4, args, optimizer, source / shards[0], output) print(f"Wrote {output}") @@ -284,14 +398,19 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument("--overwrite", action="store_true", help="Replace an existing output file") parser.add_argument("--convert-to-quant", type=Path, default=None, help="convert_to_quant checkout, if it is not installed") + parser.add_argument("--device", default="auto", help="auto, cpu, cuda, or cuda:N") parser.add_argument("--simple", action="store_true", help="Round onto the E2M1 grid without learned rounding") parser.add_argument("--optimizer", default="auto", choices=("auto", "prodigy", "adamw", "radam", "original"), help="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-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-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("--lr-threshold", type=float, default=DEFAULT_LR_THRESHOLD, + help="Minimum relative improvement to avoid stall count (default: 0.01)") parser.add_argument("--lr-schedule", default="plateau", choices=("plateau", "exponential", "adaptive")) parser.add_argument( "--scale-optimization", default=DEFAULT_SCALE_OPTIMIZATION, choices=("fixed", "iterative", "joint"), @@ -309,14 +428,24 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument("--early-stop-loss", type=float, default=5e-9) parser.add_argument( "--early-stop-stall", type=int, default=DEFAULT_EARLY_STOP_STALL, - help="Stop after this many steps without progress. At the learning-rate floor, progress means a 1%% loss drop (default: 2000)", + help="Stop after this many steps without a new best, at any learning rate (default: 2000)", ) - parser.add_argument("--calib-samples", type=int, default=3072) + parser.add_argument( + "--target-loss", type=float, default=DEFAULT_TARGET_LOSS, + help="Stop when the best loss reaches this value, at any learning rate (default: 1e-5). 0 disables it.", + ) + parser.add_argument( + "--early-stop-floor-stall", type=int, default=DEFAULT_EARLY_STOP_FLOOR_STALL, + 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("--calib-samples", type=int, default=3072, help="Random inputs used to correct a sibling bias") parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--no-bias-correction", action="store_true", help="Leave bias weights unchanged without bias correction") parser.add_argument("--use-speed", action="store_true", help="Enable Prodigy's use_speed") parser.add_argument("--low-memory", action="store_true", help="Stream each shard instead of loading it all at once") parser.add_argument("--heur", action="store_true", help="Skip layers convert_to_quant's size heuristic rejects") - parser.add_argument("--exclude-layers", default=None, help="Regex of tensor names to leave unquantized") + parser.add_argument("--exclude", "--exclude-layers", dest="exclude_layers", default=None, + help="Regex of tensor names to leave unquantized") parser.add_argument("--output-dtype", default="bfloat16", choices=("bfloat16", "float16"), help="Dtype recorded for quantized layers and used for unquantized 2D weights") parser.add_argument("--verbose", action="store_true") diff --git a/nvfp4gguf/__init__.py b/nvfp4gguf/__init__.py index 1acba59..90e89c3 100644 --- a/nvfp4gguf/__init__.py +++ b/nvfp4gguf/__init__.py @@ -7,6 +7,9 @@ from .quant import ( DEFAULT_EARLY_STOP_FLOOR_STALL, DEFAULT_TARGET_LOSS, DEFAULT_EARLY_STOP_STALL, + DEFAULT_LR_FLOOR, + DEFAULT_LR_PATIENCE, + DEFAULT_LR_THRESHOLD, DEFAULT_SCALE_OPTIMIZATION, DEFAULT_SCALE_REFINEMENT_ROUNDS, FP4_LEVELS, @@ -59,6 +62,9 @@ __all__ = [ "DEFAULT_EARLY_STOP_FLOOR_STALL", "DEFAULT_TARGET_LOSS", "DEFAULT_EARLY_STOP_STALL", + "DEFAULT_LR_FLOOR", + "DEFAULT_LR_PATIENCE", + "DEFAULT_LR_THRESHOLD", "DEFAULT_SCALE_OPTIMIZATION", "DEFAULT_SCALE_REFINEMENT_ROUNDS", "FP4_LEVELS", diff --git a/nvfp4gguf/quant.py b/nvfp4gguf/quant.py index bcb28cd..665c530 100644 --- a/nvfp4gguf/quant.py +++ b/nvfp4gguf/quant.py @@ -319,6 +319,9 @@ class GgufLearnedNVFP4: if self._settled: log.info(" Loss reached %g. Stopping early.", float(self.target_loss)) return True + if limit > 0 and self._learning_rate_floored() and worse_loss_counter >= limit: + log.info(" Loss stalled at lr-min (%d steps). Stopping early.", worse_loss_counter) + return True return super()._check_early_stop(current_loss, curr_lr, worse_loss_counter) def _compute_svd_components(self, W_float32, verbose: bool = True):