Wrap batch_sampler with _IndexBatchSamplerWrapper
State: closed | Pull request: yes
Labels: name: Run CICD

# What does this PR do ?

### NeMo `MegatronDataSampler` is not tracking the `batch_indices`.

In BioNeMo 2.0 we use `PredictionWriter` callback of `BasePredictionWriter` that supposed to capture the `batch_indices` from dataloader's batch sampler and predictions from `predict_step` of the Lightning module. However, `batch_indices` are not being tracked properly and that prevents us from mapping the predictions to input sequences.

`MegatronDataSampler` is [transforming the dataloaders](https://github.com/NVIDIA/NeMo/blob/1757ff9ed10272bf5ee7332d64fccd4bd9676f1b/nemo/lightning/pytorch/plugins/data_sampler.py#L67) by adding megatron batch sampler. But [Lightning prediction loop](https://github.com/Lightning-AI/pytorch-lightning/blob/8ad3e29816a63d8ce5c00ac104b14729a4176f4f/src/lightning/pytorch/loops/prediction_loop.py#L304) tracks `batch_indices` only if the batch sampler is an instance of `lightning.pytorch.overrides.distributed._IndexBatchSamplerWrapper`. To get this fixed I propose the following change in [nemo.lightning.data.add_megatron_sampler](https://github.com/NVIDIA/NeMo/blob/1757ff9ed10272bf5ee7332d64fccd4bd9676f1b/nemo/lightning/data.py#L218)

```python
def add_megatron_sampler(...)
    ...
    return DataLoader(
        dataloader.dataset,
        batch_sampler=_IndexBatchSamplerWrapper(batch_sampler),
        num_workers=dataloader.num_workers,
        pin_memory=dataloader.pin_memory,
        persistent_workers=dataloader.persistent_workers,
        collate_fn=dataloader.collate_fn,
    )
```

# GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

# Before your PR is "Ready for review"
**Pre checks**:
- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] 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.

CC: @jstjohn @ericharper 
BioNeMo PR: https://github.com/NVIDIA/bionemo-framework/pull/319 

## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)


> 3e3e39b

Sure, I made the changes here: [3e3e39b](https://github.com/NVIDIA/NeMo/pull/10934/commits/3e3e39b5e1a93a85431333eafc74d80f4f4942ae)
I have also removed the `drop_last` arg that I have added in https://github.com/NVIDIA/NeMo/pull/10837 since we can use the dataloader `mode` to avoid dropping the incomplete batches only during inference.

One of the commits is not signed and I wasn't able to fix that by reverting and recommitting the changes. Is there anyway to fix that?

Hi @cuichenx, thanks for the review. How do I see what caused the CI failure?

> Hi @cuichenx, thanks for the review. How do I see what caused the CI failure?

CI was cancelled yesterday to expedite some higher priority jobs. I will restart CI

@cuichenx There is an import error in the failed tests. Isn't `lightning` a dependency?
```
    from lightning.pytorch.overrides.distributed import _IndexBatchSamplerWrapper
E   ModuleNotFoundError: No module named 'lightning'
```