From a5ff7fd878d0a8663fc22760fd17f33c54f97c79 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Tue, 26 Aug 2025 15:21:35 -0500 Subject: [PATCH] Fixed one of two glaring bugs introduced by recent "improvements" --- distorch_2.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/distorch_2.py b/distorch_2.py index 22c499c..db5bdf8 100644 --- a/distorch_2.py +++ b/distorch_2.py @@ -140,12 +140,17 @@ def analyze_safetensor_loading(model_patcher, allocations_str): virtual_vram_gb = 0.0 distorch_alloc, virtual_vram_str = allocations_str.split('#') + + compute_device = virtual_vram_str.split(';')[0] + logger.info(f"[MultiGPU_DisTorch2] Compute Device: {compute_device}") + if not distorch_alloc: mode = "fraction" - distorch_alloc = calculate_safetensor_vvram_allocation(model_patcher, virtual_vram_str) - logger.info("[MultiGPU_DisTorch2] Examples:") + logger.info("[MultiGPU_DisTorch2] Expert String Examples:") logger.info(" Direct(byte) Mode - cuda:0,500mb;cuda:1,3.0g;cpu,5gb* -> '*' cpu = over/underflow device, put 0.50gb on cuda0, 3.00gb on cuda1, and 5.00gb (or the rest) on cpu") logger.info(" Ratio(%) Mode - cuda:0,8%;cuda:1,8%;cpu,4% -> 8:8:4 ratio, put 40% on cuda0, 40% on cuda1, and 20% on cpu") + distorch_alloc = calculate_safetensor_vvram_allocation(model_patcher, virtual_vram_str) + elif any(c in distorch_alloc.lower() for c in ['g', 'm', 'k', 'b']): mode = "byte" distorch_alloc = calculate_fraction_from_byte_expert_string(model_patcher, distorch_alloc) @@ -254,8 +259,6 @@ def analyze_safetensor_loading(model_patcher, allocations_str): device_assignments = {device: [] for device in DEVICE_RATIOS_DISTORCH.keys()} block_assignments = {} - compute_device = str(current_device) - # Create a memory quota for each donor device based on its calculated allocation. donor_devices = [d for d in sorted_devices if d != compute_device] donor_quotas = {