Small fixes for balanced device maps

[sgugger]

This PR adds a couple of fixes to the balanced device map API, as using it on small models has given me a couple of errors.

• First the GPU 0 should also have the buffer (contrarily to the initial implementation) because it always saves memory for the largest layer inside infer_auto_device_map (so if we don't give it the buffer, we end up with way less on the GPU 0).
• Then we need to give extra room on the last GPU to avoid accidental CPU offload if one or two parameters end up without device at the end. The easiest way is just to leave it with the max memory.

[HuggingFaceDocBuilderDev]

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

[muellerzr]

LG2M, running CUDA tests first so lets make sure those pass