diff --git a/quantize_streaming.py b/quantize_streaming.py index f38e8b3..f75eb58 100644 --- a/quantize_streaming.py +++ b/quantize_streaming.py @@ -11,7 +11,7 @@ from torch import nn def streaming_quantize(model_path, output_path): - """Quantize model by processing one shard at a time.""" + """Quantize model by processing one shard at a time, using both GPUs.""" print(f"Loading config from: {model_path}") config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) @@ -35,35 +35,64 @@ def streaming_quantize(model_path, output_path): print(" Loading shard to CPU...") shard_state_dict = torch.load(shard_file, map_location="cpu", weights_only=True) - # Quantize Linear layers in this shard - print(" Quantizing Linear layers...") + # Quantize Linear layers in this shard using both GPUs + print(" Quantizing Linear layers (both GPUs)...") quantized_keys = 0 - for key, tensor in list(shard_state_dict.items()): - if 'weight' in key and tensor.dim() == 2: - # This is a Linear layer weight - # Create a dummy Linear4bit to get the quantization format - in_features = tensor.size(1) - out_features = tensor.size(0) - - dummy_linear = Linear4bit( - in_features, - out_features, - bias=False, - compute_dtype=torch.float16, - quant_type='nf4', - ) - - # Quantize the weight - with torch.no_grad(): - dummy_linear.weight = nn.Parameter(tensor.clone()) - # Force quantization by accessing quant_state - _ = dummy_linear.weight.quant_state - - # Replace with quantized version - shard_state_dict[key] = dummy_linear.weight - quantized_keys += 1 - print(f" ✓ Quantized {quantized_keys} weights") + # Get all weight tensors + weight_keys = [k for k, v in shard_state_dict.items() if 'weight' in k and v.dim() == 2] + + # Distribute between GPUs + gpu0_keys = weight_keys[::2] # Even indices + gpu1_keys = weight_keys[1::2] # Odd indices + + # Quantize on GPU 0 + print(" GPU 0: Quantizing...", end=" ") + for key in gpu0_keys: + tensor = shard_state_dict[key] + in_features = tensor.size(1) + out_features = tensor.size(0) + + dummy_linear = Linear4bit( + in_features, + out_features, + bias=False, + compute_dtype=torch.float16, + quant_type='nf4', + ) + + with torch.no_grad(): + dummy_linear.weight = nn.Parameter(tensor.clone().to("cuda:0")) + _ = dummy_linear.weight.quant_state + + shard_state_dict[key] = dummy_linear.weight.to("cpu") + quantized_keys += 1 + print(f"✓ {len(gpu0_keys)} layers") + + # Quantize on GPU 1 + print(" GPU 1: Quantizing...", end=" ") + for key in gpu1_keys: + tensor = shard_state_dict[key] + in_features = tensor.size(1) + out_features = tensor.size(0) + + dummy_linear = Linear4bit( + in_features, + out_features, + bias=False, + compute_dtype=torch.float16, + quant_type='nf4', + ) + + with torch.no_grad(): + dummy_linear.weight = nn.Parameter(tensor.clone().to("cuda:1")) + _ = dummy_linear.weight.quant_state + + shard_state_dict[key] = dummy_linear.weight.to("cpu") + quantized_keys += 1 + print(f"✓ {len(gpu1_keys)} layers") + + print(f" ✓ Total: {quantized_keys} layers quantized") # Save quantized shard shard_name = f"model_shard_{shard_idx:05d}.safetensors" @@ -74,6 +103,7 @@ def streaming_quantize(model_path, output_path): # Free memory del shard_state_dict gc.collect() + torch.cuda.empty_cache() # Save config print(f"\n{'='*60}") @@ -86,6 +116,7 @@ def streaming_quantize(model_path, output_path): torch.cuda.empty_cache() print("\n✓ Streaming quantization complete!") + print(f" Used both GPUs in parallel for faster quantization") def main():