Skip to content

Fix n_devices max_memory caps#1501

Open
abhinav-bellapu wants to merge 2 commits into
TransformerLensOrg:devfrom
abhinav-bellapu:codex/fix-n-devices-max-memory
Open

Fix n_devices max_memory caps#1501
abhinav-bellapu wants to merge 2 commits into
TransformerLensOrg:devfrom
abhinav-bellapu:codex/fix-n-devices-max-memory

Conversation

@abhinav-bellapu

Copy link
Copy Markdown
Contributor

Fixes #1496.

Summary

  • Replace the invalid n_devices default max_memory={0: "auto", ...} with concrete CUDA byte budgets for devices 0..n_devices-1.
  • Keep user-supplied max_memory values unchanged.
  • Update Bridge type hints to reflect that HuggingFace accepts integer byte budgets as well as strings.
  • Add a mocked-CUDA regression test so the failing resolver path is covered on CPU CI.

Why

Accelerate rejects "auto" as a max_memory value with ValueError: size auto is not in a valid format. n_devices should continue to resolve to device_map="balanced", but the generated caps need to be concrete values accepted by HF/Accelerate.

Validation

  • make format
  • make check-format
  • uv run --python 3.11 mypy .
  • uv run --python 3.11 pytest tests/unit/utilities/test_multi_gpu_unit.py -q
  • uv run --python 3.11 pytest tests/acceptance/model_bridge/test_multi_gpu_bridge.py::TestResolveDeviceMap -q
  • uv run --python 3.11 pytest tests/acceptance/model_bridge/test_multi_gpu_bridge.py -q
  • make unit-test

@abhinav-bellapu abhinav-bellapu marked this pull request as ready for review July 10, 2026 20:03
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.

1 participant