cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: load shards directly to GPU, no CPU bottleneck

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 21:44:28 -04:00
parent 82677070ff
commit d41648d4fc
1 changed files with 12 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
quantize_streaming.py
View File

@@ -31,9 +31,9 @@ def streaming_quantize(model_path, output_path):
print(f"Processing shard {shard_idx + 1}/{len(safetensors_files)}: {shard_file}")
print(f"{'='*60}")
# Load shard to CPU
print(" Loading shard to CPU...")
shard_state_dict = torch.load(shard_file, map_location="cpu", weights_only=False)
# Load shard directly to GPU 0
print(" Loading shard to GPU 0...")
shard_state_dict = torch.load(shard_file, map_location="cuda:0", weights_only=False)
# Quantize Linear layers in this shard using both GPUs
print(" Quantizing Linear layers (both GPUs)...")
@@ -46,7 +46,7 @@ def streaming_quantize(model_path, output_path):
gpu0_keys = weight_keys[::2] # Even indices
gpu1_keys = weight_keys[1::2] # Odd indices
# Quantize on GPU 0
# Quantize on GPU 0 (shard already on GPU 0)
print(" GPU 0: Quantizing...", end=" ")
for key in gpu0_keys:
tensor = shard_state_dict[key]
@@ -62,13 +62,17 @@ def streaming_quantize(model_path, output_path):
)
with torch.no_grad():
dummy_linear.weight = nn.Parameter(tensor.clone().to("cuda:0"))
dummy_linear.weight = nn.Parameter(tensor.clone())
_ = dummy_linear.weight.quant_state
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")
# Move shard to GPU 1 for second half
print(" Moving to GPU 1...")
shard_state_dict = {k: v.to("cuda:1") if isinstance(v, torch.Tensor) else v for k, v in shard_state_dict.items()}
# Quantize on GPU 1
print(" GPU 1: Quantizing...", end=" ")
for key in gpu1_keys:
@@ -85,10 +89,10 @@ def streaming_quantize(model_path, output_path):
)
with torch.no_grad():
dummy_linear.weight = nn.Parameter(tensor.clone().to("cuda:1"))
dummy_linear.weight = nn.Parameter(tensor.clone())
_ = dummy_linear.weight.quant_state
shard_state_dict[key] = dummy_linear.weight.to("cpu")
shard_state_dict[key] = dummy_linear.weight
quantized_keys += 1
print(f"✓ {len(gpu1_keys)} layers")