Fix bug in MegatronPretrainingRandomSampler
State: closed | Pull request: yes
Labels: name: stale

name: NLP

# What does this PR do ?

Bug fix for `MegatronPretrainingRandomSampler`, which may affect NLP collections in NeMo
 
**Collection**: NLP

# Changelog 

**Changed Files**: `nemo/collections/nlp/data/language_modeling/megatron/data_samplers.py`

- [MegatronPretrainingRandomSampler](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/data/language_modeling/megatron/data_samplers.py#L162) requires `self.data_parallel_size` to calculate `bucket_offset`, but `self.data_parallel_size` is missing in the attributes.
- It will raise `AttributeError: 'MegatronPretrainingRandomSampler' object has no attribute 'data_parallel_size'. Did you mean: 'data_parallel_rank'?` when we tried to run llama2 models with `dataloader_type=cyclic`

# Reproduce steps

1. Use configuration `dataloader_type: cyclic` to launch llama2 models with [configuration file](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/language_modeling/conf/megatron_gpt_config.yaml#L210)

``` yaml
  data:
    data_impl: mmap
    dataloader_type: cyclic
    skip_warmup: true
    splits_string: "99990,8,2"
    seq_length: 4096
    reset_position_ids: false
    reset_attention_mask: false
    eod_mask_loss: false
    index_mapping_dir: null
    data_prefix:
    - ${data_dir}/my-llama_00_text_document
```

2. Then you will get the error msg shown below

``` bash
  File "/mnt/petrelfs/caizheng/miniconda3/envs/nemo/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1084, in __init__
    self._reset(loader, first_iter=True)
  File "/mnt/petrelfs/caizheng/miniconda3/envs/nemo/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1117, in _reset
    self._try_put_index()
  File "/mnt/petrelfs/caizheng/miniconda3/envs/nemo/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1351, in _try_put_index
    index = self._next_index()
  File "/mnt/petrelfs/caizheng/miniconda3/envs/nemo/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 623, in _next_index
    return next(self._sampler_iter)  # may raise StopIteration
  File "/mnt/petrelfs/caizheng/dev/NeMo/nemo/collections/nlp/data/language_modeling/megatron/data_samplers.py", line 163, in __iter__
    bucket_offset = current_epoch_samples // self.data_parallel_size
AttributeError: 'MegatronPretrainingRandomSampler' object has no attribute 'data_parallel_size'. Did you mean: 'data_parallel_rank'?
```

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [x] Did you write any new necessary tests?
- [x] 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
- [x] 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](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)