Skip to content

Fix memory accounting for integrated CUDA devices - #4187

Open
harshitakrishna18 wants to merge 1 commit into
huggingface:mainfrom
harshitakrishna18:fix-get-max-memory-integrated-cuda
Open

Fix memory accounting for integrated CUDA devices#4187
harshitakrishna18 wants to merge 1 commit into
huggingface:mainfrom
harshitakrishna18:fix-get-max-memory-integrated-cuda

Conversation

@harshitakrishna18

@harshitakrishna18 harshitakrishna18 commented Aug 27, 2026

Copy link
Copy Markdown

What does this PR do?

get_max_memory() currently reports CUDA memory and available CPU RAM as separate memory budgets. This is incorrect for integrated CUDA devices because both values refer to the same shared physical memory.

As a result, downstream code such as infer_auto_device_map() can treat the two values as independent capacities and plan a model placement that exceeds the machine's actual available memory.

This PR:

  • Checks torch.cuda.get_device_properties(i).is_integrated while discovering CUDA devices.
  • Omits the separate "cpu" memory budget when an integrated CUDA device is detected.
  • Uses getattr(..., False) so devices and older PyTorch versions without the is_integrated attribute retain the existing behavior.
  • Preserves the existing behavior for discrete CUDA devices.
  • Adds tests covering integrated devices, discrete devices, and device properties without the is_integrated attribute.

This follows the existing MPS behavior, where shared host and accelerator memory is exposed as a single allocation pool.

Fixes #4183

Testing

Added a parameterized unit test for get_max_memory() that runs without requiring CUDA hardware.

The test mocks CUDA device discovery, available device memory, device properties, CPU memory, and CUDA initialization. It covers:

  • An integrated CUDA device (is_integrated=True), verifying that no separate "cpu" memory budget is returned.
  • A discrete CUDA device (is_integrated=False), verifying that the existing CUDA and CPU budgets are preserved.
  • A device-properties object without is_integrated, verifying compatibility with older PyTorch versions through the getattr(..., False) fallback.

Results:

pytest tests/test_modeling_utils.py -k get_max_memory_integrated_cuda -q
3 passed, 44 deselected

pytest tests/test_modeling_utils.py -q
38 passed, 9 skipped, 4 subtests passed

The issue reproducer now returns:
{0: 1234}
instead of returning an additional independent "cpu" entry.
The reporter has also offered to validate the change on GB10/DGX Spark hardware, including the downstream device_map="auto" path.

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guideline,
    Pull Request section?
  • Was this discussed/approved via a Github issue or the forum? Please add a link
    to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant