Dataset creation

[Griffintaur]

I have tars where one the key is optional while other key is always present but when iterating through the dataset it raises a key missing error.
Is it the current design or there is way to navigate around it

[philipp-fischer]

Hi @Griffintaur,
this is a perfectly valid use case.
When you prepare your dataset with energon prepare, you need to choose not to create a simple field map, but instead a custom sample_loader.py.
The prepare wizard will create a template for you that you can modify. Please read this section of the documentation: https://nvidia.github.io/Megatron-Energon/advanced/advanced_dataformat.html

In your sample_loader.py you can write code to check if the key is present in the raw dict or not and create your sample accordingly. Please also read the docs section "Example: Interleaved Data and Arbitrary Image Count" which is a bit similar.

[Griffintaur]

So when we create a blended dataset by merging multiple datasets:
does the distribution of datasets across workers and rank happen in a continuous manner. i.e each worker will be allocated indices from dataset1 if their are indices remaining and then from dataset2...

[philipp-fischer]

Each dataset in the blend will be split across all the workers equally. So yes.

[Griffintaur]

so in that case we will never see the sample from dataset2 until dataset1 indices finishes in the worker.

[philipp-fischer]

No, each worker randomly chooses from which dataset to yield next. So it is blended as you would expect.
The magic happens here, where the dataset is picked by a probability according to its weight: https://github.com/NVIDIA/Megatron-Energon/blob/develop/src/megatron/energon/wrappers/blend_dataset.py#L111

Please also check out this section in the docs: https://nvidia.github.io/Megatron-Energon/basic/metadataset.html#metadataset