From fca04b6861b681e6e5837c9f9eefd63092fe2422 Mon Sep 17 00:00:00 2001 From: vidushi8 Date: Thu, 9 Apr 2026 17:37:06 -0500 Subject: [PATCH 1/4] add mlperf llama3.1 8b in primus --- examples/mlperf/README.md | 55 ++++++ examples/mlperf/config_MI355X_1x8x1.sh | 71 ++++++++ .../MI355X/llama3.1_8B-pretrain-FP8.yaml | 119 ++++++++++++ examples/mlperf/run_and_time.sh | 61 +++++++ examples/mlperf/src/prof_handler.py | 166 +++++++++++++++++ examples/mlperf/src/train.py | 85 +++++++++ .../te_patches/fused_bias_swiglu_patches.py | 63 +++++++ .../validation_data_sampling_patches.py | 170 ++++++++++++++++++ .../backends/megatron/training/evaluator.py | 3 + .../modules/trainer/megatron/pre_trainer.py | 8 +- primus/modules/trainer/megatron/trainer.py | 34 +++- 11 files changed, 829 insertions(+), 6 deletions(-) create mode 100644 examples/mlperf/README.md create mode 100755 examples/mlperf/config_MI355X_1x8x1.sh create mode 100755 examples/mlperf/configs/MI355X/llama3.1_8B-pretrain-FP8.yaml create mode 100755 examples/mlperf/run_and_time.sh create mode 100755 examples/mlperf/src/prof_handler.py create mode 100755 examples/mlperf/src/train.py create mode 100644 primus/backends/megatron/patches/te_patches/fused_bias_swiglu_patches.py create mode 100644 primus/backends/megatron/patches/validation_data_sampling_patches.py diff --git a/examples/mlperf/README.md b/examples/mlperf/README.md new file mode 100644 index 000000000..5703cbab0 --- /dev/null +++ b/examples/mlperf/README.md @@ -0,0 +1,55 @@ +# LLama3.1 8B MLPerf Pretraining + +MLPerf-compliant LLama3.1 8B pretraining using Primus + +## Setup + +### Start Docker Image + +```bash +export MLPERF_PAT= +docker run -it --device /dev/dri --device /dev/kfd --device /dev/infiniband --network host --ipc host --group-add video --cap-add SYS_PTRACE --security-opt seccomp=unconfined --privileged -v $HOME:$HOME --shm-size 128G --name primus_training_env rocm/primus:v26.2 + + +git clone --recurse-submodules https://github.com/AMD-AIG-AIMA/Primus.git +cd Primus +``` + + +### Configuration + +- **Model**: LLama3.1 8B (4096 hidden, 32 layers, 32 attention heads) +- **Training**: 1.2M iterations, GBS=32, MBS=2, LR=8e-4 +- **Precision**: FP8 hybrid +- **Data**: C4 dataset (tokenized) + +## Key Files + +- `configs/MI355X/llama3.1_8B-pretrain.yaml` - Model and training config + - Update `train_data_path` and `train_data_path` to your local downloaded location +- `config_MI355X_1x8x1.sh` - System config and env vars + - Update `PRIMUS_PATH` to clone Primus Repo + - Update `EXP`to `/examples/mlperf/configs/MI355X/llama3.1_8B-pretrain-FP8.yaml` +- `src/train.py` - Training entry point +- `run_and_time.sh` - Run script + +### Data + +Download preprocessed C4 dataset: + +```bash +cd /data/mlperf_llama31_8b +bash <(curl -s https://raw.githubusercontent.com/mlcommons/r2-downloader/refs/heads/main/mlc-r2-downloader.sh) \ + -d data https://training.mlcommons-storage.org/metadata/llama-3-1-8b-preprocessed-c4-dataset.uri +``` + +### How to run + +```bash +export HF_TOKEN= +source config_MI355X_1x8x1.sh +bash run_and_time.sh +``` +## Notes + +- `log_interval: 99999999` suppresses regular Primus logs diff --git a/examples/mlperf/config_MI355X_1x8x1.sh b/examples/mlperf/config_MI355X_1x8x1.sh new file mode 100755 index 000000000..ea3b0f5df --- /dev/null +++ b/examples/mlperf/config_MI355X_1x8x1.sh @@ -0,0 +1,71 @@ +#!/bin/bash +# MLPerf LLama3.1 8B Configuration for MI355X (1x8x1) + +export DGXSYSTEM=MI355X_1x8x1 +export GPUS_PER_NODE=8 +export NNODES=1 +export NODE_RANK=0 +export MASTER_ADDR=localhost +export MASTER_PORT=29502 + +export PRIMUS_PATH=/home/vidgoyal/Primus-dev/Primus/ +export PRIMUS_MLPERF=1 +export PYTHONPATH="${PRIMUS_PATH}:${PRIMUS_PATH}/third_party/Megatron-LM:${PYTHONPATH}" +export EXP=/home/vidgoyal/Primus-dev/Primus/examples/mlperf/configs/MI355X/llama3.1_8B-pretrain-FP8.yaml +export DATA_PATH=/data + +export PRIMUS_MICRO_BATCH_SIZE=2 +export PRIMUS_GLOBAL_BATCH_SIZE=32 +export PRIMUS_LR=8e-4 +export PRIMUS_TRAIN_ITERS=1200000 +export EVAL_SAMPLES_INTERVAL=12288 #12288 # Evaluate every 12,288 samples +export PRIMUS_EVAL_INTERVAL=$((EVAL_SAMPLES_INTERVAL / PRIMUS_GLOBAL_BATCH_SIZE)) # Auto-computed +export SEED=31952 + +export ROCTRACER_LOG=1 +export ROCTRACER_LOG_LEVEL=5 + +export HSA_ENABLE_INTERRUPT=0 +export HSA_TOOLS_LIB=/opt/rocm/lib/libroctracer64.so + + +export PRIMUS_APPLY_ROPE_FUSION=True +export PRIMUS_FP8_RECIPE=hybrid + +export HSA_NO_SCRATCH_RECLAIM=1 +export HSA_ENABLE_SDMA=1 +export GPU_MAX_HW_QUEUES=2 +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export NVTE_FUSED_ATTN=1 +export NVTE_FUSED_ATTN_CK=1 +export NVTE_FUSED_ATTN_AOTRITON=1 +export NVTE_CK_USES_FWD_V3=1 +export NVTE_CK_USES_BWD_V3=1 +export NVTE_CK_IS_V3_ATOMIC_FP32=0 +export USE_TE_SWIGLU=1 + +export NCCL_CHECKS_DISABLE=1 +export TORCH_NCCL_HIGH_PRIORITY=1 + +export NVTE_ASYNC_AMAX_REDUCTION=1 +export NVTE_DP_AMAX_REDUCE_INTERVAL=0 + +export ENABLE_MLLOG=1 +export MLLOG_OUTPUT_FILE=/results/mlperf_output.log +export MLLOG_TRAIN_LOSS_LOG_FREQ=100 +export MLLOG_TARGET_EVAL_LOSS=3.3 +export MLLOG_SUBMISSION_BENCHMARK=llama31_8b +export MLLOG_SUBMISSION_DIVISION=closed +export MLLOG_SUBMISSION_ORG=AMD +export MLLOG_SUBMISSION_PLATFORM=MI355X + +export TORCHPROF_OUTPUT_DIR=/home/vidgoyal/small_llm_pretraining/primus/outputs/ +export TORCHPROF_VERBOSE=0 +export TORCHPROF_MAXROWS=100 +export TORCHPROF_PROFILE_MEMORY=0 +export TORCHPROF_WITH_STACK=0 +export TORCHPROF_RECORD_SHAPES=0 +export TORCHPROF_WITH_FLOPS=0 +export PROF_WARMUP_STEPS=10 #128 +export PROF_ACTIVE_STEPS=6372 +# export HIPBLASLT_TUNING_OVERRIDE_FILE=tuning.txt diff --git a/examples/mlperf/configs/MI355X/llama3.1_8B-pretrain-FP8.yaml b/examples/mlperf/configs/MI355X/llama3.1_8B-pretrain-FP8.yaml new file mode 100755 index 000000000..bcb62daf0 --- /dev/null +++ b/examples/mlperf/configs/MI355X/llama3.1_8B-pretrain-FP8.yaml @@ -0,0 +1,119 @@ +work_group: ${TEAM:amd} +user_name: ${USER:root} +exp_name: ${EXP_NAME:llama3.1_8B-pretrain-v26.2} +workspace: ./output + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + # model to run + model: llama3.1_8B.yaml + overrides: + # --- Logging Config --- + wandb_project: "Primus-llama3.1-8B-pretrain" + disable_wandb: true + disable_tensorboard: true + stderr_sink_level: DEBUG + log_interval: 999999 + log_avg_skip_iterations: 2 + log_avg_reset_interval: 50 + + eval_iters: 32 # 32 * GBS = 1024 eval samples + eval_interval: ${PRIMUS_EVAL_INTERVAL:10} # 10 * GBS = 320 eval samples perform evaluation. + + # --- Training Config --- + train_iters: ${PRIMUS_TRAIN_ITERS:200} + micro_batch_size: 2 # grad_acc = global_batch_size / (micro_batch_size * num_gpus) = 32 / (2 * 8) = 2 + global_batch_size: 32 + + seq_length: 8192 + max_position_embeddings: 8192 + + seed: ${SEED:1234} + + lr: 0.0008 # 8e-4 + min_lr: 0.00008 # 10% of lr + lr_warmup_iters: 128 # TODO: lr warmup steps should be 128. + lr_decay_iters: 1199872 # 1200000 - 128 + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + eod_mask_loss: false + init_method_std: 0.02 + norm_epsilon: 1.0e-5 + + adam_eps: 1.0e-5 # TODO, find out what is the correct key to this parameter. + + check_for_nan_in_loss_and_grad: false # default true + check_for_spiky_loss: false # default false, but setting it here explicitly + check_for_large_grads: false # default false, but setting it here explicitly + + # --- Model Parallel Config --- + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + overlap_grad_reduce: true + overlap_param_gather: true + gradient_accumulation_fusion: true + + # --- Data Config --- + mock_data: false + train_data_path: "/data/mlperf/data/c4-train.en_6_text_document" + valid_data_path: "/data/mlperf/data/c4-validation-91205-samples.en_text_document" + test_data_path: null + seq_length: 8192 + data_cache_path: /npy_indices + mmap_bin_files: true + + # --- Profiling Config --- + profile: false + use_pytorch_profiler: true + profile_ranks: [0] # Only profile rank 0 to save disk space + profile_step_start: 8 # Start after warmup (step 8) + profile_step_end: 13 # Profile 5 iterations (8-12) + disable_profiler_activity_cpu: false # GPU kernels only (smaller files) + torch_profiler_record_shapes: false # Disable for smaller traces + torch_profiler_with_stack: false # Disable for smaller traces + torch_profiler_use_gzip: true # Compress output + + # --- Checkpointing Config --- + finetune: false + auto_continue_train: false + load: null + no_load_optim: null + no_load_rng: null + save: null + save_interval: 2000000 + no_save_optim: null + no_save_rng: null + disable_last_saving: true + ckpt_format: torch + + # --- FSDP Config --- + use_torch_fsdp2: false + use_distributed_optimizer: true # this is needed for fsdp2 + + # Cross entropy flags + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true + + # --- Mixed Precision Config --- + fp8: hybrid # e4m3, hybrid + fp8_amax_history_len: 4 + fp8_amax_compute_algo: "most_recent" + accumulate_allreduce_grads_in_fp32: false + grad_reduce_in_bf16: true + attention_softmax_in_fp32: false + # fp8_param_gather: true + + + # --- Primus Turbo Config --- + enable_primus_turbo: false + use_turbo_attention: false + use_turbo_parallel_linear: false # can't use together with delayed recipe + use_turbo_grouped_mlp: false + moe_use_fused_router_with_aux_score: false + enable_turbo_attention_float8 : false diff --git a/examples/mlperf/run_and_time.sh b/examples/mlperf/run_and_time.sh new file mode 100755 index 000000000..1707f854f --- /dev/null +++ b/examples/mlperf/run_and_time.sh @@ -0,0 +1,61 @@ +#!/bin/bash + +set -e + +mkdir -p /results + +export GPUS_PER_NODE=${GPUS_PER_NODE:-8} +export NNODES=${NNODES:-1} +export NODE_RANK=${NODE_RANK:-0} +export MASTER_ADDR=${MASTER_ADDR:-localhost} +export MASTER_PORT=${MASTER_PORT:-29502} +export EXP=${EXP:-/workspace/code/conf/llama3.1_8B-pretrain.yaml} +export DATA_PATH=${DATA_PATH:-/data} + +echo "============================================" +echo "MLPerf LLama3.1 8B Training" +echo "============================================" +echo "Config: ${EXP}" +echo "Data: ${DATA_PATH}" +echo "GPUs: ${GPUS_PER_NODE}" +echo "Nodes: ${NNODES}" +echo "Train iters: ${PRIMUS_TRAIN_ITERS}" +echo "Eval interval: ${PRIMUS_EVAL_INTERVAL}" +echo "Enable MLPerf logging: ${ENABLE_MLPERF}" +echo "MLLOG_TRAIN_LOSS_LOG_FREQ: ${MLLOG_TRAIN_LOSS_LOG_FREQ}" +echo "MLLOG_TARGET_EVAL_LOSS: ${MLLOG_TARGET_EVAL_LOSS}" +echo "MLLOG_SUBMISSION_BENCHMARK: ${MLLOG_SUBMISSION_BENCHMARK}" +echo "MLLOG_SUBMISSION_DIVISION: ${MLLOG_SUBMISSION_DIVISION}" +echo "MLLOG_SUBMISSION_ORG: ${MLLOG_SUBMISSION_ORG}" +echo "MLLOG_SUBMISSION_PLATFORM: ${MLLOG_SUBMISSION_PLATFORM}" +echo "============================================" + +start=$(date +%s) +start_fmt=$(date +%Y-%m-%d\ %r) +echo "STARTING TIMING RUN AT $start_fmt" + +torchrun \ + --nproc_per_node=${GPUS_PER_NODE} \ + --nnodes=${NNODES} \ + --node_rank=${NODE_RANK} \ + --master_addr=${MASTER_ADDR} \ + --master_port=${MASTER_PORT} \ + src/train.py + +ret_code=$? + +end=$(date +%s) +end_fmt=$(date +%Y-%m-%d\ %r) +echo "ENDING TIMING RUN AT $end_fmt" + +result=$(( end - start )) +result_name="LLAMA3.1_8B" +echo "RESULT,$result_name,,$result,AMD,$start_fmt" + +if [[ $ret_code != 0 ]]; then + echo "Training failed with exit code: $ret_code" + exit $ret_code +fi + +exit 0 + diff --git a/examples/mlperf/src/prof_handler.py b/examples/mlperf/src/prof_handler.py new file mode 100755 index 000000000..6a28e7417 --- /dev/null +++ b/examples/mlperf/src/prof_handler.py @@ -0,0 +1,166 @@ +import torch +import os +import pathlib +from dataclasses import dataclass +from torch.profiler import ProfilerActivity +import pandas as pd + + +TORCHPROF_OUTPUT_DIR = os.getenv("TORCHPROF_OUTPUT_DIR", "/results/artifacts") +TORCHPROF_OUTPUT = os.getenv("TORCHPROF_OUTPUT", "csv_handler") +TORCHPROF_VERBOSE = os.getenv("TORCHPROF_VERBOSE", 1) +TORCHPROF_DEVICES = os.getenv("TORCHPROF_DEVICES", "GPU") +TORCHPROF_MAXROWS = os.getenv("TORCHPROF_MAXROWS", 100) +TORCHPROF_PROFILE_MEMORY = bool(os.getenv("TORCHPROF_PROFILE_MEMORY", 1)) +TORCHPROF_WITH_STACK = bool(os.getenv("TORCHPROF_WITH_STACK", 0)) +TORCHPROF_RECORD_SHAPES = bool(os.getenv("TORCHPROF_RECORD_SHAPES", 1)) +TORCHPROF_WITH_FLOPS = bool(os.getenv("TORCHPROF_WITH_FLOPS", 1)) +PROF_WARMUP_STEPS = int(os.getenv("PROF_WARMUP_STEPS", 3)) +PROF_ACTIVE_STEPS = int(os.getenv("PROF_ACTIVE_STEPS", 2)) +PROF_REPITIONS = int(os.getenv("PROF_REPITIONS", 1)) + +def get_devices() -> list: + devices = TORCHPROF_DEVICES.split(",") + devices = [x.lower().strip() for x in devices] + devices_set = set() + for device in devices: + if device.lower() not in ['cpu', 'gpu']: + raise ValueError(f"Invalid Device :{device}") + if device.lower() == "gpu": + devices_set.add(torch.profiler.ProfilerActivity.CUDA) + elif device.lower() == "cpu": + devices_set.add(torch.profiler.ProfilerActivity.CPU) + + devices_list = list(devices_set) + + return devices_list + +@dataclass +class TorchProfConfig: + skip_first = 1 + wait = 0 + warmup = PROF_WARMUP_STEPS + active = PROF_ACTIVE_STEPS + repeat = PROF_REPITIONS + +TOTAL_WARMUP_STEPS = TorchProfConfig.skip_first + \ + TorchProfConfig.wait + \ + TorchProfConfig.warmup + +TOTAL_ACTIVE_STEPS = TorchProfConfig.active + + +def kernel_table_from_prof(prof, top_k=200): + rows = [] + for evt in prof.events(): + # evt has fields like name, device_type, cuda_time_total, cpu_time_total, etc. + # Kernel events show up as CUDA device events in the trace. + # We keep entries that have some CUDA time. + cuda_us = getattr(evt, "cuda_time_total", 0) or 0 + if cuda_us <= 0: + continue + + name = getattr(evt, "name", "") + + # Heuristic: keep kernel-ish entries. + # Depending on PyTorch version, true kernels often include patterns like: + # - "void " (C++ kernel signature) + # - "ampere_" / "sm" / "cutlass" / "cudnn" / "gemm" / "flash" etc. + # You can loosen/tighten this filter. + #kernelish = ( + # name.startswith("void ") + # or "kernel" in name.lower() + # or "cudnn" in name.lower() + # or "cutlass" in name.lower() + # or "gemm" in name.lower() + #) + #if not kernelish: + # continue + + rows.append((name, cuda_us)) + + if not rows: + return pd.DataFrame(columns=["kernel", "calls", "total_us", "avg_us", "pct_total"]) + + df = pd.DataFrame(rows, columns=["kernel", "cuda_us"]) + + agg = (df.groupby("kernel", as_index=False) + .agg(calls=("cuda_us", "size"), total_us=("cuda_us", "sum"))) + agg["avg_us"] = agg["total_us"] / agg["calls"] + total = agg["total_us"].sum() + agg["pct_total"] = 100.0 * agg["total_us"] / total + agg = agg.sort_values("total_us", ascending=False).head(top_k) + + # Pretty formatting helpers + agg["total_ms"] = agg["total_us"] / 1000.0 + agg["avg_ms"] = agg["avg_us"] / 1000.0 + + return agg[["kernel", "calls", "avg_ms", "total_ms", "pct_total"]] + + +def trace_handler(prof): + save_path = f"{TORCHPROF_OUTPUT_DIR}/key_avg_{prof.step_num}_{torch.distributed.get_rank()}.txt" + print(f"Saving torchprof results at: {save_path}") + with open(save_path, 'w') as f: + #output = prof.key_averages(group_by_input_shape=True).table(sort_by="self_cuda_time_total", row_limit=TORCHPROF_MAXROWS) + output = prof.key_averages().table(row_limit=TORCHPROF_MAXROWS) + f.write(output) + if TORCHPROF_VERBOSE: + print(output) + + save_path_pandas = f"{TORCHPROF_OUTPUT_DIR}/pandas_key_avg_{prof.step_num}_{torch.distributed.get_rank()}.txt" + with open(save_path_pandas, 'w') as f: + kt = kernel_table_from_prof(prof, top_k=200) + kt.to_csv(save_path_pandas, index=False) + + + + prof.export_chrome_trace(f"{TORCHPROF_OUTPUT_DIR}/trace_{prof.step_num}_{torch.distributed.get_rank()}.json") + + +def _get_torchprof(): + if TORCHPROF_OUTPUT == "csv_handler": + output_handler = trace_handler + elif TORCHPROF_OUTPUT == "tensorboard": + output_handler = torch.profiler.tensorboard_trace_handler(TORCHPROF_OUTPUT_DIR) + else: + raise ValueError("Invalid Output Handler for TorchProf.") + pathlib.Path(TORCHPROF_OUTPUT_DIR).mkdir(exist_ok=True, parents=True) + + prof = torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA], on_trace_ready=output_handler) + + #prof = torch.profiler.profile( + # activities=get_devices(), + # schedule=torch.profiler.schedule( + # skip_first=TorchProfConfig.skip_first, + # wait=TorchProfConfig.wait, + # warmup=TorchProfConfig.warmup, + # active=TorchProfConfig.active, + # repeat=TorchProfConfig.repeat + # ), + # called each time the trace is ready at the end of each cycle. + # on_trace_ready=output_handler, + # profile_memory=TORCHPROF_PROFILE_MEMORY, # adds extra overhead if True !! + # with_flops=TORCHPROF_WITH_FLOPS, + # with_stack=TORCHPROF_WITH_STACK, # adds extra overhead if True !! + # record_shapes=TORCHPROF_RECORD_SHAPES + #) + + return prof + +def get_profiler(): + if os.getenv("PROFILER", '') == 'torchprof': + return _get_torchprof() + return None + +def _get_rpd(): + from rpdTracerControl import rpdTracerControl + rpdTracerControl.setFilename(name=f"trace.rpd", append=True) + prof = rpdTracerControl() + print(f"RPD profiler initialized") + return prof + +def get_profiler_rpd(): + if os.getenv("PROFILER", '') == 'rpd': + return _get_rpd() + return None diff --git a/examples/mlperf/src/train.py b/examples/mlperf/src/train.py new file mode 100755 index 000000000..73b0b23b3 --- /dev/null +++ b/examples/mlperf/src/train.py @@ -0,0 +1,85 @@ +import os +import sys +from pathlib import Path +import argparse + +PRIMUS_PATH = os.getenv("PRIMUS_PATH", "/home/vidgoyal/Primus") +MEGATRON_PATH = os.path.join(PRIMUS_PATH, "third_party/Megatron-LM") + +if PRIMUS_PATH not in sys.path: + sys.path.insert(0, PRIMUS_PATH) +if MEGATRON_PATH not in sys.path: + sys.path.insert(0, MEGATRON_PATH) + +from primus.core.launcher.config import PrimusConfig +from primus.core.launcher.parser import load_primus_config, add_pretrain_parser +from primus.modules.trainer.megatron.pre_trainer import MegatronPretrainTrainer + + +def setup_environment(data_path: str = None): + if data_path and "HF_HOME" not in os.environ: + hf_home = os.path.join(data_path, "huggingface") + os.environ["HF_HOME"] = hf_home + print(f"[MLPerf Train] HF_HOME={hf_home}") + + +def load_config(config_path: str, overrides: list = None) -> PrimusConfig: + parser = argparse.ArgumentParser() + add_pretrain_parser(parser) + + args = parser.parse_args([ + '--config', config_path, + '--data_path', os.getenv('DATA_PATH', '/data'), + ]) + + primus_cfg = load_primus_config(args, overrides or []) + + print(f"[MLPerf Train] Loaded config from: {config_path}") + print(f"[MLPerf Train] Framework: {primus_cfg[0].get_module_config('pre_trainer').framework}") + + return primus_cfg[0] + + +def create_trainer(primus_cfg: PrimusConfig) -> MegatronPretrainTrainer: + rank = int(os.getenv("RANK", "0")) + world_size = int(os.getenv("WORLD_SIZE", "1")) + master_addr = os.getenv("MASTER_ADDR", "127.0.0.1") + master_port = int(os.getenv("MASTER_PORT", "29500")) + + trainer = MegatronPretrainTrainer( + primus_config=primus_cfg, + module_rank=rank, + module_world_size=world_size, + module_master_addr=master_addr, + module_master_port=master_port, + ) + + return trainer + + +def main(): + config_path = os.environ.get("EXP", "/workspace/code/conf/llama3.1_8B-pretrain.yaml") + + if not Path(config_path).exists(): + raise FileNotFoundError(f"Config not found: {config_path}") + + setup_environment(data_path=os.getenv('DATA_PATH', '/data')) + primus_cfg = load_config(config_path) + + trainer = create_trainer(primus_cfg) + trainer.init() + + if os.getenv("PROFILER") == "torchprof": + import torch.profiler + from prof_handler import trace_handler, TORCHPROF_OUTPUT_DIR + import pathlib + pathlib.Path(TORCHPROF_OUTPUT_DIR).mkdir(exist_ok=True, parents=True) + torch.profiler.tensorboard_trace_handler = lambda *a, **kw: trace_handler + print(f"[Profiler] torchprof enabled, output dir: {TORCHPROF_OUTPUT_DIR}") + + trainer.run() + + +if __name__ == "__main__": + main() + diff --git a/primus/backends/megatron/patches/te_patches/fused_bias_swiglu_patches.py b/primus/backends/megatron/patches/te_patches/fused_bias_swiglu_patches.py new file mode 100644 index 000000000..51dc6ead2 --- /dev/null +++ b/primus/backends/megatron/patches/te_patches/fused_bias_swiglu_patches.py @@ -0,0 +1,63 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Transformer Engine Fused Bias SwiGLU Patches + +Patches SwiGLUFunction to use TE's fused swiglu/dswiglu kernels when +USE_TE_SWIGLU=1 is set, providing better performance on ROCm GPUs. +""" + +import os + +from primus.core.patches import PatchContext, register_patch +from primus.modules.module_utils import log_rank_0 + + +@register_patch( + "megatron.te.fused_bias_swiglu", + backend="megatron", + phase="before_train", + description="Use TE fused swiglu/dswiglu kernels in SwiGLUFunction forward/backward", + condition=lambda ctx: os.getenv("USE_TE_SWIGLU", "0") == "1", +) +def patch_te_fused_bias_swiglu(ctx: PatchContext): + """ + Patch SwiGLUFunction to use Transformer Engine's fused swiglu/dswiglu + C++ extensions in forward and backward passes. + + Activated when USE_TE_SWIGLU=1 is set in the environment. + """ + from transformer_engine.pytorch.cpp_extensions import ( + dswiglu as te_dswiglu, + swiglu as te_swiglu, + ) + + from megatron.core.fusions.fused_bias_swiglu import SwiGLUFunction + + @staticmethod + def new_forward(ctx, input, fp8_input_store, cpu_offload_input): + input_for_backward = input.to(__import__("torch").float8_e4m3fn) if fp8_input_store else input + if cpu_offload_input: + input_for_backward.activation_offloading = True + ctx.save_for_backward(input_for_backward) + ctx.ori_input_dtype = input.dtype + ctx.fp8_input_store = fp8_input_store + return te_swiglu(input, None) + + @staticmethod + def new_backward(ctx, grad_output): + input = ctx.saved_tensors[0] + input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input + return te_dswiglu(grad_output, input, None), None, None + + SwiGLUFunction.forward = new_forward + SwiGLUFunction.backward = new_backward + + log_rank_0( + "[Patch:megatron.te.fused_bias_swiglu] Patched SwiGLUFunction " + "to use TE fused swiglu/dswiglu kernels (USE_TE_SWIGLU=1)" + ) diff --git a/primus/backends/megatron/patches/validation_data_sampling_patches.py b/primus/backends/megatron/patches/validation_data_sampling_patches.py new file mode 100644 index 000000000..d06ef186d --- /dev/null +++ b/primus/backends/megatron/patches/validation_data_sampling_patches.py @@ -0,0 +1,170 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Megatron Validation Data Sampling Patches + +Patches Megatron's validation data loading to use a fixed, reproducible +subset of validation samples (eval_iters * global_batch_size) starting +from offset 0, matching MLPerf evaluation protocol. + +Without this patch, Megatron: + 1. Over-allocates validation samples proportional to total training steps + 2. Advances through validation data using consumed_valid_samples offset, + making eval loss non-reproducible across runs +""" + +import os + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.modules.module_utils import log_rank_0 + + +def _is_mlperf_enabled(ctx): + return os.getenv("PRIMUS_MLPERF", "0") == "1" and getattr(get_args(ctx), "eval_iters", 0) > 0 + + +@register_patch( + "megatron.training.validation_num_samples", + backend="megatron", + phase="before_train", + description=( + "Fix validation sample count to eval_iters * global_batch_size " + "instead of scaling with train_iters" + ), + condition=_is_mlperf_enabled, +) +def patch_validation_num_samples(ctx: PatchContext): + """ + Patch get_train_valid_test_num_samples to allocate only + eval_iters * global_batch_size validation samples instead of + (train_iters // eval_interval + 1) * eval_iters * global_batch_size. + """ + import megatron.training.training as training_module + + def patched_get_train_valid_test_num_samples(): + from megatron.training import get_args + + args = get_args() + + if args.train_samples: + train_samples = args.train_samples + else: + train_samples = args.train_iters * args.global_batch_size + + if args.full_validation: + eval_samples = None + else: + eval_samples = args.eval_iters * args.global_batch_size + + test_samples = args.eval_iters * args.global_batch_size + + if hasattr(args, "phase_transition_iterations") and args.phase_transition_iterations: + phase_transition_samples = ( + [0] + + [t * args.global_batch_size for t in args.phase_transition_iterations] + + [args.train_samples] + ) + current_sample = args.iteration * args.global_batch_size + for i in range(len(phase_transition_samples) - 1): + if phase_transition_samples[i] <= current_sample < phase_transition_samples[i + 1]: + train_samples = phase_transition_samples[i + 1] - phase_transition_samples[i] + break + + return (train_samples, eval_samples, test_samples) + + training_module.get_train_valid_test_num_samples = patched_get_train_valid_test_num_samples + log_rank_0( + "[Patch:megatron.training.validation_num_samples] " + "Patched get_train_valid_test_num_samples: eval_samples = eval_iters * gbs" + ) + + +@register_patch( + "megatron.training.validation_data_loader", + backend="megatron", + phase="before_train", + description=( + "Always build validation dataloader with consumed_samples=0 and " + "cap total_samples at eval_iters * gbs for reproducible evaluation" + ), + condition=_is_mlperf_enabled, +) +def patch_validation_data_loader(ctx: PatchContext): + """ + Patch build_pretraining_data_loader to detect validation datasets and + force consumed_samples=0 with total_samples capped at eval_iters * gbs. + + Also patches build_train_valid_test_data_loaders to always pass + consumed_samples=0 for validation dataloaders (instead of + args.consumed_valid_samples). + """ + import torch.utils.data + + from megatron.core import mpu + from megatron.core.datasets.utils import Split + + try: + from megatron.training.datasets import data_samplers as samplers_module + from megatron.training.datasets.data_samplers import MegatronPretrainingSampler + except ImportError: + from megatron.legacy.data import data_samplers as samplers_module + from megatron.legacy.data.data_samplers import MegatronPretrainingSampler + + orig_build_loader = samplers_module.build_pretraining_data_loader + + def patched_build_pretraining_data_loader(dataset, consumed_samples, name=""): + """Replacement that forces validation to always start from sample 0 + with a fixed sample count.""" + if dataset is None: + return None + + from megatron.training import get_args + + args = get_args() + + if hasattr(dataset, "split"): + split = dataset.split + elif hasattr(dataset, "index_split"): + split = dataset.index_split + else: + split = None + + is_validation = (split == Split.valid) or (name == "validation") + + if is_validation: + eval_samples = args.eval_iters * args.global_batch_size + total_samples = min(len(dataset), eval_samples) + batch_sampler = MegatronPretrainingSampler( + total_samples=total_samples, + consumed_samples=0, + micro_batch_size=args.micro_batch_size, + data_parallel_rank=mpu.get_data_parallel_rank(), + data_parallel_size=mpu.get_data_parallel_world_size(), + ) + return torch.utils.data.DataLoader( + dataset, + batch_sampler=batch_sampler, + num_workers=args.num_workers, + pin_memory=True, + persistent_workers=True if args.num_workers > 0 else False, + ) + + return orig_build_loader(dataset, consumed_samples) + + # Replace in both the data_samplers module and training module so all + # call sites (including build_train_valid_test_data_loaders) use our version + samplers_module.build_pretraining_data_loader = patched_build_pretraining_data_loader + + import megatron.training.training as training_module + + training_module.build_pretraining_data_loader = patched_build_pretraining_data_loader + + log_rank_0( + "[Patch:megatron.training.validation_data_loader] " + "Patched build_pretraining_data_loader: validation always uses " + "consumed_samples=0, total_samples capped at eval_iters * gbs" + ) diff --git a/primus/backends/megatron/training/evaluator.py b/primus/backends/megatron/training/evaluator.py index 24d59cc7b..e56222134 100644 --- a/primus/backends/megatron/training/evaluator.py +++ b/primus/backends/megatron/training/evaluator.py @@ -140,6 +140,9 @@ def primus_evaluate( total_loss_dict[key] = total_loss_numerators[key] / total_loss_denominators[key] else: total_loss_dict[key] = 0.0 + if "lm loss" in total_loss_dict: + val = total_loss_dict["lm loss"] + args._eval_val_loss = val.item() if hasattr(val, "item") else float(val) collected_non_loss_data = None if non_loss_data_func is not None: diff --git a/primus/modules/trainer/megatron/pre_trainer.py b/primus/modules/trainer/megatron/pre_trainer.py index d34788b04..35e0500dd 100644 --- a/primus/modules/trainer/megatron/pre_trainer.py +++ b/primus/modules/trainer/megatron/pre_trainer.py @@ -80,7 +80,13 @@ def get_batch_func(data_iterator, vp_stage=None): else: batch = get_batch_on_this_cp_rank(batch) - return batch.values() + return ( + batch["tokens"], + batch["labels"], + batch["loss_mask"], + batch.get("attention_mask"), + batch["position_ids"], + ) class DataLoaderStore: diff --git a/primus/modules/trainer/megatron/trainer.py b/primus/modules/trainer/megatron/trainer.py index 8f0c9f70e..62af5e7b0 100644 --- a/primus/modules/trainer/megatron/trainer.py +++ b/primus/modules/trainer/megatron/trainer.py @@ -610,8 +610,6 @@ def train_valid_test_datasets_provider_func(train_val_test_num_samples, vp_stage args.do_valid, args.do_test, args.dataloader_type, - args.retro_project_dir, - args.retro_cyclic_train_iters, ) # Print setup timing. @@ -916,9 +914,6 @@ def setup_model_and_optimizer( optimizer = get_megatron_optimizer( config, model, - no_wd_decay_cond, - scale_lr_cond, - lr_mult, use_gloo_process_groups=args.enable_gloo_process_groups, ) else: @@ -1249,6 +1244,8 @@ def train( pre_hook_enabled = False should_exit = False exit_code = 0 + train_start_time = time.time() + log_rank_0(f"Training start time: {train_start_time:.3f} (wall-clock)") if args.manual_gc: # Disable the default garbage collector and perform the collection manually. @@ -1534,6 +1531,21 @@ def get_e2e_base_metrics(): if should_disable_forward_pre_hook(args): enable_forward_pre_hook(model) pre_hook_enabled = True + + target_eval_loss = float(os.environ.get("TARGET_EVAL_LOSS", "0")) + if target_eval_loss > 0 and hasattr(args, "_eval_val_loss"): + if args._eval_val_loss <= target_eval_loss: + run_duration = time.time() - train_start_time + log_rank_0( + f"[EarlyStop] Reached target! Stopping training. " + f"eval_loss: {args._eval_val_loss:.6f} (target: {target_eval_loss}) | " + f"consumed_train_samples: {args.consumed_train_samples} | " + f"duration: {run_duration:.1f} sec -> {run_duration / 60.0:.1f} minutes" + ) + args.train_iters = iteration + args.do_valid = False + args.do_test = False + timers("interval-time", log_level=0).start(barrier=True) # Miscellaneous post-training-step functions (e.g., FT heartbeats, GC). @@ -1562,6 +1574,18 @@ def get_e2e_base_metrics(): one_logger_utils.track_e2e_metrics() + run_duration = time.time() - train_start_time + target_eval_loss = float(os.environ.get("TARGET_EVAL_LOSS", "0")) + final_eval_loss = getattr(args, "_eval_val_loss", None) + status = "success" if (final_eval_loss and target_eval_loss > 0 and final_eval_loss <= target_eval_loss) else "aborted" + samples_per_sec = args.consumed_train_samples / run_duration if run_duration > 0 else 0 + log_rank_0( + f"[TimeToTrain] status: {status} | " + f"samples_count: {args.consumed_train_samples} | " + f"duration: {run_duration:.1f} sec -> {run_duration / 60.0:.1f} minutes | " + f"throughput: {samples_per_sec:.2f} samples/s" + ) + if args.dump_pp_data: from .utils import dump_pp_data From 9c29b9a35f054380611e25a3b56dd00c491715b2 Mon Sep 17 00:00:00 2001 From: vidushi8 Date: Thu, 16 Apr 2026 11:55:14 -0500 Subject: [PATCH 2/4] updated to work with latest primus --- examples/mlperf/config_MI355X_1x8x1.sh | 6 +++++- .../megatron/training/tokenizer/tokenizer.py | 12 ++++++++++++ primus/modules/trainer/megatron/trainer.py | 8 -------- primus/modules/trainer/megatron/utils.py | 14 +++++++++++--- 4 files changed, 28 insertions(+), 12 deletions(-) diff --git a/examples/mlperf/config_MI355X_1x8x1.sh b/examples/mlperf/config_MI355X_1x8x1.sh index ea3b0f5df..83912ce7d 100755 --- a/examples/mlperf/config_MI355X_1x8x1.sh +++ b/examples/mlperf/config_MI355X_1x8x1.sh @@ -43,7 +43,7 @@ export NVTE_CK_USES_FWD_V3=1 export NVTE_CK_USES_BWD_V3=1 export NVTE_CK_IS_V3_ATOMIC_FP32=0 export USE_TE_SWIGLU=1 - +export TOKENIZERS_PARALLELISM=false export NCCL_CHECKS_DISABLE=1 export TORCH_NCCL_HIGH_PRIORITY=1 @@ -54,6 +54,7 @@ export ENABLE_MLLOG=1 export MLLOG_OUTPUT_FILE=/results/mlperf_output.log export MLLOG_TRAIN_LOSS_LOG_FREQ=100 export MLLOG_TARGET_EVAL_LOSS=3.3 +export TARGET_EVAL_LOSS=3.3 export MLLOG_SUBMISSION_BENCHMARK=llama31_8b export MLLOG_SUBMISSION_DIVISION=closed export MLLOG_SUBMISSION_ORG=AMD @@ -69,3 +70,6 @@ export TORCHPROF_WITH_FLOPS=0 export PROF_WARMUP_STEPS=10 #128 export PROF_ACTIVE_STEPS=6372 # export HIPBLASLT_TUNING_OVERRIDE_FILE=tuning.txt + +export NVTE_FLASH_ATTN=0 +export NVTE_FUSED_ATTN=1 \ No newline at end of file diff --git a/primus/backends/megatron/training/tokenizer/tokenizer.py b/primus/backends/megatron/training/tokenizer/tokenizer.py index d60da4c4a..a6207947a 100644 --- a/primus/backends/megatron/training/tokenizer/tokenizer.py +++ b/primus/backends/megatron/training/tokenizer/tokenizer.py @@ -7,6 +7,7 @@ """Extra Megatron tokenizers.""" import math +from collections import OrderedDict from megatron.training.arguments import ( _add_tokenizer_args as megatron_add_tokenizer_args, @@ -25,6 +26,17 @@ HuggingFaceTokenizer as _HuggingFaceTokenizer, ) + +def _hf_unique_identifiers(self): + ids = OrderedDict() + ids["class"] = f"{type(self).__module__}.{type(self).__qualname__}" + ids["vocab_size"] = self.vocab_size + return ids + + +if not hasattr(_HuggingFaceTokenizer, "unique_identifiers"): + _HuggingFaceTokenizer.unique_identifiers = property(_hf_unique_identifiers) + from primus.modules.module_utils import log_rank_0 CUSTOM_TOKENIZER_TYPES = { diff --git a/primus/modules/trainer/megatron/trainer.py b/primus/modules/trainer/megatron/trainer.py index 62af5e7b0..2820e85ba 100644 --- a/primus/modules/trainer/megatron/trainer.py +++ b/primus/modules/trainer/megatron/trainer.py @@ -754,14 +754,6 @@ def initialize_megatron( "if args.decoder_pipeline_manual_split_list is None and " + ori_code.split("if ")[-1] ) - validate_args_modified(args, args_defaults, ori_code=ori_code, new_code=new_code) - elif args.fp4 is not None: - # TODO(ruibin): Remove it when ROCm TE upgrade to 2.7.0.dev0 - from .utils import validate_args_modified - - ori_code = """raise ValueError("--fp4-format requires Transformer Engine >= 2.7.0.dev0 for NVFP4BlockScaling support.")""" - new_code = """pass""" - validate_args_modified(args, args_defaults, ori_code=ori_code, new_code=new_code) else: validate_args(args, args_defaults) diff --git a/primus/modules/trainer/megatron/utils.py b/primus/modules/trainer/megatron/utils.py index a0f9c76d2..617f24e17 100644 --- a/primus/modules/trainer/megatron/utils.py +++ b/primus/modules/trainer/megatron/utils.py @@ -202,8 +202,9 @@ def validate_manual_split(args): def validate_args_modified(*args, **kwargs): def validate_args_modifier(func, modification): import inspect + import textwrap - source = inspect.getsource(func) + source = textwrap.dedent(inspect.getsource(func)) modified_source = modification(source) namespace = {} exec(modified_source, func.__globals__, namespace) @@ -214,9 +215,16 @@ def validate_args_modifier(func, modification): assert ori_code is not None and new_code is not None, "ori_code and new_code must be provided." - megatron.training.arguments.validate_args = validate_args_modifier( - megatron.training.arguments.validate_args, lambda s: s.replace(ori_code, new_code) + target_fn = getattr( + megatron.training.arguments, "_primus_original_validate_args", + megatron.training.arguments.validate_args, ) + modified_fn = validate_args_modifier(target_fn, lambda s: s.replace(ori_code, new_code)) + if hasattr(megatron.training.arguments, "_primus_original_validate_args"): + megatron.training.arguments._primus_original_validate_args = modified_fn + else: + megatron.training.arguments.validate_args = modified_fn + megatron.training.arguments.validate_args(*args, **kwargs) From d53c428944c3d74c41c4c96fa2c1776d7e966538 Mon Sep 17 00:00:00 2001 From: vidushi8 Date: Fri, 17 Apr 2026 11:36:20 -0500 Subject: [PATCH 3/4] add fix --- primus/backends/megatron/core/fp4_utils.py | 97 ++++++++++++++++++---- 1 file changed, 80 insertions(+), 17 deletions(-) diff --git a/primus/backends/megatron/core/fp4_utils.py b/primus/backends/megatron/core/fp4_utils.py index 2beca2df7..b336e51eb 100644 --- a/primus/backends/megatron/core/fp4_utils.py +++ b/primus/backends/megatron/core/fp4_utils.py @@ -6,7 +6,7 @@ """Utility functions related to FP4 that are used throughout Megatron core""" - +import os from contextlib import nullcontext from megatron.core import parallel_state @@ -34,7 +34,7 @@ try: import primus_turbo # pylint: disable=W0611 - HAVE_TURBO = True + HAVE_TURBO = False except (ImportError, ModuleNotFoundError): # Primus-Turbo not found pass @@ -148,27 +148,90 @@ def get_fp4_context(config: TransformerConfig, layer_no: int = -1, is_init: bool elif HAVE_TE: + def _ensure_mxfp4_recipe_support(): + """Patch TE's Recipe and RecipeState to support MXFP4BlockScaling. + + TE may ship MXFP4BlockScaling, MXFP4Tensor, and MXFP4Quantizer but + lack the RecipeState.create() branch and Recipe.mxfp4() discriminator. + This adds the missing wiring so fp8_autocast works with MXFP4. + """ + from transformer_engine.common.recipe import Recipe + + if hasattr(Recipe, "mxfp4"): + return + + from typing import Optional + + import torch + + from transformer_engine.common.recipe import MXFP4BlockScaling + from transformer_engine.pytorch import quantization as te_quant + from transformer_engine.pytorch.quantization import RecipeState, get_fp4_te_dtype + from transformer_engine.pytorch.tensor.mxfp4_tensor import MXFP4Quantizer + + Recipe.mxfp4 = lambda self: isinstance(self, MXFP4BlockScaling) + + class MXFP4BlockScalingRecipeState(RecipeState): + """Configuration for MXFP4 quantization.""" + + def __init__( + self, + recipe: MXFP4BlockScaling, + *, + mode: str, + num_quantizers: int = 1, + device: Optional[torch.device] = None, + ) -> None: + self.recipe = recipe + self.mode = mode + self.num_quantizers = num_quantizers + self.dtype = get_fp4_te_dtype(recipe) + if device is None: + device = torch.device("cuda") + + def make_quantizers(self) -> list: + return [ + MXFP4Quantizer( + self.dtype, + use_hadamard=self.recipe.use_hadamard, + ) + for _ in range(self.num_quantizers) + ] + + te_quant.MXFP4BlockScalingRecipeState = MXFP4BlockScalingRecipeState + + _original_create = RecipeState.create + + @staticmethod + def _patched_create(recipe, *, mode, num_quantizers=1, device=None): + if recipe.mxfp4(): + return MXFP4BlockScalingRecipeState( + recipe, mode=mode, num_quantizers=num_quantizers, device=device, + ) + return _original_create(recipe, mode=mode, num_quantizers=num_quantizers, device=device) + + RecipeState.create = _patched_create + def get_fp4_recipe(config: TransformerConfig): """Return fp4 recipe.""" - if is_te_min_version("2.7.0.dev0"): - if config.fp4_recipe == Fp4Recipe.nvfp4: - try: - fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling() - except AttributeError: - raise ValueError( - """NVFP4BlockScaling recipe is not available in this version of - Transformer Engine. Please make sure you are using TE version - >= 2.7.0.dev0.""" - ) - else: + if config.fp4_recipe == Fp4Recipe.nvfp4: + if not is_te_min_version("2.7.0.dev0"): + raise ValueError( + "NVFP4BlockScaling requires TransformerEngine >= 2.7.0.dev0." + ) + fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling() + elif config.fp4_recipe == Fp4Recipe.mxfp4: + if not is_te_min_version("2.8.0"): raise ValueError( - "NVFP4BlockScaling is the only supported FP4 recipe. " - "Please make sure you are using a compatible TE version >= 2.7.0.dev0." + "MXFP4BlockScaling requires TransformerEngine >= 2.8.0." ) + _ensure_mxfp4_recipe_support() + fp4_recipe = transformer_engine.common.recipe.MXFP4BlockScaling() + fp4_recipe.use_hadamard = os.environ.get("NVTE_MXFP4_USE_HADAMARD", "0") == "1" else: raise ValueError( - """FP4 support requires TransformerEngine version >= 2.7.0.dev0 - for NVFP4BlockScaling.""" + f"Unsupported FP4 recipe: {config.fp4_recipe}. " + "Supported: nvfp4, mxfp4." ) return fp4_recipe From 2dc5981aaed2280d65d9218706516a470b69a44d Mon Sep 17 00:00:00 2001 From: vidushi8 Date: Mon, 27 Apr 2026 23:08:32 -0500 Subject: [PATCH 4/4] add fix --- primus/backends/megatron/core/fp4_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/primus/backends/megatron/core/fp4_utils.py b/primus/backends/megatron/core/fp4_utils.py index b336e51eb..736a3619f 100644 --- a/primus/backends/megatron/core/fp4_utils.py +++ b/primus/backends/megatron/core/fp4_utils.py @@ -225,7 +225,7 @@ def get_fp4_recipe(config: TransformerConfig): raise ValueError( "MXFP4BlockScaling requires TransformerEngine >= 2.8.0." ) - _ensure_mxfp4_recipe_support() + # _ensure_mxfp4_recipe_support() fp4_recipe = transformer_engine.common.recipe.MXFP4BlockScaling() fp4_recipe.use_hadamard = os.environ.get("NVTE_MXFP4_USE_HADAMARD", "0") == "1" else: