device count

Author: ydshieh

src/transformers/modeling_utils.py

@@ -4304,7 +4304,26 @@ def from_pretrained(
 
         # Prepare the full device map
         if device_map is not None:
+            import psutil
+            from transformers.testing_utils import _cpu_memory_patch_info
+            _vm = psutil.virtual_memory()
+            _patch_info = (
+                f", real_total={_cpu_memory_patch_info['real_total_gb']:.2f} GB"
+                f", real_available={_cpu_memory_patch_info['real_available_gb']:.2f} GB"
+                f", capped_to={_cpu_memory_patch_info['limit_bytes'] / 1024**3:.2f} GB"
+                if _cpu_memory_patch_info else ""
+            )
+            logger.warning(
+                f"[device_map] Before _get_device_map — "
+                f"device_map={device_map!r}, "
+                f"max_memory={max_memory!r}, "
+                f"CPU RAM total={_vm.total / 1024**3:.2f} GB (as seen by psutil){_patch_info}"
+            )
             device_map = _get_device_map(model, device_map, max_memory, hf_quantizer)
+            logger.warning(
+                f"[device_map] After _get_device_map — "
+                f"device_map={device_map!r}"
+            )
 
         # Finalize model weight initialization
         load_config = LoadStateDictConfig(

src/transformers/testing_utils.py

@@ -3430,6 +3430,10 @@ def patched(*args, **kwargs):
         torch.compile = patched
 
 
+_cpu_memory_patch_info: dict = {}
+"""Stores pre-patch memory stats set by patch_psutil_cpu_memory, for later logging during test execution."""
+
+
 def patch_psutil_cpu_memory(limit_bytes: int):
     """
     Patch `psutil.virtual_memory` to cap the reported CPU memory to `limit_bytes`.
@@ -3441,15 +3445,11 @@ def patch_psutil_cpu_memory(limit_bytes: int):
     """
     import psutil
 
-    mem_before = psutil.virtual_memory()
-    logger.warning(
-        f"[patch_psutil_cpu_memory] Before patching — "
-        f"total={mem_before.total / 1024**3:.2f} GB, "
-        f"available={mem_before.available / 1024**3:.2f} GB, "
-        f"limit_bytes={limit_bytes / 1024**3:.2f} GB"
-    )
-
     _original_virtual_memory = psutil.virtual_memory
+    _real_mem = _original_virtual_memory()
+    _cpu_memory_patch_info["limit_bytes"] = limit_bytes
+    _cpu_memory_patch_info["real_total_gb"] = _real_mem.total / 1024**3
+    _cpu_memory_patch_info["real_available_gb"] = _real_mem.available / 1024**3
 
     def _capped_virtual_memory():
         mem = _original_virtual_memory()
@@ -3461,13 +3461,6 @@ def _capped_virtual_memory():
 
     psutil.virtual_memory = _capped_virtual_memory
 
-    mem_after = psutil.virtual_memory()
-    logger.warning(
-        f"[patch_psutil_cpu_memory] After patching — "
-        f"total={mem_after.total / 1024**3:.2f} GB, "
-        f"available={mem_after.available / 1024**3:.2f} GB"
-    )
-
 
 def _get_test_info():
     """