<!--- Copyright (c) 2025, NVIDIA CORPORATION.
SPDX-License-Identifier: BSD-3-Clause -->
<a name="top"></a>
<div align="center">
<h4 align="center">Megatron's multi-modal data loader</h4>
<h2 align="center">Megatron Energon</h2>
<p align="center">
<a href="https://github.com/NVIDIA/Megatron-Energon/actions/workflows/tests.yml"><img src="https://github.com/NVIDIA/Megatron-Energon/actions/workflows/tests.yml/badge.svg" alt="Tests"></a> <a href="https://nvidia.github.io/Megatron-Energon/"><img src="https://github.com/NVIDIA/Megatron-Energon/actions/workflows/documentation.yml/badge.svg" alt="Documentation"></a>
<br />
<a href="https://github.com/NVIDIA/Megatron-Energon/issues">Report Bug</a>
·
<a href="https://github.com/NVIDIA/Megatron-Energon/issues">Request Feature</a>
</p>
</div>
<br />
_**DISCLAIMER**: This package contains research code. APIs may change._
# What is this?
**Megatron Energon** is the multi-modal data loader of [Megatron](https://github.com/NVIDIA/Megatron-LM) (you can also use it independently).
It's best at
- loading large training data to train large multi-modal models
- blending many different datasets together
- distributing the work across many nodes and processes of a cluster
- ensuring reproducibility and resumability
- adapting easily to various types of data samples and processing
Try using it together with [Megatron](https://github.com/NVIDIA/Megatron-LM) Core.
# Quickstart
**Megatron Energon** is a pip-installable python package that offers
- dataset-related classes that you can import in your project
- a command line utility for data preprocessing and conversion
This document is just a quick start. Please also check out the [documentation](https://nvidia.github.io/Megatron-Energon/).
## Installation
To install the latest stable version:
```shell
pip install megatron-energon
```
Or to install the current development version:
```shell
pip install git+https://github.com/NVIDIA/Megatron-Energon.git
```
**NOTE**: We encourage you to install the package (and not just import a local copy). This will ensure you have all the needed dependencies and that you can use the command line tool.
For more details on installing this package, see [here](https://nvidia.github.io/Megatron-Energon/intro/installation.html).
## Usage of command line tool
After installation, the command `energon` will be available.
Here are some examples for things you can do:
| Command | Description |
|---|---|
| `energon prepare DATASET_ROOT` | Interactively prepare WebDataset indexes, splits, and configuration |
| `energon prepare DATASET.jsonl` | Create an index for a single JSONL file |
| `energon prepare-media --media-metadata-by-extension MEDIA_ROOT` | Collect media metadata, such as image dimensions and video duration |
| `energon mount DATASET_ROOT MOUNTPOINT --sample-folders` | Browse a prepared WebDataset as a read-only virtual filesystem, grouped by sample |
| `energon --help` | List available commands and options |
For example, prepare a crude WebDataset without interactive prompts:
```shell
energon prepare /path/to/dataset \
--non-interactive \
--sample-type=CrudeWebdataset \
--split-ratio=98,1,1
```
Its samples are processed by cookers registered in your application's `TaskEncoder`.
## Usage of the library
To get started, pick a [WebDataset](https://github.com/webdataset/webdataset)-compliant dataset and run `energon prepare DATASET_ROOT` on it, to run the interactive assistant and create the `.nv-meta` folder. As an alternative to WebDataset, Energon also supports the JSONL format, see [here](https://nvidia.github.io/Megatron-Energon/basic/data_prep.html).
Once done, try to load it from your Python program:
```python
from megatron.energon import get_train_dataset, get_loader, WorkerConfig
simple_worker_config = WorkerConfig(rank=0, world_size=1, num_workers=2)
train_ds = get_train_dataset(
'/my/dataset/path',
batch_size=2,
shuffle_buffer_size=None,
max_samples_per_sequence=None,
worker_config=simple_worker_config,
)
train_loader = get_loader(train_ds)
for batch in train_loader:
# Do something with batch
# Infer, gradient step, ...
pass
```
For more details, read the [documentation](https://nvidia.github.io/Megatron-Energon/).