Update transformer to accelerator.device with it's weight_dtype

[Viditnegi]

By default, the transformer always gets transferred to the device with float32 whether we specify the dtype as bfloat16 or float16. This fixes it.

What does this PR do?

Fixes # (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.
• 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.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

We should not be doing it as we're fine-tuning the transformer. accelerator.prepare() handles this for us (plus any additional hook placements that might be required).

[Viditnegi]

The transformer gets transferred to the gpu in float32 by default with accelerator.prepare(). Takes 40 gb memory at once!

[sayakpaul]

Well it will trained in mixed-precision through autocast. This is handled by accelerate. If the user wants to cast it to a lower precision, it should be done through an CLI argument, not by default IMO.

[Viditnegi]

Okay got it. Thanks.

[linoytsaban]

for context, we do have this line - transformer.to(accelerator.device, dtype=weight_dtype) in the LoRA training scripts, and the flag --upcast_before_saving (controls wether to upcast to fp32 before saving) that is False by default indeed due to the memory reqs

[Viditnegi]

Okay got it, thanks!

[DN6]

Can we close this since it seems we can resolve the issue with a CLI arg?

[Viditnegi]

Can we close this since it seems we can resolve the issue with a CLI arg?

Yes, just closed the pr!