Upload dequantize_to_bf16.py with huggingface_hub

dequantize_to_bf16.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+"""
+Dequantize a bitsandbytes NF4 model to standard bf16/fp16 weights.
+
+Usage:
+    python dequantize_to_bf16.py                    # uses current directory
+    python dequantize_to_bf16.py --input ./final-9b --output ./final-9b-bf16 --dtype bf16
+
+Requires: torch, transformers, bitsandbytes, safetensors
+"""
+
+import argparse
+import gc
+import os
+import torch
+import safetensors.torch
+from transformers import AutoModelForCausalLM, AutoConfig
+from bitsandbytes.nn import Linear4bit
+
+
+def dequantize_linear4bit(module: Linear4bit) -> torch.nn.Linear:
+    """Convert a Linear4bit module to a standard nn.Linear with dequantized weights."""
+    weight = module.weight
+    # Dequantize the 4-bit weight
+    if weight.quant_state is not None:
+        dequantized = module.dequantize(weight)
+    else:
+        # Fallback: extract via bnb.functional
+        import bitsandbytes.functional as bnb
+        dequantized = bnb.functional.dequantize_4bit(
+            weight.data,
+            weight.quant_state
+        )
+
+    out_features, in_features = dequantized.shape
+    new_linear = torch.nn.Linear(in_features, out_features, bias=module.bias is not None)
+    new_linear.weight.data = dequantized.contiguous()
+    if module.bias is not None:
+        new_linear.bias.data = module.bias.contiguous()
+    return new_linear
+
+
+def dequantize_model(
+    input_dir: str,
+    output_dir: str,
+    dtype: str = "bf16",
+):
+    device = "cpu"
+    torch_dtype = torch.bfloat16 if dtype == "bf16" else torch.float16
+
+    print(f"[*] Loading NF4 model from {input_dir} on CPU...")
+    model = AutoModelForCausalLM.from_pretrained(
+        input_dir,
+        device_map=device,
+        torch_dtype=torch_dtype,
+    )
+    print(f"    Loaded {sum(p.numel() for p in model.parameters()):,} total params")
+
+    print("[*] Dequantizing Linear4bit modules...")
+    for name, module in model.named_modules():
+        if isinstance(module, Linear4bit):
+            parent_name = ".".join(name.split(".")[:-1])
+            child_name = name.split(".")[-1]
+            if parent_name:
+                parent = model.get_submodule(parent_name)
+            else:
+                parent = model
+            new_linear = dequantize_linear4bit(module)
+            setattr(parent, child_name, new_linear)
+            print(f"    ✓ {name}")
+
+    print(f"[*] Collecting state dict ({dtype})...")
+    state_dict = model.state_dict()
+    state_dict = {k: v.to(dtype=torch_dtype, device="cpu") for k, v in state_dict.items()}
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    print(f"[*] Saving to {output_dir}...")
+    safetensors.torch.save_file(state_dict, os.path.join(output_dir, "model.safetensors"))
+
+    # Copy config files
+    import shutil
+    for f in ["config.json", "tokenizer.json", "tokenizer_config.json",
+              "generation_config.json", "chat_template.jinja"]:
+        src = os.path.join(input_dir, f)
+        if os.path.exists(src):
+            shutil.copy2(src, os.path.join(output_dir, f))
+
+    print(f"[✓] Done! Model saved to {output_dir}")
+    print(f"    Size: {sum(os.path.getsize(os.path.join(output_dir, f)) for f in os.listdir(output_dir) if os.path.isfile(os.path.join(output_dir, f))) / 1e9:.1f} GB")
+    print(f"    Dtype: {dtype}")
+    print(f"    Usage: AutoModelForCausalLM.from_pretrained('{output_dir}')")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Dequantize NF4 model to bf16/fp16")
+    parser.add_argument("--input", default=".", help="Input NF4 model directory")
+    parser.add_argument("--output", default="./dequantized", help="Output directory")
+    parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16",
+                        help="Target dtype (default: bf16)")
+    args = parser.parse_args()
+    dequantize_model(args.input, args.output, args.dtype)