Add CPU/GPU Combine/Reduction/Repartition/Replicate support to realm-execution backend - #1665
Open
seemamirch wants to merge 7 commits into
Open
Add CPU/GPU Combine/Reduction/Repartition/Replicate support to realm-execution backend #1665seemamirch wants to merge 7 commits into
seemamirch wants to merge 7 commits into
Conversation
added 7 commits
August 13, 2026 14:59
Generalizes the previously Replicate-only parallel-op handling in the dynamic-graph pipeline (pass_expansion, copy_insertion, shard_expansion) and Realm execution dispatch (pcg_instance) into a single, shared ParallelOpMovementKind abstraction (BROADCAST/GATHER/SUM_REDUCE/RESHUFFLE: 1->N copy, N->1 copy, N->1 sum, N->N reshuffle) that all four parallel ops now go through uniformly: Replicate FWD=BROADCAST BWD=SUM_REDUCE | Combine FWD=GATHER BWD=BROADCAST Reduction FWD=SUM_REDUCE BWD=BROADCAST | Repartition FWD/BWD=RESHUFFLE pass_expansion.cc: BWD expansion for all parallel ops is now the generic is_parallel_training_op case (was Replicate-only) — gradient flows in reverse with no forward activations needed, so BWD inputs are just the grad of FWD outputs and vice versa. copy_insertion.cc: resolves each value's ParallelTensorMapping either directly from the op's own node mapping (only possible on the side whose per-device coordinates are unique — the bidict requirement) or from an adjacent already-resolved value (source, for the non-unique side of BROADCAST; sink, for the non-unique side of GATHER/SUM_REDUCE). RESHUFFLE has no fixed fan-in/fan-out direction (Repartition can scatter, gather, or purely shuffle depending on its degree change), so which side is unique is checked dynamically per-invocation rather than assumed from the movement kind; it also accepts LOSS as a valid sink for the adjacent-value case, since a parallel op's FWD output can feed directly into the loss as the model's terminal output. shard_expansion.cc: RESHUFFLE splits into N separate per-device invocations (like a normal op), while BROADCAST/GATHER/SUM_REDUCE stay as a single invocation with the many-valued side represented as repeated same-slot-name entries distinguished by task_shard — matching how Replicate was already represented, per DynamicTensorSlot.task_shard's docstring. RESHUFFLE pairs each device's input/output coordinates by looking up its own shard binding directly by slot name (coordinate values can otherwise collide between the input and output coordinate spaces). Also removes shard_invocation_for_binding and restrict_tensor_mapping_keys_to_coord (dead code — never called; apply_dynamic_node_invocation_sharding_info is the real path) plus the #includes that became unused as a result. pcg_instance.cc: spawn_dynamic_node_invocation dispatches to issue_broadcast/issue_gather/issue_sum_reduce/issue_copy based on ParallelOpMovementKind rather than switching on each op's attrs type. issue_gather chains its N copies sequentially rather than firing them concurrently, since they all target the same destination instance. Combine/Reduction/Repartition never spawn a task (dispatch is pure Realm copies/reductions), so their now-dead task_id_t.cc entries return nullopt and their realm_task_registry.cc registrations are removed; a few unused lambda-parameter names in task_id_t.cc were tidied at the same time.
Renames redops/realm_redop_registry.cc to .cu (added to the CMake target as a CUDA-language source, mirroring lib/kernels), adds apply_cuda/fold_cuda kernel methods to each SumReduction<T>, and rebuilds registration around create_reduction_op<T>() + add_cuda_redop_kernels<T>() + register_reduction. Two secondary issues had to be worked around to get the new .cu file to compile under nvcc at all: PRealm's prealm.h (realm-execution/realm.h's namespace Realm = ::PRealm alias, used everywhere else in this codebase) fails a static_assert when parsed by nvcc, unrelated to reductions — the .cu file avoids it entirely and uses ::Realm:: fully-qualified throughout, so register_all_redops() also dropped its Realm::Runtime parameter (fetches ::Realm::Runtime::get_runtime() itself instead). And the borrowed-from- Legion SumReduction<bool>/SumReduction<int64_t> non-exclusive apply/fold paths called __uint2bool/__bool2uint/__longlong_as_ulonglong/ __ulonglong_as_longlong, which aren't real CUDA builtins (Legion-internal helpers) — implemented locally. Confirmed on real hardware (2x Tesla P100): all four parallel-op GPU e2e tests pass.
resolve_tensor_mappings previously only tested a pure-shuffle-shaped Repartition (both INPUT and OUTPUT unique); copy_insertion's RESHUFFLE handling determines which side is unique dynamically per-invocation, and that shape never exercised the OUTPUT-has-dup-coords branch. Adds a gather-shaped case that does. Adds GPU (cuda-realm-execution-tests) e2e training test cases for Replicate/Combine/Reduction/Repartition, and fixes the pre-existing (unrelated to this branch) base "GPU Model Parallelism" test, which requested only 1 GPU from Realm despite its config using two device coordinates.
Contributor
Author
|
@elliottslaughter @lockshaw - please review |
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.
Summary
Adds CPU and GPU Combine, Reduction, and Repartition and Replicate GPU parallel-operator support to the
realm-execution backend and support the latest FlexFlow master branch, previously Replicate-only CPU.
pass_expansion→copy_insertion→
shard_expansion) and Realm dispatch (pcg_instance) around a singleshared abstraction,
ParallelOpMovementKind(BROADCAST/GATHER/SUM_REDUCE/RESHUFFLE), so all four parallel ops go through the samecode paths instead of Replicate-specific special cases.
for all four ops
Add Combine/Reduction/Repartition support to realm-execution backend
Generalizes the previously Replicate-only parallel-op handling in the
dynamic-graph pipeline and Realm execution dispatch into
ParallelOpMovementKind(1→N copy, N→1 copy, N→1 sum, N→N reshuffle):pass_expansion.cc— BWD expansion for all parallel ops is now thegeneric
is_parallel_training_opcase (was Replicate-only): gradientflows in reverse with no forward activations needed, so BWD inputs are
just the grad of FWD outputs and vice versa.
copy_insertion.cc— resolves each value'sParallelTensorMappingshard_expansion.cc— RESHUFFLE splits into N separate per-deviceinvocations (same shape as a normal op), while BROADCAST/GATHER/
SUM_REDUCE stay as a single invocation with the many-valued side
represented as repeated same-slot-name entries distinguished by
task_shard— matching how Replicate was already representedpcg_instance.cc—spawn_dynamic_node_invocationdispatches toissue_broadcast/issue_gather/issue_sum_reduce/issue_copybasedon
ParallelOpMovementKindrather than switching on each op's attrstype.
issue_gatherchains its N copies since they all target thesame destination instance.
tasks/task_id_t.cc,tasks/realm_task_registry.cc—Combine/Reduction/Repartition never spawn a task (dispatch is pure Realm
copies/reductions), so their task-ID lookups return
nulloptand thecorresponding (dead) task registrations are removed.
Register CUDA-visible reduction kernels for realm-execution redops
redops/realm_redop_registry.cc- (added to the CMaketarget as a CUDA-language source, mirroring
lib/kernels), addedapply_cuda/fold_cudakernel methods to eachSumReduction<T>, andrebuilt registration around
create_reduction_op<T>()+add_cuda_redop_kernels<T>()+register_reduction.Confirmed on real hardware (2× Tesla P100): all four parallel-op GPU e2e
tests pass
Add parallel-op tests: gather-shaped RESHUFFLE, GPU e2e, base-test fix
copy_insertion.cctests — added a gather-shaped Repartition testcase.
resolve_tensor_mappingspreviously only tested a pure-shuffle-shaped Repartition (both INPUT and OUTPUT unique)
test_e2e.cc— added GPU (cuda-realm-execution-tests) e2e trainingtest cases for Replicate/Combine/Reduction/Repartition
Test plan
task-spec-testsrealm-execution-tests—cpu-realm-execution-tests(per test case)realm-execution-tests—cuda-realm-execution-testson real GPUhardware (per test case)
This change is