From dad32edc5a6894c62afd59bdaec41a5a1b3d5d36 Mon Sep 17 00:00:00 2001 From: lichenlu Date: Sat, 9 May 2026 01:55:57 +0000 Subject: [PATCH] Strengthen test_checkpoint to verify distributed checkpoint behavior The test now properly initializes model parallel with explicit sizes, places tensors on CUDA, enables gradients, and asserts forward output correctness, input reshaping after distribute, and backward pass. Co-authored-by: peibli --- .../unit_tests/tensor_parallel/test_random.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/tests/unit_tests/tensor_parallel/test_random.py b/tests/unit_tests/tensor_parallel/test_random.py index 4fa79733d55..873e6a30f00 100644 --- a/tests/unit_tests/tensor_parallel/test_random.py +++ b/tests/unit_tests/tensor_parallel/test_random.py @@ -193,10 +193,20 @@ def test_forward(*input): assert torch.equal( torch.ones(16) * 3, checkpoint(test_forward, None, torch.ones(16), torch.ones(16) * 2) ) - Utils.initialize_model_parallel() - input1 = torch.ones((4, 4)) - checkpoint(test_forward, True, input1, torch.ones((4, 4)) * 2) - assert torch.equal(torch.ones(input1.numel()).cuda(), input1) + + Utils.initialize_model_parallel(tensor_model_parallel_size=2, pipeline_model_parallel_size=1) + input1 = torch.ones((4, 4)).cuda() + input1.requires_grad_(True) + input2 = torch.ones((4, 4)).cuda() * 2 + output = checkpoint(test_forward, True, input1, input2) + + assert torch.equal(output, torch.ones((4, 4)).cuda() * 3) + assert input1.data.shape == (8,) + + output.sum().backward() + assert input1.grad is not None + assert torch.equal(input1.grad, torch.ones((4, 4)).cuda()) + Utils.destroy_model_parallel()