4-bit load_in_4bit + device_map='auto' + explicit max_memory still dispatches modules to CPU on a model that fits -> ValueError
Author: trevorgordon981
Pull request: no

**Env:** transformers 5.12.1, deepspeed 0.18.9, torch 2.12.0+cu130, peft 0.19.1, bitsandbytes 0.49.2, accelerate 1.14.0; 8x B200 (178GB) / 2TB RAM; model MiniMaxAI/MiniMax-M3 (428B sparse MoE VL, minimax_m3_vl).

Loading MiniMax-M3 in 4-bit (nf4, double-quant) with `device_map="auto"` and `max_memory={i: "165GiB" for i in range(8)}` (~1.3TB total) still computes a device_map with CPU entries, raising `ValueError: Some modules are dispatched on the CPU or the disk...` from `quantizer_bnb_4bit.validate_environment` (via `integrations/accelerate._get_device_map`). The model is far under the 1.3TB budget, yet auto-placement under-packs and spills to CPU (which 4-bit forbids).

**Ask:** respect the `max_memory` budget when packing 4-bit modules instead of conservatively spilling to CPU while GPU budget remains; or emit a clearer message + a documented recipe for forcing all modules onto GPUs.

Comments:
I dug into this and can reproduce it **CPU-only** — the placement is computed while the model is still on the `meta` device, so the failure is pure size arithmetic; the 8×B200 setup isn't needed to trigger it. A tiny meta model with scaled-down `max_memory` raises the same `ValueError` even with the total GPU budget at **6×** the model size.

Root cause, in transformers' vendored placement logic (`integrations/accelerate.py`):
1. `get_balanced_memory` caps every GPU except the last at `total/num_gpus + 1.25 × buffer`, but the buffer only measures `_no_split_modules` (a decoder layer) — not large childless modules such as `embed_tokens`/`lm_head`, which quantization does not shrink and which therefore dominate a 4-bit placement plan.
2. `infer_auto_device_map` walks devices monotonically: a module bigger than a device's cap skips that device *permanently* (no backfill).
3. So a single big leaf module — larger than the tiny balanced caps yet far below your explicit per-GPU budget — hops past GPUs 0..n-2, and everything after it piles onto the last GPU (the only one keeping the full user budget) until it's full, then spills to CPU. `validate_environment` then raises, with most GPUs literally empty. I suspect the same arithmetic bites MiniMax-M3-VL via its un-quantized embeddings/vision blocks vs caps of roughly `model_size/8`.

I have a fix working with regression tests: size `get_balanced_memory`'s buffer with the same largest-atom helper the packer itself uses for its reservation (`get_max_layer_size`), so per-device caps are never smaller than a placeable atom. My repro then places cleanly at 6×/3×/2× budgets (all previously raised), and a genuinely impossible budget (one atom larger than any single device's `max_memory`) now emits a clear warning instead of failing silently.

I'd like to take this — PR incoming shortly.

Hey @trevorgordon981! Thanks a lot for the report, made me realize that the no split module block size was computed by taking the first matching layer, but we may have layers of the same class with very very different size, as is the case here for minimax. So quite an important fix, thanks again for the report! https://github.com/huggingface/transformers/pull/47203

The scenario in this issue still reproduces on `main` after #47203 for large-vocab models. #47203 sizes the device_map buffer from the largest no-split block, but `embed_tokens`/`lm_head` (not in a no-split block, and not shrunk by 4-bit) can be a larger atom. Llama-3-8B in 4-bit still spills the ~2 GiB embedding to cpu and raises the same error even with 6x the model size in GPU budget. Opened #47211 with a one-line follow-up that also sizes the buffer from the largest leaf, plus CPU regression tests. cc @Cyrilvallez @ArthurZucker