Fix for offloading when using TorchAO >= 0.7.0#3332
Conversation
|
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. |
|
Cc: @jerryzh168. Perhaps these things could be handled better? |
BenjaminBossan
left a comment
There was a problem hiding this comment.
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).
Co-Authored-By: Benjamin Bossan <BenjaminBossan@users.noreply.github.com> Co-Authored-By: Xuehai Pan <XuehaiPan@pku.edu.cn>
|
cc @muellerzr for merge :) |
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? |
|
Gentle reminder to merge this, if complete. |
|
@sayakpaul I looked a bit, according to #3085 the issue is that AQT does not support |
| new_value.quant_max, | ||
| new_value.zero_point_domain, | ||
| ) | ||
| new_value = torch.nn.Parameter(param_cls(*args), requires_grad=old_value.requires_grad).to(device) |
There was a problem hiding this comment.
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)
There was a problem hiding this comment.
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 ?
What does this PR do?
There was a breaking change made in TorchAO v0.7.0, where
layout_tensorwas made an internal private attribute andtensor_implwas 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