Repository: pytorch/ignite

Issue:
Exponential moving average (EMA) handler
## 🚀 Feature
As discussed in the MONAI's issue thread, it would be great if EMA handler can be added to the ignite's handlers family.
EMA aims to keep an exponential moving averaged version of model weights. This smoothed version of weights are used for validation and inference.

Context: https://github.com/Project-MONAI/MONAI/issues/2514




Discussion:

Resolving patch:
diff --git a/ignite/handlers/__init__.py b/ignite/handlers/__init__.py
--- a/ignite/handlers/__init__.py
+++ b/ignite/handlers/__init__.py
@@ -4,6 +4,7 @@
 from ignite.engine.events import Events
 from ignite.handlers.checkpoint import Checkpoint, DiskSaver, ModelCheckpoint
 from ignite.handlers.early_stopping import EarlyStopping
+from ignite.handlers.ema_handler import EMAHandler
 from ignite.handlers.lr_finder import FastaiLRFinder
 from ignite.handlers.param_scheduler import (
     ConcatScheduler,
@@ -41,6 +42,7 @@
     "CyclicalScheduler",
     "create_lr_scheduler_with_warmup",
     "FastaiLRFinder",
+    "EMAHandler",
 ]
 
 
diff --git a/ignite/handlers/ema_handler.py b/ignite/handlers/ema_handler.py
new file mode 100644
--- /dev/null
+++ b/ignite/handlers/ema_handler.py
@@ -0,0 +1,214 @@
+from copy import deepcopy
+from typing import Optional, Union
+
+import torch.nn as nn
+
+from ignite.engine import CallableEventWithFilter, Engine, Events, EventsList
+
+__all__ = ["EMAHandler"]
+
+
+class EMAHandler:
+    r"""Exponential moving average (EMA) handler can be used to compute a smoothed version of model.
+    The EMA model is updated as follows:
+
+    .. math:: \theta_{\text{EMA}, t+1} = (1 - \lambda) \cdot \theta_{\text{EMA}, t} + \lambda \cdot \theta_{t}
+
+    where :math:`\theta_{\text{EMA}, t}` and :math:`\theta_{t}` are the EMA weights and online model weights at
+    :math:`t`-th iteration, respectively; :math:`\lambda` is the update momentum. The handler allows for linearly
+    warming up the momentum in the beginning when training process is not stable. Current momentum can be retrieved
+    from ``Engine.state.ema_momentum``.
+
+    Args:
+          model: the online model for which an EMA model will be computed. If ``model`` is ``DataParallel`` or
+              ``DistributedDataParallel``, the EMA smoothing will be applied to ``model.module`` .
+          momentum: the update momentum after warmup phase, should be float in range :math:`\left(0, 1 \right)`.
+          momentum_warmup: the initial update momentum during warmup phase, the value should be smaller than
+              ``momentum``. Momentum will linearly increase from this value to ``momentum`` in ``warmup_iters``
+              iterations. If ``None``, no warmup will be performed.
+          warmup_iters: iterations of warmup. If ``None``, no warmup will be performed.
+
+    Attributes:
+          ema_model: the exponential moving averaged model.
+          model: the online model that is tracked by EMAHandler. It is ``model.module`` if ``model`` in
+              the initialization method is an instance of ``DistributedDataParallel``.
+          momentum: the update momentum after warmup phase.
+          momentum_warmup: the initial update momentum.
+          warmup_iters: number of warmup iterations.
+
+    Note:
+          The EMA model is already in ``eval`` mode. If model in the arguments is an ``nn.Module`` or
+          ``DistributedDataParallel``, the EMA model is an ``nn.Module`` and it is on the same device as the online
+          model. If the model is an ``nn.DataParallel``, then the EMA model is an ``nn.DataParallel``.
+
+
+    Note:
+          It is recommended to initialize and use an EMA handler in following steps:
+
+          1. Initialize ``model`` (``nn.Module`` or ``DistributedDataParallel``) and ``ema_handler`` (``EMAHandler``).
+          2. Build ``trainer`` (``ignite.engine.Engine``).
+          3. Resume from checkpoint for ``model`` and ``ema_handler.ema_model``.
+          4. Attach ``ema_handler`` to ``trainer``.
+
+    Examples:
+          .. code-block:: python
+
+              device = torch.device("cuda:0")
+              model = nn.Linear(2, 1).to(device)
+              # update the ema every 5 iterations
+              ema_handler = EMAHandler(
+                  model, momentum=0.0002, momentum_warmup=0.0001, warmup_iters=10000)
+              # get the ema model, which is an instance of nn.Module
+              ema_model = ema_handler.ema_model
+              trainer = Engine(train_step_fn)
+              to_load = {"model": model, "ema_model", ema_model, "trainer", trainer}
+              if resume_from is not None:
+                  Checkpoint.load_objects(to_load, checkpoint=resume_from)
+
+              # update the EMA model every 5 iterations
+              ema_handler.attach(trainer, name="ema_momentum", event=Events.ITERATION_COMPLETED(every=5))
+
+              # add other handlers
+              to_save = to_load
+              ckpt_handler = Checkpoint(to_save, DiskSaver(...), ...)
+              trainer.add_event_handler(Events.EPOCH_COMPLETED, ckpt_handler)
+
+              # current momentum can be retrieved from engine.state,
+              # the attribute name is the `name` parameter used in the attach function
+              @trainer.on(Events.ITERATION_COMPLETED):
+              def print_ema_momentum(engine):
+                  print(f"current momentum: {engine.state.ema_momentum}"
+
+              # use ema model for validation
+              val_step_fn = get_val_step_fn(ema_model)
+              evaluator = Engine(val_step_fn)
+
+              @trainer.on(Events.EPOCH_COMPLETED)
+              def run_validation(engine):
+                  engine.run(val_data_loader)
+
+              trainer.run(...)
+
+          The following example shows how to attach two handlers to the same trainer:
+
+          .. code-block:: python
+
+              generator = build_generator(...)
+              discriminator = build_discriminator(...)
+
+              gen_handler = EMAHandler(generator)
+              disc_handler = EMAHandler(discriminator)
+
+              step_fn = get_step_fn(...)
+              engine = Engine(step_fn)
+
+              # update EMA model of generator every 1 iteration
+              gen_handler.attach(engine, "gen_ema_momentum", event=Events.ITERATION_COMPLETED)
+              # update EMA model of discriminator every 2 iteration
+              disc_handler.attach(engine, "dis_ema_momentum", event=Events.ITERATION_COMPLETED(every=2))
+
+              @engine.on(Events.ITERATION_COMPLETED)
+              def print_ema_momentum(engine):
+                  print(f"current momentum for generator: {engine.state.gen_ema_momentum}")
+                  print(f"current momentum for discriminator: {engine.state.disc_ema_momentum}")
+
+              engine.run(...)
+
+    .. versionadded:: 0.5.0
+
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        momentum: float = 0.0002,
+        momentum_warmup: Optional[float] = None,
+        warmup_iters: Optional[int] = None,
+    ) -> None:
+        if momentum_warmup is not None and not 0 < momentum_warmup < 1:
+            raise ValueError(f"Invalid momentum_warmup: {momentum_warmup}")
+        if not 0 < momentum < 1:
+            raise ValueError(f"Invalid momentum: {momentum}")
+        if momentum_warmup is not None and not momentum_warmup <= momentum:
+            raise ValueError(
+                f"momentum_warmup should be less than or equal to momentum, but got "
+                f"momentum_warmup: {momentum_warmup} and momentum: {momentum}"
+            )
+        if warmup_iters is not None and not (isinstance(warmup_iters, int) and warmup_iters > 0):
+            raise ValueError(f"Invalid warmup_iters: {warmup_iters}")
+        if not isinstance(model, nn.Module):
+            raise ValueError(
+                f"model should be an instance of nn.Module or its subclasses, but got"
+                f"model: {model.__class__.__name__}"
+            )
+        self.momentum_warmup = momentum_warmup
+        self.momentum = momentum
+        self.warmup_iters = warmup_iters
+
+        if isinstance(model, nn.parallel.DistributedDataParallel):
+            model = model.module
+        self.model = model
+
+        self.ema_model = deepcopy(self.model)
+        for param in self.ema_model.parameters():
+            param.detach_()
+        self.ema_model.eval()
+
+    def _get_momentum(self, curr_iter: int) -> float:
+        """Get current momentum, `curr_iter` should be 1-based. When `curr_iter = 1`, `momentum =
+        self.momentum_warmup`; when `curr_iter >= self.warmup_iters`, `momentum = self.momentum`"""
+
+        # TODO: use ignite's parameter scheduling, see also GitHub issue #2090
+        if curr_iter < 1:
+            raise ValueError(f"curr_iter should be at least 1, but got {curr_iter}.")
+
+        # no warmup
+        if self.momentum_warmup is None or self.warmup_iters is None:
+            return self.momentum
+
+        denominator = max(1, self.warmup_iters - 1)
+        momentum = self.momentum_warmup + (self.momentum - self.momentum_warmup) * (curr_iter - 1) / denominator
+        return min(self.momentum, momentum)
+
+    def _update_ema_model(self, engine: Engine, name: str) -> None:
+        """Update weights of ema model"""
+        momentum = getattr(engine.state, name)
+        for ema_p, model_p in zip(self.ema_model.parameters(), self.model.parameters()):
+            ema_p.mul_(1.0 - momentum).add_(model_p.data, alpha=momentum)
+        # assign the buffers
+        for ema_b, model_b in zip(self.ema_model.buffers(), self.model.buffers()):
+            ema_b.data = model_b.data
+
+    def _update_ema_momentum(self, engine: Engine, name: str) -> None:
+        """Update momentum in engine.state"""
+        curr_iter = engine.state.iteration
+        momentum = self._get_momentum(curr_iter)
+        setattr(engine.state, name, momentum)
+
+    def attach(
+        self,
+        engine: Engine,
+        name: str = "ema_momentum",
+        event: Union[str, Events, CallableEventWithFilter, EventsList] = Events.ITERATION_COMPLETED,
+    ) -> None:
+        """Attach the handler to engine. After the handler is attached, the ``Engine.state`` will add an new attribute
+        with name ``name``. Then, current momentum can be retrieved by from ``Engine.state`` when the engine runs.
+
+        Args:
+            engine: trainer to which the handler will be attached.
+            name: attribute name for retrieving EMA momentum from ``Engine.state``. It should be a unique name since a
+                trainer can have multiple EMA handlers.
+            event: event when the EMA momentum and EMA model are updated.
+
+        """
+        if hasattr(engine.state, name):
+            raise ValueError(
+                f"Attribute: '{name}' is already in Engine.state. Thus it might be "
+                f"overridden by other EMA handlers. Please select another name."
+            )
+
+        setattr(engine.state, name, 0.0)
+
+        # first update momentum, then update ema model
+        engine.add_event_handler(event, self._update_ema_momentum, name)
+        engine.add_event_handler(event, self._update_ema_model, name)


Test patch:
diff --git a/tests/ignite/handlers/test_ema_handler.py b/tests/ignite/handlers/test_ema_handler.py
new file mode 100644
--- /dev/null
+++ b/tests/ignite/handlers/test_ema_handler.py
@@ -0,0 +1,319 @@
+import os
+from typing import Any, Callable, Union
+
+import pytest
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DataParallel, DistributedDataParallel
+
+import ignite.distributed as idist
+from ignite.engine import Engine, Events
+from ignite.handlers import EMAHandler
+
+
+def _get_dummy_models() -> nn.Module:
+    model = nn.Linear(2, 1, bias=False)
+    model.weight.data.fill_(1)
+    return model
+
+
+@pytest.fixture(scope="module")
+def get_dummy_model():
+    """Returns a function since the fixture is needed multiple times in a single test"""
+    yield _get_dummy_models
+
+
+def _get_dummy_step_fn(model: Union[nn.Module, DataParallel, DistributedDataParallel]) -> Callable:
+    """Get a dummy step function, given model is a (wrapper of) dummy model returned from _get_dummy_model"""
+
+    def step_fn(engine, batch):
+        """Increment the weight by 1 at each iteration"""
+        model.weight.data.add_(1)
+        return 0
+
+    return step_fn
+
+
+@pytest.mark.parametrize("momentum", [-1, 2])
+def test_ema_invalid_momentum(get_dummy_model, momentum):
+    with pytest.raises(ValueError, match="Invalid momentum"):
+        EMAHandler(get_dummy_model(), momentum=momentum)
+
+
+@pytest.mark.parametrize("momentum_warmup", [-1, 2])
+def test_ema_invalid_momentum_warmup(get_dummy_model, momentum_warmup):
+    with pytest.raises(ValueError, match="Invalid momentum_warmup"):
+        EMAHandler(get_dummy_model, momentum_warmup=momentum_warmup)
+
+
+def test_ema_invalid_momentum_start_end(get_dummy_model):
+    """Test momentum_end > momentum_start"""
+    momentum = 0.001
+    momentum_warmup = 0.1
+    with pytest.raises(ValueError, match="momentum_warmup should be less than or equal to momentum"):
+        EMAHandler(get_dummy_model(), momentum_warmup=momentum_warmup, momentum=momentum)
+
+
+def test_ema_invalid_model():
+    with pytest.raises(ValueError, match="model should be an instance of nn.Module or its subclasses"):
+        model = "Invalid Model"
+        EMAHandler(model)  # type: ignore
+
+
+@pytest.mark.distributed
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
+def test_ema_ema_model_on_cuda(get_dummy_model):
+    """Test if ema_handler.ema_model is nn.Module and under eval mode"""
+    model = get_dummy_model().to(idist.device())
+    model = idist.auto_model(model)
+    ema_handler = EMAHandler(model)
+    ema_model = ema_handler.ema_model
+    assert (
+        isinstance(ema_model, nn.Module)
+        and not isinstance(ema_model, nn.parallel.DistributedDataParallel)
+        and not isinstance(ema_model, nn.parallel.DataParallel)
+    )
+    assert not ema_model.training
+
+
+def test_ema_load_state_dict(get_dummy_model):
+    model_1 = get_dummy_model()
+    model_1.weight.data.fill_(2)
+    state_dict_1 = model_1.state_dict()
+
+    model_2 = get_dummy_model()
+    ema_handler = EMAHandler(model_2)
+    ema_model = ema_handler.ema_model
+    ema_model.load_state_dict(state_dict_1)
+    torch.testing.assert_allclose(ema_model.weight.data, model_1.weight.data)
+
+
+def test_ema_no_warmup_momentum(get_dummy_model):
+    model = get_dummy_model()
+    step_fn = _get_dummy_step_fn(model)
+    engine = Engine(step_fn)
+
+    def assert_const_momentum(engine: Engine, const_momentum):
+        assert engine.state.ema_momentum == const_momentum
+
+    # no momentum_warmup
+    ema_handler = EMAHandler(model, momentum=0.002, momentum_warmup=None, warmup_iters=1)
+    ema_handler.attach(engine)
+    # attach the assertion handler after ema_handler, so the momentum is first updated and then tested
+    engine.add_event_handler(Events.ITERATION_COMPLETED, assert_const_momentum, ema_handler.momentum)
+    engine.run(range(2))
+
+    # no warmup_iters
+    engine = Engine(step_fn)
+    ema_handler = EMAHandler(model, momentum=0.002, momentum_warmup=0.001, warmup_iters=None)
+    ema_handler.attach(engine)
+    # attach the assertion handler after ema_handler, so the momentum is first updated and then tested
+    engine.add_event_handler(Events.ITERATION_COMPLETED, assert_const_momentum, ema_handler.momentum)
+    engine.run(range(2))
+
+
+def test_ema_update_ema_momentum(get_dummy_model):
+    model = get_dummy_model()
+    step_fn = _get_dummy_step_fn(model)
+    engine = Engine(step_fn)
+
+    warmup_iters = 4
+    momentum_warmup = 0.1
+    momentum = 0.2
+    ema_handler = EMAHandler(model, momentum_warmup=momentum_warmup, momentum=momentum, warmup_iters=warmup_iters)
+    ema_handler.attach(engine)
+
+    # add handlers to check momentum at each iteration
+    @engine.on(Events.ITERATION_COMPLETED)
+    def assert_momentum(engine: Engine):
+        curr_iter = engine.state.iteration
+        curr_momentum = engine.state.ema_momentum
+        if curr_iter == 1:
+            assert curr_momentum == momentum_warmup
+        elif 1 < curr_iter < warmup_iters:
+            assert momentum_warmup < curr_momentum < momentum
+        else:
+            assert curr_momentum == momentum
+
+    engine.run(range(2), max_epochs=5)
+
+
+def test_ema_buffer():
+    """Test if the tensors in buffer are also synchronized"""
+    model = nn.BatchNorm2d(2)
+    model.running_mean.data.fill_(1.5)
+    model.running_var.data.fill_(1.5)
+    ema_handler = EMAHandler(model)
+
+    def _bn_step_fn(engine, batch):
+        x = torch.rand(4, 2, 32, 32)
+        _ = model(x)
+        return 1
+
+    engine = Engine(_bn_step_fn)
+    ema_handler.attach(engine)
+
+    ema_model = ema_handler.ema_model
+
+    @engine.on(Events.ITERATION_COMPLETED)
+    def check_buffers():
+        torch.testing.assert_allclose(ema_model.running_mean, model.running_mean)
+        torch.testing.assert_allclose(ema_model.running_var, model.running_var)
+
+    # engine will run 4 iterations
+    engine.run([0, 1], max_epochs=2)
+
+    torch.testing.assert_allclose(ema_model.running_mean, model.running_mean)
+    torch.testing.assert_allclose(ema_model.running_var, model.running_var)
+
+
+def test_ema_two_handlers(get_dummy_model):
+    """Test when two EMA handlers are attached to a trainer"""
+    model_1 = get_dummy_model()
+    # momentum will be constantly 0.5
+    ema_handler_1 = EMAHandler(model_1, momentum_warmup=0.5, momentum=0.5, warmup_iters=1)
+
+    model_2 = get_dummy_model()
+    ema_handler_2 = EMAHandler(model_2, momentum_warmup=0.5, momentum=0.5, warmup_iters=1)
+
+    def _step_fn(engine: Engine, batch: Any):
+        model_1.weight.data.add_(1)
+        model_2.weight.data.add_(1)
+        return 0
+
+    engine = Engine(_step_fn)
+    assert not hasattr(engine.state, "ema_momentum_1")
+    # handler_1 update EMA model of model_1 every 1 iteration
+    ema_handler_1.attach(engine, "ema_momentum_1", event=Events.ITERATION_COMPLETED)
+    assert hasattr(engine.state, "ema_momentum_1")
+
+    # handler_2 update EMA model for model_2 every 2 iterations
+    ema_handler_2.attach(engine, "ema_momentum_2", event=Events.ITERATION_COMPLETED(every=2))
+    assert hasattr(engine.state, "ema_momentum_2")
+
+    # engine will run 4 iterations
+    engine.run(range(2), max_epochs=2)
+    ema_weight_1 = ema_handler_1.ema_model.weight.data
+    ema_weight_2 = ema_handler_2.ema_model.weight.data
+    torch.testing.assert_allclose(ema_weight_1, torch.full((1, 2), 4.0625))
+    torch.testing.assert_allclose(ema_weight_2, torch.full((1, 2), 3.5))
+
+    assert engine.state.ema_momentum_1 == 0.5
+    assert engine.state.ema_momentum_2 == 0.5
+
+    model_3 = get_dummy_model()
+    ema_handler_3 = EMAHandler(model_3)
+    with pytest.raises(ValueError, match="Please select another name"):
+        ema_handler_3.attach(engine, "ema_momentum_2")
+
+
+def _test_ema_final_weight(model, device, ddp=False, interval=1):
+    """Test if final smoothed weights are correct"""
+    if isinstance(device, str):
+        device = torch.device(device)
+    model = model.to(device)
+    if ddp:
+        model = idist.auto_model(model)
+    step_fn = _get_dummy_step_fn(model)
+    engine = Engine(step_fn)
+
+    # momentum will be constantly 0.5
+    ema_handler = EMAHandler(model, momentum_warmup=0.5, momentum=0.5, warmup_iters=1)
+    ema_handler.attach(engine, "model", event=Events.ITERATION_COMPLETED(every=interval))
+
+    # engine will run 4 iterations
+    engine.run(range(2), max_epochs=2)
+
+    ema_weight = ema_handler.ema_model.weight.data
+    model_weight = model.weight.data
+    assert ema_weight.device == device
+    assert model_weight.device == device
+    if interval == 1:
+        torch.testing.assert_allclose(ema_weight, torch.full((1, 2), 4.0625, device=device))
+    elif interval == 2:
+        torch.testing.assert_allclose(ema_weight, torch.full((1, 2), 3.5, device=device))
+    else:
+        pass
+    torch.testing.assert_allclose(model_weight, torch.full((1, 2), 5.0, device=device))
+
+
+@pytest.mark.parametrize("interval", [1, 2])
+def test_ema_final_weight_cpu(get_dummy_model, interval):
+    device = torch.device("cpu")
+    _test_ema_final_weight(get_dummy_model(), device=device, ddp=False, interval=interval)
+
+
+@pytest.mark.parametrize("interval", [1, 2])
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
+def test_ema_final_weight_cuda(get_dummy_model, interval):
+    device = torch.device("cuda:0")
+    _test_ema_final_weight(get_dummy_model(), device=device, ddp=False, interval=interval)
+
+
+@pytest.mark.distributed
+@pytest.mark.parametrize("interval", [1, 2])
+@pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
+def test_ema_final_weight_nccl_cuda(get_dummy_model, distributed_context_single_node_nccl, interval):
+    device = idist.device()
+    _test_ema_final_weight(get_dummy_model(), device=device, ddp=True, interval=interval)
+
+
+@pytest.mark.distributed
+@pytest.mark.parametrize("interval", [1, 2])
+@pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
+def test_ema_final_weight_gloo_cuda(get_dummy_model, distributed_context_single_node_gloo, interval):
+    device = idist.device()
+    _test_ema_final_weight(get_dummy_model(), device=device, ddp=True, interval=interval)
+
+
+@pytest.mark.distributed
+@pytest.mark.parametrize("interval", [1, 2])
+@pytest.mark.skipif(not idist.has_hvd_support, reason="Skip if no Horovod dist support")
+@pytest.mark.skipif("WORLD_SIZE" in os.environ, reason="Skip if launched as multiproc")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
+def test_ema_final_weight_hvd_cuda(gloo_hvd_executor, interval):
+    device = torch.device("cuda")
+    nproc = torch.cuda.device_count()
+
+    gloo_hvd_executor(_test_ema_final_weight, (device, True, interval), np=nproc, do_init=True)
+
+
+@pytest.mark.tpu
+@pytest.mark.skipif("NUM_TPU_WORKERS" in os.environ, reason="Skip if NUM_TPU_WORKERS is in env vars")
+@pytest.mark.skipif(not idist.has_xla_support, reason="Skip if no PyTorch XLA package")
+def test_ema_final_weight_single_device_xla(get_dummy_model):
+    device = idist.device()
+    _test_ema_final_weight(get_dummy_model(), device=device, ddp=True)
+
+
+@pytest.mark.tpu
+@pytest.mark.skipif("NUM_TPU_WORKERS" not in os.environ, reason="Skip if no NUM_TPU_WORKERS in env vars")
+@pytest.mark.skipif(not idist.has_xla_support, reason="Skip if no PyTorch XLA package")
+def test_ema_final_weight_xla_nprocs(get_dummy_model, xmp_executor):
+    n = int(os.environ["NUM_TPU_WORKERS"])
+
+    def _test_ema_final_weight_xla_nprocs(index):
+        device = idist.device()
+        _test_ema_final_weight(get_dummy_model(), device=device, ddp=True)
+
+    xmp_executor(_test_ema_final_weight_xla_nprocs, args=(), nprocs=n)
+
+
+@pytest.mark.multinode_distributed
+@pytest.mark.parametrize("interval", [1, 2])
+@pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
+@pytest.mark.skipif("MULTINODE_DISTRIB" not in os.environ, reason="Skip if not multi-node distributed")
+def test_ema_final_weight_multinode_gloo_cuda(get_dummy_model, distributed_context_multi_node_gloo, interval):
+    device = idist.device()
+    _test_ema_final_weight(get_dummy_model(), device=device, ddp=True, interval=interval)
+
+
+@pytest.mark.multinode_distributed
+@pytest.mark.parametrize("interval", [1, 2])
+@pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
+@pytest.mark.skipif("GPU_MULTINODE_DISTRIB" not in os.environ, reason="Skip if not multi-node distributed")
+def test_ema_final_weight_multinode_nccl_cuda(get_dummy_model, distributed_context_multi_node_nccl, interval):
+    device = idist.device()
+    _test_ema_final_weight(get_dummy_model(), device=device, ddp=True, interval=interval)