From 5018be86ec634a39e7c401c37e8f25507a2b9e01 Mon Sep 17 00:00:00 2001 From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com> Date: Fri, 3 Jul 2026 01:40:33 -0400 Subject: [PATCH] fix: use accelerate device_map for proper MoE quantization --- quantize_streaming.py | 167 ++++++++---------------------------------- 1 file changed, 30 insertions(+), 137 deletions(-) diff --git a/quantize_streaming.py b/quantize_streaming.py index 6a2e6f1..3afd486 100644 --- a/quantize_streaming.py +++ b/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 bitsandbytes.nn import Linear4bit -from torch import nn -from transformers import AutoConfig +from transformers import AutoModelForCausalLM, BitsAndBytesConfig, 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 - shards = sorted(glob.glob(f"{model_path}/*.safetensors")) - - # Rename existing 0-indexed files to 1-indexed - for existing in output_path.glob("*.safetensors"): - 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) + bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.float16, + bnb_4bit_use_double_quant=True, + ) - # Process continuously: max 4 shards per GPU, add to whichever GPU has room - max_per_gpu = 4 - 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) + print("Loading model with BnB 4-bit (shards streamed to GPUs)...") + print(" This will distribute across both GPUs\n") - config.save_pretrained(output_path) - print(f"\n✅ Done → {output_path}") + model = AutoModelForCausalLM.from_pretrained( + 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__":