fix(cuda-graphs): preserve nested static input buffers

megatron/core/full_cuda_graph.py

@@ -75,23 +75,31 @@ def copy_tensors_in_struct(src):
 def clone_tensors_in_struct(tgt, src):
     """Copy src to pre-existing tensors in tgt."""
     if isinstance(src, tuple):
-        raise Exception(f"Unsupported copy for tuple yet: {type(src)}")
+        if not isinstance(tgt, tuple) or len(tgt) != len(src):
+            return copy_tensors_in_struct(src)
+        return tuple(clone_tensors_in_struct(t, s) for t, s in zip(tgt, src))
     elif isinstance(src, list):
+        if not isinstance(tgt, list) or len(tgt) != len(src):
+            return copy_tensors_in_struct(src)
         for i in range(len(src)):
-            if isinstance(src[i], (tuple, list, dict, torch.Tensor)):
-                clone_tensors_in_struct(tgt[i], src[i])
-            else:
-                tgt[i] = src[i]
+            tgt[i] = clone_tensors_in_struct(tgt[i], src[i])
+        return tgt
     elif isinstance(src, dict):
+        if not isinstance(tgt, dict):
+            return copy_tensors_in_struct(src)
         for k in src:
-            if isinstance(src[k], (tuple, list, dict, torch.Tensor)):
-                clone_tensors_in_struct(tgt[k], src[k])
+            if k in tgt:
+                tgt[k] = clone_tensors_in_struct(tgt[k], src[k])
             else:
-                tgt[k] = src[k]
+                tgt[k] = copy_tensors_in_struct(src[k])
+        return tgt
     elif isinstance(src, torch.Tensor):
+        if not isinstance(tgt, torch.Tensor) or tgt.shape != src.shape or tgt.dtype != src.dtype:
+            return copy_tensors_in_struct(src)
         tgt.copy_(src, non_blocking=True)
+        return tgt
     else:
-        raise Exception(f"Expect top-level as container type but got: {type(src)}")
+        return src
 
 
 # Class to copy dataloader output to static CUDA tensors for CUDA graph input. This

tests/unit_tests/transformer/test_full_cuda_graph.py

@@ -6,7 +6,7 @@
 
 import megatron.core.pipeline_parallel.schedules as schedule
 from megatron.core import ModelParallelConfig
-from megatron.core.full_cuda_graph import FullCudaGraphWrapper
+from megatron.core.full_cuda_graph import FullCudaGraphWrapper, clone_tensors_in_struct
 from megatron.core.tensor_parallel.random import (
     HAVE_TE,
     initialize_rng_tracker,
@@ -18,6 +18,30 @@
 rank = Utils.rank
 
 
+def test_clone_tensors_in_nested_struct():
+    list_tensor = torch.zeros(2, device='cpu')
+    tuple_tensor = torch.zeros(2, device='cpu')
+    nested_tensor = torch.zeros(2, device='cpu')
+    target_list = [list_tensor, {'nested': nested_tensor}]
+    target = {'list': target_list, 'tuple': (tuple_tensor, 0)}
+    source = {
+        'list': [torch.ones(2, device='cpu'), {'nested': torch.full((2,), 2.0, device='cpu')}],
+        'tuple': (torch.full((2,), 3.0, device='cpu'), 4),
+    }
+
+    result = clone_tensors_in_struct(target, source)
+
+    assert result is target
+    assert result['list'] is target_list
+    assert result['list'][0] is list_tensor
+    assert result['list'][1]['nested'] is nested_tensor
+    assert result['tuple'][0] is tuple_tensor
+    assert result['tuple'][1] == 4
+    torch.testing.assert_close(result['list'][0], source['list'][0])
+    torch.testing.assert_close(result['list'][1]['nested'], source['list'][1]['nested'])
+    torch.testing.assert_close(result['tuple'][0], source['tuple'][0])
+
+
 @pytest.mark.skipif(
     not (HAVE_TE and is_te_min_version("1.5.0")),
     reason="use_te_rng_tracker requires TransformerEngine version >= 1.5",