Skip to content

fix: add MPS device support #3763

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Aanushka001 wants to merge 2 commits into from
Closed

fix: add MPS device support #3763

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
80d82e3
Update auto_dataloader pin_memory logic and docstring to include MPS …
Aanushka001 May 31, 2026
733b7ad
Add empty_mps_cache function and MPS cache clearing handler
Aanushka001 May 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 24 additions & 13 deletions ignite/contrib/engines/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@
from ignite.metrics.metric import RunningBatchWise
from ignite.utils import deprecated

import gc


def setup_common_training_handlers(
trainer: Engine,
Expand Down Expand Up @@ -83,7 +85,7 @@ def setup_common_training_handlers(
save_every_iters: saving interval. By default, `to_save` objects are stored
each 1000 iterations.
output_path: output path to indicate where `to_save` objects are stored. Argument is mutually
exclusive with ``save_handler``.
exclusive with `save_handler`.
lr_scheduler: learning rate scheduler
as native torch LRScheduler or ignite's parameter scheduler.
with_gpu_stats: if True, :class:`~ignite.metrics.GpuInfo` is attached to the
Expand All @@ -100,8 +102,8 @@ def setup_common_training_handlers(
clear_cuda_cache: if True, `torch.cuda.empty_cache()` is called every end of epoch.
Default, True.
save_handler: Method or callable
class to use to store ``to_save``. See :class:`~ignite.handlers.checkpoint.Checkpoint` for more details.
Argument is mutually exclusive with ``output_path``.
class to use to store `to_save`. See :class:`~ignite.handlers.checkpoint.Checkpoint` for more details.
Argument is mutually exclusive with `output_path`.
kwargs: optional keyword args to be passed to construct :class:`~ignite.handlers.checkpoint.Checkpoint`.
"""

Expand Down Expand Up @@ -184,6 +186,8 @@ def _setup_common_training_handlers(

if torch.cuda.is_available() and clear_cuda_cache:
trainer.add_event_handler(Events.EPOCH_COMPLETED, empty_cuda_cache)
if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache") and clear_cuda_cache:
trainer.add_event_handler(Events.EPOCH_COMPLETED, empty_mps_cache)

if to_save is not None:
if output_path is None and save_handler is None:
Expand Down Expand Up @@ -276,10 +280,17 @@ def _setup_common_distrib_training_handlers(

@trainer.on(Events.EPOCH_STARTED)
def distrib_set_epoch(engine: Engine) -> None:
# pyrefly: ignore [missing-attribute]
# pyrely: ignore [missing-attribute]
train_sampler.set_epoch(engine.state.epoch - 1)


def empty_mps_cache(_: Engine) -> None:
if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"):
torch.mps.empty_cache()
import gc
gc.collect()


def empty_cuda_cache(_: Engine) -> None:
torch.cuda.empty_cache()
import gc
Expand Down Expand Up @@ -318,21 +329,21 @@ def gen_save_best_models_by_val_score(
score_sign: float = 1.0,
**kwargs: Any,
) -> Checkpoint:
"""Method adds a handler to ``evaluator`` to save ``n_saved`` of best models based on the metric
(named by ``metric_name``) provided by ``evaluator`` (i.e. ``evaluator.state.metrics[metric_name]``).
"""Method adds a handler to `evaluator` to save `n_saved` of best models based on the metric
(named by `metric_name`) provided by `evaluator` (i.e. `evaluator.state.metrics[metric_name]`).
Models with highest metric value will be retained. The logic of how to store objects is delegated to
``save_handler``.
`save_handler`.

Args:
save_handler: Method or callable class to
use to save engine and other provided objects. Function receives two objects: checkpoint as a dictionary
and filename. If ``save_handler`` is callable class, it can
inherit of :class:`~ignite.handlers.checkpoint.BaseSaveHandler` and optionally implement ``remove`` method
and filename. If `save_handler` is callable class, it can
inherit of :class:`~ignite.handlers.checkpoint.BaseSaveHandler` and optionally implement `remove` method
to keep a fixed number of saved checkpoints. In case if user needs to save engine's checkpoint on a disk,
``save_handler`` can be defined with :class:`~ignite.handlers.DiskSaver`.
`save_handler` can be defined with :class:`~ignite.handlers.DiskSaver`.
evaluator: evaluation engine used to provide the score
models: model or dictionary with the object to save. Objects should have
implemented ``state_dict`` and ``load_state_dict`` methods.
implemented `state_dict` and `load_state_dict` methods.
metric_name: metric name to use for score evaluation. This metric should be present in
`evaluator.state.metrics`.
n_saved: number of best models to store
Expand Down Expand Up @@ -380,8 +391,8 @@ def save_best_model_by_val_score(
score_sign: float = 1.0,
**kwargs: Any,
) -> Checkpoint:
"""Method adds a handler to ``evaluator`` to save on a disk ``n_saved`` of best models based on the metric
(named by ``metric_name``) provided by ``evaluator`` (i.e. ``evaluator.state.metrics[metric_name]``).
"""Method adds a handler to `evaluator` to save on a disk `n_saved` of best models based on the metric
(named by `metric_name`) provided by `evaluator` (i.e. `evaluator.state.metrics[metric_name]`).
Models with highest metric value will be retained.

Args:
Expand Down
29 changes: 14 additions & 15 deletions ignite/distributed/auto.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> DataLoader | _MpDeviceLo
- if no sampler provided by user, a `torch DistributedSampler`_ is setup.
- if a `torch DistributedSampler`_ is provided by user, it is used without wrapping it.
- if another sampler is provided, it is wrapped by :class:`~ignite.distributed.auto.DistributedProxySampler`.
- if the default device is 'cuda', `pin_memory` is automatically set to `True`.
- if the default device is 'cuda' or 'mps', ``pin_memory`` is automatically set to ``True``.

.. warning::

Expand All @@ -50,7 +50,7 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> DataLoader | _MpDeviceLo
Examples:
.. code-block:: python

import ignite.distribted as idist
import ignite.distributed as idist

train_loader = idist.auto_dataloader(
train_dataset,
Expand All @@ -64,8 +64,7 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> DataLoader | _MpDeviceLo
.. _torch DataLoader: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader
.. _XLA MpDeviceLoader:
https://pytorch.org/xla/release/2.0/index.html#running-on-multiple-xla-devices-with-multi-processing
.. _torch DistributedSampler:
https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
.. _torch DistributedSampler: https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
.. _torch IterableDataset: https://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset
"""
rank = idist.get_rank()
Expand All @@ -76,9 +75,9 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> DataLoader | _MpDeviceLo
if "batch_size" in kwargs and kwargs["batch_size"] >= world_size:
kwargs["batch_size"] //= world_size

nproc = idist.get_nproc_per_node()
if "num_workers" in kwargs and kwargs["num_workers"] >= nproc:
kwargs["num_workers"] = (kwargs["num_workers"] + nproc - 1) // nproc
nprocs = idist.get_nprocs_per_node()
if "num_workers" in kwargs and kwargs["num_workers"] >= nprocs:
kwargs["num_workers"] = (kwargs["num_workers"] + nprocs - 1) // nprocs

if "batch_sampler" not in kwargs:
if isinstance(dataset, IterableDataset):
Expand Down Expand Up @@ -118,7 +117,7 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> DataLoader | _MpDeviceLo
)
kwargs["pin_memory"] = False
else:
kwargs["pin_memory"] = kwargs.get("pin_memory", "cuda" in idist.device().type)
kwargs["pin_memory"] = kwargs.get("pin_memory", "cuda" in idist.device().type or "mps" in idist.device().type)

logger.info(f"Use data loader kwargs for dataset '{repr(dataset)[:20].strip()}': \n\t{kwargs}")
dataloader = DataLoader(dataset, **kwargs)
Expand Down Expand Up @@ -155,7 +154,7 @@ def auto_model(model: nn.Module, sync_bn: bool = False, **kwargs: Any) -> nn.Mod
Args:
model: model to adapt.
sync_bn: if True, applies `torch convert_sync_batchnorm`_ to the model for native torch
distributed only. Default, False. Note, if using Nvidia/Apex, batchnorm conversion should be
distributed only. Default, False. Note, if using Nvidia/Apex, batch norm conversion should be
applied before calling ``amp.initialize``.
kwargs: kwargs to model's wrapping class: `torch DistributedDataParallel`_ or `torch DataParallel`_
if applicable. Please, make sure to use acceptable kwargs for given backend.
Expand All @@ -166,15 +165,15 @@ def auto_model(model: nn.Module, sync_bn: bool = False, **kwargs: Any) -> nn.Mod
Examples:
.. code-block:: python

import ignite.distribted as idist
import ignite.distributed as idist

model = idist.auto_model(model)

In addition with NVidia/Apex, it can be used in the following way:
In addition with Nvidia/Apex, it can be used in the following way:

.. code-block:: python

import ignite.distribted as idist
import ignite.distributed as idist

model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
model = idist.auto_model(model)
Expand Down Expand Up @@ -242,7 +241,7 @@ def auto_optim(optimizer: Optimizer, **kwargs: Any) -> Optimizer:
Internally, this method is no-op for non-distributed and torch native distributed configuration.

For XLA distributed configuration, we create a new class that inherits from provided optimizer.
The goal is to override the `step()` method with specific `xm.optimizer_step`_ implementation.
The goal is to override the ``step()`` method with specific `xm.optimizer_step`_ implementation.

For Horovod distributed configuration, optimizer is wrapped with Horovod Distributed Optimizer and
its state is broadcasted from rank 0 to all other processes.
Expand Down Expand Up @@ -339,7 +338,7 @@ class _MpDeviceLoader:
# From pytorch/xla if `torch_xla.distributed.parallel_loader.MpDeviceLoader` is not available
def __init__(self, loader: Any, device: torch.device, **kwargs: Any) -> None:
self._loader = loader
# pyrefly: ignore [read-only]
# pyrely: ignore [read-only]
self._device = device
self._parallel_loader_kwargs = kwargs

Expand All @@ -356,4 +355,4 @@ def __init__(self, optimizer: Optimizer) -> None:
self.wrapped_optimizer = optimizer

def step(self, closure: Any = None) -> Any:
xm.optimizer_step(self.wrapped_optimizer, barrier=True)
xm.optimizer_step(self.wrapped_optimizer, barrier=True)