Skip to content
  • ✗ AIP1 Isambard-AI Phase 1 unsupported
  • ✓ AIP2 Isambard-AI Phase 2 supported
  • ✗ I3 Isambard 3 unsupported
  • ✗ BC5 BlueCrystal 5 unsupported

LLM Pre-training

Abstract

This tutorial introduces pre-training a large language model (LLM) across multiple nodes on Isambard-AI. It provides a reference for pre-training using the Megatron-LM and Megatron Bridge frameworks from NVIDIA.

Prerequisites

We welcome people from all domain backgrounds that have experience:

  • Training AI models with PyTorch.
  • Serving AI models with vLLM.
  • Using Apptainer containers.

High Performance Computing (Slurm) knowledge is not required.

You will need a HuggingFace account.

Learning Objectives

The learning objectives of this tutorial are as follows:

  1. Be able to pre-train a Meta Llama 3.1 8B LLM in FP8 reduced precision using the NVIDIA Megatron Bridge framework.
  2. Understand the role that pre-training plays in the training process of an LLM.
  3. Understand how and why natural language is converted into tokens.
  4. Export an early model checkpoint into the HuggingFace format, then obtain a chat completion using vLLM.

Tutorial Contents

Introduction

What is an LLM?

Almost all large language models are based on an autoregressive deep neural network architecture called the Transformer. An LLM takes a stream of input tokens and outputs a single token, which is the most probable next token. To generate further tokens, the LLM is fed the original input plus each token it has already produced, repeating until the output is complete; an autoregressive process.

Diagram showing flow of LLM inference, including tokenization of input text, embedding, the forward pass and detokenization

Tokens

Tokens are the basic units of text an LLM operates on, and typically represent a word or a few characters. Text is converted into tokens using a tokenizer, which works at the level of raw bytes. In ASCII text, such as plain English, each character corresponds to exactly one byte. This is not true of all text; characters outside the ASCII range, such as accented letters, emoji, or non-Latin scripts, are commonly encoded as several bytes each, for example in UTF-8.

Without tokenization, your vocabulary is either individual bytes or complete words. Individual bytes allow for a very small vocabulary but provide very little indication to the model of the relation between them. Complete words would lead to an intractable vocabulary size with lots of rare words that may not appear in the training data. Subword tokenization combines both approaches, allowing novel words to be produced and the relation between different tokens to be learned.

Byte Pair Encoding (BPE) is the algorithm most commonly used to build a subword vocabulary. Starting from a vocabulary containing only the individual bytes, BPE repeatedly scans the training data and merges whichever pair of adjacent tokens occurs most frequently, adding the merged pair to the vocabulary as a new token. This merging step is applied recursively, so a token created in one pass can itself be merged with a neighbour in the next, until the vocabulary reaches its target size. As a result, byte sequences that occur often in the training data end up represented as single tokens, while rarer sequences remain split into smaller pieces.

What does tokenization look like?

The subword tokenization of "Supercomputing" may produce the following subwords: ["Su", "per", "comp", "ut", "ing"]. Each of these pieces are then translated to the corresponding ID in the vocabulary e.g. [1234, 5465, 876, 43567, 429]. The vocabulary of a typical LLM contains in excess of 100,000 tokens.

You can go to the Tokenizer Playground website to see how different models tokenize text.

After tokenization is complete, and a stream of tokens is produced, an embedding vector is obtained for each token. This involves fetching, for each token ID, the corresponding row from an embedding matrix. These vectors are then concatenated to form the model context, a tensor which dimensions corresponding to the model context length and the model's dimension. The model's dimension varies by model and plays a significant role in determining the number of parameters of a model.

We can then run a forward pass on the model with this input tensor. This outputs a single token which continues the stream of input tokens. This step is then repeated to produce a fixed number of output tokens, or until the special <EOS>, or end of sequence, token is returned. In each repetition, or model forward pass, the token output by the model is appended to the context and the first token of the context is removed.

Transformer Architecture

The forward pass mentioned above is performed by a Transformer, the neural network architecture behind almost all modern LLMs. A Transformer is built from a stack of identical blocks, each combining two components: a self-attention layer and a feed-forward layer.

Self-attention is the key ingredient that allows LLMs to build world knowledge across giant datasets. For every token in the context, it lets the model look at every other token and decide how much each one should influence the current one, all in a single, parallelisable step. This is how a model determines that "it" refers to "the supercomputer" three sentences earlier, or that "Clifton" means our command line tool rather than the area of Bristol, purely from the surrounding tokens. No explicit grammar rules or hand-written logic are involved; the model learns these relationships entirely from data during training.

Stacking many of these blocks lets the model build up increasingly abstract representations of the input, from simple word relationships in early layers to complex reasoning-like patterns in later ones.

How are LLMs trained?

Training a large language model is not a single process but a pipeline of many distinct phases. Most modern LLMs, broadly, follow a similar pipeline, involving pre-training, fine-tuning, and alignment. The complete, detailed pipeline is as follows:

  1. Pre-training

    The model is initialised with random weights and trained on a large corpus of text with the objective of accurately predicting the next token. This is the most computationally demanding step, in terms of compute and data, in the LLM training pipeline; the result is a general-purpose foundation model with broad language understanding and knowledge learned by its weights.

  2. Mid-training

    Towards the end of pre-training, the data mix is shifted towards smaller quantities of higher-quality, more curated content such as textbooks, code, mathematics, and reasoning traces, often combined with a decay of the learning rate (sometimes called annealing). This stage disproportionately shapes the model's downstream capabilities and is increasingly treated as a distinct phase rather than a tail-end of pre-training.

  3. Supervised fine-tuning

    This step adapts the model to follow instructions. The foundation model is trained on a curated dataset of prompt-response pairs to allow it to learn how to produce helpful, well-structured answers. This transforms the foundation model into something that more closely resembles a chatbot.

  4. Preference tuning / alignment

    This step aligns the model's behaviour with intended values and expectations. Reinforcement Learning from Human Feedback (RLHF) rewards the model for producing outputs rated highly by human reviewers; Direct Preference Optimisation (DPO) is a simpler alternative that learns directly from preference pairs without a separate reward model. Reinforcement Learning from AI Feedback (RLAIF) substitutes another model for the human rater, and Reinforcement Learning from Verifiable Rewards (RLVR), used to train reasoning models, rewards outputs that can be automatically checked for correctness, such as passing unit tests or producing the right answer to a mathematics problem.

  5. Continual pre-training (optional)

    An already-trained model can undergo further next-token-prediction training on new data to add a specialist domain (medical, legal, code), a new language, or more recent knowledge. Unlike the stages above, this typically happens after a model has already been deployed or released.

  6. Evaluation and iteration

    This is ongoing throughout every stage above to ensure the training process is as effective as possible, spanning benchmark performance, capability probes, safety evaluations, and human preference studies.

In this tutorial, we focus exclusively on the first step; pre-training.

Technical challenges

LLM training is distinctly different to conventional deep learning due to the sheer scale of computation involved. Updating those parameters from a meaningful quantity of data requires trillions of floating-point operations. This necessitates distributing the workload across many GPUs simultaneously, careful management of memory, and specialised frameworks designed to maximise hardware utilisation. Training at this scale introduces several engineering challenges that don't arise in smaller workloads:

  • Memory pressure

    The model weights, activations, gradients, and optimiser states must all fit within GPU memory simultaneously. For an 8B parameter model in full FP32 precision, the weights alone occupy ~29.8 GiB before accounting for anything else.

  • Communication overhead

    When training across multiple GPUs, the gradient updates computed on each device must be synchronised. The cost of this collective communication can dominate training time if not managed carefully.

  • Numerical stability

    Reduced precision formats accelerate computation but introduce approximation errors. Keeping training stable under these conditions requires specific techniques such as loss scaling and mixed-precision strategies.

  • Data throughput

    The model must never be left waiting for data. Efficient data loading, shuffling, and tokenization pipelines are essential.

Pre-training requires orders of magnitude more data and compute than fine-tuning. It also requires training from randomly initialised weights rather than starting from an existing checkpoint, which means the model must learn everything from first principles. This makes pre-training sensitive to hyperparameter choices, learning rate schedules, and data quality in ways that fine-tuning typically is not.

In practice, most organisations do not train frontier models from scratch due to the prohibitive compute cost. However, pre-training at smaller scale (e.g. on domain-specific corpora, or on a smaller dataset as a learning exercise) is valuable for research, for understanding model behaviour, and for producing domain-specialised base models. The workflow you will follow in this tutorial, pre-training Llama 3.1 8B on a sample of FineWeb using Megatron Bridge, is representative of frontier model pre-training.

Meta Llama 3.1 8B

Llama 3.1 8B is the smallest model in Meta's Llama 3.1 family of multilingual, pre-trained LLMs released in July 2024. It has approximately 8 billion parameters and was trained on over 15 trillion tokens of multilingual code and text. As this is a pre-trained model, it is designed to be adapted for a variety of natural language generation tasks. The Instruct variants of the Llama 3.1 family of models have been adapted for assistant-like chat.

Despite being over two years old, the underlying model architecture is representative of the current frontier models. Therefore, the steps involved in LLM training as well as the tools that we will use in this tutorial align with what is used by frontier AI laboratories today.

Tools and frameworks

NVIDIA Megatron-LM and Megatron Bridge

Megatron-LM is NVIDIA's open-source framework for training large transformer models at scale. It is used internally at NVIDIA and by many HPC centres and research labs to train frontier models. Megatron-LM is designed from the ground up for efficiency on NVIDIA hardware, and implements several parallelism strategies that are essential when a model or its training state does not fit within the memory of a single GPU.

Megatron Bridge is an NVIDIA library containing utilities and adapter code that connects Megatron-LM's training infrastructure to the broader HuggingFace ecosystem. It provides the ability to load HuggingFace-format datasets and to convert model weights between Megatron's internal sharded format and the HuggingFace checkpoint format. This is important in practice because it means you can pre-train with Megatron-LM's efficiency while still producing checkpoints compatible with the broader ecosystem for evaluation and deployment.

The training infrastructure provided by Megatron-LM is called Megatron-Core. For simplicity, from this point on we will collectively refer to these libraries as Megatron.

HuggingFace FineWeb

FineWeb is a large, carefully curated English web text dataset produced by HuggingFace. The full dataset contains approximately 18.5 trillion tokens; HuggingFace also releases FineWeb-Edu, a subset filtered for educational content quality. It was constructed from Common Crawl snapshots with extensive filtering to remove low-quality, duplicate, and harmful content. Common Crawl is a massive free and open corpus of web crawl data captured in monthly snapshots. For the purposes of this tutorial, we use a 10-billion-token sample of FineWeb designed specifically for small-scale pre-training experiments. Using a sample rather than the full corpus makes the experiment tractable on a limited compute budget while still being representative of real pre-training data characteristics.

To prepare this dataset for training, it must be tokenized and converted into Megatron's binary indexed dataset format (.bin / .idx file pairs). Megatron-LM includes preprocessing scripts for this purpose. The Llama 3.1 tokenizer uses a byte-pair encoding (BPE) vocabulary of 128,000 tokens plus an additional 256 special tokens.

Reduced-precision training

NVIDIA Hopper GPUs, such as those in Isambard-AI, introduced native support for FP8 which roughly halves the memory bandwidth requirements compared to BF16 and can deliver a near-doubling of matrix multiplication throughput. The trade-off with FP8 is its very limited numerical range, but the NVIDIA Transformer Engine has integrated mechanisms to ensure that training is numerically stable.

Mixed-precision training in the context of LLMs typically means maintaining a full-precision (FP32 or BF16) master copy of the model weights for the optimiser state and weight updates, while performing the forward and backward passes in lower precision for speed. The NVIDIA Transformer Engine library, which is integrated into Megatron-LM, automates much of this for BF16 and FP8 workloads.

In this tutorial, we will be pre-training Llama 3.1 8B in FP8 rather than BF16 to benefit from the significantly reduced per-step time, which allows larger experiments to be achievable within a given compute allocation.

Walkthrough

Now that we have covered the required prior knowledge, we can begin setting up our LLM pre-training experiment.

Authorise with HuggingFace

This step is necessary to complete this tutorial.

The Llama family of models are gated on HuggingFace. Users need to agree to the terms of the licensing agreement to access the model weights, tokenizer and other configuration files.

  1. Go to the Meta Llama 3.1 8B model card and request access.

    You will need to wait for your request to be approved. This may take some time.

  2. Create a HuggingFace access token with the permission:

    Read access to contents of all public gated repos you can access

  3. Install uv and run:

    uvx hf auth login

Set up working environment

We set the WORK_PATH environment variable, referenced throughout the scripts used in this tutorial, to define our working directory in the /projects directory. We also create this directory and move into it.

user@login43:~> export WORK_PATH=${PROJECTDIR}/${USER}/llm-pre-training
user@login43:~> mkdir -p ${WORK_PATH} && cd ${WORK_PATH}
user@login43:/projects/PROJECT/USER/llm-pre-training>

Next, we will move on to downloading the FineWeb 10BT sample dataset, processing it into the format expected by the Megatron frameworks using the Llama 3.1 tokenizer.

We have already carried out these steps and have stored the dataset and tokenizer in a location accessible to all users. To use the prepared dataset and tokenizer, click the Quick setup button.

Downloading NVIDIA Megatron-LM and Megatron Bridge

We clone both the Megatron-LM and Megatron Bridge Git repositories. We specifically target the 26.04-alpha.rc2 releases which match the NeMo NGC image that we will use to run our pre-training workload.

$ git clone --depth 1 --branch 26.04-alpha.rc2 https://github.com/NVIDIA/Megatron-LM.git ${WORK_PATH}/Megatron-LM
$ git clone --depth 1 --branch 26.04-alpha.rc2 https://github.com/NVIDIA-NeMo/Megatron-Bridge ${WORK_PATH}/Megatron-Bridge

Download the NeMo NGC image - 30 minutes

Next, we download the NVIDIA NeMo 26.04 NGC image that we will run with Apptainer. NeMo is NVIDIA's open-source, end-to-end training framework for large language models.

pull_nemo_image.sh
#!/bin/bash
#SBATCH --job-name=pull-nemo
#SBATCH --nodes=1
#SBATCH --gpus=4
#SBATCH --time=00:30:00
#SBATCH --exclusive
#SBATCH --output=out/%x.%j.out

export WORK_PATH=${WORK_PATH:-${HOME}/llm-pre-training}

srun apptainer build ${WORK_PATH}/nemo.sif docker://nvcr.io/nvidia/nemo:26.04

You can use curl to directly download the file:

$ curl -fsSL -o ${WORK_PATH}/pull_nemo_image.sh https://docs.isambard.ac.uk/user-documentation/tutorials/llms/pre-training/pull_nemo_image.sh

Submit the job to pull the NeMo NGC image as an Apptainer .sif image:

$ sbatch ${WORK_PATH}/pull_nemo_image.sh
Submitted batch job JOBID

Alternatively, as we have already downloaded this image, you can run the following command to copy it to your working directory:

$ cp /projects/public/brics/tutorials/llms/pre-training/nemo.sif ${WORK_PATH}/nemo.sif

Download the HuggingFace FineWeb 10BT sample dataset - 35 minutes

The process of acquiring the FineWeb dataset and processing it into the format expected by Megatron is split into three steps:

  1. Downloading the dataset from HuggingFace as a set of Parquet files.
  2. Converting the dataset from many Parquet files into a single JSONL file.
  3. Tokenizing the JSONL file into .bin and .idx files using Megatron's preprocessing script.
Data preparation script
data_preparation.sh
#!/bin/bash
#SBATCH --job-name=llama-data-prep
#SBATCH --nodes=1
#SBATCH --gpus=1
#SBATCH --time=01:00:00
#SBATCH --exclusive
#SBATCH --output=out/%x.%j.out

set -euo pipefail

# Check environment variables
export WORK_PATH=${WORK_PATH:-${PROJECTDIR}/${USER}/llm-pre-training}
export FINEWEB_PATH=${FINEWEB_PATH:-${WORK_PATH}/fineweb}
export FINEWEB_SUBSET=${FINEWEB_SUBSET:-sample/10BT}
export DATASET_PATH=${DATASET_PATH:-${WORK_PATH}/fineweb/processed}
export TOKENIZER_PATH=${TOKENIZER_PATH:-${WORK_PATH}/tokenizer}
export BRICS_TUTORIAL_PATH=/projects/public/brics/tutorials/llms/pre-training
cd ${WORK_PATH}

# Convert dataset to JSONL format
# - The // script frontmatter is used by `uv` to create a temporary 
#   virtual environment with the required dependencies.
cat > download_dataset.py << 'PY'
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "huggingface_hub==1.22.0",
#     "datatrove==0.8.0",
#     "pyarrow==24.0.0",
#     "regex==2026.6.28",
#     "orjson==3.11.9"
# ]
# ///

import os

from argparse import ArgumentParser

from datatrove.executor import LocalPipelineExecutor
from datatrove.pipeline.readers import ParquetReader
from datatrove.pipeline.filters import LambdaFilter
from datatrove.pipeline.writers import JsonlWriter

from huggingface_hub import snapshot_download


def main():
    parser = ArgumentParser()
    parser.add_argument("--dataset", type=str, required=True)
    parser.add_argument("--subset", type=str, required=True)
    parser.add_argument("--output", type=str, required=True)
    args = parser.parse_args()

    # Download FineWeb sample dataset
    os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"

    snapshot_download(
        repo_id="HuggingFaceFW/fineweb",
        repo_type="dataset",
        allow_patterns=f"{args.subset}/*",
        local_dir=args.dataset,
    )

    # Convert Parquet files to JSONL format using datatrove
    pipeline_exec = LocalPipelineExecutor(
        pipeline=[
            ParquetReader(f"{args.dataset}/{args.subset}/"),
            JsonlWriter(args.output, compression=None)
        ],
        tasks=64,
        workers=os.cpu_count(),
    )
    pipeline_exec.run()

if __name__ == "__main__":
    main()
PY

uv run download_dataset.py \
    --dataset ${FINEWEB_PATH} \
    --subset ${FINEWEB_SUBSET} \
    --output ${DATASET_PATH}

# Merge JSONL files into a single file and delete intermediate files
cat ${DATASET_PATH}/*.jsonl > ${DATASET_PATH}/merged.jsonl
find ${DATASET_PATH} -name "000*.jsonl" -exec rm {} +

# Download tokenizer and tokenize the dataset
# - Using a script provided by Megatron-LM
# - This needs to be run inside the NeMo NGC container
# - Generates the .bin and .idx files required by Megatron-LM
apptainer run \
    --nv \
    --mount type=bind,source=${WORK_PATH},destination=${WORK_PATH} \
    --mount type=bind,source=${BRICS_TUTORIAL_PATH},destination=${BRICS_TUTORIAL_PATH} \
    ${WORK_PATH}/nemo.sif \
    bash -c "
        # Download the tokenizer
        hf download meta-llama/Llama-3.1-8B \
            tokenizer.json tokenizer_config.json special_tokens_map.json \
            --local-dir ${TOKENIZER_PATH}

        # Tokenize the dataset using Megatron-LM's preprocessing script
        python Megatron-LM/tools/preprocess_data.py \
            --input ${DATASET_PATH}/merged.jsonl \
            --output-prefix ${DATASET_PATH}/data \
            --tokenizer-type HuggingFaceTokenizer \
            --tokenizer-model ${TOKENIZER_PATH} \
            --append-eod \
            --workers 16
    "

You can use curl to directly download the file:

$ curl -fsSL -o ${WORK_PATH}/data_preparation.sh https://docs.isambard.ac.uk/user-documentation/tutorials/llms/pre-training/data_preparation.sh

Then, we submit our data preparation job:

$ sbatch ${WORK_PATH}/data_preparation.sh
Submitted batch job JOBID

This will take some time. You can monitor the progress by viewing the logs at ${WORK_PATH}/out/llama-data-prep.JOBID.out.

You can verify that the job ran successfully by running the following command:

$ ls -lh ${DATASET_PATH}/
total 87G
-rw-r--r-- 1 user group  38G Jul  6 13:52 data_text_document.bin
-rw-r--r-- 1 user group 284M Jul  6 13:52 data_text_document.idx
-rw-r--r-- 1 user group  50G Jul  6 13:51 merged.jsonl

You should see the .bin and .idx files, which contain our tokenized dataset.

To use the pre-downloaded, pre-processed dataset and tokenizer, run the following command:

$ ln -s /projects/public/brics/tutorials/llms/pre-training/* ${WORK_PATH}/

This creates symlinks in your working directory which point to the pre-downloaded files in the /projects/public/brics directory.

Run pre-training - 40 minutes

We are almost ready to begin pre-training Llama 3.1 8B in FP8 reduced precision.

Megatron configuration

First, we set our required settings in the Megatron configuration file. This file specifies a variety of settings that describe our dataset, our model, checkpoint settings, logging and more.

Training configuration
config.yaml
dataset:
  data_path: ./fineweb/processed/data_text_document
  seq_length: 8192
train:
  train_iters: 100
  global_batch_size: 128
  micro_batch_size: 1
  eval_iters: 10
  eval_interval: 50
optimizer:
  lr: 0.00015
  min_lr: 0.00001
  weight_decay: 0.1
  use_distributed_optimizer: true
scheduler:
  lr_warmup_iters: 5
  lr_decay_style: cosine
checkpoint:
  save: ./checkpoints/llama31_8b
  load: ./checkpoints/llama31_8b
  save_interval: 50
ddp:
  overlap_param_gather: true
  overlap_grad_reduce: true
model:
  seq_length: 8192
  tensor_model_parallel_size: 1
  pipeline_model_parallel_size: 1
  context_parallel_size: 1
  recompute_activation: true
  recompute_granularity: full
  recompute_method: uniform
  recompute_num_layers: 1
  use_flash_attn: true
tokenizer:
  tokenizer_type: HuggingFaceTokenizer
  tokenizer_model: ./tokenizer
logger:
  log_interval: 10
  tensorboard_dir: ./tensorboard/llama31_8b
rng:
  seed: 2604

Run the following command to download the configuration file:

$ curl -fsSL -o ${WORK_PATH}/config.yaml https://docs.isambard.ac.uk/user-documentation/tutorials/llms/pre-training/config.yaml

This configuration is designed to pre-train for approximately 40 minutes across two nodes on Isambard-AI. We use a sequence length of 8192 in our configuration - this corresponds to the size of the model's context. This is very small, but is used to be able to fit our pre-training job across only two nodes. Current models of a similar size to Llama 3.1 8B have context lengths between 128,000-256,000 tokens, but this requires many more nodes to be used during training to provide enough GPU VRAM to store such large tensors.

Another important setting is the number of training iterations, train_iters, which determines how many global batches of tokens we train the model on. In our case, we opt for 100 iterations of 128 sequences, each composed of 8192 tokens.

You can find more information on the available configuration options in the Megatron Bridge documentation.

Scaling up pre-training

Consider the relevance of the settings within the configuration file.

  1. What does context_parallel_size: 1 mean and when would you increase it?
  2. What is the impact of increasing the number of iterations of learning rate warmup (lr_warmup_iters)?
  3. What are the impacts of increasing sequence length (seq_length)?

The Megatron Bridge performance tuning guide describes how to determine the optimal configuration for the model architecture and hardware.

Answers
  1. The sequence dimension is split across GPUs. TP splits the hidden dimension (sharding tensors), PP splits the layers. CP is needed when sequence length (model context) is extremely large and even with FlashAttention and activation recomputation, activations won't fit on a single GPU. This usually applies when context goes to 128k tokens and beyond.
  2. As weights are randomly-initialised and the optimiser and normalisation statistics are unstable, the full peak LR is likely to be too high and will cause issues. If we increase lr_warmup_iters we will have a more stable LR which will prevent gradient explosion and loss spikes. The number of LR warmup iters should be ~0.1-1% of training iters for pre-training and 3-10% for fine-tuning. Too much warmup wastes compute and can lead to a suboptimal initial LR.
  3. The model will have a longer context which allows for a larger 'history' during inference and to learn long-range dependencies during training, but training becomes more difficult due to the significantly larger VRAM usage. As Attention is O(N^2) in sequence length, doubling sequence length quadruples FLOPs. Training requires more GPUs and increased context parallelism as activation memory is quadratic, KV cache linearly increases. If you are hardware limited, with a fixed number of GPUs, increasing context parallelism means that you need to reduce data parallelism. Therefore, training on the same number of tokens will take twice as long. You will also need to adjust learning rate and batch size if you increase sequence length, and the documents in your training data need to be large enough to fill the context for maximum efficiency.

Training script

We will use the Llama 3.1 8B pre-training recipe Python script from NVIDIA's Megatron repository. This script is designed exclusively for pre-training Llama 3.1 8B. It simply loads our Megatron configuration file and runs the pre-defined training recipe.

Run the following commands to download the script and configure it for Llama 3.1 8B:

$ curl -fsSL -o ${WORK_PATH}/train.py https://raw.githubusercontent.com/NVIDIA-NeMo/Megatron-Bridge/refs/tags/26.04-alpha.rc2/tutorials/recipes/llama/02_pretrain_with_yaml.py
$ sed -i 's/llama32_1b/llama31_8b/g' ${WORK_PATH}/train.py

This script is designed for Llama 3.2 1B but works with Llama 3.1 8B, so we use sed to correct the training script.

Slurm job script

Our Slurm job script uses torchrun and PMI2 to distribute pre-training across two Isambard-AI nodes. This script starts four Apptainer containers on each node, each running the NeMo 26.04 NGC image. We use a custom entrypoint adapt.sh script, in /projects/public/brics/tutorials/llms/pre-training, to configure NCCL for RDMA over the Isambard-AI Slingshot high-speed network within our container, which is set up with CUDA 13.1 forward compatibility. For more information on multi-node training on Isambard-AI, see our Distributed PyTorch tutorial.

Training job script
pretrain.sh
#!/bin/bash
#SBATCH --job-name=llama-pretrain
#SBATCH --nodes=2
#SBATCH --gpus=8
#SBATCH --time=01:00:00
#SBATCH --exclusive
#SBATCH --output=out/%x.%j.out

module reset
module load brics/nccl
module load brics/apptainer-multi-node
module list

SERVER_ADDRESS=$(dig +short ${HOSTNAME}-hsn0)
HEAD_NODE=$(scontrol show hostnames $SLURM_NODELIST | head -n1)
WORKER_NODES=$(scontrol show hostnames $SLURM_NODELIST | tail -n+2)
export HEAD_NODE_IP=$(dig +short ${HEAD_NODE})
export MASTER_ADDR=$HEAD_NODE_IP
export MASTER_PORT=29600
export WORLD_SIZE=$SLURM_GPUS

export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}

WORK_PATH=${WORK_PATH:-${PROJECTDIR}/${USER}/llm-pre-training}
DATASET_PATH=${DATASET_PATH:-${WORK_PATH}/fineweb/processed}
BRICS_TUTORIAL_PATH=/projects/public/brics/tutorials/llms/pre-training

cd ${WORK_PATH}

srun --mpi=pmi2 \
    apptainer run --nv \
        --mount type=bind,source=${BRICS_TUTORIAL_PATH},destination=${BRICS_TUTORIAL_PATH} \
        --mount type=bind,source=${WORK_PATH},destination=${WORK_PATH} \
        --mount type=bind,source=${WORK_PATH}/Megatron-Bridge,destination=/opt/Megatron-Bridge \
        --mount type=bind,source=${WORK_PATH}/Megatron-LM,destination=/opt/Megatron-Bridge/3rdparty/Megatron-LM \
        ${WORK_PATH}/nemo.sif \
        ${BRICS_TUTORIAL_PATH}/adapt.sh \
        torchrun \
        --nproc_per_node=$SLURM_GPUS_ON_NODE \
        --nnodes=$SLURM_NNODES \
        --node_rank=\$SLURM_PROCID \
        --master_addr=$HEAD_NODE_IP \
        --master_port=29600 \
        ${WORK_PATH}/train.py \
        --config-file ${WORK_PATH}/config.yaml

Run the following command to download the job script:

$ curl -fsSL -o ${WORK_PATH}/pretrain.sh https://docs.isambard.ac.uk/user-documentation/tutorials/llms/pre-training/pretrain.sh

Launching our job

Finally, we can launch our Llama 3.1 8B pre-training job with Slurm:

$ sbatch ${WORK_PATH}/pretrain.sh
Submitted batch job JOBID
For workshop facilitators

This job takes around 40 minutes to complete, the longest wait in this tutorial. While it runs, this is a good point to take a break. Alternatively, you can work through the "Scaling up pre-training" questions from the Megatron configuration section above or preview the token-count comparison in Reflections below.

You can view the logs in the ${WORK_PATH}/out/llama-pretrain.JOBID.out file.

Over the first few minutes, PyTorch will establish communication between both nodes with NCCL before proceeding to dataset loading and indexing. Once this is complete, the model is sharded across the GPUs and nodes and training begins. The training process will begin between 5-10 minutes after the job starts.

While training is in progress, you will see logs like this:

Step Time : 12.64s GPU utilization: 533.6MODEL_TFLOP/s/GPU
Number of parameters in transformer layers in billions:  6.98
Number of parameters in embedding layers in billions: 1.05
Total number of parameters in billions: 8.03
Number of parameters in most loaded shard in billions: 8.0305
Theoretical memory footprints: weight and optimizer=68926.57 MB
[Rank 1] (after 10 iterations) memory (GB) | mem-allocated-gigabytes: 61.819 | mem-active-gigabytes: 61.819 | mem-inactive-gigabytes: 1.7522 | mem-reserved-gigabytes: 80.558 | mem-max-allocated-gigabytes: 78.473 | mem-max-active-gigabytes: 78.725 | mem-max-inactive-gigabytes: 6.2369 | mem-max-reserved-gigabytes: 80.558 | mem-alloc-retires: 0 | mem-allocated-count: 2656
 [2026-07-07 08:30:40] iteration       10/     100 | consumed samples:         1280 | elapsed time per iteration (ms): 12643.9 | learning rate: 1.490453E-04 | global batch size:   128 | lm loss: 1.580069E+01 | loss scale: 1.0 | grad norm: 37.247 | number of skipped iterations:   0 | number of nan iterations:   0 |
[Rank 0] (after 10 iterations) memory (GB) | mem-allocated-gigabytes: 61.819 | mem-active-gigabytes: 61.819 | mem-inactive-gigabytes: 1.7459 | mem-reserved-gigabytes: 80.526 | mem-max-allocated-gigabytes: 78.473 | mem-max-active-gigabytes: 78.708 | mem-max-inactive-gigabytes: 6.0628 | mem-max-reserved-gigabytes: 80.526 | mem-alloc-retires: 0 | mem-allocated-count: 2845
Step Time : 11.40s GPU utilization: 591.7MODEL_TFLOP/s/GPU
 [2026-07-07 08:32:34] iteration       20/     100 | consumed samples:         2560 | elapsed time per iteration (ms): 11401.2 | learning rate: 1.415632E-04 | global batch size:   128 | lm loss: 1.127280E+01 | loss scale: 1.0 | grad norm: 5.693 | number of skipped iterations:   0 | number of nan iterations:   0 |
Step Time : 11.38s GPU utilization: 593.1MODEL_TFLOP/s/GPU
 [2026-07-07 08:34:28] iteration       30/     100 | consumed samples:         3840 | elapsed time per iteration (ms): 11375.6 | learning rate: 1.274097E-04 | global batch size:   128 | lm loss: 8.378448E+00 | loss scale: 1.0 | grad norm: 7.799 | number of skipped iterations:   0 | number of nan iterations:   0 |
Step Time : 11.39s GPU utilization: 592.3MODEL_TFLOP/s/GPU
 [2026-07-07 08:36:22] iteration       40/     100 | consumed samples:         5120 | elapsed time per iteration (ms): 11389.9 | learning rate: 1.081187E-04 | global batch size:   128 | lm loss: 8.011566E+00 | loss scale: 1.0 | grad norm: 2.307 | number of skipped iterations:   0 | number of nan iterations:   0 |
Step Time : 11.40s GPU utilization: 591.5MODEL_TFLOP/s/GPU
 [2026-07-07 08:38:16] iteration       50/     100 | consumed samples:         6400 | elapsed time per iteration (ms): 11404.9 | learning rate: 8.578056E-05 | global batch size:   128 | lm loss: 7.857731E+00 | loss scale: 1.0 | grad norm: 1.215 | number of skipped iterations:   0 | number of nan iterations:   0 |
INFO:megatron.core.timers:(min, max) time across ranks (ms):
    evaluate .......................................: (39850.74, 39852.42)
----------------------------------------------------------------------------------------------
 validation loss at iteration 50 | lm loss value: 7.822796E+00 | lm loss PPL: 2.496877E+03 |
----------------------------------------------------------------------------------------------
saving checkpoint at iteration      50 to ./checkpoints/llama31_8b in torch_dist format
Storing distributed optimizer sharded state of type dp_reshardable
  successfully saved checkpoint from iteration      50 to ./checkpoints/llama31_8b [ t 1/1, p 1/1 ]
INFO:megatron.core.timers:(min, max) time across ranks (ms):
    save-checkpoint ................................: (33266.91, 33267.06)

The output above shows a log message every 10th iteration containing the training status at that iteration. Here, we can see the loss is decreasing at every iteration showing that our model is indeed learning. We can also see a throughput figure of ~590 TFLOP/s/GPU. GH200s have a maximum FP8 throughput of 1979 TFLOPs, but as we are using multiple forms of parallelism and a relatively small model, compared to the VRAM capacity of each GH200, we see a much reduced throughput.

Based on our Megatron configuration file, model checkpoints will be saved every 50 iterations in ${WORK_PATH}/checkpoints/llama31_8b.

$ ls -a ${WORK_PATH}/checkpoints/llama31_8b/
.  ..  iter_0000050  latest_checkpointed_iteration.txt  latest_train_state.pt

If you cancel your job, or it terminates at the job time limit, re-running it will automatically continue training from the last checkpoint.

Export to the HuggingFace format - 10 minutes

A key advantage of Megatron Bridge over Megatron-LM is its native support for converting models between the Megatron format and the HuggingFace format.

We will now export our last model checkpoint into the HuggingFace format, readable by libraries such as vLLM and transformers.

Script to convert model from Megatron to HuggingFace format
convert_to_hf.sh
#!/bin/bash
#SBATCH --job-name=llama-convert-hf
#SBATCH --nodes=1
#SBATCH --gpus=1
#SBATCH --time=00:30:00
#SBATCH --exclusive
#SBATCH --output=out/%x.%j.out

module reset
module load brics/nccl
module load brics/apptainer-multi-node
module list

WORK_PATH=${WORK_PATH:-${PROJECTDIR}/${USER}/llm-pre-training}
DATASET_PATH=${DATASET_PATH:-${WORK_PATH}/fineweb/processed}
BRICS_TUTORIAL_PATH=/projects/public/brics/tutorials/llms/pre-training

srun apptainer run --nv \
    --mount type=bind,source=${BRICS_TUTORIAL_PATH},destination=${BRICS_TUTORIAL_PATH} \
    --mount type=bind,source=${WORK_PATH},destination=${WORK_PATH} \
    --mount type=bind,source=${WORK_PATH}/Megatron-Bridge,destination=/opt/Megatron-Bridge \
    --mount type=bind,source=${WORK_PATH}/Megatron-LM,destination=/opt/Megatron-Bridge/3rdparty/Megatron-LM \
    ${WORK_PATH}/nemo.sif \
    ${BRICS_TUTORIAL_PATH}/adapt.sh \
    python /opt/Megatron-Bridge/examples/conversion/convert_checkpoints.py export \
    --hf-model meta-llama/Llama-3.1-8B \
    --megatron-path ${WORK_PATH}/checkpoints/llama31_8b/iter_0000100 \
    --hf-path ${WORK_PATH}/checkpoints/llama31_8b/out

You can use curl to directly download the file:

$ curl -fsSL -o ${WORK_PATH}/convert_to_hf.sh https://docs.isambard.ac.uk/user-documentation/tutorials/llms/pre-training/convert_to_hf.sh

We now launch the Slurm job:

$ sbatch ${WORK_PATH}/convert_to_hf.sh
Submitted batch job JOBID

This will take some time. You can monitor the progress by viewing the logs at ${WORK_PATH}/out/llama-convert-hf.JOBID.out.

When the export process has completed, you should see the following at the end of the log file:

Converting to HuggingFace ━━━━━━━━━━━━━━━━━━━ 100% 0:00:00 (195/195) LlamaBridge
✅ Successfully exported model to: /projects/PROJECT/USER/llm-pre-training/checkpoints/llama31_8b/out
📁 Export structure:
   📄 model-00003-of-00004.safetensors
   📄 model.safetensors.index.json
   📄 model-00001-of-00004.safetensors
   📄 config.json
   📄 model-00004-of-00004.safetensors
   📄 model-00002-of-00004.safetensors
🔍 You can now load this model with:
   from transformers import AutoModelForCausalLM
   model = AutoModelForCausalLM.from_pretrained('/projects/PROJECT/USER/llm-pre-training/checkpoints/llama31_8b/out')

Get a chat completion from the model - 5 minutes

The final step in this tutorial is to use vLLM to load our model, from the HuggingFace .safetensors format, and get it to produce a chat completion.

serve_checkpoint_vllm.sh
#!/bin/bash
#SBATCH --job-name=llama-serve-vllm
#SBATCH --nodes=1
#SBATCH --gpus=1
#SBATCH --time=00:15:00
#SBATCH --exclusive
#SBATCH --output=out/%x.%j.out

module reset
module load brics/nccl
module load brics/apptainer-multi-node
module list

WORK_PATH=${WORK_PATH:-${PROJECTDIR}/${USER}/llm-pre-training}
DATASET_PATH=${DATASET_PATH:-${WORK_PATH}/fineweb/processed}
BRICS_TUTORIAL_PATH=/projects/public/brics/tutorials/llms/pre-training

srun apptainer run --nv \
    --mount type=bind,source=${BRICS_TUTORIAL_PATH},destination=${BRICS_TUTORIAL_PATH} \
    --mount type=bind,source=${WORK_PATH},destination=${WORK_PATH} \
    --mount type=bind,source=${WORK_PATH}/Megatron-Bridge,destination=/opt/Megatron-Bridge \
    --mount type=bind,source=${WORK_PATH}/Megatron-LM,destination=/opt/Megatron-Bridge/3rdparty/Megatron-LM \
    ${WORK_PATH}/nemo.sif \
    ${BRICS_TUTORIAL_PATH}/adapt.sh \
    vllm serve ${WORK_PATH}/checkpoints/llama31_8b/out \
    --served-model-name meta-llama/Llama-3.1-8B \
    --tokenizer ${WORK_PATH}/tokenizer \
    --chat-template ${WORK_PATH}/tokenizer/tokenizer_config.json \
    --max-model-len 8192

wait

You can use curl to directly download the file:

$ curl -fsSL -o ${WORK_PATH}/serve_checkpoint_vllm.sh https://docs.isambard.ac.uk/user-documentation/tutorials/llms/pre-training/serve_checkpoint_vllm.sh

Now, we launch our Slurm job to serve our model using vLLM:

$ sbatch ${WORK_PATH}/serve_checkpoint_vllm.sh
Submitted batch job JOBID

This will take a few minutes. You can monitor the progress by viewing the logs at ${WORK_PATH}/out/llama-serve-vllm.JOBID.out. In this file, you will see the logs produced by vLLM while the server is starting. The vLLM server is ready when you see the Application startup complete log message, as seen below.

(APIServer pid=249008) INFO:     Started server process [249008]
(APIServer pid=249008) INFO:     Waiting for application startup.
(APIServer pid=249008) INFO:     Application startup complete.

Now that our model is being served, we can enter an interactive prompt to get chat completions from it. Replace JOBID with the ID of your serve_checkpoint_vllm.sh job in the command below.

$ srun --overlap --jobid JOBID --gpus=1 --pty apptainer run --nv ${WORK_PATH}/nemo.sif vllm complete --model meta-llama/Llama-3.1-8B
Using model: meta-llama/Llama-3.1-8B
Please enter prompt to complete:
> Bristol Centre for Supercomputing is
Cloud bathroom your downloaded claim quote6 ...

Type CTRL+C to exit the interactive CLI.

Reflections

As you can see, our model is producing unintelligible chat completions due to the limited pre-training that we have done. This model checkpoint was captured after 100 iterations, with each iteration consuming one global batch of 128 samples and each sample being composed of 8192 tokens (the sequence length). Therefore, using the equation defined in the second Megatron-LM paper, our model checkpoint was trained on 128 * 100 * 8192 = ~105M tokens. Comparing this to figure of 15 trillion tokens, stated by Meta, that Llama 3.1 8B was trained on, our model was trained on ~150,000x fewer tokens. Also, it is important to note that this covers only slightly more than 1% of the FineWeb sample dataset that we are using; an already-small pre-training dataset.

The Chinchilla scaling laws suggest that, for compute-optimal training, model size and training tokens should scale at the same rate. The rule of thumb being that a model with N parameters should be pre-trained on approximately 20N tokens. For Llama 3.1 8B, this implies that 8 * 20 = 160 billion tokens would be required. Our pre-training run of 100 iterations and 105 million tokens only covered around 6.25% of that total.

Cautions for further pre-training

In this tutorial, we only pre-trained Llama 3.1 8B on ~1% of FineWeb's smallest sample dataset - you would need to increase train_iters to 10,000 for the model to train over the entirety of this dataset just once. To pre-train a foundation model with a similar level of performance to Meta's official Llama 3.1 8B open weights, you would need to pre-train on the complete FineWeb dataset.

As jobs on Isambard-AI are limited in duration, you should use job dependencies and an appropriate model checkpoint configuration in the Megatron configuration file. Megatron Bridge will automatically continue pre-training from the latest model checkpoint when the same job script is re-launched, but please be aware that model checkpoints use a significant amount of storage; this can explode if saving many checkpoints during a single training run.

Storage cleanup

The data preparation step consumes 116GiB of storage and each model checkpoint is 105GiB. As a result, if you completed the full tutorial walkthrough including dataset preparation you will have used ~360GiB of storage. You may wish to clean up your working directory after you have finished with it.

Conclusion

In this tutorial, you've learned how to pre-train Meta's Llama 3.1 8B large language model with FP8 precision on Isambard-AI. We covered the different stages involved in training an LLM, the role played by pre-training and the challenges involved. Then, we saw how NVIDIA's Megatron family of frameworks provide a robust, accessible way to train LLMs at scale.

After obtaining a sample of HuggingFace's pre-training dataset, FineWeb, we converted the natural language samples into tokens using the Llama 3.1 tokenizer. Then, we began pre-training Llama 3.1 8B in FP8 precision on our dataset. This produced two model checkpoints, at iterations 50 and 100, and you converted the latter checkpoint into the widely compatible HuggingFace format. Finally, you obtained a chat completion from the model; the core purpose of foundation models.

Having completed this pre-training workflow, you now have a reproducible path from raw text data to a foundation Llama 3.1 checkpoint; the next natural step is to validate that checkpoint, fine-tune it on your target task, and export it for deployment with vLLM or HuggingFace-compatible serving.

Resources