Fix TP initialization

[Cyrilvallez]

What does this PR do?

As per the title. When TP was introduced in #34184, it used a context manager with torch.device to set the model device. However, the context manager actually does NOT put the tensor on the device explicitly (it simply changes the default device if using name without indices, such as "cuda" or .cuda()). This means that the model was still on CPU.

Then parallelize_module would automatically switch the model, however it uses the SAME gpu for all child processes. So it's fine for small models, but will OOM for larger ones (e.g. Llama 8B on 4x L4 GPU (24 GB) would OOM, even though it's supposed to take ~16 GB only).

It does not look like we can instantiate only the model shards from the CPU to the different devices to avoid loading the full model on each device at the start.
To avoid that, I believe we could instantiate model on meta device, then recursively loading and dispatching each layer, which would avoif loading the full model in RAM then dispatching which goes against parallelization idea. Will try to explore it/see with accelerate team how we can do it.

[ArthurZucker]

NIce I remember needing a special version of torch for this?
Also let's add a test (the current test does not cover this case)

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