Repository navigation
Allow controlnets to be loaded (from ckpt) in a parallel thread with a SD model (ckpt), and speed it up slightly - #4298
Conversation
…aded (from ckpt) in a parallel thread with a SD model (ckpt) without tensor errors (race condition)
sayakpaul
left a comment
sayakpaul
left a comment
There was a problem hiding this comment.
Thanks for the fix!
I think we purposefully avoid multi-threading for these conversion scripts to promote simplicity. So, not sure if there's much we can do here.
Seeking advice from @yiyixuxu, @williamberman, and @patrickvonplaten.
|
Hi @sayakpaul I understand. In this PR I'm just repeating what we do for unet, vae and other modules during the ckpt conversion. ctx = init_empty_weights if is_accelerate_available() else nullcontext
with ctx():
unet = UNet2DConditionModel(**unet_config)
if is_accelerate_available():
for param_name, param in converted_unet_checkpoint.items():
set_module_tensor_to_device(unet, param_name, "cpu", value=param)
else:
unet.load_state_dict(converted_unet_checkpoint)ctx = init_empty_weights if is_accelerate_available() else nullcontext
with ctx():
vae = AutoencoderKL(**vae_config)
if is_accelerate_available():
for param_name, param in converted_vae_checkpoint.items():
set_module_tensor_to_device(vae, param_name, "cpu", value=param)
else:
vae.load_state_dict(converted_vae_checkpoint)It's not a different multi-threaded solution, but it just follows the pattern used for unet, vae etc, of initializing the model class quickly, and then loading the tensors using the So maybe you could consider this PR as just a code-consistency + performance boost for Controlnet loading, with an added side benefit of allowing Controlnet models to be loaded in another thread. :) Controlnet's unet is pretty hefty as well. This PR isn't a broader fix for multi-threading. That's in the other issue (#4296). Thanks! |
|
This PR will also unblock Easy Diffusion, since I'm currently monkey-patching the controlnet loader function to allow our multi-GPU users to use controlnet. |
|
Makes sense thanks for explaining. Then I see no problems in this PR. |
|
Thanks! Training in ED - not yet. Really keen to get to it soon. |
williamberman
left a comment
williamberman
left a comment
There was a problem hiding this comment.
This is fine because they should be equivalent ways of loading weights, not because we say this code can be multi threaded :)
|
Thanks! Please let me know if anything else is needed from my end. |
…a SD model (ckpt), and speed it up slightly (huggingface#4298) Faster controlnet model instantiation, and allow controlnets to be loaded (from ckpt) in a parallel thread with a SD model (ckpt) without tensor errors (race condition)
…a SD model (ckpt), and speed it up slightly (huggingface#4298) Faster controlnet model instantiation, and allow controlnets to be loaded (from ckpt) in a parallel thread with a SD model (ckpt) without tensor errors (race condition)
5 participants
Avoids
metatensor errors (due to a race condition described in #4296).This allows a controlnet model to be loaded in parallel with an SD model (in two separate threads). As a side-effect, it slightly speeds up the controlnet model instantiation, by following the pattern used for unet and other sub-modules.
While this isn't a broader fix for the race-condition mentioned in #4296, it fixes the problem for Controlnet.
If there's a different way to address this issue, I'd be happy to alter the PR! :)
Thanks!
@patrickvonplaten @sayakpaul