From e250c835207a428ffdf353b8191e1f8979e9b766 Mon Sep 17 00:00:00 2001 From: ljj Date: Thu, 29 Oct 2026 11:21:48 +0000 Subject: [PATCH 1/7] feat(dsv4): seal TP SP DP fixed-order collective schedule Signed-off-by: ljj --- docs/contracts/dsv4-parallel-collectives.md | 41 +++++ .../contracts/operators/dsv4_collectives.py | 147 ++++++++++++++++ .../deepseek_v4/parallel_collectives.py | 162 ++++++++++++++++++ .../deepseek_v4/test_parallel_collectives.py | 156 +++++++++++++++++ 4 files changed, 506 insertions(+) create mode 100644 docs/contracts/dsv4-parallel-collectives.md create mode 100644 rl_engine/contracts/operators/dsv4_collectives.py create mode 100644 rl_engine/models/deepseek_v4/parallel_collectives.py create mode 100644 tests/models/deepseek_v4/test_parallel_collectives.py diff --git a/docs/contracts/dsv4-parallel-collectives.md b/docs/contracts/dsv4-parallel-collectives.md new file mode 100644 index 000000000..265df49cd --- /dev/null +++ b/docs/contracts/dsv4-parallel-collectives.md @@ -0,0 +1,41 @@ +# DSV4 fixed-order TP/SP/DP collectives (opt-in) + +`Dsv4CollectiveContract` seals the ordered communication phases, global rank +partitions, fixed input layouts, model version, memory limit, and loss divisor. +`Dsv4ParallelCollectives` executes that plan using the existing +`RCCLDeterministicCollective` on ROCm. It uses the repository's balanced-rank +tree arithmetic and rank-ordered AllGather transport rather than implementing +another numerical collective kernel. CPU/Gloo is available only by explicit +`allow_cpu_reference=True` for tests. + +The caller must create the TP, SP, and DP process groups in a common global +order, then pass the local group for each dimension. The supplied world group +contains exactly the plan's ranks; a PP stage can pass its stage-local group. +All ranks construct the same contract, call `begin_step(step_id, +model_version=...)`, execute each named phase in order, and call `end_step()`. +Control-plane signatures are exchanged before data collectives to reject +version/order/layout drift. A distributed failure poisons the runtime and +requires reconstruction. `close()` releases the existing collective resources. + +For loss normalization, only a DP FP32 AllReduce phase may declare a positive +denominator. Division is applied **after** fixed-tree reduction. The caller +supplies the intended global token/sequence denominator; this adapter does not +infer it from local batches. A change to token accounting requires a new plan. + +The current profile explicitly rejects CUDA Graph capture and unsupported +backend/group sizes/dtypes. Empty ReduceScatter and asynchronous overlap have +not been signed off. Control-plane object gathers cost latency; eliminating +them requires a separately validated static schedule and belongs to item 42. +This adapter is not yet wired into a full DSV4 optimizer/training step. + +## Validation (MI300X `rl-kernel-ljj` container) + +- Four-rank Gloo reference: TP AllGather, SP AllReduce, DP FP32 mean, + DP ReduceScatter, plus version/phase/shape mismatch rejection. +- One-rank ROCm NCCL process group: all three operations on MI300X. +- Related distributed collective tests and lint should be run before merge. +- Multi-GPU ROCm, 8+8 multi-node, fault recovery, and model-level throughput + are **not** validated: this container currently exposes one GPU. + +This PR is an incremental item-40 implementation, not full acceptance of KLR +#81. It targets `dsv4-rocm` after the branch absorbs its needed `main` commits. diff --git a/rl_engine/contracts/operators/dsv4_collectives.py b/rl_engine/contracts/operators/dsv4_collectives.py new file mode 100644 index 000000000..1324578c1 --- /dev/null +++ b/rl_engine/contracts/operators/dsv4_collectives.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Sealed rank/order contract for DSV4 TP, SP and DP collectives.""" + +from dataclasses import dataclass +import hashlib +import json + +import torch + + +class Dsv4CollectiveError(RuntimeError): + pass + + +def require(condition, message): + if not condition: + raise Dsv4CollectiveError(message) + + +@dataclass(frozen=True) +class CollectivePhase: + name: str + group: str + operation: str + shape: tuple[int, ...] + dtype: torch.dtype + denominator: int | None = None + + def __post_init__(self): + require(isinstance(self.name, str) and bool(self.name), "named phase required") + require(self.group in ("tp", "sp", "dp"), "unknown parallel group") + require( + self.operation in ("all_reduce", "all_gather", "reduce_scatter"), + "unknown collective operation", + ) + require( + isinstance(self.shape, tuple) + and bool(self.shape) + and all(type(n) is int and n >= 0 for n in self.shape), + "fixed shape required", + ) + require( + self.dtype in (torch.float32, torch.float16, torch.bfloat16), + "unvalidated collective dtype", + ) + require( + self.operation != "reduce_scatter" or self.shape[0] > 0, + "empty reduce-scatter requires separate signoff", + ) + require( + self.denominator is None + or ( + self.group == "dp" + and self.operation == "all_reduce" + and self.dtype == torch.float32 + and type(self.denominator) is int + and self.denominator > 0 + ), + "normalization requires FP32 DP all-reduce", + ) + + +@dataclass(frozen=True) +class Dsv4CollectiveContract: + global_ranks: tuple[int, ...] + tp_groups: tuple[tuple[int, ...], ...] + sp_groups: tuple[tuple[int, ...], ...] + dp_groups: tuple[tuple[int, ...], ...] + phases: tuple[CollectivePhase, ...] + model_version: str + graph: bool = False + max_size_bytes: int = 64 * 1024 * 1024 + + def __post_init__(self): + ranks = self.global_ranks + require( + isinstance(ranks, tuple) + and bool(ranks) + and all(type(r) is int and r >= 0 for r in ranks) + and tuple(sorted(set(ranks))) == ranks, + "global ranks must be sorted and unique", + ) + require( + isinstance(self.model_version, str) and bool(self.model_version), + "explicit model/scale version required", + ) + require(self.graph is False, "DSV4 collective Graph profile is not validated") + require( + type(self.max_size_bytes) is int and self.max_size_bytes > 0, + "positive collective workspace budget required", + ) + require( + isinstance(self.phases, tuple) + and bool(self.phases) + and all(isinstance(phase, CollectivePhase) for phase in self.phases) + and len({phase.name for phase in self.phases}) == len(self.phases), + "unique ordered collective phases required", + ) + for dimension in ("tp", "sp", "dp"): + groups = getattr(self, f"{dimension}_groups") + require(isinstance(groups, tuple) and bool(groups), f"{dimension} groups required") + members = [] + for group in groups: + require( + isinstance(group, tuple) + and len(group) in (1, 2, 4, 8) + and tuple(sorted(set(group))) == group, + f"{dimension} group must be sorted, unique and supported", + ) + members.extend(group) + require(tuple(sorted(members)) == ranks, f"{dimension} groups must partition ranks") + require( + len({len(group) for group in groups}) == 1, + f"{dimension} groups must have a uniform size", + ) + require( + all(any(phase.group == dim for phase in self.phases) for dim in ("tp", "sp", "dp")), + "all three dimensions need a phase", + ) + for phase in self.phases: + size = len(getattr(self, f"{phase.group}_groups")[0]) + require( + phase.operation != "reduce_scatter" or phase.shape[0] % size == 0, + f"{phase.name} leading dimension must divide {phase.group} size", + ) + numel = 1 + for dim in phase.shape: + numel *= dim + require( + numel * torch.empty((), dtype=phase.dtype).element_size() <= self.max_size_bytes, + f"{phase.name} exceeds collective capacity", + ) + + @property + def fingerprint(self): + values = { + "global_ranks": self.global_ranks, + "groups": {dim: getattr(self, f"{dim}_groups") for dim in ("tp", "sp", "dp")}, + "phases": [ + (p.name, p.group, p.operation, p.shape, str(p.dtype), p.denominator) + for p in self.phases + ], + "model_version": self.model_version, + "graph": self.graph, + "max_size_bytes": self.max_size_bytes, + } + return hashlib.sha256(json.dumps(values, sort_keys=True).encode()).hexdigest() diff --git a/rl_engine/models/deepseek_v4/parallel_collectives.py b/rl_engine/models/deepseek_v4/parallel_collectives.py new file mode 100644 index 000000000..51487c528 --- /dev/null +++ b/rl_engine/models/deepseek_v4/parallel_collectives.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Opt-in DSV4 TP/SP/DP fixed-order collective execution.""" + +import torch +import torch.distributed as dist + +from rl_engine.contracts.operators.dsv4_collectives import ( + Dsv4CollectiveContract, + require, +) +from rl_engine.distributed.transports.rccl import ( + RCCLDeterministicCollective, + TorchDistributedDeterministicCollective, +) + + +class Dsv4ParallelCollectives: + """Execute a sealed schedule using the repository's fixed-tree backends. + + The caller creates all process groups in a common global order. Construction + and each phase are collective over ``world_group``; an exception during a + distributed phase poisons this instance, so it cannot be resumed silently. + CPU/Gloo is an explicit executable reference, not a production backend. + """ + + def __init__( + self, + contract: Dsv4CollectiveContract, + *, + world_group: dist.ProcessGroup, + groups: dict[str, dist.ProcessGroup], + allow_cpu_reference: bool = False, + ) -> None: + require(isinstance(contract, Dsv4CollectiveContract), "sealed contract required") + require(dist.is_available() and dist.is_initialized(), "process group required") + require(set(groups) == {"tp", "sp", "dp"}, "TP/SP/DP groups required") + require( + tuple(dist.get_process_group_ranks(world_group)) == contract.global_ranks, + "world group differs from sealed ranks", + ) + rank = dist.get_rank() + require(rank in contract.global_ranks, "rank outside sealed group") + for dimension, group in groups.items(): + expected = next(g for g in getattr(contract, f"{dimension}_groups") if rank in g) + require( + tuple(dist.get_process_group_ranks(group)) == expected, + f"{dimension} rank partition differs from sealed plan", + ) + backend = str(dist.get_backend(world_group)).lower() + rocm = "nccl" in backend and bool(getattr(torch.version, "hip", None)) + cpu = "gloo" in backend and allow_cpu_reference + require(rocm or cpu, "unvalidated DSV4 collective transport") + for dimension, group in groups.items(): + require( + str(dist.get_backend(group)).lower() == backend, + f"{dimension} backend differs from world group", + ) + self.contract = contract + self.world_group = world_group + self.device = ( + torch.device("cuda", torch.cuda.current_device()) if rocm else torch.device("cpu") + ) + self._agree(("construct", contract.fingerprint)) + cls = RCCLDeterministicCollective if rocm else TorchDistributedDeterministicCollective + self._collectives = { + dimension: cls( + group=groups[dimension], device=self.device, max_size_bytes=contract.max_size_bytes + ) + for dimension in ("tp", "sp", "dp") + } + self._active = False + self._poisoned = False + self._closed = False + self._next = 0 + + def _agree(self, signature: tuple) -> None: + peers = [None] * dist.get_world_size(self.world_group) + dist.all_gather_object(peers, signature, group=self.world_group) + require( + all(peer == signature for peer in peers), + f"DSV4 rank schedule/version mismatch: {peers}", + ) + + def _usable(self) -> None: + require(not self._poisoned, "collective instance is closed or poisoned") + if self.device.type == "cuda": + require( + not torch.cuda.is_current_stream_capturing(), "DSV4 Graph capture is not validated" + ) + + def begin_step(self, step_id: int, *, model_version: str) -> None: + self._usable() + require(not self._active, "previous collective step not ended") + try: + self._agree(("begin", step_id, model_version, self.contract.fingerprint)) + require(type(step_id) is int and step_id >= 0, "nonnegative step id required") + require(model_version == self.contract.model_version, "stale model version") + except Exception: + self._poisoned = True + raise + self._active = True + self._next = 0 + + def execute(self, phase_name: str, value: torch.Tensor) -> torch.Tensor: + self._usable() + try: + layout = ( + (tuple(value.shape), str(value.dtype), str(value.device), value.is_contiguous()) + if isinstance(value, torch.Tensor) + else None + ) + self._agree(("phase", self._next, phase_name, layout)) + require( + self._active and self._next < len(self.contract.phases), + "no scheduled collective phase", + ) + phase = self.contract.phases[self._next] + require(phase_name == phase.name, f"expected phase {phase.name}, got {phase_name}") + require( + isinstance(value, torch.Tensor) + and tuple(value.shape) == phase.shape + and value.dtype == phase.dtype + and value.device == self.device + and value.is_contiguous(), + f"{phase.name} input differs from sealed layout", + ) + result = getattr(self._collectives[phase.group], phase.operation)(value) + if phase.denominator is not None: + result = result / phase.denominator + except Exception: + self._poisoned = True + raise + self._next += 1 + return result + + def end_step(self) -> None: + self._usable() + try: + self._agree(("end", self._next, self.contract.fingerprint)) + require( + self._active and self._next == len(self.contract.phases), + "collective step has unfinished phases", + ) + except Exception: + self._poisoned = True + raise + self._active = False + + def close(self) -> None: + if getattr(self, "_closed", False): + return + self._poisoned = True + self._closed = True + for collective in getattr(self, "_collectives", {}).values(): + collective.close() + + def __enter__(self): + self._usable() + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() diff --git a/tests/models/deepseek_v4/test_parallel_collectives.py b/tests/models/deepseek_v4/test_parallel_collectives.py new file mode 100644 index 000000000..f63ebfcdc --- /dev/null +++ b/tests/models/deepseek_v4/test_parallel_collectives.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +"""DSV4 sealed collective plan and executable Gloo reference tests.""" + +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from rl_engine.contracts.operators.dsv4_collectives import ( + CollectivePhase, + Dsv4CollectiveContract, + Dsv4CollectiveError, +) +from rl_engine.models.deepseek_v4.parallel_collectives import Dsv4ParallelCollectives + + +def _contract(): + return Dsv4CollectiveContract( + global_ranks=(0, 1, 2, 3), + tp_groups=((0, 1), (2, 3)), + sp_groups=((0,), (1,), (2,), (3,)), + dp_groups=((0, 2), (1, 3)), + phases=( + CollectivePhase("tp_gather", "tp", "all_gather", (1,), torch.float32), + CollectivePhase("sp_sum", "sp", "all_reduce", (1,), torch.float32), + CollectivePhase("dp_mean", "dp", "all_reduce", (1,), torch.float32, 2), + CollectivePhase("dp_shard", "dp", "reduce_scatter", (2,), torch.float32), + ), + model_version="checkpoint-42", + ) + + +def test_plan_rejects_invalid_groups_shapes_and_graph(): + plan = _contract() + assert plan.fingerprint == _contract().fingerprint + with pytest.raises(Dsv4CollectiveError, match="Graph"): + Dsv4CollectiveContract(**{**vars(plan), "graph": True}) + with pytest.raises(Dsv4CollectiveError, match="partition"): + Dsv4CollectiveContract(**{**vars(plan), "dp_groups": ((0, 2), (1, 2))}) + with pytest.raises(Dsv4CollectiveError, match="divide"): + Dsv4CollectiveContract( + **{ + **vars(plan), + "phases": ( + *plan.phases[:-1], + CollectivePhase("dp_shard", "dp", "reduce_scatter", (3,), torch.float32), + ), + } + ) + with pytest.raises(Dsv4CollectiveError, match="normalization"): + CollectivePhase("bad", "tp", "all_reduce", (1,), torch.float32, 2) + + +def _rank_worker(rank, rendezvous, scenario): + dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=4) + try: + plan = _contract() + groups = {} + for dimension in ("tp", "sp", "dp"): + for ranks in getattr(plan, f"{dimension}_groups"): + group = dist.new_group(ranks=list(ranks), backend="gloo") + if rank in ranks: + groups[dimension] = group + runtime = Dsv4ParallelCollectives( + plan, + world_group=dist.group.WORLD, + groups=groups, + allow_cpu_reference=True, + ) + try: + if scenario == "version_mismatch": + with pytest.raises(Dsv4CollectiveError, match="mismatch"): + runtime.begin_step(7, model_version="stale" if rank == 0 else "checkpoint-42") + with pytest.raises(Dsv4CollectiveError, match="poisoned"): + runtime.begin_step(8, model_version="checkpoint-42") + return + runtime.begin_step(7, model_version="checkpoint-42") + if scenario == "phase_mismatch": + with pytest.raises(Dsv4CollectiveError, match="mismatch"): + runtime.execute( + "wrong" if rank == 0 else "tp_gather", torch.tensor([rank + 1.0]) + ) + return + if scenario == "shape_mismatch": + with pytest.raises(Dsv4CollectiveError, match="mismatch"): + runtime.execute("tp_gather", torch.ones(2 if rank == 0 else 1)) + return + gathered = runtime.execute("tp_gather", torch.tensor([rank + 1.0])) + torch.testing.assert_close( + gathered, torch.tensor([rank // 2 * 2 + 1.0, rank // 2 * 2 + 2.0]) + ) + torch.testing.assert_close( + runtime.execute("sp_sum", torch.tensor([rank + 1.0])), torch.tensor([rank + 1.0]) + ) + torch.testing.assert_close( + runtime.execute("dp_mean", torch.tensor([rank + 1.0])), + torch.tensor([2.0 if rank % 2 == 0 else 3.0]), + ) + shard = runtime.execute("dp_shard", torch.tensor([rank + 1.0, rank + 11.0])) + expected = ( + (4.0 if rank == 0 else 24.0) if rank % 2 == 0 else (6.0 if rank == 1 else 26.0) + ) + torch.testing.assert_close(shard, torch.tensor([expected])) + runtime.end_step() + finally: + runtime.close() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize( + "scenario", ["valid", "version_mismatch", "phase_mismatch", "shape_mismatch"] +) +def test_four_rank_gloo_plan(tmp_path: Path, scenario: str): + rendezvous = str(tmp_path / f"{scenario}-rendezvous") + mp.spawn(_rank_worker, args=(rendezvous, scenario), nprocs=4, join=True) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not torch.version.hip, reason="ROCm GPU required" +) +def test_rocm_world_one(tmp_path: Path): + if dist.is_initialized(): + pytest.skip("requires an isolated process group") + torch.cuda.set_device(0) + dist.init_process_group( + "nccl", init_method=f"file://{tmp_path / 'rocm-rendezvous'}", rank=0, world_size=1 + ) + try: + plan = Dsv4CollectiveContract( + global_ranks=(0,), + tp_groups=((0,),), + sp_groups=((0,),), + dp_groups=((0,),), + phases=( + CollectivePhase("tp", "tp", "all_gather", (2,), torch.float32), + CollectivePhase("sp", "sp", "reduce_scatter", (2,), torch.float32), + CollectivePhase("dp", "dp", "all_reduce", (2,), torch.float32, 2), + ), + model_version="v1", + ) + with Dsv4ParallelCollectives( + plan, + world_group=dist.group.WORLD, + groups={key: dist.group.WORLD for key in ("tp", "sp", "dp")}, + ) as runtime: + runtime.begin_step(0, model_version="v1") + value = torch.tensor([2.0, 4.0], device="cuda") + torch.testing.assert_close(runtime.execute("tp", value), value) + torch.testing.assert_close(runtime.execute("sp", value), value) + torch.testing.assert_close(runtime.execute("dp", value), value / 2) + runtime.end_step() + finally: + dist.destroy_process_group() From fe93792ae606ed4055baa0ec5678bdb4b1da409d Mon Sep 17 00:00:00 2001 From: ljj Date: Thu, 29 Oct 2026 11:28:33 +0000 Subject: [PATCH 2/7] test(dsv4): handle empty synchronized collective phases Signed-off-by: ljj --- docs/contracts/dsv4-parallel-collectives.md | 3 ++- .../deepseek_v4/parallel_collectives.py | 13 ++++++++++- .../deepseek_v4/test_parallel_collectives.py | 23 ++++++++++++++++++- 3 files changed, 36 insertions(+), 3 deletions(-) diff --git a/docs/contracts/dsv4-parallel-collectives.md b/docs/contracts/dsv4-parallel-collectives.md index 265df49cd..9b2771eee 100644 --- a/docs/contracts/dsv4-parallel-collectives.md +++ b/docs/contracts/dsv4-parallel-collectives.md @@ -24,7 +24,8 @@ infer it from local batches. A change to token accounting requires a new plan. The current profile explicitly rejects CUDA Graph capture and unsupported backend/group sizes/dtypes. Empty ReduceScatter and asynchronous overlap have -not been signed off. Control-plane object gathers cost latency; eliminating +not been signed off. Empty AllGather and AllReduce use the synchronized phase +handshake and skip the data transport. Control-plane object gathers cost latency; eliminating them requires a separately validated static schedule and belongs to item 42. This adapter is not yet wired into a full DSV4 optimizer/training step. diff --git a/rl_engine/models/deepseek_v4/parallel_collectives.py b/rl_engine/models/deepseek_v4/parallel_collectives.py index 51487c528..db25ac06d 100644 --- a/rl_engine/models/deepseek_v4/parallel_collectives.py +++ b/rl_engine/models/deepseek_v4/parallel_collectives.py @@ -124,7 +124,18 @@ def execute(self, phase_name: str, value: torch.Tensor) -> torch.Tensor: and value.is_contiguous(), f"{phase.name} input differs from sealed layout", ) - result = getattr(self._collectives[phase.group], phase.operation)(value) + if value.numel() == 0: + # The world-wide phase handshake already established an empty input + # on every rank, so no backend transport is necessary. + group_size = self._collectives[phase.group].world_size + rows = value.size(0) + if phase.operation == "all_gather": + rows *= group_size + elif phase.operation == "reduce_scatter": + rows //= group_size + result = value.new_empty((rows, *value.shape[1:])) + else: + result = getattr(self._collectives[phase.group], phase.operation)(value) if phase.denominator is not None: result = result / phase.denominator except Exception: diff --git a/tests/models/deepseek_v4/test_parallel_collectives.py b/tests/models/deepseek_v4/test_parallel_collectives.py index f63ebfcdc..3bcef6544 100644 --- a/tests/models/deepseek_v4/test_parallel_collectives.py +++ b/tests/models/deepseek_v4/test_parallel_collectives.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 """DSV4 sealed collective plan and executable Gloo reference tests.""" +from dataclasses import replace from pathlib import Path import pytest @@ -57,6 +58,15 @@ def _rank_worker(rank, rendezvous, scenario): dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=4) try: plan = _contract() + if scenario == "empty": + plan = replace( + plan, + phases=( + CollectivePhase("tp_empty", "tp", "all_gather", (0,), torch.float32), + CollectivePhase("sp_empty", "sp", "all_reduce", (0,), torch.float32), + CollectivePhase("dp_empty", "dp", "all_reduce", (0,), torch.float32), + ), + ) groups = {} for dimension in ("tp", "sp", "dp"): for ranks in getattr(plan, f"{dimension}_groups"): @@ -77,6 +87,12 @@ def _rank_worker(rank, rendezvous, scenario): runtime.begin_step(8, model_version="checkpoint-42") return runtime.begin_step(7, model_version="checkpoint-42") + if scenario == "empty": + for phase in plan.phases: + result = runtime.execute(phase.name, torch.empty(0)) + assert tuple(result.shape) == (0,) + runtime.end_step() + return if scenario == "phase_mismatch": with pytest.raises(Dsv4CollectiveError, match="mismatch"): runtime.execute( @@ -111,13 +127,18 @@ def _rank_worker(rank, rendezvous, scenario): @pytest.mark.parametrize( - "scenario", ["valid", "version_mismatch", "phase_mismatch", "shape_mismatch"] + "scenario", ["valid", "empty", "version_mismatch", "phase_mismatch", "shape_mismatch"] ) def test_four_rank_gloo_plan(tmp_path: Path, scenario: str): rendezvous = str(tmp_path / f"{scenario}-rendezvous") mp.spawn(_rank_worker, args=(rendezvous, scenario), nprocs=4, join=True) +def test_empty_plan_rejects_reduce_scatter(): + with pytest.raises(Dsv4CollectiveError, match="empty reduce-scatter"): + CollectivePhase("empty_shard", "dp", "reduce_scatter", (0,), torch.float32) + + @pytest.mark.skipif( not torch.cuda.is_available() or not torch.version.hip, reason="ROCm GPU required" ) From 40d4d63bd082912c03516735e6fb2bdcf02ce39f Mon Sep 17 00:00:00 2001 From: ljj Date: Thu, 29 Oct 2026 12:38:06 +0000 Subject: [PATCH 3/7] fix(dsv4): reject rank-local phase drift before data collectives Signed-off-by: ljj --- .../deepseek_v4/parallel_collectives.py | 42 ++++++++++--------- .../deepseek_v4/test_parallel_collectives.py | 23 +++++++++- 2 files changed, 44 insertions(+), 21 deletions(-) diff --git a/rl_engine/models/deepseek_v4/parallel_collectives.py b/rl_engine/models/deepseek_v4/parallel_collectives.py index db25ac06d..133891951 100644 --- a/rl_engine/models/deepseek_v4/parallel_collectives.py +++ b/rl_engine/models/deepseek_v4/parallel_collectives.py @@ -73,13 +73,14 @@ def __init__( self._closed = False self._next = 0 - def _agree(self, signature: tuple) -> None: + def _agree(self, signature: tuple, valid: bool = True) -> None: peers = [None] * dist.get_world_size(self.world_group) - dist.all_gather_object(peers, signature, group=self.world_group) + dist.all_gather_object(peers, (signature, valid), group=self.world_group) require( - all(peer == signature for peer in peers), + all(peer[0] == signature for peer in peers), f"DSV4 rank schedule/version mismatch: {peers}", ) + require(all(peer[1] for peer in peers), "invalid DSV4 phase state or layout on a rank") def _usable(self) -> None: require(not self._poisoned, "collective instance is closed or poisoned") @@ -90,11 +91,14 @@ def _usable(self) -> None: def begin_step(self, step_id: int, *, model_version: str) -> None: self._usable() - require(not self._active, "previous collective step not ended") try: - self._agree(("begin", step_id, model_version, self.contract.fingerprint)) - require(type(step_id) is int and step_id >= 0, "nonnegative step id required") - require(model_version == self.contract.model_version, "stale model version") + valid = ( + not self._active + and type(step_id) is int + and step_id >= 0 + and model_version == self.contract.model_version + ) + self._agree(("begin", step_id, model_version, self.contract.fingerprint), valid) except Exception: self._poisoned = True raise @@ -109,21 +113,20 @@ def execute(self, phase_name: str, value: torch.Tensor) -> torch.Tensor: if isinstance(value, torch.Tensor) else None ) - self._agree(("phase", self._next, phase_name, layout)) - require( - self._active and self._next < len(self.contract.phases), - "no scheduled collective phase", + phase = ( + self.contract.phases[self._next] if self._next < len(self.contract.phases) else None ) - phase = self.contract.phases[self._next] - require(phase_name == phase.name, f"expected phase {phase.name}, got {phase_name}") - require( - isinstance(value, torch.Tensor) + valid = ( + self._active + and phase is not None + and phase_name == phase.name + and isinstance(value, torch.Tensor) and tuple(value.shape) == phase.shape and value.dtype == phase.dtype and value.device == self.device - and value.is_contiguous(), - f"{phase.name} input differs from sealed layout", + and value.is_contiguous() ) + self._agree(("phase", self._next, phase_name, layout), valid) if value.numel() == 0: # The world-wide phase handshake already established an empty input # on every rank, so no backend transport is necessary. @@ -147,10 +150,9 @@ def execute(self, phase_name: str, value: torch.Tensor) -> torch.Tensor: def end_step(self) -> None: self._usable() try: - self._agree(("end", self._next, self.contract.fingerprint)) - require( + self._agree( + ("end", self._next, self.contract.fingerprint), self._active and self._next == len(self.contract.phases), - "collective step has unfinished phases", ) except Exception: self._poisoned = True diff --git a/tests/models/deepseek_v4/test_parallel_collectives.py b/tests/models/deepseek_v4/test_parallel_collectives.py index 3bcef6544..3619de18a 100644 --- a/tests/models/deepseek_v4/test_parallel_collectives.py +++ b/tests/models/deepseek_v4/test_parallel_collectives.py @@ -80,6 +80,12 @@ def _rank_worker(rank, rendezvous, scenario): allow_cpu_reference=True, ) try: + if scenario == "begin_active": + if rank == 0: + runtime._active = True + with pytest.raises(Dsv4CollectiveError, match="invalid"): + runtime.begin_step(7, model_version="checkpoint-42") + return if scenario == "version_mismatch": with pytest.raises(Dsv4CollectiveError, match="mismatch"): runtime.begin_step(7, model_version="stale" if rank == 0 else "checkpoint-42") @@ -87,6 +93,12 @@ def _rank_worker(rank, rendezvous, scenario): runtime.begin_step(8, model_version="checkpoint-42") return runtime.begin_step(7, model_version="checkpoint-42") + if scenario == "execute_inactive": + if rank == 0: + runtime._active = False + with pytest.raises(Dsv4CollectiveError, match="invalid"): + runtime.execute("tp_gather", torch.tensor([rank + 1.0])) + return if scenario == "empty": for phase in plan.phases: result = runtime.execute(phase.name, torch.empty(0)) @@ -127,7 +139,16 @@ def _rank_worker(rank, rendezvous, scenario): @pytest.mark.parametrize( - "scenario", ["valid", "empty", "version_mismatch", "phase_mismatch", "shape_mismatch"] + "scenario", + [ + "valid", + "empty", + "version_mismatch", + "phase_mismatch", + "shape_mismatch", + "begin_active", + "execute_inactive", + ], ) def test_four_rank_gloo_plan(tmp_path: Path, scenario: str): rendezvous = str(tmp_path / f"{scenario}-rendezvous") From de0592534e58c5e613a485c71d9eb1a3a2216255 Mon Sep 17 00:00:00 2001 From: ljj Date: Fri, 30 Oct 2026 12:26:49 +0000 Subject: [PATCH 4/7] fix(dsv4): synchronize collective construction failures Signed-off-by: ljj --- .../deepseek_v4/parallel_collectives.py | 75 +++++++++++-------- 1 file changed, 44 insertions(+), 31 deletions(-) diff --git a/rl_engine/models/deepseek_v4/parallel_collectives.py b/rl_engine/models/deepseek_v4/parallel_collectives.py index 133891951..09b724a97 100644 --- a/rl_engine/models/deepseek_v4/parallel_collectives.py +++ b/rl_engine/models/deepseek_v4/parallel_collectives.py @@ -33,41 +33,54 @@ def __init__( ) -> None: require(isinstance(contract, Dsv4CollectiveContract), "sealed contract required") require(dist.is_available() and dist.is_initialized(), "process group required") - require(set(groups) == {"tp", "sp", "dp"}, "TP/SP/DP groups required") - require( - tuple(dist.get_process_group_ranks(world_group)) == contract.global_ranks, - "world group differs from sealed ranks", - ) - rank = dist.get_rank() - require(rank in contract.global_ranks, "rank outside sealed group") - for dimension, group in groups.items(): - expected = next(g for g in getattr(contract, f"{dimension}_groups") if rank in g) + self.contract = contract + self.world_group = world_group + local_error = None + rocm = False + try: + require(set(groups) == {"tp", "sp", "dp"}, "TP/SP/DP groups required") require( - tuple(dist.get_process_group_ranks(group)) == expected, - f"{dimension} rank partition differs from sealed plan", + tuple(dist.get_process_group_ranks(world_group)) == contract.global_ranks, + "world group differs from sealed ranks", ) - backend = str(dist.get_backend(world_group)).lower() - rocm = "nccl" in backend and bool(getattr(torch.version, "hip", None)) - cpu = "gloo" in backend and allow_cpu_reference - require(rocm or cpu, "unvalidated DSV4 collective transport") - for dimension, group in groups.items(): - require( - str(dist.get_backend(group)).lower() == backend, - f"{dimension} backend differs from world group", + rank = dist.get_rank() + require(rank in contract.global_ranks, "rank outside sealed group") + for dimension, group in groups.items(): + expected = next(g for g in getattr(contract, f"{dimension}_groups") if rank in g) + require( + tuple(dist.get_process_group_ranks(group)) == expected, + f"{dimension} rank partition differs from sealed plan", + ) + backend = str(dist.get_backend(world_group)).lower() + rocm = "nccl" in backend and bool(getattr(torch.version, "hip", None)) + cpu = "gloo" in backend and allow_cpu_reference + require(rocm or cpu, "unvalidated DSV4 collective transport") + for dimension, group in groups.items(): + require( + str(dist.get_backend(group)).lower() == backend, + f"{dimension} backend differs from world group", + ) + self.device = ( + torch.device("cuda", torch.cuda.current_device()) if rocm else torch.device("cpu") ) - self.contract = contract - self.world_group = world_group - self.device = ( - torch.device("cuda", torch.cuda.current_device()) if rocm else torch.device("cpu") - ) - self._agree(("construct", contract.fingerprint)) + except Exception as exc: + local_error = exc + self.device = torch.device("cpu") + self._agree(("construct", contract.fingerprint), local_error is None) + if local_error is not None: + raise local_error cls = RCCLDeterministicCollective if rocm else TorchDistributedDeterministicCollective - self._collectives = { - dimension: cls( - group=groups[dimension], device=self.device, max_size_bytes=contract.max_size_bytes - ) - for dimension in ("tp", "sp", "dp") - } + self._collectives = {} + try: + for dimension in ("tp", "sp", "dp"): + self._collectives[dimension] = cls( + group=groups[dimension], + device=self.device, + max_size_bytes=contract.max_size_bytes, + ) + except Exception: + self.close() + raise self._active = False self._poisoned = False self._closed = False From 1c314189bb3dd4c8b5309e8e9c5a1b7bf195b6eb Mon Sep 17 00:00:00 2001 From: ljj Date: Mon, 2 Nov 2026 11:42:42 +0000 Subject: [PATCH 5/7] fix(dsv4): preserve rank-local construction error Signed-off-by: ljj --- .../models/deepseek_v4/parallel_collectives.py | 9 ++++++--- .../models/deepseek_v4/test_parallel_collectives.py | 13 +++++++++++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/rl_engine/models/deepseek_v4/parallel_collectives.py b/rl_engine/models/deepseek_v4/parallel_collectives.py index 09b724a97..630130179 100644 --- a/rl_engine/models/deepseek_v4/parallel_collectives.py +++ b/rl_engine/models/deepseek_v4/parallel_collectives.py @@ -66,9 +66,12 @@ def __init__( except Exception as exc: local_error = exc self.device = torch.device("cpu") - self._agree(("construct", contract.fingerprint), local_error is None) - if local_error is not None: - raise local_error + try: + self._agree(("construct", contract.fingerprint), local_error is None) + except Exception as agreement_error: + if local_error is not None: + raise local_error from agreement_error + raise cls = RCCLDeterministicCollective if rocm else TorchDistributedDeterministicCollective self._collectives = {} try: diff --git a/tests/models/deepseek_v4/test_parallel_collectives.py b/tests/models/deepseek_v4/test_parallel_collectives.py index 3619de18a..1264107de 100644 --- a/tests/models/deepseek_v4/test_parallel_collectives.py +++ b/tests/models/deepseek_v4/test_parallel_collectives.py @@ -73,6 +73,18 @@ def _rank_worker(rank, rendezvous, scenario): group = dist.new_group(ranks=list(ranks), backend="gloo") if rank in ranks: groups[dimension] = group + if scenario == "construct_partition" and rank == 0: + groups["tp"] = groups["dp"] + if scenario == "construct_partition": + with pytest.raises(Dsv4CollectiveError) as error: + Dsv4ParallelCollectives( + plan, + world_group=dist.group.WORLD, + groups=groups, + allow_cpu_reference=True, + ) + assert ("tp rank partition" if rank == 0 else "invalid DSV4") in str(error.value) + return runtime = Dsv4ParallelCollectives( plan, world_group=dist.group.WORLD, @@ -148,6 +160,7 @@ def _rank_worker(rank, rendezvous, scenario): "shape_mismatch", "begin_active", "execute_inactive", + "construct_partition", ], ) def test_four_rank_gloo_plan(tmp_path: Path, scenario: str): From 8061e9331b5c35926ac7094ed44da7aff2b78769 Mon Sep 17 00:00:00 2001 From: ljj Date: Mon, 2 Nov 2026 11:54:40 +0000 Subject: [PATCH 6/7] fix(dsv4): close rank-local transport initialization gaps Signed-off-by: ljj --- rl_engine/distributed/transports/rccl.py | 77 +++++++++++++------ .../deepseek_v4/parallel_collectives.py | 17 ++-- ...test_transport_deterministic_collective.py | 5 +- .../deepseek_v4/test_parallel_collectives.py | 75 ++++++++++++++++++ 4 files changed, 142 insertions(+), 32 deletions(-) diff --git a/rl_engine/distributed/transports/rccl.py b/rl_engine/distributed/transports/rccl.py index a83866430..b2ac8eed2 100644 --- a/rl_engine/distributed/transports/rccl.py +++ b/rl_engine/distributed/transports/rccl.py @@ -54,20 +54,24 @@ def __init__( ) -> None: if not dist.is_available() or not dist.is_initialized(): raise RuntimeError("torch.distributed must be initialized before collectives") - if max_size_bytes <= 0: - raise ValueError("max_size_bytes must be positive") - self.group = group if group is not None else dist.group.WORLD self.rank = int(dist.get_rank(group=self.group)) self.world_size = int(dist.get_world_size(group=self.group)) - if self.world_size not in _SUPPORTED_WORLD_SIZES: - raise ValueError( - "deterministic collectives require world_size in " - f"{_SUPPORTED_WORLD_SIZES}, got {self.world_size}" - ) - - self.device = self._normalize_device(device) - self.max_size_bytes = int(max_size_bytes) + local_error = None + try: + if max_size_bytes <= 0: + raise ValueError("max_size_bytes must be positive") + if self.world_size not in _SUPPORTED_WORLD_SIZES: + raise ValueError( + "deterministic collectives require world_size in " + f"{_SUPPORTED_WORLD_SIZES}, got {self.world_size}" + ) + self.device = self._normalize_device(device) + self.max_size_bytes = int(max_size_bytes) + except Exception as exc: + local_error = exc + self.device = torch.device("cpu") + self.max_size_bytes = max_size_bytes self._backend = str(dist.get_backend(self.group)).lower() self._lock = threading.Lock() self._closed = False @@ -84,7 +88,7 @@ def __init__( # the standard collective contract that ranks call operations in the # same order. self._validated_signatures: set[tuple[Any, ...]] = set() - self._validate_matching_capacity() + self._validate_matching_capacity(local_error) @staticmethod def _normalize_device( @@ -535,12 +539,19 @@ def _validate_matching_signature(self, op_name: str, input: torch.Tensor) -> Non ) self._validated_signatures.add(signature) - def _validate_matching_capacity(self) -> None: + def _validate_matching_capacity(self, local_error: Exception | None = None) -> None: if self.world_size == 1: + if local_error is not None: + raise local_error return - capacities: list[int | None] = [None] * self.world_size - dist.all_gather_object(capacities, self.max_size_bytes, group=self.group) - if any(peer_capacity != self.max_size_bytes for peer_capacity in capacities): + states = [None] * self.world_size + dist.all_gather_object(states, (self.max_size_bytes, local_error is None), group=self.group) + if local_error is not None: + raise local_error + if not all(valid for _, valid in states): + raise ValueError("invalid deterministic collective construction on a rank") + capacities = [capacity for capacity, _ in states] + if any(capacity != self.max_size_bytes for capacity in capacities): raise ValueError(f"all ranks must use the same max_size_bytes; got {capacities}") def _all_gather_transport( @@ -713,10 +724,13 @@ def workspace_size_bytes(self) -> int: def _initialize_ipc_transport(self) -> None: if self.world_size == 1: return + native = None try: from rl_engine import _C except ImportError: - return + pass + else: + native = _C required_symbols = ( "deterministic_collective_rocm_ipc_allocate", "deterministic_collective_rocm_ipc_meta", @@ -734,21 +748,34 @@ def _initialize_ipc_transport(self) -> None: "deterministic_collective_rocm_ipc_all_gather", "deterministic_collective_rocm_ipc_all_gather_input", ) - if any(not hasattr(_C, symbol) for symbol in required_symbols): + available = native is not None and all( + hasattr(native, symbol) for symbol in required_symbols + ) + availability = [None] * self.world_size + dist.all_gather_object(availability, available, group=self.group) + if not all(availability): return - staging = _C.deterministic_collective_rocm_ipc_allocate(self.max_size_bytes) - handle, offset = _C.deterministic_collective_rocm_ipc_meta(staging) - local_metadata = (socket.gethostname(), handle, int(offset)) - gathered_metadata: list[tuple[str, list[int], int] | None] = [None] * self.world_size + local_error = None + local_metadata = None + try: + staging = native.deterministic_collective_rocm_ipc_allocate(self.max_size_bytes) + handle, offset = native.deterministic_collective_rocm_ipc_meta(staging) + local_metadata = (socket.gethostname(), handle, int(offset)) + except Exception as exc: + local_error = exc + gathered_metadata = [None] * self.world_size dist.all_gather_object(gathered_metadata, local_metadata, group=self.group) if any(metadata is None for metadata in gathered_metadata): - raise RuntimeError("failed to exchange ROCm IPC metadata") - complete_metadata = [metadata for metadata in gathered_metadata if metadata is not None] + error = RuntimeError("failed to allocate or exchange ROCm IPC metadata on a rank") + if local_error is not None: + raise error from local_error + raise error + complete_metadata = gathered_metadata if len({metadata[0] for metadata in complete_metadata}) != 1: return self._ipc_handle = int( - _C.deterministic_collective_rocm_ipc_create( + native.deterministic_collective_rocm_ipc_create( staging, [metadata[1] for metadata in complete_metadata], [metadata[2] for metadata in complete_metadata], diff --git a/rl_engine/models/deepseek_v4/parallel_collectives.py b/rl_engine/models/deepseek_v4/parallel_collectives.py index 630130179..9c584ad7c 100644 --- a/rl_engine/models/deepseek_v4/parallel_collectives.py +++ b/rl_engine/models/deepseek_v4/parallel_collectives.py @@ -74,16 +74,23 @@ def __init__( raise cls = RCCLDeterministicCollective if rocm else TorchDistributedDeterministicCollective self._collectives = {} - try: - for dimension in ("tp", "sp", "dp"): + for dimension in ("tp", "sp", "dp"): + local_error = None + try: self._collectives[dimension] = cls( group=groups[dimension], device=self.device, max_size_bytes=contract.max_size_bytes, ) - except Exception: - self.close() - raise + except Exception as exc: + local_error = exc + try: + self._agree(("transport", dimension, contract.fingerprint), local_error is None) + except Exception as agreement_error: + self.close() + if local_error is not None: + raise local_error from agreement_error + raise self._active = False self._poisoned = False self._closed = False diff --git a/tests/distributed/collectives/test_transport_deterministic_collective.py b/tests/distributed/collectives/test_transport_deterministic_collective.py index edbd0cdc6..77521670a 100644 --- a/tests/distributed/collectives/test_transport_deterministic_collective.py +++ b/tests/distributed/collectives/test_transport_deterministic_collective.py @@ -65,8 +65,9 @@ def get_backend(self, group: Any) -> str: def all_gather_object(self, output: list[Any], value: Any, *, group: Any) -> None: self.object_gather_calls += 1 - if isinstance(value, int): - values = self.peer_capacities or [value] * len(self.peer_inputs) + if isinstance(value, tuple) and len(value) == 2 and isinstance(value[1], bool): + capacities = self.peer_capacities or [value[0]] * len(self.peer_inputs) + values = [(capacity, value[1]) for capacity in capacities] else: values = self.peer_signatures or [value] * len(self.peer_inputs) output[:] = values diff --git a/tests/models/deepseek_v4/test_parallel_collectives.py b/tests/models/deepseek_v4/test_parallel_collectives.py index 1264107de..53484d982 100644 --- a/tests/models/deepseek_v4/test_parallel_collectives.py +++ b/tests/models/deepseek_v4/test_parallel_collectives.py @@ -3,6 +3,7 @@ from dataclasses import replace from pathlib import Path +from unittest.mock import patch import pytest import torch @@ -15,6 +16,7 @@ Dsv4CollectiveError, ) from rl_engine.models.deepseek_v4.parallel_collectives import Dsv4ParallelCollectives +from rl_engine.distributed.transports.rccl import RCCLDeterministicCollective def _contract(): @@ -73,6 +75,76 @@ def _rank_worker(rank, rendezvous, scenario): group = dist.new_group(ranks=list(ranks), backend="gloo") if rank in ranks: groups[dimension] = group + if scenario == "ipc_metadata_failure": + import rl_engine + + class FakeNative: + def __getattr__(self, name): + if name == "deterministic_collective_rocm_ipc_allocate": + + def allocate(_size): + if rank == 0: + raise RuntimeError("injected IPC allocation failure") + return torch.empty(1) + + return allocate + if name == "deterministic_collective_rocm_ipc_meta": + return lambda _staging: ([0], 0) + return lambda *args: 0 + + transport = RCCLDeterministicCollective.__new__(RCCLDeterministicCollective) + transport.world_size = 4 + transport.rank = rank + transport.group = dist.group.WORLD + transport.max_size_bytes = 1024 + with patch.object(rl_engine, "_C", FakeNative(), create=True): + with pytest.raises(RuntimeError, match="IPC metadata"): + transport._initialize_ipc_transport() + return + if scenario == "transport_failure": + import rl_engine.models.deepseek_v4.parallel_collectives as module + + class FailingTransport(module.TorchDistributedDeterministicCollective): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if rank == 0: + raise RuntimeError("injected transport failure") + + with patch.object(module, "TorchDistributedDeterministicCollective", FailingTransport): + error_type = RuntimeError if rank == 0 else Dsv4CollectiveError + with pytest.raises(error_type, match="injected transport|invalid DSV4"): + Dsv4ParallelCollectives( + plan, + world_group=dist.group.WORLD, + groups=groups, + allow_cpu_reference=True, + ) + return + if scenario == "transport_preflight": + import rl_engine.models.deepseek_v4.parallel_collectives as module + + original = module.TorchDistributedDeterministicCollective._normalize_device + + def normalize(device): + if rank == 0: + raise RuntimeError("injected device validation failure") + return original(device) + + with patch.object( + module.TorchDistributedDeterministicCollective, + "_normalize_device", + staticmethod(normalize), + ): + with pytest.raises((RuntimeError, ValueError, Dsv4CollectiveError)) as error: + Dsv4ParallelCollectives( + plan, + world_group=dist.group.WORLD, + groups=groups, + allow_cpu_reference=True, + ) + if rank == 0: + assert "injected device validation failure" in str(error.value) + return if scenario == "construct_partition" and rank == 0: groups["tp"] = groups["dp"] if scenario == "construct_partition": @@ -161,6 +233,9 @@ def _rank_worker(rank, rendezvous, scenario): "begin_active", "execute_inactive", "construct_partition", + "transport_failure", + "transport_preflight", + "ipc_metadata_failure", ], ) def test_four_rank_gloo_plan(tmp_path: Path, scenario: str): From 757da3456b0ee8ae3020de5351d04f70c851cab1 Mon Sep 17 00:00:00 2001 From: ljj Date: Mon, 2 Nov 2026 13:28:05 +0000 Subject: [PATCH 7/7] fix(dsv4): compare rank-local ROCm devices by type Signed-off-by: ljj --- docs/contracts/dsv4-parallel-collectives.md | 11 +++- .../deepseek_v4/parallel_collectives.py | 2 +- .../deepseek_v4/test_parallel_collectives.py | 59 +++++++++++++++++++ 3 files changed, 69 insertions(+), 3 deletions(-) diff --git a/docs/contracts/dsv4-parallel-collectives.md b/docs/contracts/dsv4-parallel-collectives.md index 9b2771eee..f8ac4ff1c 100644 --- a/docs/contracts/dsv4-parallel-collectives.md +++ b/docs/contracts/dsv4-parallel-collectives.md @@ -34,9 +34,16 @@ This adapter is not yet wired into a full DSV4 optimizer/training step. - Four-rank Gloo reference: TP AllGather, SP AllReduce, DP FP32 mean, DP ReduceScatter, plus version/phase/shape mismatch rejection. - One-rank ROCm NCCL process group: all three operations on MI300X. +- Four-rank ROCm NCCL: TP2/SP1/DP2 ordered phases, two successive steps, + FP32 DP normalization, ReduceScatter shards, and cross-rank stale-version + rejection on four MI300X cards. Run the isolated pytest with + `ROCR_VISIBLE_DEVICES=0,1,2,3` (the container's default is one card). +- The underlying IPC transport passes the mixed eager/Graph check at TP1/2/4/8 + on eight MI300X cards; this does **not** enable Graph for the DSV4 adapter, + whose object-gather control plane remains capture-incompatible. - Related distributed collective tests and lint should be run before merge. -- Multi-GPU ROCm, 8+8 multi-node, fault recovery, and model-level throughput - are **not** validated: this container currently exposes one GPU. +- The full DSV4 adapter on TP8, 8+8 multi-node, fault recovery, Graph, and + model-level throughput remain **unvalidated**. This PR is an incremental item-40 implementation, not full acceptance of KLR #81. It targets `dsv4-rocm` after the branch absorbs its needed `main` commits. diff --git a/rl_engine/models/deepseek_v4/parallel_collectives.py b/rl_engine/models/deepseek_v4/parallel_collectives.py index 9c584ad7c..9aec5b1a2 100644 --- a/rl_engine/models/deepseek_v4/parallel_collectives.py +++ b/rl_engine/models/deepseek_v4/parallel_collectives.py @@ -132,7 +132,7 @@ def execute(self, phase_name: str, value: torch.Tensor) -> torch.Tensor: self._usable() try: layout = ( - (tuple(value.shape), str(value.dtype), str(value.device), value.is_contiguous()) + (tuple(value.shape), str(value.dtype), value.device.type, value.is_contiguous()) if isinstance(value, torch.Tensor) else None ) diff --git a/tests/models/deepseek_v4/test_parallel_collectives.py b/tests/models/deepseek_v4/test_parallel_collectives.py index 53484d982..4c39c23e4 100644 --- a/tests/models/deepseek_v4/test_parallel_collectives.py +++ b/tests/models/deepseek_v4/test_parallel_collectives.py @@ -284,3 +284,62 @@ def test_rocm_world_one(tmp_path: Path): runtime.end_step() finally: dist.destroy_process_group() + + +def _rocm_four_rank_worker(rank: int, rendezvous: str): + torch.cuda.set_device(rank) + dist.init_process_group("nccl", init_method=f"file://{rendezvous}", rank=rank, world_size=4) + try: + plan = _contract() + groups = {} + for dimension in ("tp", "sp", "dp"): + for ranks in getattr(plan, f"{dimension}_groups"): + group = dist.new_group(ranks=list(ranks), backend="nccl") + if rank in ranks: + groups[dimension] = group + with Dsv4ParallelCollectives(plan, world_group=dist.group.WORLD, groups=groups) as runtime: + for step in (0, 1): + runtime.begin_step(step, model_version="checkpoint-42") + value = torch.tensor([float(rank + 1 + step)], device="cuda") + torch.testing.assert_close( + runtime.execute("tp_gather", value), + torch.tensor( + [rank // 2 * 2 + 1.0 + step, rank // 2 * 2 + 2.0 + step], device="cuda" + ), + ) + torch.testing.assert_close(runtime.execute("sp_sum", value), value) + torch.testing.assert_close( + runtime.execute("dp_mean", value), + torch.tensor([2.0 + rank % 2 + step], device="cuda"), + ) + shard = runtime.execute( + "dp_shard", torch.tensor([rank + 1.0 + step, rank + 11.0 + step], device="cuda") + ) + expected = ( + (4.0 + 2 * step if rank == 0 else 24.0 + 2 * step) + if rank % 2 == 0 + else (6.0 + 2 * step if rank == 1 else 26.0 + 2 * step) + ) + torch.testing.assert_close(shard, torch.tensor([expected], device="cuda")) + runtime.end_step() + with pytest.raises(Dsv4CollectiveError, match="mismatch"): + runtime.begin_step(2, model_version="stale" if rank == 0 else "checkpoint-42") + with pytest.raises(Dsv4CollectiveError, match="poisoned"): + runtime.begin_step(3, model_version="checkpoint-42") + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not torch.version.hip or torch.cuda.device_count() < 4, + reason="four ROCm GPUs required", +) +def test_rocm_four_rank_tp_sp_dp_plan(tmp_path: Path): + if dist.is_initialized(): + pytest.skip("requires an isolated process group") + mp.spawn( + _rocm_four_rank_worker, + args=(str(tmp_path / "rocm-four-rank-rendezvous"),), + nprocs=4, + join=True, + )