diff --git a/megatron/core/full_cuda_graph.py b/megatron/core/full_cuda_graph.py index 686b2db5f89..112d4278be5 100644 --- a/megatron/core/full_cuda_graph.py +++ b/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 diff --git a/tests/unit_tests/transformer/test_full_cuda_graph.py b/tests/unit_tests/transformer/test_full_cuda_graph.py index 312ae467304..59b4dd96e9b 100644 --- a/tests/unit_tests/transformer/test_full_cuda_graph.py +++ b/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",