test_weight_averaging.py

# Copyright The Lightning AI team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from copy import deepcopy
from pathlib import Path
from typing import Any, Optional

import pytest
import torch
from torch import Tensor, nn
from torch.optim.swa_utils import get_swa_avg_fn
from torch.utils.data import DataLoader, Dataset

from lightning.pytorch import LightningModule, Trainer
from lightning.pytorch.callbacks import EMAWeightAveraging, WeightAveraging
from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset, RandomIterableDataset
from tests_pytorch.helpers.runif import RunIf


class TestModel(BoringModel):
    def __init__(self, batch_norm: bool = True) -> None:
        super().__init__()
        layers = [nn.Linear(32, 32)]
        if batch_norm:
            layers.append(nn.BatchNorm1d(32))
        layers += [nn.ReLU(), nn.Linear(32, 2)]
        self.layer = nn.Sequential(*layers)
        self.crash_on_epoch = None

    def training_step(self, batch: Tensor, batch_idx: int) -> None:
        if self.crash_on_epoch and self.trainer.current_epoch >= self.crash_on_epoch:
            raise Exception("CRASH")
        return super().training_step(batch, batch_idx)

    def configure_optimizers(self) -> None:
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


class LargeTestModel(BoringModel):
    def __init__(self):
        super().__init__()
        self.layer = None

    def configure_model(self):
        print("XXX configure_model")
        self.layer = nn.Sequential(nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 2))

    def configure_optimizers(self):
        return torch.optim.SGD(self.parameters(), lr=0.01)


class EMAAveragingFunction:
    """EMA averaging function.

    Functionally equivalent to the closure that ``get_ema_avg_fn()`` would return. This class is needed because we
    cannot use a closure with ddp_spawn. (``Popen(process_obj)`` would fail with
    ``Can't get local object 'get_ema_avg_fn.<locals>.ema_update'``).

    """

    def __init__(self, decay: float = 0.999) -> None:
        self.decay = decay

    @torch.no_grad()
    def __call__(self, ema_param: Tensor, current_param: Tensor, num_averaged: Tensor) -> Tensor:
        return self.decay * ema_param + (1 - self.decay) * current_param


class EMATestCallback(WeightAveraging):
    def __init__(self, devices: int = 1, **kwargs: Any) -> None:
        super().__init__(avg_fn=EMAAveragingFunction(), **kwargs)
        self.devices = devices
        self.swap_calls = 0
        self.copy_calls = 0
        # Record the first epoch, as if we are resuming from a checkpoint this may not be equal to 0.
        self.first_epoch: Optional[int] = None

    def _swap_models(self, *args: Any, **kwargs: Any):
        self.swap_calls += 1
        return super()._swap_models(*args, **kwargs)

    def _copy_average_to_current(self, *args: Any, **kwargs: Any):
        self.copy_calls += 1
        return super()._copy_average_to_current(*args, **kwargs)

    def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
        super().on_train_start(trainer, pl_module)
        assert self.swap_calls == 0
        assert self.copy_calls == 0

    def on_train_epoch_start(self, trainer: Trainer, *args: Any) -> None:
        super().on_train_epoch_start(trainer, *args)
        # Since the checkpoint loaded was saved `on_train_epoch_end`, the first `FitLoop` iteration will not update the
        # model and will just call the epoch-level hooks. For that reason, we check that we are not restarting before
        # choosing the first epoch.
        if self.first_epoch is None and not trainer.fit_loop.restarting:
            self.first_epoch = trainer.current_epoch

    def on_train_epoch_end(self, trainer: Trainer, *args: Any) -> None:
        super().on_train_epoch_end(trainer, *args)
        assert self._average_model.n_averaged == trainer.global_step
        assert self.swap_calls == (trainer.current_epoch + 1 - self.first_epoch) * 2
        assert self.copy_calls == 0

    def on_train_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
        super().on_train_end(trainer, pl_module)
        # length=32, batch_size=4, accumulate_grad_batches=2
        # => Using one process we have 4 optimizer steps per epoch.
        # => Using two processes we have 2 optimizer steps per epoch.
        steps_per_epoch = 4 // self.devices
        assert self._average_model.n_averaged == trainer.max_epochs * steps_per_epoch
        assert self.swap_calls == (trainer.max_epochs - self.first_epoch) * 2
        assert self.copy_calls == 1


class SWATestCallback(WeightAveraging):
    def __init__(self, **kwargs: Any) -> None:
        super().__init__(avg_fn=get_swa_avg_fn(), **kwargs)
        self.swap_calls = 0
        self.copy_calls = 0

    def should_update(self, step_idx: Optional[int] = None, epoch_idx: Optional[int] = None) -> bool:
        return epoch_idx in (3, 5, 7)

    def _swap_models(self, *args: Any, **kwargs: Any):
        self.swap_calls += 1
        return super()._swap_models(*args, **kwargs)

    def _copy_average_to_current(self, *args: Any, **kwargs: Any):
        self.copy_calls += 1
        return super()._copy_average_to_current(*args, **kwargs)

    def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
        super().on_train_start(trainer, pl_module)
        assert self.swap_calls == 0
        assert self.copy_calls == 0

    def on_train_epoch_end(self, trainer: Trainer, *args: Any) -> None:
        super().on_train_epoch_end(trainer, *args)
        if trainer.current_epoch < 3:
            assert self._average_model.n_averaged == 0
        elif trainer.current_epoch < 5:
            assert self._average_model.n_averaged == 1
        elif trainer.current_epoch < 7:
            assert self._average_model.n_averaged == 2
        else:
            assert self._average_model.n_averaged == 3
        # The model is only swapped during validation once the average model has been updated. The first update happens
        # at the end of epoch 3, so the validation epochs that swap are those from epoch 4 onwards (two swaps each).
        assert self.swap_calls == max(0, trainer.current_epoch - 3) * 2
        assert self.copy_calls == 0

    def on_train_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
        super().on_train_end(trainer, pl_module)
        assert self._average_model.n_averaged == 3
        # Validation epochs 4 to ``max_epochs - 1`` each swap twice (see ``on_train_epoch_end``).
        assert self.swap_calls == max(0, trainer.max_epochs - 1 - 3) * 2
        assert self.copy_calls == 1


def test_weight_averaging_deepcopy(tmp_path):
    """Ensure that WeightAveraging callback doesn't deepcopy the data loaders or the data module and consume memory
    more than necessary."""

    class TestCallback(WeightAveraging):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self.setup_called = False

        def setup(self, trainer, pl_module, stage) -> None:
            super().setup(trainer, pl_module, stage)
            assert self._average_model.module.train_dataloader is not pl_module.train_dataloader
            assert self._average_model.module.train_dataloader.__self__ == self._average_model.module
            assert self._average_model.module._trainer is None
            self.setup_called = True

    callback = TestCallback()
    trainer = Trainer(default_root_dir=tmp_path, callbacks=callback, fast_dev_run=True)
    trainer.fit(BoringModel(), train_dataloaders=DataLoader(RandomDataset(32, 2)))
    assert callback.setup_called


@pytest.mark.parametrize("batch_norm", [True, False])
@pytest.mark.parametrize("iterable_dataset", [True, False])
def test_ema(tmp_path, batch_norm: bool, iterable_dataset: bool):
    model = TestModel(batch_norm=batch_norm)
    dataset = RandomIterableDataset(32, 32) if iterable_dataset else RandomDataset(32, 32)
    _train(model, dataset, tmp_path, EMATestCallback())


@pytest.mark.parametrize(
    "accelerator", [pytest.param("gpu", marks=RunIf(min_cuda_gpus=1)), pytest.param("mps", marks=RunIf(mps=True))]
)
def test_ema_accelerator(tmp_path, accelerator):
    model = TestModel()
    dataset = RandomDataset(32, 32)
    _train(model, dataset, tmp_path, EMATestCallback(), accelerator=accelerator, devices=1)


@RunIf(min_cuda_gpus=2, standalone=True)
def test_ema_ddp(tmp_path):
    model = TestModel()
    dataset = RandomDataset(32, 32)
    _train(model, dataset, tmp_path, EMATestCallback(devices=2), strategy="ddp", accelerator="gpu", devices=2)


@RunIf(min_cuda_gpus=2)
def test_ema_ddp_spawn(tmp_path):
    model = TestModel()
    dataset = RandomDataset(32, 32)
    _train(model, dataset, tmp_path, EMATestCallback(devices=2), strategy="ddp_spawn", accelerator="gpu", devices=2)


@RunIf(skip_windows=True)
def test_ema_ddp_spawn_cpu(tmp_path):
    model = TestModel()
    dataset = RandomDataset(32, 32)
    _train(model, dataset, tmp_path, EMATestCallback(devices=2), strategy="ddp_spawn", accelerator="cpu", devices=2)


@pytest.mark.parametrize("crash_on_epoch", [1, 3, 5])
def test_ema_resume(tmp_path, crash_on_epoch):
    dataset = RandomDataset(32, 32)
    model1 = TestModel()
    model2 = deepcopy(model1)

    _train(model1, dataset, tmp_path, EMATestCallback())

    model2.crash_on_epoch = crash_on_epoch
    model2 = _train_and_resume(model2, dataset, tmp_path)

    for param1, param2 in zip(model1.parameters(), model2.parameters()):
        assert torch.allclose(param1, param2)


@RunIf(skip_windows=True)
def test_ema_resume_ddp(tmp_path):
    model = TestModel()
    model.crash_on_epoch = 3
    dataset = RandomDataset(32, 32)
    _train_and_resume(model, dataset, tmp_path, strategy="ddp_spawn", devices=2)


def test_swa(tmp_path):
    model = TestModel()
    dataset = RandomDataset(32, 32)
    _train(model, dataset, tmp_path, SWATestCallback())


@pytest.mark.parametrize(
    ("strategy", "accelerator", "devices"),
    [
        ("auto", "cpu", 1),
        pytest.param("auto", "gpu", 1, marks=RunIf(min_cuda_gpus=1)),
        pytest.param("fsdp", "gpu", 1, marks=RunIf(min_cuda_gpus=1)),
    ],
)
def test_ema_configure_model(tmp_path, strategy, accelerator, devices):
    model = LargeTestModel()
    dataset = RandomDataset(32, 32)
    callback = EMATestCallback()
    _train(model, dataset, tmp_path, callback, strategy=strategy, accelerator=accelerator, devices=devices)
    assert isinstance(callback._average_model.module.layer, nn.Sequential)


def _train(
    model: BoringModel,
    dataset: Dataset,
    tmp_path: str,
    callback: WeightAveraging,
    strategy: str = "auto",
    accelerator: str = "cpu",
    devices: int = 1,
    checkpoint_path: Optional[str] = None,
    will_crash: bool = False,
) -> None:
    deterministic = accelerator == "cpu"
    trainer = Trainer(
        accelerator=accelerator,
        strategy=strategy,
        devices=devices,
        logger=False,
        callbacks=callback,
        max_epochs=8,
        num_sanity_val_steps=0,
        enable_checkpointing=will_crash,
        enable_progress_bar=False,
        enable_model_summary=False,
        accumulate_grad_batches=2,
        deterministic=deterministic,
        default_root_dir=tmp_path,
    )
    dataloader = DataLoader(dataset, batch_size=4, shuffle=False)
    if will_crash:
        with pytest.raises(Exception, match="CRASH"):
            trainer.fit(model, dataloader, ckpt_path=checkpoint_path)
    else:
        trainer.fit(model, dataloader, ckpt_path=checkpoint_path)
    assert trainer.lightning_module == model


def _train_and_resume(model: TestModel, dataset: Dataset, tmp_path: str, devices: int = 1, **kwargs) -> TestModel:
    _train(model, dataset, tmp_path, EMATestCallback(devices=devices), devices=devices, will_crash=True, **kwargs)

    checkpoint_dir = Path(tmp_path) / "checkpoints"
    checkpoint_names = os.listdir(checkpoint_dir)
    assert len(checkpoint_names) == 1
    checkpoint_path = str(checkpoint_dir / checkpoint_names[0])

    model = TestModel.load_from_checkpoint(checkpoint_path)
    callback = EMATestCallback(devices=devices)
    _train(model, dataset, tmp_path, callback, devices=devices, checkpoint_path=checkpoint_path, **kwargs)
    return model


@pytest.mark.parametrize(
    ("strategy", "accelerator", "devices"),
    [
        ("auto", "cpu", 1),
        pytest.param("auto", "gpu", 1, marks=RunIf(min_cuda_gpus=1)),
    ],
)
def test_ema_weight_averaging(tmp_path, strategy, accelerator, devices):
    """Test EMAWeightAveraging callback with various update configurations."""
    model = TestModel()
    dataset = RandomDataset(32, 32)

    # Test with default settings (update every step)
    callback = EMAWeightAveraging(decay=0.999, update_every_n_steps=1)
    _train(model, dataset, tmp_path, callback, strategy=strategy, accelerator=accelerator, devices=devices)

    # Verify the average model was created and updated
    assert callback._average_model is not None
    assert callback._average_model.n_averaged > 0


def test_ema_weight_averaging_step_frequency(tmp_path):
    """Test EMAWeightAveraging with custom step update frequency."""
    model = TestModel()
    dataset = RandomDataset(32, 32)

    # Update every 5 steps
    callback = EMAWeightAveraging(decay=0.95, update_every_n_steps=5)
    _train(model, dataset, tmp_path, callback)

    assert callback._average_model is not None


def test_ema_weight_averaging_starting_step(tmp_path):
    """Test EMAWeightAveraging with delayed start based on steps."""
    model = TestModel()
    dataset = RandomDataset(32, 32)

    # Start updating after step 10
    callback = EMAWeightAveraging(decay=0.999, update_every_n_steps=1, update_starting_at_step=10)
    _train(model, dataset, tmp_path, callback)

    assert callback._average_model is not None


def test_ema_weight_averaging_starting_epoch(tmp_path):
    """Test EMAWeightAveraging with delayed start based on epochs."""
    model = TestModel()
    dataset = RandomDataset(32, 32)

    # Start updating after epoch 3
    callback = EMAWeightAveraging(decay=0.999, update_every_n_steps=1, update_starting_at_epoch=3)
    _train(model, dataset, tmp_path, callback)

    assert callback._average_model is not None


def test_weight_averaging_no_swap_before_first_update(tmp_path):
    """Validation must use the current (trained) weights while the average model has not been updated yet.

    Before the first update, the ``AveragedModel`` only holds the copy of the initial weights made in ``setup()``.
    Swapping it in for validation during a delayed-start warmup would discard the trained weights and evaluate the
    untrained snapshot instead. See https://github.com/Lightning-AI/pytorch-lightning/issues/21724.

    """

    class SwapProbeModel(BoringModel):
        def __init__(self) -> None:
            super().__init__()
            self.layer = nn.Linear(32, 2)
            self.initial_weight: Optional[Tensor] = None
            self.validation_weight: Optional[Tensor] = None

        def on_train_start(self) -> None:
            # Snapshot of the weights the average model copies in ``setup()``.
            self.initial_weight = self.layer.weight.detach().clone()

        def on_validation_epoch_start(self) -> None:
            self.validation_weight = self.layer.weight.detach().clone()

        def configure_optimizers(self):
            # A large learning rate guarantees the weights move noticeably away from their initial values.
            return torch.optim.SGD(self.parameters(), lr=1.0)

    model = SwapProbeModel()
    dataset = RandomDataset(32, 32)
    # The update threshold is never reached during this run, so the average model is never updated.
    callback = EMAWeightAveraging(update_every_n_steps=1, update_starting_at_step=1000)
    trainer = Trainer(
        accelerator="cpu",
        devices=1,
        logger=False,
        callbacks=callback,
        max_epochs=1,
        num_sanity_val_steps=0,
        enable_checkpointing=False,
        enable_progress_bar=False,
        enable_model_summary=False,
        limit_train_batches=4,
        limit_val_batches=1,
        deterministic=True,
        default_root_dir=tmp_path,
    )
    dataloader = DataLoader(dataset, batch_size=4, shuffle=False)
    trainer.fit(model, train_dataloaders=dataloader, val_dataloaders=dataloader)

    assert callback._average_model is not None
    assert callback._average_model.n_averaged == 0
    assert model.initial_weight is not None
    assert model.validation_weight is not None
    # Validation must not have swapped in the un-updated (frozen) average model.
    assert not torch.allclose(model.validation_weight, model.initial_weight)


def test_ema_weight_averaging_should_update(tmp_path):
    """Test the should_update logic of EMAWeightAveraging."""
    # Test with step-based updates
    callback = EMAWeightAveraging(update_every_n_steps=5, update_starting_at_step=10)

    # Before starting step
    assert not callback.should_update(step_idx=5)
    assert not callback.should_update(step_idx=9)

    # At and after starting step, but not on update frequency
    assert callback.should_update(step_idx=10)  # First update
    assert not callback.should_update(step_idx=11)
    assert not callback.should_update(step_idx=14)
    assert callback.should_update(step_idx=15)  # Second update

    # Test with epoch-based updates
    callback = EMAWeightAveraging(update_starting_at_epoch=2)

    assert not callback.should_update(epoch_idx=0)
    assert not callback.should_update(epoch_idx=1)
    assert callback.should_update(epoch_idx=2)
    assert callback.should_update(epoch_idx=3)


def test_ema_weight_averaging_checkpoint_save_load(tmp_path):
    """Test that EMAWeightAveraging correctly saves and loads checkpoints."""
    model = TestModel()
    model.crash_on_epoch = 2
    dataset = RandomDataset(32, 32)

    callback = EMAWeightAveraging(decay=0.99, update_every_n_steps=2)

    # Train and create checkpoint
    _train(model, dataset, tmp_path, callback, will_crash=True)

    # Resume from checkpoint
    model2 = TestModel()
    callback2 = EMAWeightAveraging(decay=0.99, update_every_n_steps=2)
    import glob  # should be at the top

    _train(
        model2,
        dataset,
        tmp_path,
        callback2,
        checkpoint_path=glob.glob((tmp_path / "checkpoints" / "*.ckpt").as_posix())[0],
    )

    assert callback2._average_model is not None


@pytest.mark.parametrize("decay", [0.9, 0.99, 0.999, 0.9999])
def test_ema_weight_averaging_decay_values(tmp_path, decay):
    """Test EMAWeightAveraging with different decay values."""
    model = TestModel()
    dataset = RandomDataset(32, 32)

    callback = EMAWeightAveraging(decay=decay, update_every_n_steps=1)
    _train(model, dataset, tmp_path, callback)

    assert callback._average_model is not None