fix: use accelerate device_map for proper MoE quantization
This commit is contained in:
Christian Medina
@@ -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"))
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.float16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
)
|
||||
|
||||
# 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("Loading model with BnB 4-bit (shards streamed to GPUs)...")
|
||||
print(" This will distribute across both GPUs\n")
|
||||
|
||||
print(f"Found {len(shards)} shards\n")
|
||||
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,
|
||||
)
|
||||
|
||||
config = AutoConfig.from_pretrained(model_path, 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")
|
||||
|
||||
# Process continuously: max 4 shards per GPU, add to whichever GPU has room
|
||||
max_per_gpu = 4
|
||||
num_gpus = 2
|
||||
print(f"Saving quantized model to: {output_path}")
|
||||
model.save_pretrained(output_path)
|
||||
|
||||
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)
|
||||
print(f"\n✅ Done → {output_path}")
|
||||
print(f"\n✅ Quantized model saved to: {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user