Handle float limit_val_batches

[athitten]

What does this PR do ?

1. When a float limit_val_batches is passed by the user, the PR ensures that the limit_val_batches is cast to an equivalent int value and also makes sure the cast limit_val_batches is a multiple of num of microbatches as required by PTL >= 2.0.

2. Calls self._reconfigure_val_batches() after the setup of datasets, so that the original value of limit_val_batches is used to build the dataset and not the reconfigured value.

3. Returns the len(dataloader) in terms of number micro batches instead of num of global batches. This is required for 2 reasons:

• a) Since limit_val_batches is reconfigured to be in terms of number of micro batches, if the len(dataloader) is in num of global batches then we can run into situations where len(dataloader) < num_micro_batches and this can lead to one of the ranks hitting StopIteration in the midst of completing a global batch, leading to a hang if a different rank is waiting for the output of the first rank in case of PP.

• b) Another reason being ideally, the len(dataloader) should be returned in terms of the granularity of the batch size in which we fetch the data from the dataloader_iter. Since in megatron models a micro batch is fetched each time next(dataloader_iter) is called, the len(dataloader) should be returned in the same metric. Also, PTL's progress bar is such that it increments the epoch number after the num of batches extracted hit len(dataloader). So if len(dataloader) is x in terms of global batches, then we incorrectly increment the epoch after x micro batches are extracted even though, the dataloader still has microbatches and is not empty. This can be very misleading to the end users.

Add a one line overview of what this PR aims to accomplish.

Collection: [Note which collection this PR will affect]

Changelog

• Add specific line by line info of high level changes in this PR.

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

Jenkins CI

To run Jenkins, a NeMo User with write access must comment jenkins on the PR.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

• Related to # (issue)

[athitten]

jenkins

[athitten]

jenkins

[athitten]

jenkins

[athitten]

jenkins

[athitten]

jenkins

[athitten]

jenkins

[athitten]

jenkins

[athitten]

jenkins

[athitten]

Having this is currently causing an index error with blended dataset from mcore. Also I couldn't understand the purpose of having this. Waiting for @shanmugamr1992 's comments on it.

[jbaczek]

Leaving here a comment from an offline conversation with the explanation:

This is a hack that uses quirk in the dataset implementation that causes the behaviour explained in the comment.
This array is used to supply the N value in the E = argmin_e e * len(data) >= N formula. Then E is used to construct indices which are then used as an internal representation of dataset.
The dataset is not trimmed to N after the construction, but to the E*len(data) . So setting this to 1, we are sure that E is exactly 1, which gives us a single iteration over a validation split.
[...]
In plain words: the number of epochs E is the smallest integer number that when multiplied by length of the data is greater than requested number of samples

[athitten]

jenkins

[athitten]

jenkins

[jbaczek]

What is the problem here? Why building dataset with 0 samples would cause an error? In PTL setting this option to 0 is an indicator to skip validation loop entirely. The dataloader should never be invoked in such case.
Here is the logic responsible for that.

[athitten]

Building dataset with 0 samples causes this error: index -1 is out of bounds for axis 0 with size 0 in this line in blended_dataset.py in mcore. I agree with you that the dataset and the dataloaders shouldn't be constructed if the num_samples is 0. Since the dataset building happens in mcore, this needs to be done there. This can't be done in NeMo since we have one func call to build all 3 of train, val and test datasets here

[athitten]

For this error index -1 is out of bounds for axis 0 with size 0 mentioned above, the index becomes -1 in this line since self.size = 0 (if limit_val_batches=0)

[athitten]

Let me check in mcore if we can make a PR, to add a condition check to not build dataset if the size or num_samples is 0.

[jbaczek]

Can you post the full stack trace for the error you get?

[athitten]

Exception has occurred: IndexError       (note: full exception trace is shown but execution is paused at: _run_module_as_main)
index -1 is out of bounds for axis 0 with size 0
  File "/usr/local/lib/python3.10/dist-packages/megatron/core/datasets/blended_dataset.py", line 92, in __getitem__
    dataset_id = self.dataset_index[idx]
  File "/usr/local/lib/python3.10/dist-packages/megatron/core/datasets/blended_dataset.py", line 81, in __init__
    _ = self[self.size - 1]
  File "/usr/local/lib/python3.10/dist-packages/megatron/core/datasets/blended_megatron_dataset_builder.py", line 276, in build_generic_dataset
    dataset = cls(*args)
  File "/usr/local/lib/python3.10/dist-packages/megatron/core/datasets/blended_megatron_dataset_builder.py", line 117, in _build_blended_dataset_splits
    self.build_generic_dataset(
  File "/usr/local/lib/python3.10/dist-packages/megatron/core/datasets/blended_megatron_dataset_builder.py", line 61, in build
    return self._build_blended_dataset_splits()
  File "/workspace/software/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 1235, in build_train_valid_test_datasets
    ).build()
  File "/workspace/software/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 1325, in setup
    self.build_train_valid_test_datasets()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 145, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 86, in _call_setup_hook
    _call_lightning_module_hook(trainer, "setup", stage=fn)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 941, in _run
    call._call_setup_hook(self)  # allow user to setup lightning_module in accelerator environment
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 571, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 93, in launch
    return function(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 532, in fit
    call._call_and_handle_interrupt(
  File "/workspace/software/NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py", line 38, in main
    trainer.fit(model)
  File "/usr/local/lib/python3.10/dist-packages/hydra/core/utils.py", line 186, in run_job
      ret.return_value = task_function(task_cfg)
  File "/usr/local/lib/python3.10/dist-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/workspace/software/NeMo/nemo/core/config/hydra_runner.py", line 129, in wrapper
    _run_hydra(
  File "/workspace/software/NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py", line 42, in <module>
    main()
  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main (Current frame)
    return _run_code(code, main_globals, None,
IndexError: index -1 is out of bounds for axis 0 with size 0

[athitten]

It errors out at this line blended_dataset.py#L92 with the above error.

[jbaczek]

Separating arguments from program logic is a good practice. You could move all this arithmetic to the argument layer, as none of these variables has to be computed in the runtime.

[athitten]

Yes agreed, in case of NeMo this can probably be moved inside the __init__ of MegatronBaseModel.

[jbaczek]

How do infinite validation sets work in nemo? Is this even a valid approach? Shouldn't this be caught on dataset on argument layer, even before dataset construction?

[athitten]

I ported that condition check over from PTL as they check if length of val dataloader is not inf here while casting float limit_val_batches to int.
I don't think NeMo has any cases for infinite validation sets.

[athitten]

The code now does the round down logic only if limit_val_batches > 0.0 and less than 1.0.

[jbaczek]

As I understand this, we have 5 cases if limit_val_batches is a float:

1. limit_val_batches == 0.0. Then we stick with this value and let PTL skip validation
2. limit_val_batches == 1.0. Then we set self.trainer.limit_val_batches = len(self._validation_dl), even if it's not divisible by get_num_microbatches()
3. limit_val_batches * len(self._validation_dl) < 1. Then we raise an exception
4. get_num_microbaches() > limit_val_batches * len(self._validation_dl) > 1. Then we fix limit_val_batches = get_num_microbaches() to run at least one iteration.
5. limit_val_batches * len(self._validation_dl) > get_num_microbaches(). Then we round down to an integer multiple of get_num_microbatches()

@athitten do I understand this correctly? Doesn't it hang if limit_val_batches == 1.0?

Can you write it as a series of if/elif cases? It would be way easier to read instead of following nested ifs.

[athitten]

Nope, it does not hang when limit_val_batches == 1.0. When limit_val_batches == 1.0, we use the full len of the dataloader as the limit_val_batches and since the len(dataloader) is already in microbatches, it is divisible by get_num_microbatches(). The situation where it needs to be ensured to be a multiple of get_num_microbatches(), arises when 0<limit_val_batches<1.0. Basically a fraction.

[jbaczek]

Does this logic mean that we resign from the PTL Trainer's utility to handle floats?
What is the exact algorithm here?
AFAIU:

1. We have a val dataset of length L
2. It is sampled B=get_num_microbatches() at the time.
3. This implies ceil(L/B) validation iterations
4. We pad the dataset so L/B == L//B

So my questions here:
Why can't we draw a partial batch from the dataset? Does it mess too much with PP/DP? In what way?
Why do we round down in the else case, but in the if case we round up?

[athitten]

PTL handles float limit_val_batches in the same as we have done above (see: data_connector.pyL449 and data_connector.py#L457.

The only difference being, after the step of casting the float limit_val_batches to an equivalent int is done, we are ensuring that limit_val_batches is a multiple of num_of_microbatches.

This is required since in each global training_step say we extract x microbatches, and if limit_val_batches is not equal to x at the least or not equal to multiple of x, then this line raises StopIteration(cause PTL bounds val_dataloder by limit_val_batches and raises StopIteration when num of microbatches extracted reaches limit_val_batches) in between a global step leading to a hang in case of PP, if PP 1 hit a StopIteration and PP2 is waiting for the output of PP1. I hope this made sense. Let me know if it's unclear.

[athitten]

To answer your question on if and else condition: In the if case we are making sure if limit_val_batches > 0 and < num_micro_batches then it's at least equal to the num_micro_batches so that we can run 1 validation_step successfully. The else is for cases when limit_val_batches > num_micro_batches, almost like we are cutting out the incomplete batch that does not have as many microbatches as the get_num_microbatches()

[jbaczek]

Thanks! Now it's way clearer.

Regarding if/else: Is there an assumption that number of batches in the validation split is always greater than get_num_micro_batches()? Is this what allows rounding up in the if case?
So there is no way to pad dataset to desired length with samples that won't be counted to the final loss?

I don't think we are allowed to validate on a strict subset of the validation data. @ShriyaPalsamudram can you confirm that?

[ShriyaRishab]

No, we should not cut out the last incomplete batch by default. There is this flag validation_drop_last which decides if the last incomplete batch is dropped or not. This section handles validation_drop_last=False case in loss computation as well.

[athitten]

There's no explicit condition/assumption existing in NeMo currently, that checks num of micro batches in the validation dataloader is always greater than get_num_micro_batches(). Rounding up is added in the if condition, cause otherwise we wont even be able to run one global step of validation for reasons mentioned above and it can lead to a hang.

In the else condition, if limit_val_batches=1.0, then there are no partial batches at all. There would be partial batches only when limit_val_batches is a decimal < 1.0. In that case what we are doing in the else condition is not technically chopping the batch. What I mean is, we are basically allowing to run certain percentage of the val_dataset via float limit_val_batches. In order to make this happen we have to translate it to equivalent int value (which is the way even PTL handles. For example if limit_val_batches is 0.25 and num of batches is 6 then 0.25*6=1.5 and int(1.5)=1, but this does not mean we are chopping that 0.5 of the second batch right ?). The logic we have is just ensuring the correct value of limit_val_batches. Basically in each validation step you run limit_val_batches num of batches and if we skipped any batch in that final value of limit_val_batches with which we run, then that would be like cutting out a batch. Let me know your thoughts.

[jbaczek]

I agree that the dataloader shouldn't return partial batches and probably it would be better to round the number down. But I don't fully get what you mean by:

Basically in each validation step you run limit_val_batches num of batches and if we skipped any batch in that final value of limit_val_batches with which we run, then that would be like cutting out a batch. Let me know your thoughts.

[athitten]

jenkins

[athitten]

jenkins

[athitten]

jenkins

[jbaczek]

LGTM

[jbaczek]

jenkins