Important

NeMo 2.0 is an experimental feature and currently released in the dev container only: nvcr.io/nvidia/nemo:dev. Please refer to the Migration Guide for information on getting started.

Migrate Pre-Training from NeMo 1.0 to NeMo 2.0

NeMo 1.0 (Previous Release)

In NeMo 1.0, pre-training is configured using megatron_gpt_config.yaml and launched with megatron_gpt_pretaining.py.

NeMo 2.0 (New Release)

NeMo 2.0 introduces a Pythonic and modular approach to configuring experiments. For detailed instructions on migrating your NeMo 1.0 YAML configurations to NeMo 2.0, refer to the additional documents in this migration guide:

Data

Optimizer

Checkpointing

Experiment Manager

Trainer

In addition, NeMo 2.0 is compatible with NeMo-Run, which streamlines the configuration and execution of NeMo experiments. Refer to the NeMo-Run documentation for more.

The llm.train API can be used to run pre-training in NeMo 2.0, as follows:

import torch
from nemo import lightning as nl
from nemo.collections import llm
from megatron.core.optimizer import OptimizerConfig

### set up your GPT model config
gpt_config = llm.GPTConfig(
   num_layers=12,
   hidden_size=384,
   ffn_hidden_size=1536,
   num_attention_heads=6,
   seq_length=seq_length,
   init_method_std=0.023,
   hidden_dropout=0.1,
   attention_dropout=0.1,
   layernorm_epsilon=1e-5,
   make_vocab_size_divisible_by=128,

)

### other docs in this section explain how to configure these
data = llm.PreTrainingDataModule(...)
model = llm.GPTModel(gpt_config, tokenizer=data.tokenizer)
strategy = nl.MegatronStrategy(...)
opt_config = OptimizerConfig(...)
opt = nl.MegatronOptimizerModule(config=opt_config)
trainer = nl.Trainer(...)
nemo_logger = nl.NeMoLogger(dir="test_logdir")

llm.train(
   model=model,
   data=data,
   trainer=trainer,
   log=nemo_logger,
   tokenizer='data',
   optim=opt,
)

In addition to the generic GPTModel used in the example above, we also support Gemma, Llama, Mistral, and Mixtral models.

Migration Steps

Migrate your NeMo 1.0 YAML to NeMo 2.0 using the other documents in the migration guide.
Run pre-training using the llm.train API.