litgpt pretrain crashes or deadlocks when a rank's val_dataloader yields zero batches (multi-GPU FSDP)

Nobody has claimed this yet.

## Summary

When pretraining with `litgpt pretrain` across multiple GPUs (FSDP strategy) using the `TextFiles` data module, if the validation dataset is small enough that `litdata`'s `StreamingDataset` assigns some ranks zero chunks/batches for a given `validate()` call, one of two failures occurs:

1. **Crash**: `RuntimeError: stack expects a non-empty TensorList` at `pretrain.py`'s `validate()`, where `torch.stack(losses).mean()` is called on an empty `losses` list for the affected rank(s).
2. **Deadlock**, if that crash is naively patched around (e.g. `if losses: ... else: return NaN`): ranks with zero batches skip the `validate()` loop entirely and reach `fabric.barrier()` early, while ranks *with* batches keep calling `model(input_ids)` inside the loop. Since `model()` under FSDP triggers a distributed all-gather collective, the "full" ranks hang indefinitely waiting for peers that already moved past that collective — because the "empty" ranks never call it at all. This deadlocks the whole job (confirmed via `nvidia-smi` showing a subset of GPUs frozen at 100% util with static memory usage, and the rest idle, indefinitely).

## Environment

- `litgpt` 0.5.13
- `litdata` 0.2.72
- 8x GPU single-node FSDP run (`--devices 8`), Slurm-managed cluster
- Data module: `TextFiles`, with an explicit `--data.val_data_path` pointing to a directory of 16 `.txt` shard files

## Reproduction

1. Prepare a `TextFiles` validation set with N raw `.txt` files, where N is close to (or evenly divides into) the GPU/world-size count — in our case 16 source files, preprocessed via `litgpt`'s built-in `TextFiles.prepare_data()` → `litdata.optimize()` with `chunk_bytes="50MB"`. This produced **16 evenly-sized chunk files** (~92–98MB each), confirmed via `index.json` and `ls -la` on the val output dir.
2. Run `litgpt pretrain` with `--devices 8` (single node, 8-way FSDP).
3. During the initial "sanity check" validation call in `fit()` (`validate(fabric, model, val_dataloader, max_iters=2, verbose=False)`, called unconditionally regardless of `--eval.initial_validation`), some ranks (in our case, ranks 4–7 of 8) receive **zero** batches from `val_dataloader()`'s `StreamingDataset`, while others (ranks 0–3) receive at least one.
4. Ranks with zero batches crash on `torch.stack([]).mean()`.

We could not fully root-cause *why* the 16 chunks were split 0/1+ across specific ranks rather than an even ~2-per-rank split — the chunk files themselves were confirmed healthy and evenly sized, so the imbalance appears to originate in how `StreamingDataset`'s automatic rank/world-size sharding (in `litdata`, likely `subsample_streaming_dataset` / the shuffle+rank-assignment logic in `litdata/streaming/dataset.py`, which we did not have time to fully trace) divides a moderate chunk count across a moderate rank count. We checked and ruled out [litdata#233](https://github.com/Lightning-AI/litData/issues/233) as the cause, since that was fixed by [PR #237](https://github.com/Lightning-AI/litData/pull/237) well before our installed `litdata` version.
## Relevant code

`litgpt/pretrain.py`, `validate()`:
```python
@torch.no_grad()
def validate(
fabric: L.Fabric, model: nn.Module, val_dataloader: DataLoader, max_iters: int, verbose: bool = True
) -> torch.Tensor:
fabric.barrier()
if verbose:
fabric.print("Validating ...")
model.eval()

losses = []
for k, batch in enumerate(val_dataloader):
if k >= max_iters:
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() # <-- crashes if losses is empty
model.train()
fabric.barrier()
return val_loss
```

`litgpt/data/text_files.py`, `val_dataloader()`:
```python
def val_dataloader(self) -> DataLoader:
from litdata.streaming import StreamingDataLoader, StreamingDataset, TokensLoader

val_dataset = StreamingDataset(
input_dir=str(self.out_path_val),
item_loader=TokensLoader(block_size=self.max_seq_length),
shuffle=True,
)
val_dataloader = StreamingDataLoader(
val_dataset, batch_size=self.batch_size, pin_memory=True, num_workers=self.num_workers, drop_last=True
)
return val_dataloader
```

No explicit rank/world_size handling is passed here — presumably `StreamingDataset` auto-detects it from the environment, but we didn't trace far enough into `litdata` internals to confirm exactly how chunks get assigned per rank in this configuration.

## Why this is worth fixing at the `litgpt` level (not just `litdata`)

Even if the underlying `litdata` chunk-assignment behavior is "working as intended" for some definition of correct sharding, `litgpt`'s `validate()` has no guard against the case where a rank legitimately has zero local batches — and naively adding one (skip the loop early) creates a worse failure mode (silent deadlock) than the original crash, because `model()` is a collective operation under FSDP. Any fix needs to either:

- Guarantee every rank has ≥1 batch (upstream, in `litdata`'s sharding — outside our ability to fix), or
- Synchronize the loop across ranks so every rank calls `model()` the same number of times regardless of local data availability (fixable in `litgpt`).

## Our workaround (local patch, not a proposed final implementation)

We patched our local install so every rank agrees (via `fabric.all_reduce` with a MIN reduction) on whether *every* rank still has data before each `model()` call, and all ranks stop together the moment any single rank runs dry:

```python
losses = []
val_iter = iter(val_dataloader)
for k 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)

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)
```

This resolved both the crash and the deadlock for us, and training now proceeds normally on 8 GPUs. We're sharing it as a starting point, not a polished PR — happy to provide more detail on our exact setup/repro if useful.

## Also worth noting

We noticed `litgpt/pretrain.py`'s `fit()` runs this validation "sanity check" (`max_iters=2`) **unconditionally**, with no way to skip it via CLI even when `--eval.initial_validation=False`. We saw mention of a `--skip_validation` flag added in [litgpt#1228](https://github.com/Lightning-AI/litgpt/pull/1228) (referenced from [litgpt#1202](https://github.com/Lightning-AI/litgpt/issues/1202)) for a related finetuning validation-hang issue — if that flag (or an equivalent for `pretrain`) also covered this unconditional sanity-check call, it might be a simpler way to sidestep this specific failure mode too, though it wouldn't fix the underlying sharding imbalance for people who do want validation metrics.

Thanks for litgpt — happy to help reproduce or provide more logs if useful.

Start in litgpt/pretrain.py at validate() and trace its call from fit(), then inspect litgpt/data/text_files.py for val_dataloader(). Reproduce the zero-batch case with multi-GPU FSDP and verify that validation exits consistently across ranks without empty-loss crashes or deadlocks; the existing all_reduce workaround documents one possible direction but is not a final implementation.

Written by the indexing model from the issue text.