feat(checkpoint): zero-copy storage sharing in CheckpointWithoutOutput

megatron/core/tensor_parallel/random.py

@@ -14,6 +14,7 @@
 from torch.cuda import _lazy_call, _lazy_init
 from torch.cuda import device as device_ctx_manager
 from torch.utils.checkpoint import detach_variable
+from torch.utils.cpp_extension import load_inline
 from typing_extensions import TypeVarTuple, Unpack
 
 from megatron.core.parallel_state import (
@@ -23,6 +24,57 @@
 )
 from megatron.core.utils import is_te_min_version, safely_set_viewless_tensor_data
 
+# ---------------------------------------------------------------------------
+# C++ extension: zero-copy storage sharing for CheckpointWithoutOutput
+# ---------------------------------------------------------------------------
+# Makes dst's UntypedStorage point to src's data WITHOUT copying bytes.
+# Holds a refcounted reference to src's StorageImpl so the memory stays alive.
+# Operates below the Tensor / autograd layer → no version-counter bump,
+# and ALL TensorImpls that reference dst's StorageImpl (including views
+# created by reshape / split / etc. inside TE GroupedLinear) see the data.
+# ---------------------------------------------------------------------------
+
+_SHARE_STORAGE_SRC = r"""
+#include <torch/extension.h>
+
+void share_storage(at::Tensor dst, at::Tensor src) {
+    auto* dst_impl = dst.storage().unsafeGetStorageImpl();
+
+    // Copy src's c10::Storage (increments StorageImpl refcount).
+    auto* src_storage_ref = new c10::Storage(src.storage());
+
+    void*       data   = src_storage_ref->data_ptr().get();
+    size_t      nbytes = src_storage_ref->nbytes();
+    c10::Device device = src_storage_ref->device();
+
+    // Build a DataPtr whose deleter releases our StorageImpl reference.
+    c10::DataPtr shared(
+        data,
+        static_cast<void*>(src_storage_ref),
+        [](void* ctx) { delete static_cast<c10::Storage*>(ctx); },
+        device);
+
+    dst_impl->set_data_ptr(std::move(shared));
+    dst_impl->set_nbytes(nbytes);
+}
+"""
+
+_share_storage_ext = None
+
+
+def _get_share_storage():
+    """Lazily compile & cache the share_storage extension."""
+    global _share_storage_ext
+    if _share_storage_ext is None:
+        _share_storage_ext = load_inline(
+            name="share_storage_ext",
+            cpp_sources=_SHARE_STORAGE_SRC,
+            functions=["share_storage"],
+            verbose=False,
+        )
+    return _share_storage_ext.share_storage
+
+
 from .utils import gather_split_1d_tensor, split_tensor_into_1d_equal_chunks
 
 try:
@@ -728,12 +780,14 @@ def detach(t):
         if isinstance(outputs, torch.Tensor):
             outputs = (outputs,)
 
-        # restore the recomputed memory without changing the metadata
-        with torch.no_grad():
-            for output, recomputation_output in zip(self.outputs, outputs):
-                output_size = recomputation_output.untyped_storage().size()
-                output.untyped_storage().resize_(output_size)
-                output.untyped_storage().copy_(recomputation_output.untyped_storage())
+        # Zero-copy: make output's StorageImpl point to recomputation_output's data.
+        # This operates at the UntypedStorage level (below TensorImpl), so:
+        #   - ALL views / reshapes that reference output's StorageImpl see the data
+        #     (e.g. TE GroupedLinear's inp.reshape() + torch.split() saved for backward)
+        #   - No tensor version-counter bump (no autograd complaint)
+        share_storage = _get_share_storage()
+        for output, recomputation_output in zip(self.outputs, outputs):
+            share_storage(output, recomputation_output)
 
         self.ctx.outputs = outputs
         self.ctx.inputs = inputs

tests/unit_tests/tensor_parallel/test_random.py

@@ -227,4 +227,62 @@ def checkpoint_forward(input):
     output2.backward(torch.ones((4, 4)), retain_graph=True)
     assert torch.equal(input1.grad, input2.grad)
 
-    Utils.destroy_model_parallel()
+
+class _ViewSavingLinear(torch.autograd.Function):
+    """Saves view tensors in forward to mimic TE GroupedLinear-style backward inputs."""
+
+    @staticmethod
+    def forward(ctx, inp, weight):
+        inp_2d = inp.reshape(-1, inp.shape[-1])
+        inputmats = torch.tensor_split(inp_2d, 2, dim=0)
+        ctx.save_for_backward(*inputmats, weight)
+        ctx.input_shape = inp.shape
+        out_2d = inp_2d.matmul(weight.t())
+        return out_2d.reshape(*inp.shape[:-1], weight.shape[0])
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        *inputmats, weight = ctx.saved_tensors
+        for inputmat in inputmats:
+            if inputmat.numel() > 0 and inputmat.untyped_storage().size() == 0:
+                raise RuntimeError("Saved view tensor points to an empty storage.")
+
+        inp_2d = torch.cat(inputmats, dim=0)
+        grad_output_2d = grad_output.reshape(-1, grad_output.shape[-1])
+        grad_input_2d = grad_output_2d.matmul(weight)
+        grad_weight = grad_output_2d.t().matmul(inp_2d)
+        grad_input = grad_input_2d.reshape(ctx.input_shape)
+        return grad_input, grad_weight
+
+
+def test_checkpoint_without_output_view_sharing_regression():
+    def normal_forward(input_, weight):
+        x = torch.nn.functional.gelu(input_)
+        return _ViewSavingLinear.apply(x, weight)
+
+    def checkpoint_forward(input_, weight):
+        checkpoint = CheckpointWithoutOutput()
+        x = checkpoint.checkpoint(torch.nn.functional.gelu, input_)
+        y = _ViewSavingLinear.apply(x, weight)
+        checkpoint.discard_output_and_register_recompute(y)
+        return y
+
+    Utils.initialize_model_parallel()
+    try:
+        input1 = torch.randn((3, 2, 8), requires_grad=True)
+        weight1 = torch.randn((6, 8), requires_grad=True)
+
+        input2 = input1.detach().clone().requires_grad_(True)
+        weight2 = weight1.detach().clone().requires_grad_(True)
+
+        output1 = normal_forward(input1, weight1)
+        output2 = checkpoint_forward(input2, weight2)
+        assert torch.allclose(output1, output2)
+
+        grad = torch.randn_like(output1)
+        output1.backward(grad, retain_graph=True)
+        output2.backward(grad, retain_graph=True)
+        assert torch.allclose(input1.grad, input2.grad)
+        assert torch.allclose(weight1.grad, weight2.grad)
+    finally:
+        Utils.destroy_model_parallel()