Prune inherited multigpu clones when max_gpus is lowered

create_multigpu_deepclones cloned the existing 'multigpu' additional_models list verbatim and never pruned entries beyond limit_extra_devices. If a workflow was previously prepared for more GPUs, reducing max_gpus would leave stale clones attached and eligible for later scheduling. Replace the TODO block with a real prune that keeps only clones whose load_device is either the model's load_device or in limit_extra_devices, and re-match clones if anything was removed.

Amp-Thread-ID: https://ampcode.com/threads/T-019e43b8-8258-70fd-ab3a-53e4c97f85d5
Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
Jedrzej Kosinski 2026-05-20 16:46:45 -07:00
parent ba417750a7
commit dd85851efe

View File

@ -162,16 +162,16 @@ def create_multigpu_deepclones(model: ModelPatcher, max_gpus: int, gpu_options:
gpu_options.register(model) gpu_options.register(model)
else: else:
logging.info("No extra torch devices need initialization, skipping initializing MultiGPU Work Units.") logging.info("No extra torch devices need initialization, skipping initializing MultiGPU Work Units.")
# TODO: only keep model clones that don't go 'past' the intended max_gpu count # only keep model clones that don't go 'past' the intended max_gpu count;
# multigpu_models = model.get_additional_models_with_key("multigpu") # this prunes any inherited multigpu clones whose load_device is no longer allowed
# new_multigpu_models = [] # when max_gpus is lowered between runs.
# for m in multigpu_models: allowed_devices = set(limit_extra_devices)
# if m.load_device in limit_extra_devices: allowed_devices.add(model.load_device)
# new_multigpu_models.append(m) multigpu_models = model.get_additional_models_with_key("multigpu")
# model.set_additional_models("multigpu", new_multigpu_models) new_multigpu_models = [m for m in multigpu_models if m.load_device in allowed_devices]
# persist skip_devices for use in sampling code if len(new_multigpu_models) != len(multigpu_models):
# if len(skip_devices) > 0 or "multigpu_skip_devices" in model.model_options: model.set_additional_models("multigpu", new_multigpu_models)
# model.model_options["multigpu_skip_devices"] = skip_devices model.match_multigpu_clones()
return model return model