How to Use LightEval for Model Evaluation in Nanotron: A Complete Technical Guide

Nanotron integrates the LightEval benchmark suite through three core components—LightEvalConfig, LightEvalRunner, and SLURM-based job submission—to automate large language model evaluation at configurable intervals directly from training checkpoints.

The huggingface/nanotron repository provides native support for LightEval, a lightweight wrapper around the Hugging Face lighteval library that enables seamless benchmarking of LLMs during distributed training. This integration allows you to evaluate models on standard tasks like MMLU, ARC, and HellaSwag without manually managing checkpoint downloads or parallel execution configurations.

Understanding the LightEval Architecture in Nanotron

The evaluation workflow centers on three interconnected components that handle configuration, orchestration, and execution.

LightEvalConfig Dataclass

The LightEvalConfig dataclass, defined in [src/nanotron/config/lighteval_config.py](https://github.com/huggingface/nanotron/blob/main/src/nanotron/config/lighteval_config.py), stores all evaluation-related parameters including parallelism settings, SLURM allocation, task definitions, and logging preferences. This configuration object computes resource requirements automatically based on the data parallelism (dp), pipeline parallelism (pp), and tensor parallelism (tp) factors specified in the parallelism field.

LightEvalRunner Orchestration

The LightEvalRunner class in [src/nanotron/eval/one_job_runner.py](https://github.com/huggingface/nanotron/blob/main/src/nanotron/eval/one_job_runner.py) serves as the central orchestrator. It implements the eval_single_checkpoint method (lines 27-74) which checks the current training step against the eval_interval threshold and triggers evaluation jobs only when appropriate. This prevents unnecessary computational overhead by ensuring benchmarks run at specified milestones rather than every step.

SLURM Job Generation

When evaluation is triggered, the runner invokes run_slurm_one_job (also in [src/nanotron/eval/one_job_runner.py](https://github.com/huggingface/nanotron/blob/main/src/nanotron/eval/one_job_runner.py)) to construct and submit a fully-featured SLURM script. This function calculates the total GPU requirements using total_gpus_needed = dp * pp * tp, generates the batch submission header, handles checkpoint retrieval from S3 or local storage, and executes the evaluation via torchrun.

Configuring LightEval for Your Training Pipeline

To enable automated evaluation, add a lighteval block to your training YAML configuration. The system uses this definition to instantiate the LightEvalConfig dataclass (lines 94-126 of the configuration file).

# config.yaml (excerpt)

lighteval:
  slurm:
    gpus_per_node: 8
    partition: hopper-prod
    cpus_per_task: 88
    qos: low
    time: 24:00:00
    reservation: smollm
    hf_cache: ~/.cache/huggingface
  parallelism:
    dp: 1
    pp: 1
    tp: 8
    tp_linear_async_communication: true
  batch_size: 4
  tasks:
    tasks: "mmlu,arc,hellaswag"
    max_samples: 5000
  logging:
    local_output_path: ./eval_results
    push_results_to_hub: true
    hub_repo_results: my-org/mymodel-evals
  wandb:
    wandb_project: nanotron_evals
    wandb_entity: my-wandb-team
  upload_to_wandb: true
  eval_interval: 5000            # run LightEval every 5,000 training steps

  s3_save_path: s3://my-bucket/evals

Key configuration parameters:

• eval_interval: The step frequency at which evaluations trigger (e.g., every 5,000 steps)
• parallelism: Dictates GPU allocation via the product of dp, pp, and tp factors
• tasks: Comma-separated string of benchmark names registered in [src/nanotron/eval/evaluation_tasks.py](https://github.com/huggingface/nanotron/blob/main/src/nanotron/eval/evaluation_tasks.py)
• s3_save_path: Optional remote storage for evaluation artifacts

Running Evaluations During Training

Integrate the evaluation trigger into your training loop using the LightEvalRunner class. The runner automatically handles interval checking and conditional job submission.

from nanotron.config import Config
from nanotron.eval.one_job_runner import LightEvalRunner

# Load the training configuration (already contains .lighteval)

cfg = Config.load_from_yaml("config.yaml")

# Suppose we are at step 10,000 and want to evaluate the latest checkpoint

uploaded_files = [
    {"destination": "/scratch/run1/10000/config.yaml", "source": "config.yaml"},
    # ... other checkpoint shards

]

runner = LightEvalRunner(config=cfg)
job_id, log_path = runner.eval_single_checkpoint(uploaded_files)

if job_id:
    print(f"Submitted LightEval SLURM job {job_id}, logs at {log_path}")
else:
    print("Evaluation was skipped (interval condition not met).")

The eval_single_checkpoint method performs three critical checks before submission:

1. Verifies that cfg.lighteval exists and is properly configured
2. Confirms the current step meets the eval_interval criteria
3. Validates that checkpoint files are available for download

Understanding the SLURM Execution Flow

Under the hood, run_slurm_one_job constructs a bash script that handles the complete evaluation lifecycle. The generated script performs the following operations:

1. Resource Allocation: Requests GPUs and CPUs based on the parallelism configuration
2. Checkpoint Retrieval: Downloads model weights from S3 using s5cmd or transfers via rsync from local storage (excluding optimizer states to save bandwidth)
3. Distributed Launch: Executes run_evals.py via torchrun with the appropriate node and process counts
4. Result Upload: Optionally pushes metrics to Weights & Biases using the helper script

Here is an excerpt of the generated SLURM script:

#!/bin/bash
#SBATCH --job-name=eval_10000_2024-04-01_12-00-00-run1
#SBATCH --partition=hopper-prod
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=88
#SBATCH --gpus=8
#SBATCH --exclusive
#SBATCH --qos=low
#SBATCH --time=24:00:00
#SBATCH --output=eval_results/logs/run1/step-10000/%j-2024-04-01_12-00-00.out
#SBATCH --requeue

set -x
LOCAL_DOWNLOAD_CHECKPOINT_FOLDER=/scratch/run1-eval-2024-04-01_12-00-00/10000
mkdir -p $LOCAL_DOWNLOAD_CHECKPOINT_FOLDER

# Download from S3 (or rsync from local storage)

s5cmd cp --concurrency=50 --exclude "optimizer/*" s3://my-bucket/checkpoints/10000/* $LOCAL_DOWNLOAD_CHECKPOINT_FOLDER/

# Run LightEval

CUDA_DEVICE_MAX_CONNECTIONS=1 torchrun \
  --nproc_per_node 8 \
  --nnodes 1 \
  --node_rank $SLURM_PROCID \
  --master_addr $MASTER_ADDR \
  --master_port $MASTER_PORT \
  ./run_evals.py \
  --checkpoint-config-path $LOCAL_DOWNLOAD_CHECKPOINT_FOLDER/config.yaml \
  --lighteval-override ./eval_config_override.yaml

# Optional Wandb upload

python ./src/nanotron/eval/upload_to_wandb.py \
  --wandb_project nanotron_evals \
  --wandb_entity my-wandb-team \
  --model_name run1 \
  --results_path s3://my-bucket/evals/results/ \
  --train_step 10000 \
  --consumed_tokens 32000000

The script is written to a temporary file (launch_script-<timestamp>.slurm) and submitted via the sbatch command.

Logging Results to Weights & Biases

After the SLURM job completes, evaluation results reside in JSON files within the checkpoint directory hierarchy (<step>/results_x.json). To visualize these metrics in WandB, use the post-processing script located at [scripts/log_lighteval_to_wandb.py](https://github.com/huggingface/nanotron/blob/main/scripts/log_lighteval_to_wandb.py).

python scripts/log_lighteval_to_wandb.py \
  --eval-path ./eval_results/logs/run1/step-10000 \
  --wandb-project nanotron_evals \
  --wandb-name eval_run1_step10000

This utility extracts per-benchmark accuracies (e.g., mmlu:average_acc, arc:average_acc) from the LightEval output files and logs them as distinct metrics to the specified WandB project. The script expects the standard LightEval directory structure where results are organized by checkpoint step.

Summary

• LightEvalConfig in src/nanotron/config/lighteval_config.py encapsulates all evaluation parameters including parallelism, SLURM settings, and task definitions.
• LightEvalRunner in src/nanotron/eval/one_job_runner.py manages evaluation scheduling through the eval_single_checkpoint method, respecting the configured eval_interval.
• SLURM integration automatically calculates GPU requirements (dp * pp * tp) and generates submission scripts that handle checkpoint download and distributed execution.
• Task registration occurs in src/nanotron/eval/evaluation_tasks.py, which wraps lighteval configurations in LightevalTaskConfig objects.
• WandB logging is supported through both real-time upload during the SLURM job and post-hoc aggregation via scripts/log_lighteval_to_wandb.py.

Frequently Asked Questions

How does Nanotron determine when to run LightEval?

The LightEvalRunner.eval_single_checkpoint method checks the current training step against the eval_interval value specified in LightEvalConfig. According to the implementation in src/nanotron/eval/one_job_runner.py (lines 27-74), the evaluation only proceeds if current_step % eval_interval == 0, preventing unnecessary benchmark runs during intermediate training steps.

What parallelism configuration does LightEval use in Nanotron?

LightEval uses the ParallelismArgs defined in LightEvalConfig.parallelism to compute resource requirements. The total GPU count is calculated as the product of data parallelism (dp), pipeline parallelism (pp), and tensor parallelism (tp) factors. This determines both the --gpus flag in the SLURM header and the torchrun parameters for distributed evaluation.

Can I run LightEval on checkpoints stored in S3?

Yes. When s3_save_path is configured, the run_slurm_one_job function generates bash commands using s5cmd to download checkpoints from S3 to a local scratch directory before evaluation. The download specifically excludes optimizer states using the --exclude "optimizer/*" pattern to minimize transfer time and storage requirements.

How do I add custom evaluation tasks to the LightEval integration?

Custom tasks must be registered in src/nanotron/eval/evaluation_tasks.py by creating a LightevalTaskConfig wrapper around your lighteval task configuration and adding it to the TASKS_TABLE. The task name can then be referenced in the YAML configuration's lighteval.tasks.tasks field alongside standard benchmarks like MMLU or HellaSwag.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}