diff --git a/extensions/thunder/pretrain.py b/extensions/thunder/pretrain.py index 8e9f04fbf9..f5558fed4f 100644 --- a/extensions/thunder/pretrain.py +++ b/extensions/thunder/pretrain.py @@ -405,16 +405,32 @@ def validate(fabric: L.Fabric, model: nn.Module, val_dataloader: DataLoader, max fabric.print("Validating ...") model.eval() + # Keep ranks in lockstep (same rationale as litgpt.pretrain.validate / #2319). losses = [] - for k, batch in enumerate(val_dataloader): - if k >= max_iters: + val_iter = iter(val_dataloader) + for _ in range(max_iters): + try: + batch = next(val_iter) + has_batch = torch.tensor(1.0, device=fabric.device) + except StopIteration: + batch = None + has_batch = torch.tensor(0.0, device=fabric.device) + + fabric.all_reduce(has_batch, reduce_op="min") + if has_batch.item() == 0: break + input_ids = batch[:, 0 : model.max_seq_length].contiguous().long() targets = batch[:, 1 : (model.max_seq_length + 1)].contiguous().long() loss = forward_and_loss(model, input_ids, targets) losses.append(loss) - val_loss = torch.stack(losses).mean() + if losses: + val_loss = torch.stack(losses).mean() + else: + fabric.print("WARNING: no rank had validation data available for this call; reporting val_loss as NaN.") + val_loss = torch.tensor(float("nan"), device=fabric.device) + model.train() fabric.barrier() return val_loss diff --git a/litgpt/pretrain.py b/litgpt/pretrain.py index 35cf0bcb9a..d95cd4b4cb 100644 --- a/litgpt/pretrain.py +++ b/litgpt/pretrain.py @@ -431,17 +431,37 @@ def validate( fabric.print("Validating ...") model.eval() + # Keep ranks in lockstep: under FSDP, `model()` is a collective. If one rank + # runs out of validation batches early (common with sharded streaming data), + # the others must not keep calling `model()` alone or the job deadlocks. + # Agree via MIN-reduce whether every rank still has a batch before each step. + # See https://github.com/Lightning-AI/litgpt/issues/2319 losses = [] - for k, batch in enumerate(val_dataloader): - if k >= max_iters: + val_iter = iter(val_dataloader) + for _ in range(max_iters): + try: + batch = next(val_iter) + has_batch = torch.tensor(1.0, device=fabric.device) + except StopIteration: + batch = None + has_batch = torch.tensor(0.0, device=fabric.device) + + fabric.all_reduce(has_batch, reduce_op="min") + if has_batch.item() == 0: break + input_ids = batch[:, 0 : model.max_seq_length].contiguous().long() targets = batch[:, 1 : (model.max_seq_length + 1)].contiguous().long() logits = model(input_ids) loss = chunked_cross_entropy(logits, targets) losses.append(loss) - val_loss = torch.stack(losses).mean() + if losses: + val_loss = torch.stack(losses).mean() + else: + fabric.print("WARNING: no rank had validation data available for this call; reporting val_loss as NaN.") + val_loss = torch.tensor(float("nan"), device=fabric.device) + model.train() fabric.barrier() return val_loss diff --git a/tests/test_pretrain_validate_empty_rank.py b/tests/test_pretrain_validate_empty_rank.py new file mode 100644 index 0000000000..e74cd4ac12 --- /dev/null +++ b/tests/test_pretrain_validate_empty_rank.py @@ -0,0 +1,275 @@ +# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file. + +"""Regression tests for pretrain.validate with uneven / empty validation shards (#2319).""" + +from __future__ import annotations + +import threading +from typing import Any + +import torch +import torch.nn as nn +from lightning import Fabric +from torch.utils.data import DataLoader + +from litgpt.pretrain import validate + + +class _TinyLM(nn.Module): + """Minimal causal LM stand-in for validate().""" + + def __init__(self, vocab_size: int = 8, n_embd: int = 4, max_seq_length: int = 2): + super().__init__() + self.max_seq_length = max_seq_length + self.embed = nn.Embedding(vocab_size, n_embd) + self.head = nn.Linear(n_embd, vocab_size, bias=False) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.head(self.embed(input_ids)) + + +def _seq_batches(num_batches: int, seq_len: int = 3, vocab: int = 8) -> DataLoader: + """Build a DataLoader yielding ``[B, T]`` long tensors (not TensorDataset tuples).""" + + class _Rows(torch.utils.data.Dataset): + def __init__(self, data: torch.Tensor): + self.data = data + + def __len__(self) -> int: + return self.data.size(0) + + def __getitem__(self, index: int) -> torch.Tensor: + return self.data[index] + + if num_batches == 0: + data = torch.empty(0, seq_len, dtype=torch.long) + else: + data = torch.randint(0, vocab, (num_batches, seq_len), dtype=torch.long) + return DataLoader(_Rows(data), batch_size=1) + + +def test_validate_empty_dataloader_returns_nan(): + """Empty validation set must not crash with torch.stack([]).""" + fabric = Fabric(accelerator="cpu", devices=1) + fabric.launch() + model = fabric.setup(_TinyLM()) + loader = _seq_batches(0) + + val_loss = validate(fabric, model, loader, max_iters=2, verbose=False) + + assert torch.isnan(val_loss).item() + assert val_loss.device.type == fabric.device.type + + +def test_validate_with_batches_returns_finite_loss(): + fabric = Fabric(accelerator="cpu", devices=1) + fabric.launch() + model = fabric.setup(_TinyLM()) + loader = _seq_batches(3) + + val_loss = validate(fabric, model, loader, max_iters=2, verbose=False) + + assert torch.isfinite(val_loss).item() + assert val_loss.ndim == 0 + + +def test_validate_lockstep_logic_matches_issue_workaround(): + """Drive two ranks through the same MIN-reduce protocol as validate(). + + Verifies the synchronization contract that prevents deadlock when shards are + uneven, including the case a naive empty-skip mishandles. + """ + # Rank 0: 2 batches; rank 1: 0 → zero synchronized forwards. + remaining = [2, 0] + collective_forwards = 0 + forward_counts = [0, 0] + for _ in range(4): + votes = [1.0 if remaining[r] > 0 else 0.0 for r in range(2)] + if min(votes) == 0: + break + for rank in range(2): + remaining[rank] -= 1 + forward_counts[rank] += 1 + collective_forwards += 1 + + assert collective_forwards == 0 + assert forward_counts == [0, 0] + + # Partial overlap: rank0=3, rank1=1 → exactly one synchronized step. + remaining = [3, 1] + collective_forwards = 0 + forward_counts = [0, 0] + for _ in range(4): + votes = [1.0 if remaining[r] > 0 else 0.0 for r in range(2)] + if min(votes) == 0: + break + for rank in range(2): + remaining[rank] -= 1 + forward_counts[rank] += 1 + collective_forwards += 1 + + assert collective_forwards == 1 + assert forward_counts == [1, 1] + assert remaining == [2, 0] + + +class _BarrierReduceFabric: + """Minimal Fabric stand-in: all_reduce(min) + barrier across threads. + + A forward that also all_reduces deadlocks if ranks skip model() unequally — + the same failure mode as FSDP when validation shards are uneven (#2319). + """ + + def __init__(self, world_size: int, timeout: float = 5.0): + self.world_size = world_size + self.device = torch.device("cpu") + self._timeout = timeout + self._barrier = threading.Barrier(world_size, timeout=timeout) + self._lock = threading.Lock() + self._votes: list[torch.Tensor] = [] + self._reduced: torch.Tensor | None = None + self._reduce_barrier = threading.Barrier(world_size, timeout=timeout) + self.local = threading.local() + + @property + def global_rank(self) -> int: + return self.local.rank + + def barrier(self) -> None: + self._barrier.wait() + + def print(self, *args: Any, **kwargs: Any) -> None: + pass + + def all_reduce(self, tensor: torch.Tensor, reduce_op: str = "mean") -> torch.Tensor: + """Match Fabric: reduce in-place and return the same tensor.""" + assert reduce_op == "min" + with self._lock: + self._votes.append(tensor.detach().clone()) + if len(self._votes) == self.world_size: + self._reduced = torch.stack(self._votes).min() + self._votes.clear() + self._reduce_barrier.wait() + assert self._reduced is not None + # Copy under lock so a slow reader cannot see the next step's value. + with self._lock: + tensor.copy_(self._reduced) + # Second wait so every rank finishes reading before the next reduce. + self._reduce_barrier.wait() + return tensor + + +class _CollectiveTinyLM(_TinyLM): + def __init__(self, fabric: _BarrierReduceFabric): + super().__init__() + self.fabric = fabric + self.forward_calls = 0 + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + self.forward_calls += 1 + # Collective every forward — skipping this on some ranks hangs the barrier. + token = torch.ones((), device=self.fabric.device) + self.fabric.all_reduce(token, reduce_op="min") + return super().forward(input_ids) + + +def _run_validate_on_rank( + fabric: _BarrierReduceFabric, + rank: int, + num_batches: int, + max_iters: int, + results: dict[int, Any], + errors: dict[int, BaseException], +) -> None: + fabric.local.rank = rank + model = _CollectiveTinyLM(fabric) + loader = _seq_batches(num_batches) + try: + loss = validate(fabric, model, loader, max_iters=max_iters, verbose=False) + results[rank] = {"loss": loss, "calls": model.forward_calls} + except BaseException as exc: # noqa: BLE001 - surface any hang/crash to the parent + errors[rank] = exc + + +def test_validate_uneven_shards_does_not_deadlock_threaded(): + """Rank 0 has data, rank 1 empty: both finish; equal collective call counts.""" + fabric = _BarrierReduceFabric(world_size=2) + results: dict[int, Any] = {} + errors: dict[int, BaseException] = {} + threads = [ + threading.Thread( + target=_run_validate_on_rank, + args=(fabric, rank, batches, 4, results, errors), + daemon=True, + ) + for rank, batches in ((0, 2), (1, 0)) + ] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + assert not t.is_alive(), "validate() hung — ranks likely skipped a collective unequally" + + assert errors == {} + assert results[0]["calls"] == results[1]["calls"] == 0 + assert torch.isnan(results[0]["loss"]).item() + assert torch.isnan(results[1]["loss"]).item() + + +def test_validate_partial_overlap_stops_together_threaded(): + """Rank 0 has 3 batches, rank 1 has 1: both take exactly one model step.""" + fabric = _BarrierReduceFabric(world_size=2) + results: dict[int, Any] = {} + errors: dict[int, BaseException] = {} + threads = [ + threading.Thread( + target=_run_validate_on_rank, + args=(fabric, rank, batches, 5, results, errors), + daemon=True, + ) + for rank, batches in ((0, 3), (1, 1)) + ] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + assert not t.is_alive(), "validate() hung on partial-overlap shards" + + assert errors == {} + assert results[0]["calls"] == results[1]["calls"] == 1 + assert torch.isfinite(results[0]["loss"]).item() + assert torch.isfinite(results[1]["loss"]).item() + + +def test_naive_empty_skip_deadlocks_without_min_reduce(): + """Control: skipping model() on the empty rank hangs a collective forward. + + Documents why a bare ``if not losses: return nan`` is insufficient under FSDP. + """ + fabric = _BarrierReduceFabric(world_size=2, timeout=1.0) + hung = threading.Event() + + def rank0() -> None: + fabric.local.rank = 0 + model = _CollectiveTinyLM(fabric) + # Pretend we have a batch and enter forward (collective). + try: + model(torch.zeros(1, 2, dtype=torch.long)) + except threading.BrokenBarrierError: + hung.set() + + def rank1_skips() -> None: + fabric.local.rank = 1 + # Naive empty-rank path: never call model(), only barrier at the end. + try: + fabric.barrier() + except threading.BrokenBarrierError: + hung.set() + + t0 = threading.Thread(target=rank0, daemon=True) + t1 = threading.Thread(target=rank1_skips, daemon=True) + t0.start() + t1.start() + t0.join(timeout=3) + t1.join(timeout=3) + assert hung.is_set() or t0.is_alive() or t1.is_alive()