diff --git a/scripts/ci/trainer-rank-gpu-tests.sh b/scripts/ci/trainer-rank-gpu-tests.sh index 690d99db3..a96d18394 100755 --- a/scripts/ci/trainer-rank-gpu-tests.sh +++ b/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]' \ diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py index fd321a820..c99f70641 100644 --- a/src/art/trainer_rank/__init__.py +++ b/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 diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 0e0045d61..48cc35ede 100644 --- a/src/art/trainer_rank/_impl.py +++ b/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 diff --git a/src/art/trainer_rank/_rng.py b/src/art/trainer_rank/_rng.py new file mode 100644 index 000000000..cd1839c89 --- /dev/null +++ b/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 diff --git a/tests/integration/megatron/lora/test_dynamic_lora_slots.py b/tests/integration/megatron/lora/test_dynamic_lora_slots.py index cd9d47f60..8d1a1868d 100644 --- a/tests/integration/megatron/lora/test_dynamic_lora_slots.py +++ b/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 = {} diff --git a/tests/unit/test_trainer_rank_rng.py b/tests/unit/test_trainer_rank_rng.py new file mode 100644 index 000000000..383261a68 --- /dev/null +++ b/tests/unit/test_trainer_rank_rng.py @@ -0,0 +1,454 @@ +from __future__ import annotations + +from contextlib import nullcontext +from datetime import timedelta +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.utils.checkpoint import checkpoint + +from art.trainer_rank import AdamParams, ForwardInput, ForwardOutput, TrainerRank, _impl +from art.trainer_rank._impl import _CheckpointSlot, _GatherContextParallelRows +from art.trainer_rank._rng import TrainerRNG, _State, caller_group + + +def _trainer(device="cpu"): + model = torch.nn.Linear(3, 4, bias=False, device=device) + runtime = SimpleNamespace( + model=[model], + optimizer=None, + provider=SimpleNamespace(hidden_size=4, num_layers=1), + model_support_handler=SimpleNamespace( + build_gdn_execution_spec=False, + zero_internal_padding_grads=lambda _: None, + ), + ) + return TrainerRank(cast(Any, runtime)) + + +def _stub_forward(monkeypatch, trainer, execute): + monkeypatch.setattr( + trainer, "_plan_admissible_forward", lambda *a, **k: (None, None) + ) + monkeypatch.setattr(trainer, "_execute_admitted_plan", lambda *a, **k: execute()) + + +def _assert_state_equal(left, right): + assert torch.equal(left.cpu, right.cpu) + if left.cuda is not None: + assert torch.equal(left.cuda, right.cuda) + + +def test_model_stream_advances_without_advancing_caller(monkeypatch): + trainer = _trainer() + torch.manual_seed(811) + caller = torch.get_rng_state() + expected = torch.rand(3, 12) + torch.set_rng_state(caller) + observed = [] + internal_states = [] + + def execute(): + # Nested internal work shares the model stream instead of restarting it. + with trainer._rng.model(): + observed.append(torch.rand(12)) + internal_states.append(torch.get_rng_state()) + return [] + + _stub_forward(monkeypatch, trainer, execute) + trainer.dp_rank_forward([], no_grad=True) + assert torch.equal(torch.get_rng_state(), caller) + # Caller randomness advances independently between model forwards. + assert torch.equal(torch.rand(12), expected[0]) + trainer.dp_rank_forward([]) + assert not torch.equal(observed[0], expected[0]) + generator = torch.Generator().set_state(internal_states[0]) + torch.testing.assert_close(observed[1], torch.rand(12, generator=generator)) + assert torch.equal(torch.rand(12), expected[1]) + + +def test_forward_failure_restores_caller_and_advances_model(monkeypatch): + trainer = _trainer() + torch.manual_seed(981) + caller = torch.get_rng_state() + internal_states = [] + + def fail(): + torch.rand(7) + internal_states.append(torch.get_rng_state()) + raise ValueError("model failed") + + _stub_forward(monkeypatch, trainer, fail) + with pytest.raises(ValueError, match="model failed"): + trainer.dp_rank_forward([]) + assert torch.equal(torch.get_rng_state(), caller) + generator = torch.Generator().set_state(internal_states[0]) + with trainer._rng.model(): + assert torch.equal(torch.rand(7), torch.rand(7, generator=generator)) + + +@pytest.mark.parametrize("yield_empty", (False, True)) +def test_microbatch_yields_and_close_do_not_hold_rng_context(monkeypatch, yield_empty): + trainer = _trainer() + torch.manual_seed(177) + caller = torch.get_rng_state() + draws = torch.rand(5, 8) + torch.set_rng_state(caller) + observed = [] + internal_states = [] + + def batches(*args, **kwargs): + for index in range(3): + observed.append(torch.rand(8)) + internal_states.append(torch.get_rng_state()) + yield SimpleNamespace( + outputs=[] if index == 0 else [index], + stats=SimpleNamespace(global_count=1), + ) + + monkeypatch.setattr(trainer, "_forward_micro_batches", batches) + iterator = trainer.forward_micro_batches([], yield_empty=yield_empty) + next(iterator) + assert trainer._rng._depth == 0 + assert torch.equal(torch.get_rng_state(), caller) + assert torch.equal(torch.rand(8), draws[0]) + next(iterator) + assert trainer._rng._depth == 0 + iterator.close() + assert torch.equal(torch.rand(8), draws[1]) + generator = torch.Generator().set_state(internal_states[0]) + for draw in observed[1:]: + torch.testing.assert_close(draw, torch.rand(8, generator=generator)) + + +def test_uninitialized_model_parallel_group_does_not_mean_world(monkeypatch): + monkeypatch.setattr(dist, "is_initialized", lambda: False) + assert caller_group() is None + monkeypatch.setattr( + dist, "broadcast", lambda *a, **k: pytest.fail("unexpected WORLD broadcast") + ) + TrainerRNG(torch.device("cpu")).synchronize(None) + + +def test_caller_reseed_and_restore_between_forwards(monkeypatch): + trainer = _trainer() + _stub_forward(monkeypatch, trainer, lambda: (torch.rand(9), [])[1]) + trainer.dp_rank_forward([]) + torch.manual_seed(191) + saved = torch.get_rng_state() + expected = torch.rand(9) + torch.set_rng_state(saved) + trainer.dp_rank_forward([]) + assert torch.equal(torch.rand(9), expected) + torch.set_rng_state(saved) + trainer.dp_rank_forward([]) + assert torch.equal(torch.rand(9), expected) + + +@pytest.mark.parametrize("microbatches", (False, True)) +def test_input_iterators_use_caller_rng(monkeypatch, microbatches): + trainer = _trainer() + torch.manual_seed(617) + state = torch.get_rng_state() + expected = torch.rand(2, 5) + torch.set_rng_state(state) + request = ForwardInput(input_tokens=torch.arange(8), hidden_states=True) + + def inputs(): + assert torch.equal(torch.rand(5), expected[0]) + yield request + + def execute(): + torch.rand(11) + return [ForwardOutput(None, None, None, torch.ones(8, 4))] + + _stub_forward(monkeypatch, trainer, execute) + if microbatches: + + def batches(items, **kwargs): + assert items == [request] + yield SimpleNamespace( + outputs=execute(), stats=SimpleNamespace(global_count=1) + ) + + monkeypatch.setattr(trainer, "_forward_micro_batches", batches) + list(trainer.forward_micro_batches(inputs())) + else: + trainer.dp_rank_forward(inputs()) + assert torch.equal(torch.rand(5), expected[1]) + + +@pytest.mark.parametrize("dp_size", (1, 2)) +def test_replicated_caller_randomness_cpu(dp_size, tmp_path): + pytest.importorskip("megatron.core") + mp.spawn( + _distributed_worker, + args=(dp_size, "cp", "gloo", f"file://{tmp_path / 'rng'}"), + nprocs=2 * dp_size, + join=True, + ) + + +@pytest.mark.parametrize("parallelism", ("tp", "cp")) +def test_replicated_caller_randomness_cuda(parallelism, tmp_path): + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + pytest.skip("requires two CUDA devices") + pytest.importorskip("megatron.core") + mp.spawn( + _distributed_worker, + args=(1, parallelism, "nccl", f"file://{tmp_path / 'rng'}"), + nprocs=2, + join=True, + ) + + +def _distributed_worker(rank, dp_size, parallelism, backend, init_method): + from megatron.core import parallel_state as ps + + device = torch.device("cpu" if backend == "gloo" else f"cuda:{rank}") + if device.type == "cuda": + torch.cuda.set_device(device) + dist.init_process_group( + backend, + init_method=init_method, + rank=rank, + world_size=2 * dp_size, + timeout=timedelta(seconds=90), + ) + try: + replica_groups = [dist.new_group([2 * dp, 2 * dp + 1]) for dp in range(dp_size)] + dp_groups = [ + dist.new_group(list(range(replica, 2 * dp_size, 2))) for replica in range(2) + ] + dp_rank, replica_rank = divmod(rank, 2) + replica_group, dp_group = replica_groups[dp_rank], dp_groups[replica_rank] + with pytest.MonkeyPatch.context() as patch: + patch.setattr( + ps, "get_tensor_and_context_parallel_group", lambda **_: replica_group + ) + patch.setattr( + ps, + "get_tensor_model_parallel_world_size", + lambda: 2 if parallelism == "tp" else 1, + ) + patch.setattr( + ps, + "get_context_parallel_world_size", + lambda: 2 if parallelism == "cp" else 1, + ) + patch.setattr( + ps, + "get_tensor_model_parallel_group", + lambda **_: replica_group if parallelism == "tp" else None, + ) + patch.setattr( + ps, + "get_data_parallel_group", + lambda **_: dp_group if parallelism == "tp" else dist.group.WORLD, + ) + patch.setattr(ps, "get_data_parallel_rank", lambda: dp_rank) + patch.setattr(ps, "get_data_parallel_world_size", lambda: dp_size) + _gradient_oracle( + patch, + device, + rank, + dp_rank, + dp_size, + replica_rank, + replica_group, + dp_group, + parallelism, + ) + finally: + dist.destroy_process_group() + + +def _gradient_oracle( + patch, + device, + rank, + dp_rank, + dp_size, + replica_rank, + replica_group, + dp_group, + parallelism, +): + trainer = _trainer(device) + decoder = trainer.runtime.model[0].weight + with torch.no_grad(): + decoder.copy_(torch.arange(12, device=device).reshape(4, 3) / 30) + if parallelism == "tp": + decoder.grad_sync_op = "sum" + trainer._checkpoint_slots["student"] = _CheckpointSlot( + config={ + "base_model_name_or_path": "test", + "r": 1, + "lora_alpha": 1, + "target_modules": [], + }, + params=(decoder,), + ) + # Registration repairs initially different CPU/CUDA states within each DP + # worker, while the registered weights remain common to all DP workers. + torch.manual_seed(711 + 100 * dp_rank + replica_rank) + head = trainer.module( + "head", + lambda: torch.nn.Sequential(torch.nn.Dropout(0.3), torch.nn.Linear(4, 2)), + checkpoint="student", + ) + params = trainer._checkpoint_slots["student"].params + reference = tuple(torch.nn.Parameter(param.detach().clone()) for param in params) + optimizer = torch.optim.AdamW(reference, lr=0.01, weight_decay=0.0) + features = ( + torch.arange(24, device=device, dtype=torch.float32).reshape(8, 3) / 20 + + dp_rank / 10 + ) + rows = torch.arange(replica_rank * 4, (replica_rank + 1) * 4, device=device) + request = ForwardInput(input_tokens=torch.arange(8), hidden_states=True) + masks = [] + recompute_masks = [] + cpu_draws = [] + previous_caller_mask = None + tracker = None + if device.type == "cuda" and parallelism == "cp": + from megatron.core.tensor_parallel.random import get_cuda_rng_tracker + + tracker = get_cuda_rng_tracker() + tracker.add("art-test-model", 3199 + rank) + + def execute(): + # Internal consumption deliberately differs across physical ranks. It + # must not leak into caller masks or custom-head dropout. + torch.rand(rank + 1) + torch.rand(rank + 3, device=device) + records = [] + local_rows = rows + if parallelism == "cp" and not masks: + # One CP peer owns no tokens but still consumes the caller's full + # output and participates in backward and RNG synchronization. + local_rows = torch.arange(8 if replica_rank == 0 else 0, device=device) + + def model(weight): + with ( + tracker.fork("art-test-model") if tracker is not None else nullcontext() + ): + mask = torch.nn.functional.dropout( + torch.ones_like(features[local_rows]), 0.2 + ) + records.append(mask.detach().clone()) + return (features[local_rows] * mask) @ weight.T + + if tracker is not None: + from megatron.core.tensor_parallel import checkpoint as megatron_checkpoint + + hidden = megatron_checkpoint(model, False, decoder) + else: + hidden = checkpoint(model, decoder, use_reentrant=False) + recompute_masks.append(records) + full_mask = torch.zeros_like(features) + full_mask[local_rows] = records[0] + dist.all_reduce(full_mask, group=replica_group) + masks.append(full_mask) + if parallelism == "tp": + hidden = trainer._gather_sequence_parallel_hidden(hidden[:, None]) + else: + hidden = _GatherContextParallelRows.apply( + hidden, local_rows, len(features), replica_group + ) + return [ForwardOutput(None, None, None, hidden)] + + _stub_forward(patch, trainer, execute) + + def batches(*args, **kwargs): + yield SimpleNamespace( + outputs=execute(), stats=SimpleNamespace(global_count=dp_size) + ) + + patch.setattr(trainer, "_forward_micro_batches", batches) + for step in range(2): + losses = [] + expected_losses = [] + for micro in range(2): + if step == micro == 0: + # Disagree again after registration: forward return must repair + # the caller even though the model uses a separate stream. + torch.manual_seed(1231 + 100 * dp_rank + replica_rank) + before = _State.capture(device) + if step == 0: + output = trainer.dp_rank_forward([request])[0] + else: + iterator = trainer.forward_micro_batches([request]) + output = next(iterator).outputs[0] + assert trainer._rng._depth == 0 + iterator.close() + if replica_rank == 0: + _assert_state_equal(_State.capture(device), before) + caller = _State.capture(device) + cpu_draw = torch.rand(16) + cpu_draws.append(cpu_draw) + mask = torch.rand(len(features), device=device) > 0.35 + if previous_caller_mask is not None: + assert not torch.equal(cpu_draw, previous_caller_mask) + previous_caller_mask = cpu_draw + loss = head(output.hidden_states[mask]).square().sum() + losses.append(loss) + # The unsplit reference replays the exact caller stream, including + # the CPU mask draw and dropout; model dropout is taken from the + # actual shards so this checks caller/model gradient consistency. + with torch.random.fork_rng( + devices=[device.index] if device.type == "cuda" else [] + ): + caller.restore(device) + torch.testing.assert_close(torch.rand(16), cpu_draw) + expected_mask = torch.rand(len(features), device=device) > 0.35 + assert torch.equal(expected_mask, mask) + hidden = (features * masks[-1]) @ reference[0].T + dropped = torch.nn.functional.dropout(hidden[expected_mask], 0.3) + expected_loss = ( + torch.nn.functional.linear(dropped, reference[1], reference[2]) + .square() + .sum() + ) + torch.testing.assert_close(loss, expected_loss) + expected_losses.append(expected_loss) + # The collective comparison is independent of the reference replay. + copies = [torch.empty_like(cpu_draw, device=device) for _ in range(2)] + dist.all_gather(copies, cpu_draw.to(device), group=replica_group) + assert torch.equal(copies[0], copies[1]) + before_backward = _State.capture(device) + tracker_before = tracker.get_states() if tracker is not None else {} + torch.stack(losses).sum().backward() + _assert_state_equal(_State.capture(device), before_backward) + if tracker is not None: + for key, state in tracker_before.items(): + assert torch.equal(tracker.get_states()[key], state) + torch.stack(expected_losses).sum().backward() + reduced = trainer._reduce_dynamic_grads(params, scale_grads=1 / dp_size) + for actual, expected in zip(reduced, reference, strict=True): + assert expected.grad is not None + dist.all_reduce(expected.grad, group=dp_group) + expected.grad.div_(dp_size) + torch.testing.assert_close(actual, expected.grad, rtol=2e-5, atol=2e-5) + optimizer.step() + optimizer.zero_grad() + metrics = trainer.optim_step( + params=AdamParams(learning_rate=0.01, weight_decay=0.0, grad_clip_norm=0), + scale_grads=1 / dp_size, + checkpoints=["student"], + ) + assert metrics["update_successful"] == 1 + for actual, expected in zip(params, reference, strict=True): + torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-5) + for records in recompute_masks: + assert len(records) == 2 + assert torch.equal(records[0], records[1]) + assert not torch.equal(masks[0], masks[1]) + if dp_size > 1: + copies = [torch.empty_like(cpu_draws[0]) for _ in range(dp_size)] + dist.all_gather(copies, cpu_draws[0], group=dp_group) + assert not torch.equal(copies[0], copies[1]) diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index d6ed71048..88db18fcd 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -3434,7 +3434,12 @@ def _forward_yield_modes_worker(rank: int, world_size: int, init_method: str) -> from megatron.core import parallel_state except ImportError: core = ModuleType("megatron.core") - parallel_state = cast(Any, SimpleNamespace()) + parallel_state = cast( + Any, + SimpleNamespace( + get_tensor_and_context_parallel_group=lambda **_: None + ), + ) cast(Any, core).parallel_state = parallel_state monkeypatch.setitem(sys.modules, "megatron.core", core) monkeypatch.setattr(