SpeechLM2 : You're resuming from a checkpoint that ended before the epoch ended and your dataloader is not resumable.

[AudranBert]

Describe the bug

Hi @pzelasko,
When resuming a Speechlm2 training I get this log that is a bit worrying. I remember for ASR seeing a Nemo log telling to ignore it, but I do not see a similar Nemo log here. Does it correctly load where it was when it ended (even when using batch_tokens)?

nemo-2.4.0+py3.12.10/lib/python3.12/site-packages/lightning/pytorch/loops/training_epoch_loop.py:161: You're resuming from a checkpoint that ended before the epoch ended and your dataloader is not resumable. This can cause unreliable results if further training is done. Consider using an end-of-epoch checkpoint or make your dataloader resumable by implementing the `state_dict` / `load_state_dict` interface.

(Do not pay attention to nemo-2.4.0 we are using NeMo 2.8.0rc0 but to avoid NaN loss we use or 2.4 env)

Steps/Code to reproduce bug

Expected behavior

Should resume as it was when it ended.

Environment overview (please complete the following information)

Environment details

NeMo 2.8.0rc0

Additional context

Add any other context about the problem here.
H100

[zeel2104]

I would like to work on it. Could you assign it to me?

[pzelasko]

The dataloader doesn't support resumption. It relies on setting new random seed for the resumed run to iterate data in a different order. It happens automatically by default, see: https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/asr/datasets.html#seeds-and-randomness

We're planning to do a dataloader revision that will enable precise resumption in the next couple of months.

[AudranBert]

The dataloader doesn't support resumption. It relies on setting new random seed for the resumed run to iterate data in a different order. It happens automatically by default, see: https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/asr/datasets.html#seeds-and-randomness

We're planning to do a dataloader revision that will enable precise resumption in the next couple of months.

Hi @pzelasko, thanks for the quick response. That's not very good news since we are limited to 20H jobs in our case.
When resuming it changes the seed automatically (despite having seed: X) in the yaml?
For us a few months is quite a long time. Do you know if it is a complex change to implement, maybe I can work on it? If the seed is the same, it means the data would be seen in the same order maybe there is a way to store what was the last batch to skip (using IDs in the data ?) until reaching the new data?

Have a great day,

[pzelasko]

Exact resumption is unfortunately fairly complex due to in-memory buffers and RNGs sprinkled all over the dataloading pipeline. I have two designs ready but they need thorough testing, and I don't want to overpromise the timeline here.

The good news is that it doesn't seem to matter that much in practice, and we can consistently reproduce our models despite non-deterministic dataloader. In the expectation the model observes amounts of data proportional to the weights defined in input_cfg from each source, and with a fixed max_steps budgets, the outcomes are very similar across runs. We have 4H jobs where this "just works" with many resumptions. Unless you hit specific edge cases, described here (with a fix) #14558

When resuming it changes the seed automatically (despite having seed: X) in the yaml?

Yes, please read the attached doc.

[AudranBert]

Exact resumption is unfortunately fairly complex due to in-memory buffers and RNGs sprinkled all over the dataloading pipeline. I have two designs ready but they need thorough testing, and I don't want to overpromise the timeline here.

The good news is that it doesn't seem to matter that much in practice, and we can consistently reproduce our models despite non-deterministic dataloader. In the expectation the model observes amounts of data proportional to the weights defined in input_cfg from each source, and with a fixed max_steps budgets, the outcomes are very similar across runs. We have 4H jobs where this "just works" with many resumptions. Unless you hit specific edge cases, described here (with a fix) #14558

When resuming it changes the seed automatically (despite having seed: X) in the yaml?

Yes, please read the attached doc.

Ok thanks for the detailed answer.
Not related but a small question : In the meantime of our experiments I'm testing the bucketing system but the oomptimizer needs tokens for buckets, whereas there is no script to estimate buckets size with tokens, I only find this which works with durations (As a side note the script doesn't seem to support multimodal conversation format, it was an easy fix though). Am I missing something? I reread your messages on the subject but didn't find the answer.

[pzelasko]

Use this for token estimation: https://github.com/NVIDIA-NeMo/NeMo/blob/main/scripts/speech_llm/estimate_token_bins.py

And this for OOMptimizer: https://github.com/NVIDIA-NeMo/NeMo/blob/main/scripts/speechlm2/oomptimizer.py

Note to self: clean-up how scripts are organized and update docs to point to these too.

[AudranBert]

Use this for token estimation: https://github.com/NVIDIA-NeMo/NeMo/blob/main/scripts/speech_llm/estimate_token_bins.py

And this for OOMptimizer: https://github.com/NVIDIA-NeMo/NeMo/blob/main/scripts/speechlm2/oomptimizer.py

Note to self: clean-up how scripts are organized and update docs to point to these too.

Ok thanks, I don't know how I missed the speech_llm folder but yes I was using the right OOMptimizer. For your note, as a user I don't really understand the difference between the tools and the scripts folders, I don't know if there is a difference on what goes where, but I don't have it.

[pzelasko]

Good feedback. Tools seem more complex than scripts, almost like something that could live in a repo of its own. I'll add a note in our "repo cleanup tracker" about this.