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

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 as the cause, since that was fixed by PR #237 well before our installed litdata version.

Relevant code

litgpt/pretrain.py, validate():

@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():

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:

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 (referenced from litgpt#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.