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
This guide explains how to 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. You can find more information about the filtering process in the FineWeb-Edu blog post.
The following example shows the data format:
Download the FineWeb-Edu Dataset
This guide uses the FineWeb-Edu 10BT sample, which is 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 a tool that tokenizes data and saves it in the Megatron Core dataset format. Use the tool as follows:
The directory should look like this:
Set --workers to the number of CPU cores to use for parallel tokenization.
Use a Recipe for Pretraining
This example demonstrates how to pretrain an LLM by using the NeMo AutoModel library. It uses the LLM training recipe, specifically TrainFinetuneRecipeForNextTokenPrediction. This recipe orchestrates 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, such as pretraining.
It is equivalent to a Trainer class but is 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 and validation steps. -
Additional responsibilities: Checkpoint handling, logging, and RNG setup.
Recipe Config Example
The following is a complete configuration based on examples/llm_pretrain/megatron_pretrain_gpt2.yaml:
To add weights to the dataset blends, pass 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"].
Balance Sampling and Domain Objectives
Sampling weights control how often a domain appears in a batch. They do not
need to be the same as the domain’s desired contribution to the optimization
objective. Configure both values with domain_mixture:
See the complete runnable
megatron_pretrain_gpt2_domain_mixture.yaml
example for the full model, dataset, validation, optimizer, and scheduler setup.
For a domain with normalized sampling weight (q_i) and objective weight
(w_i), the recipe multiplies each supervised token loss by (w_i / q_i).
This keeps one optimizer and learning-rate schedule while making the expected
training loss equal to the requested weighted domain objective. The recipe
logs achieved supervised-token fractions for each domain, reports every
per-domain validation loss, and reports their objective-weighted aggregate as
weighted.
The sampling_weight values must match the explicit weights in
the MegatronPretraining dataset’s dataset.paths, and domain order must
match the blend order. Domain names are positional labels and are not inferred
from path text. JSON blend files and implicit unweighted blends are not supported.
Sequence packing is not supported with domain_mixture because
packed batches do not currently preserve per-token domain IDs. Pipeline
parallelism is also not currently supported because its target contract does
not carry per-token objective weights. Context parallelism is disabled until
its distributed gradient contract is covered by parity tests. The configured
loss must use reduction: sum. The built-in masked, fused-linear, and chunked
cross-entropy losses accept these objective weights.
TEParallelCrossEntropy is not supported: its Triton backward reads the
upstream gradient as a single scalar, so per-token weights would apply the
first sample’s multiplier to the whole microbatch — the reported loss would
look correct while the gradients were wrong. The recipe rejects that loss at
setup rather than weighting it incorrectly.
Only the cross-entropy term is domain-weighted. A MoE auxiliary/router loss is
scaled independently of loss magnitude, so its strength relative to the
weighted CE varies with each microbatch’s domain composition; an
aux_loss_coeff tuned without domain_mixture may need revisiting.
The logged domain_mixture/<name>_label_tokens counts exclude labels equal to
-100. Datasets that mark padding through a separate loss_mask
rather than by writing ignore_index into labels (Megatron’s GPT dataset
does this) therefore report one supervised token per position, so the logged
fraction equals the raw sample fraction.
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 example, a 70B-parameter model requires about 140 GB of memory for its parameters when using the BF16 data type (2 bytes per parameter). Because most widely used GPUs are limited to 80 GB, the full model cannot be loaded directly onto a single device.
In these scenarios, you can pass is_meta_device: true in the model config. The model is then instantiated using PyTorch’s Meta device, which loads no data but stores the parameter metadata necessary for sharding the model. After the model is sharded, only the weights required by each model shard are loaded.
Run the Pretraining Recipe
If you saved or plan to use the provided configuration at examples/llm_pretrain/megatron_pretrain_gpt2.yaml, run the following command:
Sample Output
During training, stepwise logs report loss, memory usage, and tokens per second. The recipe saves checkpoints under the configured checkpoints/ directory.
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, compare your convergence curves with the baseline benchmarks in the llm.c repository.
