NVIDIA/Megatron-Energon

Megatron's multi-modal data loader

Megatron Energon – a data‑loader for huge multi‑modal model training

What it is – A Python package (and a small CLI) that prepares, indexes and streams large‑scale multi‑modal datasets (images, video, text, etc.) for NVIDIA’s Megatron‑LM training framework. It can also be used on its own.

Why it matters – Training today’s massive vision‑language or video‑language models requires petabytes of data spread across many storage systems. Megatron Energon handles the heavy lifting:

  • builds reproducible WebDataset or JSON‑Lines indexes,
  • merges many heterogeneous sources into a single logical dataset,
  • shards the work across many nodes / processes, and
  • provides a resumable pipeline so a failed job can pick up where it left off.

Key features (as described in the README)

Feature What it does
CLI (energon) Prepare WebDataset indexes, collect media metadata, mount a dataset as a read‑only virtual filesystem, and more.
Python API Functions like get_train_dataset and get_loader return a PyTorch‑compatible dataset/loader that you can iterate over in your training loop.
Distributed support WorkerConfig lets you specify rank, world size and number of workers so the same code runs on a single GPU or a full multi‑node cluster.
Multi‑modal friendliness Works with images, video, text, etc., using “cookers” registered in a TaskEncoder to turn raw samples into the tensors your model expects.
Reproducibility & resumability Indexes are stored in a .nv-meta folder; the loader can restart without re‑processing the whole dataset.
Installation pip install megatron-energon (stable) or install directly from the repo.

Typical workflow

  1. Prepare data – Run energon prepare /path/to/dataset (or on a JSONL file) to generate the .nv-meta index and optional splits.
  2. (Optional) collect media stats – energon prepare-media … gathers image sizes, video durations, etc., which can be used for filtering or bucketing.
  3. Load in code – Use the Python API to obtain a train_ds and a train_loader, then iterate over batches inside your training script.
  4. Combine with Megatron‑LM – Feed the batches to Megatron’s model parallel training loop for large‑scale LLM or multi‑modal model training.

Who it’s for – Researchers and engineers building very large vision‑language, video‑language, or other multi‑modal models who need a scalable, reproducible data pipeline. It is research‑grade code, so the API may evolve.

Where to learn more – Full docs at https://nvidia.github.io/Megatron-Energon/ and the GitHub repo for issues, tests, and examples.

Related

  • Dispatch
    Training Language Models with Megatron-LM

    Hugging Face provides a guide on using NVIDIA's Megatron-LM framework to efficiently pre-train large language models on GPUs, including integration with the Transformers library.

  • Project
    NVIDIA/Megatron-LM

    A GPU-optimized library for training massive transformer models at scale, providing advanced parallelism strategies and optimized building blocks for distributed training.

  • Project
    baidu-baige/LoongForge

    An open-source training framework that accelerates the training of LLMs, VLMs, diffusion, and embodied models using a multi-backend architecture.

  • Project
    texttron/tevatron

    A unified toolkit for training and deploying billion-scale neural retrievers across different languages and modalities using LLMs.

  • Project
    NVIDIA-NeMo/RL

    A scalable post-training library for multimodal models that streamlines reinforcement learning and alignment using high-performance backends like Megatron Core and vLLM.