Skip to content
Merged
  •  
  •  
  •  
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

import json
import logging
import math
Expand Down Expand Up @@ -39,10 +41,10 @@ def test_inference_pipeline(golden_values_path: str, test_values_path: str) -> N
# First warmup iteration is excluded from throughput statistics.
throughput_sampled = median(output_current["throughput"][1:])

# 5 token/seconds is empirically observed to be within run variance.
# 10% is empirically observed to be within hardware variance.
assert (
throughput_sampled >= output_groundtruth["throughput"] - 5.0
), f"Throughput is slower than expected! Expected ~{output_groundtruth['throughput']} tok/s but benchmarked {output_current['throughput']} tok/s"
throughput_sampled >= 0.9 * output_groundtruth["throughput"]
), f"Throughput is slower than expected! Expected to be within 10% of ~{output_groundtruth['throughput']} tok/s but benchmarked {output_current['throughput']} tok/s"

# If throughput is significantly improved (> 20%), update golden values accordingly.
assert (
Expand Down
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

import logging
from typing import Dict, List, Optional

Expand All @@ -9,7 +11,7 @@
logger = logging.getLogger(__name__)

CHECK_THRESHOLDS = {
"iteration-time": [common.ApproximateTest(atol=0, rtol=0.25)],
"iteration-time": [common.ApproximateTest(atol=0, rtol=0.05)],
"mem-allocated-bytes": [common.ApproximateTest(atol=0, rtol=0.05)],
"mem-max-allocated-bytes": [common.ApproximateTest(atol=0, rtol=0.05)],
"lm loss": [common.DeterministicTest(), common.ApproximateTest(atol=0, rtol=0.05)],
Expand All @@ -32,9 +34,7 @@ def test_regular_pipeline(
model_config = yaml.safe_load(f)

checks_types = (
model_config["METRICS"]
if "METRICS" in model_config
else ["iteration-time", "lm loss", "num-zeros"]
model_config["METRICS"] if "METRICS" in model_config else ["lm loss", "num-zeros"]
)
checks = {metric: CHECK_THRESHOLDS[metric] for metric in checks_types}

Expand Down
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

import logging
from typing import Dict

Expand All @@ -20,9 +22,7 @@ def test_resume_checkpoint_pipeline(
model_config = yaml.safe_load(f)

checks_types = (
model_config["METRICS"]
if "METRICS" in model_config
else ["iteration-time", "lm loss", "num-zeros"]
model_config["METRICS"] if "METRICS" in model_config else ["lm loss", "num-zeros"]
)
checks = {
metric: test_pretraining_regular_pipeline.CHECK_THRESHOLDS[metric]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ DATASET_DIR=""
TIME="1:00:00"
RECIPES_DIR="tests/test_utils/recipes"
CONTAINER_MOUNTS=""
NO_GPUS_PER_TASK="FALSE"

# Declare associative array for tracking unique mounts
declare -A seen_mounts
Expand Down Expand Up @@ -53,6 +54,10 @@ while [[ $# -gt 0 ]]; do
TIME="$2"
shift 2
;;
--no-gpus-per-task)
NO_GPUS_PER_TASK="TRUE"
shift 1
;;
--help)
print_usage
exit 0
Expand All @@ -65,12 +70,6 @@ while [[ $# -gt 0 ]]; do
esac
done

# Check if yq is installed
if ! command -v yq &>/dev/null; then
echo "Error: yq is not installed. Please install it first."
exit 1
fi

# Validate required arguments
if [ -z "$PARTITION" ] || [ -z "$SLURM_ACCOUNT" ] || [ -z "$IMAGE" ] || [ -z "$DATASET_DIR" ]; then
echo "Error: Missing required arguments"
Expand All @@ -89,7 +88,7 @@ SRUN_CMD="srun \
--container-workdir=/opt/megatron-lm \
--container-mounts=$CONTAINER_MOUNTS \
--nodes=1 \
--gpus-per-task=8 \
$(if [ "$NO_GPUS_PER_TASK" = "FALSE" ]; then echo "--gpus-per-task=8"; fi) \
--time=$TIME \
--pty bash"

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -232,56 +232,56 @@
"end_step": 50,
"step_interval": 1,
"values": {
"1": 10.41177,
"2": 0.63219,
"3": 0.53615,
"4": 0.53244,
"5": 0.53041,
"6": 0.53364,
"7": 0.53797,
"8": 0.52807,
"9": 0.53172,
"10": 0.53116,
"11": 0.52906,
"12": 0.53113,
"13": 0.52796,
"14": 0.52974,
"15": 0.52875,
"16": 0.52005,
"17": 0.51948,
"18": 0.52008,
"19": 0.52456,
"20": 0.52593,
"21": 0.52988,
"22": 0.52281,
"23": 0.51971,
"24": 0.52235,
"25": 0.54145,
"26": 0.52876,
"27": 0.51926,
"28": 0.51381,
"29": 0.51526,
"30": 0.51632,
"31": 0.52532,
"32": 0.61496,
"33": 0.59949,
"34": 0.52069,
"35": 0.52649,
"36": 0.66485,
"37": 0.52497,
"38": 0.52464,
"39": 0.76801,
"40": 0.52465,
"41": 0.69091,
"42": 0.74369,
"43": 0.5242,
"44": 0.75825,
"45": 0.68331,
"46": 0.75831,
"47": 0.51724,
"48": 0.51305,
"49": 0.51686,
"50": 0.52176
"1": 10.44279,
"2": 0.55345,
"3": 0.53909,
"4": 0.52187,
"5": 0.52958,
"6": 0.5241,
"7": 0.5353,
"8": 0.51946,
"9": 0.52732,
"10": 0.52759,
"11": 0.51849,
"12": 0.52326,
"13": 0.52472,
"14": 0.52577,
"15": 0.51817,
"16": 0.51922,
"17": 0.51686,
"18": 0.5248,
"19": 0.51945,
"20": 0.74697,
"21": 0.51544,
"22": 0.52412,
"23": 0.66206,
"24": 0.51781,
"25": 0.52429,
"26": 0.52068,
"27": 0.62432,
"28": 0.52016,
"29": 0.52217,
"30": 0.51949,
"31": 0.69033,
"32": 0.52127,
"33": 0.52602,
"34": 0.6403,
"35": 0.51723,
"36": 0.52445,
"37": 0.51746,
"38": 0.52296,
"39": 0.52159,
"40": 0.6718,
"41": 0.58171,
"42": 0.7393,
"43": 0.54277,
"44": 0.81615,
"45": 0.52284,
"46": 0.71947,
"47": 0.52219,
"48": 0.51866,
"49": 0.51764,
"50": 0.51841
}
}
}
Loading
Loading