[NeMo2] Lightning Interaction with resuming within the same epoch
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

NeMo has some custom logic to resume the dataloader state with Megatron-LM custom Batch sampler based on consumed samples and microbatches calculated by mcore.
But it does not seem to account for how Lightning itself attempts to restore the data fetcher state for within the training epoch loop.

This results in a previous state from a checkpoint being accounted for in the dataloader / data fetcher iterator, and can cause StopIteration to be raised early before reaching the expected max_steps



**Steps/Code to reproduce bug**

For example, if my initial training run setup is as follows
```
trainer:
  accelerator: gpu
  devices: 1
  num_nodes: 1
  max_epochs: 1
  max_steps: 1000
  limit_val_batches: 2
  num_sanity_val_steps: 2

data:
  micro_batch_size: 1
  global_batch_size: 8
```

My data module should prepare max_steps * gbs = 1000 * 8 = 8000 samples total
I save a checkpoint at step 799, which has consumed 800 * gbz = 800 * 8 = 6400 samples

When I resume from checkpoint saved at step 799, it results in something like:
```
[NeMo I data_samplers:76] Instantiating MegatronPretrainingSampler with total_samples: 8040 and consumed_samples: 6400
```

With 200 steps remaining before 1,000 max_steps, I should expect to have 200 * 8 = 1600 samples remaining which must get loaded by the resumed dataloader.


However, pytorch lightning has the following logic in the TrainingEpochLoop
```
    def on_run_start(self, data_fetcher: _DataFetcher) -> None:
    	...

        # add the previous `fetched` value to properly track `is_last_batch` with no prefetching
        data_fetcher.fetched += self.batch_progress.current.ready
        ...
```

As a result, this is the state of the lightning data fetcher:
```
{
    '_batch': None,
    '_batch_idx': 0,
    '_dataloader_idx': 0,
    'done': False,
    'fetched': 798,
    'iterator': < pytorch_lightning.utilities.combined_loader.CombinedLoader object at 0x7fc9240ebeb0 > ,
    'iterator_wrapper': < pytorch_lightning.loops.fetchers._DataFetcherWrapper object at 0x7fc92410b730 > ,
    'length': 1640
}
```

The internal pytorch lightning tracked attributed `_batch_idx` and `fetched` are not synced, which results in a StopIteration being raised early when the following condition is met:
```
class _DataFetcher(Iterator):
    @override
    def __next__(self) -> _ITERATOR_RETURN:
        ....
        self.fetched += 1
        if self.length is not None:
            self.done = self.fetched >= self.length  <---------------
        return batch

class _DataLoaderIterDataFetcher(_DataFetcher):
    @override
    def __next__(self) -> Iterator["_DataFetcherWrapper"]:  # type: ignore[override]
        if self.done:   <---------------
            raise StopIteration
        return self.iterator_wrapper
```


If trainer.max_epochs == 1, then training stops immediately before the expected 1000 steps. If unset, then it jumps to epoch 1 unexpectedly:
```
...
[0]:Training epoch 0, iteration 901/999 | global_batch_size: 8 | val_loss: 9.243 | lr: 3e-06 | consumed_samples: 7216 | global_step: 901 | reduced_train_loss: 9.531 | train_step_timing in s: 0.4671
[0]:Training epoch 0, iteration 902/999 | global_batch_size: 8 | val_loss: 9.243 | lr: 3e-06 | consumed_samples: 7224 | global_step: 902 | reduced_train_loss: 9.362 | train_step_timing in s: 0.3907
[0]:Training epoch 0, iteration 903/999 | global_batch_size: 8 | val_loss: 9.243 | lr: 3e-06 | consumed_samples: 7232 | global_step: 903 | reduced_train_loss: 9.038 | train_step_timing in s: 0.4661
[0]:`Trainer.fit` stopped: `max_epochs=1` reached.
```


**Expected behavior**

Instead, we expect to see the full 1000k steps completed when we resume from a checkpoint in the middle of an epoch:

```
...
[0]:Training epoch 0, iteration 997/999 | global_batch_size: 8 | val_loss: 9.243 | lr: 3e-06 | consumed_samples: 7984 | global_step: 997 | reduced_train_loss: 9.112 | train_step_timing in s: 0.4716
[0]:Training epoch 0, iteration 998/999 | global_batch_size: 8 | val_loss: 9.243 | lr: 3e-06 | consumed_samples: 7992 | global_step: 998 | reduced_train_loss: 9.261 | train_step_timing in s: 0.4639
[0]:Training epoch 0, iteration 999/999 | global_batch_size: 8 | val_loss: 9.243 | lr: 3e-06 | consumed_samples: 8000 | global_step: 999 | reduced_train_loss: 9.092 | train_step_timing in s: 0.4736
[0]:`Trainer.fit` stopped: `max_steps=1000` reached.
```


A simple fix to get the expected behavior above is:

```
class CustomTrainingLoop(_TrainingEpochLoop):

    def on_run_start(self, data_fetcher: _DataFetcher) -> None:
        super().on_run_start(data_fetcher)
        data_fetcher.fetched = 0

def configure_custom_training_loop(trainer: pl.Trainer) -> None:
    if type(trainer.fit_loop.epoch_loop) is not _TrainingEpochLoop:
        print("Detected custom epoch loop. Skipping customizing training loop.", UserWarning)
        return

    loop = CustomTrainingLoop(trainer, trainer.min_steps, trainer.max_steps)
    trainer.fit_loop.epoch_loop = loop
```

Alternatively, we could also modify the logic in BaseMegatronSampler (nemo/collections/nlp/data/language_modeling/megatron/data_samplers.py) to return the total samples, but this could likely cause issues with NeMo1.
```
    def __len__(self):
        num_available_samples: int = self.total_samples - self.consumed_samples
        if self.global_batch_size is not None:
            if self.drop_last:
                num_global_batches = num_available_samples // self.global_batch_size
            else:
                num_global_batches = (num_available_samples + self.global_batch_size - 1) // self.global_batch_size
            # return len of dataloader in terms of micro batches to avoid discrepancy between len of dataloader and
            # num of batches fetched (as training step fetches in terms of micro batches)
            return num_global_batches * (self.global_batch_size // self.micro_batch_times_data_parallel_size)
        else:
            return (num_available_samples - 1) // self.micro_batch_times_data_parallel_size + 1
```


**Environment**

```
pytorch-lightning                  2.4.0
nemo-toolkit                       2.1.0rc0
megatron-core                      0.10.0rc0
```

**Additional context**

relevant: https://github.com/Lightning-AI/pytorch-lightning/issues/19764

Add any other context about the problem here.
Example: GPU model


root cause was due to usage of https://github.com/NVIDIA/NeMo/blob/9a268dd6cec33e0e7190cf9e1543e6473c32da2c/nemo/collections/nlp/data/language_modeling/megatron/data_samplers.py#L131

instead of the data samplers defined in https://github.com/NVIDIA/NeMo/blob/main/nemo/lightning/data.py