Tar codec

[nithinraok]

Add webdataset support for VocoderDataset

Collection: TTS

Changelog

• Add TarredVocderDataset class for dataset fetching
• sampler is not supported with IterableDataset
• added a conditional code to calculate correct number of max steps when webdataset is used

Known issue:

• Haven;t tested for passing multiple tarred datasets. Will be pushing a PR later to support it.

Usage

• add is_sharded=True in cfg for train_ds to load webdataset

Before your PR is Ready for review

Pre checks:

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

PR Type:

• New Feature
• Bugfix
• Documentation

[rlangman]

Is this the same outcome as defining this in TarredVocoderDataset?:

get_sampler():
    return None

[nithinraok]

Good point i will create a simpler case here, but I created this function to set the target so as to use existing config files and provide toggle option for sharding.

[rlangman]

Given that these are experimental recipes that are undergoing breaking changes and have no public models, I think it should be OK to require config updates. The tarred dataset also requires different inputs to instantiate it.

[nithinraok]

I am against creating one new config for each feature.

[rlangman]

Do you mean you are against creating separate config files for new features in general (eg. sampling rate, new architectures)? Or specifically for needing separate config for tarred vs non-tarred datasets?

The former of creating lots of different overlapping configs is expected across NeMo. The latter can be addressed in several different ways, if needed.

[nithinraok]

We don;t want to create multiple configs with same model architecture but different setting (like tarred/sampling_rate etc).

[rlangman]

For the same model architecture with different configs are main options I have heard before are:

1. Create different example configs.
2. Have one example config and provide web documentation and tutorials clearly indicating to users how to manually change fields for different valid configs (eg. sample rate, language).

From past discussions around this, generally (1) was preferable, especially for changes like sample rate which are relatively complicated and error prone if you attempt to update all needed parameters manually. A combination of (1) and (2) is needed if you have too many configs to support, such as when having to support multiple sample rates and languages at the same time.

[rlangman]

In this particular case, there are 2 problems with the existing approach.

• The model class takes the data loader class/config as input, however the constructor requires passing information global_rank and world_size which are known only to the model class and trainer.
• Different data loaders require significantly different inputs to instantiate (and this will diverge more as we add additional sampling strategies).

Perhaps the cleanest way to address these would be to change it so the model is expected to receive a subset of dataset parameters, and then add a module which maps the parameters to the logic to create each dataset.

For example:

In .yaml

train_ds:
  dataset_type: "vocoder_tarred"
  dataset_args:
    batch_size: ${batch_size}
    ... 

In vocoder_dataset.py

def create_dataset(dataset_type, dataset_args, global_rank, world_size, ...):
    if dataset_type == "vocoder":
        return VocoderDataset(*dataset_args)
    elif dataset_type == "vocoder_tarred":
        return TarrredVocoderDataset(*dataset_args, global_rank=global_rank, world_size=world_size)
     ...

[nithinraok]

Good catch, missed it. Updated to use those as well.

[nithinraok]

for logging, testing, validation we rely only on non tarred datasets.

[rlangman]

Given that these are experimental recipes that are undergoing breaking changes and have no public models, I think it should be OK to require config updates. The tarred dataset also requires different inputs to instantiate it.

[rlangman]

In general we need to figure out how to make the model class agnostic to the data loader. Especially given that the logic will need to be kept consistent across all vocoder recipes (audio codec, hifigan, and univnet at least).

If we think upgrading webdataset version is a blocker to streamline the code, then we can document that with some TODOs and leave these workarounds. Otherwise, we should redesign the contract between the dataset and the model classes so they work with both.

[rlangman]

Nitpick: it is not very important for audio only dataset, but code would be more readable and extensible using original dictionary format instead of converting to tuple.

[nithinraok]

jenkins

[nithinraok]

jenkins

[nithinraok]

jenkins

[rlangman]

For the same model architecture with different configs are main options I have heard before are:

1. Create different example configs.
2. Have one example config and provide web documentation and tutorials clearly indicating to users how to manually change fields for different valid configs (eg. sample rate, language).

From past discussions around this, generally (1) was preferable, especially for changes like sample rate which are relatively complicated and error prone if you attempt to update all needed parameters manually. A combination of (1) and (2) is needed if you have too many configs to support, such as when having to support multiple sample rates and languages at the same time.

[rlangman]

In this particular case, there are 2 problems with the existing approach.

• The model class takes the data loader class/config as input, however the constructor requires passing information global_rank and world_size which are known only to the model class and trainer.
• Different data loaders require significantly different inputs to instantiate (and this will diverge more as we add additional sampling strategies).

Perhaps the cleanest way to address these would be to change it so the model is expected to receive a subset of dataset parameters, and then add a module which maps the parameters to the logic to create each dataset.

For example:

In .yaml

train_ds:
  dataset_type: "vocoder_tarred"
  dataset_args:
    batch_size: ${batch_size}
    ... 

In vocoder_dataset.py

def create_dataset(dataset_type, dataset_args, global_rank, world_size, ...):
    if dataset_type == "vocoder":
        return VocoderDataset(*dataset_args)
    elif dataset_type == "vocoder_tarred":
        return TarrredVocoderDataset(*dataset_args, global_rank=global_rank, world_size=world_size)
     ...

[nithinraok]

jenkins

[rlangman]

If I am not mistaken, shuffle_n is the size of the shuffle buffer as the gpu reads in data from its shards. Allocating more shards to the gpu would not really reduce its dependence on shuffle_n. Intuitively, I would expect it to become more important (larger dataset requires more shuffling).

Somewhat related, wd. WebDataset() has a shardshuffle=True parameter that ensures that the shard list itself is fully shuffled.

[nithinraok]

From webdataset Docs:
https://webdataset.github.io/webdataset/gettingstarted/
shuffle(n): shuffle the dataset with a buffer of size n; also shuffles shards (see below)

[rlangman]

Based on the thread linked below, it sounds like the comment is effectively saying "larger datasets are less sensitive to how they are shuffled".

[nithinraok]

updated to have both buffer and inital arguments to have more control over shuffling

[rlangman]

Should we do something like dataset.shuffle(size=shuffle_n, initial=shuffle_n)? The initial buffer size of 100 is really small, which I saw slow convergence a lot on my small test datasets (but might be irrelevant on very large datasets).

[nithinraok]

If the set is small then yes number of samples per worker would become small and shuffle might not have an effect and wouln;t be a problem for larger sets as number of samples per worker would be large. As per this discussion, it looks like if we want to get a true shuffle of over all sets, its better to unbatch() on loader and shuffle with a larger number.

[rlangman]

Nitpick: other classes in the TTS collection either provide no default sample_rate or defaults to 22050.

[nithinraok]

Default for our audio codec models is 24000. Did we have a 22050 Hz model?

[rlangman]

No, but this is a data loader for all TTS vocoders. We have never used 24khz in TTS before (and when I brought up the question before of whether we should do so to synchronize with SpeechLM models, the concensus was to stick to 22.05khz). So I suppose I would favor not having a default sample rate.

[nithinraok]

updated

[nithinraok]

jenkins

[nithinraok]

jenkins

[rlangman]

LGTM. Thanks a lot for the feature and changes.