paddlepaddle--paddle/test/legacy_test/test_recompute_context.py

# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
#
# 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 threading
import unittest

import numpy as np

import paddle
from paddle import nn
from paddle.distributed.fleet.recompute import is_in_recompute
from paddle.distributed.fleet.recompute.recompute import _recompute_context
from paddle.distributed.fleet.utils import recompute


class TestRecomputeContext(unittest.TestCase):
    """Test _recompute_context and is_in_recompute."""

    def test_is_in_recompute_default_false(self):
        """By default is_in_recompute should return False."""
        self.assertFalse(is_in_recompute())

    def test_decorator_sets_active_true(self):
        """Inside decorated function, is_in_recompute should return True."""

        @_recompute_context
        def foo():
            self.assertTrue(is_in_recompute())
            return 42

        result = foo()
        self.assertEqual(result, 42)

    def test_decorator_resets_after_call(self):
        """After decorated function returns, is_in_recompute should be False."""

        @_recompute_context
        def foo():
            pass

        foo()
        self.assertFalse(is_in_recompute())

    def test_decorator_resets_on_exception(self):
        """Even if decorated function raises, active should be reset to False."""

        @_recompute_context
        def boom():
            self.assertTrue(is_in_recompute())
            raise RuntimeError("boom")

        with self.assertRaises(RuntimeError):
            boom()

        self.assertFalse(is_in_recompute())

    def test_decorator_preserves_return_value(self):
        """Decorator should not alter the return value."""

        @_recompute_context
        def add(a, b):
            return a + b

        self.assertEqual(add(3, 4), 7)

    def test_decorator_passes_kwargs(self):
        """Decorator should forward *args and **kwargs correctly."""

        @_recompute_context
        def greet(name, greeting="hello"):
            return f"{greeting} {name}"

        self.assertEqual(greet("world", greeting="hi"), "hi world")

    def test_thread_isolation(self):
        """is_in_recompute should be thread-local; other threads are unaffected."""

        @_recompute_context
        def blocked():
            # This function will block until the main thread signals
            barrier.wait()
            # While inside this thread's recompute context, main thread should still be False
            results["inside_thread"] = is_in_recompute()
            barrier.wait()

        results = {}
        barrier = threading.Barrier(2, timeout=5)

        t = threading.Thread(target=blocked)
        t.start()

        # Wait until the decorated function is running (active=True in that thread)
        barrier.wait()
        # Main thread is NOT in a recompute context
        results["main_thread"] = is_in_recompute()
        # Release the child thread
        barrier.wait()
        t.join()

        self.assertTrue(results["inside_thread"])
        self.assertFalse(results["main_thread"])
        # After thread finishes, main thread still False
        self.assertFalse(is_in_recompute())

    def test_context_manager_sets_active(self):
        """Using RecomputeContext as a with-statement should set is_in_recompute."""
        from paddle.distributed.fleet.recompute.recompute import (
            _recompute_context,
        )

        self.assertFalse(_recompute_context.active)

        with _recompute_context:
            self.assertTrue(_recompute_context.active)
            self.assertTrue(is_in_recompute())

        self.assertFalse(_recompute_context.active)
        self.assertFalse(is_in_recompute())

    def test_context_manager_resets_on_exception(self):
        """Context manager should reset active even if body raises."""
        from paddle.distributed.fleet.recompute.recompute import (
            _recompute_context,
        )

        with self.assertRaises(ValueError), _recompute_context:
            self.assertTrue(_recompute_context.active)
            raise ValueError("test")

        self.assertFalse(_recompute_context.active)

    def test_backward_compat_alias(self):
        """_recompute_context should still work as a decorator."""

        @_recompute_context
        def fn():
            return is_in_recompute()

        self.assertTrue(fn())
        self.assertFalse(is_in_recompute())


class SimpleRecomputeModel(nn.Layer):
    """A simple MLP that uses recompute on the middle layer."""

    def __init__(self, input_size=16, hidden_size=32):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, 1)
        self.relu = nn.ReLU()

    def forward(self, x, use_recompute=False):
        x = self.relu(self.fc1(x))
        if use_recompute:
            x = recompute(self._middle_block, x)
        else:
            x = self._middle_block(x)
        x = self.fc3(x)
        return x

    def _middle_block(self, x):
        return self.relu(self.fc2(x))


class TestRecomputeContextWithModel(unittest.TestCase):
    """Test recompute context with a real simple model forward/backward pass."""

    def setUp(self):
        paddle.seed(42)

    def test_recompute_produces_same_loss_as_normal(self):
        """Recompute and normal forward should produce identical loss values."""
        input_size, hidden_size, batch_size = 16, 32, 4

        # Run without recompute
        paddle.seed(42)
        model_normal = SimpleRecomputeModel(input_size, hidden_size)
        x = paddle.randn([batch_size, input_size])
        loss_normal = model_normal(x, use_recompute=False).mean()

        # Run with recompute using the same weights
        paddle.seed(42)
        model_recompute = SimpleRecomputeModel(input_size, hidden_size)
        loss_recompute = model_recompute(x, use_recompute=True).mean()

        np.testing.assert_allclose(
            loss_normal.numpy(), loss_recompute.numpy(), rtol=1e-5
        )

    def test_recompute_produces_same_gradients(self):
        """Gradients with recompute should match those without."""
        input_size, hidden_size, batch_size = 16, 32, 4

        paddle.seed(42)
        model_normal = SimpleRecomputeModel(input_size, hidden_size)
        x = paddle.randn([batch_size, input_size])
        loss = model_normal(x, use_recompute=False).mean()
        loss.backward()
        grads_normal = [
            p.grad.numpy().copy()
            for p in model_normal.parameters()
            if p.grad is not None
        ]

        paddle.seed(42)
        model_recompute = SimpleRecomputeModel(input_size, hidden_size)
        loss = model_recompute(x, use_recompute=True).mean()
        loss.backward()
        grads_recompute = [
            p.grad.numpy().copy()
            for p in model_recompute.parameters()
            if p.grad is not None
        ]

        self.assertEqual(len(grads_normal), len(grads_recompute))
        for g_normal, g_recompute in zip(grads_normal, grads_recompute):
            np.testing.assert_allclose(g_normal, g_recompute, rtol=1e-5)

    def test_is_in_recompute_true_during_recompute_forward(self):
        """is_in_recompute should be True inside the recomputed function."""
        observed_states = []

        class ObservingModel(nn.Layer):
            def __init__(self):
                super().__init__()
                self.fc = nn.Linear(8, 8)

            def forward(self, x):
                observed_states.append(is_in_recompute())
                return self.fc(x)

        model = ObservingModel()
        x = paddle.randn([2, 8])

        # Call via recompute — should observe True
        recompute(model, x)
        self.assertTrue(observed_states[-1])

        # After recompute finishes, context should be reset
        self.assertFalse(is_in_recompute())

    def test_training_step_with_recompute(self):
        """A full training step (forward + backward + optimizer) works with recompute."""
        input_size, hidden_size, batch_size = 16, 32, 4

        paddle.seed(42)
        model = SimpleRecomputeModel(input_size, hidden_size)
        optimizer = paddle.optimizer.SGD(
            learning_rate=0.01, parameters=model.parameters()
        )

        initial_params = [p.numpy().copy() for p in model.parameters()]

        x = paddle.randn([batch_size, input_size])
        loss = model(x, use_recompute=True).mean()
        loss.backward()
        optimizer.step()
        optimizer.clear_grad()

        # Parameters should have been updated
        for p_init, p_updated in zip(initial_params, model.parameters()):
            self.assertFalse(
                np.array_equal(p_init, p_updated.numpy()),
                "Parameters should change after optimization step",
            )


if __name__ == "__main__":
    unittest.main()