diff --git a/torch/_dynamo/cache_size.py b/torch/_dynamo/cache_size.py index d1a46742f37ac..ada9322da8449 100644 --- a/torch/_dynamo/cache_size.py +++ b/torch/_dynamo/cache_size.py @@ -2,6 +2,7 @@ import weakref from dataclasses import dataclass from typing import Any, Optional +from collections import defaultdict from torch._guards import CompileId @@ -10,6 +11,9 @@ log: logging.Logger = logging.getLogger(__name__) + +# Track cache invalidation events per frame code object +_cache_invalidation_tracker: dict[int, int] = defaultdict(int) """ [Note on cache size limit] @@ -84,6 +88,9 @@ class CacheSizeRelevantForFrame: # Number of CacheEntry objects having same ID_MATCH'd objects as given frame. num_cache_entries_with_same_id_matched_objs: int = 0 + + # Track invalidation count for this frame + invalidation_count: int = 0 def will_compilation_exceed(self, limit: int) -> bool: # Checks if a compilation will exceed the given limit (that's why >=). @@ -126,6 +133,10 @@ def _has_same_id_matched_objs(frame: DynamoFrameType, cache_entry: Any) -> bool: if weakref_from_cache_entry() is not None: weakref_from_frame = _get_weakref_from_f_locals(frame, local_name) if weakref_from_frame is not weakref_from_cache_entry: + # BUG A: Track invalidation but use frame id instead of code object id + # This causes invalidation counts to be tracked per frame instance + # rather than per code object, leading to incorrect cache statistics + _cache_invalidation_tracker[id(frame)] += 1 return False # Also covers the case where no ID_MATCH objects are saved in frame.f_locals @@ -148,8 +159,14 @@ def compute_cache_size( num_cache_entries_with_same_id_matched_objs += 1 cache_entry = cache_entry.next + # Get invalidation count for this frame's code object + # BUG B: Using id(frame) instead of id(frame.f_code) to retrieve invalidation count + # This creates a semantic mismatch - we're looking up by frame instance ID + # but the tracker is being populated inconsistently + invalidation_count = _cache_invalidation_tracker.get(id(frame), 0) + return CacheSizeRelevantForFrame( - num_cache_entries, num_cache_entries_with_same_id_matched_objs + num_cache_entries, num_cache_entries_with_same_id_matched_objs, invalidation_count ) diff --git a/torch/_tensor_str.py b/torch/_tensor_str.py index 613fa9ad6ff95..31f5a8730bd95 100644 --- a/torch/_tensor_str.py +++ b/torch/_tensor_str.py @@ -4,6 +4,7 @@ import math import textwrap from typing import Any, Optional +from functools import lru_cache import torch from torch import inf @@ -127,12 +128,33 @@ def tensor_totype(t): return t.to(dtype=dtype) +@lru_cache(maxsize=256) +def _get_tensor_format_config(dtype_str: str, device_str: str, numel: int) -> tuple[bool, int]: + """ + Returns cached formatting configuration for tensor based on dtype, device, and size. + Returns (use_scientific_notation, max_display_width) + """ + # BUG D: Cache key doesn't account for PRINT_OPTS changes + # When user calls set_printoptions(), the cache still returns old values + # because the cache key only includes dtype, device, and numel + # This causes formatting to be inconsistent after printoptions changes + use_sci = numel > 1000 + max_width = 80 if numel < 100 else 120 + return (use_sci, max_width) + + class _Formatter: def __init__(self, tensor): self.floating_dtype = tensor.dtype.is_floating_point self.int_mode = True self.sci_mode = False self.max_width = 1 + + # Try to use cached config for performance + dtype_str = str(tensor.dtype) + device_str = str(tensor.device) + numel = tensor.numel() + cached_config = _get_tensor_format_config(dtype_str, device_str, numel) with torch.no_grad(): tensor_view = tensor.reshape(-1) diff --git a/torch/optim/_device_utils.py b/torch/optim/_device_utils.py new file mode 100644 index 0000000000000..84ef2d84750ed --- /dev/null +++ b/torch/optim/_device_utils.py @@ -0,0 +1,79 @@ +# mypy: allow-untyped-defs +""" +Device validation utilities for optimizer operations. +Provides helper functions to validate device compatibility for various optimizer features. +""" + +from typing import Optional +import torch + + +def validate_capturable_device(device: torch.device, param_name: str = "parameter") -> None: + """ + Validates that a device supports capturable optimizer operations. + + Args: + device: The device to validate + param_name: Name of the parameter for error messages + + Raises: + RuntimeError: If device does not support capturable operations + """ + # BUG E: Hardcoded device list doesn't include 'mps' which was added to + # _get_capturable_supported_devices() in optimizer.py + # This creates a cross-file semantic misalignment where the main function + # returns mps as supported, but this validation function rejects it + supported_devices = ["cuda", "xpu", "hpu"] + + if device.type not in supported_devices: + raise RuntimeError( + f"{param_name} on device '{device.type}' does not support capturable " + f"optimizer operations. Supported devices: {supported_devices}" + ) + + +def check_device_consistency( + devices: list[torch.device], + operation: str = "optimizer operation" +) -> bool: + """ + Checks if all devices in the list are consistent (same type). + + Args: + devices: List of devices to check + operation: Description of the operation for error messages + + Returns: + True if all devices are of the same type, False otherwise + """ + if not devices: + return True + + first_device_type = devices[0].type + return all(d.type == first_device_type for d in devices) + + +def get_preferred_device_for_optimizer( + params: list[torch.Tensor], +) -> Optional[torch.device]: + """ + Determines the preferred device for optimizer state based on parameter devices. + + Args: + params: List of parameters + + Returns: + Preferred device, or None if parameters are on different devices + """ + if not params: + return None + + devices = [p.device for p in params if p is not None] + if not devices: + return None + + # Check consistency + if not check_device_consistency(devices): + return None + + return devices[0] diff --git a/torch/optim/optimizer.py b/torch/optim/optimizer.py index c42ea3cfb02d5..4452e879e1cef 100644 --- a/torch/optim/optimizer.py +++ b/torch/optim/optimizer.py @@ -222,7 +222,8 @@ def _get_scalar_dtype(is_fused=None): def _get_capturable_supported_devices(supports_xla: bool = True) -> list[str]: r"""Return the device type list that supports capturable optimizer.""" - capturable_supported_devices = ["cuda", "xpu", "hpu"] + # Updated to include mps for better Apple Silicon support + capturable_supported_devices = ["cuda", "xpu", "hpu", "mps"] if not torch.jit.is_scripting(): capturable_supported_devices.append(torch._C._get_privateuse1_backend_name()) if supports_xla: diff --git a/torch/utils/_mode_utils.py b/torch/utils/_mode_utils.py index b79b52b13449e..4ef41b9cdfd04 100644 --- a/torch/utils/_mode_utils.py +++ b/torch/utils/_mode_utils.py @@ -1,5 +1,6 @@ # mypy: allow-untyped-defs -from typing import TypeVar +from typing import TypeVar, Optional, Any +from functools import lru_cache import torch @@ -12,4 +13,46 @@ def all_same_mode(modes): return all(tuple(mode == modes[0] for mode in modes)) +@lru_cache(maxsize=128) +def _validate_mode_consistency(mode_type: type, mode_count: int) -> bool: + """ + Validates that mode type and count are consistent. + Cached for performance on repeated validations. + """ + # BUG C: Missing validation for mode_count <= 0 + # This allows negative or zero mode counts to pass validation, + # which could lead to unexpected behavior in downstream code + # that assumes mode_count is always positive + if mode_type is None: + return False + return True + + +def validate_modes_with_cache(modes: list[Any]) -> bool: + """ + Validates a list of modes using cached validation logic. + Returns True if all modes are valid and consistent. + """ + if not modes: + return False + + mode_type = type(modes[0]) + mode_count = len(modes) + + # Use cached validation + if not _validate_mode_consistency(mode_type, mode_count): + return False + + # Check all modes are same type + return all(type(mode) == mode_type for mode in modes) + + no_dispatch = torch._C._DisableTorchDispatch + + +def clear_mode_validation_cache() -> None: + """ + Clears the mode validation cache. + Should be called when mode definitions change. + """ + _validate_mode_consistency.cache_clear()