From 2b5f270309439c244fc1f040df66fca07f9c9456 Mon Sep 17 00:00:00 2001
From: Harshita <harshitakrishna18@gmail.com>
Date: Wed, 26 Aug 2026 20:34:56 -0700
Subject: [PATCH] Fix memory accounting for integrated CUDA devices

---
 src/accelerate/utils/modeling.py |  7 +++++--
 tests/test_modeling_utils.py     | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 2ac8be12fed..6479fc242a8 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -762,6 +762,7 @@ def get_max_memory(max_memory: Optional[dict[Union[int, str], Union[int, str]]]
 
     if max_memory is None:
         max_memory = {}
+        is_integrated_cuda = False
         # Make sure device is initialized on each device to have the right memory info.
         if is_npu_available():
             for i in range(torch.npu.device_count()):
@@ -816,13 +817,15 @@ def get_max_memory(max_memory: Optional[dict[Union[int, str], Union[int, str]]]
                 try:
                     _ = torch.tensor([0], device=i)
                     max_memory[i] = torch.cuda.mem_get_info(i)[0]
+                    device_properties = torch.cuda.get_device_properties(i)
+                    is_integrated_cuda = is_integrated_cuda or getattr(device_properties, "is_integrated", False)
                 except Exception:
                     logger.info(f"Device {i} seems unavailable, Proceeding to check subsequent devices.")
                     continue
-        # allocate everything in the mps device as the RAM is shared
+        # MPS and integrated CUDA devices share host RAM, so exposing a separate CPU budget would double-count it.
         if is_mps_available():
             max_memory["mps"] = psutil.virtual_memory().available
-        else:
+        elif not is_integrated_cuda:
             max_memory["cpu"] = psutil.virtual_memory().available
         return max_memory
 
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
index a7aa6e4c590..63089fa8bc9 100644
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -18,7 +18,9 @@
 import unittest
 import warnings
 from collections import OrderedDict
+from types import SimpleNamespace
 from typing import Optional
+from unittest.mock import patch
 
 import torch
 import torch.nn as nn
@@ -34,6 +36,7 @@
     require_non_hpu,
     torch_device,
 )
+from accelerate.utils import modeling
 from accelerate.utils.modeling import (
     align_module_device,
     check_device_map,
@@ -44,6 +47,7 @@
     dtype_byte_size,
     find_tied_parameters,
     get_balanced_memory,
+    get_max_memory,
     get_module_size_with_ties,
     get_non_persistent_buffers,
     get_state_dict_offloaded_model,
@@ -135,6 +139,34 @@ def sequential_model(num_layers):
 
 
 class ModelingUtilsTester(unittest.TestCase):
+    @parameterized.expand(
+        [
+            (True, {0: 1234}),
+            (False, {0: 1234, "cpu": 4321}),
+            (None, {0: 1234, "cpu": 4321}),
+        ]
+    )
+    def test_get_max_memory_integrated_cuda(self, is_integrated, expected):
+        properties = SimpleNamespace()
+        if is_integrated is not None:
+            properties.is_integrated = is_integrated
+
+        with (
+            patch.object(modeling, "is_npu_available", return_value=False),
+            patch.object(modeling, "is_mlu_available", return_value=False),
+            patch.object(modeling, "is_sdaa_available", return_value=False),
+            patch.object(modeling, "is_musa_available", return_value=False),
+            patch.object(modeling, "is_xpu_available", return_value=False),
+            patch.object(modeling, "is_hpu_available", return_value=False),
+            patch.object(modeling, "is_mps_available", return_value=False),
+            patch.object(torch.cuda, "device_count", return_value=1),
+            patch.object(torch.cuda, "mem_get_info", return_value=(1234, 5678)),
+            patch.object(torch.cuda, "get_device_properties", return_value=properties),
+            patch.object(torch, "tensor"),
+            patch("psutil.virtual_memory", return_value=SimpleNamespace(available=4321)),
+        ):
+            self.assertEqual(get_max_memory(), expected)
+
     def test_dtype_byte_size(self):
         self.assertEqual(dtype_byte_size(torch.bool), 1 / 8)
         self.assertEqual(dtype_byte_size(torch.float16), 2)