improve hf_to_nvfp4_safetensors

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "cmake.sourceDirectory": "/home/zefie/dev/nvfp4_gguf_quant/llama.cpp"
+}

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),
+    import convert_to_quant.formats.nvfp4_conversion as nvfp4_conversion
-        str(destination),
+
-        exclude_layers=args.exclude_layers,
+    orig_save_file = nvfp4_conversion.save_file
-        output_dtype=args.output_dtype,
+
-        simple=args.simple,
+    def patched_save_file(tensors, file_path, metadata=None):
-        num_iter=args.num_iter,
+        if getattr(args, "no_bias_correction", False):
-        heur=args.heur,
+            from safetensors import safe_open
-        calib_samples=args.calib_samples,
+            with safe_open(str(source), framework="pt", device="cpu") as src:
-        seed=args.seed,
+                for k in list(tensors.keys()):
-        optimizer=optimizer,
+                    if k.endswith(".bias") and k in src.keys():
-        lr=args.lr,
+                        tensors[k] = src.get_tensor(k)
-        lr_schedule=args.lr_schedule,
+        return orig_save_file(tensors, file_path, metadata=metadata)
-        top_p=args.top_p,
+
-        min_k=args.min_k,
+    nvfp4_conversion.save_file = patched_save_file
-        max_k=args.max_k,
+
-        full_matrix=args.full_matrix,
+    orig_cuda_avail = torch.cuda.is_available
-        early_stop_loss=args.early_stop_loss,
+    device_arg = getattr(args, "device", "auto")
-        early_stop_stall=args.early_stop_stall,
+    if device_arg == "cpu":
-        scale_refinement_rounds=args.scale_refinement_rounds,
+        torch.cuda.is_available = lambda: False
-        scale_optimization=args.scale_optimization,
+    elif str(device_arg).startswith("cuda:") and torch.cuda.is_available():
-        low_memory=args.low_memory,
+        torch.cuda.set_device(device_arg)
-        use_speed=args.use_speed,
+
-        verbose_output=args.verbose,
+    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-min", type=float, default=DEFAULT_LR_FLOOR,
-    parser.add_argument("--lr-patience", type=int, default=5, help="Patience for learning rate decay (default: 5)")
+                        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")

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",

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):