Megatron-LM: GPU-optimized transformer training at scale

import torch
from megatron import get_args, get_timers, initialize_megatron
from megatron.core import mpu
from megatron.model import GPTModel
from megatron.training import pretrain
from megatron.utils import get_ltor_masks_and_position_ids

# Initialize Megatron with arguments
args = initialize_megatron()

# Create model with tensor and pipeline parallelism
model = GPTModel(
    num_tokentypes=0,
    parallel_output=True,
    pre_process=mpu.is_pipeline_first_stage(),
    post_process=mpu.is_pipeline_last_stage()
)

# Setup optimizer and learning rate scheduler
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)

# Start pretraining
pretrain(train_valid_test_datasets_provider,
         model_provider,
         forward_step,
         args_defaults={'tokenizer_type': 'GPT2BPETokenizer'})