get_max_memory() adds a separate cpu entry on integrated CUDA GPUs, double-counting shared RAM (GB10 / DGX Spark)

[Defilan]

Summary

get_max_memory() already treats shared memory correctly for MPS, but the same
reasoning is not applied to integrated CUDA devices, so on those the returned
budget counts the same physical RAM twice.

Current code (src/accelerate/utils/modeling.py, around L822):

# allocate everything in the mps device as the RAM is shared
if is_mps_available():
    max_memory["mps"] = psutil.virtual_memory().available
else:
    max_memory["cpu"] = psutil.virtual_memory().available

The comment states the principle exactly: when RAM is shared, do not hand out a
second pool. An NVIDIA GB10 (DGX Spark) shares its RAM the same way, but it is
a CUDA device, so it takes the else branch and gets a separate cpu entry
alongside its device entry.

Measured on a GB10:

>>> import torch
>>> from accelerate.utils import get_max_memory
>>> torch.cuda.get_device_properties(0).is_integrated
1
>>> get_max_memory()
{0: 73446760448, 'cpu': 87132405760}

73.4 GB + 87.1 GB = 160.5 GB advertised, on a machine with:

$ grep MemTotal /proc/meminfo
MemTotal:       127535336 kB        # 127.5 GB

There is no separate device VRAM on this part. torch.cuda.mem_get_info(0)
reports free=74.1 GB, total=130.6 GB, and those are the same bytes
/proc/meminfo describes, reached over NVLink-C2C rather than PCIe. The two
entries alias one pool: allocating against either shrinks the other, which the
returned dict cannot express.

Why it matters

Everything downstream inherits the number. infer_auto_device_map and
device_map="auto" treat the entries as independent capacities, so a model
sized between the device entry and the sum can be planned to "fit" across a
split whose halves are the same memory.

The hardware is not niche. GB10 / DGX Spark, GH200, GB200, Jetson Thor and Orin
are all integrated-memory parts, and they are specifically the machines people
buy to run large models locally.

Prior art in this file

• get_max_memory() returns allocated memory for XPU instead of total device memory #2929 (get_max_memory() returns allocated instead of total device
  memory for XPU) is the same class of bug: a device type measured with the
  wrong quantity.
• The MPS branch above is the fix already written, for a different vendor.

Suggested fix

Happy to send the PR. Two shapes, and I would rather write the one you want
than guess:

A. Extend the existing shared-RAM rule to integrated CUDA devices. Where
the MPS check happens, also treat an integrated CUDA device as sharing host
RAM, so no independent cpu entry is advertised against it. Smallest change,
and it reuses the reasoning already in the file.

B. Leave get_max_memory() alone and fix the consumer. Teach
infer_auto_device_map that an integrated device and the host are one pool, so
it will not plan a split between two aliases of the same memory.

I lean towards A, since it corrects the value everything else derives from.

Two implementation notes either way:

• is_integrated should be read defensively, e.g.
  getattr(torch.cuda.get_device_properties(i), "is_integrated", False).
  setup.py allows torch>=2.0.0 and I have only confirmed the attribute on
  torch 2.10, so a bare access risks breaking older installs.
• This cannot be exercised on your CI, since it needs integrated-memory
  hardware. A unit test would have to fake the device property (monkeypatch
  get_device_properties) and assert the shape of the returned dict. Happy to
  write it that way.

What I am not claiming

I found this while debugging a 4-bit load that failed under device_map="auto"
with "Some modules are dispatched on the CPU or the disk", which
device_map={"": 0} resolved. I could not reproduce that failure with plain
transformers on the same machine and model, so I am not asserting the double
count caused it. The arithmetic above stands on its own.

Environment

accelerate    1.12.0
transformers  5.5.0
torch         2.10.0a0+b558c986e8.nv25.11  (torch.version.cuda 13.0)
device        NVIDIA GB10, capability (12, 1), is_integrated=1
driver        580.173.02, CUDA runtime 13.0
platform      Linux-6.17.0-1029-nvidia-aarch64-with-glibc2.39, aarch64
python        3.12.3

[harshitakrishna18]

Hello, I’d like to work on this. I’m leaning toward option A as @Defilan suggested because it matches the existing MPS shared-memory handling and keeps the change localized, but I’ll first check downstream assumptions around max_memory["cpu"]. If maintainers prefer fixing the aliasing in infer_auto_device_map() instead, I’m happy to take option B.

[Defilan]

Please do, and thank you for picking it up.

A is the one I would write, and your instinct to check the downstream max_memory["cpu"] assumptions first is the right order of work. I think that check resolves more easily than it looks: the MPS branch already omits cpu entirely.

if is_mps_available():
    max_memory["mps"] = psutil.virtual_memory().available
else:
    max_memory["cpu"] = psutil.virtual_memory().available

On an MPS machine the returned dict has no cpu key at all, and that path has been shipping for a long time. So consumers already have to tolerate its absence, and A puts integrated CUDA on a code path that is exercised today rather than inventing a new shape. If you find a consumer that does assume cpu is present, that is a live bug on Apple silicon too and worth reporting separately.

The two implementation notes from the issue are the parts I would still be careful about:

Read the property defensively, getattr(torch.cuda.get_device_properties(i), "is_integrated", False). setup.py allows torch>=2.0.0 and I have only confirmed the attribute on 2.10, so a bare attribute access risks breaking older installs on hardware that has nothing to do with this.

And the test has to fake the device property, since CI has no integrated-memory hardware. Monkeypatching get_device_properties and asserting the shape of the returned dict is the only way to pin this in a way that runs anywhere.

Offer: I have three GB10s (DGX Spark) and am happy to run your branch on real hardware before it merges, including the device_map="auto" path that actually consumes the number. That is the part CI structurally cannot cover, and it is the part where being wrong is expensive: the current behaviour advertises 160.5 GB on a 127.5 GB machine, so a model planned into that gap fails at load rather than at planning time. Push a branch and tell me what you want measured.

No maintainer has weighed in on A versus B yet, so treat my preference as a reporter's opinion rather than a decision. If someone from the team prefers B, I will happily test that instead.

[harshitakrishna18]

Please do, and thank you for picking it up.

A is the one I would write, and your instinct to check the downstream max_memory["cpu"] assumptions first is the right order of work. I think that check resolves more easily than it looks: the MPS branch already omits cpu entirely.

if is_mps_available():
max_memory["mps"] = psutil.virtual_memory().available
else:
max_memory["cpu"] = psutil.virtual_memory().available
On an MPS machine the returned dict has no cpu key at all, and that path has been shipping for a long time. So consumers already have to tolerate its absence, and A puts integrated CUDA on a code path that is exercised today rather than inventing a new shape. If you find a consumer that does assume cpu is present, that is a live bug on Apple silicon too and worth reporting separately.

The two implementation notes from the issue are the parts I would still be careful about:

Read the property defensively, getattr(torch.cuda.get_device_properties(i), "is_integrated", False). setup.py allows torch>=2.0.0 and I have only confirmed the attribute on 2.10, so a bare attribute access risks breaking older installs on hardware that has nothing to do with this.

And the test has to fake the device property, since CI has no integrated-memory hardware. Monkeypatching get_device_properties and asserting the shape of the returned dict is the only way to pin this in a way that runs anywhere.

Offer: I have three GB10s (DGX Spark) and am happy to run your branch on real hardware before it merges, including the device_map="auto" path that actually consumes the number. That is the part CI structurally cannot cover, and it is the part where being wrong is expensive: the current behaviour advertises 160.5 GB on a 127.5 GB machine, so a model planned into that gap fails at load rather than at planning time. Push a branch and tell me what you want measured.

No maintainer has weighed in on A versus B yet, so treat my preference as a reporter's opinion rather than a decision. If someone from the team prefers B, I will happily test that instead.

That would be great.... Will drop a note when I push my changes. Thanks again for the offer that would be really helpful!!

[harshitakrishna18]

@Defilan #4187 -- this is my pull request. Would you try it out on your GB10s? Thanks in advance.