From ee635a0b649e39da0c25227ebf4a47222b571cfa Mon Sep 17 00:00:00 2001 From: wood-b Date: Thu, 13 Aug 2026 23:51:32 +0000 Subject: [PATCH] gp test fixes --- src/fairchem/core/common/test_utils.py | 7 ++ .../common/parallelism/test_graph_parallel.py | 78 +++++++++++++++---- 2 files changed, 68 insertions(+), 17 deletions(-) diff --git a/src/fairchem/core/common/test_utils.py b/src/fairchem/core/common/test_utils.py index 0182b515cf..b3ceb567cc 100644 --- a/src/fairchem/core/common/test_utils.py +++ b/src/fairchem/core/common/test_utils.py @@ -109,6 +109,13 @@ def spawn_multi_process( Spawn single node, multi-rank function. Uses a shared file for process group initialization to avoid port races. + IMPORTANT: ``test_method`` must return CPU tensors or plain Python values. + Results are passed back through a ``multiprocessing.Manager()`` dict whose + server process is forked from the caller, so unpickling a CUDA tensor there + fails with "Cannot re-initialize CUDA in forked subprocess" (the CUDA IPC + rebuild needs a CUDA context the forked server cannot create). Call + ``.detach().cpu()`` on any tensor before returning it from a NCCL worker. + Args: world_size: number of processes backend: backend to use. for example, "nccl", "gloo", etc diff --git a/tests/core/common/parallelism/test_graph_parallel.py b/tests/core/common/parallelism/test_graph_parallel.py index 3dbd5cdbdf..928b32b9e1 100644 --- a/tests/core/common/parallelism/test_graph_parallel.py +++ b/tests/core/common/parallelism/test_graph_parallel.py @@ -328,13 +328,16 @@ def a2a_vs_allgather_test(atomic_numbers, edge_index): natoms = atomic_numbers.shape[0] # Partition atoms (same as gp_utils does) - node_partition = torch.tensor_split(torch.arange(natoms), world_size)[rank] + device = atomic_numbers.device + node_partition = torch.tensor_split( + torch.arange(natoms, device=device), world_size + )[rank] node_offset = node_partition.min().item() - # Create rank assignments - rank_assignments = partition_atoms_index_split( - natoms, world_size, torch.device("cpu") - ) + # Create rank assignments on the same device as the data: build_gp_context + # derives its working device from rank_assignments, so a CPU tensor here + # would make it index a CPU mask with CUDA edge indices. + rank_assignments = partition_atoms_index_split(natoms, world_size, device) # Filter edges: keep edges where target is in our partition target_in_partition = (edge_index[1] >= node_partition.min()) & ( @@ -351,10 +354,12 @@ def a2a_vs_allgather_test(atomic_numbers, edge_index): # Run all-to-all version result_a2a = _a2a_simple_layer(x_local, local_edge_index, rank_assignments, natoms) + # Results travel back through spawn_multi_process's forked Manager + # server, which cannot unpickle CUDA tensors. Move to CPU first. return { "rank": rank, - "allgather": result_ag.detach(), - "all_to_all": result_a2a.detach(), + "allgather": result_ag.detach().cpu(), + "all_to_all": result_a2a.detach().cpu(), "match": torch.allclose(result_ag, result_a2a, atol=1e-6), } @@ -414,12 +419,13 @@ def a2a_backward_test(atomic_numbers, edge_index): natoms = atomic_numbers.shape[0] # Partition atoms - node_partition = torch.tensor_split(torch.arange(natoms), world_size)[rank] + device = atomic_numbers.device + node_partition = torch.tensor_split( + torch.arange(natoms, device=device), world_size + )[rank] node_offset = node_partition.min().item() - rank_assignments = partition_atoms_index_split( - natoms, world_size, torch.device("cpu") - ) + rank_assignments = partition_atoms_index_split(natoms, world_size, device) # Filter edges target_in_partition = (edge_index[1] >= node_partition.min()) & ( @@ -454,8 +460,10 @@ def a2a_backward_test(atomic_numbers, edge_index): create_graph=False, )[0] - results[f"{method}_energy"] = energy.detach() - results[f"{method}_forces"] = forces.detach() + # .cpu() so the results survive the forked Manager server used by + # spawn_multi_process, which cannot rebuild CUDA tensors. + results[f"{method}_energy"] = energy.detach().cpu() + results[f"{method}_forces"] = forces.detach().cpu() results["rank"] = rank results["energy_match"] = torch.allclose( @@ -551,7 +559,9 @@ def a2a_spatial_partition_test(atomic_numbers, edge_index, pos): natoms = atomic_numbers.shape[0] # --- All-gather with index-based partitioning (baseline) --- - node_partition_idx = torch.tensor_split(torch.arange(natoms), world_size)[rank] + node_partition_idx = torch.tensor_split( + torch.arange(natoms, device=atomic_numbers.device), world_size + )[rank] node_offset_idx = node_partition_idx.min().item() target_in_partition_idx = (edge_index[1] >= node_partition_idx.min()) & ( @@ -589,10 +599,12 @@ def a2a_spatial_partition_test(atomic_numbers, edge_index, pos): full_ag = gather_from_model_parallel_region_sum_grad(result_ag, natoms) full_a2a = gather_from_model_parallel_region_sum_grad(result_a2a, natoms) + # .cpu() so the results survive the forked Manager server used by + # spawn_multi_process, which cannot rebuild CUDA tensors. return { "rank": rank, - "allgather_full": full_ag.detach(), - "all_to_all_full": full_a2a.detach(), + "allgather_full": full_ag.detach().cpu(), + "all_to_all_full": full_a2a.detach().cpu(), "match": torch.allclose(full_ag, full_a2a, atol=1e-5), } @@ -826,6 +838,26 @@ def test_energy_forces_stress_gp(world_size): ) +def _requires_gpus(n: int): + """ + Skip unless at least n CUDA devices are visible. + + ``_to_cuda`` places rank r on ``cuda:r``, so a test spawning n ranks + needs n distinct devices. The ``gpu`` marker only covers the + zero-device case (see ``pytest_runtest_setup`` in tests/conftest.py). + + Args: + n: Number of CUDA devices the test requires. + + Returns: + A pytest skipif marker. + """ + return pytest.mark.skipif( + torch.cuda.device_count() < n, + reason=f"requires {n} CUDA devices", + ) + + def _to_cuda(*tensors): device = torch.device(f"cuda:{gp_utils.get_gp_rank()}") return tuple(t.to(device) for t in tensors) @@ -846,6 +878,8 @@ def a2a_spatial_partition_test_gpu(atomic_numbers, edge_index, pos): return a2a_spatial_partition_test(atomic_numbers, edge_index, pos) +@pytest.mark.gpu() +@_requires_gpus(2) @_skip_if_ci @pytest.mark.parametrize( "num_atoms, edges", @@ -885,6 +919,8 @@ def test_a2a_vs_allgather_gpu(num_atoms, edges): ) +@pytest.mark.gpu() +@_requires_gpus(2) @_skip_if_ci def test_a2a_backward_gpu(): atomic_numbers = torch.tensor([2.0, 3.0, 5.0, 7.0]) @@ -915,8 +951,14 @@ def test_a2a_backward_gpu(): ) +@pytest.mark.gpu() +@_requires_gpus(2) @_skip_if_ci -@pytest.mark.parametrize("world_size", [2, 3]) +@pytest.mark.parametrize( + "world_size", + # Gate the 3-rank case separately so 2-GPU hosts still run the 2-rank one. + [2, pytest.param(3, marks=_requires_gpus(3))], +) def test_a2a_multi_rank_gpu(world_size): num_atoms = 6 src = list(range(num_atoms)) @@ -946,6 +988,8 @@ def test_a2a_multi_rank_gpu(world_size): ) +@pytest.mark.gpu() +@_requires_gpus(2) @_skip_if_ci def test_a2a_spatial_partition_gpu(): num_atoms = 8