As explained in [](sample-loading), the user has several options to choose how energon converts the raw (crude) data inside the tar files into Python sample objects (instances of a `Sample` dataclass) during loading.
When using crude datasets, this conversion happens through so-called "cookers", i.e. user-defined functions defined in the task encoder, as explained below.
In this case, the dataset on disk will specify neither the resulting sample type nor a sample loader for conversion, hence we call it "crude".
All of the conversion will happen in the user's code base.
## Setting Up a Crude Dataset with Cookers
Let's try it.
When you run `energon prepare` to prepare your dataset, you can pick "Crude sample" as the sample type.
If you already have an existing energon-compliant data set, you can modify it as follows (or create a copy of your `dataset.yaml` specifically for this use case and keep the original).
Original `dataset.yaml`:
```yaml
sample_type:
__module__:megatron.energon
__class__:OCRSample
sample_loader:sample_loader.py:sample_loader
part_filter:sample_loader.py:part_filter
```
New `dataset.yaml` (or `crude.yaml` of you like):
```yaml
__module__:megatron.energon
__class__:CrudeWebdataset
subflavors:
example_crude_selector:foobar
```
As shown above, [subflavors](sect-subflavors) can be assigned directly in the dataset yaml, however we recommend
to define them in the [Metadataset](../basic/metadataset) instead.
Subflavors can be used to differentiate the data origin when you mix different crude datasets or crude with non-crude data using a [Metadataset](../basic/metadataset).
You can then use the subflavors to determine how each sample shall be processed, as explained below.
Let's see how we set up our [Task Encoder](../basic/task_encoder) to cook the crude samples and turn them into real samples like {py:class}`TextSample <megatron.energon.TextSample>`.
A crude sample is essentially a dictionary ({py:class}`CrudeSample <megatron.energon.CrudeSample>` inherits from `dict`) and the loader will look for a {py:class}`Cooker <megatron.energon.Cooker>` that can handle and convert your crude samples.
Cooker(...)# other cookers for other crude data if needed
]
# ...
```
In the example above, the cooker acts on all crude samples that have a subflavor `example_crude_selector` set to `foobar`.
If you leave out the `has_subflavors` argument, the cooker will apply to any sample.
The cooker will convert the dictionary to a {py:class}`TextSample <megatron.energon.TextSample>` by decoding the raw bytes and decorating the text with some nice angle brackets.
Probably you noticed the {py:meth}`basic_sample_keys <megatron.energon.task_encoder.cooking.basic_sample_keys>` helper that we inserted.
All it does, is to forward the key, restore key and flavors from the dict to the real sample. You will always need to forward these, or your dataset will not be restorable.
In a real use-case you will want to do a lot more here and we recommend keeping the cook methods in separate files and importing them where you define your TaskEncoder.
(aux-data)=
## Auxiliary Data for Polylithic Datasets
Using a crude dataset allows you to benefit from two other features of energon:
* Auxiliary Data
* Cache Pools
Both of which are often used in combination. A typical use case is online packing.
An **auxiliary data source** is an additional data source that supports random access and can be used to load data on-demand using its filename.
It is typically used with polylithic datasets where you have one primary dataset that contains only the text-based sample data
and one or more additional auxiliary data sources that contain the (larger) media data such as images or videos.
An auxiliary data source can be either
* Another energon-prepared WebDataset
* A folder on the local or a remote file system
You can specify it in your [metadataset](../basic/metadataset) yaml as follows (look at the `aux:` section)
You can specify multiple aux sources each of which can be one of
* Relative or absolute path to a local prepared energon dataset
* Relative or absolute path to a local folder (use the prefix `filesystem://`)
* Path to a remote prepared energon dataset (use prefix `msc://`)
**[Planned future feature]*: Path to a remote folder (use prefix `filesystem+msc://`)
In your code, the cooker will automatically receive a {py:class}`FileStore <megatron.energon.FileStore>` reference to the data source as a keyword argument:
You can use multiple sources. You'll have to specify a cooker argument for each source that was defined in the metadataset.
For easier debugging, you should always keep track of all the sources you used. The `get` method takes care of this if you pass it the sample like this:
This will update the sample-internal `__sources__` list with the aux dataset you used.
If you want, you can even use your primary dataset as an auxiliary dataset and look up files by name, yes! If you want to do that, you specify it in the cooker decorator and add an arg:
You can then retrieve files by their names from the primary dataset.
(cache-pools)=
## Cache Pools
Cache pools allow the user to defer the data transfer if the content will be needed at some point in the future but not immediately.
This is only needed if the media data is rather large and does not reside on a local disk, but rather on a network file system (e.g. lustre) or a remote file system (e.g. object storage).
Cache pools are especially beneficial if you are using buffers in your pipeline, like a shuffle buffer or a packing buffer. For example, when using [online packing](../advanced/packing), we may need to keep a buffer of several thousand samples to optimize for the best packing,
but we cannot keep several thousand images in memory, also we don't need the actual image content to optimize the packing.
Hence we will use auxiliary datasets as explained above.
However, at the time of filling the buffer, **we already know** that we **will need the image content in the future**, so cache pools can be used to **prefetch it in the background**.
Initially we want to load some information about the sample and its image but not the actual image pixels.
Later, when the packing is computed, we need to retrieve the pixel values.
In practice, this means the cooker will use a cache pool to queue the data retrieval from an auxiliary data source and obtain a lazy object (a handle to this future data). In a later stage (like {py:meth}`pack_selected_samples <megatron.energon.TaskEncoder.pack_selected_samples>`), the lazy object can be used to retrieve the content.
Ideally, in the mean-time, the cache pool has already downloaded the data to a local SSD.
### Using a Cache Pool
When calling {py:func}`get_savable_loader <megatron.energon.get_savable_loader>`,
we pass a cache pool as an additional argument:
```python
frommegatron.energonimportFileStoreCachePool
# ...
loader=get_savable_loader(
my_ds,
...,
cache_pool=FileStoreCachePool(
parent_cache_dir="/local_scratch/cache",
num_workers=1,
),
)
```
Then we tell the cooker decorator that we need access to the cache pool and use it to get a lazy object:
There is a second option, e.g. if you want to combine a monolithic dataset with packing and caching: Use `cache.to_cache()` to move already loaded data to the cache:
Your `sample_loader` method must accept a dictionary as argument and return a dictionary. It directly operates on
the webdataset samples and the resulting dictionary keys should map to the corresponding sample class members.
In this case an {py:class}`OCRSample <megatron.energon.OCRSample>`.
With the optional `part_filter` method, you can prevent some webdataset fields from being loaded.
Given a field name, the method should return True if the field is to be kept.
(interleaved-sample-loader)=
## Example: Interleaved Data and Arbitrary Image Count
### The Webdataset Structure
If you need multiple files with an arbitrary number of data per sample, e.g. multiple image / video / audio files, this shows a blueprint for how to setup your webdataset tar files and how to load that webdataset with Energon.
The structure of the shard files could be like this:
`tar -tvf shard_0.tar`:
```python
sample_000001.2345ew.jpg
sample_000001.json
sample_000002.35tags.jpg
sample_000002.as23ds.jpg
sample_000002.gd1dtg.jpg
sample_000002.gds233.jpg
sample_000002.json
sample_000002.sdag42.jpg
sample_000003.json
sample_000004.asf234.jpg
sample_000004.json
```
where the structure of a json file is:
`tar -xf shard_0.tar sample_000001.json -O`:
```json
{
"images":[null,"2345ew.jpg",null],
"texts":["This is some text, an image is following.",null,"More text after the image."],
}
```
Note that the image path corresponds to the filename of the image after the first "." in the sample. This is all part of the extension as defined by webdataset. Everything before the first "." is part of the sample key and must be equal to match into the same group.
### Usage with Energon
To make this work with Energon, in the `energon prepare`[CLI preparation tool](energon-prepare), you can either tell the wizard to create a custom sample loader template for you, or change the files accordingly. Here is the example with the structure above:
`.nv-meta/dataset.yaml`:
```yaml
sample_type:
__module__:megatron.energon
__class__:InterleavedSample
part_filter:sample_loader.py:part_filter
sample_loader:sample_loader.py:sample_loader
```
`.nv-meta/sample_loader.py`:
```python
importtorch
defsample_loader(raw:dict)->dict:
# Note that the images are already decoded, as well as the json part.
When iterating, you'll get those samples as `InterleavedSample` which either contains the image tensor, or the strings for text. The first sample would look like this:
```python
InterleavedSample(
sequence=["This is some text, an image is following.",torch.Tensor(...)orPIL.Image.Image(),"More text after the image."]
Grouping allows for rule-based batching of samples into one batch on the fly.
Note how this is different from [packing](packing) which joins multiple samples into one (and is done before batching).
On the other hand, grouping is an alternative to standard batching.
## Example use cases
* Select samples to batch based on image resolution, so that only samples of the same size are in one batch
* Select blended samples based on their dataset origin, so that one batch does not mix different tasks or data types
## How to group
To use grouping, you need to define the method {py:meth}`batch_group_criterion <megatron.energon.TaskEncoder.batch_group_criterion>` in your custom task encoder.
This method gets a sample and returns a hashable value that will be used to cluster/group the samples
and it also returns the batch size for that group.
Samples with the same batch group criterion will be batched together. Once enough samples for one group
have been collected (reached the batch size for that group), they will be batched and pushed down the pipeline
to the next processing step.
Here's an example task encoder that batches samples based on their image aspect ratios:
Joining is needed when you have two or more corresponding datasets, each dataset having some parts to compose each sample.
```{image} ../images/joining.png
:alt: Joining Datasets
:class: bg-primary
:width: 550px
:align: center
```
Dataset joining is analogous to webdataset's "add column", i.e., merging two datasets which require the same files within the tar files.
## Example File Structure
Here is an example structure of the datasets within the tar files:
```
ds1
├── .nv-meta
│ ├── .info.json
│ ├── split.yaml
│ └── dataset.yaml
├── shard1.tar
│ ├── 0001.jpg
│ ├── 0002.jpg
│ └── 0003.jpg
├── shard1.idx
ds2
├── .nv-meta
│ ├── .info.json
│ ├── split.yaml
│ └── dataset.yaml
├── shard1.tar
│ ├── 0001.txt
│ ├── 0002.txt
│ └── 0003.txt
├── shard1.idx
```
In this example, `ds1/shard1.tar` and `ds2/shard1.tar` contain files with the same base names / keys.
When joining, the files from both datasets are joined based on their keys, which must be unique within each dataset across all the shards. Each dataset must be [prepared](../basic/data_prep)(i.e. .nv-meta created) with a recent version of energon.
## Minimal Example
Below is a minimal meta dataset example for joining.
`metadataset_minimal.yaml`:
```yaml
__module__:megatron.energon
__class__:MetadatasetV2
splits:
train:
join:
-path:ds1
-path:ds2
joiner:
__module__:megatron.energon
__class__:CaptioningSample
```
Note that the blending is omitted here, as it is optional. Joining is of course also optional.
```{warning}
All metadatasets that contain a `join` must be prepared with the `energon prepare` command.
This will compute the join index and store it next to the metadataset in a folder with a similar name.
```
## Join Modes
When joining two datasets, it may happen that the first dataset (primary dataset) has more samples or fewer samples than the secondary dataset(s).
In this case, we have to decide how to handle these samples that do not match.
The primary dataset always serves as the reference and there will never be more samples in the join result than in the primary dataset. However if a primary sample has no match in a secondary dataset, it may be skipped as explained below.
For each of the secondary datasets, the user can specify a `nonmatch` setting.
With one of the following options, the user can decide what happens, if a sample from the primary dataset is not found in the given secondary dataset:
*`error` (default): An error is raised
*`skip`: The whole sample is skipped
*`none`: The column for the current secondary dataset is filled with `None` if there's no match
Example `metadataset_nomatch.yaml`:
```yaml
__module__:megatron.energon
__class__:MetadatasetV2
splits:
train:
join:
-path:ds1
-path:ds2
nonmatch:skip
-path:ds3
nonmatch:none
joiner:
__module__:megatron.energon
__class__:CaptioningSample
```
To illustrate the effect, let's look at some example data:
*`ds1` samples: `s1`, `s2`, `s3`, `s5`, `s6`
*`ds2` samples: `s1`, `s3`, `s4`, `s6`, `s7`
*`ds3` samples: `s1`, `s2`, `s3`, `s100`
The resulting joined data would contain the following samples, one in each row:
| ds1 | ds2 | ds3 |
| --- | --- | ---- |
| s1 | s1 | s1 |
| s3 | s3 | s3 |
| s6 | s6 | None |
Explanation:
* The sample key `s1` is available in all dataset.
*`s2` is missing from `ds2` and nonmatch is set to `skip`, so the sample will not appear in the result.
*`s3` is available in all datasets.
*`s4` is not in the primary dataset. Only samples from the primary dataset will be included.
*`s5` is missing from `ds2` again, and this time also from `ds3`
*`s6` is missing from `ds3` and `ds3` has `nonmatch` set to `none`, so the sample is not skipped, but the column for `ds3` is set to `None`
## Extensive Example
Here is a more extensive example that shows multiple things at once:
* Joining can be used inside blending
* The datasets to be joined can have custom subflavors or dataset yamls specified
* A custom "joiner" can be specified to define how samples are joined and what the resulting type is
* The `nonmatch` setting is not included here, but would work just like shown above
`metadataset_extended.yaml`:
```yaml
__module__:megatron.energon
__class__:MetadatasetV2
splits:
train:
blend:
-weight:1
join:
-path:ds1
dataset_config:dataset.yaml# If override is needed
-path:ds2
dataset_config:dataset.yaml
subflavors:# If needed, will be merged(overriding) with parent subflavor
ds2_extra:2
split_config:split.yaml
joiner:
__module__:my_module
__class__:JoinedSample# Type should implement from_joined(ds1, ds2)
split_config:split.yaml# Sets this for all joined datasets
split_part:train# Sets this for all joined datasets
subflavors:# Sets this for all joined datasets (it will be merged with their individual subflavors)
source:metadataset.yaml
src:ds1
```
## Custom Join Type
To define a custom join type, you can create a Python class as shown below in `my_module.py`:
This class should implement the `from_joined` method to combine samples from `ds1` and `ds2`.
Note: It is important to use `derive_from` with the first argument being the first sample, as this will guarantee that the state can be saved and restored. It ensures that all the internal keys of the sample are retained.
Neural network parallelism can be categorized into several types:
1.**Data Parallelism** (DP): This involves splitting the data across multiple processors and performing the same operation on each subset of the data. It is commonly used to increase the global batch size.
2.**Model Parallelism**: In this approach, different parts of the model are distributed across multiple processors. This is useful when the model itself is too large to fit into the memory of a single processor.
3.**Pipeline Parallelism** (PP): This technique involves breaking down the model into different stages and processing different mini-batches of data through these stages in a pipeline fashion. It helps in improving the utilization of resources and reducing idle time.
4.**Tensor Parallelism** (TP): This method splits individual tensors (weights and activations) across multiple devices. It is particularly effective for very large models where even a single layer cannot fit into the memory of one device.
These parallelisms have different consequences for the dataloader:
-**Data Parallelism** (DP): The dataloader needs to ensure that each processor gets a different subset of the data. This is supported by Energon. The data parallel groups should be specified in the worker config.
-**Pipeline Parallelism** (PP): Data is typically only loaded on the first Pipeline Parallel rank, and propagates through the other ranks within the pipeline parallel group. This means, you only instantiate an Energon dataset and loader on the first ranks of those groups.
-**Tensor Parallelism** (TP): The dataloader will load the same input data on multiple devices. Typically, this can be ensured by either instantiating the dataloader exactly the same on the same data parallel ranks in different data parallel groups, or e.g. by loading the data only once and distributing it using torch distributed.
## Example
Example with the following ranks and worker configuration (Data Parallel = 2, Pipeline Parallel = 2, Tensor Parallel = 2):
Megatron Energon supports the use of remote datasets. Since version >5.2.0, Energon file access is based on [Multi Storage Client (MSC)](https://github.com/NVIDIA/multi-storage-client).
This means you can train or validate with your data right from any storage by simply swapping the dataset path for a so-called _MSC URL_.
## Prerequisites
For using a remote dataset, install energon with one or more of the extras:
*`s3`
*`aistore`
*`azure-blob-storage`
*`google-cloud-storage`
*`oci`
like this:
```sh
pip install megatron-energon[s3,oci]
```
Set up the msc config as described in [Multi Storage Client documentation](https://nvidia.github.io/multi-storage-client/).
You can also use the rclone config with msc, as was described prior to 5.2.0.
For fast data loading we recommend to activate MSC local caching:
```yaml
cache:
size:500G
use_etag:true
eviction_policy:
policy:"fifo"
refresh_interval:3600
cache_backend:
cache_path:/tmp/msc_cache# prefer to use local NVMe, but Lustre path also works
```
And point MSC to the config with
```sh
export MSC_CONFIG=/path/to/msc_config.yaml
```
## The URL syntax
The syntax is a simple as
```
msc://CONFIG_NAME/PATH
```
For example:
```
msc://coolstore/mainbucket/datasets/somedata
```
You can use this URL instead of paths to datasets in
* Functions like `get_train_dataset`, `get_val_dataset`
Dataset subsets allow restricting a dataset (or parts of a metadataset hierarchy) to a specific portion of the available samples.
This is useful for rapid prototyping, ablation studies, different training stages, or constructing disjoint train/validation/test splits that differ from the original dataset configuration.
A subset is defined by a two-element `range` list consisting of `[start, end]` (where `start` is inclusive, `end` exclusive).
Each element can be either
* a **percentage** string (e.g. `"0%"`, `"12.5%"`, `"100%"`) – interpreted relative to each inner
dataset size, or
* an **absolute** integer – interpreted as a sample index. Absolute indices are only allowed for
*leaf* datasets (`path` to a prepared dataset containing `.nv-meta`).
## Basic example
The snippet below keeps the first 80 % of *COYO*`train` split (as defined in the `split.yaml`) for training while
evaluating on the remaining 20 % of the `train` split. Note how the `subset` key is placed directly next to the corresponding `path`.
```yaml
__module__:megatron.energon
__class__:MetadatasetV2
splits:
train:
path:./coyo
subset:{range:[0%,80%]}
val:
path:./coyo
split:train
subset:{range:[80%,100%]}
```
## Nested subsets and merging rules
Subsets can appear at any level that ultimately yields samples
(direct `path` reference to a prepared dataset containing `.nv-meta`, `join`, `blend`, `blend_epochized`).
When multiple subsets are nested, the *inner* subset is applied first, then the portion selected by the *outer* subset is applied *within* the already selected range.
For percentages the ranges are composed multiplicatively.
Example: the outer subset `[0%, 50%]` followed by an inner subset `[25%, 75%]` results in the final
range `[25%, 50%]` of the original dataset.
Absolute indices short-circuit merging: they can **only** be specified at the leaf level and must
not be combined with another absolute range farther up the hierarchy.
## Absolute ranges
Absolute indices are handy when exact sample counts are required.
## Advanced examples
The following configuration combines the absolute ranges with the nested rules. The inner subset takes
the first **1000** samples from *COCO* train split and mixes them with the full *COYO* train split using
weight-based blending. The outer nesting then reduces the inner ranges to the first 50%, thus only taking
the first **500** samples of *COCO*, mixed with the first **50%** of the *COYO* dataset effectively.
```yaml
__module__:megatron.energon
__class__:MetadatasetV2
splits:
train:
# The subset is applied to each blended dataset separately.
# I.e. for the first, the sample range is [0, 500], for the second the range is [0%, 50%]
subset:{range:[0%,50%]}
blend:
-weight:1.0
path:./coco
subset:
# Take exactly 1000 samples (indices 0-999)
range:[0,1000]
-weight:1.0
path:./coyo
```
Absolute ranges can also be specified to run up to the end of the dataset using the `end` keyword:
```yaml
__module__:megatron.energon
__class__:MetadatasetV2
splits:
train:
path:./coyo
subset:{range:[1422,end]}
```
## Python usage
No API changes are required on the Python side – subsets are fully specified in the YAML. Simply
The above command will scan your existing off-the-shelf [web dataset](https://webdataset.github.io/webdataset/)
and add the [needed metadata](data-on-disk) to make it compatible with Energon.
Below, you can see the available sub-commands under `energon`.
```{eval-rst}
.. click:: megatron.energon.cli.main:main
:prog: energon
:nested: short
```
(energon-prepare)=
## energon prepare
An interactive tool to generate metadata for your existing webdataset (or [jsonl dataset](create-jsonl-dataset)).
This will help make the dataset compliant with our [format](data-on-disk).
The tool will ask you for a train/val/test split and how to assign the webdataset fields to the
fields of the corresponding sample type in Energon.
See [Data Preparation](../basic/data_prep) for more details on how to use this command.
## energon info
Prints information about the dataset such as overall number of samples and size.
It also prints the energon version that was used to prepare the dataset, if a recent version was used.
## energon lint
You can execute this tool on the prepared dataset to check if the data is valid and loadable.
It will report any problems such as non-readable images.
(energon-mount)=
## energon mount
Use this to mount your [prepared dataset](../basic/data_prep) as a virtual read-only filesystem and inspect it using `ls` or other file browsing tools.
It is as simple as running
```shell
energon mount /PATH/TO/DATASET ./MY_MOUNT_POINT
```
This will leave the process in the foreground and the mount will exist as long as the program is running.
If you want to detach the process to the background, use the `-d` or `--detach` flag.
| Description | All files from all shards listed at<br/>the root of the mount point. | One folder per sample key,<br/>each folder containing files<br/>named by the sample part extension |
