Pretrain Megatron Core Datasets with NeMo AutoModel
Introduction
Pretraining builds a base large language model (LLM) by training a randomly initialized model to predict the next token across massive, unlabeled datasets.
Robust pretraining establishes a foundation of linguistic competence and world knowledge that scales with data, parameters, and compute. This base model then serves as the necessary starting point for later fine-tuning or domain-specific adaptation.
NeMo AutoModel provides an end-to-end recipe to run LLM pretraining with Hugging Face–native models and Megatron-Core style datasets.
Model and Dataset Context
In this guide, we pretrain OpenAI’s GPT2-124M model on a FineWeb-Edu subset of 10 billion tokens.
About the FineWeb-Edu Dataset
FineWeb-Edu is a dataset consisting of 1.3T tokens of educational web pages filtered from the larger FineWeb dataset. The educational web pages were filtered from the main dataset using a fine-tuned BERT-like classifier. Further reading on the filtering process can be found here.
Here’s a glimpse of what the data looks like:
Download the FineWeb-Edu Dataset
For this guide, we use the FineWeb-Edu 10BT sample — a collection of approximately 10 billion tokens randomly drawn from the full FineWeb-Edu dataset. To prepare the data, run the following commands:
Set MEMORY_GB to the amount of system memory allocated to terashuf (the tool used for sample shuffling), and set DATA_DIR to the root directory where the data will be stored. For example:
The expected directory structure is as follows:
Preprocess to a Megatron Core Dataset
NeMo AutoModel provides tooling to perform the task of tokenizing and saving in the Megatron Core dataset format. You can use it as follows:
The directory should look like this:
Replace --workers with the amount of CPU cores you’d like to use to tokenize in parallel.
Use a Recipe for Pretraining
This example demonstrates how to perform pretraining on a large language model using NVIDIA’s NeMo AutoModel library. We use the LLM training recipe, specifically TrainFinetuneRecipeForNextTokenPrediction, which orchestrates the pretraining process — including loading, dataset preparation, optimizer setup, distributed training, checkpointing, and logging.
What is a Recipe?
A recipe in NeMo AutoModel is a self-contained orchestration module that wires together all components needed to perform a specific task (e.g., pretraining). Think of it as the equivalent of a Trainer class, but highly modular, stateful, and reproducible.
The TrainFinetuneRecipeForNextTokenPrediction class is one such recipe. It inherits from BaseRecipe and implements:
-
setup(): builds all training components from the config -
run_train_validation_loop(): executes training + validation steps -
Additional responsibilities: Checkpoint handling, logging, and RNG setup.
Recipe Config Example
Below is a complete configuration based on examples/llm_pretrain/megatron_pretrain_gpt2.yaml:
If you want to add weights to the dataset blends, you can do so by passing in a list. For example, paths: ["30", "fineweb_edu/megatron_gpt2/processed_data_0_text_document", "70", "fineweb_edu/megatron_gpt2/processed_data_1_text_document"].
Load Large Models
In distributed training, the typical model-loading pipeline has each GPU load the entire model and then retain only the shard it needs. This approach becomes problematic when the model size exceeds the memory capacity of a single GPU. For instance, a 70B-parameter model requires about 140GB of memory for its parameters when using the BF16 data type (2 bytes per parameter). Since most widely used GPUs are limited to 80GB, the full model cannot be directly loaded onto a single device.
In these scenarios, you can pass is_meta_device: true in the model config. The model will then be instantiated using PyTorch’s Meta device which loads no data, but stores all other parameter metadata necessary for sharding the model. Once the model is sharded, the model weights will be populated by only loading the weights required by the respective model shard.
Run the Pretraining Recipe
Assuming you saved, or plan to use, the provided config at examples/llm_pretrain/megatron_pretrain_gpt2.yaml:
Sample Output
You should see step‑wise logs reporting loss, memory usage, and tokens per second. Checkpoints will be saved under the checkpoints/ directory as configured.
For each training batch, the pretraining recipe logs the current loss, along with current peak memory usage and tokens per second (TPS).
As training progresses, you should observe the model loss beginning to converge. To verify your results, you can compare your convergence curves against the baseline benchmarks provided in the llm.c repository.
