fix: use accelerate device_map for proper MoE quantization

quantize_streaming.py

@@ -1,15 +1,11 @@
 #!/usr/bin/env python3
-"""True streaming 4-bit NF4 quantization - one shard at a time."""
+"""Streaming NF4 quantization using accelerate device_map."""
 
 import argparse
 import gc
 import torch
-import concurrent.futures
 from pathlib import Path
-from safetensors.torch import load_file, save_file
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig
-from bitsandbytes.nn import Linear4bit
-from torch import nn
-from transformers import AutoConfig
 
 
 def quantize_weight_nf4(weight: torch.Tensor):
@@ -36,139 +32,36 @@ def streaming_quantize(model_path: str, output_path: str):
     output_path = Path(output_path)
     output_path.mkdir(parents=True, exist_ok=True)
 
-    import glob
+    bnb_config = BitsAndBytesConfig(
-    shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
+        load_in_4bit=True,
-    
+        bnb_4bit_quant_type="nf4",
-    # Rename existing 0-indexed files to 1-indexed
+        bnb_4bit_compute_dtype=torch.float16,
-    for existing in output_path.glob("*.safetensors"):
+        bnb_4bit_use_double_quant=True,
-        parts = existing.name.split("-")
+    )
-        if len(parts) >= 3 and existing.name.startswith("model-00000-"):
-            num = int(parts[1])
-            new_name = f"model-{num+1:05d}-of-{len(shards):05d}.safetensors"
-            new_path = output_path / new_name
-            existing.rename(new_path)
-            print(f"Renamed: {existing.name} -> {new_name}")
-    
-    print(f"Found {len(shards)} shards\n")
-    
-    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
 
-    # Process continuously: max 4 shards per GPU, add to whichever GPU has room
+    print("Loading model with BnB 4-bit (shards streamed to GPUs)...")
-    max_per_gpu = 4
+    print("   This will distribute across both GPUs\n")
-    num_gpus = 2
-    
-    def process_shard(idx, shard_file, gpu_id):
-        """Process a single shard (called per-GPU)."""
-        shard_name = f"model-{idx+1:05d}-of-{len(shards):05d}.safetensors"
-        if (output_path / shard_name).exists():
-            return
-        
-        print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name} (GPU {gpu_id})")
-        
-        # Load to assigned GPU
-        state_dict = load_file(shard_file, device=f"cuda:{gpu_id}")
-        
-        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"   Quantizing {len(weight_keys)} tensors...")
-        
-        for key in weight_keys:
-            try:
-                weight = state_dict[key]
-                qweight, qstate = quantize_weight_nf4(weight)
-                state_dict[key] = qweight
-                del weight, qweight, qstate
-                gc.collect()
-                torch.cuda.empty_cache()
-            except Exception as e:
-                print(f"   Warning: Failed on {key}: {e}")
-                gc.collect()
-                torch.cuda.empty_cache()
-        
-        # Move to CPU and save
-        state_dict = {
-            k: v.cpu() if isinstance(v, torch.Tensor) else v
-            for k, v in state_dict.items()
-        }
-        
-        save_file(state_dict, output_path / shard_name)
-        print(f"   ✓ Saved {shard_name}")
-        
-        del state_dict
-        gc.collect()
-        torch.cuda.empty_cache()
-    
-    # Find first unsaved shard
-    start_idx = 0
-    for idx in range(len(shards)):
-        shard_name = f"model-{idx+1:05d}-of-{len(shards):05d}.safetensors"
-        if not (output_path / shard_name).exists():
-            start_idx = idx
-            break
-    
-    print(f"Starting from shard {start_idx+1}/16\n")
-    
-    # Process continuously: submit to GPU with fewer active tasks
-    import threading
-    gpu_locks = [threading.Lock() for _ in range(num_gpus)]
-    gpu_counts = [0] * num_gpus
-    
-    def get_next_gpu():
-        """Get the GPU with fewest active tasks (max 4 per GPU)."""
-        with gpu_locks[0]:
-            with gpu_locks[1]:
-                if gpu_counts[0] <= gpu_counts[1] and gpu_counts[0] < max_per_gpu:
-                    return 0
-                elif gpu_counts[1] < max_per_gpu:
-                    return 1
-                else:
-                    return None  # Both GPUs at max
-    
-    # Create a queue of unprocessed shards
-    from collections import deque
-    remaining = deque(range(start_idx, len(shards)))
-    
-    with concurrent.futures.ThreadPoolExecutor(max_workers=max_per_gpu * num_gpus) as executor:
-        futures = []
-        
-        def submit_next():
-            """Submit next shard to available GPU."""
-            gpu_id = get_next_gpu()
-            if gpu_id is None or not remaining:
-                return
-            
-            idx = remaining.popleft()
-            shard_file = shards[idx]
-            
-            with gpu_locks[gpu_id]:
-                gpu_counts[gpu_id] += 1
-            
-            future = executor.submit(process_shard, idx, shard_file, gpu_id)
-            futures.append(future)
-            
-            # When this future completes, release GPU slot and submit next
-            def on_complete(f):
-                with gpu_locks[gpu_id]:
-                    gpu_counts[gpu_id] -= 1
-                submit_next()  # Try to submit next shard
-            
-            future.add_done_callback(on_complete)
-        
-        # Start submitting
-        while remaining:
-            gpu_id = get_next_gpu()
-            if gpu_id is None:
-                break
-            submit_next()
-        
-        # Wait for all to complete
-        concurrent.futures.wait(futures)
 
-    config.save_pretrained(output_path)
+    model = AutoModelForCausalLM.from_pretrained(
-    print(f"\n✅ Done → {output_path}")
+        model_path,
+        quantization_config=bnb_config,
+        device_map="auto",
+        max_memory={0: "28GiB", 1: "28GiB", "cpu": "120GiB"},
+        low_cpu_mem_usage=True,
+        trust_remote_code=True,
+    )
+
+    # Count quantized parameters
+    total_params = sum(p.numel() for p in model.parameters())
+    bnb_params = sum(p.numel() for p in model.parameters() if hasattr(p, 'quant_state') and p.quant_state is not None)
+    print(f"✓ Model loaded and quantized")
+    print(f"   Total: {total_params / 1e9:.2f}B parameters")
+    print(f"   Quantized: {bnb_params / 1e9:.2f}B parameters ({bnb_params/total_params*100:.1f}%)\n")
+
+    print(f"Saving quantized model to: {output_path}")
+    model.save_pretrained(output_path)
+    
+    print(f"\n✅ Quantized model saved to: {output_path}")
 
 
 if __name__ == "__main__":