Synchronize caller PyTorch randomness across TrainerRank replicas

[bradhilton]

TrainerRank now returns full outputs on every TP/CP replica, but a caller's random token mask or custom-head dropout could still differ between those replicas. Backward would then combine gradients from different losses. This change keeps existing forward/loss/backward code and caller-owned DP partitioning and gradient accumulation boundaries while coordinating default PyTorch randomness internally.

Before each public forward return or microbatch yield, ART broadcasts the DP-local TP×CP leader's CPU and trainer-device CUDA RNG states within that group. Custom object registration also synchronizes before invoking its factory. ART's planning/model work uses a persistent, deterministically separated RNG stream and restores the caller state on return or error; input iterators retain caller RNG ownership. The context never spans a public yield, and native PyTorch/Megatron activation checkpointing preserves the caller state during recomputation. Megatron's dedicated parallel RNG tracker is unchanged.

The caller's live RNG state remains authoritative, including explicit seeding/restoration. Distinct DP workers are never synchronized together; intentionally identical DP seeds remain identical. The model stream advances between forwards rather than replaying the same masks.

Scope and limitations:

• Covers default PyTorch CPU and trainer-device CUDA RNGs. Python/NumPy, explicit generators, other CUDA devices, concurrent global RNG consumers, and rank-dependent control flow are outside this contract.
• Caller operations must still match across TP/CP. This does not make skipping backward on physical ranks safe.
• The private model RNG stream persists only in the live TrainerRank instance. Checkpoint save/load remains a weights/optimizer operation and does not include this RNG state. There is no public per-forward RNG snapshot/replay API or exact stochastic resume across process restarts. Native activation checkpointing still captures/replays forward RNG for backward and must preserve the caller's ambient RNG state.

Validation:

• 279 passed across the new RNG tests and TrainerRank validation, splitting, and topology suites. An earlier adjacent custom-object/splitting/topology run passed 139 tests.
• Backend-only environment: 8 new tests passed / 4 optional-Megatron tests skipped; the broader forward/boundary selection passed 71 / skipped 2. Updated that suite's fake Megatron topology to explicitly report no TP×CP group.
• CPU CP2 and DP2×CP2, plus local 2×H200 NCCL TP2/CP2 oracles: matching random draws, masks and dropout; caller/model stream progression; both forward APIs; two accumulated forwards per Adam step; model/custom-head gradients and two optimizer updates matching an unsplit reference. Includes empty CP ownership and native Megatron checkpoint/tracker restoration. These are bounded numerical oracles, not a full pretrained-model benchmark.
• Negative controls: the merged parent fails caller-state preservation; disabling only synchronization fails replica random-draw agreement.
• Hosted 2×H200 GPU validation passed on c0f877cdd: 89 passed / 2 skipped in the main suite, plus passing SFT and shared-expert checks, and Qwen3-0.6B CP2 and TP2 public-API canaries with both sharing layouts and two LoRA slots. The existing constructor-bypassing LoRA fixture now initializes its RNG state; all four affected native cases also pass locally (one four-GPU case skipped).
• Warm 2×H200 measurement: approximately 0.080 ms per model RNG context plus TP×CP synchronization over 500 iterations, excluding model computation.
• Hosted quality checks passed, including types and backend/lightweight tests. uv run prek run --all-files locally passes lint, formatting, and lock checks; local type checking reports exactly the same 23 diagnostics as the untouched parent in that environment.

The new oracles are included in the existing GPU CI entrypoint. Local canaries used idle GPUs. Hosted validation used free Kubernetes GPUs; its cluster was terminated and its pod removal verified. Follow-up to #912; draft for design/code review, not ready to merge.

[bradhilton]

McCarthy: source/CPU-scope CLEAR at c0f877cddfc58a13eda12fd7d2b7771d559721c7. I found no blocking correctness issue in the complete seven-file PR diff. Reviewed against merge-base 7496cc09252c52ec7a63ab74d11caabce173f5b0; the current target is 7ef066b533a92c7a009644cddb2221ffdf3c00ae. The draft/design hold remains; this is not merge or adoption authorization.

The separation and synchronization boundaries are sound for the stated contract:

• The model stream is derived separately from the initial live caller state, then persists and advances. Nested contexts share that stream. Ordinary model errors and fatal Python control-flow exceptions restore the caller before propagating. The stream is not reset at every forward.
• Both public forward paths materialize caller-owned iterators outside the model context. The microbatch wrapper exits the context before yielding, restores it only around advancing internal work, and closes the internal iterator outside that context. Caller loss/backward therefore runs on the caller stream. The existing empty-wave collective guard and output/yield policy remain intact.
• Synchronization uses the actual TP×CP group's rank-zero global rank, not WORLD rank zero. None does not mean WORLD. CPU and the trainer-device CUDA state are broadcast together, with CUDA payload transport for NCCL; distinct DP groups are not synchronized together. Custom-object registration synchronizes before the factory, after the existing collective validation. Existing factory-error handling and parameter initialization/broadcast remain intact.
• Megatron's dedicated RNG tracker and activation-checkpoint implementations are unchanged. The model stream is deliberately not serialized with checkpoint weights/optimizer state; an external caller reseed does not rewind an already-live private model stream. The documentation discloses that reproducibility boundary, and excludes other generators/devices, NumPy/Python RNG, concurrent global consumers and rank-dependent caller control flow.

Independent validation: 46 exact-source lightweight checks passed for CPU/trainer-device state isolation, deterministic stream separation/progression, nested contexts, default-state restoration through ValueError/KeyboardInterrupt/SystemExit, exact original exception/cause/context preservation at both public forward boundaries, input-iterator ownership, yield/close behavior, nonzero global group leaders, distinct DP groups, and no WORLD fallback. Two unchanged-parent controls reproduce the caller-state leak through direct and microbatch forwarding. These execute the committed RNG module and extracted public methods with explicit stdlib generator-state/collective facades; they are not native PyTorch, Gloo, NCCL, gradient, or CUDA qualification. Peak probe RSS was approximately 40 MiB under a 128 MiB address-space bound. Pinned Ruff 0.15.11 check/format pass on all six changed Python files.

I separately inspected and hashed 12 retained author-evidence files. Hosted request/result/status records identify this exact head, and hosted logs show its checkout and successful execution: 89 passed/2 skipped in the main GPU selection, subsequent SFT/shared-expert checks, and the Qwen3-0.6B TP2/CP2 public-API canaries. The new numerical oracle uses a small linear model with injected forward execution but real gather/checkpoint/reduction/optimizer paths, and reconstructs the reference using the actual shard dropout masks. That is meaningful bounded numerical coverage, not a general pretrained-model dropout equivalence proof. The local 279-test and backend/boundary results are author executions, not reruns by me. Current-head quality checks and GPU validation are green.

API/behavior assessment: all 20 constructor/public method signatures are unchanged; the public wrapper changes only documentation, and production art.megatron, kernels, dependency pins and optimizer code are untouched. This is a medium runtime change with substantial tests, but its stochastic behavior is significant: caller draws stop advancing with model work, follower TP/CP generators adopt the leader's state, and a live private model stream is introduced. That directly addresses #912's replicated-loss requirement, but it must not be treated as automatically eligible under the no-significant-behavior merge authority. Brad retains the design/merge decision. Public signature compatibility does not make this behavior-neutral.

The target has advanced by #914's selective-recompute refusal. A read-only three-way file merge is clean and preserves that refusal together with this RNG change; I did not treat the older branch's absence of that later guard as a deletion by this PR or claim a native qualification of a new integration head.

Evidence: /home/brad/.local/share/mccarthy/art918-review-20260917/ (pr.diff, probe-result.json, audit-result.json, evidence-read.json, CI snapshots and cleanup.json). Own review processes/groups are retired and shared worktrees remain unchanged. No delegate, watcher, native distributed/GPU rerun, provider operation, merge or deployment. Live head/base must remain the values above; changed bytes require review again.

[bradhilton]

Minsky independent full-head review — ART #918 at exact head c0f877cddfc58a13eda12fd7d2b7771d559721c7 (two commits on merge-base 7496cc09 = #912; current main 44bc5d4c). Source-correctness: CLEAR. Standing conditions, stated separately: medium (+601/−7 across seven files, of which +454 are tests; runtime is a new 107-line _rng.py plus 23 lines in _impl.py); public API: no signature or exported-name change — the public forward_micro_batches docstring gains a documented RNG contract; art.megatron: untouched (Megatron's own parallel RNG tracker is not modified; the only non-trainer-rank file is the GPU CI script, which adds the new test); behavior: significant — after every microbatch yield and every dp_rank_forward return, non-leader TP/CP ranks' default CPU and trainer-device CUDA RNG states are overwritten with the DP-local leader's, and ART's own planning/model consumption no longer advances the caller's generators at all, so any caller whose downstream random draws depended on the prior interleaving observes a different sequence. Under Brad's conditions this is therefore not merge-eligible automatically; the owner marks it draft and not ready to merge, and adoption is Brad's call. Head and OPEN state (draft=true) rechecked immediately before posting; GitHub base 7ef066b5 has moved past the head's merge-base 7496cc09 (#912); GitHub reports MERGEABLE/CLEAN. Checks at that moment: Run on 2x H200=pass quality-checks=pass trainer-rank-gpu-validation=pass (owner-managed).

Design, traced. TrainerRNG.model() captures the caller's CPU (and CUDA, when the trainer device is CUDA) state, on first use derives a separate model stream by SHA-256 of that state under a fixed salt (so model draws are not correlated with the caller's first draws), installs the model stream, and in finally saves the advanced model state and restores the caller — so failures inside planning or execution cannot leak model consumption into the caller, and the model stream advances between forwards instead of replaying masks. Re-entry is a no-op via _depth. synchronize(group) is a no-op for None or a singleton group (the docstring is explicit that None means no model-parallel group, not WORLD, so DP workers are never coupled), otherwise concatenates the fixed-size CPU/CUDA state tensors, broadcasts from dist.get_global_rank(group, 0) — the correct global-rank translation from #920 — on the device NCCL requires or CPU for Gloo, splits by the recorded sizes and restores. caller_group() returns Megatron's TP×CP group with check_initialized=False and degrades to None on any import/assertion/runtime error, so Megatron-less and uninitialized paths keep today's behavior.

Placement in _impl.py. In forward_micro_batches the collective guard now runs before inputs are consumed, inputs are materialized under the caller's RNG (a data loader's draws stay the caller's), and only next(batches) — planning and model work — runs inside model(); the context never spans the public yield, and synchronize(caller_group()) runs after the context closes and after the empty-batch continue, which is uniform across the TP×CP replicas of one DP worker because they process identical items. dp_rank_forward materializes under the caller's RNG, plans and executes under model(), then synchronizes. Custom object registration synchronizes before invoking the factory, after the existing cross-rank agreement check, so the collective is matched. Because the public method is itself a generator, materialization still occurs on the first next, as on the base.

Composition with current main. _impl.py has since been changed by #914, #915, #922 and #900; GitHub's synthetic merge composes cleanly and its _impl.py delta against main is byte-identical (as +/− lines) to the PR's own hunks, so nothing on main touches these call sites.

Limits I note (non-blocking). The contract covers randomness drawn after the first yield or return; draws the caller makes before the first forward (for example precomputed token masks) are not synchronized and still rely on identical caller seeding, which the docstring's "identical caller seeds are not changed" sentence implies but does not spell out. The model stream is derived from the caller's unsynchronized first-forward state, so it may differ across TP/CP replicas; nothing in the PR claims model-stream agreement and Megatron's tracker governs model-parallel dropout, but it is worth stating. Python/NumPy RNGs, explicit generators, other CUDA devices and rank-dependent control flow are excluded, as documented. The private stream is not checkpointed, so exact stochastic resume across restarts is out of scope, as documented.

Witness (serial, native threads 1, CUDA hidden, fresh accounting; cgroup 254 GiB of 512). At this head with Megatron present — test_trainer_rank_rng.py, the validation, split, topology and head-recompute suites and the changed LoRA-slots integration file: 287 passed, 13 skipped, 5 failed; the five failures are the checkpoint-prefetch cases in test_trainer_rank_validation.py that fail identically on the base 7496cc09 and on current main in this environment and touch nothing this PR changes. The RNG file with Megatron blocked hosted-style: 8 passed, 4 skipped (two need megatron.core, two need two CUDA devices), matching the owner's backend-only figure. Negative control: on the base runtime with the new _rng.py copied in so imports resolve, 9 of 10 runnable RNG cases fail (only the pure-module no-WORLD test passes), including both CPU Gloo replicated-randomness cases (TP×CP = 2 with DP 1 and DP 2), so the implementation change is discriminated. Ruff check/format and ty clean on the six changed Python files. The 2×H200 NCCL oracles, canaries and the 0.080 ms overhead figure are the owner's; hosted quality-checks, 2×H200 and GPU validation are green at this head.

Void if the head changes.

[bradhilton]

Review history (consolidated 2026-09-17)

Two ACK comments and one routing notification for the review of c0f877c were removed; they carried no findings. McCarthy's and Minsky's reviews of that head remain in place (both source-scope CLEAR; the PR stays a draft pending Brad's design and merge decision on the new RNG behavior).