From d0e6f01286ec8468c997a88cf7646c4edaa2f3c0 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 18 May 2025 04:39:48 +0300 Subject: [PATCH] IPEX remove GradScaler and use torch.amp instead --- modules/intel/ipex/__init__.py | 26 +---- modules/intel/ipex/gradscaler.py | 183 ------------------------------- modules/intel/ipex/hijacks.py | 56 +++++++++- 3 files changed, 55 insertions(+), 210 deletions(-) delete mode 100644 modules/intel/ipex/gradscaler.py diff --git a/modules/intel/ipex/__init__.py b/modules/intel/ipex/__init__.py index 2d82830c1..b0fcb655e 100644 --- a/modules/intel/ipex/__init__.py +++ b/modules/intel/ipex/__init__.py @@ -1,6 +1,5 @@ import os import sys -import contextlib import torch try: import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import @@ -21,7 +20,7 @@ def ipex_init(): # pylint: disable=too-many-statements try: # force xpu device on torch compile and triton # import inductor utils to get around lazy import - from torch._inductor import utils as torch_inductor_utils # pylint: disable=import-error, unused-import + from torch._inductor import utils as torch_inductor_utils # pylint: disable=import-error, unused-import # noqa: F401 torch._inductor.utils.GPU_TYPES = ["xpu"] torch._inductor.utils.get_gpu_type = lambda *args, **kwargs: "xpu" from triton import backends as triton_backends # pylint: disable=import-error @@ -161,29 +160,6 @@ def ipex_init(): # pylint: disable=too-many-statements torch.cuda.seed_all = torch.xpu.seed_all torch.cuda.initial_seed = torch.xpu.initial_seed - # AMP: - if has_ipex: - torch.xpu.amp.custom_fwd = torch.cuda.amp.custom_fwd - torch.xpu.amp.custom_bwd = torch.cuda.amp.custom_bwd - torch.cuda.amp = torch.xpu.amp - if torch_version < 2.3: - torch.is_autocast_enabled = torch.xpu.is_autocast_xpu_enabled - torch.get_autocast_gpu_dtype = torch.xpu.get_autocast_xpu_dtype - - if not hasattr(torch.cuda.amp, "common"): - torch.cuda.amp.common = contextlib.nullcontext() - torch.cuda.amp.common.amp_definitely_not_available = lambda: False - - try: - torch.cuda.amp.GradScaler = torch.xpu.amp.GradScaler - except Exception: # pylint: disable=broad-exception-caught - try: - from .gradscaler import gradscaler_init # pylint: disable=import-outside-toplevel, import-error - gradscaler_init() - torch.cuda.amp.GradScaler = torch.xpu.amp.GradScaler - except Exception: # pylint: disable=broad-exception-caught - torch.cuda.amp.GradScaler = ipex.cpu.autocast._grad_scaler.GradScaler - # C if torch_version < 2.3: torch._C._cuda_getCurrentRawStream = ipex._C._getCurrentRawStream diff --git a/modules/intel/ipex/gradscaler.py b/modules/intel/ipex/gradscaler.py deleted file mode 100644 index 0a8610095..000000000 --- a/modules/intel/ipex/gradscaler.py +++ /dev/null @@ -1,183 +0,0 @@ -from collections import defaultdict -import torch -import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import -import intel_extension_for_pytorch._C as core # pylint: disable=import-error, unused-import - -# pylint: disable=protected-access, missing-function-docstring, line-too-long - -device_supports_fp64 = torch.xpu.has_fp64_dtype() if hasattr(torch.xpu, "has_fp64_dtype") else torch.xpu.get_device_properties("xpu").has_fp64 -OptState = ipex.cpu.autocast._grad_scaler.OptState -_MultiDeviceReplicator = ipex.cpu.autocast._grad_scaler._MultiDeviceReplicator -_refresh_per_optimizer_state = ipex.cpu.autocast._grad_scaler._refresh_per_optimizer_state - -def _unscale_grads_(self, optimizer, inv_scale, found_inf, allow_fp16): # pylint: disable=unused-argument - per_device_inv_scale = _MultiDeviceReplicator(inv_scale) - per_device_found_inf = _MultiDeviceReplicator(found_inf) - - # To set up _amp_foreach_non_finite_check_and_unscale_, split grads by device and dtype. - # There could be hundreds of grads, so we'd like to iterate through them just once. - # However, we don't know their devices or dtypes in advance. - - # https://stackoverflow.com/questions/5029934/defaultdict-of-defaultdict - # Google says mypy struggles with defaultdicts type annotations. - per_device_and_dtype_grads = defaultdict(lambda: defaultdict(list)) # type: ignore[var-annotated] - # sync grad to master weight - if hasattr(optimizer, "sync_grad"): - optimizer.sync_grad() - with torch.no_grad(): - for group in optimizer.param_groups: - for param in group["params"]: - if param.grad is None: - continue - if (not allow_fp16) and param.grad.dtype == torch.float16: - raise ValueError("Attempting to unscale FP16 gradients.") - if param.grad.is_sparse: - # is_coalesced() == False means the sparse grad has values with duplicate indices. - # coalesce() deduplicates indices and adds all values that have the same index. - # For scaled fp16 values, there's a good chance coalescing will cause overflow, - # so we should check the coalesced _values(). - if param.grad.dtype is torch.float16: - param.grad = param.grad.coalesce() - to_unscale = param.grad._values() - else: - to_unscale = param.grad - - # -: is there a way to split by device and dtype without appending in the inner loop? - to_unscale = to_unscale.to("cpu") - per_device_and_dtype_grads[to_unscale.device][ - to_unscale.dtype - ].append(to_unscale) - - for _, per_dtype_grads in per_device_and_dtype_grads.items(): - for grads in per_dtype_grads.values(): - core._amp_foreach_non_finite_check_and_unscale_( - grads, - per_device_found_inf.get("cpu"), - per_device_inv_scale.get("cpu"), - ) - - return per_device_found_inf._per_device_tensors - -def unscale_(self, optimizer): - """ - Divides ("unscales") the optimizer's gradient tensors by the scale factor. - :meth:`unscale_` is optional, serving cases where you need to - :ref:`modify or inspect gradients` - between the backward pass(es) and :meth:`step`. - If :meth:`unscale_` is not called explicitly, gradients will be unscaled automatically during :meth:`step`. - Simple example, using :meth:`unscale_` to enable clipping of unscaled gradients:: - ... - scaler.scale(loss).backward() - scaler.unscale_(optimizer) - torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) - scaler.step(optimizer) - scaler.update() - Args: - optimizer (torch.optim.Optimizer): Optimizer that owns the gradients to be unscaled. - .. warning:: - :meth:`unscale_` should only be called once per optimizer per :meth:`step` call, - and only after all gradients for that optimizer's assigned parameters have been accumulated. - Calling :meth:`unscale_` twice for a given optimizer between each :meth:`step` triggers a RuntimeError. - .. warning:: - :meth:`unscale_` may unscale sparse gradients out of place, replacing the ``.grad`` attribute. - """ - if not self._enabled: - return - - self._check_scale_growth_tracker("unscale_") - - optimizer_state = self._per_optimizer_states[id(optimizer)] - - if optimizer_state["stage"] is OptState.UNSCALED: # pylint: disable=no-else-raise - raise RuntimeError( - "unscale_() has already been called on this optimizer since the last update()." - ) - elif optimizer_state["stage"] is OptState.STEPPED: - raise RuntimeError("unscale_() is being called after step().") - - # FP32 division can be imprecise for certain compile options, so we carry out the reciprocal in FP64. - assert self._scale is not None - if device_supports_fp64: - inv_scale = self._scale.double().reciprocal().float() - else: - inv_scale = self._scale.to("cpu").double().reciprocal().float().to(self._scale.device) - found_inf = torch.full( - (1,), 0.0, dtype=torch.float32, device=self._scale.device - ) - - optimizer_state["found_inf_per_device"] = self._unscale_grads_( - optimizer, inv_scale, found_inf, False - ) - optimizer_state["stage"] = OptState.UNSCALED - -def update(self, new_scale=None): - """ - Updates the scale factor. - If any optimizer steps were skipped the scale is multiplied by ``backoff_factor`` - to reduce it. If ``growth_interval`` unskipped iterations occurred consecutively, - the scale is multiplied by ``growth_factor`` to increase it. - Passing ``new_scale`` sets the new scale value manually. (``new_scale`` is not - used directly, it's used to fill GradScaler's internal scale tensor. So if - ``new_scale`` was a tensor, later in-place changes to that tensor will not further - affect the scale GradScaler uses internally.) - Args: - new_scale (float or :class:`torch.FloatTensor`, optional, default=None): New scale factor. - .. warning:: - :meth:`update` should only be called at the end of the iteration, after ``scaler.step(optimizer)`` has - been invoked for all optimizers used this iteration. - """ - if not self._enabled: - return - - _scale, _growth_tracker = self._check_scale_growth_tracker("update") - - if new_scale is not None: - # Accept a new user-defined scale. - if isinstance(new_scale, float): - self._scale.fill_(new_scale) # type: ignore[union-attr] - else: - reason = "new_scale should be a float or a 1-element torch.FloatTensor with requires_grad=False." - assert isinstance(new_scale, torch.FloatTensor), reason # type: ignore[attr-defined] - assert new_scale.numel() == 1, reason - assert new_scale.requires_grad is False, reason - self._scale.copy_(new_scale) # type: ignore[union-attr] - else: - # Consume shared inf/nan data collected from optimizers to update the scale. - # If all found_inf tensors are on the same device as self._scale, this operation is asynchronous. - found_infs = [ - found_inf.to(device="cpu", non_blocking=True) - for state in self._per_optimizer_states.values() - for found_inf in state["found_inf_per_device"].values() - ] - - assert len(found_infs) > 0, "No inf checks were recorded prior to update." - - found_inf_combined = found_infs[0] - if len(found_infs) > 1: - for i in range(1, len(found_infs)): - found_inf_combined += found_infs[i] - - to_device = _scale.device - _scale = _scale.to("cpu") - _growth_tracker = _growth_tracker.to("cpu") - - core._amp_update_scale_( - _scale, - _growth_tracker, - found_inf_combined, - self._growth_factor, - self._backoff_factor, - self._growth_interval, - ) - - _scale = _scale.to(to_device) - _growth_tracker = _growth_tracker.to(to_device) - # To prepare for next iteration, clear the data collected from optimizers this iteration. - self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state) - -def gradscaler_init(): - torch.xpu.amp.GradScaler = ipex.cpu.autocast._grad_scaler.GradScaler - torch.xpu.amp.GradScaler._unscale_grads_ = _unscale_grads_ - torch.xpu.amp.GradScaler.unscale_ = unscale_ - torch.xpu.amp.GradScaler.update = update - return torch.xpu.amp.GradScaler diff --git a/modules/intel/ipex/hijacks.py b/modules/intel/ipex/hijacks.py index e439365fb..db4d0e6f3 100644 --- a/modules/intel/ipex/hijacks.py +++ b/modules/intel/ipex/hijacks.py @@ -51,14 +51,42 @@ def return_xpu(device): # keep the device instance type, aka return string if th # Autocast original_autocast_init = torch.amp.autocast_mode.autocast.__init__ @wraps(torch.amp.autocast_mode.autocast.__init__) -def autocast_init(self, device_type, dtype=None, enabled=True, cache_enabled=None): - if device_type == "cuda" or device_type == "xpu": +def autocast_init(self, device_type=None, dtype=None, enabled=True, cache_enabled=None): + if device_type is None or check_cuda(device_type) or check_device_type(device_type, "xpu"): if dtype is None: dtype = devices.dtype return original_autocast_init(self, device_type="xpu", dtype=dtype, enabled=enabled, cache_enabled=cache_enabled) else: return original_autocast_init(self, device_type=device_type, dtype=dtype, enabled=enabled, cache_enabled=cache_enabled) + +original_grad_scaler_init = torch.amp.grad_scaler.GradScaler.__init__ +@wraps(torch.amp.grad_scaler.GradScaler.__init__) +def GradScaler_init(self, device: str = None, init_scale: float = 2.0**16, growth_factor: float = 2.0, backoff_factor: float = 0.5, growth_interval: int = 2000, enabled: bool = True): + if device is None or check_cuda(device): + return original_grad_scaler_init(self, device=return_xpu(device), init_scale=init_scale, growth_factor=growth_factor, backoff_factor=backoff_factor, growth_interval=growth_interval, enabled=enabled) + else: + return original_grad_scaler_init(self, device=device, init_scale=init_scale, growth_factor=growth_factor, backoff_factor=backoff_factor, growth_interval=growth_interval, enabled=enabled) + + +original_is_autocast_enabled = torch.is_autocast_enabled +@wraps(torch.is_autocast_enabled) +def torch_is_autocast_enabled(device_type=None): + if device_type is None or check_cuda(device_type): + return original_is_autocast_enabled(return_xpu(device_type)) + else: + return original_is_autocast_enabled(device_type) + + +original_get_autocast_dtype = torch.get_autocast_dtype +@wraps(torch.get_autocast_dtype) +def torch_get_autocast_dtype(device_type=None): + if device_type is None or check_cuda(device_type) or check_device_type(device_type, "xpu"): + return torch.bfloat16 + else: + return original_get_autocast_dtype(device_type) + + # Latent Antialias CPU Offload: # IPEX 2.5 and above has partial support but doesn't really work most of the time. original_interpolate = torch.nn.functional.interpolate @@ -404,4 +432,28 @@ def ipex_hijacks(): if not device_supports_fp64: torch.from_numpy = from_numpy torch.as_tensor = as_tensor + + # AMP: + torch.amp.grad_scaler.GradScaler.__init__ = GradScaler_init + torch.is_autocast_enabled = torch_is_autocast_enabled + torch.get_autocast_gpu_dtype = torch_get_autocast_dtype + torch.get_autocast_dtype = torch_get_autocast_dtype + + if hasattr(torch.xpu, "amp"): + if not hasattr(torch.xpu.amp, "custom_fwd"): + torch.xpu.amp.custom_fwd = torch.cuda.amp.custom_fwd + torch.xpu.amp.custom_bwd = torch.cuda.amp.custom_bwd + if not hasattr(torch.xpu.amp, "GradScaler"): + torch.xpu.amp.GradScaler = torch.amp.grad_scaler.GradScaler + torch.cuda.amp = torch.xpu.amp + else: + if not hasattr(torch.amp, "custom_fwd"): + torch.amp.custom_fwd = torch.cuda.amp.custom_fwd + torch.amp.custom_bwd = torch.cuda.amp.custom_bwd + torch.cuda.amp = torch.amp + + if not hasattr(torch.cuda.amp, "common"): + torch.cuda.amp.common = nullcontext() + torch.cuda.amp.common.amp_definitely_not_available = lambda: False + return device_supports_fp64, can_allocate_plus_4gb