Hang when near OOM in FSDP mode

[luciaquirke]

Potential issue where the collection run hangs at 0% if the VRAM usage is very close to 100%. Accepting replication details on this one.

[davidoj]

On coreweave nodes with 8xA40 GPUs

bergson build runs/data/facts_dataset_reworded.hf/index --model tmp/checkpoint-282 --dataset data/facts_dataset_reworded.hf --drop_columns False --prompt_column fact --completion_column reworded --fsdp --projection_dim 16 --skip_preconditioners --overwrite --token_batch_size 6000 --truncation

runs successfully (and in a few minutes) while

bergson build runs/data/facts_dataset_reworded.hf/index --model tmp/checkpoint-282 --dataset data/facts_dataset_reworded.hf --drop_columns False --prompt_column fact --completion_column reworded --fsdp --projection_dim 16 --skip_preconditioners --overwrite --token_batch_size 8192 --truncation

hangs indefinitely on step 0 - the difference is token_batch_size of 6000 vs 8192. nvidia-smi suggests in the latter case VRAM is close to 100% utilised.

The model is a LoRA on Qwen/Qwen3-4B

[LouisYRYJ]

Btw this doesnt need to be at step 0. The first batches could have sequences with short length and run fine and then at some point might hit a batch with bigger sequences which will cause it to hang. That happened to me iirc.
This doesn't really change how we should fix it, just wanted to flag that

[luciaquirke]

@LouisYRYJ noted, thanks. I merged the validate_batch_size util recently which runs the max number of tokens through the model upfront so if it's working correctly all the hangs should be at step 0 from now on.