Move io_same_device hook to before attach_align_device hook on cpu_offload and disk_offload.

[piEsposito]

Solves #767.

Moves the AlignDevicesHook with io_same_device from accelerate.cpu_offload and accelerate.disk_offload to before the attach_align_device_hook.

That way we can keep the changes on forward method for the whole module without deleting the hook we want to keep: the one with execution device and configurations on how to move the tensors between devices.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sgugger]

The fix is not exactly right: by doing so, the hook that ensures the input and output of the model are on the same device is now erased. In your code sample in #767, since x is on the CPU, net(x) should also be on the CPU. This is not the case with your PR. The solution would be to write a util function that will:

• just add the hook if none is present
• extract the current hook if one is present and chain it with this hook using a SequentialHook.

This is slightly more advanced than the current PR, so let me know if you'd prefer for me to do it :-)

[piEsposito]

@sgugger I would like to try if that's ok to you.

What do you think of creating an append_if_needed flag on add_hook_to_module that does what you just said?

[sgugger]

That works for me, though the name of the argument could simply be append :-)
Thanks for diving into this!

[piEsposito]

@sgugger append it is.

[piEsposito]

@sgugger it is ready for review, I've also added the tests.

[sgugger]

Very nice, thanks! Left a couple of nits, and I think you should still put the hook with io first: just tested locally and we still have the same issue of net(x) being on the wrong device since it runs second and the input was already moved.

[piEsposito]

Very nice, thanks! Left a couple of nits, and I think you should still put the hook with io first: just tested locally and we still have the same issue of net(x) being on the wrong device since it runs second and the input was already moved.

@sgugger I've just addressed your nits and moved the io hook to the top. Thanks for the review! I've tested it locally on the snippet of the bug report and it brings the tensor back to CPU after the inference. Thanks!

[sgugger]

Perfect, thanks! Re-tested locally and got the expected results for the code sample you shared in the issue.

[piEsposito]

@sgugger, there is a test step that failed due to an http error when installing a lib. I've created and empty commit to try running it again.