fix: use bitsandbytes.functional for expert layer quantization

quantize_streaming.py

@@ -6,85 +6,80 @@ import gc
 import torch
 from pathlib import Path
 from safetensors.torch import load_file, save_file
-from bitsandbytes.nn import Linear4bit
-from torch import nn
+from bitsandbytes.functional import quantize_nf4
 from transformers import AutoConfig
 
 
-def quantize_tensor_to_4bit(tensor: torch.Tensor) -> torch.Tensor:
-    """Quantize a single 2D weight tensor to 4-bit NF4."""
-    in_features = tensor.size(1)
-    out_features = tensor.size(0)
+def quantize_weight_nf4(weight: torch.Tensor):
+    """Quantize a single weight tensor to NF4 using bitsandbytes functional API."""
+    if weight.dim() != 2:
+        return weight, None
 
-    # Create a temporary Linear4bit layer
-    linear_4bit = Linear4bit(
-        in_features,
-        out_features,
-        bias=False,
-        compute_dtype=torch.bfloat16,
-        quant_type="nf4",
-    ).to(tensor.device)
-
-    with torch.no_grad():
-        linear_4bit.weight = nn.Parameter(tensor.clone())
-        # Force quantization to happen
-        _ = linear_4bit.weight.quant_state
-
-    return linear_4bit.weight
+    # quantize_nf4 returns (quantized_tensor, quant_state)
+    qweight, quant_state = quantize_nf4(
+        weight,
+        blocksize=64,
+        compress_statistics=True,
+    )
+    return qweight, quant_state
 
 
 def streaming_quantize(model_path: str, output_path: str):
-    print(f"Streaming quantization of: {model_path}")
+    print(f"Streaming NF4 quantization: {model_path}")
     output_path = Path(output_path)
     output_path.mkdir(parents=True, exist_ok=True)
 
-    # Load config
     config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
 
-    # Find all shards
     import glob
     shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
     print(f"Found {len(shards)} shards\n")
 
-    for idx, shard_path in enumerate(shards):
-        print(f"[{idx+1}/{len(shards)}] Processing {Path(shard_path).name}...")
+    for idx, shard_file in enumerate(shards):
+        print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name}")
 
-        # Load shard to GPU 0
-        state_dict = load_file(shard_path, device="cuda:0")
+        state_dict = load_file(shard_file, device="cuda:0")
 
-        # Find all 2D weight tensors (Linear layers)
-        weight_keys = [k for k, v in state_dict.items() 
-                       if "weight" in k and isinstance(v, torch.Tensor) and v.dim() == 2]
+        weight_keys = [
+            k for k, v in state_dict.items()
+            if "weight" in k and isinstance(v, torch.Tensor) and v.dim() == 2
+        ]
 
-        print(f"   Found {len(weight_keys)} weight tensors to quantize")
+        print(f"   Quantizing {len(weight_keys)} tensors...")
+
+        quant_states = {}  # We need to save these too for loading later
 
-        # Quantize
         for key in weight_keys:
             try:
-                state_dict[key] = quantize_tensor_to_4bit(state_dict[key])
+                qweight, qstate = quantize_weight_nf4(state_dict[key])
+                state_dict[key] = qweight
+                if qstate is not None:
+                    quant_states[f"{key}.quant_state"] = qstate
             except Exception as e:
-                print(f"   Warning: Failed to quantize {key}: {e}")
+                print(f"   Warning: Failed on {key}: {e}")
 
-        # Move to CPU and save
-        state_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v 
-                      for k, v in state_dict.items()}
+        # Move everything to CPU
+        state_dict = {
+            k: v.cpu() if isinstance(v, torch.Tensor) else v
+            for k, v in state_dict.items()
+        }
+        quant_states = {
+            k: v.cpu() if isinstance(v, torch.Tensor) else v
+            for k, v in quant_states.items()
+        }
 
-        out_shard = output_path / f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
-        save_file(state_dict, out_shard)
+        # Save shard + quant states
+        shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
+        save_file({**state_dict, **quant_states}, output_path / shard_name)
 
-        print(f"   Saved → {out_shard.name}")
+        print(f"   Saved {shard_name}")
 
-        # Cleanup
-        del state_dict
+        del state_dict, quant_states
         gc.collect()
         torch.cuda.empty_cache()
 
-    # Save config and index
     config.save_pretrained(output_path)
-
-    # Create model.safetensors.index.json if needed (for multi-shard)
-    print(f"\n✅ Streaming quantization complete!")
-    print(f"   Output: {output_path}")
+    print(f"\n✅ Done → {output_path}")
 
 
 if __name__ == "__main__":