Use device as context manager for init_on_device

[shingjan]

fixes #1814

cc: @sgugger

[HuggingFaceDocBuilderDev]

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

[sgugger]

Thanks for drafting this! I left a couple of comments, mainly we can't change the default of include_buffers which would break backward compatibility.

[BenjaminBossan]

Thanks for the PR. Is it possible to have a test for this or would it be too complicated to set up because of the version dependency?

[shingjan]

Thanks @sgugger and @BenjaminBossan for the reviews!

I added one test with include_buffers=True to make sure if testing with torch >= 2.0, we go with the device context manager. This PR however won't be fixing the issue I had before:

with tensor_mode():
  model = transformers.from_pretrained(...)

as init_empty_weights(...) uses include_buffers=False as default. I wonder if you have a better context on why we use include_buffers=False as default for init_empty_weights. If the ultimate goal is to move a nn.Module to meta device in order to save cpu/cuda memory, wouldn't buffer also be included in the state dict of a nn.Module and therefore better to move those to meta as well? Thanks again!

[sgugger]

We use include_buffers=False because buffers are usually not very heavy and it ends up taking more time to create them on the meta device and move them around than just creating them on the GPU.

[sgugger]

Nice new test!

[shingjan]

thanks @sgugger and @BenjaminBossan for getting this PR in!