Repository: NVIDIA/NeMo
Pull request: fix tarred dataset len when num shards is not divisible by workers
Signed-off-by: Iztok Lebar Bajec <ilb@fri.uni-lj.si>

# What does this PR do ?

Fixes #4522. When using tarred datasets and the number of shards is not divisible by `world_size`, nemo logs a warning, but training continues. However, in this case the computation of `DataLoader.len()` becomes incorrect. This leads to invalid `lr_scheduler.max_steps`, and other issues down the road. One example being `Validation check` not running at end of epoch. 

**Collection**: [nlp]
- language_modelling
- machine_translation
- text_normalization
- token_classification

# Changelog 
In case of `shard_strategy='scatter'` the `len()` is computed based on the amount of taken shards and assuming total `num_baches` are equally distributed over all shards, i.e. that each shard contains the same number of batches. In addition, following other implementations this PR adds `shard_strategy` support to NLP/PC tarred datasets.

There are two additional occasions, where tarred datasets are used  
https://github.com/NVIDIA/NeMo/blob/fa2e55e607f787d129141523c724b29dab7a411c/nemo/collections/asr/data/audio_to_text.py#L186
which based on the implementation does not seem to be affected by this issue, and
https://github.com/NVIDIA/NeMo/blob/fa2e55e607f787d129141523c724b29dab7a411c/nemo/collections/nlp/data/language_modeling/l2r_lm_dataset.py#L169-L184
which might, but I chose to leave it as is, as the computation of `len()` 
https://github.com/NVIDIA/NeMo/blob/fa2e55e607f787d129141523c724b29dab7a411c/nemo/collections/nlp/data/language_modeling/l2r_lm_dataset.py#L239-L240
greatly differs from others
  
**PR Type**:
- [ ] New Feature
- [x] Bugfix
- [x] Documentation

## 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.
note to @PeganovAnton for NLP/C

# Additional Information
* Related to #4522


Problem statement:
[NLP] punctuation_capitalization resume training fails
**Describe the bug**

When training PC from scratch, using the example configs, everything works fine, however if I try to resume an interrupted training from the last checkpoint, I get the following log entry
```
[NeMo W 2022-07-07 21:18:38 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:378: UserWarning: `ModelCheckpoint(monitor='val_loss')` could not find the monitored key in the returned metrics: ['epoch', 'step']. HINT: Did you call `log('val_loss', value)` in the `LightningModule`?
      warning_cache.warn(m)
```
and all new checkpoints get a `val_loss=0.0000`

**Environment details**

- pytorch:22.04-py3
- nemo:1.8.2 installed with ./reinstall.sh, without numba install

Hints from issue discussion:
Installed `nemo` as you described and could not reproduce the bug. Could you please describe exact setting?
@itzsimpl 
@PeganovAnton I apologise for taking so long, I was trying to debug. I haven't found the root cause, but by changing to pytorch:22.06-py3 and nemo:1.10.0, I don't see the issue anymore. I'll just continue with nemo:1.10.0. I did a couple of minor modifications to the model codebase (support for Uppercase, Lowercase and Capitalised), but these should not have any influence. Especially since I ported them to nemo:1.10.0 and everything works as it should. 
FWW. I was running the following command:
```
python examples/nlp/token_classification/punctuation_capitalization_train_evaluate.py \
  --config-path=conf \
  --config-name=punctuation_capitalization_config \
  \
  exp_manager.name=punctuation_capitalization_v4 \
  exp_manager.exp_dir=/experiments \
  +exp_manager.version=20220629-1159 \
  +exp_manager.resume_if_exists=true \
  +exp_manager.resume_ignore_no_checkpoint=true \
  +exp_manager.checkpoint_callback_params.save_best_model=true \
  \
  model.language_model.pretrained_model_name=bert-base-multilingual-cased \
  model.train_ds.ds_item=/data/punct_v4 \
  model.train_ds.text_file=text_train.txt \
  model.train_ds.labels_file=labels_train.txt \
  model.train_ds.use_tarred_dataset=True \
  model.train_ds.tar_metadata_file=tarred_train/metadata.punctuation_capitalization.tokens15000.max_seq_length512.bert-base-multilingual-cased.json \
  model.validation_ds.ds_item=/data/punct_v4 \
  model.validation_ds.text_file=text_val.txt \
  model.validation_ds.labels_file=labels_val.txt \
  model.validation_ds.use_tarred_dataset=True \
  model.validation_ds.tar_metadata_file=tarred_val/metadata.punctuation_capitalization.tokens15000.max_seq_length512.bert-base-multilingual-cased.json \
  ~model.test_ds \
  \
  trainer.devices=-1 \
  trainer.accelerator=gpu \
  trainer.num_nodes=1 \
  trainer.precision=bf16 \
  trainer.max_epochs=150 \
  trainer.max_steps=-1 \
  +trainer.num_sanity_val_steps=0 \
  \
  +exp_manager.checkpoint_callback_params.monitor=val_loss \
  +exp_manager.checkpoint_callback_params.mode=min \
  +exp_manager.checkpoint_callback_params.save_top_k=5 \
  +exp_manager.checkpoint_callback_params.always_save_nemo=true \
  +do_training=true \
  +do_testing=false \
  ```
I just noticed, that I'm still getting the same issue, however, this time it did not happen immediately upon restarting the training, but after the first epoch finished, hence, why it seemed to work as it should. This is what I got
```
Epoch 56:  84%|████████▍ | 17999/21375 [47:31<-1:57:05, -19.29it/s, loss=0.0326, v_num=1159, lr=2.78e-5]
Epoch 56:  84%|████████▍ | 18000/2[NeMo W 2022-07-11 15:44:02 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:378: UserWarning: `ModelCheckpoint(monitor='val_loss')` could not find the monitored key in the returned metrics: ['lr', 'train_loss', 'train_backward_timing', 'train_step_timing', 'epoch', 'step']. HINT: Did you call `log('val_loss', value)` in the `LightningModule`?
      warning_cache.warn(m)
    
Epoch 56, global step 4106000: 'val_loss' was not in top 5
1375 [47:31<-1:57:06, -19.29it/s, loss=0.0326, v_num=1159, lr=2.78e-5]
Epoch 56:  84%|████████▍ | 18000/21375 [47:31<-1:57:06, -19.29it/s, loss=0.0326, v_num=1159, lr=2.78e-5]
```

In the initial part of the log I also noticed this
```
SLURM auto-requeueing enabled. Setting signal handlers.
[NeMo W 2022-07-11 14:56:22 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/utilities/data.py:152: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-process data loading (when num_workers > 1), `__len__` could be inaccurate if each worker is not configured independently to avoid having duplicate data.
      rank_zero_warn(
    
[NeMo W 2022-07-11 14:56:22 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/utilities/data.py:106: UserWarning: Your `IterableDataset` has `__len__` defined. In combination with multi-process data loading (when num_workers > 1), `__len__` could be inaccurate if each worker is not configured independently to avoid having duplicate data.
      rank_zero_warn(
    
[NeMo W 2022-07-11 14:56:28 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/webdataset/dataset.py:139: UserWarning: num_workers 8 > num_shards 3
      warnings.warn(f"num_workers {num_workers} > num_shards {len(urls)}")
    

Sanity Checking: 0it [00:00, ?it/s]
```
could any of these be related?
I'm really puzzled; if I set `trainer.val_check_interval` to e.g. 0.25 or anything less than 1.0 (as set in config yaml), I see the `Validation Dataloader` start and finish, and a valid checkpoint will be saved. 

If, however, I set it to 1.0, the `Validation Dataloader` will not start, the `val_loss` UserWarning will appear, and the checkpoint will be saved with `val_loss=0.0000`.

My changes to the codebase are not at fault, as they are a two line substitute for  
https://github.com/NVIDIA/NeMo/blob/dc6bea2acbfff899f06d3d7fa6ac98aec3aa4af3/nemo/collections/nlp/models/token_classification/punctuation_capitalization_model.py#L1005

Any thoughts @PeganovAnton?
I've just noticed another strange behaviour, with `trainer.val_check_interval=0.25` the Epoch progress in the logs does not reach 100%; as if the last validation check is still not run, i.e. I see the first 3 validations of size 3125 (starting at 15%, at 39% and at 64%) in the logs, but not the last one.

```
Epoch 0:  88%|████████▊ | 27000/30750 [50:49<07:03,  8.85it/s, loss=0.117, v_num=1352, lr=5.26e-5]
Epoch 0:  88%|████████▊ | 27000/30750 [50:50<07:03,  8.85it/s, loss=0.117, v_num=1352, lr=5.26e-5]
Epoch 0:   0%|          | 0/30750 [00:00<?, ?it/s, loss=0.117, v_num=1352, lr=5.26e-5]            
Epoch 1:   0%|          | 0/30750 [00:00<?, ?it/s, loss=0.117, v_num=1352, lr=5.26e-5]
```
@itzsimpl Comments on the command that you provided:
1. It is a bad idea to use tarred dataset for validation and testing.
    - When tarred dataset is created, several batches can be lost. You may see warning `f"Discarded {new_file_num_batches} batches."` during tarred dataset creation.
    - If you use DDP with several GPUs, then validation and testing are performed in several processes. Tar files are distributed evenly between processes and excess tar files are discarded.
 2. You do not need to provide `text_file` and `labels_file` if `use_tarred_dataset=true`.
@itzsimpl Could you please share your validation dataset? Both text files and tarred dataset.
@PeganovAnton Ok, noted. Thanks for the input. 

I started off with normal datasets, and moved to tarred ones based on the discussion in #4410 (https://github.com/NVIDIA/NeMo/pull/4410#discussion_r904108591),  I guess I never bothered to remove the non necessary non tarred parameters. 

If I understand correctly, all GPUs will still receive an equal amount of data, but the excess tar files (the portion not divisible by the number of used GPUS) will be discarded (i.e. at max `n-1` tar files will be dropped when using `n` GPUs). Bad from the standpoint of using all the data, but all GPUs should still have data to work with, so this should not cause the issue that I am seeing, right?

Unfortunately I can not share the datasets.

This is what I have at the moment:
- the tarred train dataset has 146 files
- the tarred val dataset has 25 files
- the same code ran on 2 GPUs for 55 epochs (starting from scratch) with no issues
- I tried to resume it on 8 GPus, but at the end of epoch 56 there was the `UserWarining` and no validation check, so I started investigating (is it at all possible to resume with a different number of GPUs then initially started?)
- on 4 or 8 GPUs I don't see the end of epoch validation check even if I start from scratch, with `trainer.val_check_interval=0.25` I see the first three validation checks run, but not the one at the end of the epoch
- on 2 GPUs with `trainer.val_check_interval=1.0, trainer.limit_train_batches=0.10, trainer.limit_val_batches=0.10` I can train from scratch as well as resume (though in the resume I see negative numbers for it/s and remaining time, but I assume this is another issue, perhaps related to #4541)
- on 2 GPUs retesting from scratch there is no `UserWarning` and I see the validation check run at the end of the epoch 
- in a couple of hours I'll see what happens if I resume from epoch 55 with 2 GPUs 

At the moment it points out to an issue when the number of GPUs is increased above 2.

Commit message: fix style

Patch to review:
diff --git a/docs/source/asr/datasets.rst b/docs/source/asr/datasets.rst
index 364c7fea1926..cd4ebc3f9a02 100644
--- a/docs/source/asr/datasets.rst
+++ b/docs/source/asr/datasets.rst
@@ -1,7 +1,7 @@
 Datasets
 ========
 
-NeMo has scripts to convert several common ASR datasets into the format expected by the ``nemo_asr`` collection. You can get started 
+NeMo has scripts to convert several common ASR datasets into the format expected by the ``nemo_asr`` collection. You can get started
 with those datasets by following the instructions to run those scripts in the section appropriate to each dataset below.
 
 If the user has their own data and want to preprocess it to use with NeMo ASR models, refer to the `Preparing Custom ASR Data`_ section.
@@ -13,8 +13,8 @@ If the user already has a dataset that you want to convert to a tarred format, r
 LibriSpeech
 -----------
 
-Run the following scripts to download the LibriSpeech data and convert it into the format expected by `nemo_asr`. At least 250GB free 
-space is required. 
+Run the following scripts to download the LibriSpeech data and convert it into the format expected by `nemo_asr`. At least 250GB free
+space is required.
 
 .. code-block:: bash
 
@@ -37,18 +37,18 @@ Fisher English Training Speech
 
 Run these scripts to convert the Fisher English Training Speech data into a format expected by the ``nemo_asr`` collection.
 
-In brief, the following scripts convert the ``.sph`` files to ``.wav``, slices those files into smaller audio samples, matches the 
-smaller slices with their corresponding transcripts, and splits the resulting audio segments into train, validation, and test sets 
+In brief, the following scripts convert the ``.sph`` files to ``.wav``, slices those files into smaller audio samples, matches the
+smaller slices with their corresponding transcripts, and splits the resulting audio segments into train, validation, and test sets
 (with one manifest each).
 
 .. note::
   - 106 GB of space is required to run the ``.wav`` conversion
   - additional 105 GB is required for the slicing and matching
-  - ``sph2pipe`` is required in order to run the ``.wav`` conversion 
+  - ``sph2pipe`` is required in order to run the ``.wav`` conversion
 
 **Instructions**
 
-The following scripts assume that you already have the Fisher dataset from the Linguistic Data Consortium, with a directory structure 
+The following scripts assume that you already have the Fisher dataset from the Linguistic Data Consortium, with a directory structure
 that looks similar to the following:
 
 .. code-block:: bash
@@ -67,7 +67,7 @@ that looks similar to the following:
       ├── fe_03_p2_sph3
       └── ...
 
-The transcripts that will be used are located in the ``fe_03_p<1,2>_transcripts/data/trans`` directory. The audio files (``.sph``) 
+The transcripts that will be used are located in the ``fe_03_p<1,2>_transcripts/data/trans`` directory. The audio files (``.sph``)
 are located in the remaining directories in an ``audio`` subdirectory.
 
 #. Convert the audio files from ``.sph`` to ``.wav`` by running:
@@ -78,7 +78,7 @@ are located in the remaining directories in an ``audio`` subdirectory.
      python fisher_audio_to_wav.py \
        --data_root=<fisher_root> --dest_root=<conversion_target_dir>
 
-   This will place the unsliced ``.wav`` files in ``<conversion_target_dir>/LDC200[4,5]S13-Part[1,2]/audio-wav/``. It will take several 
+   This will place the unsliced ``.wav`` files in ``<conversion_target_dir>/LDC200[4,5]S13-Part[1,2]/audio-wav/``. It will take several
    minutes to run.
 
 #. Process the transcripts and slice the audio data.
@@ -90,7 +90,7 @@ are located in the remaining directories in an ``audio`` subdirectory.
        --dest_root=<processing_target_dir> \
        --remove_noises
 
-   This script splits the full dataset into train, validation, test sets, and places the audio slices in the corresponding folders 
+   This script splits the full dataset into train, validation, test sets, and places the audio slices in the corresponding folders
    in the destination directory. One manifest is written out per set, which includes each slice's transcript, duration, and path.
 
    This will likely take around 20 minutes to run. Once finished, delete the 10 minute long ``.wav`` files.
@@ -100,8 +100,8 @@ are located in the remaining directories in an ``audio`` subdirectory.
 
 Run the following script to convert the HUB5 data into a format expected by the ``nemo_asr`` collection.
 
-Similarly, to the Fisher dataset processing scripts, this script converts the ``.sph`` files to ``.wav``, slices the audio files and 
-transcripts into utterances, and combines them into segments of some minimum length (default is 10 seconds). The resulting segments 
+Similarly, to the Fisher dataset processing scripts, this script converts the ``.sph`` files to ``.wav``, slices the audio files and
+transcripts into utterances, and combines them into segments of some minimum length (default is 10 seconds). The resulting segments
 are all written out to an audio directory and the corresponding transcripts are written to a manifest JSON file.
 
 .. note::
@@ -123,7 +123,7 @@ You can optionally include ``--min_slice_duration=<num_seconds>`` if you would l
 AN4 Dataset
 -----------
 
-This is a small dataset recorded and distributed by Carnegie Mellon University. It consists of recordings of people spelling out 
+This is a small dataset recorded and distributed by Carnegie Mellon University. It consists of recordings of people spelling out
 addresses, names, etc. Information about this dataset can be found on the `official CMU site <http://www.speech.cs.cmu.edu/databases/an4/>`_.
 
 #. `Download and extract the dataset <http://www.speech.cs.cmu.edu/databases/an4/an4_sphere.tar.gz>`_ (which is labeled "NIST's Sphere audio (.sph) format (64M)".
@@ -153,14 +153,14 @@ After the script finishes, the ``data`` folder should contain a ``data_aishell``
 Aishell-2
 ---------
 
-To process the AIShell-2 dataset, in the command below, set the data folder of AIShell-2 using ``--audio_folder`` and where to push 
-these files using ``--dest_folder``. In order to generate files in the supported format of ``nemo_asr``, run: 
+To process the AIShell-2 dataset, in the command below, set the data folder of AIShell-2 using ``--audio_folder`` and where to push
+these files using ``--dest_folder``. In order to generate files in the supported format of ``nemo_asr``, run:
 
 .. code-block:: bash
 
     python process_aishell2_data.py --audio_folder=<data directory> --dest_folder=<destination directory>
 
-After the script finishes, the ``train.json``, ``dev.json``, ``test.json``, and ``vocab.txt`` files can be found in the ``dest_folder`` directory. 
+After the script finishes, the ``train.json``, ``dev.json``, ``test.json``, and ``vocab.txt`` files can be found in the ``dest_folder`` directory.
 
 Preparing Custom ASR Data
 -------------------------
@@ -171,7 +171,7 @@ The audio files can be of any format supported by `Pydub <https://github.com/jia
 WAV files as they are the default and have been most thoroughly tested.
 
 There should be one manifest file per dataset that will be passed in, therefore, if the user wants separate training and validation
-datasets, they should also have separate manifests. Otherwise, thay will be loading validation data with their training data and vice 
+datasets, they should also have separate manifests. Otherwise, thay will be loading validation data with their training data and vice
 versa.
 
 Each line of the manifest should be in the following format:
@@ -210,16 +210,22 @@ of filepaths, e.g. ``['/data/shard1.tar', '/data/shard2.tar']``, or in a single
 ``'/data/shard_{1..64}.tar'`` or ``'/data/shard__OP_1..64_CL_'`` (recommended, see note below).
 
 .. note::
-  For brace expansion, there may be cases where ``{x..y}`` syntax cannot be used due to shell interference. This occurs most commonly 
-  inside SLURM scripts. Therefore, we provide a few equivalent replacements. Supported opening braces (equivalent to ``{``) are ``(``, 
-  ``[``, ``<`` and the special tag ``_OP_``. Supported closing braces (equivalent to ``}``) are ``)``, ``]``, ``>`` and the special 
+  For brace expansion, there may be cases where ``{x..y}`` syntax cannot be used due to shell interference. This occurs most commonly
+  inside SLURM scripts. Therefore, we provide a few equivalent replacements. Supported opening braces (equivalent to ``{``) are ``(``,
+  ``[``, ``<`` and the special tag ``_OP_``. Supported closing braces (equivalent to ``}``) are ``)``, ``]``, ``>`` and the special
   tag ``_CL_``. For SLURM based tasks, we suggest the use of the special tags for ease of use.
 
-As with non-tarred datasets, the manifest file should be passed in ``manifest_filepath``. The dataloader assumes that the length 
+As with non-tarred datasets, the manifest file should be passed in ``manifest_filepath``. The dataloader assumes that the length
 of the manifest after filtering is the correct size of the dataset for reporting training progress.
 
-The ``tarred_shard_strategy`` field of the config file can be set if you have multiple shards and are running an experiment with 
+The ``tarred_shard_strategy`` field of the config file can be set if you have multiple shards and are running an experiment with
 multiple workers. It defaults to ``scatter``, which preallocates a set of shards per worker which do not change during runtime.
+Note that this strategy, on specific occasions (when the number of shards is not divisible with ``world_size``), will not sample
+the entire dataset. As an alternative the ``replicate`` strategy, will prealocat the entire set of shards to every worker and not
+change during runtime. The benefit of this strategy is that it allows each worker to sample data points from the entire dataset
+independently of others. Note, though, that more than one worker may sample the same shard, and even sample the same data points!
+As such, there is no assured guarantee that all samples in the dataset will be sampled at least once during 1 epoch. Note that
+for these reasons it is not advisable to use tarred datasets as validation and test datasets.
 
 For more information about the individual tarred datasets and the parameters available, including shuffling options,
 see the corresponding class APIs in the `Datasets <./api.html#Datasets>`__ section.
@@ -228,7 +234,7 @@ see the corresponding class APIs in the `Datasets <./api.html#Datasets>`__ secti
   If using multiple workers, the number of shards should be divisible by the world size to ensure an even
   split among workers. If it is not divisible, logging will give a warning but training will proceed, but likely hang at the last epoch.
   In addition, if using distributed processing, each shard must have the same number of entries after filtering is
-  applied such that each worker ends up with the same number of files. We currently do not check for this in any dataloader, but the user's 
+  applied such that each worker ends up with the same number of files. We currently do not check for this in any dataloader, but the user's
   program may hang if the shards are uneven.
 
 Conversion to Tarred Datasets
@@ -262,9 +268,9 @@ The files in the target directory should look similar to the following:
   ├── metadata.yaml
   └── tarred_audio_manifest.json
 
-Note that file structures are flattened such that all audio files are at the top level in each tarball. This ensures that 
-filenames are unique in the tarred dataset and the filepaths do not contain "-sub" and forward slashes in each ``audio_filepath`` are 
-simply converted to underscores. For example, a manifest entry for ``/data/directory1/file.wav`` would be ``_data_directory1_file.wav`` 
+Note that file structures are flattened such that all audio files are at the top level in each tarball. This ensures that
+filenames are unique in the tarred dataset and the filepaths do not contain "-sub" and forward slashes in each ``audio_filepath`` are
+simply converted to underscores. For example, a manifest entry for ``/data/directory1/file.wav`` would be ``_data_directory1_file.wav``
 in the tarred dataset manifest, and ``/data/directory2/file.wav`` would be converted to ``_data_directory2_file.wav``.
 
 Bucketing Datasets
@@ -325,9 +331,9 @@ Currently bucketing feature is just supported for tarred datasets.
 Upsampling Datasets
 ------------------
 
-Buckets may also be 'weighted' to allow multiple runs through a target dataset during each training epoch. This can be beneficial in cases when a dataset is composed of several component sets of unequal sizes and one desires to mitigate bias towards the larger sets through oversampling.   
+Buckets may also be 'weighted' to allow multiple runs through a target dataset during each training epoch. This can be beneficial in cases when a dataset is composed of several component sets of unequal sizes and one desires to mitigate bias towards the larger sets through oversampling.
 
-Weighting is managed with the `bucketing_weights` parameter. After passing your composite tarred datasets in the format described above for bucketing, pass a list of integers (one per bucket) to indicate how many times a manifest should be read during training. 
+Weighting is managed with the `bucketing_weights` parameter. After passing your composite tarred datasets in the format described above for bucketing, pass a list of integers (one per bucket) to indicate how many times a manifest should be read during training.
 
 For example, by passing `[2,1,1,3]` to the code below:
 
@@ -363,7 +369,7 @@ If using adaptive bucketing, note that the same batch size will be assigned to e
 	model.train_ds.bucketing_weights=[2,1,1,3]
 	model.train_ds.bucketing_batch_size=[4,4,4,2]
 
-All instances of data from `bucket4` will still be trained with a batch size of 2 while all others would have a batch size of 4. As with standard bucketing, this requires `batch_size`` to be set to 1. 
+All instances of data from `bucket4` will still be trained with a batch size of 2 while all others would have a batch size of 4. As with standard bucketing, this requires `batch_size`` to be set to 1.
 If `bucketing_batch_size` is not specified, all datasets will be passed with the same fixed batch size as specified by the `batch_size` parameter.
 
-It is recommended to set bucketing strategies to `fully_randomized` during multi-GPU training to prevent possible dataset bias during training. 
\ No newline at end of file
+It is recommended to set bucketing strategies to `fully_randomized` during multi-GPU training to prevent possible dataset bias during training.
\ No newline at end of file
diff --git a/docs/source/nlp/punctuation_and_capitalization.rst b/docs/source/nlp/punctuation_and_capitalization.rst
index be2890943acb..93c4416568b7 100755
--- a/docs/source/nlp/punctuation_and_capitalization.rst
+++ b/docs/source/nlp/punctuation_and_capitalization.rst
@@ -3,7 +3,7 @@
 Punctuation and Capitalization Model
 ====================================
 
-Automatic Speech Recognition (ASR) systems typically generate text with no punctuation and capitalization of the words. 
+Automatic Speech Recognition (ASR) systems typically generate text with no punctuation and capitalization of the words.
 There are two issues with non-punctuated ASR output:
 
 - it could be difficult to read and understand
@@ -35,7 +35,7 @@ For each word in the input text, the Punctuation and Capitalization model:
 - predicts a punctuation mark that should follow the word (if any). By default, the model supports commas, periods, and question marks.
 - predicts if the word should be capitalized or not
 
-In the Punctuation and Capitalization model, we are jointly training two token-level classifiers on top of a pre-trained 
+In the Punctuation and Capitalization model, we are jointly training two token-level classifiers on top of a pre-trained
 language model, such as `BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding <https://arxiv.org/abs/1810.04805>`__ :cite:`nlp-punct-devlin2018bert`.
 
 .. note::
@@ -85,7 +85,7 @@ NeMo Data Format
 
 The Punctuation and Capitalization model expects the data in the following format:
 
-The training and evaluation data is divided into 2 files: 
+The training and evaluation data is divided into 2 files:
 - ``text.txt``
 - ``labels.txt``
 
@@ -108,10 +108,10 @@ spaces. Each label in ``labels.txt`` file consists of 2 symbols:
 - the second symbol determines if a word needs to be capitalized or not (where ``U`` indicates that the word should be
   upper cased, and ``O`` - no capitalization needed)
 
-By default, the following punctuation marks are considered: commas, periods, and question marks; the remaining punctuation marks were 
+By default, the following punctuation marks are considered: commas, periods, and question marks; the remaining punctuation marks were
 removed from the data. This can be changed by introducing new labels in the ``labels.txt`` files.
 
-Each line of the ``labels.txt`` should follow the format: ``[LABEL] [SPACE] [LABEL] [SPACE] [LABEL]`` (for ``labels.txt``). For example, 
+Each line of the ``labels.txt`` should follow the format: ``[LABEL] [SPACE] [LABEL] [SPACE] [LABEL]`` (for ``labels.txt``). For example,
 labels for the above ``text.txt`` file should be:
 
     ::
@@ -120,7 +120,7 @@ labels for the above ``text.txt`` file should be:
         OU OO OO OO ...
         ...
 
-The complete list of all possible labels used in this tutorial are: 
+The complete list of all possible labels used in this tutorial are:
 
 - ``OO``
 - ``.O``
@@ -588,6 +588,19 @@ For convenience, items of data config are described in 4 tables:
      - ``1``
      - The size of shuffle buffer of `webdataset <https://github.com/webdataset/webdataset>`_. The number of batches
        which are permuted.
+   * - **shard_strategy**
+     - string
+     - ``scatter``
+     - Tarred dataset shard distribution strategy chosen as a str value during ddp. Accepted values are ``scatter`` and ``replicate``.
+       ``scatter``: Each node gets a unique set of shards, which are permanently pre-allocated and never changed at runtime, when the total
+       number of shards is not divisible with ``world_size``, some shards (at max ``world_size-1``) will not be used.
+       ``replicate``: Each node gets the entire set of shards available in the tarred dataset, which are permanently pre-allocated and never
+       changed at runtime. The benefit of replication is that it allows each node to sample data points from the entire dataset independently
+       of other nodes, and reduces dependence on value of ``tar_shuffle_n``.
+       Note: Replicated strategy allows every node to sample the entire set of available tarfiles, and therefore more than one node may sample
+       the same tarfile, and even sample the same data points! As such, there is no assured guarantee that all samples in the dataset will be
+       sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific occasions (when the number of shards is not
+       divisible with ``world_size``), will not sample the entire dataset.
 
 .. _pytorch-dataloader-parameters-label:
 
diff --git a/nemo/collections/asr/data/audio_to_label.py b/nemo/collections/asr/data/audio_to_label.py
index d002c0c76717..c84511ea56e2 100644
--- a/nemo/collections/asr/data/audio_to_label.py
+++ b/nemo/collections/asr/data/audio_to_label.py
@@ -700,7 +700,9 @@ class TarredAudioToClassificationLabelDataset(_TarredAudioLabelDataset):
                 Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
                 and therefore more than one node may sample the same tarfile, and even sample the same
                 data points! As such, there is no assured guarantee that all samples in the dataset will be
-                sampled at least once during 1 epoch.
+                sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific
+                occasions (when the number of shards is not divisible with ``world_size``), will not sample
+                the entire dataset.
         global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
         world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
         is_regression_task (bool): Whether it is a regression task. Defualts to False.
@@ -774,7 +776,9 @@ class TarredAudioToSpeechLabelDataset(_TarredAudioLabelDataset):
                 Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
                 and therefore more than one node may sample the same tarfile, and even sample the same
                 data points! As such, there is no assured guarantee that all samples in the dataset will be
-                sampled at least once during 1 epoch.
+                sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific
+                occasions (when the number of shards is not divisible with ``world_size``), will not sample
+                the entire dataset.
         global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
         world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
     """
diff --git a/nemo/collections/asr/data/audio_to_text.py b/nemo/collections/asr/data/audio_to_text.py
index 573dfd672fa2..7a735c2f4f73 100644
--- a/nemo/collections/asr/data/audio_to_text.py
+++ b/nemo/collections/asr/data/audio_to_text.py
@@ -582,7 +582,9 @@ class _TarredAudioToTextDataset(IterableDataset):
                 Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
                 and therefore more than one node may sample the same tarfile, and even sample the same
                 data points! As such, there is no assured guarantee that all samples in the dataset will be
-                sampled at least once during 1 epoch.
+                sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific
+                occasions (when the number of shards is not divisible with ``world_size``), will not sample
+                the entire dataset.
         global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
         world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
         return_sample_id (bool): whether to return the sample_id as a part of each sample
@@ -843,7 +845,9 @@ class TarredAudioToCharDataset(_TarredAudioToTextDataset):
                 Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
                 and therefore more than one node may sample the same tarfile, and even sample the same
                 data points! As such, there is no assured guarantee that all samples in the dataset will be
-                sampled at least once during 1 epoch.
+                sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific
+                occasions (when the number of shards is not divisible with ``world_size``), will not sample
+                the entire dataset.
         global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
         world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
         return_sample_id (bool): whether to return the sample_id as a part of each sample
@@ -970,7 +974,9 @@ class TarredAudioToBPEDataset(_TarredAudioToTextDataset):
                 Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
                 and therefore more than one node may sample the same tarfile, and even sample the same
                 data points! As such, there is no assured guarantee that all samples in the dataset will be
-                sampled at least once during 1 epoch.
+                sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific
+                occasions (when the number of shards is not divisible with ``world_size``), will not sample
+                the entire dataset.
         global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
         world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
         return_sample_id (bool): whether to return the sample_id as a part of each sample
diff --git a/nemo/collections/nlp/data/language_modeling/l2r_lm_dataset.py b/nemo/collections/nlp/data/language_modeling/l2r_lm_dataset.py
index eff0737c7fec..39f92cf8ca14 100644
--- a/nemo/collections/nlp/data/language_modeling/l2r_lm_dataset.py
+++ b/nemo/collections/nlp/data/language_modeling/l2r_lm_dataset.py
@@ -31,7 +31,7 @@
 class L2RLanguageModelingDataset(Dataset):
     """
     Dataset for training and evaluating left-to-right language models.
-    
+
     Args:
         tokenizer: tokenizer, such as WordTokenizer or CharTokenizer
         dataset: path to data
@@ -73,7 +73,7 @@ def __getitem__(self, idx):
 class TarredL2RLanguageModelingDataset(IterableDataset):
     """
     A similar Dataset to the L2RLanguageModelingDataset, but which loads tarred tokenized numpy files.
-    Accepts a single JSON metadata manifest file as well as the path(s) to the tarball(s) containing the wav files. 
+    Accepts a single JSON metadata manifest file as well as the path(s) to the tarball(s) containing the wav files.
     The manifest should contain information such as the number of shards, the number of tokens in the corpus,
     and the number of tokens contained within each shard of the tarfile(s).
 
@@ -114,10 +114,13 @@ class TarredL2RLanguageModelingDataset(IterableDataset):
                 available in the tarred dataset, which are permanently pre-allocated and never changed at runtime.
                 The benefit of replication is that it allows each node to sample data points from the entire
                 dataset independently of other nodes, and reduces dependence on value of `shuffle_n`.
+
                 Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
                 and therefore more than one node may sample the same tarfile, and even sample the same
                 data points! As such, there is no assured guarantee that all samples in the dataset will be
-                sampled at least once during 1 epoch.
+                sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific
+                occasions (when the number of shards is not divisible with ``world_size``), will not sample
+                the entire dataset.
         global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
         world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
     """
@@ -142,7 +145,11 @@ def __init__(
 
         valid_shard_strategies = ['scatter', 'replicate']
         if shard_strategy not in valid_shard_strategies:
-            raise ValueError(f"`shard_strategy` must be one of {valid_shard_strategies}")
+            raise ValueError(
+                f"Invalid shard strategy of type {type(shard_strategy)} "
+                f"{repr(shard_strategy) if len(repr(shard_strategy)) < 100 else repr(shard_strategy)[:100] + '...'}! "
+                f"Allowed values are: {valid_shard_strategies}."
+            )
 
         with open(metadata_path, 'r') as f:
             metadata = json.load(f)
diff --git a/nemo/collections/nlp/data/language_modeling/sentence_dataset.py b/nemo/collections/nlp/data/language_modeling/sentence_dataset.py
index e677fbade383..d4bee937058a 100644
--- a/nemo/collections/nlp/data/language_modeling/sentence_dataset.py
+++ b/nemo/collections/nlp/data/language_modeling/sentence_dataset.py
@@ -142,7 +142,7 @@ class TarredSentenceDataset(IterableDataset):
     """
     A similar Dataset to the SentenceDataset, but which loads tarred tokenized pickle files.
     Accepts a single JSON metadata file containing the total number of batches
-    as well as the path(s) to the tarball(s) containing the wav files. 
+    as well as the path(s) to the tarball(s) containing the wav files.
     Valid formats for the text_tar_filepaths argument include:
     (1) a single string that can be brace-expanded, e.g. 'path/to/text.tar' or 'path/to/text_{1..100}.tar.gz', or
     (2) a list of file paths that will not be brace-expanded, e.g. ['text_1.tar', 'text_2.tar', ...].
@@ -172,10 +172,13 @@ class TarredSentenceDataset(IterableDataset):
                 available in the tarred dataset, which are permanently pre-allocated and never changed at runtime.
                 The benefit of replication is that it allows each node to sample data points from the entire
                 dataset independently of other nodes, and reduces dependence on value of `shuffle_n`.
+
                 Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
                 and therefore more than one node may sample the same tarfile, and even sample the same
                 data points! As such, there is no assured guarantee that all samples in the dataset will be
-                sampled at least once during 1 epoch.
+                sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific
+                occasions (when the number of shards is not divisible with ``world_size``), will not sample
+                the entire dataset.
         global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
         world_size (int): Total number of processes, used for partitioning shards. Defaults to 0.
         reverse_lang_direction (bool): When True, swaps the source and target directions when returning minibatches.
@@ -198,7 +201,11 @@ def __init__(
 
         valid_shard_strategies = ['scatter', 'replicate']
         if shard_strategy not in valid_shard_strategies:
-            raise ValueError(f"`shard_strategy` must be one of {valid_shard_strategies}")
+            raise ValueError(
+                f"Invalid shard strategy of type {type(shard_strategy)} "
+                f"{repr(shard_strategy) if len(repr(shard_strategy)) < 100 else repr(shard_strategy)[:100] + '...'}! "
+                f"Allowed values are: {valid_shard_strategies}."
+            )
 
         with open(metadata_path, 'r') as f:
             metadata = json.load(f)
@@ -223,12 +230,14 @@ def __init__(
             text_tar_filepaths = list(braceexpand.braceexpand(text_tar_filepaths))
 
         if shard_strategy == 'scatter':
-            logging.info("All tarred dataset shards will be scattered evenly across all nodes.")
+            logging.info("Tarred dataset shards will be scattered evenly across all nodes.")
             if len(text_tar_filepaths) % world_size != 0:
                 logging.warning(
                     f"Number of shards in tarred dataset ({len(text_tar_filepaths)}) is not divisible "
-                    f"by number of distributed workers ({world_size})."
+                    f"by number of distributed workers ({world_size}). "
+                    f"Some shards will not be used ({len(text_tar_filepaths) % world_size})."
                 )
+            batches_per_tar = self.metadata['num_batches'] // len(text_tar_filepaths)
             begin_idx = (len(text_tar_filepaths) // world_size) * global_rank
             end_idx = begin_idx + (len(text_tar_filepaths) // world_size)
             logging.info('Begin Index : %d' % (begin_idx))
@@ -237,7 +246,7 @@ def __init__(
             logging.info(
                 "Partitioning tarred dataset: process (%d) taking shards [%d, %d)", global_rank, begin_idx, end_idx
             )
-            self.length = self.metadata['num_batches'] // world_size
+            self.length = batches_per_tar * len(text_tar_filepaths) * world_size
 
         elif shard_strategy == 'replicate':
             logging.info("All tarred dataset shards will be replicated across all nodes.")
diff --git a/nemo/collections/nlp/data/machine_translation/machine_translation_dataset.py b/nemo/collections/nlp/data/machine_translation/machine_translation_dataset.py
index fb253df1c098..4ad956d1aa0a 100644
--- a/nemo/collections/nlp/data/machine_translation/machine_translation_dataset.py
+++ b/nemo/collections/nlp/data/machine_translation/machine_translation_dataset.py
@@ -326,10 +326,13 @@ class TarredTranslationDataset(IterableDataset):
                 available in the tarred dataset, which are permanently pre-allocated and never changed at runtime.
                 The benefit of replication is that it allows each node to sample data points from the entire
                 dataset independently of other nodes, and reduces dependence on value of `shuffle_n`.
+
                 Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
                 and therefore more than one node may sample the same tarfile, and even sample the same
                 data points! As such, there is no assured guarantee that all samples in the dataset will be
-                sampled at least once during 1 epoch.
+                sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific
+                occasions (when the number of shards is not divisible with ``world_size``), will not sample
+                the entire dataset.
         global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
         world_size (int): Total number of processes, used for partitioning shards. Defaults to 1.
         reverse_lang_direction (bool): When True, swaps the source and target directions when returning minibatches.
@@ -360,7 +363,11 @@ def __init__(
 
         valid_shard_strategies = ['scatter', 'replicate']
         if shard_strategy not in valid_shard_strategies:
-            raise ValueError(f"`shard_strategy` must be one of {valid_shard_strategies}")
+            raise ValueError(
+                f"Invalid shard strategy of type {type(shard_strategy)} "
+                f"{repr(shard_strategy) if len(repr(shard_strategy)) < 100 else repr(shard_strategy)[:100] + '...'}! "
+                f"Allowed values are: {valid_shard_strategies}."
+            )
 
         with open(metadata_path, 'r') as f:
             metadata = json.load(f)
@@ -385,12 +392,14 @@ def __init__(
             text_tar_filepaths = list(braceexpand.braceexpand(text_tar_filepaths))
 
         if shard_strategy == 'scatter':
-            logging.info("All tarred dataset shards will be scattered evenly across all nodes.")
+            logging.info("Tarred dataset shards will be scattered evenly across all nodes.")
             if len(text_tar_filepaths) % world_size != 0:
                 logging.warning(
                     f"Number of shards in tarred dataset ({len(text_tar_filepaths)}) is not divisible "
-                    f"by number of distributed workers ({world_size})."
+                    f"by number of distributed workers ({world_size}). "
+                    f"Some shards will not be used ({len(text_tar_filepaths) % world_size})."
                 )
+            batches_per_tar = self.metadata['num_batches'] // len(text_tar_filepaths)
             begin_idx = (len(text_tar_filepaths) // world_size) * global_rank
             end_idx = begin_idx + (len(text_tar_filepaths) // world_size)
             logging.info('Begin Index : %d' % (begin_idx))
@@ -399,7 +408,7 @@ def __init__(
             logging.info(
                 "Partitioning tarred dataset: process (%d) taking shards [%d, %d)", global_rank, begin_idx, end_idx
             )
-            self.length = self.metadata['num_batches'] // world_size
+            self.length = batches_per_tar * len(text_tar_filepaths) * world_size
 
         elif shard_strategy == 'replicate':
             logging.info("All tarred dataset shards will be replicated across all nodes.")
diff --git a/nemo/collections/nlp/data/text_normalization/decoder_dataset.py b/nemo/collections/nlp/data/text_normalization/decoder_dataset.py
index 1110807dcb0b..d947fa31a185 100644
--- a/nemo/collections/nlp/data/text_normalization/decoder_dataset.py
+++ b/nemo/collections/nlp/data/text_normalization/decoder_dataset.py
@@ -452,10 +452,13 @@ class TarredTextNormalizationDecoderDataset(IterableDataset):
                 available in the tarred dataset, which are permanently pre-allocated and never changed at runtime.
                 The benefit of replication is that it allows each node to sample data points from the entire
                 dataset independently of other nodes, and reduces dependence on value of `shuffle_n`.
+
                 Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
                 and therefore more than one node may sample the same tarfile, and even sample the same
                 data points! As such, there is no assured guarantee that all samples in the dataset will be
-                sampled at least once during 1 epoch.
+                sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific
+                occasions (when the number of shards is not divisible with ``world_size``), will not sample
+                the entire dataset.
         global_rank: Worker rank, used for partitioning shards.
         world_size: Total number of processes, used for partitioning shards.
     """
@@ -473,7 +476,11 @@ def __init__(
 
         valid_shard_strategies = ['scatter', 'replicate']
         if shard_strategy not in valid_shard_strategies:
-            raise ValueError(f"`shard_strategy` must be one of {valid_shard_strategies}")
+            raise ValueError(
+                f"Invalid shard strategy of type {type(shard_strategy)} "
+                f"{repr(shard_strategy) if len(repr(shard_strategy)) < 100 else repr(shard_strategy)[:100] + '...'}! "
+                f"Allowed values are: {valid_shard_strategies}."
+            )
 
         if isinstance(text_tar_filepaths, str):
             # Replace '(', '[', '<' and '_OP_' with '{'
@@ -493,12 +500,14 @@ def __init__(
             text_tar_filepaths = list(braceexpand.braceexpand(text_tar_filepaths))
 
         if shard_strategy == 'scatter':
-            logging.info("All tarred dataset shards will be scattered evenly across all nodes.")
+            logging.info("Tarred dataset shards will be scattered evenly across all nodes.")
             if len(text_tar_filepaths) % world_size != 0:
                 logging.warning(
                     f"Number of shards in tarred dataset ({len(text_tar_filepaths)}) is not divisible "
-                    f"by number of distributed workers ({world_size})."
+                    f"by number of distributed workers ({world_size}). "
+                    f"Some shards will not be used ({len(text_tar_filepaths) % world_size})."
                 )
+            batches_per_tar = num_batches // len(text_tar_filepaths)
             begin_idx = (len(text_tar_filepaths) // world_size) * global_rank
             end_idx = begin_idx + (len(text_tar_filepaths) // world_size)
             logging.info('Begin Index : %d' % (begin_idx))
@@ -507,16 +516,17 @@ def __init__(
             logging.info(
                 "Partitioning tarred dataset: process (%d) taking shards [%d, %d)", global_rank, begin_idx, end_idx
             )
+            self.length = batches_per_tar * len(text_tar_filepaths) * world_size
 
         elif shard_strategy == 'replicate':
             logging.info("All tarred dataset shards will be replicated across all nodes.")
+            self.length = num_batches
 
         else:
             raise ValueError(f"Invalid shard strategy! Allowed values are: {valid_shard_strategies}")
 
         # Put together WebDataset
         self._dataset = wd.WebDataset(urls=text_tar_filepaths, nodesplitter=None)
-        self.length = num_batches // world_size
         if shuffle_n > 0:
             self._dataset = self._dataset.shuffle(shuffle_n)
         else:
diff --git a/nemo/collections/nlp/data/token_classification/punctuation_capitalization_dataset.py b/nemo/collections/nlp/data/token_classification/punctuation_capitalization_dataset.py
index 493624de3f65..4e0f7b0d1672 100644
--- a/nemo/collections/nlp/data/token_classification/punctuation_capitalization_dataset.py
+++ b/nemo/collections/nlp/data/token_classification/punctuation_capitalization_dataset.py
@@ -85,7 +85,7 @@ class PunctuationCapitalizationDataConfigBase:
 
     labels_file: Optional[str] = None
     """A path to a file with punctuation and capitalization labels in NeMo format. NeMo format is described in
-    `documentation 
+    `documentation
     <https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/punctuation_and_capitalization.html#nemo-data-format>`_
     """
 
@@ -112,7 +112,7 @@ class PunctuationCapitalizationDataConfigBase:
     """A path to a directory containing cache or directory where newly created cache is saved. By default, it is
     a directory containing ``text_file``. You may need this parameter if cache for a dataset is going to be created
     and the dataset directory is read-only.
-    
+
     ``cache_dir`` and ``label_info_save_dir`` are separate parameters for the case when a cache is ready and this cache
     is stored in a read only directory. In this case you will separate ``label_info_save_dir``."""
 
@@ -142,6 +142,18 @@ class PunctuationCapitalizationDataConfigBase:
     tar_shuffle_n: int = 1
     """The size of shuffle buffer of `webdataset`. The number of batches which are permuted."""
 
+    shard_strategy: Optional[str] = 'scatter'
+    """Tarred dataset shard distribution strategy chosen as a str value during ddp. Accepted values are `scatter` and `replicate`.
+    `scatter`: The default shard strategy applied by WebDataset, where each node gets a unique set of shards, which are permanently
+    pre-allocated and never changed at runtime. `replicate` is an optional shard strategy, where each node gets the entire set of shards
+    available in the tarred dataset, which are permanently pre-allocated and never changed at runtime. The benefit of replication is that
+    it allows each node to sample data points from the entire dataset independently of other nodes, and reduces dependence on value of
+    ``tar_shuffle_n``. Note: Replicated strategy allows every node to sample the entire set of available tarfiles, and therefore more than
+    one node may sample the same tarfile, and even sample the same data points! As such, there is no assured guarantee that all samples
+    in the dataset will be sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific occasions (when the
+    number of shards is not divisible with ``world_size``), will not sample the entire dataset.
+    """
+
     #################################################
     # PYTORCH DATALOADER PARAMETERS
     #################################################
diff --git a/nemo/collections/nlp/data/token_classification/punctuation_capitalization_tarred_dataset.py b/nemo/collections/nlp/data/token_classification/punctuation_capitalization_tarred_dataset.py
index 2bfcb7969b6e..fa3ee946d8ec 100644
--- a/nemo/collections/nlp/data/token_classification/punctuation_capitalization_tarred_dataset.py
+++ b/nemo/collections/nlp/data/token_classification/punctuation_capitalization_tarred_dataset.py
@@ -536,7 +536,7 @@ def repack_tar_files_with_not_enough_batches(output_dir: Path, num_batches_per_t
     ``repack_tar_files_with_not_enough_batches`` function into tar files with correct ``num_batches_per_tarfile``
     batches each. If there is no enough batches in repacked files, then up to ``num_batches_per_tarfile - 1``
     remaining batches may be discarded.
-    
+
     Args:
         output_dir: a path to the output directory which contains files to repack and where new files are saved
         num_batches_per_tarfile: a number of batches in 1 tar file. If number of batches in files matching a pattern
@@ -685,10 +685,10 @@ def create_tarred_dataset(
     `examples/nlp/token_classification/data/create_punctuation_capitalization_tarred_dataset.py
     <https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/token_classification/data/create_punctuation_capitalization_tarred_dataset.py>`_.
 
-    Tarred dataset is a directory which contains metadata file, tar files with batches, 
+    Tarred dataset is a directory which contains metadata file, tar files with batches,
     ``punct_label_vocab.csv`` and ``capit_label_vocab.csv`` files.
 
-    Metadata file is a JSON file with 4 items: ``'num_batches'``, ``'tar_files'``, ``'punct_label_vocab_file'``, 
+    Metadata file is a JSON file with 4 items: ``'num_batches'``, ``'tar_files'``, ``'punct_label_vocab_file'``,
     ``'capit_label_vocab_file'``. The item ``'num_batches'`` (``int``) is a total number of batches in tarred dataset.
     ``'tar_files'`` is a list of paths to tar files relative to directory containing the metadata file. The items
     ``'punct_label_vocab_file'`` and ``'capit_label_vocab_file'`` are correspondingly paths to punctuation and
@@ -871,6 +871,21 @@ class BertPunctuationCapitalizationTarredDataset(IterableDataset):
             be used in the current process.
         shuffle_n (:obj:`int`, `optional`, defaults to :obj:`1`): a number of shuffled batches in a buffer.
             ``shuffle_n`` batches are loaded into memory, shuffled, and then yielded by a dataset instance.
+        shard_strategy (:obj:`str`, defaults to :obj:``'scatter'``): Tarred dataset shard distribution strategy chosen as
+            a str value during ddp.
+            -   ``'scatter'``: The default shard strategy applied by WebDataset, where each node gets
+                a unique set of shards, which are permanently pre-allocated and never changed at runtime.
+            -   ``'replicate'``: Optional shard strategy, where each node gets all of the set of shards
+                available in the tarred dataset, which are permanently pre-allocated and never changed at runtime.
+                The benefit of replication is that it allows each node to sample data points from the entire
+                dataset independently of other nodes, and reduces dependence on value of :param:`shuffle_n`.
+
+                Note: Replicated strategy allows every node to sample the entire set of available tarfiles,
+                and therefore more than one node may sample the same tarfile, and even sample the same
+                data points! As such, there is no assured guarantee that all samples in the dataset will be
+                sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific
+                occasions (when the number of shards is not divisible with ``world_size``), will not sample
+                the entire dataset.
     """
 
     @property
@@ -897,8 +912,18 @@ def __init__(
         world_size: int = 1,
         global_rank: int = 0,
         shuffle_n: int = 1,
+        shard_strategy: str = "scatter",
     ) -> None:
         super().__init__()
+
+        valid_shard_strategies = ['scatter', 'replicate']
+        if shard_strategy not in valid_shard_strategies:
+            raise ValueError(
+                f"Invalid shard strategy of type {type(shard_strategy)} "
+                f"{repr(shard_strategy) if len(repr(shard_strategy)) < 100 else repr(shard_strategy)[:100] + '...'}! "
+                f"Allowed values are: {valid_shard_strategies}."
+            )
+
         self.tokenizer = tokenizer
         self.metadata_file = Path(metadata_file).expanduser()
         if label_info_save_dir is None:
@@ -922,13 +947,31 @@ def __init__(
         self.capit_label_ids = load_label_ids(self.capit_label_vocab_file)
         self.pad_label = pad_label
         self._check_pad_label()
-        begin_idx = (len(self.tar_files) // world_size) * global_rank
-        end_idx = begin_idx + (len(self.tar_files) // world_size)
-        logging.info(
-            "Partitioning tarred dataset: process (%d) taking shards [%d, %d)", global_rank, begin_idx, end_idx
-        )
-        self.tar_files = self.tar_files[begin_idx:end_idx]
-        self.length = self.metadata['num_batches'] // world_size
+
+        if shard_strategy == 'scatter':
+            logging.info("Tarred dataset shards will be scattered evenly across all nodes.")
+            if len(self.tar_files) % world_size != 0:
+                logging.warning(
+                    f"Number of shards in tarred dataset ({len(self.tar_files)}) is not divisible "
+                    f"by number of distributed workers ({world_size}). "
+                    f"Some shards will not be used ({len(self.tar_files) % world_size})."
+                )
+            begin_idx = (len(self.tar_files) // world_size) * global_rank
+            end_idx = begin_idx + (len(self.tar_files) // world_size)
+            logging.info(
+                "Partitioning tarred dataset: process (%d) taking shards [%d, %d)", global_rank, begin_idx, end_idx
+            )
+            batches_per_tar = self.metadata['num_batches'] // len(self.tar_files)
+            self.tar_files = self.tar_files[begin_idx:end_idx]
+            self.length = batches_per_tar * len(self.tar_files) * world_size
+
+        elif shard_strategy == 'replicate':
+            logging.info("All tarred dataset shards will be replicated across all nodes.")
+            self.length = self.metadata['num_batches']
+
+        else:
+            raise ValueError(f"Invalid shard strategy! Allowed values are: {valid_shard_strategies}")
+
         self._dataset = wds.WebDataset(urls=self.tar_files, nodesplitter=None).decode(
             wds.handle_extension('.pyd', decode_pyd)
         )
diff --git a/nemo/collections/nlp/models/duplex_text_normalization/duplex_decoder.py b/nemo/collections/nlp/models/duplex_text_normalization/duplex_decoder.py
index 4f602f90da8b..440640231664 100644
--- a/nemo/collections/nlp/models/duplex_text_normalization/duplex_decoder.py
+++ b/nemo/collections/nlp/models/duplex_text_normalization/duplex_decoder.py
@@ -15,6 +15,7 @@
 import json
 import os
 from collections import defaultdict
+from math import ceil
 from typing import Dict, List, Optional, Union
 
 import torch
@@ -398,6 +399,23 @@ def setup_training_data(self, train_data_config: Optional[DictConfig]):
             cfg=train_data_config, data_split="train"
         )
 
+        # Need to set this because if using an IterableDataset, the length of the dataloader is the total number
+        # of samples rather than the number of batches, and this messes up the tqdm progress bar.
+        # So we set the number of steps manually (to the correct number) to fix this.
+        if 'use_tarred_dataset' in train_data_config and train_data_config['use_tarred_dataset']:
+            # We also need to check if limit_train_batches is already set.
+            # If it's an int, we assume that the user has set it to something sane, i.e. <= # training batches,
+            # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0).
+            if self._trainer is not None and isinstance(self._trainer.limit_train_batches, float):
+                self._trainer.limit_train_batches = int(
+                    self._trainer.limit_train_batches * ceil(len(self._train_dl.dataset) / self.world_size)
+                )
+            elif self._trainer is None:
+                logging.warning(
+                    "Model Trainer was not set before constructing the dataset, incorrect number of "
+                    "training batches will be used. Please set the trainer and rebuild the dataset."
+                )
+
     def setup_validation_data(self, val_data_config: Optional[DictConfig]):
         if not val_data_config or not val_data_config.data_path:
             logging.info(
@@ -409,6 +427,23 @@ def setup_validation_data(self, val_data_config: Optional[DictConfig]):
             cfg=val_data_config, data_split="val"
         )
 
+        # Need to set this because if using an IterableDataset, the length of the dataloader is the total number
+        # of samples rather than the number of batches, and this messes up the tqdm progress bar.
+        # So we set the number of steps manually (to the correct number) to fix this.
+        if 'use_tarred_dataset' in val_data_config and val_data_config['use_tarred_dataset']:
+            # We also need to check if limit_val_batches is already set.
+            # If it's an int, we assume that the user has set it to something sane, i.e. <= # validation batches,
+            # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0).
+            if self._trainer is not None and isinstance(self._trainer.limit_val_batches, float):
+                self._trainer.limit_val_batches = int(
+                    self._trainer.limit_val_batches * ceil(len(self._validation_dl.dataset) / self.world_size)
+                )
+            elif self._trainer is None:
+                logging.warning(
+                    "Model Trainer was not set before constructing the dataset, incorrect number of "
+                    "validation batches will be used. Please set the trainer and rebuild the dataset."
+                )
+
     def setup_multiple_validation_data(self, val_data_config: Union[DictConfig, Dict] = None):
         if val_data_config is None:
             val_data_config = self._cfg.validation_ds
diff --git a/nemo/collections/nlp/models/language_modeling/transformer_lm_model.py b/nemo/collections/nlp/models/language_modeling/transformer_lm_model.py
index e7c13e539dd0..180fdad9ddc3 100644
--- a/nemo/collections/nlp/models/language_modeling/transformer_lm_model.py
+++ b/nemo/collections/nlp/models/language_modeling/transformer_lm_model.py
@@ -211,9 +211,43 @@ def setup_tokenizer(
     def setup_training_data(self, train_data_config: Optional[DictConfig]):
         self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config)
 
+        # Need to set this because if using an IterableDataset, the length of the dataloader is the total number
+        # of samples rather than the number of batches, and this messes up the tqdm progress bar.
+        # So we set the number of steps manually (to the correct number) to fix this.
+        if 'use_tarred_dataset' in train_data_config and train_data_config['use_tarred_dataset']:
+            # We also need to check if limit_train_batches is already set.
+            # If it's an int, we assume that the user has set it to something sane, i.e. <= # training batches,
+            # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0).
+            if self._trainer is not None and isinstance(self._trainer.limit_train_batches, float):
+                self._trainer.limit_train_batches = int(
+                    self._trainer.limit_train_batches * math.ceil(len(self._train_dl.dataset) / self.world_size)
+                )
+            elif self._trainer is None:
+                logging.warning(
+                    "Model Trainer was not set before constructing the dataset, incorrect number of "
+                    "training batches will be used. Please set the trainer and rebuild the dataset."
+                )
+
     def setup_validation_data(self, val_data_config: Optional[DictConfig]):
         self._validation_dl = self._setup_dataloader_from_config(cfg=val_data_config)
 
+        # Need to set this because if using an IterableDataset, the length of the dataloader is the total number
+        # of samples rather than the number of batches, and this messes up the tqdm progress bar.
+        # So we set the number of steps manually (to the correct number) to fix this.
+        if 'use_tarred_dataset' in val_data_config and val_data_config['use_tarred_dataset']:
+            # We also need to check if limit_val_batches is already set.
+            # If it's an int, we assume that the user has set it to something sane, i.e. <= # validation batches,
+            # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0).
+            if self._trainer is not None and isinstance(self._trainer.limit_val_batches, float):
+                self._trainer.limit_val_batches = int(
+                    self._trainer.limit_val_batches * math.ceil(len(self._validation_dl.dataset) / self.world_size)
+                )
+            elif self._trainer is None:
+                logging.warning(
+                    "Model Trainer was not set before constructing the dataset, incorrect number of "
+                    "validation batches will be used. Please set the trainer and rebuild the dataset."
+                )
+
     def setup_test_data(self, test_data_config: Optional[DictConfig]):
         self._test_dl = self._setup_dataloader_from_config(cfg=test_data_config)
 
diff --git a/nemo/collections/nlp/models/machine_translation/mt_enc_dec_model.py b/nemo/collections/nlp/models/machine_translation/mt_enc_dec_model.py
index 7c27aeb1bb20..ebdcdecdab2b 100644
--- a/nemo/collections/nlp/models/machine_translation/mt_enc_dec_model.py
+++ b/nemo/collections/nlp/models/machine_translation/mt_enc_dec_model.py
@@ -16,6 +16,7 @@
 import json
 import os
 import random
+from math import ceil
 from pathlib import Path
 from typing import Dict, List, Optional, Union
 
@@ -609,6 +610,23 @@ def setup_training_data(self, train_data_config: Optional[DictConfig]):
         )
         self._train_dl = MTEncDecModel._setup_dataloader_from_config(cfg=train_data_config, dataset=self._train_ds,)
 
+        # Need to set this because if using an IterableDataset, the length of the dataloader is the total number
+        # of samples rather than the number of batches, and this messes up the tqdm progress bar.
+        # So we set the number of steps manually (to the correct number) to fix this.
+        if 'use_tarred_dataset' in train_data_config and train_data_config['use_tarred_dataset']:
+            # We also need to check if limit_train_batches is already set.
+            # If it's an int, we assume that the user has set it to something sane, i.e. <= # training batches,
+            # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0).
+            if self._trainer is not None and isinstance(self._trainer.limit_train_batches, float):
+                self._trainer.limit_train_batches = int(
+                    self._trainer.limit_train_batches * ceil(len(self._train_dl.dataset) / self.world_size)
+                )
+            elif self._trainer is None:
+                logging.warning(
+                    "Model Trainer was not set before constructing the dataset, incorrect number of "
+                    "training batches will be used. Please set the trainer and rebuild the dataset."
+                )
+
     def setup_multiple_validation_data(self, val_data_config: Union[DictConfig, Dict]):
         self.setup_validation_data(val_data_config)
 
@@ -626,6 +644,24 @@ def setup_validation_data(self, val_data_config: Optional[DictConfig]):
         self._validation_dl = MTEncDecModel._setup_eval_dataloader_from_config(
             cfg=val_data_config, datasets=self._validation_ds
         )
+
+        # Need to set this because if using an IterableDataset, the length of the dataloader is the total number
+        # of samples rather than the number of batches, and this messes up the tqdm progress bar.
+        # So we set the number of steps manually (to the correct number) to fix this.
+        if 'use_tarred_dataset' in val_data_config and val_data_config['use_tarred_dataset']:
+            # We also need to check if limit_val_batches is already set.
+            # If it's an int, we assume that the user has set it to something sane, i.e. <= # validation batches,
+            # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0).
+            if self._trainer is not None and isinstance(self._trainer.limit_val_batches, float):
+                self._trainer.limit_val_batches = int(
+                    self._trainer.limit_val_batches * ceil(len(self._validation_dl.dataset) / self.world_size)
+                )
+            elif self._trainer is None:
+                logging.warning(
+                    "Model Trainer was not set before constructing the dataset, incorrect number of "
+                    "validation batches will be used. Please set the trainer and rebuild the dataset."
+                )
+
         # instantiate Torchmetric for each val dataloader
         if self._validation_dl is not None:
             for dataloader_idx in range(len(self._validation_dl)):
diff --git a/nemo/collections/nlp/models/token_classification/punctuation_capitalization_model.py b/nemo/collections/nlp/models/token_classification/punctuation_capitalization_model.py
index 5f6fa7f6164f..d8d07fcee87d 100644
--- a/nemo/collections/nlp/models/token_classification/punctuation_capitalization_model.py
+++ b/nemo/collections/nlp/models/token_classification/punctuation_capitalization_model.py
@@ -468,6 +468,24 @@ def setup_training_data(self, train_data_config: Optional[Union[Dict[str, Any],
             train_data_config = self._cfg.train_ds
 
         self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config, train=True)
+
+        # Need to set this because if using an IterableDataset, the length of the dataloader is the total number
+        # of samples rather than the number of batches, and this messes up the tqdm progress bar.
+        # So we set the number of steps manually (to the correct number) to fix this.
+        if 'use_tarred_dataset' in train_data_config and train_data_config['use_tarred_dataset']:
+            # We also need to check if limit_train_batches is already set.
+            # If it's an int, we assume that the user has set it to something sane, i.e. <= # training batches,
+            # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0).
+            if self._trainer is not None and isinstance(self._trainer.limit_train_batches, float):
+                self._trainer.limit_train_batches = int(
+                    self._trainer.limit_train_batches * ceil(len(self._train_dl.dataset) / self.world_size)
+                )
+            elif self._trainer is None:
+                logging.warning(
+                    "Model Trainer was not set before constructing the dataset, incorrect number of "
+                    "training batches will be used. Please set the trainer and rebuild the dataset."
+                )
+
         self.punct_label_ids = self._train_dl.dataset.punct_label_ids.copy()
         self.capit_label_ids = self._train_dl.dataset.capit_label_ids.copy()
         self.label_ids_are_set = True
@@ -540,6 +558,24 @@ def setup_validation_data(self, val_data_config: Optional[Union[Dict[str, Any],
             val_data_config = self._cfg.validation_ds
 
         self._validation_dl = self._setup_dataloader_from_config(cfg=val_data_config, train=False)
+
+        # Need to set this because if using an IterableDataset, the length of the dataloader is the total number
+        # of samples rather than the number of batches, and this messes up the tqdm progress bar.
+        # So we set the number of steps manually (to the correct number) to fix this.
+        if 'use_tarred_dataset' in val_data_config and val_data_config['use_tarred_dataset']:
+            # We also need to check if limit_val_batches is already set.
+            # If it's an int, we assume that the user has set it to something sane, i.e. <= # validation batches,
+            # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0).
+            if self._trainer is not None and isinstance(self._trainer.limit_val_batches, float):
+                self._trainer.limit_val_batches = int(
+                    self._trainer.limit_val_batches * ceil(len(self._validation_dl.dataset) / self.world_size)
+                )
+            elif self._trainer is None:
+                logging.warning(
+                    "Model Trainer was not set before constructing the dataset, incorrect number of "
+                    "validation batches will be used. Please set the trainer and rebuild the dataset."
+                )
+
         loss_kw, punct_kw, capit_kw = self._get_eval_metrics_kwargs()
         self.metrics['val']['loss'].append(GlobalAverageLossMetric(**loss_kw))
         self.metrics['val']['punct_class_report'].append(ClassificationReport(**punct_kw))
@@ -741,6 +777,7 @@ def _setup_dataloader_from_config(self, cfg: DictConfig, train: bool) -> torch.u
                 world_size=self.world_size,
                 global_rank=self.global_rank,
                 shuffle_n=cfg.tar_shuffle_n,
+                shard_strategy=cfg.shard_strategy,
                 label_info_save_dir=cfg.label_info_save_dir,
             )
             dataset.check_for_label_consistency_with_model_config(


Code review comments:
diff_hunk: @@ -210,16 +210,22 @@ of filepaths, e.g. ``['/data/shard1.tar', '/data/shard2.tar']``, or in a single
 ``'/data/shard_{1..64}.tar'`` or ``'/data/shard__OP_1..64_CL_'`` (recommended, see note below).
 
 .. note::
-  For brace expansion, there may be cases where ``{x..y}`` syntax cannot be used due to shell interference. This occurs most commonly 
-  inside SLURM scripts. Therefore, we provide a few equivalent replacements. Supported opening braces (equivalent to ``{``) are ``(``, 
-  ``[``, ``<`` and the special tag ``_OP_``. Supported closing braces (equivalent to ``}``) are ``)``, ``]``, ``>`` and the special 
+  For brace expansion, there may be cases where ``{x..y}`` syntax cannot be used due to shell interference. This occurs most commonly
+  inside SLURM scripts. Therefore, we provide a few equivalent replacements. Supported opening braces (equivalent to ``{``) are ``(``,
+  ``[``, ``<`` and the special tag ``_OP_``. Supported closing braces (equivalent to ``}``) are ``)``, ``]``, ``>`` and the special
   tag ``_CL_``. For SLURM based tasks, we suggest the use of the special tags for ease of use.
 
-As with non-tarred datasets, the manifest file should be passed in ``manifest_filepath``. The dataloader assumes that the length 
+As with non-tarred datasets, the manifest file should be passed in ``manifest_filepath``. The dataloader assumes that the length
 of the manifest after filtering is the correct size of the dataset for reporting training progress.
 
-The ``tarred_shard_strategy`` field of the config file can be set if you have multiple shards and are running an experiment with 
+The ``tarred_shard_strategy`` field of the config file can be set if you have multiple shards and are running an experiment with
 multiple workers. It defaults to ``scatter``, which preallocates a set of shards per worker which do not change during runtime.
+Note that this strategy, on specific occasions (when the number of shards is not divisible with ``world_size``), will not sample
+the entire dataset. As an alternative the ``replicate`` strategy, will prealocat the entire set of shards to every worker and not
path: docs/source/asr/datasets.rst
text: @user1:
```suggestion
the entire dataset. As an alternative the ``replicate`` strategy, will pre-allocate the entire set of shards to every worker and not
```

@author:
Changed to "preallocate" to be consistent with the form in the rest of the text (L222).