deepspeedai--deepspeed/tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py

# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

# TODO: add tests with model parallelism for activation partitioning and other features.

import pytest
import torch
import deepspeed
from deepspeed.pipe import PipelineModule, LayerSpec
from deepspeed.accelerator import get_accelerator
from copy import deepcopy
from unit.common import DistributedTest

ckpt = deepspeed.checkpointing.checkpoint


def _compute(module, *inputs, do_checkpoint=False):
    if do_checkpoint:
        outputs = ckpt(module, *inputs)
    else:
        outputs = module(*inputs)

    if torch.is_tensor(outputs):
        outputs = (outputs, )

    sum(o.sum() for o in outputs if torch.is_tensor(o) and o.requires_grad).backward()

    grads = [p.grad for p in module.parameters()]
    input_grads = [inp.grad for inp in inputs if torch.is_tensor(inp)]

    return {
        'outputs': outputs,
        'module_grads': grads,
        'input_grads': input_grads,
    }


def _prep_inputs(*inputs):
    _inputs = []

    for inp in inputs:
        inp = deepcopy(inp)
        if torch.is_tensor(inp):
            inp = inp.to(get_accelerator().device_name())
        _inputs.append(inp)

    return tuple(_inputs)


def _match_outputs(ref, tgt):
    assert type(ref) == type(tgt)
    if type(ref) in [list, tuple]:
        for x, y in zip(ref, tgt):
            _match_outputs(x, y)
    elif not torch.is_tensor(ref):
        assert ref == tgt
    elif ref.is_floating_point():
        assert torch.allclose(ref, tgt)
    else:
        assert torch.equal(ref, tgt)


def _test_activation_checkpoint(module, *inputs):
    if get_accelerator().device_name() == "cpu":
        pytest.skip("CPU accelerator does not support this test yet")
    # Move to device
    module.to(get_accelerator().device_name())

    # Get rid of dropouts until we fork the RNG between tests.
    module.eval()

    module_ = deepcopy(module)
    inputs_ = _prep_inputs(*inputs)
    base = _compute(module_, *inputs_, do_checkpoint=False)

    module_ = deepcopy(module)
    inputs_ = _prep_inputs(*inputs)
    test = _compute(module_, *inputs_, do_checkpoint=True)

    for group in base.keys():
        for b, t in zip(base[group], test[group]):
            _match_outputs(b, t)


def _test_activation_checkpoint_ordering(module, expected_ordering, *inputs):
    if get_accelerator().device_name() == "cpu":
        pytest.skip("CPU accelerator does not support this test yet")
    # Move to device
    module.to(get_accelerator().device_name())

    # Get rid of dropouts until we fork the RNG between tests.
    module.eval()

    module_ = deepcopy(module)
    inputs_ = _prep_inputs(*inputs)
    test = _compute(module_, *inputs_, do_checkpoint=True)

    outputs = test['outputs']
    test_ordering = []
    for item in outputs:
        if type(item) in [list, tuple]:
            test_ordering += [torch.is_tensor(t) for t in item]
        else:
            test_ordering += [torch.is_tensor(item)]

    assert expected_ordering == test_ordering


#
# Helpers
#


class MaskedLinear(torch.nn.Linear):

    def forward(self, x, mask):
        out = super().forward(x)
        if mask.is_floating_point():
            out = out * mask
        else:
            # must cast BoolTensor in older torch versions
            out = out * mask.type_as(out)
        return out


class MaskedLinearSeq(MaskedLinear):
    """Tests pipeline modules by also returning the mask."""

    def forward(self, x, mask):
        return super().forward(x, mask), mask


class MaskedLinearSeqDup(MaskedLinearSeq):
    """MaskedLinearSeq, but with more outputs than inputs and in a different order."""

    def forward(self, x, mask):
        dup = x.clone().detach() * 1.38  # just an arbitrary scaling
        x, mask = super().forward(x, mask)
        return dup, x, mask


class DropMaskLinear(torch.nn.Linear):

    def forward(self, x, mask):
        return super().forward(x)


class LinearNonTensorInput(torch.nn.Linear):

    def forward(self, x, non_tensor_input):
        return super().forward(x)


class LinearNonTensorOutput(torch.nn.Linear):

    def __init__(self, non_tensor_output):
        super().__init__(HIDDEN_DIM, HIDDEN_DIM)
        self.non_tensor_output = non_tensor_output

    def forward(self, x):
        out = super().forward(x)
        return out, self.non_tensor_output


HIDDEN_DIM = 20


def _mixed_mask(size=HIDDEN_DIM):
    entries = torch.randn(size)
    mask = torch.where(entries > 0, torch.ones(size), torch.zeros(size))
    mask = mask.bool()
    return mask


def _bool_to_float(btensor, dtype=torch.float32):
    """Converts a torch.BoolTensor to an equivalent dtype. """
    ones = torch.ones(size=btensor.size(), dtype=dtype)
    zeros = torch.zeros(size=btensor.size(), dtype=dtype)
    return torch.where(btensor, ones, zeros)


#
# Tests
#


# both bool and float are important, as bool is not differentiable
@pytest.mark.parametrize('mask', [
    _mixed_mask(),
    _bool_to_float(_mixed_mask()),
])
class TestActivationCheckpoint(DistributedTest):
    world_size = 1

    def test_ckpt_inputs1_outputs1(self, mask):
        module = torch.nn.Linear(HIDDEN_DIM, HIDDEN_DIM)
        inputs = torch.rand(HIDDEN_DIM)
        inputs.requires_grad = True
        _test_activation_checkpoint(module, inputs)

    def test_ckpt_inputs2_outputs1(self, mask):
        module = MaskedLinear(HIDDEN_DIM, HIDDEN_DIM)
        inputs = torch.rand(HIDDEN_DIM)
        inputs.requires_grad = True
        _test_activation_checkpoint(module, inputs, mask)

    def test_ckpt_inputs2_outputs2(self, mask):
        module = MaskedLinearSeq(HIDDEN_DIM, HIDDEN_DIM)
        inputs = torch.rand(HIDDEN_DIM)
        inputs.requires_grad = True
        _test_activation_checkpoint(module, inputs, mask)

    def test_ckpt_inputs2_outputs3(self, mask):
        module = MaskedLinearSeqDup(HIDDEN_DIM, HIDDEN_DIM)
        inputs = torch.rand(HIDDEN_DIM)
        inputs.requires_grad = True
        _test_activation_checkpoint(module, inputs, mask)

    def test_ckpt_arg_none(self, mask):
        module = DropMaskLinear(HIDDEN_DIM, HIDDEN_DIM)
        inputs = (torch.rand(HIDDEN_DIM), None)
        inputs[0].requires_grad = True
        _test_activation_checkpoint(module, *inputs)


@pytest.mark.parametrize('non_tensor', [None, 2, True, (None, 2.5), (None, True, torch.randn(HIDDEN_DIM))])
class TestCheckpointNonTensor(DistributedTest):
    world_size = 1

    def test_ckpt_non_tensor_input(self, non_tensor):
        module = LinearNonTensorInput(HIDDEN_DIM, HIDDEN_DIM)
        inputs = torch.rand(HIDDEN_DIM)
        inputs.requires_grad = True
        _test_activation_checkpoint(module, inputs, non_tensor)

    def test_ckpt_non_tensor_output(self, non_tensor):
        module = LinearNonTensorOutput(non_tensor)
        inputs = torch.rand(HIDDEN_DIM)
        inputs.requires_grad = True
        _test_activation_checkpoint(module, inputs)


@pytest.mark.parametrize('non_tensor_output', [
    None, (torch.randn(HIDDEN_DIM), 2.5), (None, torch.randn(HIDDEN_DIM), True), (None, True, torch.randn(HIDDEN_DIM))
])
class TestCheckpointNonTensorOutputOrdering(DistributedTest):
    world_size = 1

    def test_ckpt_non_tensor_output_ordering(self, non_tensor_output):
        module = LinearNonTensorOutput(non_tensor_output)
        inputs = torch.rand(HIDDEN_DIM)
        inputs.requires_grad = True

        # First return is a tensor
        ordering = [True]
        if type(non_tensor_output) in [list, tuple]:
            ordering += [torch.is_tensor(t) for t in non_tensor_output]
        else:
            ordering += [torch.is_tensor(non_tensor_output)]
        _test_activation_checkpoint_ordering(module, ordering, inputs)


class TestCheckpointableLayersConfig(DistributedTest):
    world_size = 1

    def test_gpt2_checkpointable_layers(self):
        if get_accelerator().device_name() == "cpu":
            pytest.skip("CPU accelerator does not support this test yet")

        # Create a simple topology for testing
        from deepspeed.runtime.pipe.topology import PipeModelDataParallelTopology
        topo = PipeModelDataParallelTopology(num_pp=1, num_mp=1, num_dp=1)

        # Create test classes that we want to checkpoint
        class TestTransformerLayer(torch.nn.Module):

            def forward(self, x):
                return x

        class ParallelTransformerLayerPipe(TestTransformerLayer):
            pass

        class GMLPBlock(TestTransformerLayer):
            pass

        # Create a mock GPT2 model with different layer types
        class TestGPT2ModelPipe(PipelineModule):

            def __init__(self):
                self.layers_spec = [
                    LayerSpec(ParallelTransformerLayerPipe),
                    LayerSpec(GMLPBlock),
                    LayerSpec(torch.nn.Linear, 10, 10),  # Should not be checkpointed
                ]

                super().__init__(layers=self.layers_spec,
                                 topology=topo,
                                 checkpointable_layers=["GMLPBlock", "ParallelTransformerLayerPipe"])

        model = TestGPT2ModelPipe()
        model.to(get_accelerator().device_name())

        # Build layers manually for testing
        layers = [spec.build() for spec in model.layers_spec]

        # Test that _is_checkpointable returns correct values
        assert model._is_checkpointable([layers[0]]) == True  # ParallelTransformerLayerPipe
        assert model._is_checkpointable([layers[1]]) == True  # GMLPBlock
        assert model._is_checkpointable([layers[2]]) == False  # Linear layer


def test_configure_with_contiguous_checkpointing_requires_num_checkpoints():
    # Regression: ``_configure_defaults`` previously initialized ``num_layers``
    # to ``False`` while the assert below uses ``is not None``; ``False is not
    # None`` is True, so the missing-config assert silently passed and a
    # cryptic ``IndexError`` surfaced later from ``range(num_layers)``. With
    # the default switched to ``None`` (matching the module-level default),
    # the helpful assert message fires at the configure() call site.
    #
    # ``configure()`` mutates module globals before raising, so snapshot and
    # restore them around the call to avoid order-dependent failures in other
    # activation-checkpointing tests sharing the same pytest worker.
    cp = deepspeed.checkpointing
    saved = (
        cp.PARTITION_ACTIVATIONS,
        cp.CONTIGUOUS_CHECKPOINTING,
        cp.num_layers,
        cp.CPU_CHECKPOINT,
        cp.SYNCHRONIZE,
        cp.PROFILE_TIME,
        cp.mpu,
        cp.deepspeed_checkpointing_enabled,
    )
    try:
        with pytest.raises(AssertionError, match="number of layers"):
            deepspeed.checkpointing.configure(
                mpu_=None,
                partition_activations=True,
                contiguous_checkpointing=True,
            )
    finally:
        (
            cp.PARTITION_ACTIVATIONS,
            cp.CONTIGUOUS_CHECKPOINTING,
            cp.num_layers,
            cp.CPU_CHECKPOINT,
            cp.SYNCHRONIZE,
            cp.PROFILE_TIME,
            cp.mpu,
            cp.deepspeed_checkpointing_enabled,
        ) = saved