Assign device to unet. Resolves #5897

[MohamadZeina]

What does this PR do?

Resolves issue #5897 - training using the example script examples/text_to_image/train_text_to_image_lora.py fails because parameters and data are on 2 devices, CPU and CUDA. This is caused by the unet LORA weights - I fix it by sending the unet to the accelerator.device after setting up the LORA weights.

Fixes #5897 (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case. Yes - LORA error when running train_text_to_image_lora.py, error Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! #5897
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@sayakpaul

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[sayakpaul]

Can you ensure you get expected quality with this change?

[MohamadZeina]

I've removed the initial unet.to(device) that's now redundant. Regarding quality - I don't have much experience training these models but the performance is roughly on par with what I expect.

I'm training a LORA with rank 16, on "stabilityai/stable-diffusion-2" with the usual resolution of 768. On a 3090 I can fit a batch size of 8, and I'm getting 1.65it/s. It's roughly 4 times faster than a LORA I trained on the same machine on SDXL using another example script that had no issues.

I hope that answers your question but let me know if you need anything else.

[sayakpaul]

Hi there. Will this change still matter in light of: #5388.

[MohamadZeina]

I'll let you know when I get a chance to test this.

I had a look at the peft code and I'm not sure - LORA weights are still added after moving the unet to device, which was causing the issue before. But I'm not sure how peft handles this, it might detect this and put the new weights on the right device.

[MohamadZeina]

@sayakpaul in my 1 test, looks like peft fixes the issue - it must be checking for device under the hood and making sure the adapter is on the same device

[sayakpaul]

That is very good to know. Feel free to continue to test and let us know about your findings.

[sayakpaul]

Is this still a problem? Note that we recently added peft as a dependency to this script.

[sayakpaul]

We recently incorporated peft in our training scripts and I believe with that, this issue shouldn't exist. Could you please recheck?

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.