Synchronize caller PyTorch randomness across TrainerRank replicas

scripts/ci/trainer-rank-gpu-tests.sh

@@ -10,6 +10,7 @@ test -x "${runtime_python}"
 
 "${runtime_python}" -m pytest --tb=short \
   tests/unit/test_trainer_rank_head_recompute.py \
+  tests/unit/test_trainer_rank_rng.py \
   tests/unit/test_trainer_rank_custom_tensors.py \
   tests/integration/megatron/cp_attn/test_attention_packed_vs_flattened.py \
   'tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_packed_correctness.py::test_gdn_cp_packed_sibling_order_matches_cp1_oracle[2]' \

src/art/trainer_rank/__init__.py

@@ -259,6 +259,20 @@ def forward_micro_batches(
         without multiplying them by the number of replicas. `dp_reduce` combines
         only distinct data-parallel batches.
 
+        ART isolates model PyTorch RNG consumption and synchronizes the default
+        CPU and trainer-device CUDA generators within TP/CP before each yield
+        (or `dp_rank_forward` return), using that DP worker's first TP/CP rank.
+        Matching caller-side random masks and custom-head dropout therefore need
+        no parallelism-specific seeding. Streams advance normally, and different
+        DP workers keep their own state; identical caller seeds are not changed.
+        Custom object registration also synchronizes these generators before
+        invoking its factory. Python/NumPy RNGs, explicit generators, other CUDA
+        devices, concurrent RNG use, and rank-dependent control flow are outside
+        this contract. Caller code must still perform matching operations across
+        TP/CP. Checkpoint saves contain weights/optimizer state, not RNG state;
+        exact replay requires caller RNG restoration and the same model execution
+        history. Model activation checkpointing must preserve RNG state.
+
         Empty local microbatches are skipped unless `yield_empty=True`. Every
         rank must use the same setting. When a wave skips ranks, TrainerRank
         collective methods raise if called from its loop body; fully populated

src/art/trainer_rank/_impl.py

@@ -67,6 +67,7 @@
     prefix_tree_layout_candidates,
     select_prefix_tree_layout,
 )
+from art.trainer_rank._rng import TrainerRNG, caller_group
 from art.trainer_rank._telemetry import phase as _telemetry_phase
 
 if TYPE_CHECKING:
@@ -1366,6 +1367,7 @@ def __init__(self, runtime: TrainingRuntime) -> None:
         # ignores sharding (conservative).
         self.runtime: TrainingRuntime = runtime
         self.device: torch.device = next(runtime.model[0].parameters()).device
+        self._rng = TrainerRNG(self.device)
         self._param_dtype_size = _dtype_size(next(runtime.model[0].parameters()).dtype)
         try:
             metadata_model = _language_model(runtime.model[0])
@@ -1580,6 +1582,7 @@ def _custom_object(
             raise TrainerRankSlotStateError(
                 "Custom checkpoint object registration differs across ranks"
             )
+        self._rng.synchronize(caller_group())
         slot = self._checkpoint_slots[checkpoint_name]
         existing = slot.custom.get(name)
         registered = None if existing is None else existing.kind
@@ -2250,20 +2253,24 @@ def forward_micro_batches(
         if not isinstance(yield_empty, bool):
             raise TypeError("yield_empty must be a bool")
         enabled = torch.is_grad_enabled() if no_grad is None else not no_grad
+        self._guard_forward_collective("forward_micro_batches")
+        with torch.set_grad_enabled(enabled):
+            items = [_materialize(item) for item in inputs]
         batches = self._forward_micro_batches(
-            inputs, checkpoint=checkpoint, yield_empty=yield_empty
+            items, checkpoint=checkpoint, yield_empty=yield_empty
         )
         token = object()
         try:
             while True:
                 self._guard_forward_collective("forward_micro_batches")
-                with torch.set_grad_enabled(enabled):
+                with torch.set_grad_enabled(enabled), self._rng.model():
                     try:
                         batch = next(batches)
                     except StopIteration:
                         return
                 if not yield_empty and not batch.outputs:
                     continue
+                self._rng.synchronize(caller_group())
                 if (
                     not yield_empty
                     and batch.stats.global_count < self._dp_rank_and_size()[1]
@@ -2293,12 +2300,11 @@ def _guard_forward_collective(self, operation: str) -> None:
 
     def _forward_micro_batches(
         self,
-        inputs: Iterable[ForwardInputs],
+        items: Sequence[ForwardInputs],
         *,
         checkpoint: AdapterSelection,
         yield_empty: bool,
     ) -> Generator[MicroBatch[ForwardInputs, ForwardOutputs], None, None]:
-        items = [_materialize(item) for item in inputs]
         requests = list(_flatten(items))
         self._validate_replicated_top_level_count(len(items), yield_empty=yield_empty)
         for _, indices in self._group_active_request_indices(
@@ -2451,17 +2457,22 @@ def dp_rank_forward(
     ) -> ForwardOutputs:
         self._guard_forward_collective("dp_rank_forward")
         enabled = torch.is_grad_enabled() if no_grad is None else not no_grad
+        # Iterating caller inputs can consume their RNG (e.g. a data loader).
+        # Only ART's planning/model work belongs to the private stream.
         with torch.set_grad_enabled(enabled):
-            self._reset_planning_telemetry()
             materialized = _materialize(inputs)
             requests = list(_flatten(materialized))
+        with torch.set_grad_enabled(enabled), self._rng.model():
+            self._reset_planning_telemetry()
             plan, check = self._plan_admissible_forward(
                 requests, checkpoint=checkpoint, context="dp_rank_forward"
             )
             tracked_outputs = self._execute_admitted_plan(
                 plan, check=check, context="dp_rank_forward"
             )
-            return _unflatten(materialized, iter(tracked_outputs))
+            outputs = _unflatten(materialized, iter(tracked_outputs))
+        self._rng.synchronize(caller_group())
+        return outputs
 
     def _execute_admitted_plan(
         self, plan: _AnyForwardPlan, *, check: _MemoryCheck, context: str

src/art/trainer_rank/_rng.py

@@ -0,0 +1,107 @@
+"""Separate model RNG consumption from the replicated caller's PyTorch stream."""
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+import hashlib
+
+import torch
+import torch.distributed as dist
+
+
+@dataclass
+class _State:
+    cpu: torch.Tensor
+    cuda: torch.Tensor | None
+
+    @classmethod
+    def capture(cls, device: torch.device) -> "_State":
+        return cls(
+            torch.get_rng_state(),
+            torch.cuda.get_rng_state(device) if device.type == "cuda" else None,
+        )
+
+    def restore(self, device: torch.device) -> None:
+        torch.set_rng_state(self.cpu)
+        if self.cuda is not None:
+            torch.cuda.set_rng_state(self.cuda, device)
+
+    def model_stream(self, device: torch.device) -> "_State":
+        def derive(state: torch.Tensor, target: torch.device | str) -> torch.Tensor:
+            digest = hashlib.sha256(b"art.trainer_rank.model" + state.numpy().tobytes())
+            seed = int.from_bytes(digest.digest()[:8], "little")
+            return torch.Generator(device=target).manual_seed(seed).get_state()
+
+        return _State(
+            derive(self.cpu, "cpu"),
+            derive(self.cuda, device) if self.cuda is not None else None,
+        )
+
+
+class TrainerRNG:
+    def __init__(self, device: torch.device) -> None:
+        self.device = device
+        self._model: _State | None = None
+        self._depth = 0
+
+    @contextmanager
+    def model(self) -> Iterator[None]:
+        """Advance a private model stream and restore the caller even on error.
+
+        Megatron's separate model-parallel RNG tracker is deliberately untouched.
+        Native torch/Megatron checkpoints capture these forward states and restore
+        the ambient caller state after backward recomputation themselves.
+        Never keep this context open across a public iterator yield.
+        """
+        if self._depth:
+            yield
+            return
+        caller = _State.capture(self.device)
+        if self._model is None:
+            # Copying the state would correlate model dropout with the first
+            # caller draws. Derive a separate deterministic default stream.
+            self._model = caller.model_stream(self.device)
+        self._model.restore(self.device)
+        self._depth += 1
+        try:
+            yield
+        finally:
+            self._depth -= 1
+            self._model = _State.capture(self.device)
+            caller.restore(self.device)
+
+    def synchronize(self, group: dist.ProcessGroup | None) -> None:
+        """Continue the TP×CP leader's stream; never synchronize across DP.
+
+        None means no model-parallel group, not the default WORLD group. The
+        caller's live state is authoritative, including explicit manual seeding
+        or RNG restoration between forwards. Identically seeded DP workers are
+        allowed to stay identical; this method does not reseed them.
+        """
+        if group is None or dist.get_world_size(group) == 1:
+            return
+        state = _State.capture(self.device)
+        sizes = [state.cpu.numel()]
+        states = [state.cpu]
+        if state.cuda is not None:
+            sizes.append(state.cuda.numel())
+            states.append(state.cuda)
+        payload = torch.cat(states).to(
+            self.device if dist.get_backend(group) == "nccl" else "cpu"
+        )
+        dist.broadcast(payload, src=dist.get_global_rank(group, 0), group=group)
+        received = payload.cpu().split(sizes)
+        _State(received[0], received[1] if len(received) == 2 else None).restore(
+            self.device
+        )
+
+
+def caller_group() -> dist.ProcessGroup | None:
+    if not (dist.is_available() and dist.is_initialized()):
+        return None
+    try:
+        from megatron.core import parallel_state as ps
+
+        return ps.get_tensor_and_context_parallel_group(check_initialized=False)
+    except (AssertionError, ImportError, RuntimeError, ValueError):
+        return None

tests/integration/megatron/lora/test_dynamic_lora_slots.py

@@ -36,6 +36,7 @@
     _vocab_parallel_target_logprobs,
     _vocab_parallel_topk_from_local,
 )
+from art.trainer_rank._rng import TrainerRNG  # noqa: E402
 
 
 class _CudaValueHead(torch.nn.Module):
@@ -601,6 +602,7 @@ def _trainer_for(lora: LoRA, device: torch.device) -> TrainerRank:
         model_support_handler=_IdentityModelSupportHandler(),
     )
     trainer.device = device
+    trainer._rng = TrainerRNG(device)
     trainer._slot_stack = []
     trainer._default_slot_ref = None
     trainer._skipped_forward_waves = {}