[Pipeline] Add Ray-backed stage placement and execution with pluggable transportsFeature/ray pipeline parallelism#8161
Conversation
Core abstraction: - PipelineExecutor/PipelineTransport ABCs with instruction_map dispatch - ProcessGroupExecutor: thin adapter for existing pipeline engine - NcclTransport: wraps p2p module with init guard - TcpTransport: connect-per-send + persistent mode with SocketPool - ShmTransport: shared memory for same-node CPU-CPU transfer - RayTransport: Ray object store with actor-side ref protocol - StageActor: Ray remote actor with model/optimizer/buffers/checkpointing - RayActorExecutor: per-stage Ray actor dispatch - RayTopology: placement group mapping with pure helper functions - SocketPool/SocketPoolManager: connection pool with health check, idle eviction, blocking acquire with timeout, PoolExhaustedError Key features: - NCCL auto-detection for colocated pipeline stages - Config-driven backend selection (executor/transport) - Compatibility validation in engine factory methods - Three-layer health check: SO_KEEPALIVE + MSG_PEEK + Condition - Condition-based blocking acquire with timeout - Pure helper functions testable without Ray Signed-off-by: SisPiao <piaoxiaoyu711@gmail.com>
~80 tests across 8 test files: - test_tcp_transport.py: persistent, pooled, multi-stage (3-stage), buffer correctness (GPU/CPU/dtype), stress (100 batches, 200 transfers), multi-process (spawn, importlib bypass) - test_socket_pool.py: PooledConnection, SocketPool lifecycle, idle eviction (7 boundary tests), health check edge cases, performance benchmarks, failure recovery (6 tests), blocking acquire timeout with PoolExhaustedError - test_shm_transport.py: validation + integration (float32/64/int64/scalar) - test_ray_transport.py: RayTransport + auto backend + pipeline flow - test_ray_topology.py: 372 parameterized pure helper tests - test_stage_actor.py: StageActor forward/backward/buffers/state - test_ray_engine_integration.py: ABC interface + import guards - test_ray_executor.py: RayActorExecutor dispatch + integration - conftest.py: Ray cluster fixtures Signed-off-by: SisPiao <piaoxiaoyu711@gmail.com>
examples/ray_pipeline/: - 01_simple_two_stage.py: homogeneous GPU stages with Ray object store - 02_gpu_cpu_hybrid.py: CPU embedding + GPU transformer with ShmTransport - 03_multi_accelerator.py: GPU + simulated NPU with TcpTransport - 04_moe_heterogeneous.py: 4-stage MoE on CPU/GPU-A/GPU-B/NPU - README.md: tutorial walkthrough with config reference - benchmark_tcp.py: socketpair-based transport latency benchmarks Signed-off-by: SisPiao <piaoxiaoyu711@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2578e18a6c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self._actors[stage_id] = StageActor.options(**options).remote( | ||
| stage_id=stage_id, | ||
| num_stages=self._topology.num_stages, | ||
| model=self._engine.module, | ||
| optimizer=self._engine.optimizer, |
There was a problem hiding this comment.
With executor='ray' and more than one pipeline stage, _initialize_actors() only inserts the local stage actor into self._actors before passing that map to RayTransport. Later send_activations() sends to stage_id + 1, and RayTransport.send() requires that destination stage to already be in _peer_refs, so stage 0 raises ValueError: No actor handle registered for stage 1 instead of transferring activations. The executor needs to create/discover/register the peer actor handles before executing the schedule.
Useful? React with 👍 / 👎.
| def forward_pass(self, buffer_id): | ||
| actor = self._get_actor() | ||
| ray.get(actor.forward_pass.remote(buffer_id)) |
There was a problem hiding this comment.
When executor='ray' runs on the last stage, this discards the actor's forward result and no other path updates PipelineEngine.total_loss or fwd_outputs; train_batch() then calls _aggregate_total_loss() with self.total_loss still None. The labels loaded into the actor are also never used to apply PipelineModule.loss_fn, so typical pipeline models that return logits fail during aggregation/backward or cannot report a training loss.
Useful? React with 👍 / 👎.
| t = tensor.cpu().detach().to(torch.float32).contiguous() | ||
| arr = t.numpy() | ||
| numel = arr.size | ||
| dtype_code = self._TORCH_DTYPE_CODES.get(tensor.dtype, 0) | ||
| header = struct.pack("!II", numel, dtype_code) |
There was a problem hiding this comment.
For TCP transfers of non-float32 or multidimensional activations, this casts the payload bytes to float32 while recording the original dtype, and the header only carries numel. recv() then decodes using the recorded dtype and returns a flat tensor, so float16/float64 values are corrupted or truncated and even float32 activations lose their shape. TCP serialization needs to send the actual dtype, shape, and matching byte count.
Useful? React with 👍 / 👎.
| if transport_type != 'ray': | ||
| raise ValueError(f"Ray executor requires Ray transport, not '{transport_type}'. " | ||
| "Set pipeline.transport='ray' in your DeepSpeed config.") |
There was a problem hiding this comment.
The new Ray pipeline README/examples configure pipeline.executor='ray' with transport='shm' and transport='tcp' for hybrid and multi-accelerator pipelines, but engine initialization rejects every Ray executor config except transport='ray'. Those documented configurations fail before any actors are created, so the newly added TCP/SHM transports cannot be used through PipelineEngine.
Useful? React with 👍 / 👎.
|
Hi @SisPiao, Thank you for submitting this PR! This is a very exciting contribution. Given the scope of the proposed changes, we would like to first discuss the requirements and overall design in a dedicated developer meeting before proceeding with a detailed review of the implementation. We will announce the meeting date and time soon on our Slack workspace and X. If you have any scheduling constraints, please let us know in Slack so that we can take them into consideration. You can join our Slack workspace using the following link: We look forward to discussing the proposal with you! |
2 participants
Summary
Based on the DeepSpeed pipeline parallelism heterogeneous resource track in the
Q3 2026 roadmap, this PR implements
Ray-backed pipeline stage placement and execution with pluggable transport backends,
enabling heterogeneous resource allocation across pipeline stages.
Motivation
DeepSpeed's current pipeline parallelism uses NCCL process groups within a single multi-GPU
process, which cannot support:
Ray's distributed actor model and placement groups solve all three by making each pipeline
stage an independently schedulable Ray actor.
Changes
deepspeed/runtime/pipe/):PipelineExecutorABC (how stagesexecute) and
PipelineTransportABC (how data moves), with config-driven factory selection inPipelineEngineNcclTransport(GPU↔GPU NCCL p2p),RayTransport(Ray object storewith NCCL auto-detection),
TcpTransport(TCP sockets with persistentSocketPool),ShmTransport(shared memory for CPU↔CPU)deepspeed/runtime/pipe/ray/):StageActor(Ray remote actor),RayActorExecutor(per-stage dispatch),RayTopology(placement group mapping with pure helperfunctions)
deepspeed/runtime/pipe/socket_pool.py):SocketPool/SocketPoolManagerwithMSG_PEEKhealth check, idle eviction, and blockingacquire(timeout)examples/ray_pipeline/): 4 tutorial examples from simple 2-stagehomogeneous to 4-stage MoE heterogeneous
Design
config.pipeline.executorandconfig.pipeline.transporttry/except ImportErrorguardsProcessGroupExecutordelegates to existingPipelineEngine._exec_*methods (zero refactoring)Testing
health check edge cases, failure recovery, concurrency, blocking acquire timeout, multi-stage
pipeline, multi-process stress
compute_pipe_buffershas 372 parameterized tests with property-based invariantsmp.spawn+ importlib bypass on macOS