Fix for offloading when using TorchAO >= 0.7.0

[a-r-r-o-w]

What does this PR do?

There was a breaking change made in TorchAO v0.7.0, where layout_tensor was made an internal private attribute and tensor_impl was added as another parameter.

This causes sequential offloading implementation to break in Diffusers: huggingface/diffusers#10470

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.

@SunMarc

[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]

Cc: @jerryzh168. Perhaps these things could be handled better?

[BenjaminBossan]

Thanks for providing this PR to handle the new argument name. I have a small comment to simplify the code.

Perhaps these things could be handled better?

I don't know the reason for the renaming, but I agree that ideally this can be avoided (PEFT also had an error because of that).

[SunMarc]

Thanks for the fix !

[BenjaminBossan]

Thanks, LGTM.

[a-r-r-o-w]

cc @muellerzr for merge :)

[jerryzh168]

Cc: @jerryzh168. Perhaps these things could be handled better?

I remember discussing with @bdhirsh before, and seems like no better way, but let me confirm with him

problem is: "how to we change the device for a nn.Parameter(tensor_subclass) object" right?

[nitinmukesh]

Gentle reminder to merge this, if complete.

[jerryzh168]

@sayakpaul I looked a bit, according to #3085 the issue is that AQT does not support requires_grad, would adding requires_grad arg be a better fix here?

[bdhirsh]

in the interest of making this easier to maintain - it seems like the most painful bit is that you are running one of AO's tensor subclass constructors, which does not have strong BC guarantees.

Is there a reason you need to call the constructor (to construct a new instance of this subclass) in the first place?

It sounds like the purpose of this logic is to take any parameters (which may be subclasses) that live on the wrong device, and in-place move them to the right device. So stepping back a bit:

(1) if you can, I would just call module.to(device) directly, which will save you from having to loop over and re-assign the params individually (hopefully this just works for you - if not, maybe we should discuss any problems with other core folks)

(2) if you need to reassign some params individually, can you do it without needing to reconstruct the subclass directly? I would probably do something like this:

# option 1:
new_value = torch.nn.Parameter(new_value.to(device=device), requires_grad=old_value.requires_grad)
# option 2:
new_value = torch.nn.Parameter(new_value, requires_grad=old_value.requires_grad).to(device=device)

[sayakpaul]

@SunMarc @BenjaminBossan

[SunMarc]

We don't really need it as you said. This was done for better compatibility with bitsandbytes as you can see in this PR https://github.com/huggingface/accelerate/pull/539/files. Would you like to check @a-r-r-o-w if the changes proposed by @bdhirsh works ?