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 = {