Validation loop hangs

[DNXie]

Bug description

Hi folks,

We found an hanging issue in the validation loop that might be related to the dataloder/sharding in FSDP.

Reproduction:

• Code pointer

• Command:

CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" NGPU=2 ./run_train.sh

• log

...
[rank1]:[titan] 2025-08-21 11:19:52,968 - root - INFO - [Validation Debug]  Rank 1 batch 87 Loss 12.241142272949219
[rank1]:[titan] 2025-08-21 11:19:52,972 - root - INFO - [Validation Debug]  Rank 1 batch 88 start
[rank1]:[titan] 2025-08-21 11:19:52,988 - root - INFO - [Validation Debug]  Rank 1 batch 88 Loss 12.209203720092773
[rank1]:[titan] 2025-08-21 11:19:52,993 - root - INFO - [Validation Debug]  Rank 1 batch 89 start
[rank0]:[titan] 2025-08-21 11:19:52,905 - root - INFO - [Validation Debug]  Rank 0 batch 84 Loss 12.281610488891602
[rank0]:[titan] 2025-08-21 11:19:52,910 - root - INFO - [Validation Debug]  Rank 0 batch 85 start
[rank0]:[titan] 2025-08-21 11:19:52,927 - root - INFO - [Validation Debug]  Rank 0 batch 85 Loss 12.214127540588379
[rank0]:[titan] 2025-08-21 11:19:52,931 - root - INFO - [Validation Debug]  Rank 0 batch 86 start
[rank0]:[titan] 2025-08-21 11:19:52,947 - root - INFO - [Validation Debug]  Rank 0 batch 86 Loss 12.226905822753906
[rank0]:[titan] 2025-08-21 11:19:52,951 - root - INFO - [Validation Debug]  Rank 0 batch 87 start
[rank0]:[titan] 2025-08-21 11:19:52,968 - root - INFO - [Validation Debug]  Rank 0 batch 87 Loss 12.206489562988281
[rank0]:[titan] 2025-08-21 11:19:52,972 - root - INFO - [Validation Debug]  Rank 0 batch 88 start
[rank0]:[titan] 2025-08-21 11:19:52,988 - root - INFO - [Validation Debug]  Rank 0 batch 88 Loss 12.239264488220215
[rank0]:[titan] 2025-08-21 11:19:52,992 - root - INFO - [Validation Debug]  Rank 0 batch 89 start
[rank1]:[titan] 2025-08-21 11:19:53,009 - root - INFO - [Validation Debug]  Rank 1 batch 89 Loss 12.284561157226562
[rank1]:[titan] 2025-08-21 11:19:53,013 - root - INFO - [Validation Debug]  Rank 1 batch 90 start
[rank1]:[titan] 2025-08-21 11:19:53,030 - root - INFO - [Validation Debug]  Rank 1 batch 90 Loss 12.25745964050293
[rank1]:[titan] 2025-08-21 11:19:53,034 - root - INFO - [Validation Debug]  Rank 1 batch 91 start
[rank1]:[titan] 2025-08-21 11:19:53,050 - root - INFO - [Validation Debug]  Rank 1 batch 91 Loss 12.285140991210938
[rank1]:[titan] 2025-08-21 11:19:53,054 - root - WARNING - Dataset alpaca_validation has run out of data
[rank1]:[titan] 2025-08-21 11:19:53,054 - root - INFO - [Validation Debug] Rank 1 Validation done at batch 91.
[rank0]:[titan] 2025-08-21 11:19:53,009 - root - INFO - [Validation Debug]  Rank 0 batch 89 Loss 12.28492546081543
[rank0]:[titan] 2025-08-21 11:19:53,014 - root - INFO - [Validation Debug]  Rank 0 batch 90 start
[rank0]:[titan] 2025-08-21 11:19:53,030 - root - INFO - [Validation Debug]  Rank 0 batch 90 Loss 12.212723731994629
[rank0]:[titan] 2025-08-21 11:19:53,034 - root - INFO - [Validation Debug]  Rank 0 batch 91 start
[rank0]:[titan] 2025-08-21 11:19:53,050 - root - INFO - [Validation Debug]  Rank 0 batch 91 Loss 12.228861808776855
[rank0]:[titan] 2025-08-21 11:19:53,055 - root - INFO - [Validation Debug]  Rank 0 batch 92 start

Rank 1 finishes validation at batch 91. Rank 0 hangs at batch 92.

If we set the validation step to a fixed number that is smaller than the total steps like this:

[validation]
enabled = true
local_batch_size=1
dataset = "alpaca_validation"
freq = 2
steps = 5

It doesn't hang. So our guess is that it is because the dataloder yields uneven number of batches in two ranks.

cc @ebsmothers

Versions

Code pointer provided above

torch 2.9.0a0+gite7cc42d

[fegin]

Yes, this looks like to be the case. In general, there is a debate whether FSDP or dataloader/trainer should be the one to handle the case. The current design is that either trainer or dataloader need to do the error handling not FSDP.

cc., @weifengpy please correct me if I'm wrong.

[ebsmothers]

Thanks @fegin, this makes sense. Also per @andrewkho's old comment it seems that there is not really a clean way to do this through IterableDataset as of today. I wonder whether it's reasonable to require a fixed number of validation steps in the validation config? I think training has already moved from epoch- to step-centric, I know it's a bit more common for validation to be epoch-centric, but with iterable datasets I don't really see a way to avoid these hangs in general (at least without some additional intrusive code)

[tianyu-l]

@ebsmothers that makes sense -- what would be good default? I think the steps should be almost total number of examples / (batch size * dp degree), but I'm not sure what's the best UI for setting this. I'm worried that if we don't provide such recommended limit, users don't have an idea how much they should set.

cc @wesleytruong

[ebsmothers]

@tianyu-l agreed. And with iterable dataset, it may not always be straightforward to get the total # of examples anyways (e.g. for streaming I don't think it's known beforehand?). Personally I would just hardcode to some value that won't take too long (say 100), but ofc this is extremely unscientific. Another option is to set to some multiple of validation.freq to control (approximately) how much time is spent on validation vs on training.

[DNXie]

@ebsmothers Can we let users to alternatively pass validation.epoch, and we can calculate the number of steps for them based on the dp degree and minus some offset to avoid the hang?

[tianyu-l]

@DNXie
I feel it'd be hard to calculate the number of steps without actually running data loading, because of the packing behavior.

Also I'm assuming the epoch would also be <= 1. There's not much point running validation beyond a single pass of the validation dataset.

[DNXie]

@tianyu-l Thanks for explaining. Makes sense.

[fegin]

@ebsmothers Yes, it is not a small change. I have done this for some internal models before. This requires using async CPU communication with an additional thread to avoid affecting the main training/validation thread.

[tianyu-l]

@ebsmothers @DNXie
Based on our discussions, I'm OK with setting some fixed steps instead of just hang. E.g. we can set a default in config / toml files, with proper comments to explain our choices.

Please feel free to submit PRs on this -- please address both https://github.com/pytorch/torchtitan/blob/main/torchtitan/components/validate.py and https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/flux/validate.py

[DNXie]

@tianyu-l I can take a look at these two files. Do you mean we no longer support validation.steps=-1?

[tianyu-l]

@DNXie
I think we can still allow it, but do the following to configs / tomls

1. clearly comment that when setting to -1 there could be hangs due to mismatch among ranks
2. change the default to something reasonable (e.g. under common setup of DP degree + sequence length + batch size, consume some amount of validation data, but clearly won't hit the end) -- I think this could be a bit challenging based on our discussions, and we could lean on the smaller side to be safe.

[DNXie]

@tianyu-l @ebsmothers
Submitted a PR for this, please check.