Also size the device_map buffer from the largest leaf module (follow-up to #47203)
Author: dhruv7477
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=47211)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47211)
<!-- ci-dashboard-badge:end -->

## What does this PR do?

Follow-up to #47203, which fixed the `device_map="auto"` under-packing from #46823 for same-class decoder layers of different sizes (MoE) by sizing the `get_balanced_memory` buffer from the largest no-split block.

A large childless leaf can be an even bigger un-splittable atom. `embed_tokens` and `lm_head` are not part of a no-split block, and are not shrunk by 4-bit quantization, so for large-vocab models they dominate. When such a leaf exceeds the balanced per-device caps, `infer_auto_device_map` skips every capped device, the model collapses onto the last device and then spills to cpu, and the 4-bit quantizer rejects the result.

### Reproduction (current `main`, meta device, no GPU needed)

Llama-3-8B loaded in 4-bit:

| module | size |
| --- | --- |
| `embed_tokens` / `lm_head` (fp16, not quantized) | ~2004 MiB each |
| one decoder layer (4-bit) | ~104 MiB |

The buffer is sized from the 104 MiB decoder layer, roughly 19x too small for the 2 GiB embedding, so `device_map="auto"` raises `Some modules are dispatched on the CPU or the disk...` even with 6x the model size in total GPU budget:

```
sum(GPU budget) = 6x model size  ->  ValueError (embedding spilled to cpu)
sum(GPU budget) = 4x model size  ->  ValueError
sum(GPU budget) = 3x model size  ->  ValueError
```

### Fix

Include the largest leaf module in the buffer, alongside the largest no-split block from #47203:

```python
buffer = int(1.25 * max(buffer, mean_leaves, max(leave_modules_sizes.values(), default=0)))
```

This keeps the #47203 behavior for the MoE case and additionally covers the large-leaf case. (`get_max_layer_size`, the helper the packer uses for its own largest-layer reservation, covers both in one call and subsumes the no-split-block max; I kept the change minimal and additive here, happy to switch to that if you prefer.)

### Tests

`tests/utils/test_device_map.py` runs on the meta device with fake GPU budgets, so no accelerator is required. Both tests fail on `main` and pass with this change.

## Who can review?

@Cyrilvallez @ArthurZucker

Comments:
<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47211)
**Latest run:** [29023881828](https://github.com/huggingface/transformers/actions/runs/29023881828)
**Result:** `success` | Grafana metrics are not available yet.

<!-- ci-dashboard-recap:end -->

cc @cyrilvallez I was going to close this as an agent PR, but after re-reading I think it's sort of real, but it only applies on models with massive embedding layers. Basically the error occurs when `embed_size > model_size/num_devices`, placement on every device fails and it all falls through to the last GPU or the CPU.

Glad the re-read landed. Your mechanism is exactly right, so I ran the numbers to put a boundary on it.

One refinement: "massive embedding layers" undersells the blast radius. 4-bit is what drags this into the mainstream. Quantization shrinks the decoder stack roughly 4x and leaves `embed_tokens`/`lm_head` in fp16, so on a perfectly ordinary model the embedding quietly becomes the biggest atom in the room. Llama-3-8B in 4-bit: ~2 GB of embeddings against ~104 MB per decoder layer. Plug that into your inequality and `embed > model/n` flips at n = 4. Verified on current `main` with generous budgets: clean placement on 2 and 3 GPUs, `ValueError` on 4 through 8. Which is precisely the `load_in_4bit` + `device_map="auto"` + `max_memory` combo from #46823.

The fix stays deliberately boring: one line, keeps the #47203 MoE handling intact, and the tests fail on `main` and pass with it.

Also, since agent PRs came up: the diagnosis, code, and every number above are mine and verified end to end on my machine. I do use AI to polish the write-ups because I'd rather they read professionally. If the polish is doing more harm than good here, say the word and you'll get my unpolished typing instead.