重现 NVIDIA MLPerf v5.0 的 LLM 基准测试训练成绩

# create a directory where the data will be stored
mkdir </path/to/dataset>
# start the docker container 
docker run -it --rm --gpus all --network=host --ipc=host --volume </path/to/dataset>:/data mlperf-nvidia:llama2_70b_lora-pyt
# now you should be inside the container in the /workspace/ft-llm directory. run the download scripts
python scripts/download_dataset.py --data_dir /data/gov_report  # download dataset
python scripts/download_model.py --model_dir /data/model  # download preprocessed model checkpoint in NeMo format used for initialization; could take up to 30 minutes

export HF_TOKEN=<your/huggingface/token>

/data
├── gov_report
│   ├── train.npy
│   └── validation.npy
└── model
    ├── context
    │   ├── io.json
    │   ├── model.yaml
    │   └── nemo_tokenizer
    └── weights
        ├── common.pt
        ├── metadata.json
        ├── module.decoder.final_layernorm._extra_state
        ├── module.decoder.final_layernorm.weight
        ├── module.decoder.layers.mlp.linear_fc1._extra_state
        ├── module.decoder.layers.mlp.linear_fc1.layer_norm_weight
        ├── module.decoder.layers.mlp.linear_fc1.weight
        ├── module.decoder.layers.mlp.linear_fc2._extra_state
        ├── module.decoder.layers.mlp.linear_fc2.weight
        ├── module.decoder.layers.self_attention.core_attention._extra_state
        ├── module.decoder.layers.self_attention.linear_proj._extra_state
        ├── module.decoder.layers.self_attention.linear_proj.weight
        ├── module.decoder.layers.self_attention.linear_qkv._extra_state
        ├── module.decoder.layers.self_attention.linear_qkv.layer_norm_weight
        ├── module.decoder.layers.self_attention.linear_qkv.weight
        ├── module.embedding.word_embeddings.weight
        └── module.output_layer.weight

#!/bin/bash
source $(dirname ${BASH_SOURCE[0]})/config_common.sh

# hyperparameters
export MAX_STEPS=800
export LR=0.0005
export MINIBS=1
export TP=1
export SP=0
export CP=8

export LAYER_CUDA_GRAPH=0
export MCORE_CUDA_GRAPH=1

# system parameters
export VBOOST_VALUE=0
export DGXNNODES=18
export DGXNGPU=4
export WALLTIME_RUNANDTIME=10
export WALLTIME=$((5 + ${NEXP:-1} * ($WALLTIME_RUNANDTIME + 5)))

export DATADIR="</path/to/dataset>/gov_report"  # set your </path/to/dataset>
export MODEL="</path/to/dataset>/model"  # set your </path/to/dataset>
export LOGDIR="</path/to/output_logdir>"  # set the place where the output logs will be saved
export CONT=mlperf-nvidia:llama2_70b_lora-pyt
source config_<system>.sh  # select config and source it
sbatch -N $DGXNNODES -t $WALLTIME run.sub  # you may be required to set --account and --partition here

:::MLLOG {"namespace": "", "time_ms": 1745066769306, "event_type": "INTERVAL_START", "key": "init_start", "value": null, "metadata": {"file": "/workspace/ft-llm/train.py", "lineno": 327}}

:::MLLOG {"namespace": "", "time_ms": 1745066917960, "event_type": "INTERVAL_END", "key": "init_stop", "value": null, "metadata": {"file": "/usr/local/lib/python3.12/dist-packages/mlperf_common/callbacks/logging.py", "lineno": 83}}
:::MLLOG {"namespace": "", "time_ms": 1745066917961, "event_type": "INTERVAL_START", "key": "run_start", "value": null, "metadata": {"file": "/usr/local/lib/python3.12/dist-packages/mlperf_common/callbacks/logging.py", "lineno": 83}}

:::MLLOG {"namespace": "", "time_ms": 1745067009412, "event_type": "POINT_IN_TIME", "key": "eval_accuracy", "value": 0.92474365234375, "metadata": {"file": "/usr/local/lib/python3.12/dist-packages/mlperf_common/callbacks/logging.py", "lineno": 303, "samples_count": 3024}}
…
:::MLLOG {"namespace": "", "time_ms": 1745067009420, "event_type": "INTERVAL_END", "key": "run_stop", "value": null, "metadata": {"file": "/usr/local/lib/python3.12/dist-packages/mlperf_common/callbacks/logging.py", "lineno": 106, "samples_count": 3024, "status": "success"}}

c4-train.en_<number>_text_document where number belongs to 0~7.
c4-validation-91205-samples

special_tokens_map.json
tokenizer.json
tokenizer.model
tokenizer.model.v1
tokenizer_config.json

bash scripts/cleanup.sh

c4-train.en_6_text_document.bin
c4-train.en_6_text_document.idx
c4-train.en_7_text_document.bin
c4-train.en_7_text_document.idx
c4-validation-91205-samples.en_text_document.bin
c4-validation-91205-samples.en_text_document.idx

<LOAD_CHECKPOINTS_PATH>
└── 405b
	├── context
	│   ├── nemo_tokenizer
	│   │   ├── special_tokens_map.json
	│   │   ├── tokenizer_config.json
	│   │   └── tokenizer.json
	│   ├── io.json
	│   └── model.yaml
	└── weights
    	├── __0_0.distcp
    	├── __0_1.distcp
    	├── .metadata
    	├── common.pt
    	└── metadata.json

source $(dirname ${BASH_SOURCE[0]})/config_common.sh 
source $(dirname ${BASH_SOURCE[0]})/config_common_blackwell.sh 
source $(dirname ${BASH_SOURCE[0]})/config_common_cg.sh

export MINIBS=112
export TENSOR_MODEL_PARALLEL=4
export PIPELINE_MODEL_PARALLEL=8
export INTERLEAVED_PIPELINE=8
export CONTEXT_PARALLEL=2
export MICRO_BATCH_SIZE=1
export MODEL_SIZE="405b"
export MAX_STEPS=450

export FP8_PARAM_GATHER=True
export TP_COMM_OVERLAP=True
export ASYM_PP_EMBED=True
export ASYM_PP_LOSS=True
export TP_PP_DP_MAPPING=True

# Binding
export BINDCMD="bindpcie --cpu=node"

export DGXNNODES=128
export DGXNGPU=4
export DGXSYSTEM=$(basename $(readlink -f ${BASH_SOURCE[0]}) | sed 's/^config_//' | sed 's/\.sh$//' )

export WALLTIME_RUNANDTIME=180
export WALLTIME=$((5 + ${NEXP:-1} * ($WALLTIME_RUNANDTIME + 5)))

export PREPROC_DATA="/path/to/your/preprocessed_c4"
export TOKENIZER="/path/to/your/tokenizer.model"
export LOAD_CHECKPOINTS_PATH="/path/to/your/downloaded/checkpoint"
export LOAD_CHECKPOINT="/load_checkpoints/405b"
export LOGDIR=</path/to/output/dir>  # set the place where the output logs will be saved
export CONT=mlperf-nvidia:large_language_model-pyt
source config_GB200_128x4x112xtp4pp8cp2_cg_dplast.sh  # select config and source it
sbatch -N ${DGXNNODES} --time=${WALLTIME} run.sub  # you may be required to set --account and --partition here