ControlNets for Multi-GPU

[ishan-modi]

Currently most of our documentation has the following way of using ControlNets

from diffusers import SanaControlNetModel, SanaControlNetPipeline
import torch

controlnet = SanaControlNetModel.from_pretrained(
    "ishan24/Sana_600M_1024px_ControlNet_diffusers",
    torch_dtype=torch.float16
)

pipe = SanaControlNetPipeline.from_pretrained(
    "Efficient-Large-Model/Sana_600M_1024px_diffusers",
    variant="fp16",
    controlnet=controlnet,
    torch_dtype=torch.float16,
)

pipe.to('cuda')
pipe.vae.to(torch.bfloat16)
pipe.text_encoder.to(torch.bfloat16)

This does work for a single GPU use case, but it is not appropriate for multi gpu, because controlnet and pipe both will have a separate device map, which causes issue related to device alignment during forward pass.

Following is how it should be documented

Single GPU

from diffusers import SanaControlNetModel, SanaControlNetPipeline
import torch

pipe = SanaControlNetPipeline.from_pretrained(
    "ishan24/Sana_600M_1024px_ControlNetPlus_diffusers",
    variant="fp16",
    torch_dtype=torch.float16
).to('cuda')

pipe.vae.to(torch.bfloat16)
pipe.text_encoder.to(torch.bfloat16)

Multi GPU

from diffusers import SanaControlNetModel, SanaControlNetPipeline
import torch

pipe = SanaControlNetPipeline.from_pretrained(
    "ishan24/Sana_600M_1024px_ControlNetPlus_diffusers",
    variant="fp16",
    torch_dtype=torch.float16,
    device_map="balanced"
)

pipe.vae.to(torch.bfloat16)
pipe.text_encoder.to(torch.bfloat16)

Let me know if my thinking is correct on this one and if we need to change Docstring
cc - @a-r-r-o-w

[ishan-modi]

@hlky, can you help here

[hlky]

@ishan-modi Looking into it. Passed models like controlnet aren't being considered in _get_final_device_map.

[ishan-modi]

Thanks for the response, I would love to help !

[hlky]

With some hacky workarounds I have it working, however, it's a lot slower than single GPU because transformer is being put to device 0 and controlnet to device 1, and the output is noisy with multi GPU. I will prepare a patch/draft PR so others can test, and I will test with another pipeline type to eliminate any issues in the Sana ControlNet PR as the cause but I suspect dtype is being messed up with the workarounds.

[ishan-modi]

hmm, I think it works with the following code and results are also good

from diffusers import SanaControlNetModel, SanaControlNetPipeline
import torch

pipe = SanaControlNetPipeline.from_pretrained(
    "ishan24/Sana_600M_1024px_ControlNetPlus_diffusers",
    variant="fp16",
    torch_dtype={"default": torch.bfloat16, "transformer": torch.float16},
    device_map="balanced"
)

I think if we try to load controlnet model separately there is a separate devicemap created for controlnet and we would already have another one with pipe. We would need to sync these devicemaps for efficient processing

[hlky]

ishan24/Sana_600M_1024px_ControlNetPlus_diffusers has the controlnet built-in so that is expected to work, it relates to this code

diffusers/src/diffusers/pipelines/pipeline_utils.py

Line 870 in 506f39a

init_dict, unused_kwargs, _ = pipeline_class.extract_init_dict(config_dict, **kwargs)

If the checkpoint Efficient-Large-Model/Sana_600M_1024px_diffusers doesn't have controlnet built-in then init_dict doesn't have controlnet, in turn it is not included in the device_map

diffusers/src/diffusers/pipelines/pipeline_utils.py

Line 910 in 506f39a

final_device_map = _get_final_device_map(

Passing controlnet=controlnet is the typical case so it is preferable to document it that way in the example, and if we can fix it with device_map we could possibly start documenting with device_map="balanced" rather than .to("cuda").

controlnet = SanaControlNetModel.from_pretrained(
    "ishan24/Sana_600M_1024px_ControlNet_diffusers",
    torch_dtype=torch.float16
)

pipe = SanaControlNetPipeline.from_pretrained(
    "Efficient-Large-Model/Sana_600M_1024px_diffusers",
    variant="fp16",
    controlnet=controlnet,
    torch_dtype=torch.float16,
    device_map="balanced",
)

For the Sana ControlNet PR feel free to use ishan24/Sana_600M_1024px_ControlNetPlus_diffusers in the example so it's not a blocker, we can always update it later.

[ishan-modi]

Yes !! that is exactly what I am saying I have created ishan24/Sana_600M_1024px_ControlNetPlus_diffusers model repo because it results into a single devicemap as opposed to having multiple device maps as in the other case.

But we need to figure out a way of merging the two device_maps or discourage using it that way. Anyway thanks I will make the change in controlnet PR. Feel free to close this PR

[yiyixuxu]

thanks @ishan-modi !
closing for now and will look to support device map for passed components in future