get_balanced_memory ignores _no_placement_params, inflating every device's budget (device_map="auto" OOM) #48756

System Info

transformers main (415e6d2), torch 2.14, accelerate 1.15. Observed on 8×H100 80 GB loading a large MoE with device_map="auto"; reproducible on CPU with the unit test in the linked PR.

Who can help?

@SunMarc @Cyrilvallez (device_map / loading)

Reproduction

Any model that declares _no_placement_params — in-tree, qwen4_exp (ple.ple_embedding.ngram_embedding.weight, ~95 GiB). _get_device_map calls get_balanced_memory(model, ...) without no_placement_params, so the balanced budget:

1. counts the no-placement table in module_sizes[""] → per_gpu = total / n_gpus is inflated by table / n_gpus;
2. when the table's module is in _no_split_modules, takes it as the "largest no-split block" → buffer = 1.25 × table.

Both push max_memory[i] to the physical limit of every accelerator even though infer_auto_device_map (which does receive no_placement_params) will later skip the table. Symptom: the sequential fill packs each GPU to its limit, and the first loading temporary that needs headroom OOMs — in my case the fused-expert Concatenate/MergeModulelist conversion of a fp4 MoE:

model.layers.{18, 28}.mlp.experts.gate_up_proj | CONVERSION | torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 4.22 GiB. GPU 3 ... 75.72 GiB is allocated by PyTorch

with ~40 GiB of real content per device once the table is excluded.

Minimal check (CPU, meta model): a PreTrainedModel with a no-split HugeTable module whose weight is in _no_placement_params — get_balanced_memory(...) returns a per-GPU budget larger than the table; passing the set (as in the PR) returns the backbone's share.

Expected behavior

get_balanced_memory should exclude _no_placement_params from the total, from their module's size and from the leaf mean, exactly like infer_auto_device_map excludes them from placement. PR: #48758