Restore gp test - #2147
Closed
rayg1234 wants to merge 61 commits into
Closed
Conversation
added 30 commits
May 12, 2026 05:48
All-to-all (A2A) communication module for graph parallelism, replacing the all-gather approach with point-to-point exchange of only boundary atoms. Key components: - GPContext: dataclass holding per-rank atom assignments and A2A metadata - build_gp_context: builds communication plan from edge connectivity - AllToAllCollect: autograd-compatible A2A embedding exchange - all_to_all_collect_compiled: torch.compile-friendly variant using functional collectives (no graph break) - _safe_all_to_all: Gloo fallback for CPU testing - _sparse_index_exchange: variable-split index exchange Tests (22 total, all CPU/Gloo): - 5 unit tests for build_gp_context (context building, global-to-local mapping, edge split indices) - 9 distributed tests comparing A2A vs all-gather correctness (forward, backward, multi-rank, spatial partition) - 8 distributed correctness tests (dense/sparse graphs, multi-dim embeddings, index_split + spatial strategies)
Model integration for all-to-all graph parallelism: eSCN-MD backbone (escn_md.py): - Add use_all_to_all_gp and gp_partition_strategy config options - Replace all-gather with A2A embedding exchange when enabled - Spatial partitioning via Morton Z-order curve - AABB halo filtering for reduced graph generation input - Support both autograd and compiled A2A collect variants - Block allgather+spatial combination (unsupported) eSCN-MD block (escn_md_block.py): - Add A2A collect integration in message passing layers - Use precomputed edge split indices for local/remote separation Graph generation (compute.py): - Extend filter_edges_by_node_partition with send_info computation - Eliminates NCCL index-exchange collective in build_gp_context - Pass rank_assignments through generate_graph Embedding/execution backends: - Thread GP context through embedding and execution layers Tests: - Add send_info optimization correctness test - Expand ParallelPredictUnit tests with A2A+spatial and A2A+index_split GP modes (CPU: workers=2, GPU: workers=1) - 48 parallelism tests + 25 predict tests all pass
Parametrize test_merge_mole_md_consistency with A2A+spatial GP mode to verify that all-to-all graph parallel with spatial partitioning works correctly across multiple MD timesteps (NVT and NPT). Skips A2A modes when workers < 2 since GP requires multi-rank. 2 new test cases pass (NVT+A2A, NPT+A2A), 2 correctly skipped.
…coverage - Add is_single_system guard in _generate_graph so AABB halo bails out for multi-system batches (prevents cell.view(3,3) crash) - Add defensive assertion in _compute_halo_graph verifying all local partition atoms appear in the halo mask - Add A2A+index_split mode to MD consistency test parametrization - Add A2A+spatial mode to batch predict test (exercises halo bail-out path for multi-system batches) - Document single-system limitation in _compute_aabb_halo docstring
Extract _resolve_send_metadata and _validate_gp_mappings from the monolithic build_gp_context (270→120 lines). Validation now uses fast-path .any().item() checks with full diagnostics only on error. Add _allgather_index_exchange: packs recv_counts + needed_atoms into one buffer for a single all-gather call (vs 2x all-to-all). Wired into eSCN-MD via gp_index_exchange_method constructor parameter. Benchmarks at GP=64 (8 nodes, 256K atoms): 2xA2A +22.9%, 1xAG +20.3% vs baseline all-gather GP. Both variants beat baseline at scale.
Benchmarks showed 2xA2A is ~1.5% faster than 1xallgather at GP=64 across 8 nodes (0.593 vs 0.584 ns/day) due to lower communication volume with variable splits. Remove the inferior allgather variant, the index_exchange_method parameter, and the dispatch logic. Add benchmark rationale to _sparse_index_exchange docstring.
Inference always requires autograd for force computation, so the no-grad compile-friendly path was never hit. Remove the function, its dispatch branch in escn_md_block, and associated tests.
The complex 13-parameter diagnostic function was hard to follow. Replaced with two inline asserts for the only error conditions: negative edge_index_local entries and out-of-bounds send_indices.
v2 radius graph does internal edge filtering so the full edge_index is never exposed, making send_info derivation in compute.py impossible. Revert compute.py to main, remove _resolve_send_metadata, and always use _sparse_index_exchange in build_gp_context.
This reverts commit 6168d23.
Remove 6 construction-only fields from GPContext (node_partition, rank_assignments, needed_atoms, needed_from_ranks, global_to_local, total_needed_atoms) that are never read after build_gp_context returns. Update tests to verify through edge_index_local and message passing instead of accessing removed intermediates.
The 10s timeout causes flaky failures in CI where runners are resource-constrained and process rendezvous takes longer.
Benchmarks at 64 GPUs showed halo adds ~5% overhead vs plain A2A. The _sparse_index_exchange already minimizes communication, and radius_graph_pbc_v2 grid indexing makes graph gen fast enough that filtering to a halo subset adds cost without measurable savings. Removes ~150 lines: _compute_aabb_halo, _compute_halo_graph, use_aabb_halo flag, and the halo call site in _generate_graph. Also adds timeout_hr to slurm benchmark config.
Resolve conflict in test_predict.py: adopt main's pretrained_checkpoint fixture and get_predict_unit_for_test helper, add back gp_mode parameter for A2A test coverage (spatial + index_split).
…lism module - Add GPMode enum (ALLGATHER, ALL_TO_ALL) with set/get/is_a2a accessors to gp_utils so non-model code can query the active GP communication pattern without model-level flags - Move _compute_a2a_partition from escn_md.py to common/parallelism/graph_parallel_a2a.py as a public function - Call set_gp_mode(ALL_TO_ALL) in eSCNMDBackbone.__init__ when use_all_to_all_gp is set - Reset _GP_MODE in cleanup_gp()
scatter_target is now set in _generate_graph for all three paths: - A2A: gp_ctx.edge_index_local[1] - Allgather: global_to_local[edge_index[1]] - No GP: edge_index[1] This removes the ternary in forward() and the 3-way if/elif/else for local_scatter_target in Edgewise.forward().
Move graph parallel configuration (group_size, mode, partition) out of the model constructor into a single GraphParallelConfig dataclass stored as a global in gp_utils. The model now reads GP config from the global instead of owning use_all_to_all_gp and gp_partition_strategy params. Deprecate graph_parallel_group_size in favor of job.graph_parallel block. Raise ValueError when both are specified with conflicting group sizes.
…size The CLI calls __post_init__ twice (once from OmegaConf.to_object, once explicitly). The first call coerces graph_parallel_group_size into graph_parallel, so the second call would see both set and error. Now skip coercion when values already match.
Callers already have gp_config in scope or are guarded by gp_utils.initialized(), so the wrapper adds no value.
setup_graph_parallel_groups was called unconditionally for all multi-worker predict units, causing initialized() to return True even when no GP was requested. This made _generate_graph access get_gp_config().mode on a None config. Now GP groups are only created when gp_config is explicitly provided.
all_to_all_collect now takes only (x_local, gp_ctx) — send_indices is read from gp_ctx internally.
# Conflicts: # tests/core/units/mlip_unit/test_predict.py
Replace raw strings with StrEnum types for mode and partition fields. StrEnum preserves backward compatibility — string comparisons and construction with string literals still work. PartitionStrategy is now an alias of GPPartition to avoid duplication.
Move the no-GP fallback assignment before the empty-edge block so scatter_target is always present, eliminating both 'in data_dict' checks.
With multiple workers and no gp_config, each worker ran a redundant full copy of inference. Now workers > 1 auto-creates a GraphParallelConfig(group_size=num_workers), and mismatched group_size vs num_workers raises ValueError.
When gp_config is provided with only mode/partition but group_size left at default (1), infer it from num_workers instead of erroring. Only raise ValueError when group_size is explicitly set to a different value.
# Conflicts: # src/fairchem/core/models/uma/escn_md_block.py
Both setup_gp() and setup_graph_parallel_groups() previously left
_GP_CONFIG untouched, so downstream code doing
gp_utils.get_gp_config().mode
could crash with AttributeError on None. This bit tests that use
setup_gp() via spawn_multi_process (test_utils.py), where callers
never explicitly set a config.
Have both entry points install a default
GraphParallelConfig(group_size=N) when none is already set. Callers
that need non-default settings (A2A, spatial partition) still call
set_gp_config() afterwards — the guard checks `if _GP_CONFIG is None`
so intentional configs are preserved.
…_workers Move the num_workers-vs-gp_config reconciliation out of predict.py into gp_utils.resolve_gp_config_for_workers so other multi-worker entry points can reuse it. Behavior unchanged; caller's config is never mutated (uses dataclasses.replace).
Co-authored-by: Luis Barroso-Luque <lbluque@users.noreply.github.com>
Re-add the full-model GP correctness test (no-GP / allgather / A2A-spatial / A2A-index_split at 1/2/4 workers vs single-GPU reference on energy/forces/ stress), which was silently dropped during an earlier main merge. Adapted to the current gp_config=GraphParallelConfig(...) API and the pretrained fixture convention. Referenced by graph_parallel_verification.md section 1d.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.