feat: use both GPUs in parallel for streaming quantization

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,13 +35,21 @@ 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
+        # Quantize Linear layers in this shard using both GPUs
-        print("   Quantizing Linear layers...")
+        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:
+        # Get all weight tensors
-                # This is a Linear layer weight
+        weight_keys = [k for k, v in shard_state_dict.items() if 'weight' in k and v.dim() == 2]
-                # Create a dummy Linear4bit to get the quantization format
+        
+        # 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)
             
@@ -53,17 +61,38 @@ def streaming_quantize(model_path, output_path):
                 quant_type='nf4',
             )
             
-                # Quantize the weight
             with torch.no_grad():
-                    dummy_linear.weight = nn.Parameter(tensor.clone())
+                dummy_linear.weight = nn.Parameter(tensor.clone().to("cuda:0"))
-                    # Force quantization by accessing quant_state
                 _ = dummy_linear.weight.quant_state
             
-                # Replace with quantized version
+            shard_state_dict[key] = dummy_linear.weight.to("cpu")
-                shard_state_dict[key] = dummy_linear.weight
             quantized_keys += 1
+        print(f"✓ {len(gpu0_keys)} layers")
         
-        print(f"   ✓ Quantized {quantized_keys} weights")
+        # 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():