diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000..7cb17adc --- /dev/null +++ b/.dockerignore @@ -0,0 +1,22 @@ +# Runtime artifacts from madengine runs. The primus dockerfiles build with the repo +# root as context and copy only scripts/Primus, so none of this belongs in the build. +# hf_cache matters most: scripts/jax-maxdiffusion/run.sh caches weights there, and +# flux/wan checkpoints push it past 100GB, which docker would otherwise send to the +# daemon on every rebuild. +.git +.venv +hf_cache/ +run_directory/ +output/ +logs/ +.*_env/ +.jax-*_env/ +*.csv +*.log +*.jsonl +perf_metrics_*.jsonl +scripts/jax-maxtext/output/ +scripts/jax-maxdiffusion/output/ +**/.jax_cache/ +**/core.* +scripts/Primus/data/ diff --git a/.gitignore b/.gitignore index ec81adf9..e7e7fc2b 100644 --- a/.gitignore +++ b/.gitignore @@ -82,3 +82,18 @@ venv.bak/ !.claude/skills/mad-profiling/assets/**/*.json __pycache__/ .pytest_cache/ +# Primus checkout for JAX/MaxText discovery + image build. Use either: +# git submodule update --init scripts/Primus +# tools/fetch_primus.sh +# See benchmark/jax_maxtext/README.md for details. +scripts/Primus/ + +# Artifacts left in the repo root by a madengine run. hf_cache matters most: the +# jax-maxdiffusion wrapper points HF_HOME here so gated FLUX/WAN weights survive between +# runs, which puts 100+GB one careless `git add -A` away from being staged. +hf_cache/ +run_directory/ +.jax-*_env/ +perf_metrics_*.jsonl +scripts/jax-maxtext/output/ +scripts/jax-maxdiffusion/output/ diff --git a/benchmark/jax_maxtext/README.md b/benchmark/jax_maxtext/README.md index 6c7b856d..f7f29210 100644 --- a/benchmark/jax_maxtext/README.md +++ b/benchmark/jax_maxtext/README.md @@ -4,496 +4,383 @@ MaxText framework for ROCm is a specialized fork from upstream MaxText, designed to enable training of large language model (LLM) on AMD GPUs. By leveraging AMD Instinct™ MI300X and MI355X GPUs, MaxText delivers great scalability, performance, and resource utilization for AI workload. See the GitHub repository at [ROCm/maxtext](https://github.com/ROCm/maxtext/). -AMD provides a ready-to-use Docker image for AMD Instinct MI300X and MI355X GPUs containing essential components, including Jax, XLA, ROCm libraries, and MaxText utilities. It contains the following software components to accelerate training workloads: +AMD provides a ready-to-use Docker image for AMD Instinct MI300X and MI355X GPUs containing essential components, including Jax, XLA, ROCm libraries, and MaxText utilities. ->[!NOTE] ->Shardy is the partitioning system in JAX. The v26.5 Docker image ships JAX 0.10.0, so you now have to set `shardy=True` during the training run. You might get related errors if it's not configured correctly. See the [migration guide](https://docs.jax.dev/en/latest/shardy_jax_migration.html) for more details. -> +> **Canonical reference:** For the full Primus JAX MaxText training guide — including detailed environment setup, all supported models, multi-node networking, and the complete `primus-cli` reference — see the [Primus JAX MaxText training documentation](https://github.com/AMD-AGI/Primus/blob/main/docs/02-user-guide/jax-maxtext-training.md). This README focuses on the MAD integration layer and quick-start workflows. ->[!NOTE] -> There is a known performance regression for Mixtral-8x7b in v26.5. This is being tracked and will be addressed in a future release. +> [!NOTE] +> Shardy is the partitioning system in JAX. The v26.6 Docker image ships JAX 0.11.0, so you now have to set `shardy=True` during the training run. You might get related errors if it's not configured correctly. See the [migration guide](https://docs.jax.dev/en/latest/shardy_jax_migration.html) for more details. ->[!NOTE] +> [!NOTE] > There is a discrepancy in loss curve if you set `packing=false`. It converges at a slightly higher value than previous docker images. We can achieve the same convergence as past docker images if you set `NVTE_CK_USES_FWD_V3=0`. (i.e. using FAv2 for forward instead of FAv3). This is being tracked and will be addressed in a future release. ->[!NOTE] -> On MI355X (gfx950), RCCL's WarpSpeed feature (`RCCL_WARP_SPEED_AUTO`) — a gfx950-only optimization that is enabled by default in gfx950 builds — can cause NaN losses during training. To avoid this, set `RCCL_WARP_SPEED_AUTO=0`. For MAD-integrated benchmarking system, this is already applied automatically in the gfx950 (MI355X) environment scripts under `scripts/jax-maxtext/env_scripts/` (e.g. `gfx950_llama3_8b_env.sh`), so the benchmark scripts handle it for you. If you launch training manually on MI355X, export `RCCL_WARP_SPEED_AUTO=0` yourself. This variable is a no-op on MI300X (gfx942). +> [!NOTE] +> On MI355X (gfx950), RCCL's WarpSpeed feature (`RCCL_WARP_SPEED_AUTO`) — a gfx950-only optimization that is enabled by default in gfx950 builds — can cause NaN losses during training. To avoid this, set `RCCL_WARP_SPEED_AUTO=0`. For the MAD-integrated benchmarking system, this is applied automatically by the Primus MaxText backend when a gfx950 (MI355X) device is detected, so the benchmark scripts handle it for you. If you launch training manually on MI355X, export `RCCL_WARP_SPEED_AUTO=0` yourself. This variable is a no-op on MI300X (gfx942). -| Software component | Version | -|--------------------|----------------| -| ROCm | 7.14.0 | -| Jax | 0.10.0 | -| Python | 3.12.3 | -| Transformer Engine | 2.15.0.dev0+rocm7.15.0a20260707.72d01a0 | -| hipBLASLt | 1.4.1+cd957402 | + +| Software component | Version | +| ------------------ | ------------------------- | +| ROCm | 7.14.0 | +| Jax | 0.11.0 | +| Python | 3.12.3 | +| Transformer Engine | 2.17.0+rocm7.14.0.50a84ad | +| hipBLASLt | 1.4.1+cd957402 | ## Supported features and models + MaxText supports the following key features to train large language models efficiently: -* Transformer Engine (TE) -* Flash Attention (FA) 3, with or without input sequence packing -* GEMM tuning -* Multi-node Support -* NANOO FP8 (for MI300X) or FP8 (for MI355X) - -The following models are pre-optimized for performance on the AMD Instinct MI300X and MI355X accelerator. - -* Llama 2 7B -* Llama 2 70B -* Llama 3/3.1 8B -* Llama 3/3.1 70B -* Llama 3.1 405B -* Llama 3.3 70B -* DeepSeek-V2-lite (16B) -* Mixtral-8x7B -* Qwen3 14B -* Qwen3 30B-A3B +- Transformer Engine (TE) +- Flash Attention (FA) 3, with or without input sequence packing +- GEMM tuning +- Multi-node Support +- NANOO FP8 (for MI300X) or FP8 (for MI355X) + +The following models are pre-optimized for performance on the AMD Instinct MI300X and MI355X accelerators. + +- Llama 2 7B +- Llama 2 70B +- Llama 3/3.1 8B +- Llama 3/3.1 70B +- Llama 3.3 70B +- DeepSeek-V2-lite (16B) +- Gemma4 26B +- Gemma4 31B +- Mixtral-8x7B +- Qwen3 14B +- Qwen3 30B-A3B Note: Some models, such as Llama 3, require an external license agreement through a third party (for example, Meta). - ## System validation -If you have already validated your system, skip this step. Otherwise, please complete the following [system validation and optimization steps](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/training/prerequisite-system-validation.html#train-a-model-system-validation) to set up your system before starting training. +If you have already validated your system, skip this step. Otherwise, please complete the following [system validation and optimization steps](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/training/prerequisite-system-validation.html#train-a-model-system-validation) to set up your system before starting training. ## Environment setup -This Docker image is optimized for specific model configurations outlined below. Performance can vary for other training workloads, as AMD doesn’t validate configurations and run conditions outside those described. -For multinode, we need to make sure we have all the packages installed based on the network device we use. You can check multi node examples on how to install these packages before running the workload. You need to only do the set up below if you are using multinode with RDMA, otherwise skip this part. +This Docker image is optimized for specific model configurations outlined below. Performance can vary for other training workloads, as AMD doesn't validate configurations and run conditions outside those described. -Install the packages below for building and installing the RDMA driver: -```bash -apt install iproute2 -y -apt install -y linux-headers-"$(uname -r)" libelf-dev -apt install -y gcc make libtool autoconf librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool libibverbs-dev rdma-core strace libibmad5 libibnetdisc5 ibverbs-providers libibumad-dev libibumad3 libibverbs1 libnl-3-dev libnl-route-3-dev -``` -Please refer to your NIC manufacturer's webpage for further steps about compiling and install the RoCE driver, e.g. for Broadcom, please refer to the section **Compiling Broadcom NIC Software from Source** in [Ethernet Networking Guide for AMD Instinct MI300X GPU Clusters](https://docs.broadcom.com/doc/957608-AN2XX) - -Set the following env variables. You can again check the multinode examples on how to set these variables. -- **Master Address:** - Change `localhost` to the master node's hostname: - ```bash - export MASTER_ADDR="${MASTER_ADDR:-localhost}" - ``` - -- **Number of Nodes:** - Set the number of nodes you want to train on (e.g., 2, 4, 8): - ```bash - export NNODES="${NNODES:-1}" - ``` - -- **Node Rank:** - Set the rank of each node (0 for master, 1 for the first worker node, etc.): - ```bash - export NODE_RANK="${NODE_RANK:-0}" - ``` -- **Network Interface** - Update the network interface in the script to match your system’s network interface. - To find your network interface, run (out of container): - ```bash - ip a - ``` - Then, update the following variables in the script: - ```bash - export NCCL_SOCKET_IFNAME=ens50f0np0 - ``` -- **RDMA Interface** - First make sure that packages above are installed on all the nodes. Then set the RDMA interfaces to use for communication. - ```bash - # If using Broadcom NIC - export NCCL_IB_HCA=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 - # If using Mellanox NIC - export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_8,mlx5_9 - ``` ->[!NOTE] ->The only models supported in this workflow are those listed in the above section. -> - -This container should not be expected to provide generalized performance across all training workloads. Users should expect the container perform in the model configurations described below, but other configurations and run conditions are not validated by AMD. -Use the following instructions to set up the environment, configure the script to train models, and reproduce the benchmark results on the MI300X, MI325X, MI350X, MI355X accelerators with the Docker image. - -Users have two choices to reproduce the benchmark results using this Automation and Dashboarding repository. - -- [MAD-integrated benchmarking](#mad-integrated-benchmarking) -- [Standalone benchmarking](#standalone-benchmarking) -- [Primus benchmarking](#using-primus-cli-to-run-training-jobs-with-jax-maxtext-backend) - -Jax MaxText has also been integrated into [Primus](https://github.com/AMD-AGI/Primus), which supports multiple backends including Megatron-LM, TorchTitan, and JAX MaxText, alongside ROCm-optimized components. Users can now use the unified `primus-cli` to run training jobs with Jax MaxText backend. +For multi-node training, `primus-cli` handles node discovery, RDMA interface +selection, and environment variable propagation automatically via its Slurm mode. +If you need to customize networking (for example, selecting specific RDMA devices +or overriding the socket interface), see the +[Primus multi-node networking guide](https://github.com/AMD-AGI/Primus/blob/main/docs/04-technical-guides/multi-node-networking.md). -## MAD-integrated benchmarking +> [!NOTE] +> The only models supported in this workflow are those listed in the above section. +> This container is optimized for the model configurations described below; other +> configurations and run conditions are not validated by AMD. -Clone the ROCm Model Automation and Dashboarding (MAD) repository to a local directory and install the required packages on the host machine. +### Quick start (single-node, MAD-integrated) ```sh +# 1. Clone MAD and install dependencies git clone https://github.com/ROCm/MAD cd MAD pip install -r requirements.txt -``` - -Run models through MAD-integrated benchmarking with the following command: - -```sh -export MAD_SECRETS_HFTOKEN="your personal Hugging Face token to access gated models" -python3 tools/run_models.py --tags --keep-model-dir --live-output --timeout 28800 -``` -For example, use this command to run a performance benchmark test of the Llama 2 7B model on one GPU with bf16 data type in the host machine. +# 2. Fetch Primus (required before discovery or docker build) +bash tools/fetch_primus.sh -```sh -export MAD_SECRETS_HFTOKEN="your personal Hugging Face token to access gated models" -python3 tools/run_models.py --tags jax_maxtext_train_llama-2-7b --keep-model-dir --live-output --timeout 28800 -``` +# 3. Discover available models (auto-filters by your GPU arch) +madengine discover --tags maxtext # MaxText models +madengine discover --tags maxdiffusion # MaxDiffusion models ->[!NOTE] ->The madengine package is now available allowing for the replacement of run_models.py. -> -```sh -export MAD_SECRETS_HFTOKEN="your personal Hugging Face token to access gated models" -python3 madengine run --tags jax_maxtext_train_llama-2-7b --keep-model-dir --live-output --timeout 28800 -``` - -ROCm MAD launches a Docker container with the name `container_ci-jax_maxtext_train_llama-2-7b`. The latency and throughput reports of the model are collected in the following path: - -```sh -~/MAD/perf.csv +# 4. Run all MaxText models +export MAD_SECRETS_HFTOKEN="" +madengine run --tags maxtext --live-output --timeout 14400 ``` -#### Available models - -| model_name | -| --------------------------------------- | -| jax_maxtext_train_llama-2-7b | -| jax_maxtext_train_llama-2-70b | -| jax_maxtext_train_llama-3.1-8b | -| jax_maxtext_train_llama-3.1-70b | -| jax_maxtext_train_llama-3.1-405b | -| jax_maxtext_train_llama-3.3-70b | -| jax_maxtext_train_deepseek-v2-lite-16b | -| jax_maxtext_train_mixtral-8x7b | -| jax_maxtext_train_qwen3-14b | -| jax_maxtext_train_qwen3-30b-a3b | +For standalone (no MAD) or multi-node usage, see the sections below. -## Standalone benchmarking +Users have three paths to reproduce the benchmark results: -Download and launch the Docker image +- [MAD-integrated benchmarking](#mad-integrated-benchmarking) — recommended, auto-discovers models and manages Docker +- [Standalone benchmarking](#standalone-benchmarking) — run training manually inside a Docker container +- [Primus benchmarking](#using-primus-cli-to-run-training-jobs-with-jax-maxtext-backend) — use primus-cli directly -Use the following command to pull the Docker image from Docker Hub. +Jax MaxText has also been integrated into [Primus](https://github.com/AMD-AGI/Primus), which supports multiple backends including Megatron-LM, TorchTitan, and JAX MaxText, alongside ROCm-optimized components. MAD launches all JAX MaxText training through Primus: the MAD-integrated path uses the `scripts/jax-maxtext/run.sh` wrapper around Primus (`examples/run_pretrain.sh` with `BACKEND=MaxText`), and you can also drive `primus-cli` directly (see [Using primus-cli](#using-primus-cli-to-run-training-jobs-with-jax-maxtext-backend)). -``` -docker pull rocm/jax-training:maxtext-v26.5 -``` -### Single Node Training examples - -#### Setup ->[!NOTE] ->Please adjust the following variables based on your environment. -> +> [!NOTE] +> There is no single Primus image that covers every backend. `rocm/primus:*` ships the torch/megatron/torchtitan stack and does **not** include JAX. JAX MaxText runs on the dedicated `rocm/jax-training:maxtext-v26.6` image, which is what MAD builds from `docker/primus_maxtext.ubuntu.amd.Dockerfile`. Primus is available as a git submodule (`git submodule update --init scripts/Primus`) or via `tools/fetch_primus.sh`, which clones the pinned branch into the same gitignored `scripts/Primus` path — see below. The `scripts/jax-maxtext/` launcher and metric parser are MaxText-only (no Megatron/TorchTitan logic). -Export variables -- MAD_SECRETS_HFTOKEN is your HuggingFace token to access models, tokenizers, data. See this [page](https://huggingface.co/docs/hub/en/security-tokens) for more info. -- HF_HOME is where huggingface_hub will store local data, please refer to [Huggingface cli Document](https://huggingface.co/docs/huggingface_hub/main/en/guides/cli#hf-download) on how to download the data. If you already have downloaded/cached huggingface artifacts, set this variable to that path. Downloaded files typically get cached to a place like this: `~/.cache/huggingface`. -``` -export MAD_SECRETS_HFTOKEN= -export HF_HOME= -``` +## MAD-integrated benchmarking -Launch the Docker container. +Clone the ROCm Model Automation and Dashboarding (MAD) repository and install the required +packages on the host machine. Primus must be checked out into `scripts/Primus` **before +discovery or build**, since the JAX models are discovered from its example configs and both +`primus_`* images bake the repo into the image. You can either initialize the git submodule +(`git submodule update --init scripts/Primus`) or use `tools/fetch_primus.sh`. -``` -docker run -it --device /dev/dri --device /dev/kfd --network host --ipc host --group-add video --cap-add SYS_PTRACE --security-opt seccomp=unconfined --privileged -v $HOME:$HOME -v $HOME/.ssh:/root/.ssh -v $HF_HOME:/hf_cache -e HF_HOME=/hf_cache -e MAD_SECRETS_HFTOKEN=$MAD_SECRETS_HFTOKEN --shm-size 64G --name training_env rocm/jax-training:maxtext-v26.5 -``` +```sh +git clone https://github.com/ROCm/MAD +cd MAD +pip install -r requirements.txt -Execute the training_env container (optional if not already in the container) -``` -docker start maxtext_training -docker exec -it maxtext_training bash +# Check Primus out into scripts/Primus. Idempotent, so it is safe to re-run. +bash tools/fetch_primus.sh ``` -Clone Model Automation and Dashboarding (MAD) repo -``` -git clone https://github.com/ROCm/MAD.git -cd MAD/scripts/jax-maxtext -``` +`tools/fetch_primus.sh` clones the pinned branch (`main`). Override `PRIMUS_URL`, +`PRIMUS_REF`, or `PRIMUS_DIR` for a fork, a different branch or commit, or another location. -Run setup scripts to install libraries and datasets needed for benchmarking -``` -./jax-maxtext_benchmark_setup.sh -m -``` +> [!NOTE] +> Cloning with `--recursive` is **not required** — `tools/fetch_primus.sh` or +> `git submodule update --init scripts/Primus` are sufficient. Both Docker images take their +> framework from the base image — `/workspace/maxtext` and `/workspace/maxdiffusion`, each at +> the same commit as Primus's pin — and pin `MAXTEXT_PATH` / `MAXDIFFUSION_PATH` to it. +> Primus's own `third_party/` submodules are not needed for MAD builds; the base image's +> pre-patched copies are used instead. -Run the benchmark in quantized or unquantized mode. +This step is **not** automatic: the checkout has to exist in the docker build context +before any image is built, so neither the dockerfiles nor `scripts/jax-*/run.sh` (which +runs inside the container) can do it for you. If `scripts/Primus` is missing, discovery +finds **zero** JAX models and prints a warning naming this script, rather than failing +in a way `madengine run` reports. -``` -# For unquantized training -./jax-maxtext_benchmark_report.sh -m +In CI pipelines, set `MAD_AUTO_FETCH_PRIMUS=1` to have discovery fetch Primus +automatically when the checkout is absent (off by default): -# Or for quantized training -./jax-maxtext_benchmark_report.sh -m -q nanoo_fp8 +```sh +MAD_AUTO_FETCH_PRIMUS=1 madengine run --tags maxtext --live-output --timeout 14400 ``` -The performance results should be written to a file in the parent folder. - -### Benchmarking examples +JAX MaxText models are **auto-discovered** from the Primus MaxText experiment configs +(`scripts/Primus/examples/maxtext/configs//.yaml`). madengine walks every +`scripts//` directory at discovery time and, for any directory containing a +`get_models_json.py`, calls its `list_models()` to get one virtual model per config — this +happens for `scripts/jax-maxtext/get_models_json.py` unconditionally, with no entry needed in +the root `models.json` (there is none). A `jax-maxtext/default` model is also always +registered, pointing at the same fallback config `scripts/jax-maxtext/run.sh` uses when no +`--config_path` is given — a stable name that doesn't require knowing a specific config. -#### Example commands -1. **Single-node training with Llama 2 7B model** +Discovered tags follow the pattern `jax-maxtext/maxtext__` (the `jax-maxtext/` +prefix is the `scripts/jax-maxtext` directory), e.g. +`jax-maxtext/maxtext_MI300X_llama2_7B-bf16-pretrain` or +`jax-maxtext/maxtext_MI355X_llama2_7B-fp8-pretrain`. All of them build the +`docker/primus_maxtext` image and run through `scripts/jax-maxtext/run.sh`. Each model also +carries tags (`maxtext`, `jax`, ``, ``, ``) so you can select a +single model by its full name or a group by a shared tag. -Setup -``` -./jax-maxtext_benchmark_setup.sh -m Llama-2-7B -``` +List the available models with madengine discovery: -For unquantized training -``` -./jax-maxtext_benchmark_report.sh -m Llama-2-7B -``` - -Or for nanoo_fp8 quantized training on MI300X -``` -./jax-maxtext_benchmark_report.sh -m Llama-2-7B -q nanoo_fp8 -``` - -Or for fp8 quantized training on MI355X -``` -./jax-maxtext_benchmark_report.sh -m Llama-2-7B -q fp8 +```sh +madengine discover --tags maxtext # all MaxText models +madengine discover --tags maxdiffusion # all MaxDiffusion models +madengine discover --tags jax # all JAX models (MaxText + MaxDiffusion) +madengine discover --tags nanoo_fp8 # all nanoo_fp8 (MI300X quantized) models ``` -2. **Single-node training with Llama 2 70B model** +> [!NOTE] +> On MI300X machines, madengine automatically skips MI355X models (and vice versa) via +> the `skip_gpu_arch` field — you do not need to add `MI300X` or `MI355X` to the tags. -Setup -``` -./jax-maxtext_benchmark_setup.sh -m Llama-2-70B -``` +Run all MaxText models, all MaxDiffusion models, or both at once: -For unquantized training -``` -./jax-maxtext_benchmark_report.sh -m Llama-2-70B +```sh +export MAD_SECRETS_HFTOKEN="your personal Hugging Face token to access gated models" +madengine run --tags maxtext --live-output --timeout 14400 # all MaxText models +madengine run --tags maxdiffusion --live-output --timeout 14400 # all MaxDiffusion models +madengine run --tags jax --live-output --timeout 14400 # all JAX models (MaxText + MaxDiffusion) ``` -Or for nanoo_fp8 quantized training on MI300X -``` -./jax-maxtext_benchmark_report.sh -m Llama-2-70B -q nanoo_fp8 -``` +Run a single model by its full discovered name: -Or for fp8 quantized training on MI355X -``` -./jax-maxtext_benchmark_report.sh -m Llama-2-70B -q fp8 +```sh +export MAD_SECRETS_HFTOKEN="your personal Hugging Face token to access gated models" +madengine run --tags jax-maxtext/maxtext_MI300X_llama2_7B-bf16-pretrain --keep-model-dir --live-output --timeout 28800 ``` -3. **Single-node training with Llama 3.1 8B model** +Or the nanoo_fp8 quantized Llama 2 7B on MI300X: -Setup -``` -./jax-maxtext_benchmark_setup.sh -m Llama-3.1-8B +```sh +madengine run --tags jax-maxtext/maxtext_MI300X_llama2_7B-nanoo_fp8-pretrain --keep-model-dir --live-output --timeout 28800 ``` -For unquantized training -``` -./jax-maxtext_benchmark_report.sh -m Llama-3.1-8B -``` +> [!NOTE] +> `tools/run_models.py` remains available as a drop-in alternative to `madengine run` for the same `--tags`. -Or for nanoo_fp8 quantized training on MI300X -``` -./jax-maxtext_benchmark_report.sh -m Llama-3.1-8B -q nanoo_fp8 -``` +MAD launches a Docker container named `container_ci-`. Performance metrics +(`tokens_per_second`, `tflops`) are parsed from the training log by +`scripts/jax-maxtext/extract_maxtext_perf.py` and collected in: -Or for fp8 quantized training on MI355X -``` -./jax-maxtext_benchmark_report.sh -m Llama-3.1-8B -q fp8 +```sh +~/MAD/perf.csv ``` -4. **Single-node training with Llama 3.1 70B model** - -Setup -``` -./jax-maxtext_benchmark_setup.sh -m Llama-3.1-70B -``` +#### Available models -For unquantized training -``` -./jax-maxtext_benchmark_report.sh -m Llama-3.1-70B -``` +Model tags are generated from the Primus MaxText configs for each device, so the exact +list tracks whatever configs ship in your `scripts/Primus` checkout. List the live +set via madengine discovery (`scripts/jax-maxtext/get_models_json.py`) or by +browsing `scripts/Primus/examples/maxtext/configs/`. + +Every listed model has a bf16 variant (`jax-maxtext/maxtext__-bf16-pretrain`). Quantization is +**device-specific**: MI300X uses **NANOO FP8** (`-nanoo_fp8`) and MI355X uses **FP8** (`-fp8`) +— there is no plain-fp8 on MI300X and no nanoo_fp8 on MI355X. Every model has a bf16 +variant; the columns below show which quantized variant is also available: + + +| Model | MI300X (bf16 + …) | MI355X (bf16 + …) | +| ---------------- | ----------------- | ----------------- | +| Llama 2 7B | `-nanoo_fp8` | `-fp8` | +| Llama 2 70B | `-nanoo_fp8` | `-fp8` | +| Llama 3/3.1 8B | `-nanoo_fp8` | `-fp8` | +| Llama 3/3.1 70B | bf16 only | `-fp8` | +| Llama 3.3 70B | bf16 only | `-fp8` | +| DeepSeek-V2-lite | `-nanoo_fp8` | `-fp8` | +| Gemma4 26B | `-nanoo_fp8` | `-fp8` | +| Gemma4 31B | `-nanoo_fp8` | `-fp8` | +| Mixtral-8x7B | `-nanoo_fp8` | `-fp8` | +| Qwen3 14B | `-nanoo_fp8` | `-fp8` | +| Qwen3 30B-A3B | `-nanoo_fp8` | `-fp8` | + + +Example tags: `jax-maxtext/maxtext_MI300X_llama2_7B-bf16-pretrain`, `jax-maxtext/maxtext_MI300X_llama2_7B-nanoo_fp8-pretrain`, `jax-maxtext/maxtext_MI355X_llama2_7B-fp8-pretrain`. +(MI350X/MI325X map to MI355X/MI300X configs respectively.) + +> [!NOTE] +> Multi-node-only models (Llama 3.1 405B, Grok-1, Mixtral-8x22B) are excluded from +> single-node `jax-maxtext/` discovery. Set `JAX_MAXTEXT_INCLUDE_MULTINODE=1` to include them. + +#### MaxDiffusion models + +MaxDiffusion models are discovered the same way as MaxText — auto-discovered from +`scripts/Primus/examples/maxdiffusion/configs//.yaml` by +`scripts/jax-maxdiffusion/get_models_json.py`, tagged +`jax-maxdiffusion/maxdiffusion__` — and run through +`scripts/jax-maxdiffusion/run.sh` with the `docker/primus_maxdiffusion` image. A +`jax-maxdiffusion/default` model is also always registered, mirroring +`jax-maxtext/default` above. Use `--tags maxdiffusion` to run all of them: -Or for fp8 quantized training on MI355X -``` -./jax-maxtext_benchmark_report.sh -m Llama-3.1-70B -q fp8 +```sh +madengine run --tags maxdiffusion --live-output --timeout 14400 ``` -5. **Single-node training with Llama 3.3 70B model** -Setup -``` -./jax-maxtext_benchmark_setup.sh -m Llama-3.3-70B -``` +| Model | MI300X | MI355X | +| ------------ | ------ | ------ | +| FLUX.1-dev | bf16 | bf16 | +| WAN 2.1 1.3B | bf16 | bf16 | +| WAN 2.1 14B | bf16 | bf16 | -For unquantized training -``` -./jax-maxtext_benchmark_report.sh -m Llama-3.3-70B -``` -Or for fp8 quantized training on MI355X -``` -./jax-maxtext_benchmark_report.sh -m Llama-3.3-70B -q fp8 -``` +Example tags: `jax-maxdiffusion/maxdiffusion_MI300X_flux_dev-pretrain`, `jax-maxdiffusion/maxdiffusion_MI300X_wan2.1_1.3b-pretrain`. -6. **Single-node training with DeepSeek2 16B model** +## Standalone benchmarking -Setup -``` -./jax-maxtext_benchmark_setup.sh -m DeepSeek-V2-lite -``` +Download and launch the Docker image -For unquantized training -``` -./jax-maxtext_benchmark_report.sh -m DeepSeek-V2-lite -``` +Use the following command to pull the Docker image from Docker Hub. -Or for nanoo_fp8 quantized training on MI300X ``` -./jax-maxtext_benchmark_report.sh -m DeepSeek-V2-lite -q nanoo_fp8 +docker pull rocm/jax-training:maxtext-v26.6 ``` -Or for fp8 quantized training on MI355X -``` -./jax-maxtext_benchmark_report.sh -m DeepSeek-V2-lite -q fp8 -``` +### Single Node Training examples -7. **Single-node training with Mixtral-8x7B model** +#### Setup -Setup -``` -./jax-maxtext_benchmark_setup.sh -m Mixtral-8x7B -``` +> [!NOTE] +> Please adjust the following variables based on your environment. -For unquantized training -``` -./jax-maxtext_benchmark_report.sh -m Mixtral-8x7B -``` +Export variables -Or for nanoo_fp8 quantized training on MI300X -``` -./jax-maxtext_benchmark_report.sh -m Mixtral-8x7B -q nanoo_fp8 -``` +- MAD_SECRETS_HFTOKEN is your HuggingFace token to access models, tokenizers, data. See this [page](https://huggingface.co/docs/hub/en/security-tokens) for more info. +- HF_HOME is where huggingface_hub will store local data, please refer to [Huggingface cli Document](https://huggingface.co/docs/huggingface_hub/main/en/guides/cli#hf-download) on how to download the data. If you already have downloaded/cached huggingface artifacts, set this variable to that path. Downloaded files typically get cached to a place like this: `~/.cache/huggingface`. -Or for fp8 quantized training on MI355X ``` -./jax-maxtext_benchmark_report.sh -m Mixtral-8x7B -q fp8 +export MAD_SECRETS_HFTOKEN= +export HF_HOME= ``` -8. **Single-node training with Qwen3 14B model** +Launch the Docker container. -Setup -``` -./jax-maxtext_benchmark_setup.sh -m Qwen3-14B +```bash +docker run -it \ + --device /dev/dri --device /dev/kfd \ + --network host --ipc host --group-add video \ + --cap-add SYS_PTRACE --security-opt seccomp=unconfined --privileged \ + -v $HOME:$HOME -v $HOME/.ssh:/root/.ssh \ + -v $HF_HOME:/hf_cache -e HF_HOME=/hf_cache \ + -e MAD_SECRETS_HFTOKEN=$MAD_SECRETS_HFTOKEN \ + --shm-size 64G --name training_env \ + rocm/jax-training:maxtext-v26.6 ``` -For unquantized training -``` -./jax-maxtext_benchmark_report.sh -m Qwen3-14B -``` +Execute the training_env container (optional if not already in the container) -Or for nanoo_fp8 quantized training on MI300X ``` -./jax-maxtext_benchmark_report.sh -m Qwen3-14B -q nanoo_fp8 +docker start training_env +docker exec -it training_env bash ``` -Or for fp8 quantized training on MI355X -``` -./jax-maxtext_benchmark_report.sh -m Qwen3-14B -q fp8 -``` +Inside the container, the Primus repository (with the MaxText backend) is available at +`/workspace/Primus`. Run training with `primus-cli`; **direct** mode runs in the current +container. Configs live under `examples/maxtext/configs//` where `` is +`MI300X` or `MI355X`. -9. **Single-node training with Qwen3 30B-A3B model (MoE)** +```bash +cd /workspace/Primus -Setup -``` -./jax-maxtext_benchmark_setup.sh -m Qwen3-30B-A3B +# Unquantized (bf16), e.g. Llama 2 7B on MI300X +# Note: RCCL_WARP_SPEED_AUTO=0 is auto-set by Primus on MI355X (gfx950). +./primus-cli direct -- train pretrain \ + --config examples/maxtext/configs/MI300X/llama2_7B-bf16-pretrain.yaml ``` -For unquantized training -``` -./jax-maxtext_benchmark_report.sh -m Qwen3-30B-A3B -``` +For quantized training, replace `-bf16-` in the config name with `-fp8-` (MI355X) +or `-nanoo_fp8-` (MI300X): -Or for nanoo_fp8 quantized training on MI300X -``` -./jax-maxtext_benchmark_report.sh -m Qwen3-30B-A3B -q nanoo_fp8 -``` +```bash +# nanoo_fp8 on MI300X +./primus-cli direct -- train pretrain \ + --config examples/maxtext/configs/MI300X/llama2_7B-nanoo_fp8-pretrain.yaml -Or for fp8 quantized training on MI355X -``` -./jax-maxtext_benchmark_report.sh -m Qwen3-30B-A3B -q fp8 +# fp8 on MI355X +./primus-cli direct -- train pretrain \ + --config examples/maxtext/configs/MI355X/llama2_7B-fp8-pretrain.yaml ``` +The same pattern applies to every supported model (`llama2_70B`, `llama3_8B`, `llama3_70B`, +`llama3.3_70B`, `deepseek_v2_16B`, `gemma4_26B`, `gemma4_31B`, +`mixtral_8x7B`, `qwen3_14B`, `qwen3_30B_A3B`). See the +[Using primus-cli](#using-primus-cli-to-run-training-jobs-with-jax-maxtext-backend) +section for container and Slurm modes. ### Multi-Node Training examples -Note: these scripts will launch the docker and execute the benchmark, so **please run them outside of any docker**. - -The examples below use Slurm for running on multiple nodes. The unified multinode benchmark script accepts a configuration file that specifies the model and training parameters. - -#### Running Multi-Node Training - -To run multi-node training, use the following command: - -```bash -sbatch -N jax_maxtext_multinode_benchmark.sh [docker_image] -``` - -**Parameters:** -- ``: Number of nodes to use for training (e.g., 2, 4, 8) -- ``: Path to the YAML configuration file containing model and training parameters -- `[docker_image]`: (Optional) Docker image to use. If not specified, defaults to `rocm/jax-training:maxtext-v26.5` -**Configuration files** are available in the `scripts/jax-maxtext/env_scripts/` directory for different models and GPU architectures: +Multi-node training is launched through the unified `primus-cli` in Slurm mode. +The standalone MAD multinode launcher and the per-model `env_scripts/*.yml` +configuration files have been retired; model/precision/parallelism settings now +live in the Primus MaxText experiment configs under +`examples/maxtext/configs//--pretrain.yaml` +(bundled in the `rocm/jax-training` image at `/workspace/Primus`). -For MI300X (gfx942): -- `llama2_7b.yml` - Llama 2 7B -- `llama2_70b.yml` - Llama 2 70B -- `llama3_8b.yml` - Llama 3 8B -- `llama3_70b.yml` - Llama 3 70B -- `qwen3_14b.yml` - Qwen3 14B -- `qwen3_30b_a3b.yml` - Qwen3 30B-A3B +See the [Using primus-cli](#using-primus-cli-to-run-training-jobs-with-jax-maxtext-backend) +section below for direct, container, and Slurm examples. The general form for a +multi-node run is: -For MI355X (gfx950): -- `gfx950_llama2_7b.yml` - Llama 2 7B -- `gfx950_llama2_70b.yml` - Llama 2 70B -- `gfx950_llama3_8b.yml` - Llama 3 8B -- `gfx950_llama3_70b.yml` - Llama 3 70B -- `gfx950_llama3.1_405b.yml` - Llama 3.1 405B -- `gfx950_qwen3_14b.yml` - Qwen3 14B -- `gfx950_qwen3_30b_a3b.yml` - Qwen3 30B-A3B - -#### Example Commands - -1. **Multi-node training with Llama 2 7B model on 2 nodes:** ```bash -sbatch -N 2 jax_maxtext_multinode_benchmark.sh env_scripts/llama2_7b.yml +# From /workspace/Primus (or a cloned Primus checkout) +# RCCL_WARP_SPEED_AUTO=0 is auto-set by Primus on MI355X (gfx950). +./primus-cli --config my_maxtext_config.yaml slurm srun -N \ + -- train pretrain --config examples/maxtext/configs//--pretrain.yaml ``` -2. **Multi-node training with Llama 2 70B model on 4 nodes with custom image:** -```bash -sbatch -N 4 jax_maxtext_multinode_benchmark.sh env_scripts/llama2_70b.yml rocm/jax-training:maxtext-v26.5 -``` - -3. **Multi-node training with Llama 3 8B model on 2 nodes:** -```bash -sbatch -N 2 jax_maxtext_multinode_benchmark.sh env_scripts/llama3_8b.yml -``` - -4. **Multi-node training with Llama 3 70B model on 8 nodes:** -```bash -sbatch -N 8 jax_maxtext_multinode_benchmark.sh env_scripts/llama3_70b.yml -``` - -5. **Multi-node training with Llama 3.1 405B model on MI355X (gfx950) with 8 nodes:** -```bash -sbatch -N 8 jax_maxtext_multinode_benchmark.sh env_scripts/gfx950_llama3.1_405b.yml -``` +where `` is `MI300X` or `MI355X`, `` is one of the MaxText +configs (e.g. `llama2_7B`, `llama2_70B`, `llama3_8B`, `llama3_70B`, +`gemma4_26B`, `gemma4_31B`, `mixtral_8x7B`, +`qwen3_14B`, `qwen3_30B_A3B`), and `` is `bf16`, +`fp8` (MI355X), or `nanoo_fp8` (MI300X), e.g. `llama2_7B-bf16-pretrain.yaml`. ## Using primus-cli to run training jobs with Jax MaxText backend **Clone the Primus repository** + ``` -git clone https://github.com/AMD-AIG-AIMA/Primus.git +git clone https://github.com/AMD-AGI/Primus.git cd Primus git checkout main git submodule update --init third_party/maxtext/ @@ -506,27 +393,26 @@ For detailed usage of primus-cli, please refer to [Primus CLI User Guide](https: Here are some examples of using primus-cli to run training jobs with Jax MaxText backend. Direct Mode: Running the training directly on current host or within an existing docker container. + ```bash -# On MI355X (gfx950), disable RCCL WarpSpeed to avoid NaN losses (no-op on MI300X) -export RCCL_WARP_SPEED_AUTO=0 -./primus-cli direct -- train pretrain --config examples/maxtext/configs/MI355X/llama2_7B-pretrain.yaml +# RCCL_WARP_SPEED_AUTO=0 is auto-set by Primus on MI355X (gfx950); no action needed. +./primus-cli direct -- train pretrain --config examples/maxtext/configs/MI355X/llama2_7B-bf16-pretrain.yaml ``` -Container Mode: execute in Docker/Podman containers +Container Mode: execute in Docker/Podman containers. You **must** pass `--image` because +the default Primus image (`rocm/primus`) is the PyTorch stack and does not include JAX. + ```bash -# On MI355X (gfx950), disable RCCL WarpSpeed to avoid NaN losses (no-op on MI300X) -export RCCL_WARP_SPEED_AUTO=0 -./primus-cli container --image rocm/jax-training:maxtext-v26.5 \ - -- train pretrain --config examples/maxtext/configs/MI355X/llama2_7B-pretrain.yaml +./primus-cli container --image rocm/jax-training:maxtext-v26.6 \ + -- train pretrain --config examples/maxtext/configs/MI355X/llama2_7B-bf16-pretrain.yaml ``` Slurm Mode: execute distributed training on a Slurm cluster + ```bash -# On MI355X (gfx950), disable RCCL WarpSpeed to avoid NaN losses (no-op on MI300X) -export RCCL_WARP_SPEED_AUTO=0 # Use a custom config file, where you can specify the docker image and set environment variables. ./primus-cli --config my_maxtext_config.yaml slurm srun -N 8 \ - -- train pretrain --config examples/maxtext/configs/MI355X/llama2_7B-pretrain.yaml + -- train pretrain --config examples/maxtext/configs/MI355X/llama2_7B-bf16-pretrain.yaml ``` ## Profiling with JAX XPlane Profiler @@ -545,13 +431,14 @@ upload_all_profiler_results=True # Save all GPU profiles (not just GPU0) ``` **Choosing step counts:** + - `steps` should be > `skip_first_n_steps_for_profiler` + `profiler_steps` (e.g., `steps=12` with skip=2, profile=5 gives 5 warmup + 5 profiled + 2 cooldown) - `skip_first_n_steps_for_profiler=2` skips step 0 (compilation) and step 1 (warmup) - `profiler_steps=5` is typically enough; more steps = larger `.xplane.pb` files ### Profiling with MAD/madengine -The model YAML configs under `scripts/jax-maxtext/env_scripts/` already include a `profiler` key (set to `""` by default). To enable profiling when running through MAD or madengine, edit the YAML config for your model and set the profiler fields: +The Primus MaxText experiment configs (`examples/maxtext/configs//--pretrain.yaml` in `/workspace/Primus`) already include a `profiler` key under `overrides` (set to `""` by default). To enable profiling when running through MAD or madengine, edit the `overrides` block of the config for your model and set the profiler fields: ```yaml profiler: "xplane" @@ -565,10 +452,10 @@ Then run the benchmark as usual: ```bash # Via madengine -python3 madengine run --tags jax_maxtext_train_llama-3.1-8b --keep-model-dir --live-output --timeout 28800 +madengine run --tags jax-maxtext/maxtext_MI300X_llama3_8B-bf16-pretrain --keep-model-dir --live-output --timeout 28800 # Or via run_models.py -python3 tools/run_models.py --tags jax_maxtext_train_llama-3.1-8b --keep-model-dir --live-output --timeout 28800 +python3 tools/run_models.py --tags jax-maxtext/maxtext_MI300X_llama3_8B-bf16-pretrain --keep-model-dir --live-output --timeout 28800 ``` Profile output will be written under the `base_output_directory` specified in the YAML (see [Output Structure](#output-structure) below). Use `--keep-model-dir` so the container's output directory is preserved after the run. @@ -579,8 +466,8 @@ Profile output will be written under the `base_output_directory` specified in th #!/bin/bash set -e -IMAGE="$1" # Docker image, e.g. rocm/jax-training:maxtext-v26.5 -TAG="$2" # Short tag for output folder, e.g. v26.5_llama2_7b +IMAGE="$1" # Docker image, e.g. rocm/jax-training:maxtext-v26.6 +TAG="$2" # Short tag for output folder, e.g. v26.6_llama2_7b PROFILE_DIR="/path/to/profiles/${TAG}" mkdir -p "${PROFILE_DIR}" @@ -647,6 +534,7 @@ tensorboard --logdir /path/to/profiles//profile --port 6006 Navigate to **Profile > Trace Viewer** in the TensorBoard UI. **Tips:** + - Zoom into a single training step (skip the first profiled step as it may have residual warmup) - Look at individual GPU streams to see compute/RCCL overlap @@ -660,9 +548,20 @@ Navigate to **Profile > Trace Viewer** in the TensorBoard UI. ## Profiling with rocprofv3 If you need to collect a trace and the JAX profiler isn't working then you can use rocprofv3 as a temporary workaround like this: + ``` rocprofv3 --hip-trace --kernel-trace --memory-copy-trace --rccl-trace --output-format pftrace -d ./v3_traces -- python3 app.py ``` -- Just replace `python3 app.py` with any command line command that you want to run such as `./jax-maxtext_benchmark_report.sh -m Llama-2-7B`. + +- Just replace `python3 app.py` with any command line command that you want to run such as `./primus-cli direct -- train pretrain --config examples/maxtext/configs/MI300X/llama2_7B-bf16-pretrain.yaml` (run from `/workspace/Primus`). - You can set the directory where you want the .json traces to be saved using `-d ` -- The resulting traces can be opened in perfetto: https://ui.perfetto.dev/ +- The resulting traces can be opened in perfetto: [https://ui.perfetto.dev/](https://ui.perfetto.dev/) + +## Related documentation + +- [Primus JAX MaxText training guide](https://github.com/AMD-AGI/Primus/blob/main/docs/02-user-guide/jax-maxtext-training.md) — canonical reference for environment setup, models, and training options +- [Primus CLI reference](https://github.com/AMD-AGI/Primus/blob/main/docs/02-user-guide/cli-reference.md) — full `primus-cli` command-line reference +- [End-to-end training recipes](https://github.com/AMD-AGI/Primus/blob/main/docs/02-user-guide/end-to-end-training-recipes.md) — complete config inventory and step-by-step recipes +- [MaxText parameters](https://github.com/AMD-AGI/Primus/blob/main/docs/03-configuration-reference/maxtext-parameters.md) — YAML config field reference +- [Multi-node networking](https://github.com/AMD-AGI/Primus/blob/main/docs/04-technical-guides/multi-node-networking.md) — RDMA, NCCL, and Slurm networking setup +- [Release notes](https://github.com/AMD-AGI/Primus/blob/main/docs/01-getting-started/release-notes.md) — full software stack details for each image tag diff --git a/docker/jax_maxtext.ubuntu.amd.Dockerfile b/docker/jax_maxtext.ubuntu.amd.Dockerfile deleted file mode 100644 index 90165c9b..00000000 --- a/docker/jax_maxtext.ubuntu.amd.Dockerfile +++ /dev/null @@ -1,33 +0,0 @@ -# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} -# MIT License -# -# Copyright (c) Advanced Micro Devices, Inc. -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# -ARG BASE_DOCKER=rocm/jax-training:maxtext-v26.5 -FROM $BASE_DOCKER - -USER root -ENV WORKSPACE_DIR=/workspace -RUN mkdir -p $WORKSPACE_DIR -WORKDIR $WORKSPACE_DIR - -# record configuration for posterity -RUN pip3 list diff --git a/docker/primus_maxdiffusion.ubuntu.amd.Dockerfile b/docker/primus_maxdiffusion.ubuntu.amd.Dockerfile new file mode 100644 index 00000000..244d48d5 --- /dev/null +++ b/docker/primus_maxdiffusion.ubuntu.amd.Dockerfile @@ -0,0 +1,84 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +############################################################################### +# +# MIT License +# +# Copyright (c) Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# + +# Primus JAX/MaxDiffusion launcher image for MAD (WAN 2.1 + FLUX.1-dev): bakes the Primus +# repo onto a JAX training base so scripts/jax-maxdiffusion/run.sh can run +# `train pretrain --config ...`. +# +# The base image owns the maxdiffusion stack: maxdiffusion is installed and patched at +# /workspace/maxdiffusion, at the same commit as Primus's third_party/maxdiffusion pin. So +# this image runs no setup_maxdiffusion_env.sh and installs no maxdiffusion deps. For a +# base without the stack, use the setup-script build from git history before this commit. +# +# Check Primus out first with tools/fetch_primus.sh. It is gitignored here and baked from +# the build context, which keeps git auth for a private repo out of the build. That script +# initializes no submodules: third_party/maxdiffusion at the same commit is unpatched, and +# run_pretrain.sh would select it over the base's tree if MAXDIFFUSION_PATH were ever unset. +# +# Build from the repo root, as madengine does for dockerfile paths containing "primus": +# docker build -f docker/primus_maxdiffusion.ubuntu.amd.Dockerfile . + +# madengine passes the base via docker_build_arg, which is how the v26.6 sweep put both +# maxtext and maxdiffusion on one unified CI image so their numbers share a toolchain. +ARG BASE_DOCKER=rocm/jax-training:maxtext-v26.6 +FROM $BASE_DOCKER + +USER root +ENV WORKSPACE_DIR=/workspace +# The Primus repo root, not /workspace: run.sh resolves examples/ relative to it. +ENV PRIMUS_ROOT=/workspace/Primus +# Pin the base's patched tree; run_pretrain.sh would otherwise default to +# $PRIMUS_ROOT/third_party/maxdiffusion and insert it at sys.path[0]. +ENV MAXDIFFUSION_PATH=/workspace/maxdiffusion +# Transformer Engine must load only its JAX extension (torch is present too). +ENV NVTE_FRAMEWORK=jax +RUN mkdir -p $WORKSPACE_DIR +WORKDIR $WORKSPACE_DIR + +LABEL mad.launcher=primus + +# The base may ship /workspace/Primus as a git clone, and COPY cannot replace a +# .git directory with a submodule checkout's .git file. +RUN rm -rf /workspace/Primus +COPY scripts/Primus/ /workspace/Primus/ + +RUN test -f /workspace/Primus/examples/run_pretrain.sh +RUN test -d /workspace/Primus/primus/backends/maxdiffusion \ + || (echo "ERROR: Primus checkout lacks primus/backends/maxdiffusion; use Primus main branch." >&2 && exit 1) + +# Prove the base's stack is really there, so a wrong base fails the build instead +# of step 0 of a training run. The patch fixes a segfault on TE import order. +RUN python3 -c "import maxdiffusion, os; print('maxdiffusion ->', os.path.dirname(maxdiffusion.__file__))" +RUN grep -q "preload before Transformer Engine" /workspace/maxdiffusion/src/maxdiffusion/train_utils.py \ + || (echo "ERROR: /workspace/maxdiffusion is missing or lacks the TF-preload patch." >&2 && exit 1) + +# Primus's own requirements, not maxdiffusion's, which the base already covers. +# Installed here rather than on every run: run.sh sets PRIMUS_SKIP_PIP=1 so a +# launch stays off the network. On this base it adds loguru. +RUN pip3 install --no-cache-dir -r /workspace/Primus/requirements-maxdiffusion.txt + +RUN pip3 list 2>/dev/null || true diff --git a/docker/primus_maxtext.ubuntu.amd.Dockerfile b/docker/primus_maxtext.ubuntu.amd.Dockerfile new file mode 100644 index 00000000..75e48934 --- /dev/null +++ b/docker/primus_maxtext.ubuntu.amd.Dockerfile @@ -0,0 +1,75 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +############################################################################### +# +# MIT License +# +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# + +# Primus JAX/MaxText launcher image for MAD: bakes the Primus repo onto a JAX training +# base so scripts/jax-maxtext/run.sh can run `train pretrain --config ...`. +# +# The base image owns the MaxText stack: it is installed at /workspace/maxtext, at the same +# commit as Primus's third_party/maxtext pin, so no submodule is needed here. This mirrors +# how primus_maxdiffusion takes maxdiffusion from the base. +# +# Check Primus out first with tools/fetch_primus.sh. It is gitignored here and baked from +# the build context, which keeps git auth for a private repo out of the build. +# +# Build from the repo root, as madengine does for dockerfile paths containing "primus": +# docker build -f docker/primus_maxtext.ubuntu.amd.Dockerfile . + +# madengine passes the base via docker_build_arg, which is how the v26.6 sweep put both +# maxtext and maxdiffusion on one unified CI image so their numbers share a toolchain. +ARG BASE_DOCKER=rocm/jax-training:maxtext-v26.6 +FROM $BASE_DOCKER + +USER root +ENV WORKSPACE_DIR=/workspace +# The Primus repo root, not /workspace: run.sh resolves examples/ relative to it. +ENV PRIMUS_ROOT=/workspace/Primus +# Pin the base's tree; prepare.py would otherwise default to +# $PRIMUS_ROOT/third_party/maxtext, which this image does not ship. +ENV MAXTEXT_PATH=/workspace/maxtext +RUN mkdir -p "$WORKSPACE_DIR" +WORKDIR $WORKSPACE_DIR + +LABEL mad.launcher=primus + +# The base may ship /workspace/Primus as a git clone, and COPY cannot replace a +# .git directory with a submodule checkout's .git file. +RUN rm -rf /workspace/Primus +COPY scripts/Primus/ /workspace/Primus/ + +RUN test -f /workspace/Primus/examples/run_pretrain.sh +RUN test -f /workspace/Primus/requirements-jax.txt + +# Prove the base's stack is really there, so a wrong base fails the build instead +# of step 0 of a training run. +RUN test -f /workspace/maxtext/pyproject.toml \ + || (echo "ERROR: no MaxText at /workspace/maxtext. Use a base that bakes it, or point MAXTEXT_PATH at a checkout." >&2 && exit 1) + +# Installed here rather than on every run: run.sh sets PRIMUS_SKIP_PIP=1 so a +# launch stays off the network. On this base it adds loguru, wandb, pre-commit. +RUN pip3 install --no-cache-dir -r /workspace/Primus/requirements-jax.txt + +RUN pip3 list 2>/dev/null || true diff --git a/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile b/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile index 74bca4d3..0f788de7 100644 --- a/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile +++ b/docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile @@ -214,6 +214,20 @@ ARG MORI_WHEEL_URL= # build MoRI for a different arch than the rest of the image. ARG MORI_GPU_ARCHS=${BUILD_GPU_TARGETS} ARG MORI_VERSION=1.2.0 +# NOTE (UMBP unit-test build guard): mori's src/umbp/CMakeLists.txt calls +# `add_subdirectory(tests)` UNCONDITIONALLY, ignoring the top-level BUILD_TESTS +# option (which setup.py already defaults to OFF, and which correctly gates +# tests/cpp). Those UMBP tests use gtest_discover_tests() in its default +# POST_BUILD discovery mode, which EXECUTES each freshly-linked test binary at +# build time to enumerate test cases. The binaries link the ROCm/HIP runtime +# (libamdhip64 / libhsa-runtime64 / librocm_smi64 / libhsakmt), so they abort on +# startup inside `docker build`, which has NO GPU (RUN steps cannot pass +# --gpus/--device, so /dev/kfd and /dev/dri are absent). The empty output then +# breaks CMake's string(JSON ...) in GoogleTest/ParseTestList.cmake and fails the +# amd_mori wheel build. Nothing is "pulled from the build host" -- CMake just +# tries to RUN the test binaries during the build. Re-gated behind BUILD_TESTS in +# the RUN block below; idempotent no-op once upstream fixes it (already fixed on +# mori main, still present at the pinned MORI_COMMIT). ARG MORI_SRC_DIR=/sgl-workspace/mori ENV MORI_GPU_ARCHS=${MORI_GPU_ARCHS} \ @@ -253,6 +267,12 @@ RUN set -e; \ git clone --depth 1 --branch "${MORI_BRANCH}" "${MORI_REPO}" "${MORI_SRC_DIR}"; \ fi; \ cd "${MORI_SRC_DIR}" && git submodule update --init --recursive || true; \ + if grep -qE '^add_subdirectory\(tests\)$' src/umbp/CMakeLists.txt; then \ + sed -i 's|^add_subdirectory(tests)$|if(BUILD_TESTS)\n add_subdirectory(tests)\nendif()|' src/umbp/CMakeLists.txt; \ + echo "MORI_UMBP_TESTS_GUARDED (add_subdirectory(tests) now gated by BUILD_TESTS)"; \ + else \ + echo "MORI_UMBP_TESTS_GUARD_SKIPPED (unguarded pattern not found; upstream may have fixed it)"; \ + fi; \ pip install --no-cache-dir "cmake<4.0" setuptools wheel pybind11 ninja; \ pip install --no-build-isolation --no-cache-dir --force-reinstall .; \ fi diff --git a/docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile b/docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile new file mode 100644 index 00000000..2143c0c6 --- /dev/null +++ b/docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile @@ -0,0 +1,351 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# +# ============================================================================= +# vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile +# GLM-5.1-FP8 (MLA + DeepSeek Sparse Attention) MoRI-EP WideEP disagg image. +# PER-MODEL image, isolated from the base vllm_disagg_inference Dockerfile +# (which stays pinned to the DeepSeek-V3 / R1 stack), so each model can pin its +# own vLLM/AITER/MoRI without disturbing the others. +# +# ALL connectors in one image: moriio (TP + MoRI-EP wideEP) + rixl (NIXL TP + +# DeepEP wideEP). = the fullsource MoRI stack, plus a UCX/RIXL/rocSHMEM/DeepEP +# transport layer gated by --build-arg WITH_NIXL (default 0 = MoRI-EP only). +# +# docker build -f docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile \ +# -t /vllm-disagg:glmv5.1 . +# export DOCKER_IMAGE_NAME=/vllm-disagg:glmv5.1 +# +# STATUS (GLM-5.1-FP8 on this stack): 2P/2D EP16 and 4P/4D EP32 both validated ON THE +# PINS BELOW, from an image built off this file unmodified. Long-context accuracy is +# fixed in-source (vLLM #47766); the 4P/4D EP32 token corruption is fixed by the MoRI +# combine() original-topk change (vLLM 623fdc946b). +# NIAH, 3 seeds each: 2P/2D 2k 9.7/10, 8k 9.7/10 | 4P/4D 2k 9.7/10, 8k 9.3/10 +# The 9/10s are seed variance, not length collapse: across 12 cells every score is 9 or +# 10 with no pattern by topology or context length. Read it as ~93-97% retrieval. +# A latency/throughput sweep (8192/1024 and 16384/1024 at con 32/64) ran on both +# topologies with no failed requests and no GPU faults. Numbers are deliberately not +# quoted here: they are one configuration on one cluster and would be read as a spec. +# Re-measure on your own hardware before relying on any figure. +# (BASE_IMAGE is a gated nightly; override --build-arg BASE_IMAGE=...; vLLM compile ~30-60 min.) +# ============================================================================= +# Builds the GLM-5.1 runtime stack by applying component pins ON TOP of a +# purpose-built ROCm/vLLM/MoRI base, cloning each overridden source from public Git +# (no local build-contexts): +# +# - BASE: rocm/vllm-dev:ci_base-dedbf6be8b1afa17a6220473b9c8c98242ac1c03 +# (ROCm + torch nightly). The stages below OVERRIDE the base's vLLM/MoRI/AITER +# with the pins we validate for GLM DSA. +# - MoRI -> built from ROCm/MoRI @ 624002c897a3 (BUILD_UMBP=OFF). Bumped from +# 42e895472b08 alongside the AITER bump; co-validated 2P/2D EP16. +# - AITER -> raviguptaamd/aiter @ 624e43586b (ROCm/aiter 1d872fa + the gfx942 gqa64 +# decode fix filed upstream as ROCm/aiter#4957) from source + flydsl 0.3.1; +# stale JIT wiped. (#47766 keeps persistent MLA ON -> aiter native gqa64 fold.) +# - vLLM -> COMPILED from raviguptaamd/vllm @ 094820b5d (branch +# glm5.1-dsa-wideEP_on_d626108b = upstream d626108b + the 10 GLM DSA commits, +# incl. the EP32 combine-topk fix). Full compile: a different commit than the +# base's, so a .py-only overlay would be ABI-mismatched. +# - RDMA fix (expandable_segments:False x2 + HSA_ENABLE_IPC_MODE_LEGACY=0) is NOT baked +# here — it lives in scripts/vllm_dissag/connectors/.env and the launcher +# forwards it via docker -e. ROCm 7.2.3 cannot dmabuf-export VMM memory, else MoRI +# RegisterRdmaMemoryRegion EFAULTs (errno 14) on the first disagg WRITE. +# - vllm-router (vllm-project/router PR#181 = DP-rank round-robin + 2P2D KV-notify +# dpfix) built in -> no external router binary needed. +# - MoRIIO disagg fixes (#39276 notify, #41751 LL split, DP-rank hash-failsafe) native. +# +# THIS IMAGE IS THE CONTRACT. The MAD side (scripts/vllm_dissag) is catalog/config only +# and ships NO runtime .py patchers, so EVERY GLM DSA source fix must be carried +# in-source HERE, by VLLM_REF below plus the MoRI/AITER pins. Serving GLM-5.1-FP8 from an +# image built off an older vLLM ref is unsupported and fails SILENTLY: it boots, then +# produces garbage output or stalls the disagg KV transfer, with nothing to fall back on +# (measured: rocm/pytorch-private:glm-dockerimage-built-09072026 scored NIAH 2k 0/10). +# If you need a fix, move the pin and rebuild — do not re-add runtime patchers to MAD. +# ============================================================================= + +ARG BASE_IMAGE=rocm/vllm-dev:ci_base-dedbf6be8b1afa17a6220473b9c8c98242ac1c03 +FROM ${BASE_IMAGE} + +ENTRYPOINT [] +WORKDIR /app + +ARG GFX_COMPILATION_ARCH="gfx942" +ARG PYTORCH_ROCM_ARCH="gfx942" +ARG MAX_JOBS=32 +# NIXL/RIXL transport for the rixl connector. GLM-5.1 is served over MoRI-EP, so this +# stack is dead weight here. Default 0 => lean MoRI-EP-only image, which is how the +# validated image was built. Set --build-arg WITH_NIXL=1 to get the rixl connector. +ARG WITH_NIXL=0 +ARG NIC_COMPILATION_ARCH="cx7" + +# ----------------------------------------------------------------------------- +# 1. MoRI: replace the base's bundled MoRI with the commit GLM-5.1 DSA wideEP was +# validated on, ROCm/MoRI @ 624002c897a3. NOTE this is NOT tag v1.2.1 that the base +# vllm_disagg_inference Dockerfile pins. It carries the EP/RDMA correctness fixes +# this recipe needs plus the ROCm-7.2.3 dmabuf registration path used by the +# connector .env (expandable_segments:False). MoRI is JIT-built, so this swaps the +# JIT sources the kernels compile from at runtime. +# BUILD CONFIG: match the cookbook build — MORI_GPU_ARCHS=gfx942, BUILD_UMBP=OFF, +# DEFAULT NIC backends. Do NOT pass USE_IONIC=OFF / USE_BNXT=OFF: disabling NIC +# backends produced a MoRI that deadlocked at the cross-node EP all-to-all init. +# ----------------------------------------------------------------------------- +ARG MORI_REPO=https://github.com/ROCm/mori.git +# 624002c897a3: validated MoRI tip for GLM DSA WideEP disagg (bumped from 42e895472b08, +# which predates the recv-sizing fixes the VLLM_MORI_* knobs need). The base's bundled +# amd_mori regressed GLM DSA (GPU fault on the aiter DSA decode kernel), so we build +# from source at this pinned commit by DEFAULT; WITH_MORI_BUILD=0 falls back for debug. +ARG WITH_MORI_BUILD=1 +ARG MORI_REF=624002c897a3 +ENV MORI_GPU_ARCHS=gfx942 +# Newer MoRI added the UMBP subsystem which requires gRPC (grpcpp/grpcpp.h) not +# present in this base; UMBP is unrelated to the EP dispatch/combine kernels, so +# disable it to avoid pulling in a gRPC build dependency. +ENV BUILD_UMBP=OFF BUILD_UMBP_SPDK=OFF +# Build/install COMMAND (not the version) matches dist-inf-cookbook +# Dockerfile.vllm.mori121_shareable: +# `BUILD_UMBP=OFF pip install .` (default build isolation). apt/pip build tooling kept +# for bases that lack it; harmless where already present. +RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ + sed -i 's|http://|https://|g' /etc/apt/sources.list.d/*.list 2>/dev/null || true && \ + apt-get update && apt-get install -y --no-install-recommends \ + git build-essential cmake ninja-build ccache libssl-dev pkg-config curl ca-certificates && \ + pip install meson==0.64.0 "pybind11[global]" tqdm prettytable && \ + mkdir -p /app && \ + if [ "${WITH_MORI_BUILD}" != "1" ]; then \ + python3 -c "import mori, mori.io, mori.ops; print('MoRI (bundled) OK at', mori.__path__[0])" && \ + echo "MORI_REF=BUNDLED (base amd_mori, WITH_MORI_BUILD=0)" >> /app/versions.txt ; \ + else \ + pip uninstall -y amd_mori amd-mori amd-mori-nightly mori 2>/dev/null || true && \ + rm -rf /tmp/mori-src && \ + git clone --recursive "${MORI_REPO}" /tmp/mori-src && \ + cd /tmp/mori-src && git checkout "${MORI_REF}" && git submodule update --init --recursive && \ + BUILD_UMBP=OFF pip install . && \ + python3 -c "import mori, mori.io, mori.ops; print('MoRI OK at', mori.__path__[0])" && \ + echo "MORI_REF=${MORI_REF}@$(git -C /tmp/mori-src rev-parse HEAD)" >> /app/versions.txt && \ + rm -rf /tmp/mori-src ; \ + fi + +# ----------------------------------------------------------------------------- +# 2. AITER: built from source at raviguptaamd/aiter @ 624e43586b (WITH_AITER_BUILD=1). +# That is ROCm/aiter 1d872fa plus a 7-line fix (filed upstream as ROCm/aiter#4957): +# newer aiter claims native gfx942 support for gqa64 fp8 decode and routes it to +# mla_a8w8_qh64_qseqlen1_gqaratio64_v3_ps, which GPU-faults; the fix lets gqa64 fall +# through to aiter's capture-safe persistent view-fold, so cudagraph decode is kept. +# The fork ref is TEMPORARY — revert AITER_REPO to ROCm/aiter once #4957 merges. +# --build-arg WITH_AITER_BUILD=0 falls back to the bundled aiter for debugging. +# ----------------------------------------------------------------------------- +ARG AITER_REPO=https://github.com/raviguptaamd/aiter.git +ARG WITH_AITER_BUILD=1 +ARG AITER_REF=624e43586b +RUN if [ "${WITH_AITER_BUILD}" != "1" ]; then \ + echo "AITER: using BUNDLED base aiter (WITH_AITER_BUILD=0)" && \ + python3 -c "import importlib.metadata as m; print('aiter (bundled)', m.version('amd-aiter'))" && \ + echo "AITER_REF=BUNDLED (base amd-aiter, WITH_AITER_BUILD=0)" >> /app/versions.txt ; \ + else \ + echo "Compiling STOCK AITER (no fork) from ${AITER_REPO}@${AITER_REF}" && \ + rm -rf /tmp/aiter-src && \ + git clone --recursive "${AITER_REPO}" /tmp/aiter-src && \ + cd /tmp/aiter-src && git checkout "${AITER_REF}" && \ + git submodule update --init --recursive && \ + (pip uninstall -y amd_aiter amd-aiter aiter 2>/dev/null || true) && \ + pip install --no-build-isolation --no-deps -v . && \ + pip install --no-deps -U "flydsl==0.3.1" && \ + echo "AITER_REF=${AITER_REF}@$(git rev-parse HEAD) (aiter + ROCm/aiter#4957 gqa64 fix)" >> /app/versions.txt && \ + rm -rf /tmp/aiter-src && \ + rm -rf /opt/vllm_cache/aiter_jit /root/.aiter && echo "cleared stale AITER JIT cache" ; \ + fi + +# ----------------------------------------------------------------------------- +# 3. vLLM: compile from source at the GLM-5.1 DSA wideEP branch. Full source compile +# (the base ships a different commit). The MoRIIO disagg fixes (#39276 notify, +# #41751 LL split, DP-rank hash-failsafe) AND the GLM DSA fixes are native in this +# branch, so no runtime patcher is needed — and none exists in MAD, which is why +# this ref is a hard requirement rather than a preference. Override VLLM_REF to +# rebuild a different commit; build only committed commits (no working-tree edits). +# ----------------------------------------------------------------------------- +# VLLM_REPO/REF are a PUBLIC GitHub repo + branch. Override to your own vLLM fork/branch. +ARG VLLM_REPO=https://github.com/raviguptaamd/vllm.git +# REPRODUCIBILITY: this default is now a COMMIT, not a branch name, so two builds ship +# the same engine. /app/versions.txt in the built image records the sha that was +# resolved. models.yaml's decode-cudagraph table predates this pin: those numbers were +# measured on d723eb305e and have not been re-measured on it. +# +# What the ref carries: the 10 GLM DSA commits (per-req-ctx metadata key #47766, DSA +# indexer KV transfer, invalid-token sentinel, MoRI EP sizing knobs, and the EP32 +# combine() original-topk fix) on top of upstream vLLM d626108b (2026-08-20). +ARG VLLM_REF=094820b5deeb1b93733586ca8942589e385a25dc +ENV VLLM_TARGET_DEVICE=rocm \ + PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} \ + MAX_JOBS=${MAX_JOBS} +RUN rm -rf /tmp/vllm-src && \ + git clone "${VLLM_REPO}" /tmp/vllm-src && \ + cd /tmp/vllm-src && git checkout "${VLLM_REF}" && \ + echo "VLLM_REF=${VLLM_REF}@$(git rev-parse HEAD)" >> /app/versions.txt && \ + pip uninstall -y vllm 2>/dev/null || true && \ + pip install --no-deps --no-build-isolation -v . && \ + python3 -c "import vllm; print('vLLM', vllm.__version__, 'from', vllm.__file__)" && \ + rm -rf /tmp/vllm-src + +# Cross-check MoRI + AITER survived the vLLM install (no silent downgrade). +RUN python3 - <<'PYEOF' +from importlib.metadata import version as v, PackageNotFoundError +def get(names): + for n in names: + try: return v(n) + except PackageNotFoundError: pass + return None +av = get(("amd-aiter", "amd_aiter", "aiter")) +# Assert presence, not a version string: we pin aiter by commit and the reported version +# varies by build. Do NOT `import aiter` here — it pulls torch->amdsmi->libamd_smi.so, +# which is not loadable in the no-GPU build sandbox. +assert av, "AITER missing after vLLM install (expected bundled 0.1.19 or source-built ref)" +import mori, mori.io, mori.ops +print("Post-vLLM check OK: AITER", av, "present + MoRI importable") +PYEOF + +# ----------------------------------------------------------------------------- +# 4. vllm-router (DP-rank round-robin + MoRIIO connector) — built in, so NO +# external vllm-router binary is needed (leave ROUTER_BINARY unset). +# Source = upstream vllm-project/router main @ 1d10e71 (2026-08-18) plus ONE +# commit: the 2P2D KV-notify fix (moriio_dp_size + effective_dp_size() + +# remote_dp_rank_override). REQUIRED: without it the 2P2D EP=16 run reproducibly +# wedges with "remote blocks never arrived" deferred-write expiries (decode +# notify targets the wrong DP rank). The round-robin DP-rank fix that used to +# ride along here is now upstream (prefill_dp_round_robin), so it is no longer +# carried locally. Pinned to a sha, not the branch name, so rebuilds are +# reproducible; branch is raviguptaamd/router:ravgupta/dp-roundrobin-on-tip. +# Pinned Rust toolchain (>=1.88: router deps time/home require rustc 1.88). +# ----------------------------------------------------------------------------- +ARG ROUTER_REPO=https://github.com/raviguptaamd/router.git +ARG ROUTER_REF=82dc9811af17412e6e24b5942a5486bc502df23a +ARG RUST_TOOLCHAIN=1.88.0 +RUN if ! command -v cargo >/dev/null 2>&1; then \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain "${RUST_TOOLCHAIN}"; \ + fi && \ + export PATH="/root/.cargo/bin:${PATH}" && \ + rm -rf /tmp/vllm-router-src && \ + git clone --filter=blob:none "${ROUTER_REPO}" /tmp/vllm-router-src && \ + cd /tmp/vllm-router-src && git checkout "${ROUTER_REF}" && \ + cargo build --release && \ + install -m 755 target/release/vllm-router /usr/local/bin/vllm-router && \ + vllm-router --help 2>&1 | grep -q moriio && \ + echo "VLLM_ROUTER_REF=${ROUTER_REPO}@${ROUTER_REF}@$(git -C /tmp/vllm-router-src rev-parse HEAD)" >> /app/versions.txt && \ + rm -rf /tmp/vllm-router-src + +# ----------------------------------------------------------------------------- +# 4b. WITH_NIXL=1: UCX + RIXL(+nixlbench) + rocSHMEM + DeepEP from source, +# so the rixl connector (NIXL TP + DeepEP wideEP) is present. Default is 0 +# (MoRI-EP only). Single guarded RUN so WITH_NIXL=0 skips it (no layers, no cost). +# ----------------------------------------------------------------------------- +ENV _ROCM_DIR=/opt/rocm \ + _UCX_SOURCE=https://github.com/ROCm/ucx.git \ + _UCX_BRANCH=da3fac2a \ + _UCX_INSTALL_DIR=/usr/local/ucx/ \ + _RIXL_SOURCE=https://github.com/ROCm/RIXL.git \ + _RIXL_BRANCH=f33a5599 \ + _RIXL_INSTALL_DIR=/usr/local/RIXL/install \ + _NIXLBENCH_INSTALL_DIR=/usr/local/RIXL +RUN if [ "${WITH_NIXL}" != "1" ]; then \ + echo "WITH_NIXL=${WITH_NIXL}: skipping UCX/RIXL/rocSHMEM/DeepEP (MoRI-EP + base DeepEP only)"; \ + else set -e && \ + echo "WITH_NIXL=1: building UCX + RIXL + rocSHMEM + DeepEP" && \ + apt-get update && apt-get install -y \ + autoconf automake libtool autogen pkg-config m4 gcc make \ + librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool \ + libibverbs-dev rdma-core strace libgflags-dev \ + libaio-dev liburing-dev libcpprest-dev libgrpc-dev libgrpc++-dev \ + libprotobuf-dev protobuf-compiler-grpc wget && \ + pip install meson==0.64.0 "pybind11[global]" pyyaml && \ + # UCX + cd /tmp && git clone "${_UCX_SOURCE}" && cd ucx && git checkout "${_UCX_BRANCH}" && \ + ./autogen.sh && mkdir -p build && cd build && \ + ../configure --prefix="${_UCX_INSTALL_DIR}" --with-rocm="${_ROCM_DIR}" \ + --disable-go --disable-java --disable-assertions --enable-mt && \ + make -j && make install && \ + # googletest (RIXL dep) + cd /tmp && wget -q https://github.com/google/googletest/archive/refs/tags/v1.14.0.tar.gz && \ + tar -xzf v1.14.0.tar.gz && cd googletest-1.14.0 && mkdir -p build && cd build && \ + cmake -DBUILD_SHARED_LIBS=on .. && make -j && make install && \ + # RIXL + python bindings + cd /tmp && git clone "${_RIXL_SOURCE}" && cd RIXL && git checkout "${_RIXL_BRANCH}" && \ + meson setup build/ --prefix="${_RIXL_INSTALL_DIR}" -Ducx_path="${_UCX_INSTALL_DIR}" \ + -Ddisable_gds_backend=true -Dcudapath_inc="${_ROCM_DIR}/include" -Dcudapath_lib="${_ROCM_DIR}/lib" && \ + cd build && ninja && ninja install && cd /tmp/RIXL && \ + pip install --config-settings=setup-args="-Dcudapath_inc=${_ROCM_DIR}/include" \ + --config-settings=setup-args="-Dcudapath_lib=${_ROCM_DIR}/lib" \ + --config-settings=setup-args="-Ducx_path=${_UCX_INSTALL_DIR}" \ + --config-settings=setup-args="-Ddisable_gds_backend=true" . && \ + # rocSHMEM (DeepEP dep) + cd /tmp && git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems.git && \ + cd rocm-systems && git sparse-checkout set --cone projects/rocshmem && git checkout develop && \ + mkdir -p /tmp/rocshmem-build && cd /tmp/rocshmem-build && \ + /tmp/rocm-systems/projects/rocshmem/scripts/build_configs/all_backends \ + -DUSE_EXTERNAL_MPI=OFF -DGPU_TARGETS="${GFX_COMPILATION_ARCH}" && \ + # DeepEP (build develop against the installed vLLM/torch) + cd /tmp && git clone https://github.com/ROCm/DeepEP.git && cd DeepEP && \ + PYTORCH_ROCM_ARCH="${GFX_COMPILATION_ARCH}" CFLAGS="-O3 -fPIC" \ + CXXFLAGS="-O3 -fPIC --offload-arch=${GFX_COMPILATION_ARCH}" HIP_CXX_FLAGS="-O3 -fPIC" \ + python3 setup.py --variant rocm --nic "${NIC_COMPILATION_ARCH}" build develop && \ + echo "WITH_NIXL build complete" >> /app/versions.txt && \ + rm -rf /tmp/ucx /tmp/googletest-1.14.0 /tmp/v1.14.0.tar.gz /tmp/rocm-systems /tmp/rocshmem-build; \ + fi +ENV LD_LIBRARY_PATH="/usr/local/ucx/lib:/usr/local/lib:/usr/local/RIXL/install/lib:${LD_LIBRARY_PATH}" \ + PATH="/usr/local/ucx/bin:${PATH}" + +# ----------------------------------------------------------------------------- +# 5. Cache locations (structural: WHERE the JIT/compile caches live in the image). +# These are the mount target for the launcher's persistent host JIT cache. +# ----------------------------------------------------------------------------- +# The image ships NO runtime recipe / tuning / platform ENV. By design, everything +# run-tunable is applied at launch, so this image stays a clean binary/library artifact +# and the same image serves any model/cluster without a rebuild: +# - model-serving recipe (KV_BLOCK_SIZE, KV_CACHE_DTYPE, *_CUDAGRAPH_MODE, *_MORI_BACKEND, +# GPU_MEMORY_UTILIZATION, KV_CACHE_MEMORY_BYTES, VLLM_ROCM_USE_AITER_MLA, ...) +# -> scripts/vllm_dissag/models.yaml (per-model env:, so dense vs MoE differ) +# - ROCm-7.2.3 GPU-RDMA platform env (expandable_segments:False x2, MORI_GPU_ARCHS, +# HSA_ENABLE_IPC_MODE_LEGACY=0, HSA_NO_SCRATCH_RECLAIM) and the MoRI/RDMA fabric +# tuning (MORI_RDMA_TC/SL, MORI_IB_GID_INDEX, MORI_NUM_QP_PER_PE, VLLM_MORIIO_*, ...) +# -> scripts/vllm_dissag/connectors/.env (cluster-editable, no rebuild) +# The slurm launcher forwards both via `docker -e` (platform env must reach PID 1 - +# PyTorch reads alloc-conf at import). Running this image WITHOUT the launcher: set the +# vars you need yourself (see connectors/moriio.env + models.yaml for the values). +ENV AITER_JIT_DIR=/opt/vllm_cache/aiter_jit \ + VLLM_CACHE_ROOT=/opt/vllm_cache/vllm \ + TRITON_CACHE_DIR=/opt/vllm_cache/triton \ + COMGR_CACHE_DIR=/opt/vllm_cache/comgr + +# ----------------------------------------------------------------------------- +# 6. CRITICAL: scrub build-time MoRI JIT state. The `import mori` verification +# steps above compile/lock MoRI EP kernels under /root/.mori/jit on THIS build +# host, leaving stale .hsaco.lock files (ep_internode_v1, ep_internode_v1ll, ...). +# At runtime on the cluster, MoriAll2AllManager finds those locks, waits on a +# build-in-progress whose owner PID is long gone, and DEADLOCKS at ep:0 init. +# A clean image ships /root/.mori empty -> runtime compiles fresh. +# Clearing these makes the from-source image boot clean on 2P2D/4P4D. +# ----------------------------------------------------------------------------- +RUN rm -rf /root/.mori /tmp/mori_jit_* && mkdir -p /root/.mori && \ + echo "JIT_SCRUBBED: /root/.mori + /tmp/mori_jit_* cleared at build end" >> /app/versions.txt + +RUN cat /app/versions.txt 2>/dev/null | tail -20 || true diff --git a/models.json b/models.json index c81bec8e..fe51488c 100644 --- a/models.json +++ b/models.json @@ -1,23 +1 @@ -[ - { - "name": "primus_train", - "dockerfile": "docker/primus", - "dockercontext": ".", - "scripts": "scripts/primus_train", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "", - "tags": [ - "training", - "primus", - "megatron", - "pretrain" - ], - "args": "", - "env": { - "HSA_NO_SCRATCH_RECLAIM": "1", - "PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32": "1", - "NVTE_CK_IS_V3_ATOMIC_FP32": "1" - } - } -] +[] diff --git a/scripts/Primus b/scripts/Primus index 30cf4514..dc3f4ba1 160000 --- a/scripts/Primus +++ b/scripts/Primus @@ -1 +1 @@ -Subproject commit 30cf451497f24aebb5e031e8d4850d5f449433b4 +Subproject commit dc3f4ba18a7314974b54e58dbe4927e9baf489a3 diff --git a/scripts/common/agentic_lib.sh b/scripts/common/agentic_lib.sh new file mode 100755 index 00000000..5799b29a --- /dev/null +++ b/scripts/common/agentic_lib.sh @@ -0,0 +1,491 @@ +#!/bin/bash +# Agentic trace-replay helpers (aiperf driver) for the disaggregated P/D launchers. +# +# Ported from SemiAnalysisAI/InferenceX benchmarks/benchmark_lib.sh (agentic +# section). Drives aiperf's `inferencex-agentx-mvp` scenario against a live +# disagg endpoint (SGLang router :2322 or vLLM proxy) and writes aiperf +# artifacts + an aggregate JSON + plots. Does NOT touch parse_to_csv.py/perf.csv. +# +# The toolkit (customized aiperf + agentx scenario + WEKA trace loaders + +# result aggregator) is installed at RUN TIME into a throwaway uv venv that is +# kept OUT of the inference server's site-packages, from pinned commits. +# +# Consumed env (set by the caller / launcher): +# MODEL_PATH path the server was launched with (required) +# MODEL served model name aiperf sends (auto-resolved if unset) +# AGENTIC_PORT endpoint port (default 2322 for sglang router) +# AGENTIC_CONC session-tree concurrency (default 16) +# DURATION measurement window seconds (default 120; scenario min 900) +# RESULT_DIR output dir (default /run_logs/$SLURM_JOB_ID) +# MODEL_PREFIX model family key for trace-loader default (optional) +# WEKA_LOADER_OVERRIDE pin a specific trace loader (optional) +# DRY_RUN=1 print the assembled command + resolved values, then exit 0 +# +# Suite mode (scripts/common/benchmark_agentic_suite.sh) additionally consumes: +# AGENTIC_CONFIG path to agentic.yaml (a serving block + workloads LIST) +# AGENTIC_WORKLOAD single-workload shorthand (run just this named entry) +# SUITE_CORPUS_DIR persistent per-node cache for generated profile corpora +# Per-workload knobs the suite sets before build_replay_cmd (all default to the +# legacy hf/inferencex behavior when unset, so the single-workload path is +# unchanged / byte-identical): +# WL_SOURCE hf (default) | profile | corpus +# CORPUS_DIR for WL_SOURCE=profile/corpus: the weka_trace corpus dir +# AGENTIC_MAX_CONTEXT_LENGTH per-workload --max-context-length (else MAX_MODEL_LEN) +# +# Pins (Phase 0 blocker): concrete commits, overridable by env. Bump by editing +# these defaults after re-validating against a fresh smoke run. +INFERENCEX_REPO="${INFERENCEX_REPO:-https://github.com/SemiAnalysisAI/InferenceX.git}" +INFERENCEX_PIN="${AGENTIC_UTILS_PIN:-ef8a17ecf0c3679dc12020eddab3c1a36d285b58}" # InferenceX main @ 2026-07-25 +AIPERF_PIN="${AIPERF_PIN:-0d2aa0572ac685943d38c580675c4a61023581d3}" # utils/aiperf submodule (cquil11/aiperf-agentx-v1.0) + +set -o pipefail + +agentic_log() { echo "[agentic] $*"; } +agentic_warn() { echo "[agentic][WARN] $*" >&2; } +agentic_err() { echo "[agentic][ERROR] $*" >&2; } +agentic_die() { agentic_err "$*"; exit 1; } + +# Location of the generic AgentX generator/verifier/loader + profiles. +AGENTX_DIR="${AGENTX_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/agentx}" +# Persistent per-node corpus cache (regenerate once per profile+seed). +SUITE_CORPUS_DIR="${SUITE_CORPUS_DIR:-${TMPDIR:-/tmp}/agentx_corpora}" + +# -------------------------------------------------------------------------- +# Runtime install (isolated uv venv, pinned sources) +# -------------------------------------------------------------------------- +AGENTIC_RUNTIME_DIR="${AGENTIC_RUNTIME_DIR:-${TMPDIR:-/tmp}/mad-agentic-${SLURM_JOB_ID:-$$}}" +INFMAX_WS="${INFMAX_WS:-${AGENTIC_RUNTIME_DIR}/InferenceX}" +AIPERF_VENV="${AIPERF_VENV:-${AGENTIC_RUNTIME_DIR}/venv}" +AIPERF_UV_INSTALL_DIR="${AIPERF_UV_INSTALL_DIR:-${AGENTIC_RUNTIME_DIR}/uv/bin}" +AIPERF_UV_CACHE_DIR="${AIPERF_UV_CACHE_DIR:-${AGENTIC_RUNTIME_DIR}/uv-cache}" +AIPERF_PYTHON="${AIPERF_VENV}/bin/python" +AIPERF_CLI="${AIPERF_VENV}/bin/aiperf" +AIPERF_HF_CLI="${AIPERF_VENV}/bin/hf" +AIPERF_DEPS_READY=0 +AIPERF_FAILED_REQUEST_THRESHOLD="${AIPERF_FAILED_REQUEST_THRESHOLD:-0.10}" + +# Pin the uv release rather than tracking "latest", and gate the remote install +# behind an explicit opt-in (supply chain: don't silently pipe an unpinned +# installer into sh). An already-installed uv on PATH is always preferred. +AGENTIC_UV_VERSION="${AGENTIC_UV_VERSION:-0.5.11}" + +ensure_agentic_uv() { + if command -v uv >/dev/null 2>&1; then + AIPERF_UV_BIN="$(command -v uv)" + return + fi + AIPERF_UV_BIN="${AIPERF_UV_INSTALL_DIR}/uv" + if [ ! -x "$AIPERF_UV_BIN" ]; then + if [ "${AGENTIC_ALLOW_UV_INSTALL:-0}" != "1" ]; then + agentic_die "uv not found on PATH and no cached uv at $AIPERF_UV_BIN. Install uv (>= $AGENTIC_UV_VERSION) or set AGENTIC_ALLOW_UV_INSTALL=1 to permit the pinned remote install." + fi + mkdir -p "$AIPERF_UV_INSTALL_DIR" + curl -LsSf "https://astral.sh/uv/${AGENTIC_UV_VERSION}/install.sh" | UV_INSTALL_DIR="$AIPERF_UV_INSTALL_DIR" sh + fi + [ -x "$AIPERF_UV_BIN" ] || agentic_die "uv installation did not create $AIPERF_UV_BIN" +} + +_clone_inferencex_pinned() { + if [ -d "$INFMAX_WS/.git" ]; then + agentic_log "InferenceX checkout already present at $INFMAX_WS" + return + fi + command -v git >/dev/null 2>&1 || { apt-get update && apt-get install -y git; } + mkdir -p "$(dirname "$INFMAX_WS")" + agentic_log "Cloning InferenceX @ $INFERENCEX_PIN (+ aiperf submodule @ $AIPERF_PIN)" + git clone --filter=blob:none "$INFERENCEX_REPO" "$INFMAX_WS" + git -C "$INFMAX_WS" checkout --quiet "$INFERENCEX_PIN" + # Pull the aiperf submodule at the commit InferenceX pins (holds the + # inferencex-agentx-mvp scenario + semianalysis_cc_traces_weka loaders). + git -C "$INFMAX_WS" submodule update --init --recursive utils/aiperf + local got + got="$(git -C "$INFMAX_WS/utils/aiperf" rev-parse HEAD 2>/dev/null)" + if [ "$got" != "$AIPERF_PIN" ]; then + agentic_log "WARN: aiperf submodule at $got, expected $AIPERF_PIN (InferenceX pin drift)" + fi +} + +install_agentic_deps() { + [ "$AIPERF_DEPS_READY" = "1" ] && return + ensure_agentic_uv + _clone_inferencex_pinned + + # aiperf must NOT share site-packages with the inference server: installing + # it into SGLang/vLLM's Python can upgrade fastapi/starlette/transformers + # under the live server. Build a throwaway venv instead. + rm -rf "$AIPERF_VENV" + mkdir -p "$AIPERF_UV_CACHE_DIR" + UV_CACHE_DIR="$AIPERF_UV_CACHE_DIR" "$AIPERF_UV_BIN" venv --python "$(command -v python3)" "$AIPERF_VENV" + UV_CACHE_DIR="$AIPERF_UV_CACHE_DIR" "$AIPERF_UV_BIN" pip install --python "$AIPERF_PYTHON" \ + -r "$INFMAX_WS/utils/agentic-benchmark/requirements.txt" \ + -e "$INFMAX_WS/utils/aiperf" \ + "datasets>=4.7.0" "huggingface_hub[cli]>=0.25.0" urllib3 requests + + [ -x "$AIPERF_CLI" ] && [ -x "$AIPERF_HF_CLI" ] || \ + agentic_die "isolated aiperf environment incomplete at $AIPERF_VENV" + AIPERF_DEPS_READY=1 +} + +# -------------------------------------------------------------------------- +# Trace source resolution (loader name is pure; download is retried) +# -------------------------------------------------------------------------- +# Sets TRACE_LOADER + TRACE_DATASET; does no I/O so DRY_RUN can call it. +# For WL_SOURCE=profile the trace comes from a locally generated weka_trace +# corpus (--custom-dataset-type weka_trace --input-file), NOT an HF download. +# For WL_SOURCE=hf (default) the behavior is unchanged / byte-identical. +resolve_trace_loader() { + if [ "${WL_SOURCE:-hf}" = "profile" ] || [ "${WL_SOURCE:-hf}" = "corpus" ]; then + TRACE_LOADER="" + TRACE_DATASET="" + TRACE_SOURCE_FLAG="--custom-dataset-type weka_trace --input-file ${CORPUS_DIR}" + return + fi + local default_loader + case "${MODEL_PREFIX:-}" in + dsv4*|deepseek*|DeepSeek*|glm5*|minimaxm3*) + default_loader="semianalysis_cc_traces_weka_062126" ;; # 1M-ctx families: full corpus + *) + default_loader="semianalysis_cc_traces_weka_062126_256k" ;; # shorter-ctx: 256k-capped + esac + TRACE_LOADER="${WEKA_LOADER_OVERRIDE:-$default_loader}" + case "$TRACE_LOADER" in + semianalysis_cc_traces_weka_062126) TRACE_DATASET="semianalysisai/cc-traces-weka-062126" ;; + semianalysis_cc_traces_weka_062126_256k) TRACE_DATASET="semianalysisai/cc-traces-weka-062126-256k" ;; + semianalysis_cc_traces_weka_061526) TRACE_DATASET="semianalysisai/cc-traces-weka-061526" ;; + semianalysis_cc_traces_weka_061526_256k) TRACE_DATASET="semianalysisai/cc-traces-weka-061526-256k" ;; + *) agentic_die "unknown WEKA_LOADER_OVERRIDE='$TRACE_LOADER' (see resolve_trace_loader)";; + esac + # Tier 2: an hf workload with a local filter is trimmed once into a weka_trace + # dir and replayed through --input-file (like the profile path); without a + # filter, keep the byte-identical --public-dataset behavior. + if _hf_filter_active; then + TRACE_SOURCE_FLAG="--custom-dataset-type weka_trace --input-file ${CORPUS_DIR}" + else + TRACE_SOURCE_FLAG="--public-dataset $TRACE_LOADER" + fi +} + +# True when a Tier 2 local filter is configured for the current hf workload. +_hf_filter_active() { + [ -n "${WL_FILTER_MAX_ISL:-}" ] || [ -n "${WL_FILTER_MAX_TURNS:-}" ] || [ -n "${WL_FILTER_SAMPLE:-}" ] +} + +# Deterministic cache-key dir name for a filtered hf corpus: hf__. +_hf_corpus_key() { + local loader="$1" filter_json + filter_json="$(printf '{"max_isl": %s, "max_turns": %s, "sample": %s}' \ + "${WL_FILTER_MAX_ISL:-null}" "${WL_FILTER_MAX_TURNS:-null}" "${WL_FILTER_SAMPLE:-null}")" + printf 'hf_%s_%s' "$loader" "$(printf '%s' "$filter_json" | sha1sum | cut -c1-8)" +} + +# Download the dataset into the shared HF cache with retries (3 attempts, +# 900s each, backoff). Fails the run only after all attempts. +resolve_trace_source() { + resolve_trace_loader + agentic_log "Trace loader: $TRACE_LOADER ($TRACE_DATASET)" + local attempts="${AGENTIC_TRACE_DL_ATTEMPTS:-3}" + local per_timeout="${AGENTIC_TRACE_DL_TIMEOUT:-900}" + local i backoff=30 + for ((i = 1; i <= attempts; i++)); do + agentic_log "trace download attempt $i/$attempts (timeout ${per_timeout}s)" + if timeout "$per_timeout" "$AIPERF_HF_CLI" download --repo-type dataset "$TRACE_DATASET"; then + return 0 + fi + agentic_log "attempt $i failed" + [ "$i" -lt "$attempts" ] && { sleep "$backoff"; backoff=$((backoff * 2)); } + done + agentic_die "trace download failed after $attempts attempts ($TRACE_DATASET)" +} + +# -------------------------------------------------------------------------- +# Profile corpus materialization (WL_SOURCE=profile) + context compatibility +# -------------------------------------------------------------------------- +# Generate a weka_trace corpus for one workload profile into +# $SUITE_CORPUS_DIR/ if absent, then verify it as a PRE-GATE (die unless +# N/N axes within band). Sets CORPUS_DIR to the materialized corpus. +# $1 = workload name $2 = resolved profile JSON file (WL_PROFILE_FILE) +materialize_corpus() { + local name="$1" profile_json="$2" + local py="${AIPERF_PYTHON:-python3}" + [ -f "$profile_json" ] || agentic_die "materialize_corpus($name): profile JSON not found: $profile_json" + CORPUS_DIR="${SUITE_CORPUS_DIR}/${name}" + if [ -d "$CORPUS_DIR" ] && [ -n "$(ls -A "$CORPUS_DIR" 2>/dev/null)" ] && [ "${SUITE_CORPUS_FORCE:-0}" != "1" ]; then + agentic_log "corpus for '$name' already present at $CORPUS_DIR (SUITE_CORPUS_FORCE=1 to regen)" + else + agentic_log "generating corpus for '$name' -> $CORPUS_DIR" + rm -rf "$CORPUS_DIR" + "$py" "$AGENTX_DIR/gen_agentx_profile.py" --profile "$profile_json" --out-dir "$CORPUS_DIR" \ + || agentic_die "corpus generation failed for '$name'" + fi + agentic_log "verifying corpus '$name' against its profile (pre-gate)" + local out + out="$("$py" "$AGENTX_DIR/verify_agentx_profile.py" --profile "$profile_json" --corpus "$CORPUS_DIR")" || { + echo "$out"; agentic_die "corpus '$name' failed conformance pre-gate (not N/N)"; } + echo "$out" +} + +# Tier 2: download an hf weka corpus once, filter/trim it locally to fit the +# model, and materialize a per-session weka_trace dir for --input-file replay. +# Sets CORPUS_DIR. Dies on download exhaustion or an empty (over-aggressive) +# filter result. $1 = workload name $2 = resolved loader (WL_LOADER) +materialize_hf_corpus() { + local name="$1" loader="$2" + local py="${AIPERF_PYTHON:-python3}" + CORPUS_DIR="${SUITE_CORPUS_DIR}/$(_hf_corpus_key "$loader")" + if [ -d "$CORPUS_DIR" ] && [ -n "$(ls -A "$CORPUS_DIR" 2>/dev/null)" ] && [ "${SUITE_CORPUS_FORCE:-0}" != "1" ]; then + agentic_log "filtered hf corpus for '$name' already present at $CORPUS_DIR (SUITE_CORPUS_FORCE=1 to regen)" + return 0 + fi + agentic_log "downloading + filtering hf corpus for '$name' ($loader) -> $CORPUS_DIR" + resolve_trace_source # 3x retry/backoff download into the shared HF cache; sets TRACE_DATASET + local raw + raw="$("$py" - "$TRACE_DATASET" <<'PY' +import sys +from huggingface_hub import snapshot_download +print(snapshot_download(sys.argv[1], repo_type="dataset")) +PY +)" || agentic_die "could not locate downloaded corpus for $TRACE_DATASET" + rm -rf "$CORPUS_DIR" + "$py" "$AGENTX_DIR/filter_weka_corpus.py" --input "$raw" --out-dir "$CORPUS_DIR" \ + ${WL_FILTER_MAX_ISL:+--max-isl "$WL_FILTER_MAX_ISL"} \ + ${WL_FILTER_MAX_TURNS:+--max-turns "$WL_FILTER_MAX_TURNS"} \ + ${WL_FILTER_SAMPLE:+--sample "$WL_FILTER_SAMPLE"} \ + || agentic_die "filter too aggressive or corpus unreadable for '$name' (0 sessions?)" +} + +# Smallest power of two >= n (used to size --max-context-length from the ISL tail). +_next_pow2() { + local n="$1" p=1 + while [ "$p" -lt "$n" ]; do p=$((p * 2)); done + echo "$p" +} + +# Compare the workload's ISL tail against the served max_model_len and set the +# per-workload --max-context-length (AGENTIC_MAX_CONTEXT_LENGTH). WARNs (or skips +# when AGENTIC_STRICT_CONTEXT=1) if the model window cannot hold the ISL tail. +# $1 = workload name $2 = ISL tail (clamp-hi / P99) $3 = served max_model_len +# Sets AGENTIC_MAX_CONTEXT_LENGTH and CONTEXT_VERDICT (OK|WARN|SKIP). +context_compat_check() { + local name="$1" tail="$2" mml="$3" + CONTEXT_VERDICT="OK" + if [ -z "$tail" ] || [ "$tail" = "0" ]; then + # Genuine "unknown" fallback only (tail==0): no ISL estimate available, so + # keep the served window as-is. hf workloads now carry a loader-derived + # tail (>0) and flow through the WARN/SKIP/auto-size path below. + AGENTIC_MAX_CONTEXT_LENGTH="${mml:-0}" + agentic_log "context[$name]: unknown ISL tail, --max-context-length=${AGENTIC_MAX_CONTEXT_LENGTH}" + return 0 + fi + local needed + needed="$(_next_pow2 "$tail")" + if [ -n "$mml" ] && [ "$mml" != "0" ] && [ "$needed" -gt "$mml" ]; then + CONTEXT_VERDICT="WARN" + AGENTIC_MAX_CONTEXT_LENGTH="$mml" + agentic_warn "context[$name]: ISL tail $tail needs >= $needed but served max_model_len=$mml." + agentic_warn " Requests beyond $mml will be truncated. Serve with --max-model-len >= $needed for '$name'." + if [ "${AGENTIC_STRICT_CONTEXT:-0}" = "1" ]; then + CONTEXT_VERDICT="SKIP" + agentic_err "context[$name]: AGENTIC_STRICT_CONTEXT=1 -> SKIP" + fi + else + AGENTIC_MAX_CONTEXT_LENGTH="$needed" + agentic_log "context[$name]: ISL tail $tail -> --max-context-length=$needed (served max_model_len=${mml:-unset})" + fi + return 0 +} + +# -------------------------------------------------------------------------- +# Endpoint helpers: model-name alignment + router readiness (CRITICAL) +# -------------------------------------------------------------------------- +# aiperf --model MUST equal the server's registered served-model name or every +# request 404s. Prefer the router's advertised id; fall back to basename. +resolve_served_model_name() { + local base="http://127.0.0.1:${AGENTIC_PORT}" + local name="" + name="$(curl -sf "${base}/v1/models" 2>/dev/null \ + | "${AIPERF_PYTHON:-python3}" -c 'import sys,json; +d=json.load(sys.stdin); print((d.get("data") or [{}])[0].get("id",""))' 2>/dev/null)" + if [ -z "$name" ]; then + name="$(basename "${MODEL_PATH:-}")" + agentic_log "served-model name not advertised; falling back to basename: $name" + fi + [ -n "$name" ] || agentic_die "could not resolve a served model name (set MODEL explicitly)" + MODEL="$name" + agentic_log "aiperf --model resolved to: $MODEL" +} + +# Best-effort auto-detect of the served max_model_len (context window) so hf +# gating works model-agnostically when serving.max_model_len / MAX_MODEL_LEN is +# unset. SGLang (>=PR #4809) and vLLM both expose it in /v1/models +# data[0].max_model_len; /server_info is the old-sglang fallback. Config value +# always wins (caller only calls this when MAX_MODEL_LEN is empty/0). Retries 3x +# (server may be warming); prints "0" + WARN if all attempts fail. +# NB: the /server_info fallback reads context_length ONLY. Do not fall back to +# max_total_num_tokens -- that is the KV-cache token-pool budget (GPU mem / KV +# size across all concurrent seqs), not the per-request context window, so it +# would mis-gate the ISL tail (usually under-warn). +resolve_served_max_model_len() { + local base="http://127.0.0.1:${AGENTIC_PORT}" i v="" + for i in 1 2 3; do + v="$(curl -sf "$base/v1/models" 2>/dev/null \ + | "${AIPERF_PYTHON:-python3}" -c 'import sys,json;d=json.load(sys.stdin);print((d.get("data") or [{}])[0].get("max_model_len") or "")' 2>/dev/null)" + [ -n "$v" ] && { echo "$v"; return 0; } + v="$(curl -sf "$base/server_info" 2>/dev/null \ + | "${AIPERF_PYTHON:-python3}" -c 'import sys,json;d=json.load(sys.stdin);print(d.get("context_length") or "")' 2>/dev/null)" + [ -n "$v" ] && { echo "$v"; return 0; } + sleep 2 + done + agentic_err "could not auto-detect max_model_len from ${base}; set serving.max_model_len / MAX_MODEL_LEN" + echo "0" +} + +wait_for_router_ready() { + local base="http://127.0.0.1:${AGENTIC_PORT}" + local max="${AGENTIC_ROUTER_READY_TIMEOUT:-600}" + local waited=0 step=5 + agentic_log "waiting for endpoint readiness at $base (max ${max}s)" + while (( waited < max )); do + if curl -sf "${base}/v1/models" >/dev/null 2>&1; then + agentic_log "endpoint ready after ${waited}s"; return 0 + fi + sleep "$step"; waited=$((waited + step)) + done + agentic_die "endpoint not ready after ${max}s at $base" +} + +# -------------------------------------------------------------------------- +# aiperf command assembly +# -------------------------------------------------------------------------- +build_replay_cmd() { + local result_dir="$1" + local duration="${DURATION:-120}" + local conc="${AGENTIC_CONC:-16}" + # Model-size-aware cache warmup: DeepSeek/large families need a longer warm. + local cache_warmup + case "${MODEL_PREFIX:-}${MODEL:-}${MODEL_PATH:-}" in + *[Dd]eep[Ss]eek*|*dsv4*|*[Kk]imi*|*[Gg][Ll][Mm]*) cache_warmup="${AGENTIC_CACHE_WARMUP_DURATION:-300}" ;; + *) cache_warmup="${AGENTIC_CACHE_WARMUP_DURATION:-60}" ;; + esac + + export AIPERF_DATASET_CONFIGURATION_TIMEOUT=1800 + export AIPERF_SERVICE_PROFILE_CONFIGURE_TIMEOUT=1800 + + REPLAY_CMD="$AIPERF_CLI profile --scenario ${AGENTIC_SCENARIO:-inferencex-agentx-mvp}" + REPLAY_CMD+=" --url http://localhost:${AGENTIC_PORT}" + REPLAY_CMD+=" --endpoint /v1/chat/completions --endpoint-type chat --streaming" + REPLAY_CMD+=" --model ${MODEL:-auto}" + REPLAY_CMD+=" --concurrency $conc" + REPLAY_CMD+=" --benchmark-duration $duration" + REPLAY_CMD+=" --random-seed 42" + REPLAY_CMD+=" --failed-request-threshold $AIPERF_FAILED_REQUEST_THRESHOLD" + # Trajectory start window: hf captured traces resume mid-conversation (0.25/0.75); + # generated profile corpora replay near-complete sessions (0.90/0.98). A + # workload may override via WL_TRAJ_MIN/MAX (Tier 1); unset keeps today's values. + local traj_min="0.25" traj_max="0.75" + if [ "${WL_SOURCE:-hf}" = "profile" ]; then traj_min="0.90"; traj_max="0.98"; fi + [ -n "${WL_TRAJ_MIN:-}" ] && traj_min="$WL_TRAJ_MIN" + [ -n "${WL_TRAJ_MAX:-}" ] && traj_max="$WL_TRAJ_MAX" + REPLAY_CMD+=" --trajectory-start-min-ratio $traj_min --trajectory-start-max-ratio $traj_max" + REPLAY_CMD+=" --agentic-cache-warmup-duration $cache_warmup" + REPLAY_CMD+=" --warmup-grace-period ${AGENTIC_WARMUP_GRACE_PERIOD:-1800}" + REPLAY_CMD+=" --use-server-token-count --tokenizer-trust-remote-code" + REPLAY_CMD+=" --no-gpu-telemetry --slice-duration 1.0" + if [ -n "${AGENTIC_SERVER_METRICS:-}" ]; then + REPLAY_CMD+=" --server-metrics ${AGENTIC_SERVER_METRICS}" + fi + # --num-dataset-entries only applies to hf downloads (how many trace files to + # pull); a generated profile corpus is consumed whole. Per-workload + # WL_NUM_DATASET_ENTRIES (Tier 1) wins, then the AGENTIC_* env, then 393. + if [ "${WL_SOURCE:-hf}" = "hf" ]; then + REPLAY_CMD+=" --num-dataset-entries ${WL_NUM_DATASET_ENTRIES:-${AGENTIC_NUM_DATASET_ENTRIES:-393}}" + fi + # Per-workload context length: context_compat_check sets AGENTIC_MAX_CONTEXT_LENGTH + # (profile ISL tail rounded up, capped at max_model_len). Falls back to + # MAX_MODEL_LEN for the legacy/hf path (byte-identical). This also fixes the + # #173 bug where Case-B was capped at 262144 instead of its required 524288. + local ctx_len="${AGENTIC_MAX_CONTEXT_LENGTH:-${MAX_MODEL_LEN:-}}" + if [ -n "$ctx_len" ] && [ "$ctx_len" != "0" ]; then + REPLAY_CMD+=" --max-context-length $ctx_len" + fi + REPLAY_CMD+=" --output-artifact-dir $result_dir/aiperf_artifacts" + # The scenario enforces a 900s minimum; smoke runs opt into --unsafe-override + # (marks submission_valid=false, expected for non-canonical runs). + if [ "$duration" -lt 900 ] || [ "${AIPERF_UNSAFE_OVERRIDE:-false}" = "true" ]; then + REPLAY_CMD+=" --unsafe-override" + fi + REPLAY_CMD+=" $TRACE_SOURCE_FLAG" +} + +# -------------------------------------------------------------------------- +# Run + aggregate + rollback +# -------------------------------------------------------------------------- +write_agentic_result_json() { + local result_dir="$1" + # process_agentic_result reads aiperf artifacts from RESULT_DIR and writes + # $AGENTIC_OUTPUT_DIR/$RESULT_FILENAME.json (RESULT_FILENAME is required). + local result_filename="${AGENTIC_RESULT_FILENAME:-agentic_${SLURM_JOB_ID:-0}_xP${xP:-1}_yD${yD:-1}_${MODEL_NAME:-model}}" + AGENTIC_RESULT_JSON="${AGENTIC_OUTPUT_DIR:-$result_dir}/${result_filename}.json" + # process_agentic_result requires KV_OFFLOADING; "none" is the no-offload case + # (and requires KV_OFFLOAD_BACKEND to be empty). All other metadata env vars + # default cleanly, so the aggregate JSON's metrics are unaffected. + ( cd "$INFMAX_WS" && \ + RESULT_DIR="$result_dir" AGENTIC_OUTPUT_DIR="${AGENTIC_OUTPUT_DIR:-$result_dir}" \ + RESULT_FILENAME="$result_filename" \ + KV_OFFLOADING="${KV_OFFLOADING:-none}" \ + "$AIPERF_PYTHON" -m utils.agentic.aggregation.process_agentic_result ) + "$AIPERF_PYTHON" "$INFMAX_WS/utils/generate_aiperf_plots.py" "$result_dir" 2>&1 || true + agentic_log "aggregate JSON: $AGENTIC_RESULT_JSON" +} + +run_agentic_replay_and_write_outputs() { + local result_dir="$1" + local replay_rc + mkdir -p "$result_dir" + echo "$REPLAY_CMD" > "$result_dir/benchmark_command.txt" + + $REPLAY_CMD 2>&1 | tee "$result_dir/benchmark.log" + replay_rc=${PIPESTATUS[0]} + + write_agentic_result_json "$result_dir" + + # Best-effort post-benchmark health check (PASS/WARN on error + cache-hit). + local _validator="$(dirname "${BASH_SOURCE[0]}")/validate_agentic_result.sh" + if [ -f "$_validator" ] && [ -n "${AGENTIC_RESULT_JSON:-}" ]; then + AIPERF_PYTHON="$AIPERF_PYTHON" bash "$_validator" "$AGENTIC_RESULT_JSON" || true + fi + + if [ "$replay_rc" -ne 0 ]; then + # Automated rollback: mark the run invalid, leave logs for triage. + echo '{"submission_valid": false, "reason": "replay_rc='"$replay_rc"'"}' \ + > "$result_dir/RUN_INVALID.json" + agentic_err "agentic replay exited $replay_rc (results written, run marked invalid)" + return "$replay_rc" + fi + agentic_log "agentic replay complete -> $result_dir" +} + +# -------------------------------------------------------------------------- +# DRY_RUN: resolve everything possible without contacting a server, print, exit +# -------------------------------------------------------------------------- +agentic_dry_run() { + local result_dir="$1" + resolve_trace_loader + : "${AIPERF_CLI:=aiperf}" + if [ -z "${MODEL:-}" ]; then MODEL="$(basename "${MODEL_PATH:-}")"; fi + build_replay_cmd "$result_dir" + cat <} + AGENTIC_PORT : ${AGENTIC_PORT} + AGENTIC_CONC : ${AGENTIC_CONC:-16} + DURATION : ${DURATION:-120} + trace loader / dataset : ${TRACE_LOADER} / ${TRACE_DATASET} + RESULT_DIR : ${result_dir} + InferenceX pin : ${INFERENCEX_PIN} + aiperf pin : ${AIPERF_PIN} + +[agentic][DRY_RUN] assembled command: +${REPLAY_CMD} +EOF +} diff --git a/scripts/common/agentx/HOW_IT_WORKS.md b/scripts/common/agentx/HOW_IT_WORKS.md new file mode 100644 index 00000000..9d36e1bf --- /dev/null +++ b/scripts/common/agentx/HOW_IT_WORKS.md @@ -0,0 +1,358 @@ +# How AgentX works (and why it is accurate) + +This is the *how/why* companion to the AgentX docs. The other three answer +different questions: + +- [README.md](README.md) - the config schema and Tier 1 / Tier 2 knobs. +- [profiles/README.md](profiles/README.md) - how to author a profile/preset. +- [README.md Examples](README.md#examples) - copy-paste `agentic.yaml` recipes. + +This doc explains the **mechanism**: what the trace actually contains, why a +synthetic corpus measures a real engine faithfully, and how each knob moves +prefill/decode/cache/latency. Diagrams use a deliberately tiny **toy profile** +so the block and cache math is countable by hand, then Section 9 maps the toy +numbers back to the real `conformance_256k` preset. + +**Toy profile (used throughout):** + +```yaml +block_size: 64 +seed: 42 +n_sessions: 4 +isl_p: [128, 256, 384] # P50 / P90 / P99 input tokens +osl_p: [16, 48, 96] # P50 / P90 / P99 output tokens +delay_p: [1, 3, 8] # P50 / P90 / P99 inter-turn delay (s) +turns: + values: [1, 2, 3] + weights: [2, 2, 1] +cache_hit: [0.88, 0.90] # per-turn prefix-reuse band -> ~89% +``` + +## 1. Mental model + +AgentX offers a **fixed, reproducible load** to a serving engine and measures the +**achieved metrics** that engine produces. The *offered* side - how many input +and output tokens each turn carries, how many turns per session, how much of each +turn's prompt is a reused prefix, and how long the client "thinks" between turns - +is fully determined by the profile plus its `seed`, so it is **byte-identical on +every run and every engine**. The *achieved* side - throughput, time-to-first-token +(TTFT), inter-token latency, and the GPU-observed cache-hit rate - is what varies +by engine, topology, and configuration. Holding offered load fixed is exactly what +makes cross-engine comparisons fair: any difference in the achieved metrics is a +property of the engine, not of the workload. + +## 2. Core idea: content-independence + +Serving cost is driven by token **counts** and cache **structure**, not by token +**content**. A request of a given ISL and OSL with a given prefix-reuse pattern +does the same prefill and decode work whether the tokens spell real English or are +synthetic. AgentX exploits this: each turn in the trace carries a **token count** +(`in` / `out`) and a **hash_ids** list (the block-level structure that tells the +engine which KV blocks are shared prefix vs new). The **actual token text** is not +in the trace at all. Two different actors materialize text at replay time: + +- **Input (prompt) text** is synthesized **client-side by aiperf**, not the server. + aiperf's `PromptGenerator` is keyed by the turn's `hash_ids` (shared ids emit the + **same** tokens across turns, new ids emit fresh tokens), so the prompt is + deterministic filler whose block structure matches the trace. This is what makes + prefix reuse real: identical `hash_ids` produce byte-identical prompt blocks the + engine can serve from KV cache. +- **Output text** is produced by the **server** during decode; its timing is + measured (TTFT, inter-token latency) and then the text itself is discarded. + +```mermaid +flowchart TD + turn["One turn in the trace"] --> count["Token COUNT (in / out)"] + turn --> struct["hash_ids STRUCTURE (which blocks are shared vs new)"] + count --> fixed["Fixed and deterministic (from profile + seed)"] + struct --> fixed + fixed --> cost["Determines prefill + decode + cache work"] + fixed --> inText["Input text: aiperf PromptGenerator synthesizes it client-side, keyed by hash_ids"] + inText --> sent["Sent to the server as the prompt (shared hash_ids -> identical blocks -> cache hit)"] + sent --> outText["Output text: produced by the server at replay time"] + outText --> measured["Measured for timing"] + measured --> discarded["Then discarded (never fed back into the trace)"] +``` + +Because the parts that determine serving cost (counts + structure) are fixed +offline, a synthetic corpus exercises the engine identically to a captured trace +of the same shape - while *guaranteeing* the target distribution. + +## 3. End-to-end pipeline + +The suite resolves one `agentic.yaml` into N workloads and, per workload, either +generates or downloads a corpus, verifies it as a hard pre-gate, assembles the +replay command, runs aiperf, and rolls the result into a suite summary. Each stage +below is annotated with *what it decides*. + +```mermaid +flowchart TD + config["agentic.yaml"] --> resolve["resolve_config: merge presets, apply env overrides"] + resolve --> src{"source?"} + src -->|profile| gen["gen_agentx_profile.py: synthesize corpus (decides token counts + hash_ids)"] + src -->|hf| dl["download (+ filter_weka_corpus.py): fetch/trim captured corpus"] + gen --> verify["verify_agentx_profile.py: 13-axis pre-gate (decides run/abort)"] + dl --> verify + verify -->|"N/N axes within band"| gate["context_compat_check: decides max-context-length (OK / WARN / SKIP)"] + gate --> assemble["assemble aiperf command (decides flags: concurrency, duration, seed, trajectory)"] + assemble --> execute["aiperf execute: replay corpus, measure achieved metrics"] + execute --> aggregate["aggregate per-workload result JSON"] + aggregate --> summary["suite_summary.json"] +``` + +The entire corpus is generated **and** verified before any server contact - the +verify pre-gate aborts the run unless the corpus matches the profile's own targets. + +## 4. Corpus generation internals + +`generate_corpus()` in [gen_agentx_profile.py](gen_agentx_profile.py) draws every +random value from a single `random.Random(seed)` in a **fixed call order**. Per +session it first draws the turn count, then walks the turns; within each turn it +draws ISL, then OSL, then (only for turns after the first) the reuse fraction and +the think delay: + +```mermaid +flowchart TD + seedstep["random.Random(seed)"] --> turns["draw turn count (rng.choices)"] + turns --> loop["for each turn"] + loop --> isl["draw ISL (gauss)"] + isl --> osl["draw OSL (gauss)"] + osl --> first{"turn == 0 ?"} + first -->|yes| skip["no reuse / no delay draw (all blocks new)"] + first -->|no| reuse["draw reuse fraction (uniform in cache_hit band)"] + reuse --> delay["draw think delay (gauss)"] + skip --> emit["emit hash_ids + counts for this turn"] + delay --> emit + emit --> nextturn["next turn (accumulate prefix)"] + nextturn --> writefile["write session_XXXXX.json"] +``` + +Fixing `seed` makes the corpus **byte-identical** because every choice is a pure +function of that seed and the fixed call order: the same `random.Random(42)` +produces the same sequence of draws, `hash_ids` are `blake2b` hashes of +`":::"`, and the session id is a `blake2b` hash of +`"--"`. No wall-clock, no server, no floating-point +nondeterminism from the engine enters the corpus. Change `seed` (or `id_prefix`) +and you get a different-but-still-deterministic corpus; keep them and you can +regenerate the exact same bytes anywhere. + +## 5. Multi-turn replay mechanics (the key visual) + +This is the part most people picture incorrectly. Take one toy session of 3 turns +where the ISL draws land on the P50 / P90 / P99 values (`128`, `256`, `384` +tokens). With `block_size 64`: + +- **Turn 0** - ISL `128` = **2 blocks**. Turn 0 has no prior context, so all + **2 blocks are new** (`0 shared + 2 new`). The accumulated prefix is now 2 blocks. +- **Turn 1** - ISL `256` = **4 blocks**. Wanted reuse `floor(4 * 0.89) = 3`, but + only 2 prefix blocks exist, so reuse is capped at `min(2, 3) = 2`: + **2 shared + 2 new = 4**. Prefix grows to 4 blocks. +- **Turn 2** - ISL `384` = **6 blocks**. Wanted reuse `floor(6 * 0.89) = 5`, capped + at `min(4, 5) = 4`: **4 shared + 2 new = 6**. Prefix grows to 6 blocks. + +Every turn's counts add up exactly (`0+2=2`, `2+2=4`, `4+2=6`). The per-turn reuse +is capped by the available prefix (`min(len(prefix_blocks), floor(total * reuse))`), +so short early turns sit below the `0.88-0.90` band; over long sessions the prefix +saturates and the corpus-wide reuse converges into that band (Section 6). + +```mermaid +flowchart TD + t0["Turn 0 (2 blocks): 0 shared, 2 new"] -->|"prefix = 2 blocks"| t1 + t1["Turn 1 (4 blocks): 2 shared, 2 new"] -->|"prefix = 4 blocks"| t2 + t2["Turn 2 (6 blocks): 4 shared, 2 new"] -->|"prefix = 6 blocks"| done["session_00000.json"] + t0 --> n0(("new x2")) + t1 --> s1["shared x2"] + t1 --> n1(("new x2")) + t2 --> s2["shared x4"] + t2 --> n2(("new x2")) +``` + +Shared blocks (rectangles labeled `shared`) reuse the *same* `hash_ids` as prior +turns, so the engine can serve them from KV cache; new blocks (circles labeled +`new`) carry fresh `hash_ids` that force fresh prefill. + +### The code that does this + +All of the above is one loop in `make_session()` in +[gen_agentx_profile.py](gen_agentx_profile.py). Note how `prefix_blocks` +accumulates across turns and each turn's `hash_ids` is `reuse_slice + new_ids`: + +```python +prefix_blocks = [] # accumulated shared prefix (block hashes) +salt = f"{seed}:{idx}" +for turn in range(n_turns): + isl = samp(ISL_mu, ISL_sig, isl_lo, isl_hi) + osl = samp(OSL_mu, OSL_sig, osl_lo, osl_hi) + total_blocks = max(1, isl // block) + if turn == 0: + new_blocks = total_blocks # turn 0: everything is new + else: + reuse = min(len(prefix_blocks), # cap reuse at prefix we have + int(total_blocks * rng.uniform(cache_lo, cache_hi))) + new_blocks = max(1, total_blocks - reuse) # always >=1 fresh block + reuse_slice = prefix_blocks[:total_blocks - new_blocks] # SHARED ids (same as prior turns) + new_ids = [] + for b in range(new_blocks): + h = int(hashlib.blake2b(f"{salt}:{turn}:{b}".encode(), + digest_size=8).hexdigest(), 16) & 0x7FFFFFFFFFFFFFFF + new_ids.append(h) # fresh ids -> force prefill + hash_ids = reuse_slice + new_ids + prefix_blocks = hash_ids # grow the prefix for next turn +``` + +Three details make the mechanics deterministic and structural: + +- **`prefix_blocks = hash_ids`** at the end of each turn is what grows the shared + prefix; turn `N` can only reuse the blocks turns `0..N-1` laid down, which is why + early short turns fall below the `cache_hit` band. +- **`reuse_slice = prefix_blocks[:...]`** takes the *first* K prior block ids + verbatim, so shared blocks carry the identical `hash_ids` the engine already has + cached - reuse is by structure, not by content. +- **`new_ids`** are `blake2b` of `":::"`, so they are fresh + (force prefill) yet fully reproducible from `seed` - no server, no wall-clock. + +### What people think vs what actually happens + +The common misconception is that turn 2's prompt is assembled from the server's +*actual* turn-1 output text - i.e. that replay is a live conversation. It is not. + +```mermaid +flowchart TD + myth["MYTH: turn 2 prompt = server's real turn-1 output"] --> m1["turn 1 runs on server"] + m1 --> m2["read back generated text"] + m2 --> m3["build turn 2 from that text"] + reality["REALITY: turn 2 prompt is predetermined offline in the trace"] --> r1["turn 2 ISL + prefix-reuse structure fixed at generation time"] + r1 --> r2["shared hash_ids emit the SAME synthetic tokens; new ids emit fresh tokens"] + r2 --> r3["server's turn-1 text is measured for timing, then discarded"] + r3 --> r4["corpus is fully generated + verified BEFORE any server contact"] +``` + +Because the whole corpus exists and passes the verify pre-gate **before** the +server is ever contacted, turn 2 cannot depend on the server's turn-1 output. The +prefix reuse is structural (shared `hash_ids`), not semantic. + +## 6. Theoretical vs achieved cache-hit + +Two different cache-hit numbers show up in results, and they mean different things: + +- **theoretical_cache_hit_rate** is computed purely from the **trace structure** + (the ratio of shared to total blocks the profile lays down). It is a property of + the corpus, so it is **engine-invariant** - the same corpus yields the same + theoretical rate everywhere. For `conformance_256k` this sits at ~`0.89` + (`cache_hit: [0.88, 0.90]`, `cache_target: 89`). +- **gpu_cache_hit_rate** is what the engine *actually* achieves at runtime. It + depends on engine, KV-cache capacity, routing, and prefill/decode topology, so it + **varies** and is typically well below the theoretical ceiling. + +```mermaid +flowchart TD + trace["Trace structure (shared vs total blocks)"] --> theo["theoretical_cache_hit_rate ~0.89 (fixed, engine-invariant)"] + engine["Engine + topology + KV capacity + routing"] --> gpu["gpu_cache_hit_rate (varies by deployment)"] + theo --> gap["Gap = headroom the engine leaves on the table"] + gpu --> gap + gpu --> obs1["SGLang disagg 1P1D MoRI: ~0.367 (job 215077, Qwen3-32B)"] + gpu --> obs2["vLLM disagg NIXL 1P1D rixl+TP: ~0.369 (job 215125, Qwen3-32B)"] +``` + +Both anchor runs above have `error_rate 0.0` against the same theoretical ~`0.89` +corpus, yet each achieves only ~`0.37` GPU cache-hit - the gap is an engine/topology +property surfaced by holding the offered load fixed. + +## 7. Parameter -> effect map + +Each bullet reads "**increase** the knob -> effect on prefill / decode / cache / +latency / measurement." Grouped by what the knob controls. + +**Offered load** + +- **`isl_p`** - increase -> more input tokens per turn -> more prefill work and + more blocks -> higher TTFT; grows the prefix so later turns have more to reuse. +- **`osl_p`** - increase -> more output tokens per turn -> more decode steps -> + higher end-to-end latency and inter-token time; little effect on prefill/cache. +- **`turns`** - shift weight to higher values -> longer sessions -> prefix + saturates, so corpus-wide reuse climbs toward the `cache_hit` band. +- **`delay_p`** - increase -> longer client think-time between turns -> lower + request pressure per session; can let cached prefixes age out on capacity-bound + engines (lowering achieved cache-hit). +- **`n_sessions`** - increase -> more concurrent session trees and tighter + percentiles (closer match to targets) -> bigger corpus and longer generation. + +**Cache structure** + +- **`block_size`** - increase -> fewer, coarser blocks per turn -> coarser reuse + granularity (a single changed token invalidates a larger block); changes the + block accounting behind both cache-hit numbers. +- **`cache_hit`** - raise the `[lo, hi]` band -> more shared blocks per turn -> + higher theoretical cache-hit and less new prefill per late turn. +- **`hash_ids`** - not a user knob; it is the emitted per-block structure. Shared + ids across turns are exactly what makes prefix reuse measurable. + +**Determinism** + +- **`seed`** - change it -> a completely different but still reproducible corpus; + keep it -> byte-identical regeneration everywhere. +- **`id_prefix`** - change it -> different session ids (corpus identity) with the + same distribution; used to byte-match a committed corpus. + +**Validation (verify bands)** + +- **ISL band `0.80-1.20`** - measured ISL percentiles must land within +/-20% of + target or the axis is `off`. +- **OSL band `0.70-1.40`** - wider tolerance for the longer-tailed output lengths. +- **Turns band `0.60-1.60`** - widest, since the discrete turn distribution is + coarse at small `n_sessions`. +- **Delay band `0.50-2.00`** - very wide; delay does not affect served token work. +- **Cache band `0.97-1.03`** - tightest; the cache-hit P50 must be within +/-3% of + target, so structural reuse is held to close tolerance. + +**Cosmetic / routing** + +- **`model_tag`** - written into `requests[].model` and `models[]`; retag per + served model. No effect on token work. +- **Tier-2 `max_isl`** - drop sessions with any turn over N input tokens (trims the + ISL tail of a downloaded corpus). +- **Tier-2 `max_turns`** - truncate each session to its first N turns (caps session + length / prefix growth). +- **Tier-2 `sample`** - randomly keep N sessions (fixed `seed=42`) to shrink a + downloaded corpus. + +## 8. Example metrics output + +An illustrative per-workload result block (values are illustrative, not from a +specific run) with each field mapped to the knob that drives it: + +``` +theoretical_cache_hit_rate : 0.89 # trace structure: cache_hit + turns + block_size (fixed) +gpu_cache_hit_rate : 0.37 # engine + topology + KV capacity + routing (varies) +error_rate : 0.00 # request health; gated by --failed-request-threshold +output_token_throughput : 1850 tok/s # decode capacity: osl_p + concurrency + engine +time_to_first_token_p50 : 640 ms # prefill cost: isl_p + block_size + achieved cache-hit +inter_token_latency_p50 : 18 ms # decode step cost: osl_p + engine + concurrency +``` + +Reading it: the two cache-hit lines are the theoretical (fixed by the corpus) vs +achieved (engine-dependent) pair from Section 6; `error_rate` should be `0.00` +for a valid run; throughput and inter-token latency track output work (`osl_p`, +concurrency, engine); TTFT tracks input work (`isl_p`, `block_size`) minus whatever +the engine reuses from cache. + +## 9. Mapping the toy to `conformance_256k` + +The toy is the same shape as the shipped [conformance_256k.yaml](profiles/conformance_256k.yaml) +preset - just scaled down so the block math is hand-countable. To connect them: + +- **`isl_p`** toy `[128, 256, 384]` -> real `[74000, 155000, 235000]` (hundreds of + blocks per turn instead of a handful). +- **`osl_p`** toy `[16, 48, 96]` -> real `[320, 3300, 17000]` (long, heavy-tailed + outputs). +- **`delay_p`** toy `[1, 3, 8]` -> real `[4, 31, 240]` seconds. +- **`turns`** toy `values [1,2,3]` -> real long-tail `values [1,2,3,4,6,10,20,45,103]`, + so real sessions run long enough for prefix reuse to saturate into the band. +- **`cache_hit`** `[0.88, 0.90]` - identical in both (the reuse band that yields the + ~`0.89` theoretical cache-hit). +- **`seed`** `42` and **`block_size`** `64` - identical in both, so the same + regeneration guarantees apply. + +Everything the toy demonstrates - content-independence, deterministic generation, +structural prefix reuse, and the theoretical-vs-achieved cache-hit split - holds +unchanged at conformance scale; only the counts get bigger. diff --git a/scripts/common/agentx/README.md b/scripts/common/agentx/README.md new file mode 100644 index 00000000..22c2e121 --- /dev/null +++ b/scripts/common/agentx/README.md @@ -0,0 +1,846 @@ +# AgentX benchmarking core + +AgentX runs a **list of agentic trace-replay workloads** against **one** served +endpoint. For each workload the suite either **generates** a synthetic +`weka_trace` corpus from a distribution profile or **downloads** a real captured +HuggingFace (HF) trace, **verifies** the corpus as a hard pre-gate, **replays** +it with aiperf, and writes a per-workload result dir plus a combined suite +summary. The whole run is described by one `agentic.yaml` (a `serving:` block, a +`run:` block, and a `workloads:` list); adding "N cases" just means adding N list +entries. + +```mermaid +flowchart LR + config["agentic.yaml"] --> resolve["agentx_config.py (resolve_config)"] + resolve --> src{"workload source?"} + src -->|profile| gen["gen_agentx_profile.py (generate corpus)"] + src -->|hf| dl["download (+ filter_weka_corpus.py)"] + gen --> verify["verify_agentx_profile.py (pre-gate)"] + verify -->|"13/13 axes within band"| replay["aiperf replay"] + dl --> replay + replay --> results["RESULT_DIR// results"] +``` + +This document covers the **agentx core**: the config schema, the profile/preset +model, the generate/verify tools, and the Tier 1 / Tier 2 knobs. The launcher / +disaggregated-serving integration and the deep env/run-path reference are out of +scope here (see [See also](#see-also)). + +## Onboarding new users + +New to AgentX? Follow this path: + +1. **Read your launcher's entry point** — the agentic-replay section of + [sglang_disagg/README.MD](../../sglang_disagg/README.MD) or + [vllm_dissag/README.MD](../../vllm_dissag/README.MD) for how a real run is + invoked end-to-end (both share this same AgentX core), then this README for + the config/profile model. +2. **Copy-paste a scenario** — [Examples](#examples) has a ready `agentic.yaml` + for every workload shape. Start from the one closest to your goal. +3. **Understand the mechanism** (optional) — [HOW_IT_WORKS.md](HOW_IT_WORKS.md) explains + why the replay is reproducible and how each knob maps to prefill/decode/cache. + +Two commands to try first: + +- **Preview with no server** — resolve the plan and print every assembled aiperf + command in under a second, without contacting a server, downloading, or replaying: + +```bash +DRY_RUN=1 AGENTIC_CONFIG=agentic.yaml bash scripts/common/benchmark_agentic_suite.sh +``` + +- **Fast smoke run** — `AGENTIC_WORKLOAD=small` uses the shipped + [profiles/small.yaml](profiles/small.yaml): a generated corpus with tiny ISL + (`2k-6k` tokens) so it materializes, verifies, and replays quickly. Note it is + still 200 sessions (fast because the prompts are small, not because the corpus + is tiny): + +```bash +AGENTIC_WORKLOAD=small bash scripts/common/benchmark_agentic_suite.sh +``` + +## What AgentX adds on top of aiperf + +aiperf's `inferencex-agentx-mvp` scenario provides the base: an agentic trace +replay engine, the `weka_trace` corpus format, HuggingFace loaders, and the core +metrics (TTFT, E2E latency, throughput, cache-hit). It does **not** provide +multi-workload orchestration, corpus generation, corpus verification, declarative +config, or backend portability. AgentX is the orchestration + generation layer +that closes those gaps by **wrapping** aiperf (not forking it). + +| Capability | aiperf (base) | AgentX adds | +| --- | --- | --- | +| Replay engine | Single-workload CLI | Multi-workload orchestration over one config | +| Corpus source | HF captured traces only | HF traces + seed-deterministic generated profiles | +| Corpus verification | None | 13-axis pre-gate, hard abort on drift | +| Config format | CLI flags only | Declarative YAML (serving + workloads[]) | +| Preset reuse | Copy-paste | Inheritance with per-entry override | +| Context sizing | Hardcoded 256k/1M | Auto-detect from /v1/models, any window | +| Corpus filtering | Download full corpus | Tier 2 trim (max_isl/max_turns/sample), cached | +| Backend integration | vLLM-specific paths | OpenAI-API core, thin per-backend hook | +| Multi-workload runs | N CLI calls | One config -> N result dirs + unified summary | +| Preview | None | DRY_RUN=1, no server, sub-second | + +The nine additions, one line each: + +1. **Suite driver** (`benchmark_agentic_suite.sh`) — loops `workloads[]`, materializes each corpus, runs aiperf, writes one `suite_summary.json`. +2. **Profile generator** (`gen_agentx_profile.py`) — distribution targets -> byte-identical corpus via fixed seed. +3. **13-axis verifier** (`verify_agentx_profile.py`) — measures ISL/OSL/Turns/Delay P50/P90/P99 + Cache-hit P50 and aborts unless all axes pass. +4. **Declarative config** (`agentx_config.py` + `agentic.yaml`) — serving / run / workloads[] blocks; add a workload in one entry. +5. **Tier 1 / Tier 2 separation** — Tier 1 (concurrency, num_dataset_entries, trajectory) steers replay; Tier 2 (max_isl, max_turns, sample) produces a cached subset corpus. +6. **Model-agnostic context gating** — auto-detect served window from `/v1/models`, cap `--max-context-length` at `min(next_pow2(ISL tail), window)`, WARN/cap or SKIP under `AGENTIC_STRICT_CONTEXT=1`. +7. **Preset inheritance** — ship a shape once in `profiles/.yaml`, reference via `preset:`, override per entry; circular chains raise `ValueError`. +8. **Backend-agnostic core** (`agentic_lib.sh`, ~480 lines) — speaks only the OpenAI API; a new backend is a thin (~60-line) hook. Currently integrated: SGLang disaggregated (`scripts/sglang_disagg/benchmark_agentic.sh`, port 2322); vLLM/TRT-LLM/TGI are future hooks. +9. **DRY_RUN preview** — `DRY_RUN=1` prints the resolved plan and every assembled aiperf command with no server/download/replay. + +## File map + +``` +scripts/common/agentx/ + agentx_config.py # config loader: parses agentic.yaml + profiles, resolves per-workload params, emits JSON/shell + gen_agentx_profile.py # seed-deterministic corpus generator (profile JSON -> session_XXXXX.json files) + verify_agentx_profile.py# corpus verifier: 13-axis conformance table + "N/N axes within band" pre-gate + filter_weka_corpus.py # Tier 2 filter: trim a downloaded hf corpus (max_isl / max_turns / sample) + agentic.example.yaml # annotated canonical config; copy to agentic.yaml and edit + profiles/ # shipped presets + authoring guide + README.md # profile/preset authoring guide (see profiles/README.md) + conformance_256k.yaml # Case-A generated conformance profile (ExplainX targets) + conformance_512k.yaml # Case-B longer-context conformance profile (ISL tail to 500k) + inferencex_256k.yaml # reusable source=hf preset (loader + bundled sweep + Tier 1 knobs) + small.yaml # tiny generated profile (fast smoke) + custom.example.yaml # annotated template for a user-defined profile +``` + +The two bash drivers that consume this core live one level up: +`scripts/common/benchmark_agentic_suite.sh` (the suite loop) and +`scripts/common/agentic_lib.sh` (corpus materialization + replay assembly). + +## Quick start + +A minimal valid `agentic.yaml` is just a workloads list; `serving:` and `run:` +fall back to their defaults: + +```yaml +workloads: + - { name: quick, preset: conformance_256k } +``` + +With no `serving:`/`run:` blocks, `resolve_config()` supplies `serving.model: +auto`, `serving.max_model_len: 0` (auto-detect), `serving.port: auto` (`auto` +resolves to the recipe default — `2322` for the sglang router — so `auto` and +`2322` name the same port), `serving.server_metrics: auto`, `run.concurrency: +16`, and `run.duration: 900`. See [Examples](#examples) for copy-paste +examples of every workload shape. + +### Run it + +**Prerequisite:** an OpenAI-compatible endpoint must already be served on +`AGENTIC_PORT` (default `2322`). This suite does **not** start a server — on a +cluster the endpoint is brought up by the launcher recipe (below; see +[../../sglang_disagg/README.MD](../../sglang_disagg/README.MD)), and for a +direct run you must start or point to your own endpoint first. To wire up a new +backend, see the integration guide at +[./templates/README.md](./templates/README.md). + +```bash +# Single preset, config-less (driver synthesizes a one-entry config) +AGENTIC_WORKLOAD=conformance_256k bash scripts/common/benchmark_agentic_suite.sh + +# Multi-workload suite from a config +AGENTIC_CONFIG=agentic.yaml bash scripts/common/benchmark_agentic_suite.sh + +# Preview first — no server needed (see "Preview / debug" below) +DRY_RUN=1 AGENTIC_CONFIG=agentic.yaml bash scripts/common/benchmark_agentic_suite.sh +``` + +See [Preview / debug](#preview--debug) for what `DRY_RUN=1` prints. + +On a cluster you normally invoke the suite **indirectly** via `sbatch +scripts/sglang_disagg/run_xPyD_models.slurm` (see +[../../sglang_disagg/README.MD](../../sglang_disagg/README.MD)); run +`benchmark_agentic_suite.sh` directly only for local or `DRY_RUN` use, and in +that case set `RESULT_DIR` yourself since it is otherwise launcher-provided. + +### What to expect / run timing + +Per workload, wall-time is dominated by a **cache-warmup** phase followed by the +**measured replay window**: + +- **Measured window** — `run.duration` (default `900`s = 15 min); this is the + aiperf `--benchmark-duration`. +- **Cache warmup** — `--agentic-cache-warmup-duration` (default `60`s; `300`s + for large model families such as DeepSeek/Kimi/GLM) runs *before* measurement, + bounded by `--warmup-grace-period` (default `1800`s). +- **Corpus generation** — for `source: profile` workloads this is a one-time + step cached under `SUITE_CORPUS_DIR` (fast; scales with `n_sessions`) and + reused on later runs unless `SUITE_CORPUS_FORCE=1`. `source: hf` workloads + download instead of generate. +- **Concurrency sweeps** — a `concurrency` **list** runs each value in sequence, + so it multiplies wall-time (see [Examples](#examples)). + +## Glossary + +- **Profile** — a set of distribution targets (`isl_p`, `osl_p`, `delay_p`, + `turns`, `cache_hit`, `clamps`) plus a `seed`/`n_sessions` that + `gen_agentx_profile.py` turns into a reproducible corpus. See + [profiles/README.md](profiles/README.md). +- **Preset** — a shipped profile file in `profiles/.yaml` that a workload + inherits via `preset: `. A preset can carry distribution params + (`source: profile`), an hf `loader` + Tier 1/Tier 2 knobs (`source: hf`), + and/or run knobs (`concurrency`/`duration`). +- **Workload** — one entry in the `workloads:` list; one result subdir. +- **Corpus** — the per-session `session_XXXXX.json` files aiperf replays. Lives + under `SUITE_CORPUS_DIR/` (a reusable cache), **not** under `RESULT_DIR`. +- **ISL** — input sequence length (input tokens per request). +- **OSL** — output sequence length (output tokens per request). + +**Tier 1 (replay-level knobs)** steer how the trace is replayed and are stripped +from the generator profile dict (`_CONTROL_KEYS` in `agentx_config.py`): + +- `concurrency` — session-tree concurrency (scalar or list; a list sweeps). +- `num_dataset_entries` — how many hf trace sessions to pull (hf only). +- `trajectory: { min, max }` — start-window ratio for captured traces. + +**Tier 2 (corpus filter)** trims a *downloaded* hf corpus before replay +(`filter_weka_corpus.py`), applied in this order: + +- `max_turns` — truncate each session to its first N turns. +- `max_isl` — drop a session if any (post-truncation) turn's input exceeds N. +- `sample` — randomly keep N sessions (fixed `seed=42`). + +## Config schema + +The schema is defined by `resolve_config()` in `agentx_config.py`. It reads three +top-level keys. + +### `serving:` — one endpoint for the whole run + +| key | default | notes | +| --- | --- | --- | +| `model` | `auto` | `auto` resolves the served-model id from `/v1/models`; or set explicitly. | +| `max_model_len` | `0` | `0` auto-detects the served window; a set value always wins. See below. | +| `port` | `auto` | `auto` -> recipe default (sglang router `2322` / vLLM shim port). | +| `server_metrics` | `auto` | `auto` -> recipe host:port list; or space-separated endpoints. | + +### `run:` — default replay knobs + +| key | default | notes | +| --- | --- | --- | +| `concurrency` | `16` | scalar or list; a list sweeps per workload. | +| `duration` | `900` | measured window (s); the scenario minimum for a valid submission is 900. | + +### `workloads:` — a list of entries + +Each entry is merged over its `preset:` chain (`_merge_preset()`; entry keys win) +and resolved by `_resolve_workload_entry()`. Recognized control keys +(`_CONTROL_KEYS`, stripped from any generator profile): + +| key | applies to | meaning | +| --- | --- | --- | +| `name` | all | workload name; becomes the result subdir. | +| `source` | all | `profile` (generate) or `hf` (download). Defaults to `profile`. | +| `preset` | all | inherit `profiles/.yaml`. | +| `loader` | `hf` | aiperf `--public-dataset` id (sets the context-gating ISL tail). | +| `filter` | `hf` | Tier 2 map: `max_isl` / `max_turns` / `sample`. | +| `num_dataset_entries` | `hf` | Tier 1: trace sessions to pull. | +| `trajectory` | all | Tier 1: `{ min, max }` start-window ratio (`0.0 <= min <= max <= 1.0`). | +| `concurrency` | all | per-entry override of `run.concurrency`. | +| `duration` | all | per-entry override of `run.duration`. | + +For `source: profile`, the entry additionally carries (or inherits) the profile +distribution fields (`model_tag`, `id_prefix`, `seed`, `n_sessions`, +`block_size`, `isl_p`, `osl_p`, `delay_p`, `turns`, `cache_hit`, `clamps`, +`verify`) documented in [profiles/README.md](profiles/README.md). + +## Environment overrides + +Environment variables override file values (applied in `resolve_config()`): + +- `MODEL` -> `serving.model` +- `MAX_MODEL_LEN` -> `serving.max_model_len` +- `AGENTIC_PORT` -> `serving.port` +- `AGENTIC_SERVER_METRICS` -> `serving.server_metrics` +- `AGENTIC_CONC` -> `run.concurrency` +- `DURATION` -> `run.duration` +- `AGENTIC_WORKLOAD=` — restrict the run to that single named entry (also + enables the config-less shorthand; see below). + +`AGENTIC_WORKLOAD` has **two distinct meanings** depending on whether +`AGENTIC_CONFIG` is set — keep the two straight: + +**Config-less shorthand (true minimum, no `AGENTIC_CONFIG`).** `AGENTIC_WORKLOAD` +*names the workload to synthesize*: + +```bash +AGENTIC_WORKLOAD=conformance_256k MAX_MODEL_LEN=262144 AGENTIC_CONC=4 \ + bash scripts/common/benchmark_agentic_suite.sh +``` + +**Filter an existing config to one entry (with `AGENTIC_CONFIG`).** Here +`AGENTIC_WORKLOAD` *selects* the single entry named `quick` from the config +(via `resolve_config()`); it does **not** use the shorthand: + +```bash +AGENTIC_CONFIG=agentic.yaml AGENTIC_WORKLOAD=quick \ + bash scripts/common/benchmark_agentic_suite.sh +``` + +With no `--config`, `_synth_config_from_env()` synthesizes a one-entry config +from `AGENTIC_WORKLOAD`: `inferencex` maps to the shipped hf loader (`_HF_PRESETS`), +and any other name is treated as `preset: ` (so `conformance_256k` / +`conformance_512k` resolve to their shipped profiles). + +### Minimal required + +The true minimum for each entry point — everything else auto-defaults (model +`auto`, `max_model_len` auto-detect, port `2322`, concurrency `16`, duration +`900`). Prerequisite in every case: an OpenAI-compatible endpoint must already be +served on `AGENTIC_PORT` (default `2322`) — launcher-provided on a cluster, or +your own for a direct run. + +- **Multi-workload suite** — point at a config, nothing else required: + +```bash +AGENTIC_CONFIG=agentic.yaml bash scripts/common/benchmark_agentic_suite.sh +``` + +- **Single preset (config-less)** — name a preset; the driver synthesizes a + one-entry config from that name: + +```bash +AGENTIC_WORKLOAD=conformance_256k bash scripts/common/benchmark_agentic_suite.sh +``` + +## Environment variable reference + +Every user-facing environment variable, grouped by role. Defaults shown are the +values applied when the variable is unset. + +All variables are **optional** except that you must set exactly **one** +entry-point variable (`AGENTIC_CONFIG` or `AGENTIC_WORKLOAD`) and have a served +endpoint (see [Minimal required](#minimal-required)). + +### Serving / selection + +| Variable | Default | Meaning | +| --- | --- | --- | +| `MODEL` | `auto` | served model id (maps to `serving.model`; auto-discovers from `/v1/models`) | +| `MAX_MODEL_LEN` | `0` (auto) | pin served context window; `0` = auto-detect (`serving.max_model_len`) | +| `AGENTIC_PORT` | `2322` | endpoint/router port (`serving.port`) | +| `AGENTIC_SERVER_METRICS` | `auto` | aiperf `--server-metrics` endpoints; on disaggregated serving its first `host:port` is also probed for the served `max_model_len` | + +### Run / replay + +| Variable | Default | Meaning | +| --- | --- | --- | +| `AGENTIC_CONC` | `16` | replay concurrency (`run.concurrency`) | +| `DURATION` | `900` | measured window in seconds (`run.duration`) | +| `AGENTIC_NUM_DATASET_ENTRIES` | `393` | default `--num-dataset-entries` for hf workloads (global default; per-workload YAML `num_dataset_entries` overrides) | + +### Corpus & HuggingFace + +| Variable | Default | Meaning | +| --- | --- | --- | +| `SUITE_CORPUS_DIR` | `${TMPDIR:-/tmp}/agentx_corpora` | corpus cache root | +| `SUITE_CORPUS_FORCE` | `0` | set `1` to regenerate a cached corpus after editing a profile | +| `WEKA_LOADER_OVERRIDE` | recipe default | override the hf trace loader id outside the YAML `loader:` path | +| `AGENTIC_HF_ISL_TAIL` | derived from loader | override the loader-derived ISL tail used for context gating | +| `AGENTIC_TRACE_DL_ATTEMPTS` | `3` | hf trace download retry count | +| `AGENTIC_TRACE_DL_TIMEOUT` | `900` | per-attempt hf download timeout (seconds) | + +### Context gating + +| Variable | Default | Meaning | +| --- | --- | --- | +| `AGENTIC_MAX_CONTEXT_LENGTH` | derived (falls back to `MAX_MODEL_LEN`) | force the `--max-context-length` cap | +| `AGENTIC_STRICT_CONTEXT` | `0` | set `1` to SKIP (instead of WARN/cap) a workload whose ISL tail exceeds the served window | + +### Timing / warmup + +| Variable | Default | Meaning | +| --- | --- | --- | +| `AGENTIC_CACHE_WARMUP_DURATION` | `60` (`300` for DeepSeek/Kimi/GLM families) | `--agentic-cache-warmup-duration` | +| `AGENTIC_WARMUP_GRACE_PERIOD` | `1800` | `--warmup-grace-period` | +| `AGENTIC_ROUTER_READY_TIMEOUT` | `600` | seconds to wait for the router/endpoint to become ready before failing | + +### Entry points / preview + +| Variable | Default | Meaning | +| --- | --- | --- | +| `AGENTIC_CONFIG` | (unset) | Required (choose one of CONFIG/WORKLOAD): path to an `agentic.yaml`; runs the multi-workload suite driver | +| `AGENTIC_WORKLOAD` | (unset) | Required (choose one of CONFIG/WORKLOAD): run one workload by name: config-less preset shorthand, or select a single entry from `AGENTIC_CONFIG` | +| `DRY_RUN` | `0` | set `1` to print the resolved plan and each assembled aiperf command without contacting a server | +| `AGENTIC_RESOLVE_ONLY` | `0` | set `1` to resolve the served `max_model_len` and exit without running; diagnostic | + +### Output / labeling + +| Variable | Default | Meaning | +| --- | --- | --- | +| `RESULT_DIR` | `/run_logs/${SLURM_JOB_ID:-0}` | root for per-workload result dirs and `suite_summary.json`; the hooks and suite driver default it to `/run_logs/${SLURM_JOB_ID:-0}` (so a launcher supplies `SLURM_JOB_ID`). For a direct run without the launcher, set `RESULT_DIR` yourself (else results land in `/run_logs/0`). | +| `AGENTIC_RESULT_FILENAME` | `agentic_${SLURM_JOB_ID}_xP..._yD..._${MODEL_NAME}` | aggregate result JSON basename | +| `KV_OFFLOADING` | `none` | label threaded into result aggregation | + +Advanced / maintainer overrides default to sane values and normally need no +change: `AIPERF_PIN`, `AGENTIC_UTILS_PIN`, `INFERENCEX_REPO`, +`AGENTIC_RUNTIME_DIR`, `INFMAX_WS`, `AIPERF_VENV`, +`AIPERF_FAILED_REQUEST_THRESHOLD`, `AIPERF_UNSAFE_OVERRIDE`, +`AGENTX_YAML_FALLBACK`, `AGENTIC_OUTPUT_DIR`, `AGENTIC_LIB`. + +## Examples + +Copy-paste `agentic.yaml` recipes for common AgentX cases. Each scenario shows +the config snippet (reusing [agentic.example.yaml](agentic.example.yaml) +entries), what happens step by step, and the resulting output tree. + +**Where things land (read this once):** + +- **Corpus** (the `session_XXXXX.json` files) lives under + `SUITE_CORPUS_DIR/` — default `/tmp/agentx_corpora/` — a reusable + cache. It is **not** written under `RESULT_DIR`. +- **`RESULT_DIR//`** holds only replay output: `aiperf_artifacts/`, + `benchmark_command.txt`, `benchmark.log`, the aggregate `*.json` (name from + `AGENTIC_RESULT_FILENAME`), and `RUN_INVALID.json` on failure. A concurrency + **list** adds `conc/` subdirs; a single concurrency writes flat. +- The suite also writes `RESULT_DIR/suite_summary.json`. + +Run any of these with: + +```bash +AGENTIC_CONFIG=agentic.yaml bash scripts/common/benchmark_agentic_suite.sh +``` + +### Scenario 1: single generated profile (`conformance_256k`) + +The canonical single-workload run — the same example walked through in +[Quick start](#quick-start). The minimal config: + +```yaml +serving: + model: auto + max_model_len: 0 # auto-detect the served window +run: + concurrency: 16 + duration: 900 +workloads: + - name: conformance_256k + source: profile + preset: conformance_256k +``` + +The entry resolves against `profiles/conformance_256k.yaml`, the corpus is +generated into `/tmp/agentx_corpora/conformance_256k/` (cached) and verified as a +pre-gate (`13/13 axes within band`), then aiperf replays it at concurrency `16` +with results written flat under `RESULT_DIR/conformance_256k/`. See +[Quick start](#quick-start) for the config-less shorthand and the run commands. + +### Scenario 2: multiple profiles in one run + +```yaml +workloads: + - name: conformance_256k + source: profile + preset: conformance_256k + - name: conformance_512k + source: profile + preset: conformance_512k +``` + +**What happens:** the two workloads run **in sequence** against the same +endpoint. Each is generated + verified + replayed independently; both roll up +into one `suite_summary.json`. + +**Output tree:** + +``` +RESULT_DIR/ + conformance_256k/ + aiperf_artifacts/ benchmark_command.txt benchmark.log .json + conformance_512k/ + aiperf_artifacts/ benchmark_command.txt benchmark.log .json + suite_summary.json +``` + +### Scenario 3: real HF trace (`source: hf` + `loader:`) + +```yaml +workloads: + - name: inferencex + source: hf + loader: semianalysis_cc_traces_weka_062126_256k +``` + +**What happens:** + +1. No corpus is generated. `resolve_trace_loader()` maps the loader to + `--public-dataset semianalysis_cc_traces_weka_062126_256k`; the dataset is + downloaded at run time (3× retry/backoff) into the shared HF cache. +2. The `_256k` loader suffix sets the context-gating ISL tail to `262144`. +3. With **no** Tier 2 `filter:`, the replay uses the byte-identical + `--public-dataset` path (no local `filter_weka_corpus.py` step). + +**Output tree:** same layout as Scenario 1 under `RESULT_DIR/inferencex/`. There +is no `SUITE_CORPUS_DIR/inferencex/` dir because the unfiltered hf path streams +from the HF cache rather than a materialized `weka_trace` dir. + +### Scenario 4: reusable HF preset with a bundled sweep (`inferencex_256k`) + +```yaml +workloads: + - name: inferencex_preset + preset: inferencex_256k +``` + +**What happens:** [profiles/inferencex_256k.yaml](profiles/inferencex_256k.yaml) +bundles `source: hf`, the loader, `concurrency: [2, 4, 8]`, `duration: 900`, and +Tier 1 knobs (`num_dataset_entries: 393`, `trajectory: {min: 0.25, max: 0.75}`). +Because `concurrency` is a **list**, the workload is swept and results land in +per-concurrency subdirs. The sweep runs the concurrency values **sequentially** +(not in parallel), each into its own `conc/` subdir, so total wall-time is +roughly `N x (cache-warmup + duration)`. + +**Output tree:** + +``` +RESULT_DIR/ + inferencex_preset/ + conc2/ aiperf_artifacts/ benchmark_command.txt benchmark.log .json + conc4/ aiperf_artifacts/ benchmark_command.txt benchmark.log .json + conc8/ aiperf_artifacts/ benchmark_command.txt benchmark.log .json + suite_summary.json +``` + +### Scenario 5: Tier 1 knobs (`num_dataset_entries`, `trajectory`) + +```yaml +workloads: + - name: inferencex_light + preset: inferencex_256k + num_dataset_entries: 50 # pull fewer trace sessions + trajectory: { min: 0.30, max: 0.80 } # start-window ratio for captured traces +``` + +**What happens:** entry keys win over the preset (`_merge_preset()`). Tier 1 +knobs steer the replay only: `--num-dataset-entries 50` (hf downloads only) and +`--trajectory-start-min-ratio 0.30 --trajectory-start-max-ratio 0.80`. No corpus +is filtered on disk. `trajectory` is validated as `0.0 <= min <= max <= 1.0`. + +**Output tree:** as Scenario 4 (a sweep, since the preset's `[2, 4, 8]` is +inherited) under `RESULT_DIR/inferencex_light/conc/`. + +### Scenario 6: Tier 2 filter (`max_isl`, `max_turns`, `sample`) + +```yaml +workloads: + - name: inferencex_small + preset: inferencex_256k + concurrency: [2] # entry overrides the preset's [2, 4, 8] sweep + num_dataset_entries: 50 # Tier 1: pull fewer sessions + trajectory: { min: 0.30, max: 0.80 } + filter: # Tier 2: local subset/trim (download once, then filter) + max_isl: 200000 # drop sessions with any turn over 200k input tokens + max_turns: 40 # truncate each session to its first 40 turns + sample: 100 # randomly keep 100 sessions (seed=42) +``` + +**What happens:** a `filter:` on an hf workload triggers `materialize_hf_corpus()`: +download once, then `filter_weka_corpus.py` applies **max_turns, then max_isl, +then sample** (in that order) and writes a materialized `weka_trace` dir. The +replay then uses `--custom-dataset-type weka_trace --input-file `. The +filtered corpus is cached under a content-addressed key +`hf__`. An empty filter result fails loudly (exit 1). + +**Output tree:** + +``` +/tmp/agentx_corpora/hf_semianalysis_cc_traces_weka_062126_256k_/ # filtered corpus (cache) + session_00000.json ... + +RESULT_DIR/ + inferencex_small/ + aiperf_artifacts/ benchmark_command.txt benchmark.log .json # single conc => flat + suite_summary.json +``` + +### Scenario 7: custom inline workload (`my_case`) + +Define a profile inline (no preset) — copy the fields from +[profiles/custom.example.yaml](profiles/custom.example.yaml): + +```yaml +workloads: + - name: my_case + source: profile + model_tag: GLM-5.2-MXFP4 + id_prefix: my_case + seed: 42 + n_sessions: 150 + block_size: 64 + isl_p: [48000, 120000, 200000] + osl_p: [256, 2000, 9000] + delay_p: [3, 20, 180] + turns: + values: [2, 3, 4, 6, 10, 20, 45, 103] + weights: [20, 24, 20, 12, 8, 6, 7, 3] + cache_hit: [0.88, 0.90] + clamps: + isl: [1200, 205000] + osl: [8, 20000] + delay: [1, 600] +``` + +**What happens:** identical to Scenario 1, but the distribution targets come from +the entry itself instead of a shipped preset. With no `verify:` block, the +verifier derives the turns targets from the `turns` distribution and the cache +target from `mean(cache_hit) * 100`, so the profile still round-trips to +`13/13 axes within band`. Output tree matches Scenario 1 under +`RESULT_DIR/my_case/`. + +### Scenario 8: context-window behavior (auto vs pin) + +Case-B's ISL tail (`clamps.isl` hi `520000`) rounds up to a `524288` window. + +**Auto-detect (recommended):** + +```yaml +serving: + max_model_len: 0 # auto-detect from /v1/models +workloads: + - { name: conformance_512k, preset: conformance_512k } +``` + +- If the served window is `>= 524288`, `context_compat_check()` sets + `--max-context-length 524288`, verdict `OK`. +- If the served window is smaller (say `262144 < 524288`), it **WARNs** and caps + `--max-context-length` at the served window (late turns get truncated). +- With `AGENTIC_STRICT_CONTEXT=1`, that same case is **SKIPPED** instead + (recorded as `SKIP(context)` in `suite_summary.json`, no replay). + +**Pin explicitly** (e.g. when `/v1/models` doesn't expose the window, or to force +a cap): + +```yaml +serving: + max_model_len: 524288 # Case-B needs 524288 +``` + +A pinned value always wins over auto-detect (and over the `0` default). See also +[`max_model_len` guidance](#max_model_len-guidance). + +### Scenario 9: `DRY_RUN=1` preview + +Preview the resolved plan and every assembled aiperf command with no server — +the same mechanism documented in [Preview / debug](#preview--debug): + +```bash +DRY_RUN=1 AGENTIC_CONFIG=agentic.yaml bash scripts/common/benchmark_agentic_suite.sh +``` + +For the Scenario 1 config, the driver prints the resolved plan and each assembled +command **without contacting a server** (no download, no generate, no replay): + +``` +[agentic][DRY_RUN] resolved suite plan + config : agentic.yaml + serving.model : auto + serving.max_model_len : 0 + serving.port : auto (AGENTIC_PORT=2322) + serving.server_metrics : auto + run.concurrency : 16 + run.duration : 900 + workloads (conformance_256k) + RESULT_DIR : /run_logs/0 + SUITE_CORPUS_DIR : /tmp/agentx_corpora + +[agentic][DRY_RUN] workload='conformance_256k' source='profile' conc=16 duration=900 + context verdict : OK (--max-context-length 262144) + trace source : --custom-dataset-type weka_trace --input-file /tmp/agentx_corpora/conformance_256k + result dir : /run_logs/0/conformance_256k + command: +aiperf profile --scenario inferencex-agentx-mvp --url http://localhost:2322 --endpoint /v1/chat/completions --endpoint-type chat --streaming --model auto --concurrency 16 --benchmark-duration 900 --random-seed 42 --failed-request-threshold 0.10 --trajectory-start-min-ratio 0.90 --trajectory-start-max-ratio 0.98 --agentic-cache-warmup-duration 60 --warmup-grace-period 1800 --use-server-token-count --tokenizer-trust-remote-code --no-gpu-telemetry --slice-duration 1.0 --max-context-length 262144 --output-artifact-dir /run_logs/0/conformance_256k/aiperf_artifacts --custom-dataset-type weka_trace --input-file /tmp/agentx_corpora/conformance_256k +``` + +Notes: `source: profile` replays near-complete sessions +(`--trajectory-start-*-ratio 0.90/0.98`); hf workloads default to `0.25/0.75`. +The leading `aiperf` is the isolated venv's aiperf CLI path at run time. The +`--scenario` value defaults to `inferencex-agentx-mvp` and is configurable via +`run.scenario` (or the `AGENTIC_SCENARIO` env, with an optional per-workload +`scenario:` override). + +### Verify a corpus by hand (`verify_agentx_profile.py`) + +The suite runs the 13-axis verifier automatically as a pre-gate +(`materialize_corpus()`), but you can also run it directly against an +already-materialized corpus — handy after a `DRY_RUN` plan, when triaging a +`not N/N` abort, or when byte-matching a corpus. The verifier consumes the +**resolved profile JSON** (not the YAML), so first emit the JSON with +`agentx_config.py`, then point `--corpus` at the generated `weka_trace` dir under +`SUITE_CORPUS_DIR/`: + +```bash +# 1. resolve the YAML profile to the JSON the verifier consumes +python3 agentx_config.py --profile profiles/conformance_256k.yaml --emit-json \ + > /tmp/conformance_256k.json + +# 2. verify a materialized corpus against that profile +python3 verify_agentx_profile.py --profile /tmp/conformance_256k.json \ + --corpus /tmp/agentx_corpora/conformance_256k +``` + +**Success** — every axis lands within its band, the footer reads `13/13 axes +within band`, and the exit code is `0` (values below are illustrative): + +``` +corpus=/tmp/agentx_corpora/conformance_256k profile=conformance_256k sessions=200 requests=1974 + +axis measured target verdict +------------------------------------------------------ +Input ISL P50 73,984 74,000 PASS +Input ISL P90 156,000 155,000 PASS +Input ISL P99 232,000 235,000 PASS +Output OSL P50 315 320 PASS +Output OSL P90 3,280 3,300 PASS +Output OSL P99 16,500 17,000 PASS +Turns P50 3 3 PASS +Turns P90 19 20 PASS +Turns P99 103 103 PASS +Delay P50 (s) 4 4 PASS +Delay P90 (s) 30 31 PASS +Delay P99 (s) 236 240 PASS +Cache hit P50 % 89 89 PASS +------------------------------------------------------ +13/13 axes within band +``` + +**Failure** — any axis outside its band prints `off`, the footer reports fewer +than `13/13`, and the exit code is `1`. For example, an OSL P99 that drifts high: + +``` +Output OSL P99 25,600 17,000 off +------------------------------------------------------ +12/13 axes within band +``` + +Inside the suite this same non-`N/N` result aborts the run before any server +contact: `materialize_corpus()` prints the table, then +`corpus 'conformance_256k' failed conformance pre-gate (not N/N)`. Two other +validation errors the verifier emits directly: a profile missing distribution +fields fails with `[verify_agentx_profile] profile missing required field(s): +isl_p, osl_p, delay_p` (exit `2`), and omitting an argument prints `usage: +verify_agentx_profile.py --profile P.json --corpus DIR` (exit `2`). See +[Troubleshooting](#troubleshooting) for what to change when a real corpus fails +the pre-gate. + +### Scenario 10: anti-pattern — circular preset + +```yaml +# profiles/a.yaml -> preset: b +# profiles/b.yaml -> preset: a +workloads: + - { name: loop, preset: a } +``` + +**What happens:** `_merge_preset()` tracks visited presets and raises immediately: + +``` +ValueError: circular preset: a +``` + +The config load fails (`config load failed`) before any corpus work. Break the +cycle so each preset chain terminates at a base profile. + +## `max_model_len` guidance + +- **Prefer `0` / auto** (the default). The suite auto-detects the served window + from `/v1/models` (`resolve_served_max_model_len` in `agentic_lib.sh`) and + never over-estimates it. +- **Pin a value** only when either (a) auto-detect returns `0` because the + endpoint doesn't expose it (e.g. the vLLM disagg `/v1/models` shim), or + (b) you want to force a smaller cap than the model actually supports. +- **ISL-tail interaction:** if a workload's ISL tail exceeds the served window, + `context_compat_check()` WARNs and caps `--max-context-length` at the window; + set `AGENTIC_STRICT_CONTEXT=1` to SKIP that workload instead. + +## Preview / debug + +Two distinct mechanisms: + +**Config-level (Python, no server).** Inspect what the loader resolves: + +```bash +# resolve a YAML profile to JSON (what gen/verify consume) +python3 agentx_config.py --profile profiles/conformance_256k.yaml --emit-json + +# dump the fully-resolved config (serving/run/workloads) +python3 agentx_config.py --config agentic.yaml --dump-json + +# emit SUITE_* globals the bash driver eval's +python3 agentx_config.py --config agentic.yaml --emit-config-shell + +# emit WL_* for one workload (writes the resolved profile JSON to P) +python3 agentx_config.py --config agentic.yaml --workload conformance_256k \ + --profile-out /tmp/p.json --emit-workload-shell +``` + +**Runtime (bash, no server).** `DRY_RUN=1` prints the resolved N-workload plan +and each assembled aiperf command without contacting an endpoint: + +```bash +DRY_RUN=1 AGENTIC_CONFIG=agentic.yaml bash scripts/common/benchmark_agentic_suite.sh +``` + +See [Scenario 9: `DRY_RUN=1` preview](#scenario-9-dry_run1-preview) for the exact +`DRY_RUN` output shape. + +## Troubleshooting + +- **Verify pre-gate fails (`not N/N`).** `materialize_corpus()` aborts the run + when the corpus doesn't match the profile's own `verify:` targets. Check the + profile's `verify.turns_p` / `cache_target` / `band_overrides`, and if you just + edited the profile, regenerate with `SUITE_CORPUS_FORCE=1` (see below). +- **Stale cached corpus.** Corpora are cached at `SUITE_CORPUS_DIR/` + (default `/tmp/agentx_corpora`). Editing a profile does **not** invalidate the + cache; set `SUITE_CORPUS_FORCE=1` to regenerate. +- **Unknown preset name.** `preset: ` loads `profiles/.yaml`; a + missing/misspelled name yields an empty base merge (or a `FileNotFoundError`). + Confirm the file exists under `profiles/`. +- **`AGENTIC_CONFIG` (or a profile/corpus path) not found inside the container.** + On a cluster the launcher runs the suite **inside a container** and only + bind-mounts a fixed set of host paths — `$HOME`, `/shared_inference`, + `/mnt/m2m_nobackup`, the log dir (as `/run_logs`), and the repo checkout (as + `/opt/mooncake-cookbook`) — while passing `AGENTIC_CONFIG` straight through as + an env var (see `scripts/sglang_disagg/run_xPyD_models.slurm`). If the value + points at a host path outside those mounts, the file does not exist in the + container and the suite driver aborts early with `AGENTIC_CONFIG not found: + ` (the same applies to a `source: corpus` `input_dir`). **Fix:** put the + config on a mounted path — typically under `$HOME`, e.g. + `$HOME/MAD/scripts/common/agentx/agentic.example.yaml`, exactly as the + `models.json` agentic entries reference it. + +### Error reference + +Common failure modes across the AgentX layer and its launcher hook, with the +message you will see and the fix. Only failures substantiated by the current +scripts are listed. + +| Symptom / message | Cause | Fix | +| --- | --- | --- | +| `AGENTIC_CONFIG not found: ` (in-container) | config/corpus path is outside the container bind mounts | put it under a mounted path (`$HOME/MAD/...`, `/shared_inference`, `/mnt/m2m_nobackup`); see the entry above | +| `corpus '' failed conformance pre-gate (not N/N)` | generated/edited corpus drifts outside a verify band | check `verify.turns_p` / `cache_target` / `band_overrides`; after editing a profile regenerate with `SUITE_CORPUS_FORCE=1` | +| `could not resolve a served model name (set MODEL explicitly)`; or every request 404s | `--model` doesn't match the server's registered served-model id (`/v1/models` not advertising an id) | let `resolve_served_model_name()` auto-detect, or set `MODEL` / `serving.model` to the exact served id | +| `could not auto-detect max_model_len from ; set serving.max_model_len / MAX_MODEL_LEN` (or, disagg: `could not resolve served max_model_len from sglang worker ...`) | endpoint doesn't expose the context window (e.g. router front-end / vLLM shim) | pin `MAX_MODEL_LEN` / `serving.max_model_len` (see [`max_model_len` guidance](#max_model_len-guidance)) | +| `uv not found on PATH and no cached uv ... set AGENTIC_ALLOW_UV_INSTALL=1 to permit the pinned remote install` | the isolated aiperf venv can't be built because `uv` is missing and the remote install is gated off | install `uv` (>= the pinned version) or set `AGENTIC_ALLOW_UV_INSTALL=1` (the agentic hook already opts in by default) | +| `endpoint not ready after s at ` | no OpenAI-compatible endpoint is up on `AGENTIC_PORT` before the suite runs | start/point to the server first; raise `AGENTIC_ROUTER_READY_TIMEOUT` if it is merely warming | +| `trace download failed after attempts ()` | hf trace corpus could not be fetched | check network/HF auth; raise `AGENTIC_TRACE_DL_ATTEMPTS` / `AGENTIC_TRACE_DL_TIMEOUT`, or pin a different `WEKA_LOADER_OVERRIDE` | +| `filter too aggressive or corpus unreadable for '' (0 sessions?)` | a Tier 2 `filter:` (`max_isl`/`max_turns`/`sample`) removed every session | relax the filter thresholds so at least one session survives | +| `✗ FATAL ERROR: Model '' not found on ALL allocated nodes ...` (launcher) | model weights are missing from the expected paths on the allocated nodes | stage weights under `/mnt/m2m_nobackup/models_blog/` or `/shared_inference/models_blog/` (or set `MODEL_DIR`) on every node | + +## See also + +- How/why the replay mechanism works + accuracy: [HOW_IT_WORKS.md](HOW_IT_WORKS.md). + +Intentionally out of scope here (pointers only): + +- Launcher / disaggregated-serving integration. +- Full env / run-path reference (`agentic_lib.sh` install + endpoint helpers). +- Result-JSON interpretation (`suite_summary.json`, aggregate metrics). +- Post-run health check: `scripts/common/validate_agentic_result.sh`. diff --git a/scripts/common/agentx/agentic.example.yaml b/scripts/common/agentx/agentic.example.yaml new file mode 100644 index 00000000..ba5400f1 --- /dev/null +++ b/scripts/common/agentx/agentic.example.yaml @@ -0,0 +1,92 @@ +# AgentX suite config: one served model, a LIST of workloads run in sequence. +# +# Copy this to agentic.yaml (or point AGENTIC_CONFIG at it) and edit. The suite +# driver (scripts/common/benchmark_agentic_suite.sh) generates+verifies a corpus +# for each source=profile entry (or downloads the dataset for source=hf), runs +# the aiperf replay against the one endpoint, and writes a per-workload result +# dir plus a combined suite summary. Adding "N cases" = adding N list entries. +# +# Environment variables override file values (MODEL, MAX_MODEL_LEN, AGENTIC_PORT, +# AGENTIC_SERVER_METRICS, AGENTIC_CONC, DURATION). AGENTIC_WORKLOAD= runs +# just that one entry. + +serving: + model: auto # 'auto' resolves the served-model id from /v1/models; or set explicitly + max_model_len: 524288 # must be >= each workload's ISL tail (Case-B needs 524288). + # Set 0 (or leave unset / MAX_MODEL_LEN=0) to AUTO-DETECT the served + # window from /v1/models; a set value always wins. If a workload's ISL + # tail exceeds the window the suite WARNs and caps --max-context-length + # (set AGENTIC_STRICT_CONTEXT=1 to SKIP that workload instead). + port: auto # 'auto' -> recipe default (sglang router 2322 / vLLM shim port) + server_metrics: auto # 'auto' -> recipe-provided host:port list; or space-separated endpoints + +run: + concurrency: [2, 4, 8] # scalar or list; >=2 (warmup-credit floor). Swept per workload. + duration: 900 # measured window (s); scenario minimum is 900 for a valid submission + scenario: inferencex-agentx-mvp # aiperf --scenario; override per run (or env AGENTIC_SCENARIO) + +workloads: + # Case-A: generated conformance corpus (ExplainX targets), inherits the shipped preset. + - name: conformance_256k + source: profile + preset: conformance_256k + + # Case-B: longer-context conformance corpus; needs max_model_len 524288. + - name: conformance_512k + source: profile + preset: conformance_512k + + # inferencex: download a real captured trace corpus instead of generating one. + # The loader name sets the context-gating ISL tail (_256k -> 262144; full + # corpus -> conservative ~1M, override with AGENTIC_HF_ISL_TAIL). + - name: inferencex + source: hf + loader: semianalysis_cc_traces_weka_062126_256k + + # inferencex via a reusable HF preset (profiles/inferencex_256k.yaml): inherits + # the loader + bundled concurrency sweep + duration + Tier 1 knobs. + - name: inferencex_preset + preset: inferencex_256k + + # Same preset, but override the sweep and trim the corpus to fit a smaller model. + - name: inferencex_small + preset: inferencex_256k + concurrency: [2] # entry overrides the preset's [2, 4, 8] sweep + num_dataset_entries: 50 # Tier 1: pull fewer sessions + trajectory: { min: 0.30, max: 0.80 } + filter: # Tier 2: local subset/trim (download once, then filter) + max_isl: 200000 # drop sessions with any turn over 200k input tokens + max_turns: 40 # truncate each session to its first 40 turns + sample: 100 # randomly keep 100 sessions (seed=42) + + # Replay an EXISTING on-disk weka_trace corpus as-is (no generate, no download). + # `input_dir` must be a directory of session_XXXXX.json files. Uncomment a + # `preset:` (or inline distribution fields) to run verify_agentx_profile.py as + # an optional pre-gate; with none, the corpus is replayed without verification. +# - name: my_corpus +# source: corpus +# input_dir: /tmp/agentx_corpora/conformance_256k +# # preset: conformance_256k # optional: verify the corpus against this profile +# # isl_tail: 262144 # optional: override the context-gating ISL tail + + # A user-added custom case, defined inline (no preset). Copy profiles/custom.example.yaml + # fields here, or reference your own profile via `preset: ` after dropping it in profiles/. +# - name: my_case +# source: profile +# model_tag: GLM-5.2-MXFP4 +# id_prefix: my_case +# seed: 42 +# n_sessions: 150 +# block_size: 64 +# isl_p: [48000, 120000, 200000] +# osl_p: [256, 2000, 9000] +# delay_p: [3, 20, 180] +# turns: +# values: [2, 3, 4, 6, 10, 20, 45, 103] +# weights: [20, 24, 20, 12, 8, 6, 7, 3] +# cache_hit: [0.88, 0.90] +# clamps: +# isl: [1200, 205000] +# osl: [8, 20000] +# delay: [1, 600] +# # per-workload run override (optional): concurrency: [2] diff --git a/scripts/common/agentx/agentx_config.py b/scripts/common/agentx/agentx_config.py new file mode 100644 index 00000000..54f85a33 --- /dev/null +++ b/scripts/common/agentx/agentx_config.py @@ -0,0 +1,580 @@ +#!/usr/bin/env python3 +"""AgentX suite config loader. + +Parses the suite config (agentic.yaml) and the per-workload profile files, then +yields fully-resolved, per-workload parameter sets to the bash suite driver. + +The config carries a `serving:` block (one served endpoint for the whole run), a +`run:` block (default concurrency/duration), and a `workloads:` LIST. Each +workload entry is either: + - source: profile -> carries the distribution params inline, or `preset: conformance_256k` + to inherit scripts/common/agentx/profiles/conformance_256k.yaml + - source: hf -> carries a `loader` name (an aiperf --public-dataset id) + - source: corpus -> replays an existing on-disk weka_trace corpus (`input_dir`) + as-is; optional profile/preset enables a verify pre-gate + +Design notes: + * Runs in the aiperf venv. Uses PyYAML if importable; otherwise falls back to a + small pure-stdlib parser for the restricted YAML subset the shipped files use + (block maps/seqs, inline [..]/{..} flows, scalars, comments). No hard third- + party dependency. Set AGENTX_YAML_FALLBACK=1 to force the fallback (tests). + * Environment variables OVERRIDE file values: + MODEL -> serving.model MAX_MODEL_LEN -> serving.max_model_len + AGENTIC_PORT -> serving.port AGENTIC_SERVER_METRICS -> serving.server_metrics + AGENTIC_CONC -> run.concurrency DURATION -> run.duration + * Single-workload shorthand: AGENTIC_WORKLOAD= restricts the run to that + one entry (a 1-entry list). With no --config, conformance_256k/conformance_512k/inferencex are + synthesized from the shipped presets so the shorthand works standalone. + +CLI: + agentx_config.py --profile --emit-json + Resolve a single profile file (yaml/json) and print it as JSON. + agentx_config.py --config --emit-config-shell + Print SUITE_* globals + SUITE_WORKLOAD_NAMES (eval-able in bash). + agentx_config.py --config --workload [--profile-out P] --emit-workload-shell + Resolve one workload; print WL_* (eval-able). For source=profile, write the + resolved profile JSON to P (for gen/verify). + agentx_config.py --config --dump-json + Print the fully-resolved config (serving/run/workloads) as JSON. +""" +import json, os, sys + +HERE = os.path.dirname(os.path.abspath(__file__)) +PROFILES_DIR = os.path.join(HERE, "profiles") + + +def _argval(it, flag): + try: + return next(it) + except StopIteration: + sys.stderr.write(f"[agentx_config] {flag} requires a value\n") + raise SystemExit(2) + +_RUN_KEYS = ("concurrency", "duration") +# Keys that steer resolution / run knobs but are NOT part of a generator profile +# dict (so they are stripped when building the source=profile profile JSON). +_CONTROL_KEYS = ("source", "preset", "loader", "filter", + "num_dataset_entries", "trajectory", + "input_dir", "isl_tail", "scenario") + _RUN_KEYS + + +# -------------------------------------------------------------------------- +# YAML loading: PyYAML if available, else a small restricted-subset fallback. +# -------------------------------------------------------------------------- +def _yaml_load(text): + if os.environ.get("AGENTX_YAML_FALLBACK", "") != "1": + try: + import yaml # type: ignore + return yaml.safe_load(text) + except ImportError: + pass + return _fallback_load(text) + + +def _scalar(s): + s = s.strip() + if s == "": + return None + if (s[0] == '"' and s[-1] == '"') or (s[0] == "'" and s[-1] == "'"): + return s[1:-1] + low = s.lower() + if low in ("null", "~"): + return None + if low == "true": + return True + if low == "false": + return False + try: + return int(s) + except ValueError: + pass + try: + return float(s) + except ValueError: + pass + return s + + +def _split_top(s): + """Split a flow-collection body on top-level commas (respects nesting/quotes).""" + parts = [] + depth = 0 + inq = None + cur = "" + for ch in s: + if inq: + cur += ch + if ch == inq: + inq = None + elif ch in "\"'": + inq = ch + cur += ch + elif ch in "[{": + depth += 1 + cur += ch + elif ch in "]}": + depth -= 1 + cur += ch + elif ch == "," and depth == 0: + parts.append(cur) + cur = "" + else: + cur += ch + if cur.strip() != "" or parts: + parts.append(cur) + return parts + + +def _parse_node(s): + s = s.strip() + if s.startswith("["): + inner = s[1:-1].strip() + return [] if inner == "" else [_parse_node(p) for p in _split_top(inner)] + if s.startswith("{"): + inner = s[1:-1].strip() + d = {} + if inner: + for p in _split_top(inner): + k, _, v = p.partition(":") + d[str(_scalar(k))] = _parse_node(v) + return d + return _scalar(s) + + +def _strip_comment(line): + inq = None + out = "" + for i, ch in enumerate(line): + if inq: + out += ch + if ch == inq: + inq = None + elif ch in "\"'": + inq = ch + out += ch + elif ch == "#" and (i == 0 or line[i - 1] == " "): + break + else: + out += ch + return out + + +def _parse_block(lines, i, indent): + _, content = lines[i] + if content.startswith("-"): + seq = [] + while i < len(lines): + ind, c = lines[i] + if ind != indent or not c.startswith("-"): + break + rest = c[1:].lstrip() + item_indent = indent + (len(c) - len(c[1:].lstrip())) + if rest == "": + i += 1 + if i < len(lines) and lines[i][0] > indent: + val, i = _parse_block(lines, i, lines[i][0]) + else: + val = None + seq.append(val) + else: + sub = [(item_indent, rest)] + i += 1 + while i < len(lines) and lines[i][0] > indent: + sub.append(lines[i]) + i += 1 + val, _ = _parse_block(sub, 0, item_indent) + seq.append(val) + return seq, i + d = {} + while i < len(lines): + ind, c = lines[i] + if ind != indent: + break + key, _, rest = c.partition(":") + key = str(_scalar(key.strip())) + rest = rest.strip() + if rest == "": + i += 1 + if i < len(lines) and lines[i][0] > indent: + val, i = _parse_block(lines, i, lines[i][0]) + elif (i < len(lines) and lines[i][0] == indent + and lines[i][1].startswith("-")): + val, i = _parse_block(lines, i, indent) # same-indent block sequence + else: + val = None + d[key] = val + else: + d[key] = _parse_node(rest) + i += 1 + return d, i + + +def _fallback_load(text): + lines = [] + for raw in text.splitlines(): + line = _strip_comment(raw) + if line.strip() == "": + continue + indent = len(line) - len(line.lstrip(" ")) + lines.append((indent, line.strip())) + if not lines: + return None + val, _ = _parse_block(lines, 0, lines[0][0]) + return val + + +def _load_file(path): + with open(path) as f: + text = f.read() + if path.endswith(".json"): + return json.loads(text) + return _yaml_load(text) + + +# -------------------------------------------------------------------------- +# Resolution +# -------------------------------------------------------------------------- +def load_profile_file(path): + return _load_file(path) + + +def _merge_preset(entry, _visited): + """Return `entry` merged over its `preset:` chain (entry keys win). + + Works for any source: a profiles/.yaml may declare distribution + params (source=profile), an hf loader + Tier 1/Tier 2 knobs (source=hf), + and/or run knobs (concurrency/duration). Circular references raise. + """ + name = entry.get("preset") + if not name: + return dict(entry) + if name in _visited: + raise ValueError(f"circular preset: {name}") + _visited.add(name) + preset_path = os.path.join(PROFILES_DIR, f"{name}.yaml") + if not os.path.isfile(preset_path): + raise ValueError(f"preset not found: {name}") + base = _merge_preset( + load_profile_file(preset_path) or {}, _visited) + for k, v in entry.items(): + if k == "preset": + continue + base[k] = v + return base + + +def _profile_from_merged(merged, name): + """Build a generator profile dict from a merged workload dict.""" + profile = {k: v for k, v in merged.items() if k not in _CONTROL_KEYS} + profile.setdefault("name", name) + return profile + + +def _validate_tier1(name, nde, tmin, tmax): + if nde is not None and int(nde) < 1: + raise ValueError(f"workload '{name}': num_dataset_entries must be >= 1") + if tmin is not None or tmax is not None: + lo = 0.0 if tmin is None else float(tmin) + hi = 1.0 if tmax is None else float(tmax) + if not (0.0 <= lo <= hi <= 1.0): + raise ValueError( + f"workload '{name}': trajectory requires 0.0 <= min <= max <= 1.0 " + f"(got min={tmin}, max={tmax})") + + +def _hf_isl_tail(loader): + """ISL tail (max input tokens) for an hf loader, for context gating. + + Option A explicit matching: the `_256k` suffix is definitional and checked + FIRST (wins over the date substring); the full-corpus loaders use a + conservative HIGH default. The gate caps --max-context-length at the served + window, so over-estimation only over-WARNs. Override with AGENTIC_HF_ISL_TAIL. + """ + env = os.environ.get("AGENTIC_HF_ISL_TAIL") + if env: + return int(env) + if loader.endswith("_256k"): + return 262144 # definitional: 256k cap + if "062126" in loader or "061526" in loader: + # full corpus: conservative ~1M. Measured max per-turn ISL is 989824 for + # both 062126 and 061526 (in-container, all sessions); 1048576 (2^20) is a + # safe over-estimate and rounds to the same power-of-two window as 989824. + return 1048576 + return 1048576 # unknown loader -> conservative default (errs to WARN) + + +def _resolve_workload_entry(entry, _visited=None): + """Merge a workload entry with its preset (any source) and resolve it.""" + if _visited is None: + _visited = set() + merged = _merge_preset(entry, _visited) + src = merged.get("source", "profile") + name = entry.get("name") or entry.get("preset") or "workload" + nde = merged.get("num_dataset_entries") + traj = merged.get("trajectory") or {} + tmin = traj.get("min") + tmax = traj.get("max") + _validate_tier1(name, nde, tmin, tmax) + wl = { + "name": name, + "source": src, + "concurrency": _norm_concurrency(merged.get("concurrency")), + "duration": merged.get("duration"), + "num_dataset_entries": nde, + "traj_min": tmin, + "traj_max": tmax, + "scenario": merged.get("scenario"), + } + if src == "profile": + prof = _profile_from_merged(merged, name) + wl["profile"] = prof + wl["isl_tail"] = _isl_tail(prof) + elif src == "hf": + wl["loader"] = merged.get("loader", "") + wl["isl_tail"] = _hf_isl_tail(wl["loader"]) + wl["filter"] = merged.get("filter") or {} + elif src == "corpus": + input_dir = merged.get("input_dir") + if not input_dir: + raise ValueError(f"workload '{name}': source=corpus requires input_dir") + wl["input_dir"] = input_dir + # Optional verification profile: build one only if the entry supplies + # distribution fields (inline) or inherited a preset. Otherwise replay + # the corpus as-is with no pre-gate. + prof = _profile_from_merged(merged, name) + has_profile = any(k != "name" for k in prof) or bool(entry.get("preset")) + if has_profile: + wl["profile"] = prof + # ISL tail (context gating): explicit key > profile-derived > conservative + # default (same fallback the hf-unknown path uses; the gate caps to the + # served window, so over-estimation only over-WARNs). + if merged.get("isl_tail") is not None: + wl["isl_tail"] = int(merged["isl_tail"]) + elif has_profile: + wl["isl_tail"] = _isl_tail(prof) + else: + wl["isl_tail"] = 1048576 + else: + raise ValueError(f"workload '{name}': unknown source '{src}'") + return wl + + +def _isl_tail(profile): + clamps = profile.get("clamps", {}) or {} + if "isl" in clamps: + return int(clamps["isl"][1]) + if "isl_p" not in profile: + raise ValueError( + "profile missing 'isl_p' (or clamps.isl) required to derive ISL tail") + return int(profile["isl_p"][2]) + + +def _norm_concurrency(v): + if v is None: + return None + if isinstance(v, (list, tuple)): + return " ".join(str(int(x)) for x in v) + return " ".join(str(int(x)) for x in str(v).replace(",", " ").split()) + + +def resolve_config(config): + serving = dict(config.get("serving", {}) or {}) + run = dict(config.get("run", {}) or {}) + + env = os.environ + if env.get("MODEL"): + serving["model"] = env["MODEL"] + if env.get("MAX_MODEL_LEN"): + serving["max_model_len"] = int(env["MAX_MODEL_LEN"]) + if env.get("AGENTIC_PORT"): + serving["port"] = env["AGENTIC_PORT"] + if env.get("AGENTIC_SERVER_METRICS"): + serving["server_metrics"] = env["AGENTIC_SERVER_METRICS"] + if env.get("AGENTIC_CONC"): + run["concurrency"] = env["AGENTIC_CONC"] + if env.get("DURATION"): + run["duration"] = int(env["DURATION"]) + if env.get("AGENTIC_SCENARIO"): + run["scenario"] = env["AGENTIC_SCENARIO"] + + serving.setdefault("model", "auto") + serving.setdefault("max_model_len", 0) + serving.setdefault("port", "auto") + serving.setdefault("server_metrics", "auto") + run.setdefault("concurrency", 16) + run.setdefault("duration", 900) + run.setdefault("scenario", "inferencex-agentx-mvp") + + workloads = [] + for entry in config.get("workloads", []) or []: + workloads.append(_resolve_workload_entry(entry)) + + want = env.get("AGENTIC_WORKLOAD") + if want: + filtered = [w for w in workloads if w["name"] == want] + if not filtered: + raise ValueError(f"AGENTIC_WORKLOAD='{want}' not found in workloads list") + workloads = filtered + + return {"serving": serving, "run": run, "workloads": workloads} + + +# Presets used for the config-less single-workload shorthand. +_HF_PRESETS = {"inferencex": "semianalysis_cc_traces_weka_062126_256k"} + + +def _synth_config_from_env(): + want = os.environ.get("AGENTIC_WORKLOAD") + if not want: + raise SystemExit("no --config and no AGENTIC_WORKLOAD set") + if want in _HF_PRESETS: + wl = {"name": want, "source": "hf", "loader": _HF_PRESETS[want]} + else: + wl = {"name": want, "source": "profile", "preset": want} + return {"serving": {}, "run": {}, "workloads": [wl]} + + +def _load_config_arg(path): + if path: + return resolve_config(_load_file(path)) + return resolve_config(_synth_config_from_env()) + + +# -------------------------------------------------------------------------- +# Shell emitters +# -------------------------------------------------------------------------- +def _sh(v): + return "'" + str(v).replace("'", "'\\''") + "'" + + +def emit_config_shell(resolved): + s = resolved["serving"] + r = resolved["run"] + names = [w["name"] for w in resolved["workloads"]] + out = [] + out.append(f"SUITE_SERVING_MODEL={_sh(s['model'])}") + out.append(f"SUITE_MAX_MODEL_LEN={_sh(s['max_model_len'])}") + out.append(f"SUITE_PORT={_sh(s['port'])}") + out.append(f"SUITE_SERVER_METRICS={_sh(s['server_metrics'])}") + out.append(f"SUITE_CONCURRENCY={_sh(_norm_concurrency(r['concurrency']))}") + out.append(f"SUITE_DURATION={_sh(r['duration'])}") + out.append(f"SUITE_SCENARIO={_sh(r['scenario'])}") + out.append(f"SUITE_WORKLOAD_NAMES={_sh(' '.join(names))}") + return "\n".join(out) + + +def emit_workload_shell(resolved, name, profile_out): + wl = next((w for w in resolved["workloads"] if w["name"] == name), None) + if wl is None: + raise SystemExit(f"workload '{name}' not in resolved config") + r = resolved["run"] + conc = wl["concurrency"] or _norm_concurrency(r["concurrency"]) + dur = wl["duration"] if wl["duration"] is not None else r["duration"] + def _opt(v): + return "" if v is None else v + + out = [ + f"WL_NAME={_sh(wl['name'])}", + f"WL_SOURCE={_sh(wl['source'])}", + f"WL_CONCURRENCY={_sh(conc)}", + f"WL_DURATION={_sh(dur)}", + f"WL_ISL_TAIL={_sh(wl.get('isl_tail', 0))}", + f"WL_NUM_DATASET_ENTRIES={_sh(_opt(wl.get('num_dataset_entries')))}", + f"WL_TRAJ_MIN={_sh(_opt(wl.get('traj_min')))}", + f"WL_TRAJ_MAX={_sh(_opt(wl.get('traj_max')))}", + f"WL_SCENARIO={_sh(_opt(wl.get('scenario')))}", + ] + if wl["source"] == "hf": + out.append(f"WL_LOADER={_sh(wl.get('loader', ''))}") + f = wl.get("filter") or {} + out.append(f"WL_FILTER_MAX_ISL={_sh(_opt(f.get('max_isl')))}") + out.append(f"WL_FILTER_MAX_TURNS={_sh(_opt(f.get('max_turns')))}") + out.append(f"WL_FILTER_SAMPLE={_sh(_opt(f.get('sample')))}") + out.append("WL_PROFILE_FILE=''") + out.append("WL_INPUT_DIR=''") + elif wl["source"] == "corpus": + out.append("WL_LOADER=''") + out.append("WL_FILTER_MAX_ISL=''") + out.append("WL_FILTER_MAX_TURNS=''") + out.append("WL_FILTER_SAMPLE=''") + out.append(f"WL_INPUT_DIR={_sh(wl['input_dir'])}") + # Only write/point at a profile JSON if the entry supplied one (optional + # verification pre-gate); otherwise replay the corpus as-is. + if profile_out and wl.get("profile"): + with open(profile_out, "w") as f: + json.dump(wl["profile"], f) + out.append(f"WL_PROFILE_FILE={_sh(profile_out)}") + else: + out.append("WL_PROFILE_FILE=''") + else: + out.append("WL_LOADER=''") + out.append("WL_FILTER_MAX_ISL=''") + out.append("WL_FILTER_MAX_TURNS=''") + out.append("WL_FILTER_SAMPLE=''") + out.append("WL_INPUT_DIR=''") + if profile_out: + with open(profile_out, "w") as f: + json.dump(wl["profile"], f) + out.append(f"WL_PROFILE_FILE={_sh(profile_out)}") + else: + out.append("WL_PROFILE_FILE=''") + out.append(f"WL_MODEL_TAG={_sh(wl['profile'].get('model_tag', ''))}") + return "\n".join(out) + + +def main(argv): + config_path = None + profile_path = None + workload = None + profile_out = None + mode = None + it = iter(argv) + for a in it: + if a == "--config": + config_path = _argval(it, a) + elif a == "--profile": + profile_path = _argval(it, a) + elif a == "--workload": + workload = _argval(it, a) + elif a == "--profile-out": + profile_out = _argval(it, a) + elif a in ("--emit-json", "--emit-config-shell", "--emit-workload-shell", "--dump-json"): + mode = a + elif a in ("-h", "--help"): + print(__doc__) + return 0 + else: + sys.stderr.write(f"[agentx_config] unknown arg: {a}\n") + return 2 + + if mode == "--emit-json": + if not profile_path: + sys.stderr.write("--emit-json requires --profile\n") + return 2 + prof = _load_file(profile_path) + json.dump(prof, sys.stdout) + sys.stdout.write("\n") + return 0 + + resolved = _load_config_arg(config_path) + + if mode == "--dump-json": + json.dump(resolved, sys.stdout, indent=2) + sys.stdout.write("\n") + return 0 + if mode == "--emit-config-shell": + print(emit_config_shell(resolved)) + return 0 + if mode == "--emit-workload-shell": + if not workload: + sys.stderr.write("--emit-workload-shell requires --workload\n") + return 2 + print(emit_workload_shell(resolved, workload, profile_out)) + return 0 + + sys.stderr.write("no mode selected (see --help)\n") + return 2 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/scripts/common/agentx/filter_weka_corpus.py b/scripts/common/agentx/filter_weka_corpus.py new file mode 100644 index 00000000..c564b52d --- /dev/null +++ b/scripts/common/agentx/filter_weka_corpus.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +"""Filter/trim a downloaded weka_trace corpus to fit a model, then re-emit it as +per-session JSON files aiperf can replay via `--input-file `. + +The HuggingFace weka corpora ship as a single `traces.jsonl` (one session per +line) whose per-session schema matches what gen_agentx_profile.py emits: + {id, models, block_size, hash_id_scope, requests:[{t, in, out, hash_ids, ...}]} +This reads that download (a directory holding traces.jsonl and/or per-session +*.json files, or a single .jsonl/.json file) and writes filtered +session_XXXXX.json files into --out-dir. + +Filters, applied IN ORDER: + 1. --max-turns N: TRUNCATE each session to its first N requests (keeps the + growing-prefix structure intact). + 2. --max-isl N : PER-TURN drop -- discard a session if ANY of its (already + truncated) turns' `in` exceeds N (it won't fit the window). + Applied AFTER truncation so "trim to first N turns, then keep + what fits" works for growing-prefix corpora whose late turns + always exceed a small window. No token trim. + 3. --sample N : RANDOM subset of N sessions with a fixed seed=42 + (reproducible) when more than N remain; else keep all. + +An empty result is an error (exit 1) so an over-aggressive filter fails loudly. + +Pure Python stdlib only (json/os/sys/random/glob). +""" +import json, os, sys, random, glob + + +def _argval(it, flag): + try: + return next(it) + except StopIteration: + sys.stderr.write(f"[filter_weka_corpus] {flag} requires a value\n") + raise SystemExit(2) + + +def _iter_file(path): + if path.endswith(".jsonl"): + with open(path) as fh: + for line in fh: + line = line.strip() + if line: + yield json.loads(line) + else: + with open(path) as fh: + obj = json.load(fh) + if isinstance(obj, list): + for s in obj: + yield s + else: + yield obj + + +def _iter_sessions(path): + if os.path.isdir(path): + files = sorted(glob.glob(os.path.join(path, "*.jsonl"))) + \ + sorted(glob.glob(os.path.join(path, "*.json"))) + for f in files: + yield from _iter_file(f) + else: + yield from _iter_file(path) + + +def _max_turn_isl(session): + m = 0 + for req in session.get("requests", []) or []: + v = req.get("in") + if isinstance(v, int) and v > m: + m = v + return m + + +def filter_corpus(sessions, max_isl=None, max_turns=None, sample=None): + out = [] + for s in sessions: + if max_turns is not None: + s = dict(s) + s["requests"] = list(s.get("requests", []) or [])[:max_turns] + if max_isl is not None and _max_turn_isl(s) > max_isl: + continue + out.append(s) + if sample is not None and len(out) > sample: + out = random.Random(42).sample(out, sample) + return out + + +def main(argv): + inp = out_dir = None + max_isl = max_turns = sample = None + it = iter(argv) + for a in it: + if a == "--input": + inp = _argval(it, a) + elif a == "--out-dir": + out_dir = _argval(it, a) + elif a == "--max-isl": + max_isl = int(_argval(it, a)) + elif a == "--max-turns": + max_turns = int(_argval(it, a)) + elif a == "--sample": + sample = int(_argval(it, a)) + elif a in ("-h", "--help"): + print(__doc__) + return 0 + else: + sys.stderr.write(f"[filter_weka_corpus] unknown arg: {a}\n") + return 2 + if not inp or not out_dir: + sys.stderr.write("usage: filter_weka_corpus.py --input SRC --out-dir DIR " + "[--max-isl N] [--max-turns N] [--sample N]\n") + return 2 + + sessions = list(_iter_sessions(inp)) + kept = filter_corpus(sessions, max_isl, max_turns, sample) + if not kept: + sys.stderr.write( + f"[filter_weka_corpus] filter too aggressive: 0 sessions from {len(sessions)} " + f"(max_isl={max_isl}, max_turns={max_turns}, sample={sample})\n") + return 1 + + os.makedirs(out_dir, exist_ok=True) + for i, s in enumerate(kept): + with open(os.path.join(out_dir, f"session_{i:05d}.json"), "w") as fh: + json.dump(s, fh) + print(f"wrote {len(kept)}/{len(sessions)} sessions -> {out_dir} " + f"(max_isl={max_isl}, max_turns={max_turns}, sample={sample})") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/scripts/common/agentx/gen_agentx_profile.py b/scripts/common/agentx/gen_agentx_profile.py new file mode 100644 index 00000000..f61e721d --- /dev/null +++ b/scripts/common/agentx/gen_agentx_profile.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python3 +"""Generic, seed-deterministic AgentX WekaTrace corpus synthesizer. + +ONE generator that reads a single workload PROFILE (a set of distribution +targets + a seed) and emits a reproducible `weka_trace` corpus: one +`session_XXXXX.json` per session, in the exact schema aiperf's +`inferencex-agentx-mvp` scenario consumes: + + {id, models, block_size, hash_id_scope, + requests:[{t, type, model, in, out, hash_ids, api_time, think_time, stop}]} + +This is a straight refactor of ROCm/MAD #173's `gen_caseA_conformance.py` / +`gen_caseB_conformance.py` into a single parameterized generator. The SAMPLING +ALGORITHM is preserved byte-for-byte (lognormal-from-percentiles fit, +weighted-choice turns, growing-prefix hash_ids reuse, identical RNG call order) +so that invoking it with the Case-A / Case-B preset parameters reproduces #173's +committed corpora exactly and passes the conformance verifier 13/13. + +Constructed (not engine-captured) because serving cost depends on token COUNTS +and cache STRUCTURE, not token content: a request of ISL=N, OSL=M with a given +prefix-reuse pattern does the same prefill+decode work regardless of whether the +tokens are "real". A conformance trace therefore exercises the engine identically +to a captured one of the same shape while guaranteeing the target distribution. + +Profile schema (JSON/dict): + { + "name": "conformance_256k", # informational + "model_tag": "GLM-5.2-MXFP4", # written to requests[].model + models[] + "id_prefix": "caseA", # session-id salt prefix (keep "caseA" to + # byte-match #173, which used it for both) + "seed": 42, + "n_sessions": 200, + "block_size": 64, + "isl_p": [74000, 155000, 235000], # ISL P50/P90/P99 + "osl_p": [320, 3300, 17000], # OSL P50/P90/P99 + "delay_p": [4, 31, 240], # inter-turn delay P50/P90/P99 (s) + "turns": {"values": [1,2,3,...], "weights": [22,24,20,...]}, + "cache_hit": [0.88, 0.90], # per-turn prefix-reuse band + "clamps": {"isl": [1200, 245000], "osl": [8, 20000], "delay": [1, 600]} + } + +Pure Python stdlib only (json/os/sys/math/random/hashlib) so it runs anywhere +without a third-party install. YAML profiles are resolved to JSON by the config +loader (scripts/common/agentx/agentx_config.py) before being handed here. + +Usage: + gen_agentx_profile.py --profile --out-dir [overrides] + gen_agentx_profile.py --profile-json '{...}' --out-dir [overrides] + overrides: --n-sessions N --seed S --model-tag TAG --id-prefix P --block-size B +""" +import json, os, sys, math, random, hashlib + +DEFAULT_MODEL_TAG = "GLM-5.2-MXFP4" +DEFAULT_ID_PREFIX = "caseA" # #173 used the literal "caseA" salt for BOTH cases +DEFAULT_BLOCK = 64 + +REQUIRED_PROFILE_FIELDS = ("isl_p", "osl_p", "delay_p", "turns", "cache_hit") + + +def _require_fields(profile): + missing = [k for k in REQUIRED_PROFILE_FIELDS if k not in profile] + if missing: + raise SystemExit( + "[gen_agentx_profile] profile missing required field(s): " + + ", ".join(missing)) + + +def _argval(it, flag): + try: + return next(it) + except StopIteration: + sys.stderr.write(f"[gen_agentx_profile] {flag} requires a value\n") + raise SystemExit(2) + + +def lognorm_from_p(p50, p90, p99): + """Return (mu, sigma) of a lognormal matched to a (P50,P90,P99) triple. + mu=ln(p50); sigma averages the p90- and p99-implied estimates + (z90=1.2816, z99=2.3263).""" + mu = math.log(p50) + s90 = (math.log(p90) - mu) / 1.2816 + s99 = (math.log(p99) - mu) / 2.3263 + sigma = (s90 + s99) / 2.0 + return mu, sigma + + +def generate_corpus(profile, out_dir): + """Materialize a weka_trace corpus for one workload profile into out_dir. + + The RNG call order (turns choice, then per-turn ISL gauss, OSL gauss, + reuse uniform [t>0], delay gauss [t>0]) is identical to #173 so preset + profiles reproduce the committed corpora byte-for-byte. Returns n_sessions. + """ + _require_fields(profile) + seed = int(profile.get("seed", 42)) + n = int(profile.get("n_sessions", 200)) + block = int(profile.get("block_size", DEFAULT_BLOCK)) + if block < 1: + raise SystemExit("[gen] block_size must be >= 1") + model_tag = str(profile.get("model_tag", DEFAULT_MODEL_TAG)) + id_prefix = str(profile.get("id_prefix", DEFAULT_ID_PREFIX)) + + isl_p = profile["isl_p"] + osl_p = profile["osl_p"] + delay_p = profile["delay_p"] + turns = profile["turns"] + turns_values = list(turns["values"]) + turns_weights = list(turns["weights"]) + if not turns_values or not turns_weights: + raise SystemExit("[gen] turns must have non-empty values and weights") + if len(turns_values) != len(turns_weights): + raise SystemExit("[gen] turns values and weights must have equal length") + if any(w <= 0 for w in turns_weights): + raise SystemExit("[gen] turns weights must all be > 0") + cache_lo, cache_hi = profile["cache_hit"] + clamps = profile.get("clamps", {}) + isl_lo, isl_hi = clamps.get("isl", [1200, 245000]) + osl_lo, osl_hi = clamps.get("osl", [8, 20000]) + dly_lo, dly_hi = clamps.get("delay", [1, 600]) + + rng = random.Random(seed) + + ISL_mu, ISL_sig = lognorm_from_p(*isl_p) + OSL_mu, OSL_sig = lognorm_from_p(*osl_p) + DLY_mu, DLY_sig = lognorm_from_p(*delay_p) + + def samp(mu, sig, lo, hi): + return int(min(hi, max(lo, math.exp(mu + sig * rng.gauss(0, 1))))) + + def sample_turns(): + return rng.choices(turns_values, weights=turns_weights)[0] + + def make_session(idx): + sid = hashlib.blake2b(f"{id_prefix}-{seed}-{idx:05d}".encode(), digest_size=18).hexdigest() + n_turns = sample_turns() + prefix_blocks = [] # accumulated shared prefix (block hashes) + salt = f"{seed}:{idx}" + reqs = [] + t_clock = 0.0 + for turn in range(n_turns): + isl = samp(ISL_mu, ISL_sig, isl_lo, isl_hi) + osl = samp(OSL_mu, OSL_sig, osl_lo, osl_hi) + # Floor division is deliberate: it keeps generated corpora byte-for-byte + # identical to the #173 committed corpora / pre-gate verify. Do NOT + # change to rounding/ceil — it would break reproducibility. + total_blocks = max(1, isl // block) + if turn == 0: + new_blocks = total_blocks + else: + reuse = min(len(prefix_blocks), int(total_blocks * rng.uniform(cache_lo, cache_hi))) + new_blocks = max(1, total_blocks - reuse) + reuse_slice = prefix_blocks[:total_blocks - new_blocks] + new_ids = [] + for b in range(new_blocks): + h = int(hashlib.blake2b(f"{salt}:{turn}:{b}".encode(), digest_size=8).hexdigest(), 16) & 0x7FFFFFFFFFFFFFFF + new_ids.append(h) + hash_ids = reuse_slice + new_ids + prefix_blocks = hash_ids + think = 0.0 if turn == 0 else float(round(samp(DLY_mu, DLY_sig, dly_lo, dly_hi), 2)) + t_clock += think + reqs.append({ + "t": round(t_clock, 3), + "type": "n", + "model": model_tag, + "in": isl, + "out": osl, + "hash_ids": hash_ids, + "api_time": 0.0, + "think_time": think, + "stop": "stop", + }) + return {"id": sid, "models": [model_tag], "block_size": block, + "hash_id_scope": "local", "requests": reqs} + + os.makedirs(out_dir, exist_ok=True) + for i in range(n): + s = make_session(i) + with open(os.path.join(out_dir, f"session_{i:05d}.json"), "w") as fh: + json.dump(s, fh) + return n + + +def _load_profile(path): + with open(path) as f: + return json.load(f) + + +def main(argv): + profile = None + out_dir = None + overrides = {} + it = iter(argv) + for a in it: + if a in ("--profile", "-p"): + profile = _load_profile(_argval(it, a)) + elif a == "--profile-json": + profile = json.loads(_argval(it, a)) + elif a in ("--out-dir", "-o"): + out_dir = _argval(it, a) + elif a == "--n-sessions": + overrides["n_sessions"] = int(_argval(it, a)) + elif a == "--seed": + overrides["seed"] = int(_argval(it, a)) + elif a == "--model-tag": + overrides["model_tag"] = _argval(it, a) + elif a == "--id-prefix": + overrides["id_prefix"] = _argval(it, a) + elif a == "--block-size": + overrides["block_size"] = int(_argval(it, a)) + elif a in ("-h", "--help"): + print(__doc__) + return 0 + else: + sys.stderr.write(f"[gen_agentx_profile] unknown arg: {a}\n") + return 2 + if profile is None or out_dir is None: + sys.stderr.write("usage: gen_agentx_profile.py --profile P.json --out-dir DIR [overrides]\n") + return 2 + profile.update(overrides) + n = generate_corpus(profile, out_dir) + print(f"wrote {n} sessions -> {out_dir} " + f"(seed={profile.get('seed', 42)}, block={profile.get('block_size', DEFAULT_BLOCK)}, " + f"model={profile.get('model_tag', DEFAULT_MODEL_TAG)})") + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/scripts/common/agentx/profiles/README.md b/scripts/common/agentx/profiles/README.md new file mode 100644 index 00000000..7af9bf2f --- /dev/null +++ b/scripts/common/agentx/profiles/README.md @@ -0,0 +1,181 @@ +# Authoring AgentX profiles & presets + +A **profile** is ~6 distribution targets plus a `seed` and a session count. +`gen_agentx_profile.py` turns it into a reproducible `weka_trace` corpus, and +`verify_agentx_profile.py` gates that corpus against the same targets. There is +**no per-case code** — you add a workload by adding a profile file here. + +- New to the core concepts? Start at [../README.md](../README.md). +- Want worked, copy-paste examples? See [../README.md#examples](../README.md#examples). + +## Anatomy of a profile + +Fields below are shown against [conformance_256k.yaml](conformance_256k.yaml) (the +generalized ROCm/MAD #173 Case-A profile). They are consumed by +`generate_corpus()` in `gen_agentx_profile.py` and `verify()` in +`verify_agentx_profile.py`. + +```yaml +name: conformance_256k # informational workload name +model_tag: GLM-5.2-MXFP4 # written into requests[].model + models[]; retag per served model +id_prefix: caseA # session-id salt (see note below) +seed: 42 # fixed seed => byte-identical corpus every run +n_sessions: 200 # sessions to generate (more => tighter percentiles, bigger corpus) +block_size: 64 # KV block size (tokens/block) + +# Distribution targets, each a P50 / P90 / P99 triple: +isl_p: [74000, 155000, 235000] # input tokens per request +osl_p: [320, 3300, 17000] # output tokens per request +delay_p: [4, 31, 240] # inter-turn think delay (seconds) + +# Turns-per-session discrete distribution (parallel lists): +turns: + values: [1, 2, 3, 4, 6, 10, 20, 45, 103] + weights: [22, 24, 20, 12, 6, 5, 6, 3, 2] + +cache_hit: [0.88, 0.90] # per-turn prefix-reuse band [lo, hi] + +clamps: # post-draw sampling clamps [lo, hi] + isl: [1200, 245000] + osl: [8, 20000] + delay: [1, 600] + +# Optional verifier block: +verify: + turns_p: [3, 20, 103] # Turns P50/P90/P99 targets + cache_target: 89 # Cache-hit P50 % target + # band_overrides: # widen a per-axis tolerance band, e.g.: + # "Input ISL P99": [0.75, 1.25] +``` + +### What the verifier checks (13 axes) + +`verify_agentx_profile.py` measures the corpus and prints a per-axis table whose +verdict tokens are `PASS` / `off`, then a `N/N axes within band` summary. The 13 +axes are: **ISL** P50/P90/P99, **OSL** P50/P90/P99, **Turns** P50/P90/P99, +**Delay** P50/P90/P99, and **Cache hit P50 %**. + +Default tolerance bands (from `DEFAULT_BANDS`): + +| group | band (lo–hi multipliers) | +| --- | --- | +| isl | 0.80 – 1.20 | +| osl | 0.70 – 1.40 | +| turns | 0.60 – 1.60 | +| delay | 0.50 – 2.00 | +| cache | 0.97 – 1.03 | + +Override a single axis with `verify.band_overrides` keyed by the exact axis label +(e.g. Case-B in [conformance_512k.yaml](conformance_512k.yaml) widens +`"Input ISL P99"` to `[0.75, 1.25]`). Targets for turns/cache come from +`verify.turns_p` / `verify.cache_target` when present; otherwise turns targets are +derived from the `turns` distribution and the cache target from `mean(cache_hit) * 100`. + +## Field constraints + +- **Percentiles monotonic:** each `*_p` triple should satisfy `P50 <= P90 <= P99` + (the lognormal fit in `lognorm_from_p()` assumes an increasing triple). +- **Equal-length turns arrays:** `turns.values` and `turns.weights` must be the + same length (they are zipped in `rng.choices(...)` and the weighted-percentile + derivation). +- **`cache_hit` is a `[lo, hi]` band** with `0 <= lo <= hi <= 1` (used as + `rng.uniform(cache_lo, cache_hi)` per turn). +- **`clamps` are `[lo, hi]` pairs** for `isl` / `osl` / `delay`; each draw is + clamped into `[lo, hi]`. + +## Create a profile from scratch + +Use the **flag-based** CLIs (not any positional form). Round-trip needs no GPU: + +```bash +# 1. Start from the annotated template. +cp profiles/custom.example.yaml profiles/my_case.yaml +# edit name/id_prefix/targets to taste + +# 2. Resolve YAML -> JSON (what gen/verify consume). +python3 agentx_config.py --profile profiles/my_case.yaml --emit-json > /tmp/my.json + +# 3. Generate a corpus (overrides optional). +python3 gen_agentx_profile.py --profile /tmp/my.json --out-dir /tmp/my_corpus \ + [--n-sessions N --seed S --model-tag TAG --id-prefix P --block-size B] + +# 4. Verify until it passes. +python3 verify_agentx_profile.py --profile /tmp/my.json --corpus /tmp/my_corpus +# -> ends with "13/13 axes within band" (exit 0) when all axes PASS +``` + +Then reference it from `agentic.yaml`: + +```yaml +workloads: + - { name: my_case, preset: my_case } +``` + +If the verifier reports an `off` axis, nudge the offending `*_p` target (or widen +that axis via `verify.band_overrides`) and re-run steps 3–4. + +## Preset inheritance + +`_merge_preset()` (in `agentx_config.py`) merges an entry over its `preset:` +chain, **entry keys win** over inherited ones. A preset may bundle distribution +params (`source: profile`), an hf `loader` + Tier 1/Tier 2 knobs +(`source: hf`, see [inferencex_256k.yaml](inferencex_256k.yaml)), and/or run +knobs (`concurrency` / `duration`). A **circular** preset chain raises +`ValueError: circular preset: `. + +```yaml +# preset carries the sweep; the entry overrides just concurrency +- name: inferencex_small + preset: inferencex_256k + concurrency: [2] # wins over the preset's [2, 4, 8] +``` + +## The verify pre-gate + +At run time `materialize_corpus()` (in `../agentic_lib.sh`) generates the corpus +then runs `verify_agentx_profile.py`, and **aborts the run** unless the corpus is +`N/N axes within band` against the profile's own `verify:` targets. Corpora are +cached at `SUITE_CORPUS_DIR/` (default `/tmp/agentx_corpora`); editing a +profile does not invalidate the cache, so regenerate with `SUITE_CORPUS_FORCE=1`. + +The verifier prints a `corpus=... profile=... sessions=... requests=...` header, +then a per-axis table (columns: **measured** before **target**, then verdict), +then the `N/N axes within band` summary. Below is `conformance_256k` (real +targets from [conformance_256k.yaml](conformance_256k.yaml); the **measured** +values are illustrative). A full run prints all 13 axes — only 4 are shown here: + +``` +corpus=/tmp/agentx_corpora/conformance_256k profile=conformance_256k sessions=200 requests=1180 + +axis measured target verdict +------------------------------------------------------ +Input ISL P50 73,842 74,000 PASS +Input ISL P90 154,110 155,000 PASS +Input ISL P99 233,900 235,000 PASS +Cache hit P50 % 89 89 PASS +------------------------------------------------------ +13/13 axes within band +``` + +## Placement rule + +Drop the file at `profiles/.yaml` so that `preset: ` (and the +config-less `AGENTIC_WORKLOAD=` shorthand) resolves it. + +## Common mistakes + +- **Circular preset chain** — `A` presets `B` presets `A` -> `ValueError: + circular preset`. +- **Mismatched `turns` lengths** — `values` and `weights` must be equal length. +- **Non-monotonic percentiles** — a `*_p` triple that isn't increasing skews the + lognormal fit and fails verification. +- **Editing a profile without `SUITE_CORPUS_FORCE=1`** — the stale cached corpus + is reused and your edits appear to have no effect. + +## Note on `id_prefix` + +The conformance presets pin `id_prefix: caseA` (the literal salt #173 used for +**both** Case-A and Case-B) so regeneration is byte-identical to the committed +ROCm/MAD #173 corpora. For a new workload, use a distinct `id_prefix` — it is +just a session-id salt that changes corpus identity. Start from +[custom.example.yaml](custom.example.yaml). diff --git a/scripts/common/agentx/profiles/conformance_256k.yaml b/scripts/common/agentx/profiles/conformance_256k.yaml new file mode 100644 index 00000000..49c3c173 --- /dev/null +++ b/scripts/common/agentx/profiles/conformance_256k.yaml @@ -0,0 +1,34 @@ +# Case-A conformance profile (ROCm/MAD #173, generalized). +# ExplainX Case-A targets. Regenerating with these params via +# gen_agentx_profile.py reproduces #173's committed corpus byte-for-byte and +# verifies 13/13 through verify_agentx_profile.py. +name: conformance_256k +model_tag: GLM-5.2-MXFP4 # written into requests[].model + models[]; retag per served model +id_prefix: caseA # session-id salt; keep "caseA" to byte-match #173 +seed: 42 +n_sessions: 200 +block_size: 64 + +# Distribution targets: P50 / P90 / P99 +isl_p: [74000, 155000, 235000] # input tokens per request +osl_p: [320, 3300, 17000] # output tokens per request +delay_p: [4, 31, 240] # inter-turn think delay (seconds) + +# Turns-per-session discrete distribution (long-tail agentic). +turns: + values: [1, 2, 3, 4, 6, 10, 20, 45, 103] + weights: [22, 24, 20, 12, 6, 5, 6, 3, 2] + +# Per-turn prefix-reuse band -> ~88-90% prefix cache hit. +cache_hit: [0.88, 0.90] + +# Sampling clamps (lo, hi). +clamps: + isl: [1200, 245000] + osl: [8, 20000] + delay: [1, 600] + +# Verifier targets + tolerance bands (defaults match #173). +verify: + turns_p: [3, 20, 103] # Turns P50/P90/P99 + cache_target: 89 # Cache-hit P50 % diff --git a/scripts/common/agentx/profiles/conformance_512k.yaml b/scripts/common/agentx/profiles/conformance_512k.yaml new file mode 100644 index 00000000..9be4d543 --- /dev/null +++ b/scripts/common/agentx/profiles/conformance_512k.yaml @@ -0,0 +1,37 @@ +# Case-B conformance profile (ROCm/MAD #173, generalized). +# Longer-context / heavier-tail variant of Case-A. Regenerating with these +# params via gen_agentx_profile.py reproduces #173's committed corpus +# byte-for-byte and verifies 13/13 through verify_agentx_profile.py. +name: conformance_512k +model_tag: GLM-5.2-MXFP4 # written into requests[].model + models[]; retag per served model +id_prefix: caseA # #173 used the literal "caseA" salt for Case-B too; keep to byte-match +seed: 42 +n_sessions: 300 +block_size: 64 + +# Distribution targets: P50 / P90 / P99 +isl_p: [62000, 220000, 500000] # input tokens per request (tail to 500k) +osl_p: [180, 1400, 7000] # output tokens per request +delay_p: [3.6, 23, 240] # inter-turn think delay (seconds) + +# Turns-per-session discrete distribution (heavier tail than Case-A). +turns: + values: [2, 3, 5, 8, 20, 50, 82, 110, 144] + weights: [20, 20, 20, 10, 8, 8, 7, 4, 3] + +# Per-turn prefix-reuse band -> ~88-90% prefix cache hit. +cache_hit: [0.88, 0.90] + +# Sampling clamps (lo, hi). ISL tail runs to 520k for this case. +clamps: + isl: [1200, 520000] + osl: [8, 20000] + delay: [1, 600] + +# Verifier targets + tolerance bands. Case-B widens the ISL P99 band to +# 0.75-1.25 (the 500k tail is noisier), matching #173's verify_caseB.py. +verify: + turns_p: [5, 82, 144] # Turns P50/P90/P99 + cache_target: 89 # Cache-hit P50 % + band_overrides: + "Input ISL P99": [0.75, 1.25] diff --git a/scripts/common/agentx/profiles/custom.example.yaml b/scripts/common/agentx/profiles/custom.example.yaml new file mode 100644 index 00000000..495befb8 --- /dev/null +++ b/scripts/common/agentx/profiles/custom.example.yaml @@ -0,0 +1,50 @@ +# Annotated template for a CUSTOM AgentX workload profile. +# +# Copy this file, edit the targets to your workload, then reference it from +# agentic.yaml either inline or as a profile file. A profile is just ~6 +# distribution targets + a session count + a seed; the generator turns it into +# a reproducible weka_trace corpus and the verifier gates it against these same +# targets. There is no per-case code -- add a case by adding a profile. +# +# Round-trip (no GPU needed): +# python3 agentx_config.py --profile profiles/custom.example.yaml --emit-json > /tmp/c.json +# python3 gen_agentx_profile.py --profile /tmp/c.json --out-dir /tmp/custom_corpus +# python3 verify_agentx_profile.py --profile /tmp/c.json --corpus /tmp/custom_corpus +# # -> "N/N axes within band" + +name: my_case # unique workload name (becomes the result subdir) +model_tag: GLM-5.2-MXFP4 # tag written into the trace; retag to your served model if you like +id_prefix: my_case # session-id salt (any string; changes the corpus identity) +seed: 42 # fixed seed => byte-identical corpus every run +n_sessions: 150 # more sessions => tighter percentiles (and a bigger corpus) +block_size: 64 # KV block size (tokens/block); keep 64 unless your engine differs + +# Distribution targets, each a P50 / P90 / P99 triple. +isl_p: [48000, 120000, 200000] # input tokens per request +osl_p: [256, 2000, 9000] # output tokens per request +delay_p: [3, 20, 180] # inter-turn think delay (seconds) + +# Turns-per-session discrete distribution: parallel values/weights lists. +# Keep enough multi-turn mass that the prefix-cache-hit median lands ~88-90%. +turns: + values: [2, 3, 4, 6, 10, 20, 45, 103] + weights: [20, 24, 20, 12, 8, 6, 7, 3] + +# Per-turn prefix-reuse band -> target prefix cache-hit fraction. +cache_hit: [0.88, 0.90] + +# Sampling clamps (lo, hi) applied after the lognormal draw. +clamps: + isl: [1200, 205000] + osl: [8, 20000] + delay: [1, 600] + +# Optional verifier block. If omitted, turns targets are derived from the turns +# distribution and the cache target from the midpoint of cache_hit, so a custom +# profile round-trips to N/N without you spelling the targets out. Provide it to +# pin exact targets or widen a band on a noisy tail axis. +# verify: +# turns_p: [3, 30, 60] +# cache_target: 89 +# band_overrides: +# "Input ISL P99": [0.75, 1.25] diff --git a/scripts/common/agentx/profiles/inferencex_256k.yaml b/scripts/common/agentx/profiles/inferencex_256k.yaml new file mode 100644 index 00000000..a2c4d736 --- /dev/null +++ b/scripts/common/agentx/profiles/inferencex_256k.yaml @@ -0,0 +1,20 @@ +# Reusable HF (weka) workload preset: the canonical InferenceX AgentX-MVP +# corpus, 256k-capped. Reference it from agentic.yaml with `preset: inferencex_256k`. +# A source=hf preset bundles the loader + Tier 1 replay knobs (+ optional Tier 2 +# local filter) and may carry its own run knobs (concurrency sweep / duration). +# Entry fields override these; env (AGENTIC_CONC/DURATION/MAX_MODEL_LEN) still wins. +source: hf +loader: semianalysis_cc_traces_weka_062126_256k # -> ISL tail 262144 for context gating + +# Bundled run knobs (preset-level; an entry's concurrency/duration override these). +concurrency: [2, 4, 8] +duration: 900 + +# Tier 1 replay knobs (per-workload; defaults match today's hardcoded behavior). +num_dataset_entries: 393 # how many trace sessions to pull +trajectory: { min: 0.25, max: 0.75 } # start-window ratio for captured traces + +# Tier 2 (optional): trim the downloaded corpus locally to fit a smaller model. +# Uncomment to drop sessions with any turn over max_isl, truncate to max_turns, +# and/or randomly sample down (seed=42). Omit for the full --public-dataset path. +# filter: { max_isl: 200000, max_turns: 40, sample: 100 } diff --git a/scripts/common/agentx/profiles/small.yaml b/scripts/common/agentx/profiles/small.yaml new file mode 100644 index 00000000..0b6a6013 --- /dev/null +++ b/scripts/common/agentx/profiles/small.yaml @@ -0,0 +1,15 @@ +name: small +model_tag: small-model +id_prefix: small +seed: 42 +n_sessions: 200 +block_size: 64 +isl_p: [2000, 4000, 6000] +osl_p: [128, 256, 512] +delay_p: [1, 3, 8] +turns: { values: [1,2,3,4,6], weights: [30,30,20,12,8] } +cache_hit: [0.88, 0.90] +clamps: { isl: [512,7000], osl: [8,1024], delay: [1,60] } +verify: + band_overrides: + "Cache hit P50 %": [0.94, 1.06] diff --git a/scripts/common/agentx/templates/README.md b/scripts/common/agentx/templates/README.md new file mode 100644 index 00000000..bd1cf828 --- /dev/null +++ b/scripts/common/agentx/templates/README.md @@ -0,0 +1,76 @@ +# AgentX backend integration guide + +The AgentX core ([../README.md](../README.md)) speaks only the OpenAI API, so +all backend hook logic now lives in ONE shared, `--backend`-parameterized script +([../../benchmark_agentic.sh](../../benchmark_agentic.sh)). Adding a backend is +two small edits — a `case` arm plus a thin shim — instead of copying and drifting +a full template. + +## Adding a backend + +1. **Add a `case "$backend"` arm** in + [`scripts/common/benchmark_agentic.sh`](../../benchmark_agentic.sh) that sets + the backend's `AGENTIC_PORT` default and its ctx-window resolver endpoint list + (`ctx_endpoints=("path|kind" ...)`, where `kind` is `models` to read + `data[0].max_model_len` from an OpenAI `/v1/models` ModelCard, or anything + else to read `max_model_len`/`context_length`/`server_args.*` — e.g. sglang's + `/get_server_info`). Single-endpoint backends list one entry; add more only if + your framework serves the window elsewhere. +2. **Add the `$backend` value** to the validation and dispatch `case`s (mirror + the existing `sglang|vllm` arms). +3. **Add a ~7-line shim** at `scripts/_disagg/benchmark_agentic.sh` that + locates the shared script (its `../common/benchmark_agentic.sh` sibling, plus + any launcher repo-dir env var if the in-container mount hides `../common`, plus + an `AGENTIC_LIB` override) and `exec bash "$_cand" --backend "$@"`. + Copy an existing shim + ([sglang](../../../sglang_disagg/benchmark_agentic.sh), + [vllm](../../../vllm_dissag/benchmark_agentic.sh)). +4. **Wire it into the launcher:** it is selected via + `export BENCHMARK_SCRIPT_FILE=benchmark_agentic.sh` (typically exposed to users + as `BENCHMARK_SCRIPT=agentic`). +5. Preview with `DRY_RUN=1` (no server needed), then run against a live endpoint. + +## Backend contract + +The core talks to your endpoint over two OpenAI-compatible routes: + +- `POST /v1/chat/completions` — streaming chat, used for the trace replay. +- `GET /v1/models` — model discovery and readiness gating + (`wait_for_router_ready` and `resolve_served_model_name` both poll this; + `/v1/models` `data[0].max_model_len` also drives context auto-detection). + +If your framework does not serve `/v1/models` (or returns 503 while workers +register), front it with a tiny shim that answers `/v1/models` once the upstream +is healthy and proxies everything else. See the working example +[../../../vllm_dissag/agentic_models_shim.py](../../../vllm_dissag/agentic_models_shim.py) +and point `AGENTIC_PORT` at the shim. + +## Disaggregated context-window resolution + +Disaggregated P/D front-ends (a router, proxy, or `/v1/models` shim) often do +**not** advertise `max_model_len`, so the library's front-end auto-detect returns +`0`. The shared script ships a `# === agentx:BEGIN resolve served context window +(disagg) ===` block that instead probes the prefill **WORKER** — the first +`host:port` in `AGENTIC_SERVER_METRICS` — over each entry in the backend's +`ctx_endpoints`. On the shipped launchers `AGENTIC_SERVER_METRICS` is +auto-derived in-container, so users normally never set it. Single-node / monolith +backends need no worker probe — an empty `AGENTIC_SERVER_METRICS` falls through +to the library's front-end `/v1/models` auto-detect. + +`AGENTIC_RESOLVE_ONLY=1` resolves the served `max_model_len`, prints it, and +exits without running — a diagnostic for checking the probe. It is **not** +forwarded through the launchers, so use it in a direct/local run. + +**Intentional probe divergence:** sglang probes `/v1/models` **and** +`/get_server_info` (older builds only expose it there), while vLLM probes only +`/v1/models`. This lives in each backend's `ctx_endpoints` — keep per-backend +lists rather than forcing one shared list. + +## Required vs optional env + +Set exactly **one** entry-point variable — `AGENTIC_CONFIG` (a config path) or +`AGENTIC_WORKLOAD` (a single-workload name) — and have a live endpoint on +`AGENTIC_PORT`. Everything else is optional and auto-defaults. See the core +[env reference](../README.md#environment-variable-reference) and +[Minimal required](../README.md#minimal-required), and +[../README.md#examples](../README.md#examples) for copy-paste configs. diff --git a/scripts/common/agentx/tests/_stub_server.py b/scripts/common/agentx/tests/_stub_server.py new file mode 100644 index 00000000..e94b9686 --- /dev/null +++ b/scripts/common/agentx/tests/_stub_server.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python3 +"""Minimal stdlib HTTP stub of an sglang API worker for offline ctx-resolve tests. + +Usage: python3 _stub_server.py [empty] + +When the 2nd arg is "empty" the window fields are omitted so the resolver's +fail-fast path can be exercised. +""" +import json +import sys +from http.server import BaseHTTPRequestHandler, HTTPServer + +EMPTY = len(sys.argv) > 2 and sys.argv[2] == "empty" + + +class Handler(BaseHTTPRequestHandler): + def _send(self, obj): + body = json.dumps(obj).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if self.path == "/v1/models": + if EMPTY: + self._send({"data": [{"id": "stub"}]}) + else: + self._send({"data": [{"id": "stub", "max_model_len": 131072}]}) + elif self.path == "/get_server_info": + if EMPTY: + self._send({"server_args": {}}) + else: + self._send({"server_args": {"context_length": 131072}}) + else: + self.send_response(404) + self.end_headers() + + def log_message(self, *args): + pass + + +if __name__ == "__main__": + port = int(sys.argv[1]) + HTTPServer(("127.0.0.1", port), Handler).serve_forever() diff --git a/scripts/common/agentx/tests/harvest.sh b/scripts/common/agentx/tests/harvest.sh new file mode 100644 index 00000000..99f09e21 --- /dev/null +++ b/scripts/common/agentx/tests/harvest.sh @@ -0,0 +1,107 @@ +#!/bin/bash +# AgentX campaign harvester (read-only). For each job result dir, parse +# suite_summary.json and print one matrix row per workload: +# cell workload error_rate gpu_cache_hit_rate theoretical_cache_hit_rate verdict +# +# Verdict (same thresholds as validate_agentic_result.sh / the READMEs): +# FAIL suite_summary.json missing/unreadable, status INVALID, or +# error_rate missing or > AGENTIC_MAX_ERROR_RATE (default 0.10) +# WARN gpu_cache_hit_rate < AGENTIC_MIN_CACHE_HIT (default 0.30) +# PASS otherwise +# +# Inputs: JOBIDs (or full result dirs) as args. A JOBID resolves to +# ${RESULT_ROOT:-/run_logs}/. If a JOB_MAP tsv (JOBIDcell) is given, +# the cell label is looked up from it; otherwise the jobid/dir is the cell label. +# +# Usage: +# bash harvest.sh 12345 12346 # /run_logs/12345 /run_logs/12346 +# RESULT_ROOT=/run_logs bash harvest.sh 12345 +# JOB_MAP=agentx_jobmap.tsv bash harvest.sh $(cut -f1 agentx_jobmap.tsv) +# bash harvest.sh /path/to/run_logs/12345 # explicit dir +set -uo pipefail + +RESULT_ROOT="${RESULT_ROOT:-/run_logs}" +JOB_MAP="${JOB_MAP:-}" +MAX_ERROR_RATE="${AGENTIC_MAX_ERROR_RATE:-0.10}" +MIN_CACHE_HIT="${AGENTIC_MIN_CACHE_HIT:-0.30}" + +if [ "$#" -eq 0 ]; then + echo "usage: bash harvest.sh [JOBID|result_dir ...]" >&2 + exit 2 +fi + +PY="python3" +command -v "$PY" >/dev/null 2>&1 || { echo "python3 not found" >&2; exit 2; } + +_cell_label() { # $1=jobid/dirname -> cell label via JOB_MAP or identity + local key="$1" + if [ -n "$JOB_MAP" ] && [ -f "$JOB_MAP" ]; then + local hit + hit="$(awk -F'\t' -v k="$key" '$1==k{print $2; exit}' "$JOB_MAP")" + [ -n "$hit" ] && { echo "$hit"; return; } + fi + echo "$key" +} + +printf '%-28s %-18s %10s %12s %12s %s\n' \ + "cell" "workload" "err_rate" "cache_hit" "theo_hit" "verdict" +printf '%s\n' "----------------------------------------------------------------------------------------------" + +overall_fail=0 +for arg in "$@"; do + if [ -d "$arg" ]; then + dir="$arg"; key="$(basename "$arg")" + else + key="$arg"; dir="${RESULT_ROOT}/${arg}" + fi + cell="$(_cell_label "$key")" + summary="${dir}/suite_summary.json" + + if [ ! -f "$summary" ]; then + printf '%-28s %-18s %10s %12s %12s %s\n' \ + "$cell" "-" "-" "-" "-" "FAIL(no summary)" + overall_fail=1 + continue + fi + + # Parse + verdict per workload in python3; print TSV rows, exit 1 if any FAIL. + rows="$(SUMMARY="$summary" CELL="$cell" MAXERR="$MAX_ERROR_RATE" MINCACHE="$MIN_CACHE_HIT" \ + "$PY" - <<'PY' +import json, os, sys +summary = os.environ["SUMMARY"]; cell = os.environ["CELL"] +maxerr = float(os.environ["MAXERR"]); mincache = float(os.environ["MINCACHE"]) +try: + d = json.load(open(summary)) +except Exception as e: + print(f"{cell}\t-\t-\t-\t-\tFAIL(bad json)") + sys.exit(1) +wls = d.get("workloads") or [] +if not wls: + print(f"{cell}\t-\t-\t-\t-\tFAIL(empty)") + sys.exit(1) +any_fail = 0 +def pct(v): return "-" if v is None else f"{v*100:.1f}%" +for w in wls: + name = w.get("workload", "?") + err = w.get("error_rate"); ch = w.get("gpu_cache_hit_rate"); th = w.get("theoretical_cache_hit_rate") + status = w.get("status") + if status == "INVALID" or err is None or err > maxerr: + verdict = "FAIL"; any_fail = 1 + elif ch is None or ch < mincache: + verdict = "WARN" + else: + verdict = "PASS" + print(f"{cell}\t{name}\t{pct(err)}\t{pct(ch)}\t{pct(th)}\t{verdict}") +sys.exit(any_fail) +PY + )" || overall_fail=1 + + while IFS=$'\t' read -r c wl er chit thit verdict; do + [ -n "$c" ] || continue + printf '%-28s %-18s %10s %12s %12s %s\n' "$c" "$wl" "$er" "$chit" "$thit" "$verdict" + done <<< "$rows" +done + +printf '%s\n' "----------------------------------------------------------------------------------------------" +[ "$overall_fail" -eq 0 ] && echo "harvest: all cells PASS/WARN" || echo "harvest: one or more cells FAIL/INVALID" +exit 0 diff --git a/scripts/common/agentx/tests/resolve_ctx_offline.sh b/scripts/common/agentx/tests/resolve_ctx_offline.sh new file mode 100755 index 00000000..a17f2aca --- /dev/null +++ b/scripts/common/agentx/tests/resolve_ctx_offline.sh @@ -0,0 +1,122 @@ +#!/bin/bash +# Backend-agnostic OFFLINE test for the disagg served-context-window resolver in +# the per-backend benchmark_agentic.sh hooks. Spins up a stdlib http stub in +# place of the prefill worker and, for EACH backend (sglang, vllm), exercises 4 +# paths: SUCCESS, FAIL-FAST, FALLTHROUGH, PROPAGATION. No GPU, no network, no +# dep install (RESOLVE_ONLY exits early). The shared stub is a superset that +# serves both /v1/models (max_model_len) and /get_server_info +# (server_args.context_length); each backend's resolver queries whichever it +# needs. Prints PASS/FAIL per assertion; exits non-zero if any fail. +set -uo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# tests -> agentx -> common -> scripts -> repo root +REPO_ROOT="$(cd "$HERE/../../../.." && pwd)" +STUB="$HERE/_stub_server.py" + +if [ ! -f "$REPO_ROOT/scripts/sglang_disagg/benchmark_agentic.sh" ]; then + echo "could not locate repo root (missing scripts/sglang_disagg/benchmark_agentic.sh under $REPO_ROOT)" >&2 + exit 1 +fi + +fails=0 +STUB_PID="" + +cleanup() { [ -n "$STUB_PID" ] && kill "$STUB_PID" 2>/dev/null; } +trap cleanup EXIT + +pick_port() { + python3 -c 'import socket +s=socket.socket(); s.bind(("127.0.0.1",0)); print(s.getsockname()[1]); s.close()' +} + +start_stub() { # $1=port $2=optional "empty" + STUB_PID="" + python3 "$STUB" "$1" ${2:-} & + STUB_PID=$! + for _ in $(seq 1 50); do + if curl -sf "http://127.0.0.1:$1/v1/models" >/dev/null 2>&1; then return 0; fi + sleep 0.1 + done + echo " stub failed to accept on port $1" >&2 + return 1 +} + +stop_stub() { [ -n "$STUB_PID" ] && kill "$STUB_PID" 2>/dev/null; wait "$STUB_PID" 2>/dev/null; STUB_PID=""; } + +report() { # $1=name $2=0/1 pass + if [ "$2" = "1" ]; then echo "PASS: $1"; else echo "FAIL: $1"; fails=$((fails+1)); fi +} + +run_backend() { # $1=backend name $2=backend dir under scripts/ + local bk="$1" + local dir="$2" + local HOOK="$REPO_ROOT/scripts/${dir}/benchmark_agentic.sh" + local before=$fails + + if [ ! -f "$HOOK" ]; then + report "[$bk] hook exists" 0 + return + fi + + # (a) SUCCESS: resolver exports MAX_MODEL_LEN from the worker. + local PORT out rc + PORT="$(pick_port)" + if start_stub "$PORT"; then + out="$(AGENTIC_RESOLVE_ONLY=1 AGENTIC_SERVER_METRICS="127.0.0.1:$PORT" MAX_MODEL_LEN= DRY_RUN=0 \ + bash "$HOOK" 2>&1)" + echo "$out" | grep -q "MAX_MODEL_LEN=131072" && report "[$bk] (a) SUCCESS resolves 131072" 1 || { echo "$out"; report "[$bk] (a) SUCCESS resolves 131072" 0; } + else + report "[$bk] (a) SUCCESS resolves 131072" 0 + fi + stop_stub + + # (b) FAIL-FAST: worker set but no window field anywhere -> non-zero exit. + PORT="$(pick_port)" + if start_stub "$PORT" empty; then + AGENTIC_RESOLVE_ONLY=1 AGENTIC_SERVER_METRICS="127.0.0.1:$PORT" MAX_MODEL_LEN= DRY_RUN=0 \ + bash "$HOOK" >/dev/null 2>&1 + rc=$? + [ "$rc" -ne 0 ] && report "[$bk] (b) FAIL-FAST exits non-zero" 1 || report "[$bk] (b) FAIL-FAST exits non-zero" 0 + else + report "[$bk] (b) FAIL-FAST exits non-zero" 0 + fi + stop_stub + + # (c) FALLTHROUGH: no worker (non-disagg) -> exit 0, empty MAX_MODEL_LEN. + out="$(AGENTIC_RESOLVE_ONLY=1 AGENTIC_SERVER_METRICS= bash "$HOOK" 2>&1)" + rc=$? + if [ "$rc" -eq 0 ] && echo "$out" | grep -q "MAX_MODEL_LEN="; then + report "[$bk] (c) FALLTHROUGH exit 0 + empty MAX_MODEL_LEN" 1 + else + echo "rc=$rc"; echo "$out"; report "[$bk] (c) FALLTHROUGH exit 0 + empty MAX_MODEL_LEN" 0 + fi + + # (d) PROPAGATION: pinned MAX_MODEL_LEN flows through the DRY_RUN suite plan. + out="$(MAX_MODEL_LEN=131072 DRY_RUN=1 AGENTIC_WORKLOAD=small bash "$HOOK" 2>&1)" + if echo "$out" | grep -E "max_model_len" | grep -q "131072"; then + report "[$bk] (d) PROPAGATION max_model_len=131072 in suite plan" 1 + else + echo "$out"; report "[$bk] (d) PROPAGATION max_model_len=131072 in suite plan" 0 + fi + + local bk_fails=$((fails - before)) + echo "---- [$bk] summary: $([ "$bk_fails" -eq 0 ] && echo "ALL PASS" || echo "$bk_fails FAILED")" +} + +# backend name -> its directory under scripts/ (note the vllm dir spelling). +for pair in "sglang:sglang_disagg" "vllm:vllm_dissag"; do + bk="${pair%%:*}"; dir="${pair##*:}" + echo "=== backend: $bk ===" + run_backend "$bk" "$dir" + echo "" +done + +echo "======================================================" +if [ "$fails" -eq 0 ]; then + echo "resolve_ctx_offline: ALL PASS" + exit 0 +else + echo "resolve_ctx_offline: $fails assertion(s) FAILED" + exit 1 +fi diff --git a/scripts/common/agentx/tests/run_offline.sh b/scripts/common/agentx/tests/run_offline.sh new file mode 100644 index 00000000..2a2f0c8f --- /dev/null +++ b/scripts/common/agentx/tests/run_offline.sh @@ -0,0 +1,329 @@ +#!/bin/bash +# Backend-agnostic OFFLINE gate for the AgentX agentic suite. No server, no GPU, +# no cluster, no network. Runs the REAL code paths where possible so a future +# edit to the config loader / generator / verifier / suite driver is caught here. +# +# Checks: +# 1. bash -n every backend hook + connector + the suite driver + agentic_lib. +# 2. Every workload in agentic.example.yaml resolves (--emit-workload-shell). +# 3. Deterministic gen+verify smoke on the tiny `small` profile (13/13 axes). +# 4. DRY_RUN suite driver prints a plan line for each workload and exits 0. +# +# Usage: bash scripts/common/agentx/tests/run_offline.sh (exit 0 = all pass) +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +AGENTX_DIR="$(cd "$HERE/.." && pwd)" +COMMON_DIR="$(cd "$AGENTX_DIR/.." && pwd)" +REPO_ROOT="$(cd "$COMMON_DIR/../.." && pwd)" +# Use the test config (includes my_corpus); agentic.example.yaml is for user reference. +CONFIG="$HERE/test_offline.yaml" +SUITE_DRIVER="$COMMON_DIR/benchmark_agentic_suite.sh" +PY="python3" + +pass=0; fail=0 +_pass() { printf " PASS %s\n" "$1"; pass=$((pass+1)); } +_fail() { printf " FAIL %s\n" "$1"; fail=$((fail+1)); } + +# Isolated tmp workspace (corpus + profile JSON), cleaned up on exit. +TMP="$(mktemp -d "${TMPDIR:-/tmp}/agentx_offline.XXXXXX")" +cleanup() { rm -rf "$TMP"; } +trap cleanup EXIT + +# --------------------------------------------------------------------------- +echo "=== 1. bash -n syntax check (hooks + connectors + driver + lib) ===" +SYNTAX_TARGETS=( + "$COMMON_DIR/benchmark_agentic.sh" + "$REPO_ROOT/scripts/sglang_disagg/benchmark_agentic.sh" + "$REPO_ROOT/scripts/vllm_dissag/benchmark_agentic.sh" + "$REPO_ROOT/scripts/vllm_dissag/connectors/rixl.sh" + "$REPO_ROOT/scripts/vllm_dissag/connectors/moriio.sh" + "$SUITE_DRIVER" + "$COMMON_DIR/agentic_lib.sh" + "$HERE/resolve_ctx_offline.sh" +) +for f in "${SYNTAX_TARGETS[@]}"; do + rel="${f#"$REPO_ROOT/"}" + if [ ! -f "$f" ]; then + _fail "missing: $rel" + elif bash -n "$f" 2>/tmp/agentx_offline_syn.$$; then + _pass "bash -n $rel" + else + _fail "bash -n $rel" + sed 's/^/ /' /tmp/agentx_offline_syn.$$ || true + fi + rm -f /tmp/agentx_offline_syn.$$ +done +# The shared ctx-resolve stub is Python, so py_compile it (bash -n won't do). +STUB_PY="$HERE/_stub_server.py" +stub_rel="${STUB_PY#"$REPO_ROOT/"}" +if [ ! -f "$STUB_PY" ]; then + _fail "missing: $stub_rel" +elif "$PY" -m py_compile "$STUB_PY" 2>/tmp/agentx_offline_pyc.$$; then + _pass "py_compile $stub_rel" +else + _fail "py_compile $stub_rel" + sed 's/^/ /' /tmp/agentx_offline_pyc.$$ || true +fi +rm -f /tmp/agentx_offline_pyc.$$ + +# --------------------------------------------------------------------------- +echo "" +echo "=== 2. per-workload config resolution (agentic.example.yaml) ===" +# Enumerate workloads from the resolved config rather than hardcoding. +mapfile -t WORKLOADS < <("$PY" "$AGENTX_DIR/agentx_config.py" --config "$CONFIG" --dump-json \ + | "$PY" -c 'import sys,json; [print(w["name"]) for w in json.load(sys.stdin)["workloads"]]') +if [ "${#WORKLOADS[@]}" -eq 0 ]; then + _fail "enumerate workloads from --dump-json" +else + _pass "enumerated ${#WORKLOADS[@]} workloads: ${WORKLOADS[*]}" +fi +for name in "${WORKLOADS[@]}"; do + if "$PY" "$AGENTX_DIR/agentx_config.py" --config "$CONFIG" --workload "$name" \ + --profile-out "$TMP/${name}.profile.json" --emit-workload-shell >/dev/null 2>"$TMP/wl.err"; then + _pass "resolve workload '$name'" + else + _fail "resolve workload '$name'" + sed 's/^/ /' "$TMP/wl.err" || true + fi +done + +# --------------------------------------------------------------------------- +echo "" +echo "=== 3. deterministic gen+verify smoke (profiles/small.yaml, seed=42) ===" +SMALL_JSON="$TMP/small.profile.json" +SMALL_CORPUS="$TMP/small_corpus" +if "$PY" "$AGENTX_DIR/agentx_config.py" --profile "$AGENTX_DIR/profiles/small.yaml" \ + --emit-json > "$SMALL_JSON" 2>"$TMP/small.err"; then + _pass "emit small profile JSON" +else + _fail "emit small profile JSON"; sed 's/^/ /' "$TMP/small.err" || true +fi +if "$PY" "$AGENTX_DIR/gen_agentx_profile.py" --profile "$SMALL_JSON" --seed 42 \ + --out-dir "$SMALL_CORPUS" >/dev/null 2>"$TMP/gen.err"; then + _pass "generate small corpus" +else + _fail "generate small corpus"; sed 's/^/ /' "$TMP/gen.err" || true +fi +# verify exits 0 iff all axes pass; its final line is "N/N axes within band". +if verify_out="$("$PY" "$AGENTX_DIR/verify_agentx_profile.py" --profile "$SMALL_JSON" \ + --corpus "$SMALL_CORPUS" 2>&1)"; then + band_line="$(echo "$verify_out" | tail -n1)" + _pass "verify small corpus (${band_line})" +else + _fail "verify small corpus" + echo "$verify_out" | sed 's/^/ /' +fi + +# --------------------------------------------------------------------------- +echo "" +echo "=== 4. DRY_RUN suite driver (no server) ===" +# DRY_RUN must not need the aiperf venv or network. If it does, degrade to WARN. +dry_out=""; dry_rc=0 +dry_out="$(DRY_RUN=1 AGENTIC_CONFIG="$CONFIG" bash "$SUITE_DRIVER" 2>&1)" || dry_rc=$? +if [ "$dry_rc" -ne 0 ]; then + if echo "$dry_out" | grep -qiE 'network|download|uv |venv|pip|install'; then + printf " WARN DRY_RUN suite driver exited %s (looks env/network related)\n" "$dry_rc" + echo "$dry_out" | tail -n 20 | sed 's/^/ /' + else + _fail "DRY_RUN suite driver exited $dry_rc" + echo "$dry_out" | tail -n 20 | sed 's/^/ /' + fi +else + _pass "DRY_RUN suite driver exit 0" + for name in "${WORKLOADS[@]}"; do + if echo "$dry_out" | grep -q "workload='${name}'"; then + _pass "plan line for '$name'" + else + _fail "no plan line for '$name'" + fi + done +fi + +# --------------------------------------------------------------------------- +echo "" +echo "=== 5. source=corpus resolution + configurable scenario ===" +# corpus workload resolves to an --input-file replay (no download/generate) and +# skips verification when it carries no profile/preset. +corpus_wl="$("$PY" "$AGENTX_DIR/agentx_config.py" --config "$CONFIG" --workload my_corpus \ + --emit-workload-shell 2>/dev/null)" || true +if echo "$corpus_wl" | grep -q "WL_SOURCE='corpus'" \ + && echo "$corpus_wl" | grep -Eq "WL_INPUT_DIR='.+'"; then + _pass "source=corpus resolves input_dir" +else + _fail "source=corpus resolves input_dir" +fi +if echo "$corpus_wl" | grep -q "WL_PROFILE_FILE=''"; then + _pass "source=corpus verification optional (no profile -> no pre-gate)" +else + _fail "source=corpus verification optional (no profile -> no pre-gate)" +fi +# scenario: defaults to inferencex-agentx-mvp, overridable via AGENTIC_SCENARIO. +ovr_sc="$(AGENTIC_SCENARIO=my-scenario "$PY" "$AGENTX_DIR/agentx_config.py" \ + --config "$CONFIG" --emit-config-shell | grep '^SUITE_SCENARIO=')" +if [ "$ovr_sc" = "SUITE_SCENARIO='my-scenario'" ]; then + _pass "AGENTIC_SCENARIO overrides run.scenario" +else + _fail "AGENTIC_SCENARIO overrides run.scenario (got $ovr_sc)" +fi +if echo "$dry_out" | grep -Eq "^ run.scenario +: inferencex-agentx-mvp"; then + _pass "DRY_RUN plan shows default scenario" +else + _fail "DRY_RUN plan shows default scenario" +fi +sc_out="$(DRY_RUN=1 AGENTIC_SCENARIO=my-scenario AGENTIC_CONFIG="$CONFIG" bash "$SUITE_DRIVER" 2>&1)" || true +if echo "$sc_out" | grep -q -- "--scenario my-scenario"; then + _pass "DRY_RUN replay command honors AGENTIC_SCENARIO" +else + _fail "DRY_RUN replay command honors AGENTIC_SCENARIO" +fi + +# --------------------------------------------------------------------------- +echo "" +echo "=== 6. offline ctx-window resolver (stub server) ===" +# Runs the shared backend-looped resolver test once (sglang + vllm x 4 paths). +# Guard the invocation so its non-zero exit doesn't abort us under `set -e`. +CTX_TEST="$HERE/resolve_ctx_offline.sh" +ctx_rc=0 +ctx_out="$(bash "$CTX_TEST" 2>&1)" || ctx_rc=$? +echo "$ctx_out" | sed 's/^/ /' +for bk in sglang vllm; do + if echo "$ctx_out" | grep -q "^---- \[$bk\] summary: ALL PASS"; then + _pass "ctx resolver [$bk] all paths" + else + _fail "ctx resolver [$bk] all paths" + fi +done +if [ "$ctx_rc" -ne 0 ]; then + echo " (resolve_ctx_offline.sh exited $ctx_rc)" +fi + +# --------------------------------------------------------------------------- +echo "" +# --------------------------------------------------------------------------- +echo "" +echo "=== 7. N1 parser + robustness guards ===" + +# 7.1 N1: workloads: at the SAME indent as its `-` items must resolve to a +# NON-EMPTY workload list via the fallback YAML loader (regression for the +# silently-dropped same-indent block sequence). +cat > "$TMP/n1_same_indent.yaml" <<'YAML' +serving: + model: auto + max_model_len: 524288 +run: + concurrency: [2] + duration: 900 + scenario: inferencex-agentx-mvp +workloads: +- name: samelevel + source: profile + preset: conformance_256k +YAML +n1_count="$(AGENTX_YAML_FALLBACK=1 "$PY" "$AGENTX_DIR/agentx_config.py" \ + --config "$TMP/n1_same_indent.yaml" --dump-json 2>/dev/null \ + | "$PY" -c 'import sys,json; print(len(json.load(sys.stdin).get("workloads") or []))')" || n1_count=0 +if [ "${n1_count:-0}" -ge 1 ]; then + _pass "N1 fallback parser resolves same-indent workloads ($n1_count)" +else + _fail "N1 fallback parser resolves same-indent workloads (got '${n1_count}')" +fi + +# 7.2 gen guard: block_size 0 -> non-zero exit + explicit message. +"$PY" -c 'import json,sys; d=json.load(open(sys.argv[1])); d["block_size"]=0; json.dump(d,open(sys.argv[2],"w"))' \ + "$SMALL_JSON" "$TMP/bs0.json" +bs0_out="$("$PY" "$AGENTX_DIR/gen_agentx_profile.py" --profile "$TMP/bs0.json" --seed 42 \ + --out-dir "$TMP/bs0_corpus" 2>&1)" && bs0_rc=0 || bs0_rc=$? +if [ "$bs0_rc" -ne 0 ] && echo "$bs0_out" | grep -q '\[gen\] block_size must be >= 1'; then + _pass "gen guard: block_size 0 rejected" +else + _fail "gen guard: block_size 0 rejected (rc=$bs0_rc)" + echo "$bs0_out" | sed 's/^/ /' +fi + +# 7.3 gen guard: empty turns.values -> non-zero exit + explicit message. +"$PY" -c 'import json,sys; d=json.load(open(sys.argv[1])); d["turns"]["values"]=[]; json.dump(d,open(sys.argv[2],"w"))' \ + "$SMALL_JSON" "$TMP/tv0.json" +tv0_out="$("$PY" "$AGENTX_DIR/gen_agentx_profile.py" --profile "$TMP/tv0.json" --seed 42 \ + --out-dir "$TMP/tv0_corpus" 2>&1)" && tv0_rc=0 || tv0_rc=$? +if [ "$tv0_rc" -ne 0 ] && echo "$tv0_out" | grep -q '\[gen\] turns must have non-empty'; then + _pass "gen guard: empty turns.values rejected" +else + _fail "gen guard: empty turns.values rejected (rc=$tv0_rc)" + echo "$tv0_out" | sed 's/^/ /' +fi + +# 7.4 verify guard: session JSON lacking 'requests' -> non-zero exit + message. +mkdir -p "$TMP/badcorpus" +echo '{}' > "$TMP/badcorpus/s1.json" +vf_out="$("$PY" "$AGENTX_DIR/verify_agentx_profile.py" --profile "$SMALL_JSON" \ + --corpus "$TMP/badcorpus" 2>&1)" && vf_rc=0 || vf_rc=$? +if [ "$vf_rc" -ne 0 ] && echo "$vf_out" | grep -q "missing 'requests'"; then + _pass "verify guard: missing 'requests' rejected" +else + _fail "verify guard: missing 'requests' rejected (rc=$vf_rc)" + echo "$vf_out" | sed 's/^/ /' +fi + +# 7.5 preset guard: unknown preset name -> reported, not silently loaded. +cat > "$TMP/preset_missing.yaml" <<'YAML' +serving: + model: auto + max_model_len: 524288 +run: + concurrency: [2] + duration: 900 +workloads: + - name: bad + preset: __missing__ +YAML +pm_out="$("$PY" "$AGENTX_DIR/agentx_config.py" --config "$TMP/preset_missing.yaml" --dump-json 2>&1)" || true +if echo "$pm_out" | grep -q "preset not found"; then + _pass "preset guard: missing preset reported" +else + _fail "preset guard: missing preset reported" + echo "$pm_out" | sed 's/^/ /' +fi + +# --------------------------------------------------------------------------- +# --------------------------------------------------------------------------- +echo "" +echo "=== 8. shell driver static guards (fixes 2/3/4) ===" +SLURM_LAUNCHER="$REPO_ROOT/scripts/sglang_disagg/run_xPyD_models.slurm" + +# 8.1 suite driver must no longer swallow replay failures with `|| true` +# (exclude comment lines to avoid false positives). +if grep -vE '^[[:space:]]*#' "$SUITE_DRIVER" \ + | grep -q 'run_agentic_replay_and_write_outputs.*|| true'; then + _fail "suite driver: active '|| true' on replay call removed" +else + _pass "suite driver: active '|| true' on replay call removed" +fi + +# 8.2 sglang launcher must set pipefail before the tee pipe. +if grep -vE '^[[:space:]]*#' "$SLURM_LAUNCHER" | grep -q 'set -o pipefail'; then + _pass "sglang launcher: 'set -o pipefail' present" +else + _fail "sglang launcher: 'set -o pipefail' present" +fi + +# 8.3 sglang launcher must expand $HOME in AGENTIC_CONFIG before docker forward. +if grep -vE '^[[:space:]]*#' "$SLURM_LAUNCHER" | grep -q 'AGENTIC_CONFIG="\$(eval echo'; then + _pass "sglang launcher: AGENTIC_CONFIG \$HOME expansion present" +else + _fail "sglang launcher: AGENTIC_CONFIG \$HOME expansion present" +fi + +# 8.4 vLLM launcher must expand $HOME in AGENTIC_CONFIG before docker forward. +VLLM_LAUNCHER="$REPO_ROOT/scripts/vllm_dissag/run_xPyD_models.slurm" +if grep -vE '^[[:space:]]*#' "$VLLM_LAUNCHER" | grep -q 'AGENTIC_CONFIG=.*/#\$HOME'; then + _pass "vllm launcher: AGENTIC_CONFIG \$HOME expansion present" +else + _fail "vllm launcher: AGENTIC_CONFIG \$HOME expansion present" +fi + +# --------------------------------------------------------------------------- +echo "======================================================" +echo " run_offline: ${pass} passed, ${fail} failed" +echo "======================================================" +[ "$fail" -eq 0 ] diff --git a/scripts/common/agentx/tests/submit_matrix.sh b/scripts/common/agentx/tests/submit_matrix.sh new file mode 100644 index 00000000..b777c7eb --- /dev/null +++ b/scripts/common/agentx/tests/submit_matrix.sh @@ -0,0 +1,130 @@ +#!/bin/bash +# AgentX campaign driver: parameterized sbatch fan-out over a matrix of cells. +# Intended to run FROM THE OCI LOGIN NODE (submit-only). Each cell submits one +# run_xPyD_models.slurm job (sglang_disagg or vllm_dissag entrypoint) with the +# AgentX env for that (backend x connector x mode x workload) combination and +# records JOBID -> cell in a job-map file for later harvesting. +# +# DRY_RUN=1 prints the sbatch commands WITHOUT submitting (validate on banff). +# +# Common env (shared by every cell): +# MODEL_NAME model to serve (default DeepSeek-V3) +# PARTITION sbatch -p partition (default amd-rccl) +# DOCKER_IMAGE_NAME vLLM image (used by backend=vllm cells) +# SGLANG_IMAGE sglang image (used by backend=sglang cells) +# TIME sbatch --time minutes (default 90) +# JOB_MAP job-map output file (default ./agentx_jobmap..tsv) +# +# Matrix: CELLS is a newline/semicolon list of cells, each a ':'-delimited tuple: +# backend:connector:mode:workload:max_model_len:nodes +# backend sglang | vllm +# connector rixl | moriio +# mode 0 (TP) | 1 (wideEP); maps to DP_MODE (sglang) / WIDE_EP (vllm) +# workload AGENTIC_WORKLOAD name (e.g. conformance_256k) +# max_model_len MAX_MODEL_LEN (0 = auto-detect served window) +# nodes total nodes -> sbatch -N/-n +# +# Usage: +# DRY_RUN=1 bash scripts/common/agentx/tests/submit_matrix.sh # preview only +# bash scripts/common/agentx/tests/submit_matrix.sh # submit (login node) +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$HERE/../../../.." && pwd)" + +MODEL_NAME="${MODEL_NAME:-DeepSeek-V3}" +PARTITION="${PARTITION:-amd-rccl}" +DOCKER_IMAGE_NAME="${DOCKER_IMAGE_NAME:-}" +SGLANG_IMAGE="${SGLANG_IMAGE:-}" +TIME="${TIME:-90}" +JOB_MAP="${JOB_MAP:-$PWD/agentx_jobmap.$(date +%s).tsv}" + +_is_dry=0 +[ "${DRY_RUN:-0}" = "1" ] && _is_dry=1 + +# Default matrix (override by exporting CELLS). One cell per line: +# backend:connector:mode:workload:max_model_len:nodes +CELLS="${CELLS:-$(cat <<'EOF' +sglang:moriio:1:conformance_256k:262144:3 +sglang:moriio:1:conformance_512k:524288:3 +vllm:rixl:0:conformance_256k:262144:2 +vllm:moriio:1:conformance_256k:262144:2 +EOF +)}" + +echo "=== agentx submit matrix (DRY_RUN=${_is_dry}) ===" +echo " model=${MODEL_NAME} partition=${PARTITION} time=${TIME}" +echo " vllm_image=${DOCKER_IMAGE_NAME} sglang_image=${SGLANG_IMAGE}" +echo " job_map=${JOB_MAP}" +echo "" + +[ "$_is_dry" = "1" ] || : > "$JOB_MAP" + +_submit_cell() { + local backend="$1" connector="$2" mode="$3" workload="$4" mml="$5" nodes="$6" + + local slurm_dir image jobname + local -a envs + envs=(BENCHMARK_SCRIPT=agentic "MODEL_NAME=${MODEL_NAME}" "AGENTIC_WORKLOAD=${workload}" "MAX_MODEL_LEN=${mml}") + + case "$backend" in + sglang) + slurm_dir="$REPO_ROOT/scripts/sglang_disagg" + image="$SGLANG_IMAGE" + envs+=("DOCKER_IMAGE_NAME=${image}" "DP_MODE=${mode}") + [ "$connector" = "moriio" ] && [ "$mode" = "1" ] && envs+=("RUN_MORI=1") + ;; + vllm) + slurm_dir="$REPO_ROOT/scripts/vllm_dissag" + image="$DOCKER_IMAGE_NAME" + envs+=("DOCKER_IMAGE_NAME=${image}" "CONNECTOR=${connector}" "WIDE_EP=${mode}") + [ "$connector" = "moriio" ] && [ "$mode" = "1" ] && envs+=("RUN_MORI=1") + ;; + *) + echo " SKIP unknown backend '$backend' in cell" >&2 + return 0 + ;; + esac + + jobname="agx_${backend}_${connector}_ep${mode}_${workload}" + + local -a cmd + cmd=(env "${envs[@]}" sbatch -N "$nodes" -n "$nodes" -p "$PARTITION" \ + --time="$TIME" -J "$jobname" "$slurm_dir/run_xPyD_models.slurm") + + echo "# cell: ${backend}:${connector}:mode${mode}:${workload}:mml${mml}:N${nodes}" + printf '%q ' "${cmd[@]}"; echo + + if [ "$_is_dry" = "1" ]; then + return 0 + fi + # --parsable makes sbatch print just the JOBID; record JOBID -> cell. + local jobid + jobid="$(env "${envs[@]}" sbatch --parsable -N "$nodes" -n "$nodes" -p "$PARTITION" \ + --time="$TIME" -J "$jobname" "$slurm_dir/run_xPyD_models.slurm")" + printf '%s\t%s:%s:mode%s:%s:mml%s:N%s\n' \ + "$jobid" "$backend" "$connector" "$mode" "$workload" "$mml" "$nodes" >> "$JOB_MAP" + echo " submitted JOBID=${jobid} -> ${JOB_MAP}" + echo "" +} + +while IFS= read -r line; do + line="${line%%#*}" # strip trailing comments + line="$(echo "$line" | tr ';' '\n')" # allow ';'-separated cells too + while IFS= read -r cell; do + cell="$(echo "$cell" | xargs)" # trim whitespace + [ -n "$cell" ] || continue + IFS=':' read -r backend connector mode workload mml nodes <<< "$cell" + if [ -z "${nodes:-}" ]; then + echo " SKIP malformed cell (need 6 ':'-fields): '$cell'" >&2 + continue + fi + _submit_cell "$backend" "$connector" "$mode" "$workload" "$mml" "$nodes" + done <<< "$line" +done <<< "$CELLS" + +if [ "$_is_dry" = "1" ]; then + echo "=== DRY_RUN: nothing submitted ===" +else + echo "=== submitted. job map: ${JOB_MAP} ===" +fi diff --git a/scripts/common/agentx/tests/test_offline.yaml b/scripts/common/agentx/tests/test_offline.yaml new file mode 100644 index 00000000..fa42f9ff --- /dev/null +++ b/scripts/common/agentx/tests/test_offline.yaml @@ -0,0 +1,17 @@ +# Test config for run_offline.sh - includes my_corpus for testing source=corpus path. +# This is NOT the example config used in production; see agentic.example.yaml for that. +run: + scenario: inferencex-agentx-mvp + concurrency: [8, 16] + timeout_sec: 900 + +workloads: + # Testing source=corpus resolution (no preset -> no pre-gate verification) + - name: my_corpus + source: corpus + input_dir: /tmp/agentx_corpora/conformance_256k + + # Testing source=profile with preset + - name: test_profile + source: profile + preset: conformance_256k diff --git a/scripts/common/agentx/verify_agentx_profile.py b/scripts/common/agentx/verify_agentx_profile.py new file mode 100644 index 00000000..c8437e90 --- /dev/null +++ b/scripts/common/agentx/verify_agentx_profile.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +"""Verify a WekaTrace corpus against ITS workload profile's targets + bands. + +Generalization of ROCm/MAD #173's `verify_caseA.py` / `verify_caseB.py`: instead +of hard-coding the Case-A/Case-B numbers, the targets come from the profile so +any workload (conformance_256k, conformance_512k, or a user-defined custom case) verifies against its +own distribution. Prints a per-axis conformance table and "N/N axes within band". + +Targets read from the profile: + ISL P50/P90/P99 <- profile.isl_p + OSL P50/P90/P99 <- profile.osl_p + Delay P50/P90/P99 <- profile.delay_p + Turns P50/P90/P99 <- profile.verify.turns_p (else derived from turns dist) + Cache hit P50 % <- profile.verify.cache_target (else mean(cache_hit)*100) + +Tolerance bands (same multipliers as #173), overridable per profile: + ISL 0.80-1.20 OSL 0.70-1.40 Turns 0.60-1.60 Delay 0.50-2.00 Cache 0.97-1.03 +Per-axis overrides via profile.verify.band_overrides, e.g. Case-B widens +"Input ISL P99" to 0.75-1.25. + +Pure Python stdlib only. YAML profiles are resolved to JSON by the config loader +before being handed here. + +Usage: + verify_agentx_profile.py --profile --corpus + verify_agentx_profile.py --profile-json '{...}' --corpus +""" +import json, glob, os, sys + +DEFAULT_BANDS = { + "isl": (0.80, 1.20), + "osl": (0.70, 1.40), + "turns": (0.60, 1.60), + "delay": (0.50, 2.00), + "cache": (0.97, 1.03), +} + + +def _argval(it, flag): + try: + return next(it) + except StopIteration: + sys.stderr.write(f"[verify_agentx_profile] {flag} requires a value\n") + raise SystemExit(2) + + +def p(a, q): + a = sorted(a) + return a[min(len(a) - 1, int(q * len(a)))] if a else 0 + + +def _weighted_percentiles(values, weights, qs): + """Percentiles of a discrete weighted distribution (cumulative-mass method), + matching how #173 derived turns P50/P90/P99 from the turns values+weights.""" + total = float(sum(weights)) + out = [] + for q in qs: + thresh = q * total + cum = 0.0 + chosen = values[-1] + for v, w in zip(values, weights): + cum += w + if cum >= thresh: + chosen = v + break + out.append(chosen) + return out + + +def measure(corpus): + ai = []; oa = []; tu = []; dl = []; hit = [] + for f in glob.glob(os.path.join(corpus, "*.json")): + with open(f) as fh: + b = json.load(fh) + r = b.get("requests") + if r is None: + raise SystemExit(f"[verify] {f}: session JSON missing 'requests'") + seen = set(); tu.append(len(r)) + for x in r: + if x.get("in"): ai.append(x["in"]) + if x.get("out"): oa.append(x["out"]) + if x.get("think_time") and x["t"] > 0: dl.append(x["think_time"]) + h = x.get("hash_ids") or [] + if h: + nw = sum(1 for z in h if z not in seen); tt = len(h); [seen.add(z) for z in h] + if tt and x["t"] > 0: hit.append(100 * (tt - nw) / tt) + return ai, oa, tu, dl, hit + + +def verify(profile, corpus): + verify_cfg = profile.get("verify", {}) or {} + bands = dict(DEFAULT_BANDS) + for k, v in (verify_cfg.get("bands", {}) or {}).items(): + bands[k] = tuple(v) + overrides = verify_cfg.get("band_overrides", {}) or {} + + missing = [k for k in ("isl_p", "osl_p", "delay_p") if k not in profile] + if missing: + raise SystemExit( + "[verify_agentx_profile] profile missing required field(s): " + + ", ".join(missing)) + isl_p = profile["isl_p"] + osl_p = profile["osl_p"] + delay_p = profile["delay_p"] + turns_p = verify_cfg.get("turns_p") + if turns_p is None: + t = profile["turns"] + turns_p = _weighted_percentiles(list(t["values"]), list(t["weights"]), (0.5, 0.9, 0.99)) + cache_target = verify_cfg.get("cache_target") + if cache_target is None: + lo, hi = profile["cache_hit"] + cache_target = round((lo + hi) / 2.0 * 100) + + ai, oa, tu, dl, hit = measure(corpus) + + rows = [ + ("Input ISL P50", p(ai, .5), isl_p[0], "isl"), + ("Input ISL P90", p(ai, .9), isl_p[1], "isl"), + ("Input ISL P99", p(ai, .99), isl_p[2], "isl"), + ("Output OSL P50", p(oa, .5), osl_p[0], "osl"), + ("Output OSL P90", p(oa, .9), osl_p[1], "osl"), + ("Output OSL P99", p(oa, .99), osl_p[2], "osl"), + ("Turns P50", p(tu, .5), turns_p[0], "turns"), + ("Turns P90", p(tu, .9), turns_p[1], "turns"), + ("Turns P99", p(tu, .99), turns_p[2], "turns"), + ("Delay P50 (s)", p(dl, .5), delay_p[0], "delay"), + ("Delay P90 (s)", p(dl, .9), delay_p[1], "delay"), + ("Delay P99 (s)", p(dl, .99), delay_p[2], "delay"), + ("Cache hit P50 %", p(hit, .5), cache_target, "cache"), + ] + + name = profile.get("name", "?") + print(f"corpus={corpus} profile={name} sessions={len(tu)} requests={len(ai)}\n") + print(f"{'axis':<20}{'measured':>12}{'target':>10}{' verdict'}") + print("-" * 54) + npass = 0 + for axis, meas, tgt, group in rows: + lo, hi = overrides.get(axis, bands[group]) + verdict = "PASS" if (tgt * lo) <= meas <= (tgt * hi) else "off" + npass += verdict == "PASS" + print(f"{axis:<20}{meas:>12,.0f}{tgt:>10,.0f} {verdict}") + print("-" * 54) + print(f"{npass}/{len(rows)} axes within band") + return npass, len(rows) + + +def _load_profile(path): + with open(path) as f: + return json.load(f) + + +def main(argv): + profile = None + corpus = None + it = iter(argv) + for a in it: + if a in ("--profile", "-p"): + profile = _load_profile(_argval(it, a)) + elif a == "--profile-json": + profile = json.loads(_argval(it, a)) + elif a in ("--corpus", "-c"): + corpus = _argval(it, a) + elif a in ("-h", "--help"): + print(__doc__) + return 0 + else: + sys.stderr.write(f"[verify_agentx_profile] unknown arg: {a}\n") + return 2 + if profile is None or corpus is None: + sys.stderr.write("usage: verify_agentx_profile.py --profile P.json --corpus DIR\n") + return 2 + npass, total = verify(profile, corpus) + return 0 if npass == total else 1 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/scripts/common/benchmark_agentic.sh b/scripts/common/benchmark_agentic.sh new file mode 100755 index 00000000..8f95c061 --- /dev/null +++ b/scripts/common/benchmark_agentic.sh @@ -0,0 +1,152 @@ +#!/bin/bash +# Shared agentic trace-replay benchmark hook for disaggregated P/D launchers. +# +# Backend-parameterized (--backend sglang|vllm, or AGENTIC_BACKEND) single source +# of truth for what used to be two byte-identical-except-three-axes hooks. Thin +# per-backend shims (scripts/_disagg/benchmark_agentic.sh) exec this +# script with the right --backend. Drop-in alternative to benchmark_xPyD.sh: +# instead of the random ISL/OSL concurrency sweep, it replays real Claude Code +# agentic traces via aiperf's inferencex-agentx-mvp scenario against the backend +# router/proxy and writes aiperf artifacts + an aggregate JSON + plots. Selected +# by the launcher via +# export BENCHMARK_SCRIPT_FILE=benchmark_agentic.sh (BENCHMARK_SCRIPT=agentic) +# +# The three backend axes live in the `case "$backend"` below: (1) AGENTIC_PORT +# default, (2) the ctx-window resolver endpoint list, and (3) the lib-locator +# repo-dir candidates (superset, harmless when an env var is unset). +# +# Testable standalone (Phase 1/2) against a running server, and with DRY_RUN=1 +# without any server. Env knobs are documented in scripts/common/agentic_lib.sh. +set -uo pipefail + +while [[ $# -gt 0 ]]; do + case "$1" in + --backend) backend="$2"; shift 2;; + *) break;; + esac +done +backend="${backend:-${AGENTIC_BACKEND:-}}" +case "$backend" in + sglang|vllm) : ;; + "") echo "[agentic][ERROR] --backend or AGENTIC_BACKEND required (sglang|vllm)" >&2; exit 2;; + *) echo "[agentic][ERROR] unknown backend '$backend'" >&2; exit 2;; +esac + +_here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# Locate the shared lib. In-container the launcher dir may be mounted somewhere +# whose sibling common/ does not exist, so search a superset of candidates and +# allow AGENTIC_LIB to override. This file lives in scripts/common/, so +# agentic_lib.sh is a SIBLING (normal case); the repo-dir candidates cover the +# in-container mounts of both sglang (MOONCAKE_REPO_DIR) and vllm +# (NIXL_COOKBOOK_PATH / NIXL_REPO_DIR) launchers. +_agentic_lib="" +for _cand in \ + "${AGENTIC_LIB:-}" \ + "${_here}/agentic_lib.sh" \ + "${_here}/../common/agentic_lib.sh" \ + "${MOONCAKE_REPO_DIR:-}/../common/agentic_lib.sh" \ + "${NIXL_COOKBOOK_PATH:-}/../common/agentic_lib.sh" \ + "${NIXL_REPO_DIR:-}/../common/agentic_lib.sh"; do + if [ -n "$_cand" ] && [ -f "$_cand" ]; then _agentic_lib="$_cand"; break; fi +done +[ -n "$_agentic_lib" ] || { echo "[agentic][ERROR] agentic_lib.sh not found (set AGENTIC_LIB)" >&2; exit 1; } +# shellcheck source=/dev/null +source "$_agentic_lib" + +# Agentic benchmarking is an explicit opt-in path, so permit the pinned uv +# install by default (overridable with AGENTIC_ALLOW_UV_INSTALL=0). Non-agentic +# launcher paths keep the gate off. +: "${AGENTIC_ALLOW_UV_INSTALL:=1}" +export AGENTIC_ALLOW_UV_INSTALL + +# Backend axes: serve-port default + ctx-window resolver endpoint list +# ("path|kind"; kind drives the parser below). sglang's disagg front-end +# (sglang_router :2322) exposes the window only on the worker, sometimes only via +# /get_server_info on older builds -> two endpoints; vLLM's worker advertises it +# on /v1/models -> one endpoint. +case "$backend" in + sglang) + : "${AGENTIC_PORT:=2322}" # sglang router port + ctx_endpoints=("/v1/models|models" "/get_server_info|serverinfo") + ;; + vllm) + : "${AGENTIC_PORT:=${BENCHMARK_PORT:-${PROXY_PORT:-8000}}}" # vLLM router/proxy port (BENCHMARK_PORT at runtime) + ctx_endpoints=("/v1/models|models") + ;; +esac +RESULT_DIR="${RESULT_DIR:-/run_logs/${SLURM_JOB_ID:-0}}" +# MODEL_PREFIX feeds the trace-loader default; derive from MODEL_NAME if unset. +: "${MODEL_PREFIX:=${MODEL_NAME:-}}" + +# === agentx:BEGIN resolve served context window (disagg) === +# The disagg front-end (router/proxy/shim) does not advertise max_model_len, so +# resolve it from the prefill WORKER (first host:port in AGENTIC_SERVER_METRICS). +# Skipped when pinned (>0) or DRY_RUN; non-disagg (no AGENTIC_SERVER_METRICS) +# falls through to existing auto-detect. +if [ "${DRY_RUN:-0}" != "1" ] && ! { [ -n "${MAX_MODEL_LEN:-}" ] && [ "${MAX_MODEL_LEN}" -gt 0 ] 2>/dev/null; }; then + _worker="${AGENTIC_SERVER_METRICS%% *}" + if [ -z "$_worker" ]; then + agentic_log "AGENTIC_SERVER_METRICS unset (non-disagg); skipping worker max_model_len auto-detect" + else + [[ "$_worker" =~ ^[^[:space:]]+:[0-9]+$ ]] \ + || agentic_die "malformed worker endpoint '$_worker' (expected host:port); pin MAX_MODEL_LEN" + _tried=""; _mml=""; _won="" + for _cand in "${ctx_endpoints[@]}"; do + _path="${_cand%%|*}"; _kind="${_cand##*|}" + _tried="${_tried:+$_tried, }${_path}" + for _i in 1 2 3; do + _mml="$(curl -sf "http://${_worker}${_path}" 2>/dev/null \ + | python3 -c "$(cat <<'PY' +import sys, json +kind = sys.argv[1] +try: + d = json.load(sys.stdin) +except Exception: + print(""); sys.exit() +v = "" +if kind == "models": + data = d.get("data") or [] + if data: + v = data[0].get("max_model_len") or "" +else: + sa = d.get("server_args") or {} + v = d.get("max_model_len") or d.get("context_length") \ + or sa.get("max_model_len") or sa.get("context_length") or "" +print(v or "") +PY +)" "$_kind")" + if [ -n "$_mml" ] && [ "$_mml" != "0" ]; then _won="$_path"; break 2; fi + sleep 2 + done + done + [ -n "$_mml" ] && [ "$_mml" != "0" ] \ + || agentic_die "could not resolve served max_model_len from ${backend} worker ${_worker} (tried ${_tried}); pin MAX_MODEL_LEN" + export MAX_MODEL_LEN="$_mml" + agentic_log "resolved MAX_MODEL_LEN=${MAX_MODEL_LEN} from ${backend} worker ${_worker} (${_won})" + fi +fi +[ "${AGENTIC_RESOLVE_ONLY:-0}" = "1" ] && { echo "MAX_MODEL_LEN=${MAX_MODEL_LEN:-}"; exit 0; } +# === agentx:END resolve served context window === + +# Suite mode: a workloads config (AGENTIC_CONFIG) or a single-workload shorthand +# (AGENTIC_WORKLOAD) runs the generic multi-workload driver. Without either, the +# legacy single hf/inferencex replay below runs UNCHANGED (byte-identical). +if [ -n "${AGENTIC_CONFIG:-}" ] || [ -n "${AGENTIC_WORKLOAD:-}" ]; then + _agentic_suite="$(dirname "$_agentic_lib")/benchmark_agentic_suite.sh" + [ -f "$_agentic_suite" ] || { echo "[agentic][ERROR] suite driver not found: $_agentic_suite" >&2; exit 1; } + # shellcheck source=/dev/null + source "$_agentic_suite" + exit $? +fi + +if [ "${DRY_RUN:-0}" = "1" ]; then + agentic_dry_run "$RESULT_DIR" + exit 0 +fi + +install_agentic_deps +resolve_trace_source +wait_for_router_ready +[ -n "${MODEL:-}" ] || resolve_served_model_name +build_replay_cmd "$RESULT_DIR" +run_agentic_replay_and_write_outputs "$RESULT_DIR" diff --git a/scripts/common/benchmark_agentic_suite.sh b/scripts/common/benchmark_agentic_suite.sh new file mode 100644 index 00000000..bfc57290 --- /dev/null +++ b/scripts/common/benchmark_agentic_suite.sh @@ -0,0 +1,261 @@ +#!/bin/bash +# AgentX suite driver: run a LIST of workloads in sequence against ONE served +# endpoint. Reads a config (agentic.yaml: serving + run + workloads[]) via the +# config loader, then per workload: +# source=profile -> materialize_corpus() (generate + verify N/N pre-gate) +# source=hf -> resolve the --public-dataset loader (download at run time) +# -> context_compat_check() -> build_replay_cmd() -> run into // +# (optionally sweeping concurrency), then a combined suite summary. +# +# Serving is GLOBAL: one model/endpoint, N workloads. Env vars override the file; +# AGENTIC_WORKLOAD= runs a single entry. DRY_RUN=1 prints the resolved +# N-workload plan + each per-workload command + context verdicts (no server). +set -uo pipefail + +_here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# Reuse an already-sourced lib (hook path) or source it ourselves (standalone). +if ! declare -F build_replay_cmd >/dev/null 2>&1; then + _lib="" + for _cand in "${AGENTIC_LIB:-}" "${_here}/agentic_lib.sh" "${_here}/common/agentic_lib.sh"; do + if [ -n "$_cand" ] && [ -f "$_cand" ]; then _lib="$_cand"; break; fi + done + [ -n "$_lib" ] || { echo "[agentic][ERROR] agentic_lib.sh not found (set AGENTIC_LIB)" >&2; exit 1; } + # shellcheck source=/dev/null + source "$_lib" +fi + +SUITE_PY="python3" +command -v "$SUITE_PY" >/dev/null 2>&1 || SUITE_PY="${AIPERF_PYTHON:-python3}" +CONFIG_TOOL="$AGENTX_DIR/agentx_config.py" + +: "${AGENTIC_PORT:=2322}" +RESULT_DIR="${RESULT_DIR:-/run_logs/${SLURM_JOB_ID:-0}}" + +# Config arg: an explicit file, else rely on the AGENTIC_WORKLOAD synth. +_cfg_args=() +if [ -n "${AGENTIC_CONFIG:-}" ]; then + [ -f "$AGENTIC_CONFIG" ] || agentic_die "AGENTIC_CONFIG not found: $AGENTIC_CONFIG" + _cfg_args=(--config "$AGENTIC_CONFIG") +elif [ -z "${AGENTIC_WORKLOAD:-}" ]; then + agentic_die "suite driver needs AGENTIC_CONFIG= or AGENTIC_WORKLOAD=" +fi + +# Resolve global serving + run + the workload name list. +_cfg_shell="$("$SUITE_PY" "$CONFIG_TOOL" "${_cfg_args[@]}" --emit-config-shell)" \ + || agentic_die "config load failed" +eval "$_cfg_shell" + +# Serving is global: apply non-'auto' serving values to the per-workload env. +if [ "${SUITE_SERVING_MODEL}" != "auto" ]; then MODEL="${MODEL:-$SUITE_SERVING_MODEL}"; fi +if [ "${SUITE_PORT}" != "auto" ]; then AGENTIC_PORT="$SUITE_PORT"; fi +if [ "${SUITE_SERVER_METRICS}" != "auto" ]; then AGENTIC_SERVER_METRICS="$SUITE_SERVER_METRICS"; fi +MAX_MODEL_LEN="${SUITE_MAX_MODEL_LEN}" +export AGENTIC_SCENARIO="$SUITE_SCENARIO" + +_is_dry=0 +[ "${DRY_RUN:-0}" = "1" ] && _is_dry=1 + +mkdir -p "$SUITE_CORPUS_DIR" + +if [ "$_is_dry" = "1" ]; then + cat <} + serving.model : ${SUITE_SERVING_MODEL} + serving.max_model_len : ${SUITE_MAX_MODEL_LEN} + serving.port : ${SUITE_PORT} (AGENTIC_PORT=${AGENTIC_PORT}) + serving.server_metrics : ${SUITE_SERVER_METRICS} + run.concurrency : ${SUITE_CONCURRENCY} + run.duration : ${SUITE_DURATION} + run.scenario : ${SUITE_SCENARIO} + workloads (${SUITE_WORKLOAD_NAMES}) + RESULT_DIR : ${RESULT_DIR} + SUITE_CORPUS_DIR : ${SUITE_CORPUS_DIR} +EOF +else + install_agentic_deps + wait_for_router_ready + if [ -z "${MODEL:-}" ] || [ "${MODEL:-}" = "auto" ]; then resolve_served_model_name; fi + # Auto-detect the served context window when not pinned (config value wins). + if [ -z "${MAX_MODEL_LEN:-}" ] || [ "${MAX_MODEL_LEN}" = "0" ]; then + MAX_MODEL_LEN="$(resolve_served_max_model_len)" + agentic_log "auto-detected served max_model_len: ${MAX_MODEL_LEN}" + fi +fi + +# -------------------------------------------------------------------------- +# Per-workload loop +# -------------------------------------------------------------------------- +SUITE_SUMMARY_JSON="${RESULT_DIR}/suite_summary.json" +_summary_rows=() +_suite_failed=0 + +for name in $SUITE_WORKLOAD_NAMES; do + _profile_json="${SUITE_CORPUS_DIR}/${name}.profile.json" + _wl_shell="$("$SUITE_PY" "$CONFIG_TOOL" "${_cfg_args[@]}" --workload "$name" \ + --profile-out "$_profile_json" --emit-workload-shell)" \ + || agentic_die "workload resolve failed: $name" + # Resets WL_* for this iteration. + WL_LOADER=""; WL_PROFILE_FILE=""; WL_MODEL_TAG="" + WL_NUM_DATASET_ENTRIES=""; WL_TRAJ_MIN=""; WL_TRAJ_MAX="" + WL_FILTER_MAX_ISL=""; WL_FILTER_MAX_TURNS=""; WL_FILTER_SAMPLE="" + WL_INPUT_DIR=""; WL_SCENARIO="" + eval "$_wl_shell" + + # Scenario: suite default, overridden per-workload (re-derived each iteration + # so a per-workload override doesn't leak into the next one). + AGENTIC_SCENARIO="$SUITE_SCENARIO" + [ -n "$WL_SCENARIO" ] && AGENTIC_SCENARIO="$WL_SCENARIO" + export AGENTIC_SCENARIO + + # Per-workload trace source + env. + CORPUS_DIR="" + if [ "$WL_SOURCE" = "hf" ]; then + WEKA_LOADER_OVERRIDE="$WL_LOADER" + elif [ "$WL_SOURCE" = "corpus" ]; then + CORPUS_DIR="$WL_INPUT_DIR" + if [ "$_is_dry" != "1" ]; then + if ! { [ -n "$CORPUS_DIR" ] && [ -d "$CORPUS_DIR" ] && [ -n "$(ls -A "$CORPUS_DIR" 2>/dev/null)" ]; }; then + agentic_err "[$name] source=corpus input_dir not found or empty: $CORPUS_DIR" + _suite_failed=1 + continue + fi + fi + fi + export WL_SOURCE CORPUS_DIR + + # Context compatibility vs the served window. + context_compat_check "$name" "$WL_ISL_TAIL" "$MAX_MODEL_LEN" + if [ "${CONTEXT_VERDICT}" = "SKIP" ]; then + agentic_err "[$name] skipped (context)" + _summary_rows+=("$name|SKIP(context)|-|-") + continue + fi + export AGENTIC_MAX_CONTEXT_LENGTH + + # Materialize (generate + verify) for profile workloads. + _verify_out="" + if [ "$WL_SOURCE" = "profile" ]; then + if [ "$_is_dry" = "1" ]; then + CORPUS_DIR="${SUITE_CORPUS_DIR}/${name}" + else + # materialize_corpus runs in a subshell (command substitution), so set + # CORPUS_DIR in THIS shell too (same deterministic path) — otherwise the + # subshell's assignment is lost and --input-file is built empty. + CORPUS_DIR="${SUITE_CORPUS_DIR}/${name}" + _verify_out="$(materialize_corpus "$name" "$WL_PROFILE_FILE")" + echo "$_verify_out" + fi + elif [ "$WL_SOURCE" = "hf" ] && _hf_filter_active; then + # Tier 2: download + trim once into a deterministic filtered corpus dir. + if [ "$_is_dry" = "1" ]; then + CORPUS_DIR="${SUITE_CORPUS_DIR}/$(_hf_corpus_key "$WL_LOADER")" + else + materialize_hf_corpus "$name" "$WL_LOADER" + fi + export CORPUS_DIR + elif [ "$WL_SOURCE" = "corpus" ]; then + # Existing on-disk corpus replayed as-is (CORPUS_DIR already set above). + # Optional pre-gate: verify only if the entry supplied a profile/preset. + if [ "$_is_dry" != "1" ] && [ -n "$WL_PROFILE_FILE" ]; then + agentic_log "verifying corpus '$name' against supplied profile (pre-gate)" + _verify_out="$("${AIPERF_PYTHON:-python3}" "$AGENTX_DIR/verify_agentx_profile.py" \ + --profile "$WL_PROFILE_FILE" --corpus "$CORPUS_DIR")" || { + echo "$_verify_out"; agentic_die "corpus '$name' failed conformance pre-gate (not N/N)"; } + echo "$_verify_out" + fi + fi + + resolve_trace_loader + + # Concurrency sweep (single value => flat result dir; list => per-conc subdirs). + _conc_list="$WL_CONCURRENCY" + _n_conc=$(echo "$_conc_list" | wc -w) + for conc in $_conc_list; do + AGENTIC_CONC="$conc" + DURATION="$WL_DURATION" + if [ "$_n_conc" -gt 1 ]; then + _rdir="${RESULT_DIR}/${name}/conc${conc}" + else + _rdir="${RESULT_DIR}/${name}" + fi + build_replay_cmd "$_rdir" + if [ "$_is_dry" = "1" ]; then + cat < 0: + err = drop / (prof + drop) + if os.path.exists(os.path.join(wl_dir, "RUN_INVALID.json")): + rec["status"] = "INVALID" + rec["theoretical_cache_hit_rate"] = conformance + rec["gpu_cache_hit_rate"] = cache + rec["error_rate"] = err + rows.append(rec) +json.dump({"workloads": rows}, open(out_path, "w"), indent=2) +print(f"{'workload':<16}{'err_rate':>10}{'cache_hit':>12}{'theo_hit':>12}") +print("-" * 50) +for r in rows: + er = "-" if r["error_rate"] is None else f"{r['error_rate']*100:.1f}%" + ch = "-" if r["gpu_cache_hit_rate"] is None else f"{r['gpu_cache_hit_rate']*100:.1f}%" + th = "-" if r["theoretical_cache_hit_rate"] is None else f"{r['theoretical_cache_hit_rate']*100:.1f}%" + print(f"{r['workload']:<16}{er:>10}{ch:>12}{th:>12}") +print("-" * 50) +print(f"suite summary JSON -> {out_path}") +PY + agentic_log "suite complete -> $RESULT_DIR" +fi + +exit "$_suite_failed" diff --git a/scripts/common/preflight_weights.sh b/scripts/common/preflight_weights.sh new file mode 100755 index 00000000..9cb1d883 --- /dev/null +++ b/scripts/common/preflight_weights.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# Phase 0 pre-flight: verify the model weights exist on every target node BEFORE +# sbatch. Local NVMe (/mnt/m2m_nobackup/models_blog) is per-node and NON-UNIFORM +# on OCI amd-rccl, so a disagg job whose nodelist includes a node missing the +# weights will fail deep into bring-up. Run from the login node. +# +# Usage: +# MODEL_DIR=/mnt/m2m_nobackup/models_blog MODEL_NAME=DeepSeek-V3 \ +# NODELIST=useocpm2m-097-083,useocpm2m-097-087 \ +# bash scripts/common/preflight_weights.sh +# +# Confirmed local-NVMe DeepSeek-V3/R1 node set (survey 2026-07-25): +# 008 030 038 083 087 099 119 122 (ABSENT on 137) +set -uo pipefail + +MODEL_DIR="${MODEL_DIR:-/mnt/m2m_nobackup/models_blog}" +MODEL_NAME="${MODEL_NAME:-}" +NODELIST="${NODELIST:-}" +PARTITION="${PARTITION:-amd-rccl}" +[ -n "$MODEL_NAME" ] || { echo "[preflight][ERROR] set MODEL_NAME" >&2; exit 2; } +[ -n "$NODELIST" ] || { echo "[preflight][ERROR] set NODELIST (comma-separated)" >&2; exit 2; } + +target="$MODEL_DIR/$MODEL_NAME" +echo "[preflight] checking $target on: $NODELIST" + +# One task per node; each prints PRESENT/MISSING with its hostname. +# Do NOT mask srun's stderr: a failed srun (or empty output) must be a hard +# failure, otherwise "no MISSING lines" would falsely PASS. +out="$(srun -p "$PARTITION" --nodelist="$NODELIST" \ + --ntasks-per-node=1 --gres=gpu:1 --time=3 --overcommit bash -c \ + "if [ -d '$target' ] && [ -n \"\$(ls -A '$target' 2>/dev/null)\" ]; then \ + echo \"\$(hostname) PRESENT\"; else echo \"\$(hostname) MISSING\"; fi")" +rc=$? + +echo "$out" | sort +if [ "$rc" -ne 0 ]; then + echo "[preflight][FAIL] srun failed (exit $rc) while checking $target on $NODELIST; cannot confirm weights." >&2 + exit 1 +fi +if [ -z "${out//[[:space:]]/}" ]; then + echo "[preflight][FAIL] srun produced no output checking $target on $NODELIST; cannot confirm weights." >&2 + exit 1 +fi +if echo "$out" | grep -q MISSING; then + echo "[preflight][FAIL] some nodes lack $target. Pick from the confirmed set (008 030 038 083 087 099 119 122) or use MODEL_DIR=/shared_inference/models_blog (NFS, uniform)." >&2 + exit 1 +fi +echo "[preflight][OK] all nodes have $target" diff --git a/scripts/common/validate_agentic_result.sh b/scripts/common/validate_agentic_result.sh new file mode 100755 index 00000000..e957f2e7 --- /dev/null +++ b/scripts/common/validate_agentic_result.sh @@ -0,0 +1,75 @@ +#!/bin/bash +# Phase 4 post-benchmark health check: read the agentic aggregate JSON and print +# a PASS/WARN summary for error rate and prefix-cache hit rate. A low GPU cache +# hit rate on a multi-turn agentic replay usually means the disagg router is not +# keeping a conversation's turns on the worker that holds its prefix (missing +# session stickiness) - see the session-affinity note in the SGLang README. +# +# Usage: +# scripts/common/validate_agentic_result.sh +# # or point at a run dir and it picks the newest *.json (excluding sidecars) +# scripts/common/validate_agentic_result.sh /run_logs/ +# +# Thresholds (override via env): +# AGENTIC_MAX_ERROR_RATE (default 0.10) +# AGENTIC_MIN_CACHE_HIT (default 0.30) # warn-only; workload/affinity dependent +set -uo pipefail + +arg="${1:-}" +[ -n "$arg" ] || { echo "[validate][ERROR] usage: validate_agentic_result.sh " >&2; exit 2; } + +json="$arg" +if [ -d "$arg" ]; then + json="$(ls -t "$arg"/*.json 2>/dev/null | grep -v -E 'RUN_INVALID|profile_export|server_metrics' | head -1)" +fi +[ -n "$json" ] && [ -f "$json" ] || { echo "[validate][ERROR] no aggregate JSON found at $arg" >&2; exit 2; } + +MAX_ERR="${AGENTIC_MAX_ERROR_RATE:-0.10}" MIN_HIT="${AGENTIC_MIN_CACHE_HIT:-0.30}" \ +"${AIPERF_PYTHON:-python3}" - "$json" <<'PY' +import json, os, sys +p = sys.argv[1] +d = json.load(open(p)) +max_err = float(os.environ.get("MAX_ERR", "0.10")) +min_hit = float(os.environ.get("MIN_HIT", "0.30")) + +# Error rate is over *measured* requests only. Warmup records are intentionally +# dropped and must not count as failures, so prefer request_accounting: errors are +# records_error_dropped over (records_profiled + records_error_dropped). +acct = d.get("request_accounting", {}) +if acct: + ok = acct.get("records_profiled", 0) + errs = acct.get("records_error_dropped", 0) + total = ok + errs +else: + total = d.get("num_requests_total") or 0 + ok = d.get("num_requests_successful") or 0 + errs = total - ok +err_rate = (errs / total) if total else 0.0 + +sm_cache = d.get("server_metrics", {}).get("cache", {}) +hit = sm_cache.get("gpu_cache_hit_rate") +if hit is None: + hit = d.get("request_metrics", {}).get("cache", {}).get("theoretical_cache_hit_rate") + +tput = d.get("request_metrics", {}).get("throughput", {}) +per_gpu = tput.get("per_gpu", {}).get("total_tput_tps") + +print(f"[validate] file: {p}") +print(f"[validate] requests: {ok}/{total} ok error_rate={err_rate:.1%}") +print(f"[validate] gpu_cache_hit_rate: {hit if hit is None else f'{hit:.1%}'}") +if per_gpu is not None: + print(f"[validate] throughput_per_gpu: {per_gpu:.0f} tok/s") + +status = 0 +if total == 0: + print("[validate][WARN] no requests recorded"); status = 1 +elif err_rate > max_err: + print(f"[validate][WARN] error_rate {err_rate:.1%} exceeds {max_err:.0%}"); status = 1 +else: + print(f"[validate][PASS] error_rate within {max_err:.0%}") + +if hit is not None and hit < min_hit: + print(f"[validate][WARN] cache hit {hit:.1%} < {min_hit:.0%} - check router session affinity (xP>1)") + +sys.exit(status) +PY diff --git a/scripts/common/verify_agentic_image.sh b/scripts/common/verify_agentic_image.sh new file mode 100755 index 00000000..ba09133f --- /dev/null +++ b/scripts/common/verify_agentic_image.sh @@ -0,0 +1,133 @@ +#!/bin/bash +# Phase 0 pre-flight: verify a disagg Docker image can support agentic replay +# BEFORE spending a Slurm allocation. Fails fast with actionable messages. +# +# Checks: +# 1. SGLang exposes the OpenAI chat route aiperf needs (/v1/chat/completions). +# A router that only serves /generate will make aiperf 404 on every turn. +# 2. The Mooncake disaggregation transfer backend is importable (Variant B +# uses RUN_MORI=0 -> KV_TRANSFER_BACKEND=mooncake). If absent, use a +# MoRI-built image or add Mooncake. +# +# Usage: +# DOCKER_IMAGE_NAME= bash scripts/common/verify_agentic_image.sh +# +# -------------------------------------------------------------------------- +# TIMEOUT FORMULA (size Slurm --time from this; agentic replay adds phases the +# random sweep does not have): +# +# --time (s) >= T_pull # docker pull per node (skip if cached) +# + T_load # server weight load: DeepSeek-V3 671B + 642G +# # local-NVMe read dominates; allow >=1200s +# + T_venv # aiperf uv venv build (~300s) +# + N_retry*T_dl # HF trace download, 3 x up to 900s worst case +# + T_warmup # cache warmup + grace (<=1800s) +# + DURATION # measurement window +# + T_agg # aggregation + plots (~120s) +# +# Example DeepSeek-V3 smoke (DURATION=120, cached image, one clean download): +# ~1200 + 300 + 900 + 900(grace) + 120 + 120 ~= 3540s -> request --time>=3600. +# Raise SGLang server-ready/watchdog timeouts for 671B (ROUTER_READY_TIMEOUT_SECONDS). +# -------------------------------------------------------------------------- +set -uo pipefail + +IMG="${DOCKER_IMAGE_NAME:-}" +[ -n "$IMG" ] || { echo "[verify][ERROR] set DOCKER_IMAGE_NAME" >&2; exit 2; } + +fail() { echo "[verify][FAIL] $*" >&2; exit 1; } +ok() { echo "[verify][OK] $*"; } + +echo "[verify] image: $IMG" + +# ============================================================================= +# vLLM branch (AGENTIC_ENGINE=vllm or RUN_VLLM=1). Verifies a vLLM disagg image +# can support agentic replay: KV transfer backend importable + the vLLM OpenAI +# API server module present (serves /v1/chat/completions + /v1/models). Leaves +# the SGLang/RUN_MORI path below untouched. +# ============================================================================= +if [[ "${AGENTIC_ENGINE:-}" == "vllm" || "${RUN_VLLM:-0}" == "1" ]]; then + _conn="${CONNECTOR:-rixl}" + echo "[verify] engine=vllm connector=${_conn}" + + # 1. KV transfer backend importable in the image. + if [[ "$_conn" == "moriio" ]]; then + if docker run --rm --entrypoint bash "$IMG" -lc 'python3 - </dev/null'; then + ok "NIXL (rixl) transfer backend importable (import nixl)" + else + fail "'import nixl' failed in image. Build with WITH_NIXL=1 (the Dockerfile default) or use a MoRIIO image with CONNECTOR=moriio." + fi + fi + + # 2. vLLM OpenAI API server module present (locate, do not import — importing + # pulls the GPU engine and would false-negative in a GPU-less pre-flight). + if docker run --rm --entrypoint bash "$IMG" -lc \ + 'python3 - </dev/null 2>&1'; then + ok "vllm-router binary on PATH" + else + echo "[verify][note] vllm-router not on PATH (ok if using the toy proxy or ROUTER_BINARY override)" + fi + + ok "image pre-flight passed" + exit 0 +fi + + +# 1. Mooncake transfer backend importable inside the image. The canonical import +# for SGLang's --disaggregation-transfer-backend mooncake is +# `from mooncake.engine import TransferEngine` (see +# sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py and +# MAD-private scripts/kvcache_transfer_bench/backends/mooncake/*.py). +# RUN_MORI=1 uses the MoRI IO backend (matches the launcher's RUN_MORI gate), +# which does not require Mooncake, so skip this check there. +if [[ "${RUN_MORI:-0}" == "1" ]]; then + echo "[verify][note] RUN_MORI=1 (MoRI IO backend): skipping Mooncake import check" +elif docker run --rm --entrypoint bash "$IMG" -lc \ + 'python3 -c "from mooncake.engine import TransferEngine" 2>/dev/null'; then + ok "Mooncake transfer backend importable (mooncake.engine.TransferEngine)" +else + fail "mooncake.engine.TransferEngine not importable in image. Use RUN_MORI=1 with a MoRI-built image, or add Mooncake (mooncake-transfer-engine)." +fi + +# 2. SGLang serves the OpenAI chat endpoint (served by sglang.launch_server / +# sglang_router). Use importlib.util.find_spec to LOCATE the http_server / +# openai serving_chat modules without executing them -- importing them pulls +# in the GPU engine, which fails in a GPU-less pre-flight container and would +# give a false negative. Module presence is sufficient to confirm the route. +if docker run --rm --entrypoint bash "$IMG" -lc \ + 'python3 - < + ``train_utils.write_metrics``). Each line is a dict, e.g.:: + + {"perf/step_time_seconds": 0.83, "perf/per_device_tflops": 174.7, + "perf/per_device_tflops_per_sec": 210.4, "learning/loss": 0.123, + "step": 12.0, "run_name": "wan2.1_1.3b_pretrain"} + + This bypasses stdout entirely and is the PREFERRED source. run.sh points + ``PERF_METRICS_FILE`` (-> config metrics_file) at a path in the persisted + run dir and passes it here via --metrics-file. + +Throughput is derived per the retired jax-maxdiffusion_benchmark_report.py: + fps_per_gpu = per_device_batch_size / avg_seconds_per_step + images_per_sec_per_gpu = per_device_batch_size * num_frames / avg_seconds_per_step + TFLOPS_per_gpu = avg TFLOP/s/device + +batch size and frame count are read from the training log's config dump +(both the "Config param : " and the Primus +" : ()" formats are recognized). Averages skip warmup steps. + +Output CSV format (model, performance, metric) — matches +scripts/jax-maxtext/extract_maxtext_perf.py so both feed madengine +multiple_results (primus_perf_output.csv) identically: + model,performance,metric + wan2.1_1.3b-pretrain,7.23,fps_per_gpu + wan2.1_1.3b-pretrain,585.6,images_per_sec_per_gpu + wan2.1_1.3b-pretrain,210.4,TFLOPS_per_gpu +""" +import argparse +import csv +import json +import re +import sys + +# Trailing per-step samples: skip the first SKIP_WARMUP steps, then average. +SKIP_WARMUP = 2 + +_STEP_RE = re.compile( + r"completed step:\s*(\d+),\s*seconds:\s*([0-9][0-9.eE+-]*),\s*TFLOP/s/device:\s*([0-9][0-9.eE+-]*)" +) + + +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") + + +def _config_param(content: str, name: str): + # MaxDiffusion "Config param : " (raw print) ... + m = re.search(rf"Config param {re.escape(name)}:\s*(.+)", content) + if m: + return _ANSI_RE.sub("", m.group(1)).strip() + # ... or the Primus config dump " : ()" format. + m = re.search(rf"(?:^|\]|\s){re.escape(name)}\s*:\s*(.+?)\s*\((?:bool|int|float|str|list|NoneType|tuple|dict)\)", content, re.MULTILINE) + return _ANSI_RE.sub("", m.group(1)).strip() if m else None + + +def _parse_frames(raw): + if raw is None: + return None + t = str(raw).strip().lower() + if t in ("", "none", "null"): + return None + try: + return int(float(t)) + except ValueError: + return None + + +def _effective_num_frames(content: str) -> float: + """Frames used for throughput. FLUX (image) = 1; WAN uses synthetic-override + logic (synthetic_override_num_frames when dataset_type=synthetic, else + num_frames / data_frames).""" + model_name = (_config_param(content, "model_name") or "").lower() + pretrained = (_config_param(content, "pretrained_model_name_or_path") or "").lower() + if "flux" in f"{model_name} {pretrained}": + return 1.0 + + dataset_type = (_config_param(content, "dataset_type") or "").strip().lower() + override = _parse_frames(_config_param(content, "synthetic_override_num_frames")) + num_frames = _parse_frames(_config_param(content, "num_frames")) + data_frames = _parse_frames(_config_param(content, "data_frames")) + + if dataset_type == "synthetic" and override is not None: + chosen = override + elif num_frames is not None: + chosen = num_frames + elif data_frames is not None: + chosen = data_frames + else: + chosen = None + return float(chosen) if chosen is not None else 1.0 + + +def _samples_from_metrics_file(metrics_file: str): + """Return (seconds[], tflops[]) parsed from the JSON-lines metrics file, or + ([], []) if the file is missing/empty/unparseable.""" + seconds, tflops = [], [] + try: + with open(metrics_file, "r", encoding="utf-8", errors="ignore") as f: + for raw in f: + raw = raw.strip() + if not raw: + continue + try: + d = json.loads(raw) + except ValueError: + continue + s = d.get("perf/step_time_seconds") + t = d.get("perf/per_device_tflops_per_sec") + if s is None or t is None: + continue + try: + seconds.append(float(s)) + tflops.append(float(t)) + except (TypeError, ValueError): + continue + except OSError: + return [], [] + return seconds, tflops + + +def _samples_from_log(content: str): + matches = _STEP_RE.findall(content) + seconds = [float(m[1]) for m in matches] + tflops = [float(m[2]) for m in matches] + return seconds, tflops + + +def extract_metrics(log_path: str, metrics_file: str = "") -> dict: + try: + with open(log_path, "r", encoding="utf-8", errors="ignore") as f: + content = f.read() + except OSError as e: + print(f"Error reading log {log_path}: {e}", file=sys.stderr) + content = "" + + # Prefer the JSON-lines metrics file; fall back to the stdout log line. + seconds, tflops = ([], []) + source = "" + if metrics_file: + seconds, tflops = _samples_from_metrics_file(metrics_file) + if seconds: + source = "metrics_file" + if not seconds: + seconds, tflops = _samples_from_log(content) + if seconds: + source = "log" + + if not seconds: + return {} + + # Drop warmup (compile) steps, then average. + v_seconds = seconds[SKIP_WARMUP:] or seconds + v_tflops = tflops[SKIP_WARMUP:] or tflops + avg_seconds = sum(v_seconds) / len(v_seconds) + avg_tflops = sum(v_tflops) / len(v_tflops) + + batch_raw = _config_param(content, "per_device_batch_size") + try: + batch = float(batch_raw) if batch_raw is not None else 1.0 + except ValueError: + batch = 1.0 + frames = _effective_num_frames(content) + + fps = batch / avg_seconds if avg_seconds > 0 else 0.0 + images_per_sec = batch * frames / avg_seconds if avg_seconds > 0 else 0.0 + return { + "fps": f"{fps:.4f}", + "images_per_sec": f"{images_per_sec:.4f}", + "tflops": f"{avg_tflops:.4f}", + "_source": source, + "_nsteps": str(len(v_seconds)), + } + + +def main(): + parser = argparse.ArgumentParser(description="Extract MaxDiffusion perf metrics to multiple_results CSV") + parser.add_argument("log_path", help="Path to training log") + parser.add_argument("output_csv", help="Path to output CSV (primus_perf_output.csv)") + parser.add_argument("--model-id", default="maxdiffusion_run", help="Model id for the CSV rows") + parser.add_argument( + "--metrics-file", + default="", + help="Path to the JSON-lines metrics file written by the trainer (config.metrics_file). Preferred over the log.", + ) + args = parser.parse_args() + + metrics = extract_metrics(args.log_path, args.metrics_file) + if not metrics: + print( + "Error: no MaxDiffusion perf metrics found. Looked in metrics-file " + f"'{args.metrics_file}' and for 'completed step: ..., TFLOP/s/device: ...' " + f"lines in log {args.log_path}.", + file=sys.stderr, + ) + sys.exit(1) + + rows = [ + {"model": args.model_id, "performance": metrics["fps"], "metric": "fps_per_gpu"}, + {"model": args.model_id, "performance": metrics["images_per_sec"], "metric": "images_per_sec_per_gpu"}, + {"model": args.model_id, "performance": metrics["tflops"], "metric": "TFLOPS_per_gpu"}, + ] + + with open(args.output_csv, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=["model", "performance", "metric"]) + writer.writeheader() + writer.writerows(rows) + + print( + f"Wrote {args.output_csv}: {len(rows)} rows from {metrics.get('_source', '?')} " + f"({metrics.get('_nsteps', '?')} steps; fps_per_gpu={rows[0]['performance']}, " + f"images_per_sec_per_gpu={rows[1]['performance']}, TFLOPS_per_gpu={rows[2]['performance']})" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/jax-maxdiffusion/get_models_json.py b/scripts/jax-maxdiffusion/get_models_json.py new file mode 100644 index 00000000..88d36232 --- /dev/null +++ b/scripts/jax-maxdiffusion/get_models_json.py @@ -0,0 +1,139 @@ +""" +Discover Primus JAX/MaxDiffusion example configs as madengine models. + +MaxDiffusion-only: globs examples/maxdiffusion/configs/**/*.yaml from the Primus +submodule (scripts/Primus). These run through the Primus `maxdiffusion` (JAX) +backend — Google's MaxDiffusion WAN/FLUX trainers launched via primus/cli, the +same way jax-maxtext runs MaxText. New MaxDiffusion configs added under +examples/maxdiffusion/configs// are picked up automatically. + +All discovered models build docker/primus_maxdiffusion (rocm/jax-training based + +maxdiffusion installed) and run through run.sh; args pass --config_path . +Mirrors scripts/jax-maxtext/get_models_json.py. +""" +import os +import glob +import subprocess +import sys + +try: + from madengine.utils.discover_models import CustomModel # madengine v2 +except ImportError: + from madengine.tools.discover_models import CustomModel # madengine v1 + +# This file lives in scripts/jax-maxdiffusion; Primus submodule is scripts/Primus. +THIS_DIR = os.path.dirname(os.path.abspath(__file__)) +PRIMUS_ROOT = os.path.normpath(os.path.join(THIS_DIR, "..", "Primus")) +FETCH_SCRIPT = os.path.normpath(os.path.join(THIS_DIR, "..", "..", "tools", "fetch_primus.sh")) +CONFIGS_GLOB = os.path.join(PRIMUS_ROOT, "examples", "maxdiffusion", "configs", "**", "*.yaml") + +# JAX/MaxDiffusion image, relative to scripts/jax-maxdiffusion. +DOCKERFILE = "../../docker/primus_maxdiffusion" + +# Multi-node-only models (matched against the base model token of a config +# filename, e.g. "wan2.1_14b-pretrain" -> "wan2.1_14b"). None yet: the current +# WAN/FLUX benchmark configs run single-node on 8 GPUs. Override with +# JAX_MAXDIFFUSION_INCLUDE_MULTINODE=1 to discover any listed here. +MULTINODE_MODELS = set() + +# Device -> GPU arch that should SKIP that device's configs (madengine skip_gpu_arch). +# Mirrors jax-maxtext: a single discovery works on both host types; only the +# host-appropriate configs run, the others are recorded as SKIPPED. +ARCH_SKIP_GPU = {"MI300X": "gfx950", "MI355X": "gfx942"} + + +def _precision_from_name(short_name: str) -> str: + """Infer training precision from a config filename (…-fp8-…, …-nanoo_fp8-…, else bf16).""" + lowered = short_name.lower() + if "nanoo_fp8" in lowered: + return "nanoo_fp8" + if "fp8" in lowered: + return "fp8" + return "bf16" + + +def _have_primus(): + """Report whether the Primus checkout these models come from is usable. + + Mirrors scripts/jax-maxtext/get_models_json.py. Discovery is the only host-side + hook that runs before the image build, which is where the checkout has to exist + (both primus_* dockerfiles COPY it from the build context). Fetching is opt-in + even so: cloning over the network is a surprising side effect of listing models, + and it would fire on every madengine invocation. + """ + if os.path.isdir(PRIMUS_ROOT): + return True + if os.environ.get("MAD_AUTO_FETCH_PRIMUS", "") not in ("", "0"): + print("MAD_AUTO_FETCH_PRIMUS is set: fetching Primus into %s" % PRIMUS_ROOT, file=sys.stderr) + rc = subprocess.call(["bash", FETCH_SCRIPT], stdout=sys.stderr.fileno()) + if rc == 0 and os.path.isdir(PRIMUS_ROOT): + return True + print("ERROR: %s failed (exit %d); no JAX/MaxDiffusion models discovered." % (FETCH_SCRIPT, rc), file=sys.stderr) + return False + print( + "WARNING: no Primus checkout at %s, so no JAX/MaxDiffusion models can be discovered. " + "Run tools/fetch_primus.sh, or set MAD_AUTO_FETCH_PRIMUS=1 to fetch it here." % PRIMUS_ROOT, + file=sys.stderr, + ) + return False + + +def list_models(): + # Default/smoke-test entry -> "jax-maxdiffusion/default". Reachable only via the scoped + # name (--tags jax-maxdiffusion/default); tags is ["default"] with no family/arch/ + # precision tags so it never appears in sweeps like --tags maxdiffusion or --tags jax + # and cannot duplicate the per-yaml entry for the same config. + models = [ + CustomModel( + name="default", + dockerfile=DOCKERFILE, + dockercontext=".", + scripts="run.sh", + data="", + n_gpus="8", + owner="mad.support@amd.com", + timeout=86400, + training_precision="bf16", + tags=["default"], + args="--config_path examples/maxdiffusion/configs/MI355X/wan2.1_1.3b-pretrain.yaml", + multiple_results="primus_perf_output.csv", + skip_gpu_arch=ARCH_SKIP_GPU["MI355X"], + ) + ] + if not _have_primus(): + return models + include_multinode = os.environ.get("JAX_MAXDIFFUSION_INCLUDE_MULTINODE", "") not in ("", "0") + for yaml_path in sorted(glob.glob(CONFIGS_GLOB)): + rel_path = os.path.relpath(yaml_path, PRIMUS_ROOT) + # Path shape: examples/maxdiffusion/configs//.yaml + parts = rel_path.split(os.sep) + if len(parts) < 5: + continue + arch = parts[3] # MI300X, MI355X, etc. + short_name = os.path.splitext(os.path.basename(yaml_path))[0] + base_model = short_name.split("-")[0] + if base_model in MULTINODE_MODELS and not include_multinode: + continue + precision = _precision_from_name(short_name) + # discover_models prefixes discovered names with this dir (jax-maxdiffusion/), + # so the final madengine tag is jax-maxdiffusion/maxdiffusion__. + name = f"maxdiffusion_{arch}_{short_name}" + tags = ["maxdiffusion", "jax", arch, short_name, precision] + models.append( + CustomModel( + name=name, + dockerfile=DOCKERFILE, + dockercontext=".", + scripts="run.sh", + data="", + n_gpus="8", + owner="mad.support@amd.com", + timeout=86400, + training_precision=precision, + tags=tags, + args=f"--config_path {rel_path}", + multiple_results="primus_perf_output.csv", + skip_gpu_arch=ARCH_SKIP_GPU.get(arch, ""), + ) + ) + return models diff --git a/scripts/jax-maxdiffusion/run.sh b/scripts/jax-maxdiffusion/run.sh new file mode 100755 index 00000000..185da6a5 --- /dev/null +++ b/scripts/jax-maxdiffusion/run.sh @@ -0,0 +1,141 @@ +#!/usr/bin/env bash +############################################################################### +# +# MIT License +# +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# + +# Wrapper for Primus JAX/MaxDiffusion pretrain when run via madengine (local, SLURM, K8s). +# Sets EXP from PRIMUS_CONFIG_PATH or --config_path, runs Primus examples/run_pretrain.sh +# with BACKEND=MaxDiffusion, then extracts fps/tflops into primus_perf_output.csv for +# madengine multiple_results. Same shape as scripts/jax-maxtext/run.sh. +set -e + +# madengine invokes this as `cd run_directory && bash run.sh ...`. +RUN_DIR="$(pwd)" + +# Primus root: repo checkout, then image COPY / K8s ConfigMap extract, then env, then legacy paths. +script_dir="$(cd "$(dirname "$0")" && pwd)" +if [[ -f "$script_dir/../Primus/examples/run_pretrain.sh" ]]; then + export PRIMUS_ROOT="$(cd "$script_dir/../Primus" && pwd)" +elif [[ -f "/workspace/Primus/examples/run_pretrain.sh" ]]; then + export PRIMUS_ROOT="/workspace/Primus" +elif [[ -n "${PRIMUS_ROOT:-}" ]]; then + : +elif [[ -f "/opt/primus/examples/run_pretrain.sh" ]]; then + export PRIMUS_ROOT="/opt/primus" +elif [[ -f "/workspace/examples/run_pretrain.sh" ]]; then + export PRIMUS_ROOT="/workspace" +else + echo "ERROR: Could not find Primus run_pretrain.sh. Set PRIMUS_ROOT or use a repo with scripts/Primus submodule." >&2 + exit 1 +fi + +# EXP is required by run_pretrain.sh. --config_path must also be stripped from the +# forwarded args: run_pretrain.sh appends leftovers to the training command and it is +# not a valid MaxDiffusion flag. +forward_args=() +if [[ -n "${PRIMUS_CONFIG_PATH:-}" ]]; then + export EXP="$PRIMUS_CONFIG_PATH" + forward_args=("$@") +else + export EXP="" + args=("$@") + i=0 + while [[ $i -lt ${#args[@]} ]]; do + if [[ "${args[i]}" == "--config_path" && -n "${args[i+1]:-}" ]]; then + export EXP="${args[i+1]}" + i=$((i + 2)) + continue + fi + forward_args+=("${args[i]}") + i=$((i + 1)) + done +fi + +if [[ -z "$EXP" ]]; then + echo "ERROR: --config_path or PRIMUS_CONFIG_PATH required." >&2 + exit 1 +fi + +# Makes run_pretrain.sh launch primus/cli train pretrain rather than torchrun. +export BACKEND="MaxDiffusion" + +export MAXDIFFUSION_PATH="${MAXDIFFUSION_PATH:-/workspace/maxdiffusion}" +export BACKEND_PATH="${BACKEND_PATH:-$MAXDIFFUSION_PATH}" + +# The image already satisfies requirements-maxdiffusion.txt and owns the pinned +# maxdiffusion stack (patched source, specific transformers/torch), so the per-run pip +# install can only clobber it. PRIMUS_SKIP_PIP=0 restores it. +export PRIMUS_SKIP_PIP="${PRIMUS_SKIP_PIP:-1}" + +# HF_TOKEN for Primus prepare: explicit, then MAD convention, then madengine v2. +if [[ -n "${HF_TOKEN:-}" ]]; then + export HF_TOKEN +elif [[ -n "${MAD_SECRETS_HFTOKEN:-}" ]]; then + export HF_TOKEN="$MAD_SECRETS_HFTOKEN" +elif [[ -n "${MAD_SECRET_HFTOKEN:-}" ]]; then + export HF_TOKEN="$MAD_SECRET_HFTOKEN" +fi + +# Cache weights on the mounted checkout, not Primus's default /workspace/hf_cache in the +# container's writable layer: flux_dev pulls ~58GB and this host's root filesystem also +# holds /var/lib/docker. A re-run then reuses the download instead of refetching. +export HF_HOME="${HF_HOME:-/myworkspace/hf_cache}" + +# This wrapper deliberately exports no perf/arch env. All XLA_FLAGS and NVTE/HIP/HSA +# tunables travel with each config's top-level env: block, and the arch-gated ones are +# applied in-process before JAX init by primus/backends/maxdiffusion/env_spec.py. + +# I/O contract, not a knob: tells Primus where to write the log this wrapper parses. +mkdir -p "$RUN_DIR/output" +export TRAIN_LOG="$RUN_DIR/output/log_mp_pretrain_$(basename "$EXP" .yaml).txt" + +# The trainer writes per-step JSON metrics here (configs bind metrics_file to it). This is +# the reliable perf source: the per-step stdout line does not survive the Primus launcher's +# stdout handling. Parent of run_directory, so it outlives madengine's cleanup. +export PERF_METRICS_FILE="$RUN_DIR/../perf_metrics_$(basename "$EXP" .yaml).jsonl" +rm -f "$PERF_METRICS_FILE" + +# Without these, a hard exit during trainer teardown (a fatal HIP/JAX abort in +# cleanup on_error) discards block-buffered stdout and the traceback, leaving only +# "launcher exited with code 1". The fault handler covers SIGSEGV/SIGABRT/SIGFPE. +export PYTHONUNBUFFERED=1 +export PYTHONFAULTHANDLER=1 + +# EXP paths are relative to PRIMUS_ROOT. No exec: the perf extractor runs after training. +# The `||` is what keeps set -e from exiting here, so a failed run still gets parsed. +cd "$PRIMUS_ROOT" +exitcode=0 +bash "$PRIMUS_ROOT/examples/run_pretrain.sh" "${forward_args[@]}" || exitcode=$? + +# madengine resolves multiple_results against its own CWD (the parent of run_directory) +# and deletes run_directory before parsing perf, so the CSV must go to the parent. +PERF_OUT="$RUN_DIR/../primus_perf_output.csv" +if [[ -f "$TRAIN_LOG" ]]; then + extract_script="${script_dir}/extract_maxdiffusion_perf.py" + [[ -f "$RUN_DIR/extract_maxdiffusion_perf.py" ]] && extract_script="$RUN_DIR/extract_maxdiffusion_perf.py" + python3 "$extract_script" "$TRAIN_LOG" "$PERF_OUT" --model-id "$(basename "$EXP" .yaml)" \ + --metrics-file "$PERF_METRICS_FILE" || true +fi +exit "$exitcode" diff --git a/scripts/jax-maxtext/env_scripts/deepseek2_16b.yml b/scripts/jax-maxtext/env_scripts/deepseek2_16b.yml deleted file mode 100644 index dd426d56..00000000 --- a/scripts/jax-maxtext/env_scripts/deepseek2_16b.yml +++ /dev/null @@ -1,68 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "deepseek2_16b_1node" -base_output_directory: "./" -hardware: "gpu" -steps: 50 -model_name: "deepseek2-16b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -log_period: 100 -#inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -#intra-node parallelism strategy -ici_fsdp_parallelism: 1 -ici_data_parallelism: 1 -ici_expert_parallelism: -1 -remat_policy: "minimal_flash" -use_iota_embed: True -scan_layers: True -async_checkpointing: False -logits_dot_in_fp32: False -profiler: "" -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -max_target_length: 4096 -dataset_type: "synthetic" -per_device_batch_size: 8 -megablox: False -capacity_factor: 1.25 -sparse_matmul: False -sharding_tolerance: 0.05 - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/deepseek2_env_16b.sh b/scripts/jax-maxtext/env_scripts/deepseek2_env_16b.sh deleted file mode 100644 index 10e6e75a..00000000 --- a/scripts/jax-maxtext/env_scripts/deepseek2_env_16b.sh +++ /dev/null @@ -1,43 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export HSA_NO_SCRATCH_RECLAIM=1 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/gfx950_deepseek2_16b.yml b/scripts/jax-maxtext/env_scripts/gfx950_deepseek2_16b.yml deleted file mode 100644 index 991770f0..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_deepseek2_16b.yml +++ /dev/null @@ -1,68 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "deepseek2_16b_1node" -base_output_directory: "./" -hardware: "gpu" -steps: 50 -model_name: "deepseek2-16b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -log_period: 100 -#inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -#intra-node parallelism strategy -ici_fsdp_parallelism: 1 -ici_data_parallelism: 1 -ici_expert_parallelism: -1 -remat_policy: "minimal_flash" -use_iota_embed: True -scan_layers: True -async_checkpointing: False -logits_dot_in_fp32: False -profiler: "" -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -max_target_length: 4096 -dataset_type: "synthetic" -per_device_batch_size: 8 -megablox: False -capacity_factor: 1.25 -sparse_matmul: False -sharding_tolerance: 0.05 - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/gfx950_deepseek2_env_16b.sh b/scripts/jax-maxtext/env_scripts/gfx950_deepseek2_env_16b.sh deleted file mode 100644 index 059a88f1..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_deepseek2_env_16b.sh +++ /dev/null @@ -1,44 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export NVTE_CK_USES_BWD_V3=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export RCCL_WARP_SPEED_AUTO=0 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/gfx950_llama2_70b.yml b/scripts/jax-maxtext/env_scripts/gfx950_llama2_70b.yml deleted file mode 100644 index aa40bfd6..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_llama2_70b.yml +++ /dev/null @@ -1,76 +0,0 @@ -#!/bin/bash -############################################################################### -# -# MIT License -# -# Copyright (c) Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "llama2_70b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 30 -model_name: "llama2-70b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -dcn_pipeline_parallelism: 1 -dcn_tensor_parallelism: 1 -dcn_sequence_parallelism: 1 -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -ici_sequence_parallelism: 1 -ici_tensor_parallelism: 1 -ici_pipeline_parallelism: 1 -ici_tensor_sequence_parallelism: -1 - -remat_policy: 'full' -optimizer_memory_host_offload: False -param_scan_axis: 1 - -use_iota_embed: True -scan_layers: True - -profiler: "" - -async_checkpointing: False -logits_dot_in_fp32: False -megablox: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -per_device_batch_size: 15 -max_target_length: 4096 -dataset_type: "synthetic" -tokenizer_path: "meta-llama/Llama-2-70b" - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/gfx950_llama2_70b_env.sh b/scripts/jax-maxtext/env_scripts/gfx950_llama2_70b_env.sh deleted file mode 100644 index 3d2c553e..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_llama2_70b_env.sh +++ /dev/null @@ -1,44 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export NVTE_CK_USES_BWD_V3=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export RCCL_WARP_SPEED_AUTO=0 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/gfx950_llama2_7b.yml b/scripts/jax-maxtext/env_scripts/gfx950_llama2_7b.yml deleted file mode 100644 index 436da3ba..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_llama2_7b.yml +++ /dev/null @@ -1,65 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "llama2_7b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 50 -model_name: "llama2-7b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -log_period: 100 - #inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 - #intra-node parallelism strategy -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -ici_tensor_sequence_parallelism: -1 -remat_policy: "minimal_flash" -use_iota_embed: True -scan_layers: True -async_checkpointing: False -logits_dot_in_fp32: False -profiler: "" -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -max_target_length: 4096 -per_device_batch_size: 12 -dataset_type: "synthetic" -tokenizer_path: "meta-llama/Llama-2-7b" - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/gfx950_llama2_7b_env.sh b/scripts/jax-maxtext/env_scripts/gfx950_llama2_7b_env.sh deleted file mode 100644 index 3d2c553e..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_llama2_7b_env.sh +++ /dev/null @@ -1,44 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export NVTE_CK_USES_BWD_V3=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export RCCL_WARP_SPEED_AUTO=0 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/gfx950_llama3.1_405b.yml b/scripts/jax-maxtext/env_scripts/gfx950_llama3.1_405b.yml deleted file mode 100644 index 85fc662b..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_llama3.1_405b.yml +++ /dev/null @@ -1,46 +0,0 @@ -base_config: "base.yml" -run_name: "llama3_405B_training" -hardware: "gpu" -steps: 20 -model_name: "llama3.1-405b" -enable_checkpointing: False -attention: "cudnn_flash_te" -dcn_data_parallelism: 1 -dcn_fsdp_parallelism: -1 -dcn_pipeline_parallelism: 1 -dcn_tensor_parallelism: 1 -dcn_sequence_parallelism: 1 -ici_fsdp_parallelism: -1 -ici_data_parallelism: 1 -ici_sequence_parallelism: 1 -ici_tensor_parallelism: 1 -ici_pipeline_parallelism: 1 - -remat_policy: 'full' -optimizer_memory_host_offload: False -param_scan_axis: 1 - -use_iota_embed: True -scan_layers: True - -profiler: "" #"xplane" -skip_first_n_steps_for_profiler: 3 -profiler_steps: 1 - -async_checkpointing: False -logits_dot_in_fp32: False -megablox: False -dtype: "bfloat16" -quantization: "fp8" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -per_device_batch_size: 5 # 3 for 4 nodes, 5 for 8 nodes -max_target_length: 8192 - -dataset_type: "synthetic" -tokenizer_path: "meta-llama/Llama-3.3-70B-Instruct" -max_segments_per_seq: 32 -shardy: True \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/gfx950_llama3.3_70b.yml b/scripts/jax-maxtext/env_scripts/gfx950_llama3.3_70b.yml deleted file mode 100644 index 53e39ee9..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_llama3.3_70b.yml +++ /dev/null @@ -1,75 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "llama3.3_70b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 30 -model_name: "llama3.3-70b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -dcn_pipeline_parallelism: 1 -dcn_tensor_parallelism: 1 -dcn_sequence_parallelism: 1 -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -ici_sequence_parallelism: 1 -ici_tensor_parallelism: 1 -ici_pipeline_parallelism: 1 -ici_tensor_sequence_parallelism: -1 - -remat_policy: 'full' -optimizer_memory_host_offload: False -param_scan_axis: 1 - -use_iota_embed: True -scan_layers: True - -profiler: "" - -async_checkpointing: False -logits_dot_in_fp32: False -megablox: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -per_device_batch_size: 10 -max_target_length: 8192 -dataset_type: "synthetic" -tokenizer_path: "meta-llama/Llama-3.3-70B-Instruct" - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/gfx950_llama3.3_70b_env.sh b/scripts/jax-maxtext/env_scripts/gfx950_llama3.3_70b_env.sh deleted file mode 100644 index 3d2c553e..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_llama3.3_70b_env.sh +++ /dev/null @@ -1,44 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export NVTE_CK_USES_BWD_V3=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export RCCL_WARP_SPEED_AUTO=0 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/gfx950_llama3_70b.yml b/scripts/jax-maxtext/env_scripts/gfx950_llama3_70b.yml deleted file mode 100644 index 2aed92ec..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_llama3_70b.yml +++ /dev/null @@ -1,75 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "llama3_70b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 30 -model_name: "llama3-70b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -dcn_pipeline_parallelism: 1 -dcn_tensor_parallelism: 1 -dcn_sequence_parallelism: 1 -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -ici_sequence_parallelism: 1 -ici_tensor_parallelism: 1 -ici_pipeline_parallelism: 1 -ici_tensor_sequence_parallelism: -1 - -remat_policy: 'full' -optimizer_memory_host_offload: False -param_scan_axis: 1 - -use_iota_embed: True -scan_layers: True - -profiler: "" - -async_checkpointing: False -logits_dot_in_fp32: False -megablox: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -per_device_batch_size: 10 -max_target_length: 8192 -dataset_type: "synthetic" -tokenizer_path: "meta-llama/Meta-Llama-3-70B" - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/gfx950_llama3_70b_env.sh b/scripts/jax-maxtext/env_scripts/gfx950_llama3_70b_env.sh deleted file mode 100644 index 3d2c553e..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_llama3_70b_env.sh +++ /dev/null @@ -1,44 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export NVTE_CK_USES_BWD_V3=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export RCCL_WARP_SPEED_AUTO=0 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/gfx950_llama3_8b.yml b/scripts/jax-maxtext/env_scripts/gfx950_llama3_8b.yml deleted file mode 100644 index 0d1d6e2a..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_llama3_8b.yml +++ /dev/null @@ -1,65 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "llama3_8b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 50 -model_name: "llama3-8b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -log_period: 100 - #inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 - #intra-node parallelism strategy -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -ici_tensor_sequence_parallelism: -1 -remat_policy: "minimal_flash" -use_iota_embed: True -scan_layers: True -async_checkpointing: False -logits_dot_in_fp32: False -profiler: "" -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -max_target_length: 8192 -per_device_batch_size: 9 -dataset_type: "synthetic" -tokenizer_path: "meta-llama/Meta-Llama-3-8B" - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/gfx950_llama3_8b_env.sh b/scripts/jax-maxtext/env_scripts/gfx950_llama3_8b_env.sh deleted file mode 100644 index 3d2c553e..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_llama3_8b_env.sh +++ /dev/null @@ -1,44 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export NVTE_CK_USES_BWD_V3=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export RCCL_WARP_SPEED_AUTO=0 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/gfx950_mixtral_8x7b.yml b/scripts/jax-maxtext/env_scripts/gfx950_mixtral_8x7b.yml deleted file mode 100644 index f9ae5db7..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_mixtral_8x7b.yml +++ /dev/null @@ -1,71 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "mixtral_8x7b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 40 -model_name: "mixtral-8x7b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -log_period: 100 - #inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 - #intra-node parallelism strategy -ici_fsdp_parallelism: 1 -ici_data_parallelism: 1 -ici_expert_parallelism: -1 -remat_policy: "minimal" -use_iota_embed: True -scan_layers: True -moe_dispatch_no_expert_sharding: True - -async_checkpointing: False -logits_dot_in_fp32: False -dtype: "bfloat16" -quantization: "" #"nanoo_fp8" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint - -per_device_batch_size: 11 -max_target_length: 4096 -sparse_matmul: False -megablox: False -capacity_factor: 1 - -dataset_type: "synthetic" -tokenizer_path: "mistralai/Mixtral-8x7B-v0.1" - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/gfx950_mixtral_8x7b_env.sh b/scripts/jax-maxtext/env_scripts/gfx950_mixtral_8x7b_env.sh deleted file mode 100644 index 059a88f1..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_mixtral_8x7b_env.sh +++ /dev/null @@ -1,44 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export NVTE_CK_USES_BWD_V3=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export RCCL_WARP_SPEED_AUTO=0 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/gfx950_qwen3_14b.yml b/scripts/jax-maxtext/env_scripts/gfx950_qwen3_14b.yml deleted file mode 100644 index 3b57bace..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_qwen3_14b.yml +++ /dev/null @@ -1,62 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "qwen3_14b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 20 -model_name: "qwen3-14b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -log_period: 100 - #inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 - #intra-node parallelism strategy -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -remat_policy: "minimal_flash" -use_iota_embed: True -scan_layers: True -async_checkpointing: False -logits_dot_in_fp32: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False -per_device_batch_size: 4 -max_target_length: 8192 -dataset_type: "synthetic" -tokenizer_path: "Qwen/Qwen3-14B" -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/gfx950_qwen3_14b_env.sh b/scripts/jax-maxtext/env_scripts/gfx950_qwen3_14b_env.sh deleted file mode 100644 index 059a88f1..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_qwen3_14b_env.sh +++ /dev/null @@ -1,44 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export NVTE_CK_USES_BWD_V3=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export RCCL_WARP_SPEED_AUTO=0 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/gfx950_qwen3_30b_a3b.yml b/scripts/jax-maxtext/env_scripts/gfx950_qwen3_30b_a3b.yml deleted file mode 100644 index 21350b68..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_qwen3_30b_a3b.yml +++ /dev/null @@ -1,65 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "qwen3_30b_a3b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 20 -model_name: "qwen3-30b-a3b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -log_period: 100 -# Parallelism — MoE model, so use expert parallelism -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -ici_fsdp_parallelism: 1 -ici_data_parallelism: 1 -ici_expert_parallelism: -1 -remat_policy: "minimal" -use_iota_embed: True -scan_layers: True -async_checkpointing: False -logits_dot_in_fp32: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False -per_device_batch_size: 4 -max_target_length: 8192 -sparse_matmul: False -megablox: False -capacity_factor: 1 -dataset_type: "synthetic" -tokenizer_path: "Qwen/Qwen3-30B-A3B" -max_segments_per_seq: 32 -shardy: True \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/gfx950_qwen3_30b_a3b_env.sh b/scripts/jax-maxtext/env_scripts/gfx950_qwen3_30b_a3b_env.sh deleted file mode 100644 index 059a88f1..00000000 --- a/scripts/jax-maxtext/env_scripts/gfx950_qwen3_30b_a3b_env.sh +++ /dev/null @@ -1,44 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export NVTE_CK_USES_BWD_V3=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export RCCL_WARP_SPEED_AUTO=0 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/llama2_70b.yml b/scripts/jax-maxtext/env_scripts/llama2_70b.yml deleted file mode 100644 index 56996fbd..00000000 --- a/scripts/jax-maxtext/env_scripts/llama2_70b.yml +++ /dev/null @@ -1,75 +0,0 @@ -#!/bin/bash -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "llama2_70b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 30 -model_name: "llama2-70b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -dcn_pipeline_parallelism: 1 -dcn_tensor_parallelism: 1 -dcn_sequence_parallelism: 1 -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -ici_sequence_parallelism: 1 -ici_tensor_parallelism: 1 -ici_pipeline_parallelism: 1 - -remat_policy: 'full' -optimizer_memory_host_offload: False -param_scan_axis: 1 - -use_iota_embed: True -scan_layers: True - -profiler: "" - -async_checkpointing: False -logits_dot_in_fp32: False -megablox: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -per_device_batch_size: 15 -max_target_length: 4096 -dataset_type: "synthetic" -tokenizer_path: "meta-llama/Llama-2-70b" - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/llama2_70b_env.sh b/scripts/jax-maxtext/env_scripts/llama2_70b_env.sh deleted file mode 100644 index 10e6e75a..00000000 --- a/scripts/jax-maxtext/env_scripts/llama2_70b_env.sh +++ /dev/null @@ -1,43 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export HSA_NO_SCRATCH_RECLAIM=1 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/llama2_7b.yml b/scripts/jax-maxtext/env_scripts/llama2_7b.yml deleted file mode 100644 index fba522e1..00000000 --- a/scripts/jax-maxtext/env_scripts/llama2_7b.yml +++ /dev/null @@ -1,64 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "llama2_7b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 50 -model_name: "llama2-7b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -log_period: 100 - #inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 - #intra-node parallelism strategy -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -remat_policy: "minimal_flash" -use_iota_embed: True -scan_layers: True -async_checkpointing: False -logits_dot_in_fp32: False -profiler: "" -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -max_target_length: 4096 -per_device_batch_size: 12 -dataset_type: "synthetic" -tokenizer_path: "meta-llama/Llama-2-7b" - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/llama2_7b_env.sh b/scripts/jax-maxtext/env_scripts/llama2_7b_env.sh deleted file mode 100644 index 10e6e75a..00000000 --- a/scripts/jax-maxtext/env_scripts/llama2_7b_env.sh +++ /dev/null @@ -1,43 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export HSA_NO_SCRATCH_RECLAIM=1 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/llama3.3_70b.yml b/scripts/jax-maxtext/env_scripts/llama3.3_70b.yml deleted file mode 100644 index 979ce3eb..00000000 --- a/scripts/jax-maxtext/env_scripts/llama3.3_70b.yml +++ /dev/null @@ -1,74 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "llama3.3_70b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 30 -model_name: "llama3.3-70b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -dcn_pipeline_parallelism: 1 -dcn_tensor_parallelism: 1 -dcn_sequence_parallelism: 1 -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -ici_sequence_parallelism: 1 -ici_tensor_parallelism: 1 -ici_pipeline_parallelism: 1 - -remat_policy: 'full' -optimizer_memory_host_offload: False -param_scan_axis: 1 - -use_iota_embed: True -scan_layers: True - -profiler: "" - -async_checkpointing: False -logits_dot_in_fp32: False -megablox: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -per_device_batch_size: 7 -max_target_length: 8192 -dataset_type: "synthetic" -tokenizer_path: "meta-llama/Llama-3.3-70B-Instruct" - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/llama3.3_70b_env.sh b/scripts/jax-maxtext/env_scripts/llama3.3_70b_env.sh deleted file mode 100644 index 10e6e75a..00000000 --- a/scripts/jax-maxtext/env_scripts/llama3.3_70b_env.sh +++ /dev/null @@ -1,43 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export HSA_NO_SCRATCH_RECLAIM=1 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/llama3_70b.yml b/scripts/jax-maxtext/env_scripts/llama3_70b.yml deleted file mode 100644 index 1f98a6df..00000000 --- a/scripts/jax-maxtext/env_scripts/llama3_70b.yml +++ /dev/null @@ -1,74 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "llama3_70b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 30 -model_name: "llama3-70b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -dcn_pipeline_parallelism: 1 -dcn_tensor_parallelism: 1 -dcn_sequence_parallelism: 1 -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -ici_sequence_parallelism: 1 -ici_tensor_parallelism: 1 -ici_pipeline_parallelism: 1 - -remat_policy: 'full' -optimizer_memory_host_offload: False -param_scan_axis: 1 - -use_iota_embed: True -scan_layers: True - -profiler: "" - -async_checkpointing: False -logits_dot_in_fp32: False -megablox: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -per_device_batch_size: 7 -max_target_length: 8192 -dataset_type: "synthetic" -tokenizer_path: "meta-llama/Meta-Llama-3-70B" - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/llama3_70b_env.sh b/scripts/jax-maxtext/env_scripts/llama3_70b_env.sh deleted file mode 100644 index 10e6e75a..00000000 --- a/scripts/jax-maxtext/env_scripts/llama3_70b_env.sh +++ /dev/null @@ -1,43 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export HSA_NO_SCRATCH_RECLAIM=1 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/llama3_8b.yml b/scripts/jax-maxtext/env_scripts/llama3_8b.yml deleted file mode 100644 index 40c6db52..00000000 --- a/scripts/jax-maxtext/env_scripts/llama3_8b.yml +++ /dev/null @@ -1,64 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "llama3_8b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 50 -model_name: "llama3-8b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -log_period: 100 - #inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 - #intra-node parallelism strategy -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -remat_policy: "minimal_flash" -use_iota_embed: True -scan_layers: True -async_checkpointing: False -logits_dot_in_fp32: False -profiler: "" -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -max_target_length: 8192 -per_device_batch_size: 4 -dataset_type: "synthetic" -tokenizer_path: "meta-llama/Meta-Llama-3-8B" - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/llama3_8b_env.sh b/scripts/jax-maxtext/env_scripts/llama3_8b_env.sh deleted file mode 100644 index 10e6e75a..00000000 --- a/scripts/jax-maxtext/env_scripts/llama3_8b_env.sh +++ /dev/null @@ -1,43 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export HSA_NO_SCRATCH_RECLAIM=1 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/mixtral_8x7b.yml b/scripts/jax-maxtext/env_scripts/mixtral_8x7b.yml deleted file mode 100644 index 5221ec8f..00000000 --- a/scripts/jax-maxtext/env_scripts/mixtral_8x7b.yml +++ /dev/null @@ -1,71 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "mixtral_8x7b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 40 -model_name: "mixtral-8x7b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -log_period: 100 - #inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 - #intra-node parallelism strategy -ici_fsdp_parallelism: 1 -ici_data_parallelism: 1 -ici_expert_parallelism: -1 -remat_policy: "save_dot_with_context_except_mlp" -use_iota_embed: True -scan_layers: True -moe_dispatch_no_expert_sharding: True - -async_checkpointing: False -logits_dot_in_fp32: False -dtype: "bfloat16" -quantization: "" #"nanoo_fp8" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint - -per_device_batch_size: 12 -max_target_length: 4096 -sparse_matmul: False -megablox: False -capacity_factor: 1 - -dataset_type: "synthetic" -tokenizer_path: "mistralai/Mixtral-8x7B-v0.1" - -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/mixtral_8x7b_env.sh b/scripts/jax-maxtext/env_scripts/mixtral_8x7b_env.sh deleted file mode 100644 index 533b2734..00000000 --- a/scripts/jax-maxtext/env_scripts/mixtral_8x7b_env.sh +++ /dev/null @@ -1,43 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export HSA_NO_SCRATCH_RECLAIM=1 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/qwen3_14b.yml b/scripts/jax-maxtext/env_scripts/qwen3_14b.yml deleted file mode 100644 index c0e0382e..00000000 --- a/scripts/jax-maxtext/env_scripts/qwen3_14b.yml +++ /dev/null @@ -1,62 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "qwen3_14b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 20 -model_name: "qwen3-14b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -log_period: 100 - #inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 - #intra-node parallelism strategy -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -remat_policy: "minimal_flash" -use_iota_embed: True -scan_layers: True -async_checkpointing: False -logits_dot_in_fp32: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False -per_device_batch_size: 2 -max_target_length: 8192 -dataset_type: "synthetic" -tokenizer_path: "Qwen/Qwen3-14B" -max_segments_per_seq: 32 -shardy: True diff --git a/scripts/jax-maxtext/env_scripts/qwen3_14b_env.sh b/scripts/jax-maxtext/env_scripts/qwen3_14b_env.sh deleted file mode 100644 index 533b2734..00000000 --- a/scripts/jax-maxtext/env_scripts/qwen3_14b_env.sh +++ /dev/null @@ -1,43 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export HSA_NO_SCRATCH_RECLAIM=1 \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/qwen3_30b_a3b.yml b/scripts/jax-maxtext/env_scripts/qwen3_30b_a3b.yml deleted file mode 100644 index eaaccb64..00000000 --- a/scripts/jax-maxtext/env_scripts/qwen3_30b_a3b.yml +++ /dev/null @@ -1,65 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -base_config: "base.yml" -run_name: "qwen3_30b_a3b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 20 -model_name: "qwen3-30b-a3b" -enable_goodput_recording: False -monitor_goodput: False -enable_checkpointing: False -attention: "cudnn_flash_te" -packing: True -log_period: 100 -# Parallelism — MoE model, so use expert parallelism -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -ici_fsdp_parallelism: 1 -ici_data_parallelism: 1 -ici_expert_parallelism: -1 -remat_policy: "minimal" -use_iota_embed: True -scan_layers: True -async_checkpointing: False -logits_dot_in_fp32: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False -per_device_batch_size: 2 -max_target_length: 8192 -sparse_matmul: False -megablox: False -capacity_factor: 1 -dataset_type: "synthetic" -tokenizer_path: "Qwen/Qwen3-30B-A3B" -max_segments_per_seq: 32 -shardy: True \ No newline at end of file diff --git a/scripts/jax-maxtext/env_scripts/qwen3_30b_a3b_env.sh b/scripts/jax-maxtext/env_scripts/qwen3_30b_a3b_env.sh deleted file mode 100644 index 533b2734..00000000 --- a/scripts/jax-maxtext/env_scripts/qwen3_30b_a3b_env.sh +++ /dev/null @@ -1,43 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# - -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH -export NVTE_USE_HIPBLASLT=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_command_buffer='' --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=4 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export HSA_NO_SCRATCH_RECLAIM=1 \ No newline at end of file diff --git a/scripts/jax-maxtext/extract_maxtext_perf.py b/scripts/jax-maxtext/extract_maxtext_perf.py new file mode 100644 index 00000000..1a655a75 --- /dev/null +++ b/scripts/jax-maxtext/extract_maxtext_perf.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python3 +""" +Extract JAX/MaxText performance metrics from a training log and write a madengine +multiple_results CSV (one row per metric). + +MaxText-only: this parser handles the MaxText per-step log format and nothing else +(no Megatron/TorchTitan formats). MaxText prints per-step lines such as: + + completed step: 50, seconds: 0.85, TFLOP/s/device: 421.3, Tokens/s/device: 12345.6, ... + +Values are averaged over the last N steps, matching the retired JAX report script. + +Output CSV format (model, performance, metric) — one row per metric. The source values are +per-device (Tokens/s/device, TFLOP/s/device), i.e. per-GPU, so the metric names match the +per-GPU convention used by the existing MAD JAX/MaxText perf CSVs: + model,performance,metric + maxtext_run,12345.6,tok_per_s_per_gpu + maxtext_run,421.3,TFLOPS_per_gpu +""" +import argparse +import csv +import re +import sys + +# Number of trailing per-step samples to average (matches the old JAX report). +AVG_WINDOW = 10 + + +def extract_metrics(log_path: str) -> dict: + """Parse a MaxText log and return averaged tps/tflops from the trailing steps.""" + tps_re = re.compile(r'Tokens/s/device:\s*([0-9][0-9.eE+-]*)') + tflops_re = re.compile(r'TFLOP/s/device:\s*([0-9][0-9.eE+-]*)') + tps_samples = [] + tflops_samples = [] + + try: + with open(log_path, "r", encoding="utf-8", errors="ignore") as f: + for line in f: + m = tps_re.search(line) + if m: + try: + tps_samples.append(float(m.group(1))) + except ValueError: + pass + m = tflops_re.search(line) + if m: + try: + tflops_samples.append(float(m.group(1))) + except ValueError: + pass + except OSError as e: + print(f"Error reading log {log_path}: {e}", file=sys.stderr) + return {} + + tps = tflops = None + if tps_samples: + window = tps_samples[-AVG_WINDOW:] + tps = f"{sum(window) / len(window):.4f}" + if tflops_samples: + window = tflops_samples[-AVG_WINDOW:] + tflops = f"{sum(window) / len(window):.4f}" + + return {"tps": tps, "tflops": tflops} + + +def main(): + parser = argparse.ArgumentParser(description="Extract MaxText perf metrics to multiple_results CSV") + parser.add_argument("log_path", help="Path to training log (e.g. output/log_mp_pretrain_*.txt)") + parser.add_argument("output_csv", help="Path to output CSV (e.g. run_directory/primus_perf_output.csv)") + parser.add_argument("--model-id", default="maxtext_run", help="Model id for the CSV rows") + args = parser.parse_args() + + metrics = extract_metrics(args.log_path) + if not metrics or metrics.get("tps") is None: + print(f"Error: No 'Tokens/s/device:' metric found in log {args.log_path}", file=sys.stderr) + sys.exit(1) + + rows = [ + {"model": args.model_id, "performance": metrics.get("tps") or "", "metric": "tok_per_s_per_gpu"}, + {"model": args.model_id, "performance": metrics.get("tflops") or "", "metric": "TFLOPS_per_gpu"}, + ] + + with open(args.output_csv, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=["model", "performance", "metric"]) + writer.writeheader() + writer.writerows(rows) + + print( + f"Wrote {args.output_csv}: {len(rows)} rows " + f"(tok_per_s_per_gpu={rows[0]['performance']}, TFLOPS_per_gpu={rows[1]['performance']})" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/jax-maxtext/get_models_json.py b/scripts/jax-maxtext/get_models_json.py new file mode 100644 index 00000000..0443043d --- /dev/null +++ b/scripts/jax-maxtext/get_models_json.py @@ -0,0 +1,144 @@ +""" +Discover Primus JAX/MaxText example configs as madengine models. + +MaxText-only: globs examples/maxtext/configs/**/*.yaml from the Primus submodule +(scripts/Primus). No Megatron/TorchTitan configs are discovered here — those backends +use rocm/primus:* images and their own MAD integration. New MaxText configs added under +examples/maxtext/configs// are picked up automatically. + +All discovered models build docker/primus_maxtext (rocm/jax-training:maxtext-*, the only +image that ships JAX) and run through run.sh; args pass --config_path . +For SLURM/K8s, supply distributed settings via additional_context. +""" +import os +import glob +import subprocess +import sys + +try: + from madengine.utils.discover_models import CustomModel # madengine v2 +except ImportError: + from madengine.tools.discover_models import CustomModel # madengine v1 + +# This file lives in scripts/jax-maxtext; Primus submodule is scripts/Primus +THIS_DIR = os.path.dirname(os.path.abspath(__file__)) +PRIMUS_ROOT = os.path.normpath(os.path.join(THIS_DIR, "..", "Primus")) +FETCH_SCRIPT = os.path.normpath(os.path.join(THIS_DIR, "..", "..", "tools", "fetch_primus.sh")) +CONFIGS_GLOB = os.path.join(PRIMUS_ROOT, "examples", "maxtext", "configs", "**", "*.yaml") + +# JAX/MaxText image (rocm/jax-training:maxtext-*), relative to scripts/jax-maxtext. +DOCKERFILE = "../../docker/primus_maxtext" + +# Multi-node-only models. The MAD JAX/MaxText suite mirrors the single-node +# env_scripts set (see ROCm/MAD scripts/jax-maxtext/env_scripts); these large +# models require multiple nodes and are intentionally NOT discovered as +# single-node madengine models. Matched against the base model token of a config +# filename (the part before the first '-', e.g. "llama3.1_405B-fp8-pretrain" -> +# "llama3.1_405B"). Override with JAX_MAXTEXT_INCLUDE_MULTINODE=1 to discover them. +MULTINODE_MODELS = {"grok1", "llama3.1_405B", "mixtral_8x22B"} + +# Device -> GPU arch that should SKIP that device's configs (madengine skip_gpu_arch). +# MI300X configs are tuned for gfx942 and skipped on gfx950; MI355X configs are tuned +# for gfx950 and skipped on gfx942. So a single discovery works on both host types: +# only the host-appropriate configs run, the others are recorded as SKIPPED. +ARCH_SKIP_GPU = {"MI300X": "gfx950", "MI355X": "gfx942"} + + +def _precision_from_name(short_name: str) -> str: + """Infer training precision from a config filename (…-fp8-…, …-nanoo_fp8-…, else bf16).""" + lowered = short_name.lower() + if "nanoo_fp8" in lowered: + return "nanoo_fp8" + if "fp8" in lowered: + return "fp8" + return "bf16" + + +def _have_primus(): + """Report whether the Primus checkout these models come from is usable. + + Discovery is the only host-side hook that runs before the image build, which is + where the checkout has to exist (both primus_* dockerfiles COPY it from the build + context). Fetching is opt-in even so: cloning over the network is a surprising + side effect of listing models, and it would fire on every madengine invocation. + Everything goes to stderr to keep discovery's stdout clean. + """ + if os.path.isdir(PRIMUS_ROOT): + return True + if os.environ.get("MAD_AUTO_FETCH_PRIMUS", "") not in ("", "0"): + print("MAD_AUTO_FETCH_PRIMUS is set: fetching Primus into %s" % PRIMUS_ROOT, file=sys.stderr) + rc = subprocess.call(["bash", FETCH_SCRIPT], stdout=sys.stderr.fileno()) + if rc == 0 and os.path.isdir(PRIMUS_ROOT): + return True + print("ERROR: %s failed (exit %d); no JAX/MaxText models discovered." % (FETCH_SCRIPT, rc), file=sys.stderr) + return False + # Say something rather than returning an empty list, which reads as "no MaxText + # models exist" instead of "the checkout they are discovered from is missing". + print( + "WARNING: no Primus checkout at %s, so no JAX/MaxText models can be discovered. " + "Run tools/fetch_primus.sh, or set MAD_AUTO_FETCH_PRIMUS=1 to fetch it here." % PRIMUS_ROOT, + file=sys.stderr, + ) + return False + + +def list_models(): + # Default/smoke-test entry -> "jax-maxtext/default". Reachable only via the scoped + # name (--tags jax-maxtext/default); tags is ["default"] with no family/arch/precision + # tags so it never appears in sweeps like --tags maxtext or --tags jax and cannot + # duplicate the per-yaml entry for the same config. + models = [ + CustomModel( + name="default", + dockerfile=DOCKERFILE, + dockercontext=".", + scripts="run.sh", + data="", + n_gpus="8", + owner="mad.support@amd.com", + timeout=86400, + training_precision="bf16", + tags=["default"], + args="--config_path examples/maxtext/configs/MI300X/llama2_7B-bf16-pretrain.yaml", + multiple_results="primus_perf_output.csv", + skip_gpu_arch=ARCH_SKIP_GPU["MI300X"], + ) + ] + if not _have_primus(): + return models + include_multinode = os.environ.get("JAX_MAXTEXT_INCLUDE_MULTINODE", "") not in ("", "0") + for yaml_path in sorted(glob.glob(CONFIGS_GLOB)): + rel_path = os.path.relpath(yaml_path, PRIMUS_ROOT) + # Path shape: examples/maxtext/configs//.yaml + parts = rel_path.split(os.sep) + if len(parts) < 5: + continue + arch = parts[3] # MI300X, MI355X, etc. + short_name = os.path.splitext(os.path.basename(yaml_path))[0] + # Skip multi-node-only models unless explicitly requested. + base_model = short_name.split("-")[0] + if base_model in MULTINODE_MODELS and not include_multinode: + continue + precision = _precision_from_name(short_name) + # discover_models prefixes discovered names with this dir (jax-maxtext/), so the + # final madengine tag is jax-maxtext/maxtext__. No prefix here. + name = f"maxtext_{arch}_{short_name}" + tags = ["maxtext", "jax", arch, short_name, precision] + models.append( + CustomModel( + name=name, + dockerfile=DOCKERFILE, + dockercontext=".", + scripts="run.sh", + data="", + n_gpus="8", + owner="mad.support@amd.com", + timeout=86400, + training_precision=precision, + tags=tags, + args=f"--config_path {rel_path}", + multiple_results="primus_perf_output.csv", + skip_gpu_arch=ARCH_SKIP_GPU.get(arch, ""), + ) + ) + return models diff --git a/scripts/jax-maxtext/gpu-rocm/deepseek_v2_16b.sh b/scripts/jax-maxtext/gpu-rocm/deepseek_v2_16b.sh deleted file mode 100644 index 0091203b..00000000 --- a/scripts/jax-maxtext/gpu-rocm/deepseek_v2_16b.sh +++ /dev/null @@ -1,75 +0,0 @@ -#!/bin/bash - -set -x -# Create output dir -OUTPUT_DIR="$HOME/output" -mkdir -p $OUTPUT_DIR - -# Environment variables -echo ' export XLA_FLAGS=" --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_latency_hiding_scheduler=TRUE --xla_gpu_enable_cublaslt=True --xla_gpu_graph_level=0 --xla_gpu_autotune_level=5 --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_all_gather_combine_by_dim=FALSE --xla_gpu_memory_limit_slop_factor=95" -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=0.967 -export LD_LIBRARY_PATH=/opt/rocm/lib:$LD_LIBRARY_PATH' > $OUTPUT_DIR/deepseek2_env_16b.sh - -# Model Configuration -echo 'base_config: "base.yml" -run_name: "deepseek2_16b_1node" -base_output_directory: "./" -hardware: "gpu" -steps: 50 -model_name: "deepseek2-16b" -enable_checkpointing: False -attention: "dot_product" -log_period: 100 -#inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -#intra-node parallelism strategy -ici_fsdp_parallelism: 1 -ici_data_parallelism: 1 -ici_expert_parallelism: -1 -remat_policy: "minimal_flash" -use_iota_embed: True -scan_layers: True -async_checkpointing: False -logits_dot_in_fp32: False -profiler: "" -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -max_target_length: 4096 -dataset_type: "synthetic" -per_device_batch_size: 8 -megablox: False -capacity_factor: 1.25 -sparse_matmul: False -sharding_tolerance: 0.05' > $OUTPUT_DIR/deepseek2_16b_gpu.yml - - - -#If podman is available instead of docker, then you need this export otherwise -#comment the below line and uncomment the line after that -#export docker=podman -docker=docker - -#docker run --rm --privileged --network host --device /dev/dri --device /dev/kfd --cap-add=IPC_LOCK --volume /dev/infiniband:/dev/infiniband -w /workspace/maxtext $IMAGE /bin/bash -#--cap-add=IPC_LOCK --volume /dev/infiniband:/dev/infiniband --tmpfs /dev/shm:size=50G -w /workspace/maxtext '${IMAGE}' /bin/bash - -$docker run --rm --privileged --network host --device /dev/dri --device /dev/kfd \ - --cap-add=IPC_LOCK --volume /dev/infiniband:/dev/infiniband \ - -v $HOME:$HOME -v $HOME/data:/home/amd/data --tmpfs /dev/shm:size=50G \ - --mount type=bind,source=$OUTPUT_DIR,target=/workspace/maxtext/output \ - -w /workspace/maxtext $IMAGE /bin/bash -c " - set -e - echo \"Running Deepseek-v2-16b\" - echo '${IMAGE}' - cp $OUTPUT_DIR/deepseek2_env_16b.sh . - mkdir -p configs - cp $OUTPUT_DIR/deepseek2_16b_gpu.yml configs/. - source deepseek2_env_16b.sh - python -m maxtext.trainers.pre_train.train configs/deepseek2_16b_gpu.yml base_output_directory=output 2>&1 |& tee -a deepseek_v2_16b.log - " diff --git a/scripts/jax-maxtext/gpu-rocm/llama2_70b.sh b/scripts/jax-maxtext/gpu-rocm/llama2_70b.sh deleted file mode 100644 index bc48f670..00000000 --- a/scripts/jax-maxtext/gpu-rocm/llama2_70b.sh +++ /dev/null @@ -1,87 +0,0 @@ -#!/bin/bash - -set -x -# Create output dir -OUTPUT_DIR="$HOME/output" -mkdir -p $OUTPUT_DIR - -# Environment variables -echo 'export XLA_FLAGS="--xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_graph_level=0 --xla_gpu_autotune_level=0 --xla_gpu_enable_reduce_scatter_combine_by_dim=false --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_all_reduce_combine_threshold_bytes=8589934592 --xla_gpu_all_gather_combine_threshold_bytes=137438953472 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export XLA_PYTHON_CLIENT_MEM_FRACTION=0.975 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH' > $OUTPUT_DIR/maxtext_env_70b.sh - - -# Model Configuration -echo 'base_config: "base.yml" -run_name: "llama2_70b_training" -hardware: "gpu" -steps: 30 -model_name: "llama2-70b" -enable_checkpointing: False -attention: "cudnn_flash_te" -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -dcn_pipeline_parallelism: 1 -dcn_tensor_parallelism: 1 -dcn_sequence_parallelism: 1 -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -ici_sequence_parallelism: 1 -ici_tensor_parallelism: 1 -ici_pipeline_parallelism: 1 - -remat_policy: 'full' -optimizer_memory_host_offload: False -param_scan_axis: 1 - -use_iota_embed: True -scan_layers: True - -profiler: "" - -async_checkpointing: False -logits_dot_in_fp32: False -megablox: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -per_device_batch_size: 15 -max_target_length: 4096 -hf_path: "parquet" -hf_train_files: "/hf_cache/hub/datasets--legacy-datasets--c4/snapshots/5abe0d085aa23dd9db2a6c1e86cfce4e4db6f0c3/en/partial-train/000*.parquet" -dataset_type: "hf" -tokenizer_path: "meta-llama/Llama-2-70b"' > $OUTPUT_DIR/llama2_70b_gpu.yml - -#If podman is available instead of docker, then you need this export otherwise -#comment the below line and uncomment the line after that -#export docker=podman -docker=docker - - - -# get the test data -echo "For downloading data, we will mount \$HF_HOME to the docker and try to get llama tokenizer directly from there" -echo "Please set \$HF_HOME when calling this script, your HF_HOME is set as" -echo $HF_HOME -hf download legacy-datasets/c4 --include "*.parquet" --repo-type dataset --revision refs/convert/parquet - - -$docker run --rm --privileged --network host --device /dev/dri --device /dev/kfd \ - --cap-add=IPC_LOCK --volume /dev/infiniband:/dev/infiniband \ - -v $HOME:$HOME -v $HOME/data:/home/amd/data -v $HF_HOME:/hf_cache -e HF_HOME=/hf_cache --tmpfs /dev/shm:size=50G \ - --mount type=bind,source=$OUTPUT_DIR,target=/workspace/maxtext/output \ - -w /workspace/maxtext $IMAGE /bin/bash -c " - set -e - echo \"Running Llama-2-70b\" - echo '${IMAGE}' - cp $OUTPUT_DIR/maxtext_env_70b.sh . - mkdir -p configs - cp $OUTPUT_DIR/llama2_70b_gpu.yml configs/llama2_70b_gpu.yml - source maxtext_env_70b.sh - python -m maxtext.trainers.pre_train.train configs/llama2_70b_gpu.yml base_output_directory=output 2>&1 |& tee -a llama2_70b.real.log - " - diff --git a/scripts/jax-maxtext/gpu-rocm/llama2_70b_multinode.sh b/scripts/jax-maxtext/gpu-rocm/llama2_70b_multinode.sh deleted file mode 100644 index e3a3013c..00000000 --- a/scripts/jax-maxtext/gpu-rocm/llama2_70b_multinode.sh +++ /dev/null @@ -1,177 +0,0 @@ -#!/bin/bash - -# ======================================================== # -# SLURM HEADERS # -# ======================================================== # - -#SBATCH --job-name=training_llama2_70B -#SBATCH --output=logs/multinode-job-llama2-70b.%j.out -#SBATCH --time=3:00:00 -#SBATCH --nodes=2 -#SBATCH --ntasks-per-node=1 # setting this to 8 would launch 8 dockers on the single node with 8 GPU -#SBATCH --exclusive -#SBATCH --reservation=gpu-40_gpu-41_gpu-43_gpu-44_gpu-46_gpu-47_gpu-50_gpu-55_reservation - -# SLURM_NNODES -# Total number of nodes in the job's resource allocation. See SLURM_JOB_NUM_NODES. Included for backwards compatibility. - -# SLURM_NODEID -# ID of the nodes allocated. - -# SLURM_NODELIST -# List of nodes allocated to the job. See SLURM_JOB_NODELIST. Included for backwards compatibility. - -# srun echo $SLURM_NNODES -# echo $SLURM_LOCALID -# srun echo $SLURM_LOCALID -# srun -N $SLURM_JOB_NUM_NODES -n $SLURM_JOB_NUM_NODES echo $SLURM_LOCALID -# echo 'echo $SLURM_NODEID' > script.sh -# srun bash ./script.sh -# srun echo $SLURM_NODEID -# srun -N $SLURM_JOB_NUM_NODES -n $SLURM_JOB_NUM_NODES echo $SLURM_NODEID -# srun echo $SLURM_NODELIST -# srun echo $SLURM_JOB_NODELIST -# see https://gist.github.com/TengdaHan/1dd10d335c7ca6f13810fff41e809904 - -set -x -OUTPUT_DIR="$HOME/output" - -# Install required packages -echo ' -apt install iproute2 -y -apt install rdma-core -y -apt install -y linux-headers-"$(uname -r)" libelf-dev -apt install -y gcc make libtool autoconf librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool libibverbs-dev rdma-core strace libibmad5 libibnetdisc5 ibverbs-providers libibumad-dev libibumad3 libibverbs1 libnl-3-dev libnl-route-3-dev -' > $OUTPUT_DIR/install_packages.sh - - -# Environment variables -echo ' -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export NVTE_USE_HIPBLASLT=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_graph_level=0 --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=0 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NCCL_IB_TIMEOUT=20 -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -' > $OUTPUT_DIR/maxtext_env_70b.sh - - -# Model Configuration -echo 'base_config: "base.yml" -run_name: "llama2_70b_training" -hardware: "gpu" -steps: 30 -model_name: "llama2-70b" -enable_checkpointing: False -attention: "cudnn_flash_te" -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 2 -dcn_pipeline_parallelism: 1 -dcn_tensor_parallelism: 1 -dcn_sequence_parallelism: 1 -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -ici_sequence_parallelism: 1 -ici_tensor_parallelism: 1 -ici_pipeline_parallelism: 1 - -remat_policy: 'full' -optimizer_memory_host_offload: False -param_scan_axis: 1 - -use_iota_embed: True -scan_layers: True - -profiler: "" - -async_checkpointing: False -logits_dot_in_fp32: False -megablox: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -per_device_batch_size: 15 -max_target_length: 4096 -dataset_type: "synthetic" -enable_goodput_recording: False -monitor_goodput: False -shardy: False -' > $OUTPUT_DIR/llama2_70b_gpu.yml - - -srun hostname -# srun master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) -# srun export MASTER_ADDR=$master_addr -export MASTER_NAME=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) -export MASTER_ADDR=$(cat /etc/hosts | grep $MASTER_NAME | awk '{print $1}' ) -# MASTER_ADDR=$(cat /etc/hosts | grep gpu-14 | awk '{print $1}' ) -srun echo "MASTER_ADDR="$MASTER_ADDR - -# srun ping $MASTER_ADDR - -#If podman is available instead of docker, then you need this export otherwise -#comment the below line and uncomment the line after that -export docker=podman -#docker=docker - -export NNODES=$SLURM_NNODES -export JAX_COORDINATOR_IP=$MASTER_ADDR -export JAX_COORDINATOR_PORT=1234 - -#Change this to one of the IP interfaces used for communication -export NCCL_SOCKET_IFNAME=ens8np0 -echo $NCCL_SOCKET_IFNAME - -# For Mellanox NIC -export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_8,mlx5_9 -# For Broadcom Thor NIC, uncomment the line below and comment the line above -#export NCCL_IB_HCA=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 -echo $NCCL_IB_HCA - -srun --nodes=$SLURM_JOB_NUM_NODES --ntasks=$SLURM_JOB_NUM_NODES \ - --export=ALL \ - bash -c '\ - NODE_RANK=$SLURM_PROCID; \ - NNODES=$SLURM_JOB_NUM_NODES; \ - $docker run --rm --privileged --network host \ - --device /dev/dri --device /dev/kfd \ - --cap-add=IPC_LOCK \ - --volume /dev/infiniband:/dev/infiniband \ - -v $HOME:$HOME \ - -v $HOME/araina/data:/home/amd/data \ - --tmpfs /dev/shm:size=50G \ - --mount type=bind,source='${OUTPUT_DIR}',target=/workspace/maxtext/output \ - -e NNODES=$NNODES \ - -e NODE_RANK=$NODE_RANK \ - -e JAX_COORDINATOR_IP='${JAX_COORDINATOR_IP}' \ - -e JAX_COORDINATOR_PORT='${JAX_COORDINATOR_PORT}' \ - -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ - -w /workspace/maxtext \ - '${IMAGE}' \ - /bin/bash -c " - set -e - echo \"Running Llama-2-70b\" - echo '${IMAGE}' - echo \"Coordinator IP: \$JAX_COORDINATOR_IP\" - cp '${OUTPUT_DIR}'/install_packages.sh . - cp '${OUTPUT_DIR}'/maxtext_env_70b.sh . - mkdir -p configs - cp '${OUTPUT_DIR}'/llama2_70b_gpu.yml configs/llama2_70b_gpu.yml - source install_packages.sh - source maxtext_env_70b.sh - python -m maxtext.trainers.pre_train.train configs/llama2_70b_gpu.yml 2>&1 |& tee -a llama2_70b.synthetic.log - "' \ No newline at end of file diff --git a/scripts/jax-maxtext/gpu-rocm/llama2_7b.sh b/scripts/jax-maxtext/gpu-rocm/llama2_7b.sh deleted file mode 100644 index 33909b48..00000000 --- a/scripts/jax-maxtext/gpu-rocm/llama2_7b.sh +++ /dev/null @@ -1,76 +0,0 @@ -#!/bin/bash - -set -x -# Create output dir -OUTPUT_DIR="$HOME/output" -mkdir -p $OUTPUT_DIR - -# Environment variables -echo 'export XLA_FLAGS="--xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_graph_level=0 --xla_gpu_autotune_level=0 --xla_gpu_enable_latency_hiding_scheduler=TRUE --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_all_gather_combine_by_dim=FALSE --xla_gpu_memory_limit_slop_factor=95" -export XLA_PYTHON_CLIENT_MEM_FRACTION=0.967 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH' > $OUTPUT_DIR/maxtext_env_7b.sh - -# Model Configuration -echo 'base_config: "base.yml" -run_name: "llama2_7b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 50 -model_name: "llama2-7b" -enable_checkpointing: False -attention: "cudnn_flash_te" -log_period: 100 - #inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 - #intra-node parallelism strategy -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -remat_policy: "minimal_flash" -use_iota_embed: True -scan_layers: False -async_checkpointing: False -logits_dot_in_fp32: False -profiler: "" -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -max_target_length: 4096 -per_device_batch_size: 10 -hf_path: "parquet" -hf_train_files: "/hf_cache/hub/datasets--legacy-datasets--c4/snapshots/5abe0d085aa23dd9db2a6c1e86cfce4e4db6f0c3/en/partial-train/000*.parquet" -dataset_type: "hf" -tokenizer_path: "meta-llama/Llama-2-7b"' > $OUTPUT_DIR/llama2_7b_gpu.yml - -#If podman is available instead of docker, then you need this export otherwise -#comment the below line and uncomment the line after that -#export docker=podman -docker=docker - - - -# get the test data -echo "For downloading data, we will mount \$HF_HOME to the docker and try to get llama tokenizer directly from there" -echo "Please set \$HF_HOME when calling this script, your HF_HOME is set as" -echo $HF_HOME -hf download legacy-datasets/c4 --include "*.parquet" --repo-type dataset --revision refs/convert/parquet - -$docker run --rm --privileged --network host --device /dev/dri --device /dev/kfd \ - --cap-add=IPC_LOCK --volume /dev/infiniband:/dev/infiniband \ - -v $HOME:$HOME -v $HOME/data:/home/amd/data -v $HF_HOME:/hf_cache -e HF_HOME=/hf_cache --tmpfs /dev/shm:size=50G \ - --mount type=bind,source=$OUTPUT_DIR,target=/workspace/maxtext/output \ - -w /workspace/maxtext $IMAGE /bin/bash -c " - set -e - echo \"Running Llama-2-7b\" - echo '${IMAGE}' - cp $OUTPUT_DIR/maxtext_env_7b.sh . - mkdir -p configs - cp $OUTPUT_DIR/llama2_7b_gpu.yml configs/llama2_7b_gpu.yml - source maxtext_env_7b.sh - python -m maxtext.trainers.pre_train.train configs/llama2_7b_gpu.yml base_output_directory=output 2>&1 |& tee -a llama2_7b.real.log - " - diff --git a/scripts/jax-maxtext/gpu-rocm/llama2_7b_multinode.sh b/scripts/jax-maxtext/gpu-rocm/llama2_7b_multinode.sh deleted file mode 100644 index bbfb5b53..00000000 --- a/scripts/jax-maxtext/gpu-rocm/llama2_7b_multinode.sh +++ /dev/null @@ -1,165 +0,0 @@ -#!/bin/bash - -# ======================================================== # -# SLURM HEADERS # -# ======================================================== # - -#SBATCH --job-name=training_llama2_7B -#SBATCH --output=logs/multinode-job-llama2-7b.%j.out -#SBATCH --time=3:00:00 -#SBATCH --nodes=2 -#SBATCH --ntasks-per-node=1 # setting this to 8 would launch 8 dockers on the single node with 8 GPU -#SBATCH --exclusive -#SBATCH --reservation=gpu-40_gpu-41_gpu-43_gpu-44_gpu-46_gpu-47_gpu-50_gpu-55_reservation - -# SLURM_NNODES -# Total number of nodes in the job's resource allocation. See SLURM_JOB_NUM_NODES. Included for backwards compatibility. - -# SLURM_NODEID -# ID of the nodes allocated. - -# SLURM_NODELIST -# List of nodes allocated to the job. See SLURM_JOB_NODELIST. Included for backwards compatibility. - -# srun echo $SLURM_NNODES -# echo $SLURM_LOCALID -# srun echo $SLURM_LOCALID -# srun -N $SLURM_JOB_NUM_NODES -n $SLURM_JOB_NUM_NODES echo $SLURM_LOCALID -# echo 'echo $SLURM_NODEID' > script.sh -# srun bash ./script.sh -# srun echo $SLURM_NODEID -# srun -N $SLURM_JOB_NUM_NODES -n $SLURM_JOB_NUM_NODES echo $SLURM_NODEID -# srun echo $SLURM_NODELIST -# srun echo $SLURM_JOB_NODELIST -# see https://gist.github.com/TengdaHan/1dd10d335c7ca6f13810fff41e809904 - -set -x -OUTPUT_DIR="$HOME/output" - -# Install required packages -echo ' -apt install iproute2 -y -apt install -y linux-headers-"$(uname -r)" libelf-dev -apt install -y gcc make libtool autoconf librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool libibverbs-dev rdma-core strace libibmad5 libibnetdisc5 ibverbs-providers libibumad-dev libibumad3 libibverbs1 libnl-3-dev libnl-route-3-dev -' > $OUTPUT_DIR/install_packages.sh - -# Environment variables -echo ' -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export NVTE_USE_HIPBLASLT=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_graph_level=0 --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=0 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NCCL_IB_TIMEOUT=20 -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -' > $OUTPUT_DIR/maxtext_env_7b.sh - -# Model Configuration -echo 'base_config: "base.yml" -run_name: "llama2_7b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 50 -model_name: "llama2-7b" -enable_checkpointing: False -attention: "cudnn_flash_te" -log_period: 100 - #inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 2 - #intra-node parallelism strategy -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -remat_policy: "minimal_flash" -use_iota_embed: True -scan_layers: False -async_checkpointing: False -logits_dot_in_fp32: False -profiler: "" -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -max_target_length: 4096 -per_device_batch_size: 10 -dataset_type: "synthetic" -enable_goodput_recording: False -monitor_goodput: False -shardy: False -' > $OUTPUT_DIR/llama2_7b_gpu.yml - - -srun hostname -# srun master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) -# srun export MASTER_ADDR=$master_addr -export MASTER_NAME=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) -export MASTER_ADDR=$(cat /etc/hosts | grep $MASTER_NAME | awk '{print $1}' ) -# MASTER_ADDR=$(cat /etc/hosts | grep gpu-14 | awk '{print $1}' ) -srun echo "MASTER_ADDR="$MASTER_ADDR - -# srun ping $MASTER_ADDR - -#If podman is available instead of docker, then you need this export otherwise -#comment the below line and uncomment the line after that -export docker=podman -#docker=docker - -export NNODES=$SLURM_NNODES -export JAX_COORDINATOR_IP=$MASTER_ADDR -export JAX_COORDINATOR_PORT=1234 - -#Change this to one of the IP interfaces used for communication -export NCCL_SOCKET_IFNAME=ens8np0 -echo $NCCL_SOCKET_IFNAME - -# For Mellanox NIC -export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_8,mlx5_9 -# For Broadcom Thor NIC, uncomment the line below and comment the line above -#export NCCL_IB_HCA=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 -echo $NCCL_IB_HCA - -srun --nodes=$SLURM_JOB_NUM_NODES --ntasks=$SLURM_JOB_NUM_NODES \ - --export=ALL \ - bash -c '\ - NODE_RANK=$SLURM_PROCID; \ - NNODES=$SLURM_JOB_NUM_NODES; \ - $docker run --rm --privileged --network host \ - --device /dev/dri --device /dev/kfd \ - --cap-add=IPC_LOCK \ - --volume /dev/infiniband:/dev/infiniband \ - -v $HOME:$HOME \ - -v $HOME/araina/data:/home/amd/data \ - --tmpfs /dev/shm:size=50G \ - --mount type=bind,source='${OUTPUT_DIR}',target=/workspace/maxtext/output \ - -e NNODES=$NNODES \ - -e NODE_RANK=$NODE_RANK \ - -e JAX_COORDINATOR_IP='${JAX_COORDINATOR_IP}' \ - -e JAX_COORDINATOR_PORT='${JAX_COORDINATOR_PORT}' \ - -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ - -w /workspace/maxtext \ - '${IMAGE}' \ - /bin/bash -c " - set -e - echo \"Running Llama-2-7b\" - echo '${IMAGE}' - echo \"Coordinator IP: \$JAX_COORDINATOR_IP\" - cp '${OUTPUT_DIR}'/install_packages.sh . - cp '${OUTPUT_DIR}'/maxtext_env_7b.sh . - mkdir -p configs - cp '${OUTPUT_DIR}'/llama2_7b_gpu.yml configs/llama2_7b_gpu.yml - source install_packages.sh - source maxtext_env_7b.sh - python -m maxtext.trainers.pre_train.train configs/llama2_7b_gpu.yml 2>&1 |& tee -a llama2_7b.synthetic.log - "' \ No newline at end of file diff --git a/scripts/jax-maxtext/gpu-rocm/llama3.3_70b.sh b/scripts/jax-maxtext/gpu-rocm/llama3.3_70b.sh deleted file mode 100644 index 36e9a6af..00000000 --- a/scripts/jax-maxtext/gpu-rocm/llama3.3_70b.sh +++ /dev/null @@ -1,86 +0,0 @@ -#!/bin/bash - -set -x -# Create output dir -OUTPUT_DIR="$HOME/output" -mkdir -p $OUTPUT_DIR - -# Environment variables -echo 'export XLA_FLAGS="--xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_graph_level=0 --xla_gpu_autotune_level=0 --xla_gpu_enable_reduce_scatter_combine_by_dim=false --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_all_reduce_combine_threshold_bytes=8589934592 --xla_gpu_all_gather_combine_threshold_bytes=137438953472 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export XLA_PYTHON_CLIENT_MEM_FRACTION=0.975 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH' > $OUTPUT_DIR/maxtext_env_70b.sh - - -# Model Configuration -echo 'base_config: "base.yml" -run_name: "llama3.3_70b_training" -hardware: "gpu" -steps: 30 -model_name: "llama3.3-70b" -enable_checkpointing: False -attention: "cudnn_flash_te" -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -dcn_pipeline_parallelism: 1 -dcn_tensor_parallelism: 1 -dcn_sequence_parallelism: 1 -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -ici_sequence_parallelism: 1 -ici_tensor_parallelism: 1 -ici_pipeline_parallelism: 1 - -remat_policy: 'full' -optimizer_memory_host_offload: False -param_scan_axis: 1 - -use_iota_embed: True -scan_layers: True - -profiler: "" - -async_checkpointing: False -logits_dot_in_fp32: False -megablox: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -per_device_batch_size: 7 -max_target_length: 8192 -hf_path: "parquet" -hf_train_files: "/hf_cache/hub/datasets--legacy-datasets--c4/snapshots/5abe0d085aa23dd9db2a6c1e86cfce4e4db6f0c3/en/partial-train/000*.parquet" -dataset_type: "hf" -tokenizer_path: "meta-llama/Llama-3.3-70B-Instruct"' > $OUTPUT_DIR/llama3.3_70b_gpu.yml - -#If podman is available instead of docker, then you need this export otherwise -#comment the below line and uncomment the line after that -#export docker=podman -docker=docker - - - -# get the test data -echo "For downloading data, we will mount \$HF_HOME to the docker and try to get llama tokenizer directly from there" -echo "Please set \$HF_HOME when calling this script, your HF_HOME is set as" -echo $HF_HOME -hf download legacy-datasets/c4 --include "*.parquet" --repo-type dataset --revision refs/convert/parquet - -$docker run --rm --privileged --network host --device /dev/dri --device /dev/kfd \ - --cap-add=IPC_LOCK --volume /dev/infiniband:/dev/infiniband \ - -v $HOME:$HOME -v $HOME/data:/home/amd/data -v $HF_HOME:/hf_cache -e HF_HOME=/hf_cache --tmpfs /dev/shm:size=50G \ - --mount type=bind,source=$OUTPUT_DIR,target=/workspace/maxtext/output \ - -w /workspace/maxtext $IMAGE /bin/bash -c " - set -e - echo \"Running Llama-3-70b\" - echo '${IMAGE}' - cp $OUTPUT_DIR/maxtext_env_70b.sh . - mkdir -p configs - cp $OUTPUT_DIR/llama3.3_70b_gpu.yml configs/llama3.3_70b_gpu.yml - source maxtext_env_70b.sh - python -m maxtext.trainers.pre_train.train configs/llama3.3_70b_gpu.yml base_output_directory=output 2>&1 |& tee -a llama3_70b.real.log - " - diff --git a/scripts/jax-maxtext/gpu-rocm/llama3_70b.sh b/scripts/jax-maxtext/gpu-rocm/llama3_70b.sh deleted file mode 100644 index 3db367b1..00000000 --- a/scripts/jax-maxtext/gpu-rocm/llama3_70b.sh +++ /dev/null @@ -1,86 +0,0 @@ -#!/bin/bash - -set -x -# Create output dir -OUTPUT_DIR="$HOME/output" -mkdir -p $OUTPUT_DIR - -# Environment variables -echo 'export XLA_FLAGS="--xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_graph_level=0 --xla_gpu_autotune_level=0 --xla_gpu_enable_reduce_scatter_combine_by_dim=false --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_all_reduce_combine_threshold_bytes=8589934592 --xla_gpu_all_gather_combine_threshold_bytes=137438953472 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export XLA_PYTHON_CLIENT_MEM_FRACTION=0.975 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH' > $OUTPUT_DIR/maxtext_env_70b.sh - - -# Model Configuration -echo 'base_config: "base.yml" -run_name: "llama3_70b_training" -hardware: "gpu" -steps: 30 -model_name: "llama3-70b" -enable_checkpointing: False -attention: "cudnn_flash_te" -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 -dcn_pipeline_parallelism: 1 -dcn_tensor_parallelism: 1 -dcn_sequence_parallelism: 1 -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -ici_sequence_parallelism: 1 -ici_tensor_parallelism: 1 -ici_pipeline_parallelism: 1 - -remat_policy: 'full' -optimizer_memory_host_offload: False -param_scan_axis: 1 - -use_iota_embed: True -scan_layers: True - -profiler: "" - -async_checkpointing: False -logits_dot_in_fp32: False -megablox: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -per_device_batch_size: 7 -max_target_length: 8192 -hf_path: "parquet" -hf_train_files: "/hf_cache/hub/datasets--legacy-datasets--c4/snapshots/5abe0d085aa23dd9db2a6c1e86cfce4e4db6f0c3/en/partial-train/000*.parquet" -dataset_type: "hf" -tokenizer_path: "meta-llama/Meta-Llama-3-70B"' > $OUTPUT_DIR/llama3_70b_gpu.yml - -#If podman is available instead of docker, then you need this export otherwise -#comment the below line and uncomment the line after that -#export docker=podman -docker=docker - - - -# get the test data -echo "For downloading data, we will mount \$HF_HOME to the docker and try to get llama tokenizer directly from there" -echo "Please set \$HF_HOME when calling this script, your HF_HOME is set as" -echo $HF_HOME -hf download legacy-datasets/c4 --include "*.parquet" --repo-type dataset --revision refs/convert/parquet - -$docker run --rm --privileged --network host --device /dev/dri --device /dev/kfd \ - --cap-add=IPC_LOCK --volume /dev/infiniband:/dev/infiniband \ - -v $HOME:$HOME -v $HOME/data:/home/amd/data -v $HF_HOME:/hf_cache -e HF_HOME=/hf_cache --tmpfs /dev/shm:size=50G \ - --mount type=bind,source=$OUTPUT_DIR,target=/workspace/maxtext/output \ - -w /workspace/maxtext $IMAGE /bin/bash -c " - set -e - echo \"Running Llama-3-70b\" - echo '${IMAGE}' - cp $OUTPUT_DIR/maxtext_env_70b.sh . - mkdir -p configs - cp $OUTPUT_DIR/llama3_70b_gpu.yml configs/llama3_70b_gpu.yml - source maxtext_env_70b.sh - python -m maxtext.trainers.pre_train.train configs/llama3_70b_gpu.yml base_output_directory=output 2>&1 |& tee -a llama3_70b.real.log - " - diff --git a/scripts/jax-maxtext/gpu-rocm/llama3_70b_multinode.sh b/scripts/jax-maxtext/gpu-rocm/llama3_70b_multinode.sh deleted file mode 100644 index 1736c62b..00000000 --- a/scripts/jax-maxtext/gpu-rocm/llama3_70b_multinode.sh +++ /dev/null @@ -1,177 +0,0 @@ -#!/bin/bash - -# ======================================================== # -# SLURM HEADERS # -# ======================================================== # - -#SBATCH --job-name=training_llama3_70B -#SBATCH --output=logs/multinode-job-llama3-70b.%j.out -#SBATCH --time=3:00:00 -#SBATCH --nodes=2 -#SBATCH --ntasks-per-node=1 # setting this to 8 would launch 8 dockers on the single node with 8 GPU -#SBATCH --exclusive -#SBATCH --reservation=gpu-40_gpu-41_gpu-43_gpu-44_gpu-46_gpu-47_gpu-50_gpu-55_reservation - -# SLURM_NNODES -# Total number of nodes in the job's resource allocation. See SLURM_JOB_NUM_NODES. Included for backwards compatibility. - -# SLURM_NODEID -# ID of the nodes allocated. - -# SLURM_NODELIST -# List of nodes allocated to the job. See SLURM_JOB_NODELIST. Included for backwards compatibility. - -# srun echo $SLURM_NNODES -# echo $SLURM_LOCALID -# srun echo $SLURM_LOCALID -# srun -N $SLURM_JOB_NUM_NODES -n $SLURM_JOB_NUM_NODES echo $SLURM_LOCALID -# echo 'echo $SLURM_NODEID' > script.sh -# srun bash ./script.sh -# srun echo $SLURM_NODEID -# srun -N $SLURM_JOB_NUM_NODES -n $SLURM_JOB_NUM_NODES echo $SLURM_NODEID -# srun echo $SLURM_NODELIST -# srun echo $SLURM_JOB_NODELIST -# see https://gist.github.com/TengdaHan/1dd10d335c7ca6f13810fff41e809904 - -set -x -OUTPUT_DIR="$HOME/output" - -# Install required packages -echo ' -apt install iproute2 -y -apt install rdma-core -y -apt install -y linux-headers-"$(uname -r)" libelf-dev -apt install -y gcc make libtool autoconf librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool libibverbs-dev rdma-core strace libibmad5 libibnetdisc5 ibverbs-providers libibumad-dev libibumad3 libibverbs1 libnl-3-dev libnl-route-3-dev -' > $OUTPUT_DIR/install_packages.sh - - -# Environment variables -echo ' -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export NVTE_USE_HIPBLASLT=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_graph_level=0 --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=0 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NCCL_IB_TIMEOUT=20 -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -' > $OUTPUT_DIR/maxtext_env_70b.sh - - -# Model Configuration -echo 'base_config: "base.yml" -run_name: "llama3_70b_training" -hardware: "gpu" -steps: 30 -model_name: "llama3-70b" -enable_checkpointing: False -attention: "cudnn_flash_te" -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 2 -dcn_pipeline_parallelism: 1 -dcn_tensor_parallelism: 1 -dcn_sequence_parallelism: 1 -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -ici_sequence_parallelism: 1 -ici_tensor_parallelism: 1 -ici_pipeline_parallelism: 1 - -remat_policy: 'full' -optimizer_memory_host_offload: False -param_scan_axis: 1 - -use_iota_embed: True -scan_layers: True - -profiler: "" - -async_checkpointing: False -logits_dot_in_fp32: False -megablox: False -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -per_device_batch_size: 7 -max_target_length: 8192 -dataset_type: "synthetic" -enable_goodput_recording: False -monitor_goodput: False -shardy: False -' > $OUTPUT_DIR/llama3_70b_gpu.yml - - -srun hostname -# srun master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) -# srun export MASTER_ADDR=$master_addr -export MASTER_NAME=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) -export MASTER_ADDR=$(cat /etc/hosts | grep $MASTER_NAME | awk '{print $1}' ) -# MASTER_ADDR=$(cat /etc/hosts | grep gpu-14 | awk '{print $1}' ) -srun echo "MASTER_ADDR="$MASTER_ADDR - -# srun ping $MASTER_ADDR - -#If podman is available instead of docker, then you need this export otherwise -#comment the below line and uncomment the line after that -export docker=podman -#docker=docker - -export NNODES=$SLURM_NNODES -export JAX_COORDINATOR_IP=$MASTER_ADDR -export JAX_COORDINATOR_PORT=1234 - -#Change this to one of the IP interfaces used for communication -export NCCL_SOCKET_IFNAME=ens8np0 -echo $NCCL_SOCKET_IFNAME - -# For Mellanox NIC -export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_8,mlx5_9 -# For Broadcom Thor NIC, uncomment the line below and comment the line above -#export NCCL_IB_HCA=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 -echo $NCCL_IB_HCA - -srun --nodes=$SLURM_JOB_NUM_NODES --ntasks=$SLURM_JOB_NUM_NODES \ - --export=ALL \ - bash -c '\ - NODE_RANK=$SLURM_PROCID; \ - NNODES=$SLURM_JOB_NUM_NODES; \ - $docker run --rm --privileged --network host \ - --device /dev/dri --device /dev/kfd \ - --cap-add=IPC_LOCK \ - --volume /dev/infiniband:/dev/infiniband \ - -v $HOME:$HOME \ - -v $HOME/araina/data:/home/amd/data \ - --tmpfs /dev/shm:size=50G \ - --mount type=bind,source='${OUTPUT_DIR}',target=/workspace/maxtext/output \ - -e NNODES=$NNODES \ - -e NODE_RANK=$NODE_RANK \ - -e JAX_COORDINATOR_IP='${JAX_COORDINATOR_IP}' \ - -e JAX_COORDINATOR_PORT='${JAX_COORDINATOR_PORT}' \ - -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ - -w /workspace/maxtext \ - '${IMAGE}' \ - /bin/bash -c " - set -e - echo \"Running Llama-3-70b\" - echo '${IMAGE}' - echo \"Coordinator IP: \$JAX_COORDINATOR_IP\" - cp '${OUTPUT_DIR}'/install_packages.sh . - cp '${OUTPUT_DIR}'/maxtext_env_70b.sh . - mkdir -p configs - cp '${OUTPUT_DIR}'/llama3_70b_gpu.yml configs/llama3_70b_gpu.yml - source install_packages.sh - source maxtext_env_70b.sh - python -m maxtext.trainers.pre_train.train configs/llama3_70b_gpu.yml 2>&1 |& tee -a llama3_70b.synthetic.log - "' \ No newline at end of file diff --git a/scripts/jax-maxtext/gpu-rocm/llama3_70b_multinode_metal.sh b/scripts/jax-maxtext/gpu-rocm/llama3_70b_multinode_metal.sh deleted file mode 100644 index 08dce2b7..00000000 --- a/scripts/jax-maxtext/gpu-rocm/llama3_70b_multinode_metal.sh +++ /dev/null @@ -1,204 +0,0 @@ -#!/bin/bash -# -# This is a version of training script for bare metal runs, -# without slurm or k8s, all you need is ssh access to the gpu servers -# -# This script is to be invoked like this: -# -# for x in $(cat host_ip_file); \ -# do \ -# ssh root@$x "docker exec jax_train /workspace/maxtext/output/llama3_70b_multinode_metal.sh" & \ -# done -# -# on all the gpu servers, run the jax container and create -# jax_train like this: -# -# -# docker run -d -it --name jax_train --network host --ipc host \ -# --privileged --shm-size 64G --tmpfs /dev/shm:size=200G \ -# --cap-add IPC_LOCK --cap-add SYS_PTRACE --security-opt seccomp=unconfined \ -# --group-add video --device /dev/kfd --device /dev/dri --device /dev/infiniband \ -# --volume /dev/infiniband:/dev/infiniband -v /etc/libibverbs.d:/etc/libibverbs.d:ro \ -# -v /usr/lib/x86_64-linux-gnu:/usr/lib/x86_64-linux-gnu:ro \ -# -v /usr/local/lib:/usr/local/lib:ro -v /root/.ssh:/root/.ssh:ro \ -# -v /mnt/testvfs/jax/output:/workspace/maxtext/output \ -# -w /workspace/maxtext rocm/jax-training:maxtext-v25.5 bash -# -# this script llama3_70b_multinode_metal.sh needs to be placed in output directory of -# the docker, furthermore, output directory should be NFS mounted which is mounted -# across all the containers, that way runjax script is the same across -# all the containers. based on the host_info array and matching the $(hostname) -# it will compute its own node rank -# - - -# creating the host_info array (create a plain text file with ip and name) -# and pass it to following one-line bash script to generate host_info array -# -# rank=0; \ -# while read -r ip name; \ -# do \ -# printf "host_info[\"%s\"]=\"%s %d\"\n" $name $ip $rank \ -# rank=$((rank+1)) \ -# done < /tmp/v16.names -# - -declare -A host_info - -# host_info array is indexed by hostname, and has a tuple of (ip, rank) as value -host_info["node1"]="10.10.0.1 0" -host_info["node2"]="10.10.0.2 1" -host_info["node3"]="10.10.0.3 2" -host_info["node4"]="10.10.0.4 3" -host_info["node5"]="10.10.0.5 4" -host_info["node6"]="10.10.0.6 5" -host_info["node7"]="10.10.0.7 6" -host_info["node8"]="10.10.0.8 7" -host_info["node9"]="10.10.0.9 8" -host_info["node10"]="10.10.0.10 9" -host_info["node11"]="10.10.0.11 10" -host_info["node12"]="10.10.0.12 11" -host_info["node13"]="10.10.0.13 12" -host_info["node14"]="10.10.0.14 13" -host_info["node15"]="10.10.0.15 14" -host_info["node16"]="10.10.0.16 15" - - -export NNODES=2 - -# Get the current hostname -current_hostname=$(hostname) - -if [[ -n "${host_info[$current_hostname]}" ]]; then - read -r ip rank <<< "${host_info[$current_hostname]}" - echo "Current Hostname: $current_hostname" - echo "IP Address: $ip" - echo "Rank: $rank" - NODE_RANK=$rank -else - echo "Hostname '$current_hostname' not found in the host dictionary." - exit 1 -fi - -# change this to MASTER_ADDR -export JAX_COORDINATOR_IP='10.10.0.1' -export JAX_COORDINATOR_PORT=12345 -export JAX_DISTRIBUTED_INITIALIZATION_TIMEOUT_SECONDS=1800 -export JAX_PROCESS_COUNT=${NNODES} -export JAX_PROCESS_INDEX=${rank} -export NODE_RANK=$rank - -set -e -echo "Starting node $NODE_RANK of $NNODES" -echo "Coordinator IP: $JAX_COORDINATOR_IP" - -apt update -apt install iproute2 -y -apt install rdma-core -y -apt install apt-utils -y -apt install -y linux-headers-"$(uname -r)" libelf-dev -apt install -y gcc make libtool autoconf librdmacm-dev rdmacm-utils \ - infiniband-diags ibverbs-utils perftest ethtool libibverbs-dev \ - rdma-core strace libibmad5 libibnetdisc5 ibverbs-providers \ - libibumad-dev libibumad3 libibverbs1 libnl-3-dev libnl-route-3-dev - - -mkdir -p /workspace/maxtext/output/configs -cat > /workspace/maxtext/output/configs/llama3_70b_gpu.yml <&1 |& tee -a llama3_70b.real.log" - -# Run the training -python -m maxtext.trainers.pre_train.train /workspace/maxtext/output/configs/llama3_70b_gpu.yml base_output_directory=${BASE_OUTPUT_DIRECTORY} 2>&1 | tee >(grep ".") > ${OUT_FILE_NAME} diff --git a/scripts/jax-maxtext/gpu-rocm/llama3_8b.sh b/scripts/jax-maxtext/gpu-rocm/llama3_8b.sh deleted file mode 100644 index 617b6fca..00000000 --- a/scripts/jax-maxtext/gpu-rocm/llama3_8b.sh +++ /dev/null @@ -1,77 +0,0 @@ -#!/bin/bash - -set -x -# Create output dir -OUTPUT_DIR="$HOME/output" -mkdir -p $OUTPUT_DIR - -# Environment variables -echo 'export XLA_FLAGS="--xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_graph_level=0 --xla_gpu_autotune_level=0 --xla_gpu_enable_latency_hiding_scheduler=TRUE --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_all_gather_combine_by_dim=FALSE --xla_gpu_memory_limit_slop_factor=95" -export XLA_PYTHON_CLIENT_MEM_FRACTION=0.967 -export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH' > $OUTPUT_DIR/maxtext_env_8b.sh - -# Model Configuration -echo 'base_config: "base.yml" -run_name: "llama3_8b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 50 -model_name: "llama3-8b" -enable_checkpointing: False -attention: "cudnn_flash_te" -log_period: 100 - #inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 1 - #intra-node parallelism strategy -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -remat_policy: "minimal_flash" -use_iota_embed: True -scan_layers: False -async_checkpointing: False -logits_dot_in_fp32: False -profiler: "" -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -max_target_length: 8192 -per_device_batch_size: 4 -hf_path: "parquet" -hf_train_files: "/hf_cache/hub/datasets--legacy-datasets--c4/snapshots/5abe0d085aa23dd9db2a6c1e86cfce4e4db6f0c3/en/partial-train/000*.parquet" -dataset_type: "hf" -tokenizer_path: "meta-llama/Meta-Llama-3-8B"' > $OUTPUT_DIR/llama3_8b_gpu.yml - -#If podman is available instead of docker, then you need this export otherwise -#comment the below line and uncomment the line after that -#export docker=podman -docker=docker - - - -# get the test data -echo "For downloading data, we will mount \$HF_HOME to the docker and try to get llama tokenizer directly from there" -echo "Please set \$HF_HOME when calling this script, your HF_HOME is set as" -echo $HF_HOME -hf download legacy-datasets/c4 --include "*.parquet" --repo-type dataset --revision refs/convert/parquet - - -$docker run --rm --privileged --network host --device /dev/dri --device /dev/kfd \ - --cap-add=IPC_LOCK --volume /dev/infiniband:/dev/infiniband \ - -v $HOME:$HOME -v $HOME/data:/home/amd/data -v $HF_HOME:/hf_cache -e HF_HOME=/hf_cache --tmpfs /dev/shm:size=50G \ - --mount type=bind,source=$OUTPUT_DIR,target=/workspace/maxtext/output \ - -w /workspace/maxtext $IMAGE /bin/bash -c " - set -e - echo \"Running Llama-3-8b\" - echo '${IMAGE}' - cp $OUTPUT_DIR/maxtext_env_8b.sh . - mkdir -p configs - cp $OUTPUT_DIR/llama3_8b_gpu.yml configs/llama3_8b_gpu.yml - source maxtext_env_8b.sh - python -m maxtext.trainers.pre_train.train configs/llama3_8b_gpu.yml base_output_directory=output 2>&1 |& tee -a llama3_8b.real.log - " - diff --git a/scripts/jax-maxtext/gpu-rocm/llama3_8b_multinode.sh b/scripts/jax-maxtext/gpu-rocm/llama3_8b_multinode.sh deleted file mode 100644 index f98619df..00000000 --- a/scripts/jax-maxtext/gpu-rocm/llama3_8b_multinode.sh +++ /dev/null @@ -1,165 +0,0 @@ -#!/bin/bash - -# ======================================================== # -# SLURM HEADERS # -# ======================================================== # - -#SBATCH --job-name=training_llama3_8B -#SBATCH --output=logs/multinode-job-llama3-8b.%j.out -#SBATCH --time=3:00:00 -#SBATCH --nodes=2 -#SBATCH --ntasks-per-node=1 # setting this to 8 would launch 8 dockers on the single node with 8 GPU -#SBATCH --exclusive -#SBATCH --reservation=gpu-40_gpu-41_gpu-43_gpu-44_gpu-46_gpu-47_gpu-50_gpu-55_reservation - -# SLURM_NNODES -# Total number of nodes in the job's resource allocation. See SLURM_JOB_NUM_NODES. Included for backwards compatibility. - -# SLURM_NODEID -# ID of the nodes allocated. - -# SLURM_NODELIST -# List of nodes allocated to the job. See SLURM_JOB_NODELIST. Included for backwards compatibility. - -# srun echo $SLURM_NNODES -# echo $SLURM_LOCALID -# srun echo $SLURM_LOCALID -# srun -N $SLURM_JOB_NUM_NODES -n $SLURM_JOB_NUM_NODES echo $SLURM_LOCALID -# echo 'echo $SLURM_NODEID' > script.sh -# srun bash ./script.sh -# srun echo $SLURM_NODEID -# srun -N $SLURM_JOB_NUM_NODES -n $SLURM_JOB_NUM_NODES echo $SLURM_NODEID -# srun echo $SLURM_NODELIST -# srun echo $SLURM_JOB_NODELIST -# see https://gist.github.com/TengdaHan/1dd10d335c7ca6f13810fff41e809904 - -set -x -OUTPUT_DIR="$HOME/output" - -# Install required packages -echo ' -apt install iproute2 -y -apt install -y linux-headers-"$(uname -r)" libelf-dev -apt install -y gcc make libtool autoconf librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool libibverbs-dev rdma-core strace libibmad5 libibnetdisc5 ibverbs-providers libibumad-dev libibumad3 libibverbs1 libnl-3-dev libnl-route-3-dev -' > $OUTPUT_DIR/install_packages.sh - -# Environment variables -echo ' -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export XLA_PYTHON_CLIENT_MEM_FRACTION=.97 -export NVTE_USE_HIPBLASLT=1 -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_graph_level=0 --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=8589934592 --xla_gpu_enable_triton_gemm=False --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=0 --xla_gpu_enable_all_gather_combine_by_dim=FALSE" -export GPU_MAX_HW_QUEUES=2 -export HIP_FORCE_DEV_KERNARG=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export NVTE_FUSED_ATTN=1 -export NCCL_DEBUG=VERSION -export NCCL_IB_TIMEOUT=20 -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -' > $OUTPUT_DIR/maxtext_env_8b.sh - -# Model Configuration -echo 'base_config: "base.yml" -run_name: "llama3_8b_training" -base_output_directory: "./" -hardware: "gpu" -steps: 50 -model_name: "llama3-8b" -enable_checkpointing: False -attention: "cudnn_flash_te" -log_period: 100 - #inter-node parallelism strategy -dcn_data_parallelism: -1 -dcn_fsdp_parallelism: 2 - #intra-node parallelism strategy -ici_fsdp_parallelism: 8 -ici_data_parallelism: 1 -remat_policy: "minimal_flash" -use_iota_embed: True -scan_layers: False -async_checkpointing: False -logits_dot_in_fp32: False -profiler: "" -dtype: "bfloat16" -quantization: "" -quantize_kvcache: False -kv_quant_axis: "heads_and_dkv" -kv_quant_dtype: "int8" -weight_dtype: bfloat16 -checkpoint_is_quantized: False # Set to True if reading from a saved aqt quantized checkpoint -max_target_length: 8192 -per_device_batch_size: 5 -dataset_type: "synthetic" -enable_goodput_recording: False -monitor_goodput: False -shardy: False -' > $OUTPUT_DIR/llama3_8b_gpu.yml - - -srun hostname -# srun master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) -# srun export MASTER_ADDR=$master_addr -export MASTER_NAME=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) -export MASTER_ADDR=$(cat /etc/hosts | grep $MASTER_NAME | awk '{print $1}' ) -# MASTER_ADDR=$(cat /etc/hosts | grep gpu-14 | awk '{print $1}' ) -srun echo "MASTER_ADDR="$MASTER_ADDR - -# srun ping $MASTER_ADDR - -#If podman is available instead of docker, then you need this export otherwise -#comment the below line and uncomment the line after that -export docker=podman -#docker=docker - -export NNODES=$SLURM_NNODES -export JAX_COORDINATOR_IP=$MASTER_ADDR -export JAX_COORDINATOR_PORT=1234 - -#Change this to one of the IP interfaces used for communication -export NCCL_SOCKET_IFNAME=ens8np0 -echo $NCCL_SOCKET_IFNAME - -# For Mellanox NIC -export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_8,mlx5_9 -# For Broadcom Thor NIC, uncomment the line below and comment the line above -#export NCCL_IB_HCA=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 -echo $NCCL_IB_HCA - -srun --nodes=$SLURM_JOB_NUM_NODES --ntasks=$SLURM_JOB_NUM_NODES \ - --export=ALL \ - bash -c '\ - NODE_RANK=$SLURM_PROCID; \ - NNODES=$SLURM_JOB_NUM_NODES; \ - $docker run --rm --privileged --network host \ - --device /dev/dri --device /dev/kfd \ - --cap-add=IPC_LOCK \ - --volume /dev/infiniband:/dev/infiniband \ - -v $HOME:$HOME \ - -v $HOME/araina/data:/home/amd/data \ - --tmpfs /dev/shm:size=50G \ - --mount type=bind,source='${OUTPUT_DIR}',target=/workspace/maxtext/output \ - -e NNODES=$NNODES \ - -e NODE_RANK=$NODE_RANK \ - -e JAX_COORDINATOR_IP='${JAX_COORDINATOR_IP}' \ - -e JAX_COORDINATOR_PORT='${JAX_COORDINATOR_PORT}' \ - -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ - -w /workspace/maxtext \ - '${IMAGE}' \ - /bin/bash -c " - set -e - echo \"Running Llama-3-8b\" - echo '${IMAGE}' - echo \"Coordinator IP: \$JAX_COORDINATOR_IP\" - cp '${OUTPUT_DIR}'/install_packages.sh . - cp '${OUTPUT_DIR}'/maxtext_env_8b.sh . - mkdir -p configs - cp '${OUTPUT_DIR}'/llama3_8b_gpu.yml configs/llama3_8b_gpu.yml - source install_packages.sh - source maxtext_env_8b.sh - python -m maxtext.trainers.pre_train.train configs/llama3_8b_gpu.yml 2>&1 |& tee -a llama3_8b.synthetic.log - "' \ No newline at end of file diff --git a/scripts/jax-maxtext/gpu-rocm/readme.md b/scripts/jax-maxtext/gpu-rocm/readme.md deleted file mode 100644 index 531c27d7..00000000 --- a/scripts/jax-maxtext/gpu-rocm/readme.md +++ /dev/null @@ -1,39 +0,0 @@ -# ROCM benchmarking -Scripts under this folder are used to benchmark rocm docker for maxtext-jax with different models. They will launch docker and run the benchmark. **Please run them on host instead of inside any docker** - -All the scripts without the _multinode suffix can be launched on single node like: -``` -IMAGE="rocm/jax-maxtext-training:xxx" HF_HOME=/home/amd-shared-home/.cache/huggingface bash ./deepseek_v2_16b.sh -``` -Please adjust the $HF_HOME and $IMAGE to your environment. - -HF_HOME is where huggingface_hub will store local data, please refer to [Huggingface cli Document](https://huggingface.co/docs/huggingface_hub/main/en/guides/cli#hf-download) on how to download the data. - -For the multinode one, they were written for AMD internal cluster, and will need to be adjusted for other cluster setting. They can be launched via slurm like: -``` -sbatch -N llama3_70b_multinode.sh -``` -## Tokenizer download -For single node scripts, they will use the $HF_HOME folder on the host. The script will mount the host HF folder to the docker. Please make sure that the data already got downloaded to $HF_HOME folder / your HF token is saved in the config file before running the script. The tokenizer of corresponding models will be used for the training. - -| Model tag | Huggingface webpage | -|---|---| -| meta-llama/Llama-2-7b | https://huggingface.co/meta-llama/Llama-2-7b | -| meta-llama/Llama-2-70b | https://huggingface.co/meta-llama/Llama-2-70b | -| meta-llama/Meta-Llama-3-8B | https://huggingface.co/meta-llama/Meta-Llama-3-8B | -| meta-llama/Meta-Llama-3-70B | https://huggingface.co/meta-llama/Meta-Llama-3-70B | - -Example command for downloading the llama model tokenizer -``` -hf login --token=hf_xxxx -hf download meta-llama/Llama-2-7b --include "**token**" -hf download meta-llama/Llama-2-70b --include "**token**" -hf download meta-llama/Meta-Llama-3-8B --include "**token**" -hf download meta-llama/Meta-Llama-3-70B --include "**token**" -``` -## Dataset download -Please run this command for downloading the c4 dataset -``` -hf download legacy-datasets/c4 --include "*.parquet" --repo-type dataset --revision refs/convert/parquet -``` -Please check this path and see if data got downloaded to $HF_HOME/hub/datasets--legacy-datasets--c4/snapshots/5abe0d085aa23dd9db2a6c1e86cfce4e4db6f0c3/en/partial-train/000*.parquet \ No newline at end of file diff --git a/scripts/jax-maxtext/jax-maxtext_benchmark_report.py b/scripts/jax-maxtext/jax-maxtext_benchmark_report.py deleted file mode 100644 index 606bfebf..00000000 --- a/scripts/jax-maxtext/jax-maxtext_benchmark_report.py +++ /dev/null @@ -1,100 +0,0 @@ -############################################################################### -# -# MIT License -# -# Copyright (c) 2024 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# -import pandas as pd -import numpy as np -import argparse -import csv -import re - -# parse arguments -parser = argparse.ArgumentParser(description='Convert pytorch train output format to MAD csv output format') -parser.add_argument("--mode", - type=str, - help="pretrain or finetune") -parser.add_argument("--model", - type=str, - help="model name") -parser.add_argument("--quantization", - type=str, - default="bf16", - help="quantization type, e.g. bf16, nanoo_fp8, etc.") -parser.add_argument("--input", - type=str, - help="path to input file") -parser.add_argument("--output", - type=str, - help="path to output file") -parser.add_argument("--batch_size", - type=str, - help="batch size") -parser.add_argument("--seq_len", - type=str, - help="sequence length") -parser.add_argument("--device", - type=str, - help="device name") -parser.add_argument("--num_gpus", - type=str, - help="number of GPUs") - -# read arguments -args = parser.parse_args() -input_file = args.input -output_file = args.output -quantization = args.quantization -print("Input file path: ", input_file) -print("Output file path: ", output_file) -print("Quantization: ", quantization) - -def find_match(file, search_string, num_iters): - with open(file, 'r') as file: - content = file.read() - pattern = fr"{re.escape(search_string)}\s*(\d+\.\d+|\d+)" - matches = re.findall(pattern, content) - perf_nums = [float(num) for num in matches][-num_iters:] - avg = np.average(perf_nums) - return str("{:.2f}".format(avg)) - -if args.model == "Llama-3.1-8B" or args.model == "Llama-3.1-70B" or \ - args.model == "Llama-3.3-70B" or \ - args.model == "Llama-2-7B" or args.model == "Llama-2-70B" or \ - args.model == "DeepSeek-V2-lite" or args.model == "Mixtral-8x7B" or\ - args.model == "Qwen3-14B" or\ - args.model == "Qwen3-30B-A3B": - tok_per_s_per_gpu = find_match(input_file, "Tokens/s/device:", 10) - TFLOPS_per_gpu = find_match(input_file, "TFLOP/s/device:", 10) - data = [ - {'model': args.model, 'performance': tok_per_s_per_gpu, 'metric': 'tok_per_s_per_gpu', 'mode': args.mode, 'precision': args.quantization, 'batch_size': args.batch_size, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus}, - {'model': args.model, 'performance': TFLOPS_per_gpu, 'metric': 'TFLOPS_per_gpu', 'mode': args.mode, 'precision': args.quantization, 'batch_size': args.batch_size, 'seq_len': args.seq_len, 'device': args.device, 'num_gpus': args.num_gpus} - ] - -with open(output_file, mode='w', newline='') as file: - print("Preparing to write performance data...") - print("Data: ", data) - writer = csv.DictWriter(file, fieldnames=['model','performance','metric','mode','precision','batch_size','seq_len','device','num_gpus']) - writer.writeheader() - writer.writerows(data) - print("Completed writing to output file") diff --git a/scripts/jax-maxtext/jax-maxtext_benchmark_report.sh b/scripts/jax-maxtext/jax-maxtext_benchmark_report.sh deleted file mode 100755 index 944b4b1d..00000000 --- a/scripts/jax-maxtext/jax-maxtext_benchmark_report.sh +++ /dev/null @@ -1,167 +0,0 @@ -#!/bin/bash -############################################################################### -# -# MIT License -# -# Copyright (c) 2024-2025 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# -## Usage: -#./jax-maxtext_benchmark_report.sh -m $model_name -q $quantization - - -# Parse command-line arguments -while getopts "m:q:" opt; do - case "$opt" in - m) MODEL_REPO="$OPTARG" ;; - q) QUANTIZATION="$OPTARG" ;; - *) usage ;; - esac -done - -# Set default values for additional parameters -MODE="pretrain" -NNODES=1 # default to 1 node -GPUS_PER_NODE=8 # default to 8 GPUs per node -NUM_GPUS=$((NNODES*GPUS_PER_NODE)) - -echo "=hyper params start=" -echo $MODEL_REPO -echo $QUANTIZATION -echo "=hyper params end=" - -if [ -z "$QUANTIZATION" ]; then - PERF_LOG="$(pwd)/../perf_${MODEL_REPO}.csv" -else - PERF_LOG="$(pwd)/../perf_${MODEL_REPO}_${QUANTIZATION}.csv" -fi -perf_script="$(pwd)/jax-maxtext_benchmark_report.py" - -# Run rocminfo and grep for "AMD Instinct" -DEVICE=$(/opt/rocm/bin/rocminfo | grep "AMD Instinct" | head -n1 | awk '{print $5}') -if [ -z "$DEVICE" ]; then - ARCH=$(/opt/rocm/bin/rocminfo | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]') - case "$ARCH" in - "gfx942") DEVICE="MI300X" ;; - "gfx950") DEVICE="MI355X" ;; - *) DEVICE="" ;; - esac -fi -echo "GPU DEVICE name: $DEVICE" - -MAXTEXT="/workspace/maxtext" -MAXTEXT_DIR="/workspace/maxtext/src/maxtext" -ENV_SCRIPT_DIR="$(pwd)/env_scripts" - -cd $MAXTEXT - - -execute_training(){ - gpu_architecture=$(rocminfo | grep -o -m 1 'gfx.*' | xargs ) - env_file=$ENV_SCRIPT_DIR/$1 - if test -e $ENV_SCRIPT_DIR/$gpu_architecture"_"$1; then - env_file=$ENV_SCRIPT_DIR/$gpu_architecture"_"$1 - fi - config_file=$ENV_SCRIPT_DIR/$2 - if test -e $ENV_SCRIPT_DIR/$gpu_architecture"_"$2; then - config_file=$ENV_SCRIPT_DIR/$gpu_architecture"_"$2 - fi - - # output for logging - echo "Using env file:" - echo $env_file - cat $env_file - - echo "Using yaml config file:" - echo $config_file - cat $config_file - - yaml() { - python3 -c "import yaml;print(yaml.safe_load(open('$1'))$2)" - } - - per_device_batch_size=$(yaml $config_file "['per_device_batch_size']") - max_target_length=$(yaml $config_file "['max_target_length']") - echo $per_device_batch_size - echo $max_target_length - - # execute - source $env_file - python -m maxtext.trainers.pre_train.train $config_file \ - quantization=$3 2>&1 |& tee -a $2.log - if [ -z "$3" ]; then - python3 $perf_script --model $MODEL_REPO --input $MAXTEXT/$2.log --output $PERF_LOG --mode $MODE --quantization bf16 --batch_size $per_device_batch_size --seq_len $max_target_length --device $DEVICE --num_gpus $NUM_GPUS - else - python3 $perf_script --model $MODEL_REPO --input $MAXTEXT/$2.log --output $PERF_LOG --mode $MODE --quantization $3 --batch_size $per_device_batch_size --seq_len $max_target_length --device $DEVICE --num_gpus $NUM_GPUS - fi - -} - - -if [[ "$MODEL_REPO" == "Llama-2-7B" ]]; then - echo "[INFO] LLAMA 2 7B TRAINING with following parameters" - echo " QUANTIZATION: $QUANTIZATION" - execute_training llama2_7b_env.sh llama2_7b.yml $QUANTIZATION - -elif [[ "$MODEL_REPO" == "Llama-2-70B" ]]; then - echo "[INFO] LLAMA 2 70B TRAINING with following parameters" - echo " QUANTIZATION: $QUANTIZATION" - execute_training llama2_70b_env.sh llama2_70b.yml $QUANTIZATION - -elif [[ "$MODEL_REPO" == "Llama-3.1-8B" ]]; then - echo "[INFO] LLAMA 3.1 8B TRAINING with following parameters" - echo " QUANTIZATION: $QUANTIZATION" - execute_training llama3_8b_env.sh llama3_8b.yml $QUANTIZATION - -elif [[ "$MODEL_REPO" == "Llama-3.1-70B" ]]; then - echo "[INFO] LLAMA 3.1 70B TRAINING with following parameters" - echo " QUANTIZATION: $QUANTIZATION" - execute_training llama3_70b_env.sh llama3_70b.yml $QUANTIZATION - -elif [[ "$MODEL_REPO" == "Llama-3.3-70B" ]]; then - echo "[INFO] LLAMA 3.3 70B TRAINING with following parameters" - echo " QUANTIZATION: $QUANTIZATION" - execute_training llama3.3_70b_env.sh llama3.3_70b.yml $QUANTIZATION - -elif [[ "$MODEL_REPO" == "DeepSeek-V2-lite" ]]; then - echo "[INFO] DEEPSEEK V2 LITE TRAINING with following parameters" - echo " QUANTIZATION: $QUANTIZATION" - execute_training deepseek2_env_16b.sh deepseek2_16b.yml $QUANTIZATION - -elif [[ "$MODEL_REPO" == "Mixtral-8x7B" ]]; then - echo "[INFO] MIXTRAL-8x7B TRAINING with following parameters" - echo " QUANTIZATION: $QUANTIZATION" - execute_training mixtral_8x7b_env.sh mixtral_8x7b.yml $QUANTIZATION - -elif [[ "$MODEL_REPO" == "Qwen3-14B" ]]; then - echo "[INFO] QWEN3-14B TRAINING with following parameters" - echo " QUANTIZATION: $QUANTIZATION" - execute_training qwen3_14b_env.sh qwen3_14b.yml $QUANTIZATION - -elif [[ "$MODEL_REPO" == "Qwen3-30B-A3B" ]]; then - echo "[INFO] QWEN3-30B-A3B TRAINING with following parameters" - echo " QUANTIZATION: $QUANTIZATION" - execute_training qwen3_30b_a3b_env.sh qwen3_30b_a3b.yml $QUANTIZATION - -else - echo "Error: Unsupported training mode." - exit 1 -fi diff --git a/scripts/jax-maxtext/jax-maxtext_benchmark_setup.sh b/scripts/jax-maxtext/jax-maxtext_benchmark_setup.sh deleted file mode 100755 index 77536a4d..00000000 --- a/scripts/jax-maxtext/jax-maxtext_benchmark_setup.sh +++ /dev/null @@ -1,90 +0,0 @@ -#!/bin/bash -############################################################################### -# -# MIT License -# -# Copyright (c) 2024-2025 Advanced Micro Devices, Inc. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -# -################################################################################# -## Usage: -#./jax-maxtext_benchmark_report.sh -m $model_name - - -# Parse command-line arguments -while getopts "m:" opt; do - case "$opt" in - m) MODEL_REPO="$OPTARG" ;; - *) usage ;; - esac -done - -echo "=hyper params start=" -echo $MODEL_REPO -echo "=hyper params end=" - - -cd $MAXTEXT -echo "Building dependencies for $MODEL_REPO" - -set -x -export HF_HOME=/hf_cache -mkdir /hf_cache -hf auth login --token $HF_TOKEN --add-to-git-credential - -# always download and use the real dataset -# hf download legacy-datasets/c4 \ -# --include "en/partial-train/000*.parquet" \ -# --repo-type dataset \ -# --revision refs/convert/parquet - -# debug -# apt install tree -y > /dev/null -# tree /hf_cache/ -# tree /hf_cache/hub/datasets--legacy-datasets--c4/snapshots/5abe0d085aa23dd9db2a6c1e86cfce4e4db6f0c3/ -# ls /hf_cache/hub/datasets--legacy-datasets--c4/snapshots/5abe0d085aa23dd9db2a6c1e86cfce4e4db6f0c3/en/ -# ls /hf_cache/hub/datasets--legacy-datasets--c4/snapshots/5abe0d085aa23dd9db2a6c1e86cfce4e4db6f0c3/en/partial-train/ - -download_tokenizer(){ - hf download $1 --include "**token**" -} - -if [[ "$MODEL_REPO" == "Llama-2-7B" ]]; then - download_tokenizer "meta-llama/Llama-2-7b" -elif [[ "$MODEL_REPO" == "Llama-2-70B" ]]; then - download_tokenizer "meta-llama/Llama-2-70b" -elif [[ "$MODEL_REPO" == "Llama-3.1-8B" ]]; then - download_tokenizer "meta-llama/Meta-Llama-3-8B" -elif [[ "$MODEL_REPO" == "Llama-3.1-70B" ]]; then - download_tokenizer "meta-llama/Meta-Llama-3-70B" -elif [[ "$MODEL_REPO" == "Llama-3.3-70B" ]]; then - download_tokenizer "meta-llama/Llama-3.3-70B-Instruct" -elif [[ "$MODEL_REPO" == "DeepSeek-V2-lite" ]]; then - echo "No tokenizer for download" -elif [[ "$MODEL_REPO" == "Mixtral-8x7B" ]]; then - download_tokenizer "mistralai/Mixtral-8x7B-v0.1" -elif [[ "$MODEL_REPO" == "Qwen3-14B" ]]; then - download_tokenizer "Qwen/Qwen3-14B" -elif [[ "$MODEL_REPO" == "Qwen3-30B-A3B" ]]; then - download_tokenizer "Qwen/Qwen3-30B-A3B" -else - echo "Error: Unsupported training mode." - exit 1 -fi diff --git a/scripts/jax-maxtext/jax_maxtext_multinode_benchmark.sh b/scripts/jax-maxtext/jax_maxtext_multinode_benchmark.sh deleted file mode 100644 index 4706184b..00000000 --- a/scripts/jax-maxtext/jax_maxtext_multinode_benchmark.sh +++ /dev/null @@ -1,389 +0,0 @@ -#!/bin/bash -#SBATCH --nodes=8 -#SBATCH --exclusive # exclusive node access -#SBATCH --gpus-per-task=8 -#SBATCH --mem=0 # all mem avail -#SBATCH --ntasks-per-node=1 # n tasks per machine (one task per gpu) -#SBATCH --overcommit - -set -e - -################################################################################# -# Unified MaxText Training Script -# -# Usage: sbatch jax_maxtext_multinode_benchmark.sh [docker_image] -# -# Arguments: -# config_file.yml - Required. Path to model config YAML file (absolute or relative) -# Examples: llama2-7b.gpu.yml, /path/to/model.gpu.yml, ../configs/my-model.yml -# docker_image - Optional. Docker image to use (default: rocm/jax-training:latest) -# -# Examples: -# sbatch jax_maxtext_multinode_benchmark.sh llama2-7b.gpu.yml -# sbatch jax_maxtext_multinode_benchmark.sh /mnt/vast/araina/configs/llama3-70b.gpu.yml -# sbatch jax_maxtext_multinode_benchmark.sh ../models/custom-model.gpu.yml my-docker-image:tag -################################################################################# - -mkdir -p -v outputs; chmod a+w outputs - -LOOKUP_USER="${USER:-}" - -# ------- Parse command line arguments ------- -DEFAULT_DOCKER_IMAGE="rocm/jax-training:latest" - -if [[ $# -eq 0 ]]; then - echo "ERROR: No config file provided!" - echo "" - echo "Usage: sbatch jax_maxtext_multinode_benchmark.sh [docker_image]" - echo "" - echo "Arguments:" - echo " config_file.yml - Required. Path to model config YAML file (absolute or relative)" - echo " docker_image - Optional. Default: $DEFAULT_DOCKER_IMAGE" - echo "" - echo "Examples:" - echo " sbatch jax_maxtext_multinode_benchmark.sh llama2-7b.gpu.yml" - echo " sbatch jax_maxtext_multinode_benchmark.sh /path/to/model.gpu.yml" - echo " sbatch jax_maxtext_multinode_benchmark.sh ../configs/my-model.yml" - exit 1 -fi - -CONFIG_FILE="$1" -DOCKER_IMAGE="${2:-$DEFAULT_DOCKER_IMAGE}" -EXP_TAG="" - -# Convert to absolute path for consistency -if [[ "$CONFIG_FILE" != /* ]]; then - # Relative path - convert to absolute - CONFIG_FILE="$(cd "$(dirname "$CONFIG_FILE")" && pwd)/$(basename "$CONFIG_FILE")" -fi - -# Validate config file exists -if [[ ! -f "$CONFIG_FILE" ]]; then - echo "ERROR: Config file not found: $CONFIG_FILE" >&2 - exit 1 -fi - -# Extract model name from config filename (remove path and extension) -CONFIG_BASENAME=$(basename "$CONFIG_FILE") -MODEL_NAME="${CONFIG_BASENAME%.gpu.yml}" -MODEL_NAME="${MODEL_NAME%.yml}" - -# Get the directory containing the config file (for mounting) -CONFIG_DIR=$(dirname "$CONFIG_FILE") - -echo "CONFIG_FILE=$CONFIG_FILE" -echo "CONFIG_DIR=$CONFIG_DIR" -echo "CONFIG_BASENAME=$CONFIG_BASENAME" -echo "MODEL_NAME=$MODEL_NAME" -echo "DOCKER_IMAGE=$DOCKER_IMAGE" - -# ------- date command detection (GNU date or gdate) ------- -DATE_CMD="date" -if ! date -d '1970-01-01 00:00:00' +%s >/dev/null 2>&1; then - if command -v gdate >/dev/null 2>&1; then - DATE_CMD="gdate" - else - echo "WARNING: Your 'date' doesn't support -d. Reservation time parsing may fail; falling back to first active match." >&2 - fi -fi - -# ------- function to pick reservation ------- -get_reservation_for_user() { - local uname="${1}" - local datecmd="${2}" - local now_epoch - now_epoch="$("$datecmd" +%s)" - scontrol show reservation -o 2>/dev/null | \ - awk -v user="$uname" -v now="$now_epoch" -v datecmd="$datecmd" ' - function to_epoch(ts, cmd, epoch_str) { - gsub(/T/, " ", ts) - if (ts == "" || ts == "Unknown") return 0 - cmd = datecmd " -d \"" ts "\" +%s" - epoch_str = "" - cmd | getline epoch_str - close(cmd) - if (epoch_str ~ /^[0-9]+$/) return epoch_str + 0 - return 0 - } - { - name=""; users=""; start_s=""; end_s="" - if (match($0, /ReservationName=([^ ]+)/, m)) name=m[1] - if (match($0, /Users=([^ ]+)/, mu)) users=mu[1] - if (match($0, /StartTime=([^ ]+)/, ms)) start_s=ms[1] - if (match($0, /EndTime=([^ ]+)/, me)) end_s=me[1] - n = split(users, arr, ",") - ok=0 - for (i=1; i<=n; i++) if (arr[i] == user) { ok=1; break } - if (!ok) next - start = to_epoch(start_s) - end = to_epoch(end_s) - if (start==0 || end==0) { - start=1; end=now+1 - } - if (start <= now && now <= end) { - printf("%d\t%s\n", start, name) - } - } - ' | sort -nr | awk 'NR==1 { print $2 }' -} - -RESERVATION_NAME="$(get_reservation_for_user "${LOOKUP_USER}" "${DATE_CMD}")" -if [[ -n "${RESERVATION_NAME}" ]]; then - echo "Using reservation for user '${LOOKUP_USER}': ${RESERVATION_NAME}" -else - echo "No active reservation found for user '${LOOKUP_USER}'. Submitting without --reservation." -fi - -# Config file already validated above -echo "EXP_TAG=$EXP_TAG" - -# ------- Build job name ------- -JOB_NAME="JAX-${MODEL_NAME}" -if [[ -n "$EXP_TAG" ]]; then - JOB_NAME="${JOB_NAME}-${EXP_TAG}" -fi -echo "JOB_NAME=$JOB_NAME" - -# ------- Setup for distributed execution ------- -echo "SLURM_JOB_NUM_NODES=$SLURM_JOB_NUM_NODES" - -# Determine coordinator IP (first node in the job) -COORDINATOR_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) -echo "Coordinator node: $COORDINATOR_NODE" - -# Get the IP of the coordinator node -if [[ -n "${SLURM_LAUNCH_NODE_IPADDR:-}" ]]; then - JAX_COORDINATOR_IP=$SLURM_LAUNCH_NODE_IPADDR - echo "Using JAX_COORDINATOR_IP=$JAX_COORDINATOR_IP (from SLURM_LAUNCH_NODE_IPADDR)" -else - # Resolve the first node's IP address - JAX_COORDINATOR_IP=$(srun --nodes=1 --ntasks=1 -w "$COORDINATOR_NODE" hostname -I | awk '{print $1}') - echo "Using JAX_COORDINATOR_IP=$JAX_COORDINATOR_IP (resolved from $COORDINATOR_NODE)" -fi - -JAX_PORT=$((20000 + $RANDOM % 40000)) -echo "JAX_PORT=$JAX_PORT" - -# ------- Git summary ------- -echo "=== GIT SUMMARY BEGIN ===" -echo "[BRANCH]" -git status --branch --short 2>/dev/null || echo "Not a git repository" -echo -echo "[LAST COMMIT]" -git --no-pager log -1 --pretty=format:"%h %s (%ad) <%an>" 2>/dev/null || echo "No commits" -echo -echo "=== GIT SUMMARY END ===" - -# ============================================================================ -# Node Setup Script (heredoc for readability - no escaping needed) -# This function outputs the script that runs on each node before docker run. -# Handles: docker detection, GPU cleanup, image pull, NCCL setup -# ============================================================================ -write_node_setup_script() { -cat << 'NODE_SETUP_EOF' -#!/bin/bash -set -e - -echo "=== Node $SLURM_NODEID: Starting ===" - -# ------- Get docker binary ------- -if command -v podman >/dev/null 2>&1; then - runtime_dir="${XDG_RUNTIME_DIR:-/run/user/$(id -u)}" - if [ -d "$runtime_dir" ] && podman info >/dev/null 2>&1; then - DOCKER_BIN="podman" - fi -fi -if [ -z "${DOCKER_BIN:-}" ] && command -v docker >/dev/null 2>&1; then - if docker info >/dev/null 2>&1; then - DOCKER_BIN="docker" - else - DOCKER_BIN="sudo docker" - fi -fi -if [ -z "${DOCKER_BIN:-}" ]; then - echo "ERROR: No docker/podman found" >&2 - exit 1 -fi -echo "DOCKER_BIN=$DOCKER_BIN" - -# ------- GPU cleanup ------- -echo "=== GPU cleanup ===" -GPU_PIDS=$(rocm-smi --showpids 2>/dev/null | grep -oP "^\d+" | grep -v "^$" || true) -if [ -n "$GPU_PIDS" ]; then - echo "Found GPU processes: $GPU_PIDS" - ALL_CONTAINERS=$($DOCKER_BIN ps -q 2>/dev/null || true) - if [ -n "$ALL_CONTAINERS" ]; then - echo "Stopping containers..." - echo "$ALL_CONTAINERS" | xargs $DOCKER_BIN stop -t 10 || true - sleep 20 - fi - USE_SUDO="" - [[ "$DOCKER_BIN" == sudo* ]] && USE_SUDO="sudo" - for PID in $GPU_PIDS; do - $USE_SUDO kill -9 $PID 2>/dev/null || true - done - sleep 10 -fi -# ------- Docker image pull ------- -# DOCKER_IMAGE is passed via environment variable -if ! $DOCKER_BIN image inspect "$DOCKER_IMAGE" >/dev/null 2>&1; then - echo "[INFO] Pulling $DOCKER_IMAGE ..." - $DOCKER_BIN pull "$DOCKER_IMAGE" -fi - -# ------- NCCL setup ------- -NCCL_IB_HCA=$(ls /sys/class/infiniband 2>/dev/null | tr "\n" "," | sed "s/,$//" || true) -candidates=$(ip -o -4 addr show scope global 2>/dev/null | awk '{print $2, $4}' | sort -k1,1) -NCCL_SOCKET_IFNAME=$(echo "$candidates" | awk '$2 ~ /^10\./ {print $2, $1}' | sort -V -k1,1 | head -n1 | awk '{print $2}') -[ -z "$NCCL_SOCKET_IFNAME" ] && NCCL_SOCKET_IFNAME=$(echo "$candidates" | awk '$2 ~ /^172\.(1[6-9]|2[0-9]|3[0-1])\./ {print $2, $1}' | sort -V -k1,1 | head -n1 | awk '{print $2}') -[ -z "$NCCL_SOCKET_IFNAME" ] && NCCL_SOCKET_IFNAME=$(echo "$candidates" | awk '$2 ~ /^192\.168\./ {print $2, $1}' | sort -V -k1,1 | head -n1 | awk '{print $2}') - -# Export variables for the caller (DOCKER_IMAGE comes from parent environment) -export DOCKER_BIN NCCL_IB_HCA NCCL_SOCKET_IFNAME -NODE_SETUP_EOF -} - -# ============================================================================ -# Inner Docker Script (heredoc for readability - no escaping needed) -# This function outputs the script that runs inside the container. -# ============================================================================ -write_inner_script() { -cat << INNER_SCRIPT_EOF -#!/bin/bash -set -ex - -cd /workspace/maxtext -MAXTEXT_SRC_DIR=. -[[ -d ./src ]] && MAXTEXT_SRC_DIR=./src -export PYTHONPATH="${MAXTEXT_SRC_DIR}:${PYTHONPATH}" -cd "${MAXTEXT_SRC_DIR}" - -# ------- Output directory setup ------- -export OUTPUT_PATH="/dockerx/outputs/\${JOB_ID_AND_NAME}" -mkdir -p \$OUTPUT_PATH - -# ------- NCCL Configuration ------- -export NCCL_CHECKS_DISABLE=1 -export NCCL_DEBUG=INFO -export TF_CPP_MIN_LOG_LEVEL=2 - -# ------- JAX/XLA Configuration ------- -export XLA_PYTHON_CLIENT_MEM_FRACTION=.93 -export JAX_HIP_GRAPH_LOWERING=false -# Note: XLA_FLAGS is passed via docker --env to override container's baked-in value - -# ------- NCCL Performance Tuning ------- -export NCCL_CROSS_NIC=2 -export NCCL_NCHANNELS_PER_NET_PEER=4 -export NCCL_NSOCKS_PERTHREAD=4 -export NCCL_SOCKET_NTHREADS=8 -export NCCL_IB_QPS_PER_CONNECTION=4 -export CUDA_DEVICE_MAX_CONNECTIONS=1 -export GPU_MAX_HW_QUEUES=2 - -# ------- HIP/ROCm Configuration ------- -export HIP_FORCE_DEV_KERNARG=1 -export HSA_ENABLE_IPC_MODE_LEGACY=1 -export HSA_FORCE_FINE_GRAIN_PCIE=1 -export HSA_NO_SCRATCH_RECLAIM=1 - -# ------- Transformer Engine Configuration ------- -export NVTE_CK_USES_BWD_V3=1 -export NVTE_CK_USES_FWD_V3=1 -export NVTE_FRAMEWORK=jax -export NVTE_FUSED_ATTN=1 -export NVTE_FUSED_ATTN_AOTRITON=0 -export NVTE_FUSED_ATTN_CK=1 -export NVTE_USE_CAST_TRANSPOSE_TRITON=0 -export NVTE_USE_HIPBLASLT=1 -export NVTE_USE_ROCM=1 -export CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=2 -export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -export NVTE_CK_HOW_V3_BF16_CVT=2 -export NVTE_CK_IS_V3_ATOMIC_FP32=0 - -# ------- RCCL/NCCL IB Tuning ------- -export IONIC_LOCKFREE=all -export NCCL_GDR_COPY_ENABLE=1 -export NCCL_GDR_FLUSH_DISABLE=1 -export NCCL_IB_ECE_ENABLE=0 -export NCCL_IB_FIFO_TC=184 -export NCCL_IB_GID_INDEX=1 -export NCCL_IB_PCI_RELAXED_ORDERING=1 -export NCCL_IB_TC=96 -export NCCL_IB_USE_INLINE=1 -export NCCL_IGNORE_CPU_AFFINITY=1 -export NCCL_PXN_DISABLE=0 -export NET_OPTIONAL_RECV_COMPLETION=1 -export RCCL_GDR_FLUSH_GPU_MEM_NO_RELAXED_ORDERING=0 -export RCCL_LL128_FORCE_ENABLE=1 - -# ------- Run MaxText Training ------- -echo "Starting MaxText training with model: \${MODEL_NAME}" -echo "Output path: \${OUTPUT_PATH}" - -# Use config file path - if it's in a different directory, it will be mounted separately -python3 -m maxtext.trainers.pre_train.train /configs/${CONFIG_BASENAME} base_output_directory=\${OUTPUT_PATH} -INNER_SCRIPT_EOF -} - -# ============================================================================ -# Launch on all nodes via srun -# ============================================================================ -echo "==STARTING JOBS ON ALL NODES==" - -# Export variables that srun needs -export JAX_COORDINATOR_IP -export JAX_PORT -export MODEL_NAME -export CONFIG_BASENAME -export CONFIG_DIR -export DOCKER_IMAGE - -# XLA_FLAGS must be passed via docker --env to override the container's baked-in value -# (the container's profile scripts set XLA_FLAGS with 'FALSE' which fails parsing) -export XLA_FLAGS="--xla_gpu_memory_limit_slop_factor=95 \ ---xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 \ ---xla_gpu_enable_command_buffer='' \ ---xla_gpu_enable_latency_hiding_scheduler=true \ ---xla_gpu_all_gather_combine_threshold_bytes=8589934592 \ ---xla_gpu_enable_triton_gemm=false \ ---xla_gpu_enable_cublaslt=true \ ---xla_gpu_autotune_level=4 \ ---xla_gpu_enable_all_gather_combine_by_dim=false" - -# Write the setup scripts (defined above via heredoc for readability) -NODE_SETUP_SCRIPT=".maxtext_node_setup_$$.sh" -INNER_SCRIPT=".maxtext_inner_$$.sh" -write_node_setup_script > "$NODE_SETUP_SCRIPT" -write_inner_script > "$INNER_SCRIPT" -chmod +x "$NODE_SETUP_SCRIPT" "$INNER_SCRIPT" -trap "rm -f '$NODE_SETUP_SCRIPT' '$INNER_SCRIPT'" EXIT - -srun -l bash -c ' -# Source the node setup script (docker detection, GPU cleanup, image pull, NCCL setup) -source "'"$NODE_SETUP_SCRIPT"'" - -# Docker run -echo "==Starting container on node $SLURM_NODEID==" -EXTRA_GIDS=$(id -G) -GROUP_ADD_ARGS="" -for gid in $EXTRA_GIDS; do GROUP_ADD_ARGS="$GROUP_ADD_ARGS --group-add $gid"; done - -$DOCKER_BIN run --rm --cap-add=SYS_PTRACE --ipc=host --network=host \ - --device=/dev/kfd --device=/dev/dri --device=/dev/infiniband \ - --env JAX_COORDINATOR_IP=$JAX_COORDINATOR_IP \ - --env JAX_COORDINATOR_PORT=$JAX_PORT \ - --env JOB_ID_AND_NAME="${SLURM_JOB_ID}-${SLURM_JOB_NAME}" \ - --env MODEL_NAME=$MODEL_NAME \ - --env NCCL_IB_HCA=$NCCL_IB_HCA \ - --env NCCL_SOCKET_IFNAME=$NCCL_SOCKET_IFNAME \ - --env NNODES=$SLURM_NNODES \ - --env NODE_RANK=$SLURM_NODEID \ - --env XLA_FLAGS="$XLA_FLAGS" \ - --security-opt seccomp=unconfined --privileged $GROUP_ADD_ARGS \ - -v /boot:/boot:ro -v $PWD:/dockerx -v $CONFIG_DIR:/configs:ro -w /dockerx \ - $DOCKER_IMAGE /bin/bash -lc "source /dockerx/'"$INNER_SCRIPT"'" -' - -echo "==DONE==" diff --git a/scripts/jax-maxtext/models.json b/scripts/jax-maxtext/models.json deleted file mode 100644 index 15008566..00000000 --- a/scripts/jax-maxtext/models.json +++ /dev/null @@ -1,468 +0,0 @@ -[ - { - "name": "jax_maxtext_train_llama-2-7b", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "", - "multiple_results": "perf_Llama-2-7B.csv", - "tags": [ - "jax", - "llama2", - "training" - ], - "timeout": -1, - "args": "--model_repo jax_maxtext_train_llama-2-7b" - }, - { - "name": "jax_maxtext_train_llama-2-7b_nanoo_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "nanoo_fp8", - "multiple_results": "perf_Llama-2-7B_nanoo_fp8.csv", - "tags": [ - "jax", - "llama2", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx950", - "args": "--model_repo jax_maxtext_train_llama-2-7b --quantization nanoo_fp8" - }, - { - "name": "jax_maxtext_train_llama-2-7b_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "fp8", - "multiple_results": "perf_Llama-2-7B_fp8.csv", - "tags": [ - "jax", - "llama2", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx942", - "args": "--model_repo jax_maxtext_train_llama-2-7b --quantization fp8" - }, - { - "name": "jax_maxtext_train_llama-2-70b", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "", - "multiple_results": "perf_Llama-2-70B.csv", - "tags": [ - "jax", - "llama2", - "training" - ], - "timeout": -1, - "args": "--model_repo jax_maxtext_train_llama-2-70b" - }, - { - "name": "jax_maxtext_train_llama-2-70b_nanoo_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "nanoo_fp8", - "multiple_results": "perf_Llama-2-70B_nanoo_fp8.csv", - "tags": [ - "jax", - "llama2", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx950", - "args": "--model_repo jax_maxtext_train_llama-2-70b --quantization nanoo_fp8" - }, - { - "name": "jax_maxtext_train_llama-2-70b_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "fp8", - "multiple_results": "perf_Llama-2-70B_fp8.csv", - "tags": [ - "jax", - "llama2", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx942", - "args": "--model_repo jax_maxtext_train_llama-2-70b --quantization fp8" - }, - { - "name": "jax_maxtext_train_llama-3.1-8b", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "", - "multiple_results": "perf_Llama-3.1-8B.csv", - "tags": [ - "jax", - "llama3", - "training" - ], - "timeout": -1, - "args": "--model_repo jax_maxtext_train_llama-3.1-8b" - }, - { - "name": "jax_maxtext_train_llama-3.1-8b_nanoo_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "nanoo_fp8", - "multiple_results": "perf_Llama-3.1-8B_nanoo_fp8.csv", - "tags": [ - "jax", - "llama3", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx950", - "args": "--model_repo jax_maxtext_train_llama-3.1-8b --quantization nanoo_fp8" - }, - { - "name": "jax_maxtext_train_llama-3.1-8b_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "fp8", - "multiple_results": "perf_Llama-3.1-8B_fp8.csv", - "tags": [ - "jax", - "llama3", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx942", - "args": "--model_repo jax_maxtext_train_llama-3.1-8b --quantization fp8" - }, - { - "name": "jax_maxtext_train_llama-3.1-70b", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "", - "multiple_results": "perf_Llama-3.1-70B.csv", - "tags": [ - "jax", - "llama3", - "training" - ], - "timeout": -1, - "args": "--model_repo jax_maxtext_train_llama-3.1-70b" - }, - { - "name": "jax_maxtext_train_llama-3.1-70b_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "fp8", - "multiple_results": "perf_Llama-3.1-70B_fp8.csv", - "tags": [ - "jax", - "llama3", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx942", - "args": "--model_repo jax_maxtext_train_llama-3.1-70b --quantization fp8" - }, - { - "name": "jax_maxtext_train_llama-3.3-70b", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "", - "multiple_results": "perf_Llama-3.3-70B.csv", - "tags": [ - "jax", - "llama3", - "training" - ], - "timeout": -1, - "args": "--model_repo jax_maxtext_train_llama-3.3-70b" - }, - { - "name": "jax_maxtext_train_llama-3.3-70b_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "fp8", - "multiple_results": "perf_Llama-3.3-70B_fp8.csv", - "tags": [ - "jax", - "llama3", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx942", - "args": "--model_repo jax_maxtext_train_llama-3.3-70b --quantization fp8" - }, - { - "name": "jax_maxtext_train_deepseek-v2-lite-16b", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "", - "multiple_results": "perf_DeepSeek-V2-lite.csv", - "tags": [ - "jax", - "deepseekv2", - "training" - ], - "timeout": -1, - "args": "--model_repo jax_maxtext_train_deepseek-v2-lite-16b" - }, - { - "name": "jax_maxtext_train_deepseek-v2-lite-16b_nanoo_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "nanoo_fp8", - "multiple_results": "perf_DeepSeek-V2-lite_nanoo_fp8.csv", - "tags": [ - "jax", - "deepseekv2", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx950", - "args": "--model_repo jax_maxtext_train_deepseek-v2-lite-16b --quantization nanoo_fp8" - }, - { - "name": "jax_maxtext_train_deepseek-v2-lite-16b_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "fp8", - "multiple_results": "perf_DeepSeek-V2-lite_fp8.csv", - "tags": [ - "jax", - "deepseekv2", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx942", - "args": "--model_repo jax_maxtext_train_deepseek-v2-lite-16b --quantization fp8" - }, - { - "name": "jax_maxtext_train_mixtral-8x7b", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "", - "multiple_results": "perf_Mixtral-8x7B.csv", - "tags": [ - "jax", - "mixtral", - "training" - ], - "timeout": -1, - "args": "--model_repo jax_maxtext_train_mixtral-8x7b" - }, - { - "name": "jax_maxtext_train_mixtral-8x7b_nanoo_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "nanoo_fp8", - "multiple_results": "perf_Mixtral-8x7B_nanoo_fp8.csv", - "tags": [ - "jax", - "mixtral", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx950", - "args": "--model_repo jax_maxtext_train_mixtral-8x7b --quantization nanoo_fp8" - }, - { - "name": "jax_maxtext_train_mixtral-8x7b_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "fp8", - "multiple_results": "perf_Mixtral-8x7B_fp8.csv", - "tags": [ - "jax", - "mixtral", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx942", - "args": "--model_repo jax_maxtext_train_mixtral-8x7b --quantization fp8" - }, - { - "name": "jax_maxtext_train_qwen3-14b", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "", - "multiple_results": "perf_Qwen3-14B.csv", - "tags": [ - "jax", - "qwen3", - "training" - ], - "timeout": -1, - "args": "--model_repo jax_maxtext_train_qwen3-14b" - }, - { - "name": "jax_maxtext_train_qwen3-14b_nanoo_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "nanoo_fp8", - "multiple_results": "perf_Qwen3-14B_nanoo_fp8.csv", - "tags": [ - "jax", - "qwen3", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx950", - "args": "--model_repo jax_maxtext_train_qwen3-14b --quantization nanoo_fp8" - }, - { - "name": "jax_maxtext_train_qwen3-14b_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "fp8", - "multiple_results": "perf_Qwen3-14B_fp8.csv", - "tags": [ - "jax", - "qwen3", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx942", - "args": "--model_repo jax_maxtext_train_qwen3-14b --quantization fp8" - }, - { - "name": "jax_maxtext_train_qwen3-30b-a3b", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "", - "multiple_results": "perf_Qwen3-30B-A3B.csv", - "tags": [ - "jax", - "qwen3", - "training" - ], - "timeout": -1, - "args": "--model_repo jax_maxtext_train_qwen3-30b-a3b" - }, - { - "name": "jax_maxtext_train_qwen3-30b-a3b_nanoo_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "nanoo_fp8", - "multiple_results": "perf_Qwen3-30B-A3B_nanoo_fp8.csv", - "tags": [ - "jax", - "qwen3", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx950", - "args": "--model_repo jax_maxtext_train_qwen3-30b-a3b --quantization nanoo_fp8" - }, - { - "name": "jax_maxtext_train_qwen3-30b-a3b_fp8", - "dockerfile": "../../docker/jax_maxtext", - "scripts": "run.sh", - "url": "", - "data": "huggingface", - "n_gpus": "-1", - "owner": "mad.support@amd.com", - "training_precision": "fp8", - "multiple_results": "perf_Qwen3-30B-A3B_fp8.csv", - "tags": [ - "jax", - "qwen3", - "training" - ], - "timeout": -1, - "skip_gpu_arch": "gfx942", - "args": "--model_repo jax_maxtext_train_qwen3-30b-a3b --quantization fp8" - } -] diff --git a/scripts/jax-maxtext/run.sh b/scripts/jax-maxtext/run.sh old mode 100644 new mode 100755 index 3b5146d7..d0c670f4 --- a/scripts/jax-maxtext/run.sh +++ b/scripts/jax-maxtext/run.sh @@ -1,9 +1,9 @@ -#!/bin/bash +#!/usr/bin/env bash ############################################################################### # # MIT License # -# Copyright (c) 2025 Advanced Micro Devices, Inc. +# Copyright (c) 2026 Advanced Micro Devices, Inc. # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal @@ -25,45 +25,103 @@ # ################################################################################# -export HF_TOKEN=$MAD_SECRETS_HFTOKEN +# Wrapper for Primus JAX/MaxText pretrain when run via madengine (local, SLURM, K8s). +# Sets EXP from PRIMUS_CONFIG_PATH or --config_path, runs Primus examples/run_pretrain.sh +# with BACKEND=MaxText, then extracts tps/tflops into primus_perf_output.csv for +# madengine multiple_results. MaxText-only: no Megatron/TorchTitan logic here. +set -e -# Parse named arguments -while [[ "$#" -gt 0 ]]; do - case $1 in - --model_repo) MODEL_REPO="$2"; shift ;; - *) echo "Unknown parameter passed: $1"; usage ;; - esac - shift - case $1 in - --quantization) QUANTIZATION="$2"; shift ;; - *) echo "Unknown parameter passed: $1"; usage ;; - esac - shift -done +# madengine invokes this as `cd run_directory && bash run.sh ...`. +RUN_DIR="$(pwd)" -echo "Model repo: $MODEL_REPO" +# Primus root: repo checkout, then image COPY / K8s ConfigMap extract, then env, then legacy paths. +script_dir="$(cd "$(dirname "$0")" && pwd)" +if [[ -f "$script_dir/../Primus/examples/run_pretrain.sh" ]]; then + export PRIMUS_ROOT="$(cd "$script_dir/../Primus" && pwd)" +elif [[ -f "/workspace/Primus/examples/run_pretrain.sh" ]]; then + export PRIMUS_ROOT="/workspace/Primus" +elif [[ -n "${PRIMUS_ROOT:-}" ]]; then + : +elif [[ -f "/opt/primus/examples/run_pretrain.sh" ]]; then + export PRIMUS_ROOT="/opt/primus" +elif [[ -f "/workspace/examples/run_pretrain.sh" ]]; then + export PRIMUS_ROOT="/workspace" +else + echo "ERROR: Could not find Primus run_pretrain.sh. Set PRIMUS_ROOT or use a repo with scripts/Primus submodule." >&2 + exit 1 +fi + +# EXP is required by run_pretrain.sh. --config_path must also be stripped from the +# forwarded args: run_pretrain.sh appends leftovers to the training command and it is +# not a valid MaxText flag. +forward_args=() +if [[ -n "${PRIMUS_CONFIG_PATH:-}" ]]; then + export EXP="$PRIMUS_CONFIG_PATH" + forward_args=("$@") +else + export EXP="" + args=("$@") + i=0 + while [[ $i -lt ${#args[@]} ]]; do + if [[ "${args[i]}" == "--config_path" && -n "${args[i+1]:-}" ]]; then + export EXP="${args[i+1]}" + i=$((i + 2)) + continue + fi + forward_args+=("${args[i]}") + i=$((i + 1)) + done +fi -if [[ "$MODEL_REPO" == "jax_maxtext_train_llama-3.1-8b" ]]; then - model="Llama-3.1-8B" -elif [[ "$MODEL_REPO" == "jax_maxtext_train_llama-3.1-70b" ]]; then - model="Llama-3.1-70B" -elif [[ "$MODEL_REPO" == "jax_maxtext_train_llama-3.3-70b" ]]; then - model="Llama-3.3-70B" -elif [[ "$MODEL_REPO" == "jax_maxtext_train_llama-2-7b" ]]; then - model="Llama-2-7B" -elif [[ "$MODEL_REPO" == "jax_maxtext_train_llama-2-70b" ]]; then - model="Llama-2-70B" -elif [[ "$MODEL_REPO" == "jax_maxtext_train_deepseek-v2-lite-16b" ]]; then - model="DeepSeek-V2-lite" -elif [[ "$MODEL_REPO" == "jax_maxtext_train_mixtral-8x7b" ]]; then - model="Mixtral-8x7B" -elif [[ "$MODEL_REPO" == "jax_maxtext_train_qwen3-14b" ]]; then - model="Qwen3-14B" -elif [[ "$MODEL_REPO" == "jax_maxtext_train_qwen3-30b-a3b" ]]; then - model="Qwen3-30B-A3B" +if [[ -z "$EXP" ]]; then + echo "ERROR: --config_path or PRIMUS_CONFIG_PATH required." >&2 + exit 1 fi -./jax-maxtext_benchmark_setup.sh -m $model -./jax-maxtext_benchmark_report.sh -m $model -q $QUANTIZATION +# Makes run_pretrain.sh launch primus/cli train pretrain rather than torchrun. +export BACKEND="MaxText" -echo "performance: 1 pass" +# Use the image-baked MaxText (/workspace/maxtext) rather than Primus's +# third_party/maxtext checkout. The image build installs matching XLA/JAX/TE +# wheels alongside MaxText; using a mismatched third_party copy can trigger +# hipblaslt Tensile kernel failures on gfx950. +export MAXTEXT_PATH="${MAXTEXT_PATH:-/workspace/maxtext}" +export BACKEND_PATH="${BACKEND_PATH:-$MAXTEXT_PATH}" + +# The image already satisfies requirements-jax.txt (installed at build time), so the +# per-run pip install has nothing to do. Skipping it keeps launches off the network and +# stops a resolve from moving pinned versions under a benchmark. PRIMUS_SKIP_PIP=0 restores it. +export PRIMUS_SKIP_PIP="${PRIMUS_SKIP_PIP:-1}" + +# HF_TOKEN for Primus prepare: explicit, then MAD convention, then madengine v2. +if [[ -n "${HF_TOKEN:-}" ]]; then + export HF_TOKEN +elif [[ -n "${MAD_SECRETS_HFTOKEN:-}" ]]; then + export HF_TOKEN="$MAD_SECRETS_HFTOKEN" +elif [[ -n "${MAD_SECRET_HFTOKEN:-}" ]]; then + export HF_TOKEN="$MAD_SECRET_HFTOKEN" +fi + +# This wrapper deliberately exports no perf/arch env. All XLA_FLAGS and NVTE/HIP/HSA +# tunables, including the arch-gated ones, are applied in-process before JAX init by +# primus/backends/maxtext/env_spec.py. MAD only picks the config and finds the log. + +# I/O contract, not a knob: tells Primus where to write the log this wrapper parses. +mkdir -p "$RUN_DIR/output" +export TRAIN_LOG="$RUN_DIR/output/log_mp_pretrain_$(basename "$EXP" .yaml).txt" + +# EXP paths are relative to PRIMUS_ROOT. No exec: the perf extractor runs after training. +# The `||` is what keeps set -e from exiting here, so a failed run still gets parsed. +cd "$PRIMUS_ROOT" +exitcode=0 +bash "$PRIMUS_ROOT/examples/run_pretrain.sh" "${forward_args[@]}" || exitcode=$? + +# madengine resolves multiple_results against its own CWD (the parent of run_directory) +# and deletes run_directory before parsing perf, so the CSV must go to the parent. +PERF_OUT="$RUN_DIR/../primus_perf_output.csv" +if [[ -f "$TRAIN_LOG" ]]; then + extract_script="${script_dir}/extract_maxtext_perf.py" + [[ -f "$RUN_DIR/extract_maxtext_perf.py" ]] && extract_script="$RUN_DIR/extract_maxtext_perf.py" + python3 "$extract_script" "$TRAIN_LOG" "$PERF_OUT" || true +fi +exit "$exitcode" diff --git a/scripts/primus_train/get_models_json.py b/scripts/primus_train/get_models_json.py index 906cbf8e..2906deee 100644 --- a/scripts/primus_train/get_models_json.py +++ b/scripts/primus_train/get_models_json.py @@ -21,9 +21,30 @@ # One glob for all launchers: examples//configs/**/*.yaml CONFIGS_GLOB = os.path.join(PRIMUS_ROOT, "examples", "*", "configs", "**", "*.yaml") +# JAX backends have their own dedicated discovery (scripts/jax-maxtext/get_models_json.py, +# scripts/jax-maxdiffusion/get_models_json.py) with correct dockerfiles and arch filtering. +# Discovering them here too would create duplicates on the wrong base image. +JAX_BACKENDS = {"maxtext", "maxdiffusion"} + def list_models(): - models = [] + # Default/smoke-test entry -> "primus_train/default". Lives here (not root models.json) + # so this directory has one registration file, per madengine's models.json vs. + # get_models_json.py rule. HSA_NO_SCRATCH_RECLAIM etc. are not modeled here: madengine + # has no per-model env field, so pass them via --additional-context docker_env_vars + # (see benchmark/primus/README.md) instead. + models = [ + CustomModel( + name="default", + dockerfile="../../docker/primus", + dockercontext=".", + scripts="run.sh", + n_gpus="-1", + owner="mad.support@amd.com", + tags=["training", "primus", "megatron", "pretrain"], + args="", + ) + ] if not os.path.isdir(PRIMUS_ROOT): return models for yaml_path in sorted(glob.glob(CONFIGS_GLOB)): @@ -33,6 +54,8 @@ def list_models(): if len(parts) < 5: continue launcher = parts[1] # megatron, torchtitan, megatron_bridge, etc. + if launcher in JAX_BACKENDS: + continue arch = parts[3] # MI300X, MI355X, etc. short_name = os.path.splitext(os.path.basename(yaml_path))[0] # discover_models prefixes with dirname (primus_train/), so no prefix here diff --git a/scripts/primus_train/run.sh b/scripts/primus_train/run.sh index 69419ad7..c28dc8b9 100755 --- a/scripts/primus_train/run.sh +++ b/scripts/primus_train/run.sh @@ -62,6 +62,8 @@ fi exp_lower="$(echo "$EXP" | tr '[:upper:]' '[:lower:]')" if [[ "$exp_lower" == *"/maxtext/"* ]]; then export BACKEND="MaxText" +elif [[ "$exp_lower" == *"/maxdiffusion/"* ]]; then + export BACKEND="MaxDiffusion" elif [[ "$exp_lower" == *"/torchtitan/"* ]]; then export BACKEND="torchtitan" elif [[ "$exp_lower" == *"/megatron_bridge/"* ]]; then diff --git a/scripts/sglang_disagg/README.MD b/scripts/sglang_disagg/README.MD index f7414f1c..1ef736f4 100644 --- a/scripts/sglang_disagg/README.MD +++ b/scripts/sglang_disagg/README.MD @@ -185,6 +185,203 @@ curl -X POST http://127.0.0.1:2322/generate \ -d '{"text": "Let me tell you a story", "sampling_params": {"temperature": 0.3}}' ``` +## Agentic Replay Benchmark + +The agentic replay benchmark replays real Claude Code agentic coding traces using the `inferencex-agentx-mvp` scenario from [SemiAnalysisAI aiperf](https://github.com/SemiAnalysisAI/aiperf), driven against the sglang router (`:2322`). This workload is highly prefix-repetitive (each turn re-sends the accumulated context), so it exercises the radix prefix cache and reports a **measured** `gpu_cache_hit_rate` alongside the dataset's theoretical ceiling. + +> **New to AgentX? Start here:** [scripts/common/agentx/README.md](../common/agentx/README.md) has the onboarding path, config schema, and copy-paste [examples](../common/agentx/README.md#examples); [HOW_IT_WORKS.md](../common/agentx/HOW_IT_WORKS.md) explains the replay mechanism. + +It is a drop-in alternative to the default `benchmark_xPyD.sh` concurrency sweep, selected via `BENCHMARK_SCRIPT=agentic`. No GitHub Actions / CI is involved — it runs through the same `run_xPyD_models.slurm` entrypoint as the other benchmarks. + +> **Breaking change:** the legacy `AGENTIC` env shorthand has been removed. Use `BENCHMARK_SCRIPT=agentic` instead. + +The four validated agentic configurations are declared as first-class, `agentic`-tagged `models.json` entries. Discover them with: + +```bash +madengine run --tags agentic +``` + +#### Validated agentic configurations + +| Backend | Model | Connector | Entry Name | +|---------|-------|-----------|------------| +| sglang | Qwen3-32B | mori_io | `pyt_sglang_disagg_mori_io_agentic_qwen3-32b` | +| sglang | DeepSeek-V3 | mori_dp | `pyt_sglang_disagg_mori_dp_agentic_deepseek-v3` | +| vllm | DeepSeek-V3 | mori | `pyt_vllm_disagg_mori_agentic_deepseek-v3` | +| vllm | DeepSeek-V3-5layer | mori | `pyt_vllm_disagg_mori_agentic_deepseek-v3-5layer` | + +### Quick start (canonical DeepSeek-V3 1P/1D) + +```bash +cd scripts/sglang_disagg + +export DOCKER_IMAGE_NAME= # e.g. rocm/pytorch-private:sglang-...-mori-YYYYMMDD +export BENCHMARK_SCRIPT=agentic # selects benchmark_agentic.sh +export RUN_MORI=1 DP_MODE=1 xP=1 yD=1 MODEL_NAME=DeepSeek-V3 +export DURATION=900 AGENTIC_CONC=16 AGENTIC_CACHE_WARMUP_DURATION=300 MAX_MODEL_LEN=160000 + +# 3 same-subnet (10.158.x) nodes with the model on local NVMe +sbatch -N 3 -n 3 -p amd-rccl --time=90 \ + --nodelist=,, \ + run_xPyD_models.slurm +``` + +`BENCHMARK_SCRIPT=agentic` selects `benchmark_agentic.sh`; the in-container launcher then auto-enables server-side Prometheus metrics and strips `--disable-radix-cache` (cache ON). + +### Environment variables + +Full AgentX env reference: [scripts/common/agentx/README.md#environment-variable-reference](../common/agentx/README.md#environment-variable-reference). +The table below lists only the launcher-specific knobs; the shared AgentX +variables (`AGENTIC_CONC`, `AGENTIC_CACHE_WARMUP_DURATION`, `MAX_MODEL_LEN`, +`AGENTIC_CONFIG`, `AGENTIC_WORKLOAD`, `AGENTIC_STRICT_CONTEXT`, +`AGENTIC_HF_ISL_TAIL`, …) are documented centrally. + +| Variable | Default | Description | +|----------|---------|-------------| +| `BENCHMARK_SCRIPT` | `sweep` | Set to `agentic` to run the agentic replay (`sweep` -> `benchmark_xPyD.sh`, `agentic` -> `benchmark_agentic.sh`) | +| `DURATION` | `120` | Benchmark duration in seconds (use `>=900` for a canonical/valid submission) | +| `AGENTIC_PORT` | `2322` | sglang router port aiperf targets | +| `ENABLE_SERVER_METRICS` | auto | `1` forces server Prometheus metrics on (auto-on for agentic) | +| `ENABLE_RADIX_CACHE` | auto | `1` forces radix cache on i.e. strips `--disable-radix-cache` (auto-on for agentic) | +| `SUITE_CORPUS_FORCE` | `0` | Set `1` to regenerate a cached corpus after editing a profile. | + +Note: `AGENTIC_SERVER_METRICS` is set automatically by the launcher to the space-separated server metrics endpoints (`:3000 :3000`) and consumed by aiperf's `--server-metrics`; it is not a boolean toggle. + +For config schema, profile authoring, and Tier 1/Tier 2 knobs, see: +- [AgentX core README](../common/agentx/README.md) +- [Profile authoring guide](../common/agentx/profiles/README.md) +- [Config examples cookbook](../common/agentx/README.md#examples) + +### Running HF (weka) trace datasets on your own model + +The agentic suite (`scripts/common/agentx/`, driven by `benchmark_agentic_suite.sh` via an `AGENTIC_CONFIG=agentic.yaml`) can replay the SemiAnalysis HF weka corpora against *any* served model, not just the 256k/1M reference targets: + +- **Model-agnostic context gating.** Leave `serving.max_model_len` unset or `0` (or `MAX_MODEL_LEN=0`) and the suite resolves the served window automatically. On disaggregated serving the router front-end may not advertise `max_model_len`, so the window is resolved from the prefill worker (the first `host:port` in `AGENTIC_SERVER_METRICS`, auto-derived by the recipe) via `/v1/models` (`data[0].max_model_len`, with a `/get_server_info` fallback); a monolith front-end that does advertise it is auto-detected directly. Each hf workload gets a loader-derived ISL tail (a `_256k` loader → 262144, a full corpus → a conservative ~1M, override with `AGENTIC_HF_ISL_TAIL`). If the tail exceeds the served window the suite **WARNs** and caps `--max-context-length` at the window; `AGENTIC_STRICT_CONTEXT=1` **SKIP**s that workload instead. An explicit `max_model_len` / `MAX_MODEL_LEN` always wins. +- **Tier 1 replay knobs (per workload).** `num_dataset_entries` (how many trace sessions to pull) and `trajectory: {min, max}` (start-window ratio) are per-workload fields; unset keeps today's defaults (`393`, `0.25`/`0.75`). +- **Tier 2 local trim (fit a smaller model).** Add a `filter: {max_isl, max_turns, sample}` block to an hf workload/preset: the corpus is downloaded once, trimmed locally (drop sessions with any turn over `max_isl`, truncate to `max_turns`, randomly `sample` with seed 42), and replayed through `--input-file` — no server-side changes. +- **Reusable presets.** Drop a `source: hf` file in `scripts/common/agentx/profiles/` (see `profiles/inferencex_256k.yaml`) and reference it with `preset: `; it can bundle the loader, Tier 1/Tier 2 knobs, and a concurrency sweep + duration. Entry fields override the preset; env (`AGENTIC_CONC`/`DURATION`/`MAX_MODEL_LEN`) still wins. + +See `scripts/common/agentx/agentic.example.yaml` for a worked config. + +### Supported models and modes + +Agentic replay works for every model in the recipe allowlist. `DP_MODE=1` (MoRI wideEP) is only valid for the DeepSeek models; all others run `DP_MODE=0` (TP-only). + +| Model | `DP_MODE` | Notes | +|-------|-----------|-------| +| DeepSeek-V3 | 1 (wideEP) | Canonical agentic target | +| DeepSeek-R1 | 1 (wideEP) | Same path as DeepSeek-V3 | +| Qwen3-32B | 0 (TP) | | +| Mixtral-8x7B-Instruct-v0.1 | 0 (TP) | | +| Llama-3.1-8B-Instruct | 0 (TP) | Good quick smoke | +| amd-Llama-3.3-70B-Instruct-FP8-KV | 0 (TP) | | +| Llama-3.1-405B-Instruct-FP8-KV | 0 (TP) | Needs larger topology | + +Topology is `xP` prefill + `yD` decode + 1 router node (minimum 3 nodes, i.e. 1P/1D). + +### Invocation examples (agentic suite) + +**Note:** `AGENTIC_CONFIG` paths must be visible inside the container. The `$HOME:$HOME` bind mount means `$HOME/MAD/...` works if MAD is cloned at `$HOME/MAD`. Alternatively use `/opt/mooncake-cookbook/scripts/...` (the in-container `MOONCAKE_REPO_DIR` mount). + +#### 1. Multi-workload suite config + +```bash +cd scripts/sglang_disagg + +export DOCKER_IMAGE_NAME= +export BENCHMARK_SCRIPT=agentic +export RUN_MORI=1 DP_MODE=1 xP=1 yD=1 MODEL_NAME=DeepSeek-V3 +export AGENTIC_CONFIG=$HOME/MAD/scripts/common/agentx/agentic.example.yaml + +sbatch -N 3 -n 3 -p amd-rccl --time=90 \ + --nodelist=,, \ + run_xPyD_models.slurm +``` + +#### 2. Single preset shorthand (no config file) + +```bash +cd scripts/sglang_disagg + +export DOCKER_IMAGE_NAME= +export BENCHMARK_SCRIPT=agentic +export RUN_MORI=1 DP_MODE=1 xP=1 yD=1 MODEL_NAME=DeepSeek-V3 +export AGENTIC_WORKLOAD=conformance_256k # config-less shorthand: runs just this preset + +sbatch -N 3 -n 3 -p amd-rccl --time=90 \ + --nodelist=,, \ + run_xPyD_models.slurm +``` + +#### 3. Longer-context baseline (Case-B 512k window) + +```bash +cd scripts/sglang_disagg + +export DOCKER_IMAGE_NAME= +export BENCHMARK_SCRIPT=agentic +export RUN_MORI=1 DP_MODE=1 xP=1 yD=1 MODEL_NAME=DeepSeek-V3 +export AGENTIC_WORKLOAD=conformance_512k +export MAX_MODEL_LEN=524288 # Case-B ISL tail needs the 512k window + +sbatch -N 3 -n 3 -p amd-rccl --time=90 \ + --nodelist=,, \ + run_xPyD_models.slurm +``` + +#### 4. Single entry selected from a config (smaller model) + +```bash +cd scripts/sglang_disagg + +export DOCKER_IMAGE_NAME= +export BENCHMARK_SCRIPT=agentic +export RUN_MORI=1 DP_MODE=0 xP=1 yD=1 MODEL_NAME=Llama-3.1-8B-Instruct +export AGENTIC_CONFIG=$HOME/MAD/scripts/common/agentx/agentic.example.yaml +export AGENTIC_WORKLOAD=inferencex_small # preset bundles filter: {max_isl, max_turns, sample} + +sbatch -N 3 -n 3 -p amd-rccl --time=90 \ + --nodelist=,, \ + run_xPyD_models.slurm +``` + +#### 5. DRY_RUN preview from the login node (no server) + +```bash +cd $HOME/MAD + +# Prints the resolved N-workload plan + per-workload commands + context verdicts; no server needed. +DRY_RUN=1 AGENTIC_CONFIG=$HOME/MAD/scripts/common/agentx/agentic.example.yaml \ + bash scripts/common/benchmark_agentic_suite.sh +``` + +### Output files + +Written to `/run_logs/$SLURM_JOB_ID/` (host `LOG_PATH//`): + +``` +agentic__xP_yD_.json # aggregate metrics (parsed by validate_agentic_result.sh) +aiperf_artifacts/ # per-request logs, timeslices, server_metrics_export.{csv,json} +``` + +Key fields in the aggregate JSON: +- `request_metrics.cache.theoretical_cache_hit_rate` — dataset-derived ceiling. +- `server_metrics.present` / `server_metrics.cache.gpu_cache_hit_rate` — **measured** runtime GPU cache hit rate. +- `request_accounting` — `records_profiled` vs `records_error_dropped` (error rate is computed over measured requests, excluding warmup). + +### Validation + +`scripts/common/validate_agentic_result.sh` runs automatically at the end of the agentic benchmark and prints, e.g.: + +``` +[validate] requests: 70/72 ok error_rate=2.8% +[validate] gpu_cache_hit_rate: 32.4% +[validate] throughput_per_gpu: 6247 tok/s +[validate][PASS] error_rate within 10% +``` + +Thresholds: `AGENTIC_MAX_ERROR_RATE` (default `0.10`, gates PASS/WARN) and `AGENTIC_MIN_CACHE_HIT` (default `0.30`, warn-only). A low cache hit rate on `xP>1` usually indicates the router is not keeping a conversation's turns on the worker holding its prefix (session-affinity). + ## Known Issues For larger models, such as DeepSeekV3 and Llama-3.1-405B-Instruct-FP8-KV and higher concurrency(512+), errors with below signature is observed:
_' diff --git a/scripts/sglang_disagg/benchmark_agentic.sh b/scripts/sglang_disagg/benchmark_agentic.sh new file mode 100755 index 00000000..be593421 --- /dev/null +++ b/scripts/sglang_disagg/benchmark_agentic.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# Thin shim: execs the shared scripts/common/benchmark_agentic.sh --backend sglang. +set -uo pipefail +_here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +for _cand in \ + "${_here}/../common/benchmark_agentic.sh" \ + "${MOONCAKE_REPO_DIR:-}/../common/benchmark_agentic.sh" \ + "${AGENTIC_LIB:+$(dirname "$AGENTIC_LIB")/benchmark_agentic.sh}"; do + if [ -n "$_cand" ] && [ -f "$_cand" ]; then exec bash "$_cand" --backend sglang "$@"; fi +done +echo "[agentic][ERROR] shared benchmark_agentic.sh not found (set AGENTIC_LIB)" >&2 +exit 1 diff --git a/scripts/sglang_disagg/models.json b/scripts/sglang_disagg/models.json index 31a9fb2f..dfc7caef 100644 --- a/scripts/sglang_disagg/models.json +++ b/scripts/sglang_disagg/models.json @@ -619,5 +619,73 @@ ], "timeout": 14400, "args": "" + }, + { + "name": "pyt_sglang_disagg_mori_io_agentic_qwen3-32b", + "dockerfile": "../../docker/sglang_disagg_inference", + "scripts": "run_xPyD_models.slurm", + "url": "", + "data": "huggingface", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "tags": [ + "pyt", + "sglang", + "sglang_disagg", + "mori_io", + "inference", + "agentic" + ], + "timeout": -1, + "distributed": { + "launcher": "slurm_multi" + }, + "env_vars": { + "DOCKER_IMAGE_NAME": "", + "MODEL_NAME": "Qwen3-32B", + "xP": "1", + "yD": "1", + "DP_MODE": "0", + "RUN_MORI": "1", + "USE_CX7_NICS": "0", + "BENCHMARK_SCRIPT": "agentic", + "AGENTIC_WORKLOAD": "conformance_256k" + }, + "args": "-N 2 -n 2" + }, + { + "name": "pyt_sglang_disagg_mori_dp_agentic_deepseek-v3", + "dockerfile": "../../docker/sglang_disagg_inference", + "scripts": "run_xPyD_models.slurm", + "url": "", + "data": "huggingface", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "tags": [ + "pyt", + "sglang", + "sglang_disagg", + "mori_dp", + "inference", + "agentic" + ], + "timeout": -1, + "distributed": { + "launcher": "slurm_multi" + }, + "env_vars": { + "DOCKER_IMAGE_NAME": "", + "MODEL_NAME": "DeepSeek-V3", + "xP": "1", + "yD": "1", + "DP_MODE": "1", + "RUN_MORI": "1", + "USE_CX7_NICS": "0", + "BENCHMARK_SCRIPT": "agentic", + "AGENTIC_WORKLOAD": "conformance_256k" + }, + "args": "-N 2 -n 2" } ] diff --git a/scripts/sglang_disagg/models.yaml b/scripts/sglang_disagg/models.yaml index 6c16ca62..5a12b999 100644 --- a/scripts/sglang_disagg/models.yaml +++ b/scripts/sglang_disagg/models.yaml @@ -9,7 +9,16 @@ # - base_flags: always applied (prefill + decode) # - tp_flags / dp_flags: mode-level flags applied to BOTH prefill and decode (omit if empty) # - prefill/decode.: role + mode specific flags -# - experimental_flags: optional extra CLI flags (omit if empty) +# - experimental_flags: optional extra CLI flags appended after role/mode flags on BOTH prefill and decode workers. +# Use for try-outs that are independent of TP vs DP (PARALLEL_MODE); omit or "" when unused. +# +# Radix prefix cache: base_flags below carry "--disable-radix-cache" for the default +# xPyD performance sweeps (clean, cache-free throughput). The agentic replay benchmark +# needs the cache ON, so the launcher (sglang_disagg_mori_io_ep.sh) STRIPS +# "--disable-radix-cache" from PREFILL/DECODE configs when BENCHMARK_SCRIPT=agentic +# (or ENABLE_RADIX_CACHE=1). That strip is exercised only by the agentic-tagged +# models.json entries (BENCHMARK_SCRIPT=agentic), not by any perf-sweep entry. Do not +# remove it here, or the perf sweeps lose their cache-free baseline. # --- Dense Llama / Qwen --- diff --git a/scripts/sglang_disagg/run_xPyD_models.slurm b/scripts/sglang_disagg/run_xPyD_models.slurm index b9f8d6b9..59f5e85e 100755 --- a/scripts/sglang_disagg/run_xPyD_models.slurm +++ b/scripts/sglang_disagg/run_xPyD_models.slurm @@ -121,6 +121,20 @@ DP_MODE="${DP_MODE:-0}" #-> mori_io_ep: 1=dp flags + --dp-size/--ep-size; 0=tp f SKIP_BENCHMARK="${SKIP_BENCHMARK:-0}" SKIP_CURL_TEST="${SKIP_CURL_TEST:-0}" +# Benchmark script selector: BENCHMARK_SCRIPT tag -> file run by the launcher. +# sweep (default) -> benchmark_xPyD.sh (random ISL/OSL concurrency sweep) +# agentic -> benchmark_agentic.sh (aiperf agentic trace replay) +# The in-container launcher (sglang_disagg_mori_io_ep.sh) auto-enables server +# metrics + radix prefix cache when BENCHMARK_SCRIPT=agentic. +BENCHMARK_SCRIPT="${BENCHMARK_SCRIPT:-sweep}" +case "$BENCHMARK_SCRIPT" in + sweep) BENCHMARK_SCRIPT_FILE="benchmark_xPyD.sh" ;; + agentic) BENCHMARK_SCRIPT_FILE="benchmark_agentic.sh" ;; + *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, agentic)" >&2; exit 1 ;; +esac +export BENCHMARK_SCRIPT BENCHMARK_SCRIPT_FILE +echo "BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' -> $BENCHMARK_SCRIPT_FILE" + MODEL_NAME="${MODEL_NAME:-None}" MODEL_DIR="${MODEL_DIR:-"/shared_inference/models_blog/"}" @@ -316,14 +330,49 @@ echo "" USER_NAME=$(whoami) MASTER_NODE=$(echo "$SELECTED_NODES" | head -n 1) MASTER_ADDR=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$MASTER_NODE" bash -c 'hostname -I') -MASTER_ADDR=$(echo "$MASTER_ADDR" | awk 'NR==1 {print $1}') +# Prefer the fabric-subnet IP: nodes are multi-homed and picking the first +# hostname -I address can land on a different subnet, which hangs the socket +# barrier on mixed-subnet coordination. The preferred subnet prefix is +# overridable via FABRIC_SUBNET_PREFIX (default "10.158." for the OCI +# useocpm2m fabric); setting it empty is not intended (an empty prefix would +# match every address), so when it is empty we skip the prefix match and go +# straight to the first-address fallback. Fall back to the first address if no +# matching address is present; fail fast if empty. +FABRIC_SUBNET_PREFIX="${FABRIC_SUBNET_PREFIX:-10.158.}" +# Escape regex-special chars so the prefix is matched literally, anchored at start. +FABRIC_SUBNET_RE="^$(printf '%s' "$FABRIC_SUBNET_PREFIX" | sed 's/[.[\*^$]/\\&/g')" +MASTER_ADDR_RAW="$MASTER_ADDR" +if [ -n "$FABRIC_SUBNET_PREFIX" ]; then + MASTER_ADDR=$(echo "$MASTER_ADDR_RAW" | tr ' ' '\n' | grep -E "$FABRIC_SUBNET_RE" | head -1) +else + MASTER_ADDR="" +fi +if [ -z "$MASTER_ADDR" ]; then + MASTER_ADDR=$(echo "$MASTER_ADDR_RAW" | awk '{print $1}') +fi +if [ -z "$MASTER_ADDR" ]; then + echo "ERROR: could not resolve MASTER_ADDR for node '$MASTER_NODE' from 'hostname -I' output: '$MASTER_ADDR_RAW'" >&2 + exit 1 +fi MASTER_PORT=39566 # Choose an open port IPS=() for NODE in $SELECTED_NODES; do IP=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$NODE" bash -c 'hostname -I') - IP=$(echo "$IP" | awk 'NR==1 {print $1}') + IP_RAW="$IP" + if [ -n "$FABRIC_SUBNET_PREFIX" ]; then + IP=$(echo "$IP_RAW" | tr ' ' '\n' | grep -E "$FABRIC_SUBNET_RE" | head -1) + else + IP="" + fi + if [ -z "$IP" ]; then + IP=$(echo "$IP_RAW" | awk '{print $1}') + fi + if [ -z "$IP" ]; then + echo "ERROR: could not resolve IP for node '$NODE' from 'hostname -I' output: '$IP_RAW'" >&2 + exit 1 + fi IPS+=("$IP") done @@ -369,6 +418,9 @@ export SKIP_CURL_TEST export DOCKER_CONT_NAME="container_${MODEL_NAME}_${SLURM_JOB_ID}" export RUN_FILE_FULL="$MOONCAKE_COOKBOOK_PATH/${RUN_FILE}" +# AGENTIC_CONFIG may use $HOME (e.g. from models.json); expand at submit time. +[ -n "${AGENTIC_CONFIG:-}" ] && AGENTIC_CONFIG="$(eval echo "$AGENTIC_CONFIG")" + # Use only the selected nodes for srun execution SELECTED_NODELIST_SRUN=$(echo "$SELECTED_NODES" | paste -sd,) @@ -410,14 +462,40 @@ docker run --rm \ -e BENCHMARK_ITR=$BENCHMARK_ITR \ -e SKIP_BENCHMARK=$SKIP_BENCHMARK \ -e SKIP_CURL_TEST=$SKIP_CURL_TEST \ - -e KV_TRANSFER_BACKEND=${KV_TRANSFER_BACKEND:-} \ + ${KV_TRANSFER_BACKEND:+-e KV_TRANSFER_BACKEND=$KV_TRANSFER_BACKEND} \ -e BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-1024/1024 8192/1024}" \ -e DOCKER_IMAGE_NAME=${DOCKER_IMAGE_NAME:-unknown} \ -e USE_CX7_NICS=${USE_CX7_NICS:-0} \ + ${BENCHMARK_SCRIPT:+-e BENCHMARK_SCRIPT=$BENCHMARK_SCRIPT} \ + ${BENCHMARK_SCRIPT_FILE:+-e BENCHMARK_SCRIPT_FILE=$BENCHMARK_SCRIPT_FILE} \ + -e MOONCAKE_REPO_DIR="$MOONCAKE_REPO_DIR" \ + ${AGENTIC_CONC:+-e AGENTIC_CONC=$AGENTIC_CONC} \ + ${DURATION:+-e DURATION=$DURATION} \ + ${AGENTIC_PORT:+-e AGENTIC_PORT=$AGENTIC_PORT} \ + ${MODEL_PREFIX:+-e MODEL_PREFIX=$MODEL_PREFIX} \ + ${WEKA_LOADER_OVERRIDE:+-e WEKA_LOADER_OVERRIDE=$WEKA_LOADER_OVERRIDE} \ + ${AIPERF_UNSAFE_OVERRIDE:+-e AIPERF_UNSAFE_OVERRIDE=$AIPERF_UNSAFE_OVERRIDE} \ + ${AGENTIC_CACHE_WARMUP_DURATION:+-e AGENTIC_CACHE_WARMUP_DURATION=$AGENTIC_CACHE_WARMUP_DURATION} \ + ${MAX_MODEL_LEN:+-e MAX_MODEL_LEN=$MAX_MODEL_LEN} \ + ${AGENTIC_CONFIG:+-e AGENTIC_CONFIG=$AGENTIC_CONFIG} \ + ${AGENTIC_WORKLOAD:+-e AGENTIC_WORKLOAD=$AGENTIC_WORKLOAD} \ + ${SUITE_CORPUS_DIR:+-e SUITE_CORPUS_DIR=$SUITE_CORPUS_DIR} \ + ${AIPERF_PIN:+-e AIPERF_PIN=$AIPERF_PIN} \ + ${AGENTIC_UTILS_PIN:+-e AGENTIC_UTILS_PIN=$AGENTIC_UTILS_PIN} \ + ${SUITE_CORPUS_FORCE:+-e SUITE_CORPUS_FORCE=$SUITE_CORPUS_FORCE} \ + ${AGENTIC_STRICT_CONTEXT:+-e AGENTIC_STRICT_CONTEXT=$AGENTIC_STRICT_CONTEXT} \ + ${AGENTIC_HF_ISL_TAIL:+-e AGENTIC_HF_ISL_TAIL=$AGENTIC_HF_ISL_TAIL} \ + ${HF_HOME:+-e HF_HOME=$HF_HOME} \ + ${ENABLE_SERVER_METRICS:+-e ENABLE_SERVER_METRICS=$ENABLE_SERVER_METRICS} \ + ${ENABLE_RADIX_CACHE:+-e ENABLE_RADIX_CACHE=$ENABLE_RADIX_CACHE} \ + ${DRY_RUN:+-e DRY_RUN=$DRY_RUN} \ + ${AGENTIC_ALLOW_UV_INSTALL:+-e AGENTIC_ALLOW_UV_INSTALL=$AGENTIC_ALLOW_UV_INSTALL} \ + ${AGENTIC_UV_VERSION:+-e AGENTIC_UV_VERSION=$AGENTIC_UV_VERSION} \ --ulimit nofile=1048576:1048576 \ --name $DOCKER_CONT_NAME \ --entrypoint /bin/bash \ $DOCKER_IMAGE_NAME -c " + set -o pipefail mkdir -p /run_logs/${SLURM_JOB_ID} $RUN_FILE_FULL 2>&1 | tee /run_logs/${SLURM_JOB_ID}/pd_sglang_bench_serving.sh_NODE${SLURM_PROCID}.log " diff --git a/scripts/sglang_disagg/salloc_launch.sh b/scripts/sglang_disagg/salloc_launch.sh index 5064e221..95f42e2e 100644 --- a/scripts/sglang_disagg/salloc_launch.sh +++ b/scripts/sglang_disagg/salloc_launch.sh @@ -15,3 +15,18 @@ export xP=2; export yD=2; export MODEL_NAME=amd-Llama-3.3-70B-Instruct-FP8-KV; #export DOCKER_IMAGE_NAME= #export xP=; export yD=; export MODEL_NAME=Llama-3.1-8B-Instruct; sbatch -N -n --nodelist= run_xPyD_models.slurm +# === Agentic replay benchmark (aiperf inferencex-agentx-mvp) === +# Selected via BENCHMARK_SCRIPT=agentic. +# BENCHMARK_SCRIPT=agentic auto-enables server metrics (gpu_cache_hit_rate) + radix prefix cache. +# +# DeepSeek-V3 1P/1D (canonical, DP_MODE=1 wideEP): +#export DOCKER_IMAGE_NAME= +#export BENCHMARK_SCRIPT=agentic RUN_MORI=1 DP_MODE=1 xP=1 yD=1 MODEL_NAME=DeepSeek-V3 +#export DURATION=900 AGENTIC_CONC=16 AGENTIC_CACHE_WARMUP_DURATION=300 MAX_MODEL_LEN=160000 +#sbatch -N 3 -n 3 -p amd-rccl --nodelist=<3-nodes> run_xPyD_models.slurm +# +# Dense model 1P/1D (TP-only, DP_MODE=0) quick smoke: +#export BENCHMARK_SCRIPT=agentic RUN_MORI=1 DP_MODE=0 xP=1 yD=1 MODEL_NAME=Llama-3.1-8B-Instruct +#export DURATION=120 AGENTIC_CONC=8 +#sbatch -N 3 -n 3 -p amd-rccl --nodelist=<3-nodes> run_xPyD_models.slurm + diff --git a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh index e8ab7e4c..cfd043fc 100755 --- a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh +++ b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh @@ -258,7 +258,26 @@ PREFILL_MODEL_CONFIG="${MODEL_BASE_FLAGS} ${MODEL_ARCH_FLAGS} ${MODEL_MODE_FLAGS DECODE_MODEL_CONFIG="${MODEL_BASE_FLAGS} ${MODEL_ARCH_FLAGS} ${MODEL_MODE_FLAGS} ${MODEL_DECODE_FLAGS} ${MODEL_EXPERIMENTAL_FLAGS}" echo "Using model-specific configuration for: $MODEL_NAME (mode=${PARALLEL_MODE})" -export PREFILL_MODEL_CONFIG DECODE_MODEL_CONFIG MODEL_EXPERIMENTAL_FLAGS +# Agentic gating: the default concurrency sweep keeps the perf-tuned config +# (models.yaml base_flags, radix cache off). Only the agentic trace-replay path +# (BENCHMARK_SCRIPT_FILE=benchmark_agentic.sh) — or an explicit override — turns +# on the radix prefix cache and server-side Prometheus metrics. This keeps the +# default path byte-for-byte unchanged. +AGENTIC_METRICS_ENABLED=0 +if [[ "${BENCHMARK_SCRIPT:-}" == "agentic" || "${ENABLE_SERVER_METRICS:-0}" == "1" ]]; then + AGENTIC_METRICS_ENABLED=1 +fi +SERVER_METRICS_FLAGS="" +if [[ "${AGENTIC_METRICS_ENABLED}" == "1" ]]; then + SERVER_METRICS_FLAGS="--enable-metrics --enable-metrics-for-all-schedulers" +fi +if [[ "${BENCHMARK_SCRIPT:-}" == "agentic" || "${ENABLE_RADIX_CACHE:-0}" == "1" ]]; then + PREFILL_MODEL_CONFIG="${PREFILL_MODEL_CONFIG//--disable-radix-cache/}" + DECODE_MODEL_CONFIG="${DECODE_MODEL_CONFIG//--disable-radix-cache/}" + echo "[radix] radix prefix cache ENABLED (stripped --disable-radix-cache) for agentic/ENABLE_RADIX_CACHE" +fi + +export PREFILL_MODEL_CONFIG DECODE_MODEL_CONFIG MODEL_EXPERIMENTAL_FLAGS SERVER_METRICS_FLAGS # shellcheck disable=SC1091 source "${SCRIPT_DIR}/mori_ep_env.sh" @@ -768,15 +787,32 @@ PY fi benchmark_status=0 + # Server-side Prometheus metrics for the agentic replay: point aiperf at the + # sglang servers' :3000/metrics directly (the router exposes Prometheus on a + # separate port and lacks gpu_cache_hit_rate). Gated on the agentic path so + # the default sweep is unaffected. AGENTIC_SERVER_METRICS is consumed by + # scripts/common/agentic_lib.sh (build_replay_cmd -> aiperf --server-metrics). + if [[ "${AGENTIC_METRICS_ENABLED}" == "1" \ + && -n "${IP_FIRST_PREFILL:-}" && -n "${IP_FIRST_DECODE:-}" ]]; then + export AGENTIC_SERVER_METRICS="${AGENTIC_SERVER_METRICS:-${IP_FIRST_PREFILL}:3000 ${IP_FIRST_DECODE}:3000}" + echo "[metrics] AGENTIC_SERVER_METRICS=${AGENTIC_SERVER_METRICS}" + echo "=== server /metrics reachability check ===" + curl -sf "http://${IP_FIRST_PREFILL}:3000/metrics" | head -3 || echo "PREFILL metrics UNREACHABLE" + curl -sf "http://${IP_FIRST_DECODE}:3000/metrics" | head -3 || echo "DECODE metrics UNREACHABLE" + fi if [[ "${SKIP_BENCHMARK:-0}" != "1" ]] && [[ -n "${MOONCAKE_COOKBOOK_PATH:-}" ]]; then - if [[ -f "${MOONCAKE_COOKBOOK_PATH}/benchmark_xPyD.sh" ]]; then - echo "Running ${MOONCAKE_COOKBOOK_PATH}/benchmark_xPyD.sh" + # Benchmark hook is selectable: default random-sweep benchmark_xPyD.sh, + # or set BENCHMARK_SCRIPT_FILE=benchmark_agentic.sh for aiperf agentic + # trace replay (see scripts/common/agentic_lib.sh). + _bench_file="${BENCHMARK_SCRIPT_FILE:-benchmark_xPyD.sh}" + if [[ -f "${MOONCAKE_COOKBOOK_PATH}/${_bench_file}" ]]; then + echo "Running ${MOONCAKE_COOKBOOK_PATH}/${_bench_file}" ( cd "${MOONCAKE_COOKBOOK_PATH}" || exit 1 - bash benchmark_xPyD.sh + bash "${_bench_file}" ) || benchmark_status=$? else - echo "WARN: benchmark_xPyD.sh not found under MOONCAKE_COOKBOOK_PATH=${MOONCAKE_COOKBOOK_PATH}" >&2 + echo "WARN: ${_bench_file} not found under MOONCAKE_COOKBOOK_PATH=${MOONCAKE_COOKBOOK_PATH}" >&2 benchmark_status=1 fi fi @@ -828,6 +864,7 @@ elif [[ "$NODE_RANK" -ge 1 && "$NODE_RANK" -lt "$xP" ]]; then PREFILL_CMD+=" \ --decode-log-interval 1 \ ${PREFILL_MODEL_CONFIG} \ + ${SERVER_METRICS_FLAGS} \ --log-level-http warning" export PREFILL_CMD PREFILL_NODE_RANK @@ -920,6 +957,7 @@ elif [[ "$NODE_RANK" -ge $xP && "$NODE_RANK" -le $((xP + yD - 1)) ]]; then DECODE_CMD+=" \ --decode-log-interval 1 \ ${DECODE_MODEL_CONFIG} \ + ${SERVER_METRICS_FLAGS} \ --log-level-http warning" export DECODE_CMD DECODE_NODE_RANK diff --git a/scripts/vllm/configs/default.yaml b/scripts/vllm/configs/default.yaml index 887d3e9a..e9dd9662 100644 --- a/scripts/vllm/configs/default.yaml +++ b/scripts/vllm/configs/default.yaml @@ -95,6 +95,7 @@ VLLM_USE_BREAKABLE_CUDAGRAPH: 0 extra_args: --kv-cache-dtype: fp8 + --max-num-batched-tokens: 8192 - benchmark: serving model: @@ -164,4 +165,4 @@ VLLM_ROCM_USE_AITER: 1 arch_overrides: gfx942: - dtype: float16 \ No newline at end of file + dtype: float16 diff --git a/scripts/vllm/configs/extended.yaml b/scripts/vllm/configs/extended.yaml index 95897795..57fcbd79 100644 --- a/scripts/vllm/configs/extended.yaml +++ b/scripts/vllm/configs/extended.yaml @@ -60,6 +60,8 @@ max_concurrency: 1 env: VLLM_ROCM_USE_AITER: 1 + extra_args: + --enable-expert-parallel: True - benchmark: serving model: @@ -77,4 +79,4 @@ --attention-backend: ROCM_AITER_FA arch_overrides: gfx942: - dtype: float16 \ No newline at end of file + dtype: float16 diff --git a/scripts/vllm_dissag/README.MD b/scripts/vllm_dissag/README.MD index 17660abe..7b10df30 100644 --- a/scripts/vllm_dissag/README.MD +++ b/scripts/vllm_dissag/README.MD @@ -40,8 +40,9 @@ neither → combo 1. Use these exact `MODEL_NAME` values (defined in `models.yaml` + the slurm allowlists). The **Combos** column lists which of the 4 valid combos above each model may run in; the -slurm gate rejects any other pairing. Dense models are TP-only; the DeepSeek family is -**wideEP-only** (TP is rejected — its recipe needs the wideEP serve path). +slurm gate rejects any other pairing. Dense models are TP-only; the DeepSeek family and +`GLM-5.1-FP8` are **wideEP-only** (TP is rejected — their recipes need the wideEP serve +path). | Model | Type | Combos (by #) | Notes | |-------|------|---------------|-------| @@ -53,6 +54,7 @@ slurm gate rejects any other pairing. Dense models are TP-only; the DeepSeek fam | `DeepSeek-V3` | MoE | 3, 4 (wideEP only) | validated serving (moriio+MoRI-EP) | | `DeepSeek-V3-5layer` | MoE | 3, 4 (wideEP only) | smoke/dev variant | | `DeepSeek-R1` | MoE | 3, 4 (wideEP only) | needs a co-versioned AITER image (see Caveats) | +| `GLM-5.1-FP8` | MoE + DSA | 3 (moriio wideEP only) | own image (`glmv5.1` Dockerfile); 2P/2D EP16 and 4P/4D EP32 validated on the current pins (NIAH ~93-97%) | ## Prerequisites @@ -64,7 +66,7 @@ slurm gate rejects any other pairing. Dense models are TP-only; the DeepSeek fam ## Building the Docker image -No public prebuilt image ships — build your own from the one fullsource Dockerfile (MoRI + +No public prebuilt image ships — build your own from the fullsource Dockerfile (MoRI + AITER + vLLM + vllm-router + optional UCX/RIXL/rocSHMEM/DeepEP, from Git pins on an open ROCm ci_base): @@ -74,9 +76,24 @@ docker build -t /vllm-disagg:local \ -f docker/vllm_disagg_inference.ubuntu.amd.Dockerfile . ``` -One image serves all connectors. `WITH_NIXL=1` (default) builds every transport (moriio -TP/wideEP + rixl NIXL TP + DeepEP wideEP); add `--build-arg WITH_NIXL=0` for a lean -MoRI-EP-only image (moriio + deepep-from-base, faster build). +That image serves all connectors for every model in the table above **except** +`GLM-5.1-FP8`. `WITH_NIXL=1` (default) builds every transport (moriio TP/wideEP + rixl +NIXL TP + DeepEP wideEP); add `--build-arg WITH_NIXL=0` for a lean MoRI-EP-only image +(moriio + deepep-from-base, faster build). + +`GLM-5.1-FP8` needs its own per-model image, because its vLLM/AITER/MoRI pins differ +from the DeepSeek-V3/R1 stack the shared Dockerfile is pinned to: + +```bash +docker build -t /vllm-disagg:glmv5.1 \ + -f docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile . +``` + +That Dockerfile defaults to `WITH_NIXL=0` (MoRI-EP only, which is all GLM-5.1 uses) and +its `VLLM_REF` is a hard requirement, not a preference — MAD ships no runtime patchers, +so every GLM DSA fix has to be in-source in that image. Add a +`vllm_disagg_inference..ubuntu.amd.Dockerfile` per future model that needs its +own pins rather than repinning the shared one. Then `export DOCKER_IMAGE_NAME=` (or a published image) before submitting. The slurm `docker pull`s the image on every node, so a local-only tag must be pushed to a registry the nodes @@ -158,6 +175,11 @@ Env precedence (lowest → highest): connector default → `models.yaml` `env:` `-e VAR=...`. So a model can override a connector default (e.g. a `MORI_*` knob) without editing any script, and a per-run/per-cluster `-e` still wins. +**GPU_MEMORY_UTILIZATION**: If not set in yaml or submit-time, the launcher defaults to **0.70** at +EP≥32 (to accommodate MoRI buffer scaling at high EP width) and **0.8** otherwise. Models needing +topology-dependent values should omit this key from yaml and rely on the fallback; models needing a +static override can specify it in their `env:` block. + ## Node topology (all modes) ``` @@ -234,13 +256,184 @@ lengths and the model is scored on how many it retrieves. Select it in place of sweep via the launcher's benchmark hook: ```bash -export BENCHMARK_SCRIPT_FILE=benchmark_niah.sh +export BENCHMARK_SCRIPT=niah export NIAH_WORDS="2000,8000,20000,35000" # context sizes (words); optional ``` Or run `benchmark_niah.py` standalone against any live OpenAI-compatible endpoint (`NIAH_URL`, `NIAH_MODEL`). +## Agentic Replay Benchmark + +The agentic replay benchmark replays real Claude Code agentic coding traces using the `inferencex-agentx-mvp` scenario from [SemiAnalysisAI aiperf](https://github.com/SemiAnalysisAI/aiperf), driven against the vLLM router/proxy. This workload is highly prefix-repetitive (each turn re-sends the accumulated context), so it exercises the vLLM prefix cache and reports a **measured** `gpu_cache_hit_rate` alongside the dataset's theoretical ceiling. It shares the endpoint-agnostic harness (`scripts/common/agentic_lib.sh` + the AgentX suite driver) with the [sglang integration](../sglang_disagg/README.MD#agentic-replay-benchmark). + +> **New to AgentX? Start here:** [scripts/common/agentx/README.md](../common/agentx/README.md) has the onboarding path, config schema, and copy-paste [examples](../common/agentx/README.md#examples); [HOW_IT_WORKS.md](../common/agentx/HOW_IT_WORKS.md) explains the replay mechanism. + +It is a drop-in alternative to the default `benchmark_xPyD.sh` concurrency sweep, selected via `BENCHMARK_SCRIPT=agentic`. No GitHub Actions / CI is involved — it runs through the same `run_xPyD_models.slurm` entrypoint as the other benchmarks. + +> **Breaking change:** the legacy `AGENTIC` env shorthand has been removed. Use `BENCHMARK_SCRIPT=agentic` instead. + +The four validated agentic configurations are declared as first-class, `agentic`-tagged `models.json` entries. Discover them with: + +```bash +madengine run --tags agentic +``` + +#### Validated agentic configurations + +| Backend | Model | Connector | Entry Name | +|---------|-------|-----------|------------| +| sglang | Qwen3-32B | mori_io | `pyt_sglang_disagg_mori_io_agentic_qwen3-32b` | +| sglang | DeepSeek-V3 | mori_dp | `pyt_sglang_disagg_mori_dp_agentic_deepseek-v3` | +| vllm | DeepSeek-V3 | mori | `pyt_vllm_disagg_mori_agentic_deepseek-v3` | +| vllm | DeepSeek-V3-5layer | mori | `pyt_vllm_disagg_mori_agentic_deepseek-v3-5layer` | + +**vLLM specifics.** aiperf targets `AGENTIC_PORT`, which the launcher auto-resolves to the proxy/router port `BENCHMARK_PORT` (set by `connector_start_proxy`: `ROUTER_PORT` for `vllm_router`, `PROXY_PORT` for the toy proxy). Because the vLLM PD router 503s on `GET /v1/models` under MoRIIO service discovery (empty HTTP worker registry) while the shared harness gates readiness/served-model on `/v1/models`, the launcher starts a tiny stdlib shim (`agentic_models_shim.py`) on `BENCHMARK_PORT+1` that answers `/v1/models` (gated on the router's `/health`) and stream-proxies everything else; `AGENTIC_PORT` is then pointed at the shim. Both are byte-for-byte no-ops for the default sweep (guarded by `BENCHMARK_SCRIPT=agentic`). + +### Quick start (canonical DeepSeek-V3 1P/1D, MoRI-EP) + +```bash +cd scripts/vllm_dissag + +export DOCKER_IMAGE_NAME= +export BENCHMARK_SCRIPT=agentic # selects benchmark_agentic.sh +export RUN_MORI=1 xP=1 yD=1 MODEL_NAME=DeepSeek-V3 +export DURATION=900 AGENTIC_CONC=16 AGENTIC_CACHE_WARMUP_DURATION=300 MAX_MODEL_LEN=160000 + +# 3 nodes with the model on local NVMe (router+prefill co-located on node 0, decode on node 1) +sbatch -N 2 -n 2 -p amd-rccl --time=90 \ + --nodelist=, \ + run_xPyD_models.slurm +``` + +`BENCHMARK_SCRIPT=agentic` selects `benchmark_agentic.sh`; the connectors then auto-enable server-side Prometheus metrics and strip `--no-enable-prefix-caching` (prefix cache ON). + +### Environment variables + +Full AgentX env reference: [scripts/common/agentx/README.md#environment-variable-reference](../common/agentx/README.md#environment-variable-reference). +The table below lists only the launcher-specific knobs; the shared AgentX +variables (`AGENTIC_CONC`, `AGENTIC_CACHE_WARMUP_DURATION`, `MAX_MODEL_LEN`, +`AGENTIC_CONFIG`, `AGENTIC_WORKLOAD`, `AGENTIC_STRICT_CONTEXT`, +`AGENTIC_HF_ISL_TAIL`, …) are documented centrally. On this disaggregated +serving the shim front-end may not advertise `max_model_len`, so the window is +resolved from the prefill worker (the first `host:port` in +`AGENTIC_SERVER_METRICS`, auto-derived by the recipe) via `/v1/models`; pin +`MAX_MODEL_LEN` if autodetect returns 0. + +| Variable | Default | Description | +|----------|---------|-------------| +| `BENCHMARK_SCRIPT` | `sweep` | Set to `agentic` to run the agentic replay (`sweep`, `long_context`, `agentic`) | +| `DURATION` | `120` | Benchmark duration in seconds (use `>=900` for a canonical/valid submission) | +| `AGENTIC_PORT` | `${BENCHMARK_PORT}` | Endpoint aiperf targets (auto-resolved to the proxy/router port, then the shim port) | +| `ENABLE_SERVER_METRICS` | auto | `1` forces server Prometheus metrics on (auto-on for agentic) | +| `ENABLE_PREFIX_CACHE` | auto | `1` forces prefix cache on i.e. strips `--no-enable-prefix-caching` (auto-on for agentic) | +| `SUITE_CORPUS_FORCE` | `0` | Set `1` to regenerate a cached corpus after editing a profile. | + +Note: `AGENTIC_SERVER_METRICS` is set automatically by the connector to the space-separated backend server `/metrics` endpoints (`: :`, e.g. `:20005` for moriio, `:2584`/`SERVER_PORT` for rixl) and consumed by aiperf's `--server-metrics`; it is not a boolean toggle. Its first `host:port` (the prefill worker) is also probed for the served `max_model_len` when the front-end doesn't advertise it. + +For config schema, profile authoring, and Tier 1/Tier 2 knobs, see: +- [AgentX core README](../common/agentx/README.md) +- [Profile authoring guide](../common/agentx/profiles/README.md) +- [Config examples cookbook](../common/agentx/README.md#examples) + +### Supported models and modes + +Agentic replay works for every model in the recipe allowlist, in the same combos each model supports (see **Supported models** above). The DeepSeek family is wideEP-only; dense/MoE models such as Qwen3-32B run TP. + +| Model | Combos | Notes | +|-------|--------|-------| +| DeepSeek-V3 | 3, 4 (wideEP) | Canonical agentic target | +| DeepSeek-R1 | 3, 4 (wideEP) | Same path as DeepSeek-V3 | +| Qwen3-32B | 1, 2 (TP) | Ships a chat_template; good smoke target | +| Qwen3-30B-A3B | 1, 2 (TP) | MoE, needs a co-versioned AITER image | +| amd-Llama-3.3-70B-Instruct-FP8-KV | 1, 2 (TP) | | +| Llama-3.1-405B-Instruct-FP8-KV | 1, 2 (TP) | Needs larger topology | + +> aiperf drives `/v1/chat/completions`, so the served model must ship a `chat_template` in its `tokenizer_config.json` (Qwen2.5/Qwen3-Instruct do; base Llama-3.1-8B does not). + +### Invocation examples (agentic suite) + +**Note:** `AGENTIC_CONFIG` paths must be visible inside the container. The `$HOME:$HOME` bind mount means `$HOME/MAD/...` works if MAD is cloned at `$HOME/MAD`. Alternatively use the in-container `NIXL_COOKBOOK_PATH` mount (`$NIXL_REPO_DIR` on the host). + +#### 1. Multi-workload suite config + +```bash +cd scripts/vllm_dissag + +export DOCKER_IMAGE_NAME= +export BENCHMARK_SCRIPT=agentic RUN_MORI=1 xP=1 yD=1 MODEL_NAME=DeepSeek-V3 +export AGENTIC_CONFIG=$HOME/MAD/scripts/common/agentx/agentic.example.yaml + +sbatch -N 2 -n 2 -p amd-rccl --time=90 --nodelist=<2-nodes> run_xPyD_models.slurm +``` + +#### 2. Single preset shorthand (no config file) + +```bash +cd scripts/vllm_dissag + +export DOCKER_IMAGE_NAME= +export BENCHMARK_SCRIPT=agentic CONNECTOR=rixl WIDE_EP=0 xP=1 yD=1 MODEL_NAME=Qwen3-32B +export AGENTIC_WORKLOAD=conformance_256k # config-less shorthand: runs just this preset + +sbatch -N 2 -n 2 -p amd-rccl --time=90 --nodelist=<2-nodes> run_xPyD_models.slurm +``` + +#### 3. Single entry selected from a config (smaller model, Tier-2 filter) + +```bash +cd scripts/vllm_dissag + +export DOCKER_IMAGE_NAME= +export BENCHMARK_SCRIPT=agentic CONNECTOR=rixl WIDE_EP=0 xP=1 yD=1 MODEL_NAME=Qwen3-32B +export AGENTIC_CONFIG=$HOME/MAD/scripts/common/agentx/agentic.example.yaml +export AGENTIC_WORKLOAD=inferencex_small # preset bundles filter: {max_isl, max_turns, sample} + +sbatch -N 2 -n 2 -p amd-rccl --time=90 --nodelist=<2-nodes> run_xPyD_models.slurm +``` + +#### 4. DRY_RUN preview from the login node (no server) + +```bash +cd $HOME/MAD + +# Prints the resolved N-workload plan + per-workload commands + context verdicts; no server needed. +DRY_RUN=1 AGENTIC_CONFIG=$HOME/MAD/scripts/common/agentx/agentic.example.yaml \ + bash scripts/common/benchmark_agentic_suite.sh +``` + +### Output files + +Written to `/run_logs/$SLURM_JOB_ID/` (host `LOG_PATH//`): + +``` +agentic__xP_yD_.json # single-replay aggregate (validate_agentic_result.sh) +suite_summary.json # multi-workload suite roll-up (workloads[]) +/[conc/] # per-workload (and per-concurrency) result dirs +aiperf_artifacts/ # per-request logs, timeslices, server_metrics_export.{csv,json} +agentic_models_shim_NODE0.log # side-port /v1/models shim log +``` + +Key fields in the aggregate JSON: +- `request_metrics.cache.theoretical_cache_hit_rate` — dataset-derived ceiling. +- `server_metrics.present` / `server_metrics.cache.gpu_cache_hit_rate` — **measured** runtime GPU cache hit rate. +- `request_accounting` — `records_profiled` vs `records_error_dropped` (error rate is computed over measured requests, excluding warmup). + +The suite roll-up (`suite_summary.json`) has a `workloads[]` array; each entry carries `workload`, `result_dir`, `status`, `theoretical_cache_hit_rate`, `gpu_cache_hit_rate`, and `error_rate`. A workload that exits non-zero is marked `INVALID` (with a `RUN_INVALID.json` sidecar in its dir) and the suite **continues** to the next workload. + +### Validation + +`scripts/common/validate_agentic_result.sh` runs automatically at the end of the agentic benchmark and prints, e.g.: + +``` +[validate] requests: 70/72 ok error_rate=2.8% +[validate] gpu_cache_hit_rate: 32.4% +[validate] throughput_per_gpu: 6247 tok/s +[validate][PASS] error_rate within 10% +``` + +Thresholds: `AGENTIC_MAX_ERROR_RATE` (default `0.10`, gates PASS/WARN) and `AGENTIC_MIN_CACHE_HIT` (default `0.30`, warn-only). A low cache hit rate on `xP>1` usually indicates the router is not keeping a conversation's turns on the worker holding its prefix (session-affinity). + ## Testing the launcher (no GPUs) Offline suites validate the launcher without a cluster — run after any change to the diff --git a/scripts/vllm_dissag/agentic_models_shim.py b/scripts/vllm_dissag/agentic_models_shim.py new file mode 100755 index 00000000..018cf241 --- /dev/null +++ b/scripts/vllm_dissag/agentic_models_shim.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +"""Agentic models shim for the vLLM disaggregated PD router. + +The vLLM production PD router (vllm-router) serves /v1/chat/completions and +/v1/completions but returns 503 "No prefill servers available" on /v1/models +when workers register via MoRIIO service discovery (the HTTP worker registry +stays empty). The shared agentic harness (scripts/common/agentic_lib.sh) gates +readiness + served-model resolution on GET /v1/models, so it never starts. + +This tiny shim (stdlib only, agentic path only) sits on a side port and: + * GET /v1/models -> 200 with the served model id, ONLY once the upstream + router answers GET /health 200 (so it doubles as the + readiness gate the harness expects). + * GET /health -> mirror upstream /health. + * everything else -> stream-proxied verbatim to the upstream router + (POST /v1/chat/completions etc., SSE-safe). + +Env: + AGENTIC_SHIM_PORT listen port (required) + AGENTIC_SHIM_UPSTREAM upstream router host:port (default 127.0.0.1:30000) + AGENTIC_SHIM_MODEL served model id to advertise on /v1/models (required) +""" +import http.client +import json +import os +import sys +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +LISTEN_PORT = int(os.environ["AGENTIC_SHIM_PORT"]) +UPSTREAM = os.environ.get("AGENTIC_SHIM_UPSTREAM", "127.0.0.1:30000") +MODEL = os.environ.get("AGENTIC_SHIM_MODEL", "") +PREFILL = os.environ.get("AGENTIC_SHIM_PREFILL", "").strip() # host:port of prefill backend (diag only) +UP_HOST, UP_PORT = UPSTREAM.split(":") +UP_PORT = int(UP_PORT) +_HOP = {"connection", "keep-alive", "proxy-authenticate", "proxy-authorization", + "te", "trailers", "transfer-encoding", "upgrade", "host", "content-length"} + + +def _health_ok(host, port): + try: + c = http.client.HTTPConnection(host, port, timeout=5) + c.request("GET", "/health") + r = c.getresponse() + r.read() + c.close() + return r.status == 200 + except Exception: + return False + + +def _upstream_health_ok(): + # Readiness signal: the production vllm-router serves GET /health (200 when + # workers are registered). The MoRIIO toy proxy does NOT implement /health + # (404), so fall back to the prefill backend's /health (a real vLLM OpenAI + # server, 200 when the engine is up) when a PREFILL backend is configured. + if _health_ok(UP_HOST, UP_PORT): + return True + if PREFILL: + try: + h, pt = PREFILL.split(":") + return _health_ok(h, int(pt)) + except Exception: + return False + return False + + +class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *a): # quiet + pass + + def _send_json(self, code, obj): + body = json.dumps(obj).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.send_header("Connection", "close") + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if self.path.rstrip("/") == "/v1/models": + if _upstream_health_ok(): + self._send_json(200, {"object": "list", "data": [ + {"id": MODEL, "object": "model", "owned_by": "vllm"}]}) + else: + self._send_json(503, {"error": "router not ready"}) + return + if self.path.rstrip("/") == "/health": + self._send_json(200 if _upstream_health_ok() else 503, {"status": "ok"}) + return + self._proxy("GET") + + def do_POST(self): + self._proxy("POST") + + def _replay_backend(self, method, body, headers): + if not PREFILL: + return + try: + h, pt = PREFILL.split(":") + c = http.client.HTTPConnection(h, int(pt), timeout=60) + c.request(method, self.path, body=body, headers=headers) + r = c.getresponse(); b = r.read(); c.close() + print(f"[agentic-shim][diag] backend {PREFILL} {self.path} -> {r.status}: " + f"{b[:800].decode('utf-8','replace')}", flush=True) + except Exception as e: + print(f"[agentic-shim][diag] backend replay failed: {e}", flush=True) + + def _proxy(self, method): + length = int(self.headers.get("Content-Length", 0) or 0) + body = self.rfile.read(length) if length else b"" + headers = {k: v for k, v in self.headers.items() if k.lower() not in _HOP} + try: + conn = http.client.HTTPConnection(UP_HOST, UP_PORT, timeout=3600) + conn.request(method, self.path, body=body, headers=headers) + resp = conn.getresponse() + except Exception as e: + self._send_json(502, {"error": f"upstream proxy failed: {e}"}) + return + # Surface upstream error bodies (e.g. backend 400s) for diagnosis; these are + # small non-streaming JSON responses, so read fully, log, and relay verbatim. + if resp.status >= 400: + err = resp.read() + print(f"[agentic-shim] upstream {method} {self.path} -> {resp.status} " + f"(req_bytes={length}): {err[:600].decode('utf-8', 'replace')}", flush=True) + if resp.status >= 500 and self.path.rstrip('/').endswith('/chat/completions'): + self._replay_backend(method, body, headers) + self.send_response(resp.status) + for k, v in resp.getheaders(): + if k.lower() in _HOP: + continue + self.send_header(k, v) + self.send_header("Content-Length", str(len(err))) + self.send_header("Connection", "close") + self.end_headers() + self.wfile.write(err) + conn.close() + return + # The NIXL toy_proxy relays streamed SSE chunks but labels them + # Content-Type: application/json, so aiperf never strips the "data:" SSE + # framing and marks every request invalid. Force text/event-stream when the + # client asked for a stream (JSON body "stream": true) OR the upstream is + # already SSE; leave non-streaming JSON responses untouched. + want_stream = False + try: + want_stream = bool(json.loads(body or b"{}").get("stream")) + except Exception: + want_stream = False + up_ctype = (resp.getheader("Content-Type", "") or "").lower() + force_sse = want_stream or "text/event-stream" in up_ctype + self.send_response(resp.status) + for k, v in resp.getheaders(): + if k.lower() in _HOP: + continue + if force_sse and k.lower() == "content-type": + continue + self.send_header(k, v) + if force_sse: + self.send_header("Content-Type", "text/event-stream") + self.send_header("Connection", "close") + self.end_headers() + try: + while True: + chunk = resp.read(8192) + if not chunk: + break + self.wfile.write(chunk) + self.wfile.flush() + except Exception: + pass + finally: + conn.close() + + +if __name__ == "__main__": + if not MODEL: + print("[agentic-shim][ERROR] AGENTIC_SHIM_MODEL must be set", file=sys.stderr) + sys.exit(2) + srv = ThreadingHTTPServer(("0.0.0.0", LISTEN_PORT), Handler) + print(f"[agentic-shim] listening :{LISTEN_PORT} -> {UPSTREAM} (model={MODEL})", flush=True) + srv.serve_forever() diff --git a/scripts/vllm_dissag/benchmark_agentic.sh b/scripts/vllm_dissag/benchmark_agentic.sh new file mode 100755 index 00000000..f3619cc5 --- /dev/null +++ b/scripts/vllm_dissag/benchmark_agentic.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# Thin shim: execs the shared scripts/common/benchmark_agentic.sh --backend vllm. +set -uo pipefail +_here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +for _cand in \ + "${_here}/../common/benchmark_agentic.sh" \ + "${NIXL_COOKBOOK_PATH:-}/../common/benchmark_agentic.sh" \ + "${NIXL_REPO_DIR:-}/../common/benchmark_agentic.sh" "${AGENTIC_LIB:+$(dirname "$AGENTIC_LIB")/benchmark_agentic.sh}"; do + if [ -n "$_cand" ] && [ -f "$_cand" ]; then exec bash "$_cand" --backend vllm "$@"; fi +done +echo "[agentic][ERROR] shared benchmark_agentic.sh not found (set AGENTIC_LIB)" >&2 +exit 1 diff --git a/scripts/vllm_dissag/benchmark_niah.py b/scripts/vllm_dissag/benchmark_niah.py index 0cdd027e..43d81850 100755 --- a/scripts/vllm_dissag/benchmark_niah.py +++ b/scripts/vllm_dissag/benchmark_niah.py @@ -8,7 +8,15 @@ # NIAH_MODEL model name/tag the server serves (required — the served path) # NIAH_WORDS comma list of context sizes in words (default 2000,8000,20000,35000) # NIAH_MAXTOK max_tokens for the answer (default 2048) +# NIAH_SEEDS comma list of needle-layout seeds (default 0,1,2); summary reports +# mean/min/max across seeds to separate real accuracy from variance # NIAH_TIMEOUT per-request timeout seconds (default 1800) +# NIAH_WARMUP 1 (default) = send one throwaway request per context length BEFORE +# scoring, so the first-hit JIT/kernel-autotune compile happens outside +# the scored/gated window. On a freshly-booted node the first request of +# a shape can take minutes to compile; without warmup that lands on the +# first scored request -> false 0/10 or timeout. Warmup failures are +# tolerated (logged, not fatal). Set 0 to disable. import os, sys, json, random, urllib.request URL = os.environ.get("NIAH_URL", "http://127.0.0.1:30000/v1/chat/completions") @@ -16,6 +24,14 @@ WORDS = [int(x) for x in os.environ.get("NIAH_WORDS", "2000,8000,20000,35000").split(",") if x.strip()] MAXTOK = int(os.environ.get("NIAH_MAXTOK", "2048")) TIMEOUT = float(os.environ.get("NIAH_TIMEOUT", "1800")) +# Needle layout is seeded, so a single run is deterministic (bit-exact repro on the +# same stack). Run multiple seeds to distinguish real accuracy from single-needle +# variance; the summary reports mean/min/max across seeds. Default 0,1,2. +SEEDS = [int(x) for x in os.environ.get("NIAH_SEEDS", "0,1,2").split(",") if x.strip()] +WARMUP = os.environ.get("NIAH_WARMUP", "1") == "1" +# Warmup uses a generous timeout (cold compile of a long-context shape can take minutes) +# and never fails the run — its only job is to trigger compilation before scoring. +WARMUP_TIMEOUT = max(TIMEOUT, 1800.0) FILLER = ( "table chair window bottle pencil garden river mountain coffee planet " @@ -41,27 +57,53 @@ def make_haystack(n_words, seed=0): return " ".join(words) -def run(n_words): +def _request(n_words, seed, max_tokens, timeout): + """POST one NIAH request; return (message_dict, error_str). Exactly one is non-None.""" body = { "model": MODEL, "messages": [ {"role": "system", "content": SYSTEM}, - {"role": "user", "content": "Find the animals in this list:\n\n" + make_haystack(n_words)}, + {"role": "user", "content": "Find the animals in this list:\n\n" + make_haystack(n_words, seed)}, ], "temperature": 0.0, - "max_tokens": MAXTOK, + "max_tokens": max_tokens, + # Thinking models (e.g. GLM-5.1) emit chain-of-thought into a separate + # reasoning field and leave `content` empty until the final answer; with a + # small max_tokens the answer never appears in `content` and the score is a + # false 0/10. Disable thinking so the answer lands in `content` directly. + "chat_template_kwargs": {"enable_thinking": False}, } data = json.dumps(body).encode() req = urllib.request.Request(URL, data=data, headers={"Content-Type": "application/json"}) try: - with urllib.request.urlopen(req, timeout=TIMEOUT) as r: - msg = json.loads(r.read())["choices"][0]["message"] + with urllib.request.urlopen(req, timeout=timeout) as r: + return json.loads(r.read())["choices"][0]["message"], None except Exception as e: - print("words=%6d ERROR %s" % (n_words, e), flush=True) + return None, str(e) + + +def warmup(n_words): + """One throwaway request per length so first-hit compile happens off the scored path. + Never fatal: a warmup timeout just means the shape is still compiling; the scored + request will pay whatever remains (bounded by NIAH_TIMEOUT).""" + _, err = _request(n_words, seed=0, max_tokens=8, timeout=WARMUP_TIMEOUT) + status = "ok" if err is None else ("timeout/err: %s" % err) + print("words=%6d [warmup] %s" % (n_words, status), flush=True) + + +def run(n_words, seed=0): + # Sentinel: None = timeout/transport error (NOT a wrong answer); int = score 0..10. + msg, err = _request(n_words, seed, MAXTOK, TIMEOUT) + if err is not None: + print("words=%6d seed=%d TIMEOUT/ERROR %s" % (n_words, seed, err), flush=True) return None - text = ((msg.get("content") or "") + " " + (msg.get("reasoning_content") or "")).lower() + # Score content plus any reasoning field (some servers surface CoT as + # `reasoning` or `reasoning_content`) so a thinking model is never mis-scored. + text = ((msg.get("content") or "") + " " + + (msg.get("reasoning_content") or "") + " " + + (msg.get("reasoning") or "")).lower() found = sorted(a for a in ANIMALS if a in text) - print("words=%6d found=%2d/10 %s" % (n_words, len(found), found), flush=True) + print("words=%6d seed=%d found=%2d/10 %s" % (n_words, seed, len(found), found), flush=True) return len(found) @@ -70,14 +112,29 @@ def main(): print("NIAH_MODEL must be set (the served model path/name)", file=sys.stderr) sys.exit(2) print("=== NIAH retrieval test ===", flush=True) - print("url=%s model=%s sizes=%s" % (URL, MODEL, WORDS), flush=True) - results = {} + print("url=%s model=%s sizes=%s seeds=%s warmup=%s" % (URL, MODEL, WORDS, SEEDS, WARMUP), flush=True) + # Warmup pass: compile every shape once before scoring, so cold JIT never lands on a + # scored/gated request (the common cause of false 0/10 or timeout on a fresh boot). + if WARMUP: + print("=== NIAH warmup (one throwaway request per length) ===", flush=True) + for n in WORDS: + warmup(n) + results = {} # n_words -> list of scores across seeds (None = timeout/error, not a wrong answer) for n in WORDS: - results[n] = run(n) - print("=== NIAH summary ===", flush=True) + results[n] = [run(n, s) for s in SEEDS] + print("=== NIAH summary (mean/min/max across %d seed(s)) ===" % len(SEEDS), flush=True) for n in WORDS: - v = results[n] - print(" words=%6d found=%s/10" % (n, "ERR" if v is None else v), flush=True) + scored = results[n] + vals = [v for v in scored if v is not None] + n_to = sum(1 for v in scored if v is None) # timeouts/errors, excluded from mean + if not vals: + print(" words=%6d NO-RESULT (%d/%d timed out or errored — likely cold compile; " + "raise NIAH_TIMEOUT or keep NIAH_WARMUP=1)" % (n, n_to, len(scored)), flush=True) + continue + mean = sum(vals) / len(vals) + extra = (" [%d timeout/err excluded]" % n_to) if n_to else "" + print(" words=%6d mean=%.1f/10 min=%d max=%d (n=%d)%s" + % (n, mean, min(vals), max(vals), len(vals), extra), flush=True) if __name__ == "__main__": diff --git a/scripts/vllm_dissag/benchmark_niah.sh b/scripts/vllm_dissag/benchmark_niah.sh index ba49a359..247152e1 100755 --- a/scripts/vllm_dissag/benchmark_niah.sh +++ b/scripts/vllm_dissag/benchmark_niah.sh @@ -14,15 +14,35 @@ LOG="/run_logs/${SLURM_JOB_ID}/niah_${SLURM_JOB_ID}_${timestamp}_xP${xP}_yD${yD} echo "==== NIAH long-context retrieval test ====" echo "port=${BENCHMARK_PORT} model=${MODEL_PATH} sizes=${NIAH_WORDS:-2000,8000,20000,35000}" -# Give the router a moment to be fully ready for chat completions. -sleep 10 +# Wait until the router actually serves before starting (replaces a blind sleep). On a +# fresh boot the router may register a few seconds after the workers report ready; poll +# /v1/models until it answers, up to ~5 min. Non-fatal: fall through if the probe can't +# confirm (the harness's own warmup + timeout still protect the run). +_ready=0 +for _i in $(seq 1 60); do + if curl -s -o /dev/null -w '%{http_code}' --max-time 5 \ + "http://127.0.0.1:${BENCHMARK_PORT}/v1/models" 2>/dev/null | grep -q '^200$'; then + _ready=1; echo "[niah] router ready after ~$((_i*5))s"; break + fi + sleep 5 +done +[ "$_ready" = 1 ] || echo "[niah] WARN: router readiness not confirmed in 300s; proceeding (warmup will absorb)" # The server registers the model under its path (served_model_name = MODEL_PATH). +# NIAH_WARMUP=1 (harness default): first-hit JIT compiles off the scored path so a cold +# boot does not produce false 0/10 or timeouts on the first scored request. NIAH_URL="http://127.0.0.1:${BENCHMARK_PORT}/v1/chat/completions" \ NIAH_MODEL="${MODEL_PATH}" \ NIAH_WORDS="${NIAH_WORDS:-2000,8000,20000,35000}" \ +NIAH_SEEDS="${NIAH_SEEDS:-0,1,2}" \ NIAH_MAXTOK="${NIAH_MAXTOK:-2048}" \ NIAH_TIMEOUT="${NIAH_TIMEOUT:-1800}" \ +NIAH_WARMUP="${NIAH_WARMUP:-1}" \ python3 "${DIR}/benchmark_niah.py" 2>&1 | tee -a "${LOG}" +# Generate madengine perf.csv rows from NIAH results (mirrors benchmark_xPyD.sh) +python3 "$NIXL_COOKBOOK_PATH/parse_to_csv.py" "${LOG}" --niah \ + --perf-csv /run_logs/${SLURM_JOB_ID}/perf.csv --model-name "${MODEL_NAME}" \ + 2>&1 | tee -a "${LOG}" + echo "NIAH results -> ${LOG}" diff --git a/scripts/vllm_dissag/benchmark_xPyD.sh b/scripts/vllm_dissag/benchmark_xPyD.sh index b8851d24..068181d9 100755 --- a/scripts/vllm_dissag/benchmark_xPyD.sh +++ b/scripts/vllm_dissag/benchmark_xPyD.sh @@ -40,6 +40,33 @@ for i in $(seq 1 $BENCHMARK_ITR); do echo "Running the benchserving script for iter: $i" | tee -a ${LOG}_CONCURRENCY.log >/dev/null for combo in "${COMBINATIONS[@]}"; do IFS="/" read -r isl osl <<< "$combo" + # Per-shape warmup at the REAL isl/osl, low concurrency: the global warmup above is + # isl=osl=32/con=1, so it never exercises this shape's prefill path, its Triton/aiter + # kernel variants or the decode cudagraph batch sizes. DEFAULT OFF -- every model + # shares this sweep path, and an A/B at 1024/1024 con=8 measured it neutral, + # so recipes validated without it must not be shifted for + # no gain. GLM opts in via its models.yaml env:, which is what its published latency + # numbers were measured under. Enable per-run with SHAPE_WARMUP=1. + if [[ "${SHAPE_WARMUP:-0}" == "1" ]]; then + _w_con="${SHAPE_WARMUP_CON:-4}" + _w_prompts="${SHAPE_WARMUP_PROMPTS:-8}" + echo "[WARMUP] shape isl $isl osl $osl con ${_w_con} prompts ${_w_prompts}" \ + | tee -a ${LOG}_CONCURRENCY.log >/dev/null + timeout "${SHAPE_WARMUP_TIMEOUT:-2400}" vllm bench serve \ + --model $MODEL_PATH \ + --backend vllm \ + --host 127.0.0.1 \ + --port $BENCHMARK_PORT \ + --dataset-name "random" \ + --random-input-len $isl \ + --random-output-len $osl \ + --random-prefix-len 0 \ + --num-prompts ${_w_prompts} \ + --request-rate "inf" \ + --ignore-eos \ + --max-concurrency ${_w_con} \ + 2>&1 | tee -a ${LOG}_SHAPEWARMUP.log >/dev/null + fi for con in $CON; do p_con=$(($con * 2)) if [ "$p_con" -lt 16 ]; then diff --git a/scripts/vllm_dissag/connectors/moriio.sh b/scripts/vllm_dissag/connectors/moriio.sh index 44c64d68..3eb1094f 100644 --- a/scripts/vllm_dissag/connectors/moriio.sh +++ b/scripts/vllm_dissag/connectors/moriio.sh @@ -162,12 +162,18 @@ connector_launch_worker() { else _cudagraph_mode="${PREFILL_CUDAGRAPH_MODE:-$_cudagraph_mode}" fi + # use_inductor_graph_partition=true moves graph partitioning from Dynamo/FX to + # inductor codegen, splitting at cudagraph_unsafe ops (incl. the MLA KV-update) so + # they run as eager boundaries. Default OFF: enabling it here would change + # --compilation-config for EVERY model. GLM opts in via its models.yaml env:. + local _igp_json="" + [[ "${USE_INDUCTOR_GRAPH_PARTITION:-0}" == "1" ]] && _igp_json=',"use_inductor_graph_partition":true' if [[ -n "$_cudagraph_mode" && "$_cudagraph_mode" != "NONE" ]]; then local _capture_sizes="${CUDAGRAPH_CAPTURE_SIZES:-1 2 4 8 16 32 64 128 256}" - exec_args+=(--compilation-config '{"cudagraph_mode":"'"${_cudagraph_mode}"'","custom_ops":["+quant_fp8"]}') + exec_args+=(--compilation-config '{"cudagraph_mode":"'"${_cudagraph_mode}"'","custom_ops":["+quant_fp8"]'"${_igp_json}"'}') exec_args+=(--cudagraph-capture-sizes ${_capture_sizes}) else - exec_args+=(--compilation-config '{"cudagraph_mode":"NONE","custom_ops":["+quant_fp8"]}') + exec_args+=(--compilation-config '{"cudagraph_mode":"NONE","custom_ops":["+quant_fp8"]'"${_igp_json}"'}') fi # Per-model flags from models.yaml (driver-exported; empty if none). @@ -175,6 +181,16 @@ connector_launch_worker() { local _mc; if [[ "$log_prefix" == "prefill" ]]; then _mc="${MODEL_CONFIG_PREFILL:-}"; else _mc="${MODEL_CONFIG_DECODE:-}"; fi [[ -n "$_mc" ]] && eval "model_args=(${_mc})" + # Agentic gating: the default sweep keeps prefix caching OFF (clean, cache-free + # throughput) via the hardcoded --no-enable-prefix-caching below. The agentic + # trace-replay path (BENCHMARK_SCRIPT_FILE=benchmark_agentic.sh) — or an explicit + # ENABLE_PREFIX_CACHE=1 — STRIPS that flag so prefix caching is ON. Gated so the + # default (non-agentic) sweep argv is byte-for-byte unchanged. + local _prefix_cache_flag="--no-enable-prefix-caching" + if [[ "${BENCHMARK_SCRIPT:-}" == "agentic" || "${ENABLE_PREFIX_CACHE:-0}" == "1" ]]; then + _prefix_cache_flag="" + fi + if parallelism_is_wide_ep; then # ---- WIDE_EP=1 (MoriEP) ---- # Per-role all2all: prefill=high_throughput, decode=low_latency. The @@ -215,11 +231,11 @@ connector_launch_worker() { "${mem_args[@]}" \ --kv-cache-dtype "${_kvdtype}" \ --block-size "${_block}" \ - --no-enable-prefix-caching \ + ${_prefix_cache_flag} \ --all2all-backend "${_all2all}" \ --trust-remote-code \ --distributed-timeout-seconds "${DISTRIBUTED_TIMEOUT_SECONDS:-7200}" \ - "${exec_args[@]}" "${extra_args[@]}" "${kv_args[@]}" + "${exec_args[@]}" "${extra_args[@]}" "${kv_args[@]}" "${model_args[@]}" WORKER_PID=0; return 0 fi @@ -235,13 +251,14 @@ connector_launch_worker() { "${mem_args[@]}" \ --kv-cache-dtype "${_kvdtype}" \ --block-size "${_block}" \ - --no-enable-prefix-caching \ + ${_prefix_cache_flag} \ --all2all-backend "${_all2all}" \ --trust-remote-code \ --distributed-timeout-seconds ${DISTRIBUTED_TIMEOUT_SECONDS:-7200} \ "${exec_args[@]}" \ "${extra_args[@]}" \ "${kv_args[@]}" \ + "${model_args[@]}" \ 2>&1 | tee /run_logs/${SLURM_JOB_ID}/${log_prefix}_NODE${NODE_RANK}.log >/dev/null & WORKER_PID=$! return 0 @@ -313,6 +330,14 @@ connector_start_proxy() { # moriio_toy: in-image toy proxy; resolves the script across the online_serving/ # -> disaggregated/ path move. # Sets BENCHMARK_PORT (router->ROUTER_PORT, toy->PROXY_PORT) for the driver. + # Agentic replay: point aiperf at the backend vLLM servers' /metrics (SERVE_PORT) + # for prefill+decode masters so it can scrape gpu cache-hit / throughput. Gated on + # the agentic path so the default sweep is unaffected. Consumed by + # scripts/common/agentic_lib.sh (build_replay_cmd -> aiperf --server-metrics). + if [[ "${BENCHMARK_SCRIPT:-}" == "agentic" || "${ENABLE_SERVER_METRICS:-0}" == "1" ]]; then + export AGENTIC_SERVER_METRICS="${AGENTIC_SERVER_METRICS:-${PREFILL_MASTER_ADDR}:${SERVE_PORT} ${DECODE_MASTER_ADDR}:${SERVE_PORT}}" + echo "[metrics] AGENTIC_SERVER_METRICS=${AGENTIC_SERVER_METRICS}" + fi sleep 10 if [ "$PROXY_TYPE" == "vllm_router" ]; then local PREFILL_URL="http://${PREFILL_MASTER_ADDR}:${SERVE_PORT}" diff --git a/scripts/vllm_dissag/connectors/rixl.sh b/scripts/vllm_dissag/connectors/rixl.sh index e6dda9b9..4302c562 100644 --- a/scripts/vllm_dissag/connectors/rixl.sh +++ b/scripts/vllm_dissag/connectors/rixl.sh @@ -29,7 +29,12 @@ connector_init() { # rixl/TP ports SERVER_PORT=2584; SERVE_PORT="${SERVER_PORT}" KV_PORT=14600 - CONTAINER_BARRIER_PORT=5000 + # Container-creation barrier port. Env-overridable (BARRIER_PORT) so it can + # be moved off the collision-prone default 5000: the launcher's `fuser -k` + # cleanup targets this port on the host (host networking), so a stale host + # service on 5000 would otherwise be killed. Residual risk: the host-side + # fuser still kills whatever holds this port for the launching user. + CONTAINER_BARRIER_PORT="${BARRIER_PORT:-5000}" fi PROXY_TYPE="${PROXY_TYPE:-vllm_router}" @@ -273,6 +278,16 @@ _rixl_launch_deepep() { connector_setup_env "${backend}" + # Agentic gating: the default sweep keeps prefix caching OFF via the hardcoded + # --no-enable-prefix-caching below. The agentic trace-replay path + # (BENCHMARK_SCRIPT_FILE=benchmark_agentic.sh) — or ENABLE_PREFIX_CACHE=1 — + # STRIPS that flag so prefix caching is ON. Gated so the default (non-agentic) + # sweep argv is byte-for-byte unchanged. + local _prefix_cache_flag="--no-enable-prefix-caching" + if [[ "${BENCHMARK_SCRIPT:-}" == "agentic" || "${ENABLE_PREFIX_CACHE:-0}" == "1" ]]; then + _prefix_cache_flag="" + fi + local extra_args=() if [[ "$role" == "master" ]]; then extra_args+=(--api-server-count=8 --data-parallel-start-rank 0) @@ -309,7 +324,7 @@ _rixl_launch_deepep() { --data-parallel-rpc-port "${RPC_PORT}" \ --master-addr "${dp_addr}" \ "${compile_args[@]}" \ - --no-enable-prefix-caching --block-size 1 \ + ${_prefix_cache_flag} --block-size 1 \ --gpu-memory-utilization 0.8 \ --kv-cache-dtype fp8 \ --enable-expert-parallel \ @@ -331,7 +346,7 @@ _rixl_launch_deepep() { --data-parallel-rpc-port "${RPC_PORT}" \ --master-addr "${dp_addr}" \ "${compile_args[@]}" \ - --no-enable-prefix-caching --block-size 1 \ + ${_prefix_cache_flag} --block-size 1 \ --gpu-memory-utilization 0.8 \ --kv-cache-dtype fp8 \ --enable-expert-parallel \ @@ -358,6 +373,13 @@ connector_wait_workers_ready() { connector_start_proxy() { local PREFILL_ARGS="" DECODE_ARGS="" PREFILL_PORTS="" DECODE_PORTS="" i + # Agentic replay: point aiperf at the backend vLLM servers' /metrics (SERVER_PORT) + # for prefill+decode masters. Gated on the agentic path so the default sweep is + # unaffected. Consumed by scripts/common/agentic_lib.sh (aiperf --server-metrics). + if [[ "${BENCHMARK_SCRIPT:-}" == "agentic" || "${ENABLE_SERVER_METRICS:-0}" == "1" ]]; then + export AGENTIC_SERVER_METRICS="${AGENTIC_SERVER_METRICS:-${PREFILL_MASTER_ADDR}:${SERVER_PORT} ${DECODE_MASTER_ADDR}:${SERVER_PORT}}" + echo "[metrics] AGENTIC_SERVER_METRICS=${AGENTIC_SERVER_METRICS}" + fi for ((i=0; i", + "MODEL_NAME": "GLM-5.1-FP8", + "xP": "1", + "yD": "1", + "RUN_MORI": "1", + "RUN_DEEPEP": "0", + "BENCHMARK_COMBINATIONS": "1024/1024" + }, + "args": "-N 2 -n 2" + }, + { + "name": "pyt_vllm_disagg_mori_glm-5.1-fp8_niah", + "dockerfile": "../../docker/vllm_disagg_inference.glmv5.1", + "scripts": "run_xPyD_models.slurm", + "url": "", + "data": "huggingface", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "tags": [ + "pyt", + "vllm", + "vllm_disagg", + "mori_ep", + "inference", + "niah" + ], + "timeout": -1, + "distributed": { + "launcher": "slurm_multi" + }, + "env_vars": { + "DOCKER_IMAGE_NAME": "", + "MODEL_NAME": "GLM-5.1-FP8", + "xP": "1", + "yD": "1", + "RUN_MORI": "1", + "RUN_DEEPEP": "0", + "BENCHMARK_SCRIPT": "niah" + }, + "args": "-N 2 -n 2" + }, { "name": "pyt_vllm_disagg_deepep_deepseek-v3", "dockerfile": "../../docker/vllm_disagg_inference", @@ -370,5 +433,71 @@ "BENCHMARK_COMBINATIONS": "1024/1024" }, "args": "-N 2 -n 2" + }, + { + "name": "pyt_vllm_disagg_mori_agentic_deepseek-v3", + "dockerfile": "../../docker/vllm_disagg_inference", + "scripts": "run_xPyD_models.slurm", + "url": "", + "data": "huggingface", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "tags": [ + "pyt", + "vllm", + "vllm_disagg", + "mori_ep", + "inference", + "agentic" + ], + "timeout": -1, + "distributed": { + "launcher": "slurm_multi" + }, + "env_vars": { + "DOCKER_IMAGE_NAME": "", + "MODEL_NAME": "DeepSeek-V3", + "xP": "1", + "yD": "1", + "RUN_MORI": "1", + "RUN_DEEPEP": "0", + "BENCHMARK_SCRIPT": "agentic", + "AGENTIC_WORKLOAD": "conformance_256k" + }, + "args": "-N 2 -n 2" + }, + { + "name": "pyt_vllm_disagg_mori_agentic_deepseek-v3-5layer", + "dockerfile": "../../docker/vllm_disagg_inference", + "scripts": "run_xPyD_models.slurm", + "url": "", + "data": "huggingface", + "n_gpus": "-1", + "owner": "mad.support@amd.com", + "training_precision": "", + "tags": [ + "pyt", + "vllm", + "vllm_disagg", + "mori_ep", + "inference", + "agentic" + ], + "timeout": -1, + "distributed": { + "launcher": "slurm_multi" + }, + "env_vars": { + "DOCKER_IMAGE_NAME": "", + "MODEL_NAME": "DeepSeek-V3-5layer", + "xP": "1", + "yD": "1", + "RUN_MORI": "1", + "RUN_DEEPEP": "0", + "BENCHMARK_SCRIPT": "agentic", + "AGENTIC_WORKLOAD": "conformance_256k" + }, + "args": "-N 2 -n 2" } ] diff --git a/scripts/vllm_dissag/models.yaml b/scripts/vllm_dissag/models.yaml index 23d66059..2321fe13 100644 --- a/scripts/vllm_dissag/models.yaml +++ b/scripts/vllm_dissag/models.yaml @@ -13,11 +13,13 @@ # overrides connector defaults but still yields to a submit-time # `-e VAR=...` (which wins). This is the SINGLE HOME for the model # serving recipe (KV_BLOCK_SIZE, KV_CACHE_DTYPE, per-role -# PREFILL/DECODE_CUDAGRAPH_MODE, *_MORI_BACKEND, GPU_MEMORY_UTILIZATION, -# KV_CACHE_MEMORY_BYTES, VLLM_ROCM_USE_AITER_MLA, ...). The Docker image -# ships NONE of these — it is a clean binary/library artifact, so the -# same image serves any model/cluster. (Platform + MoRI fabric env lives -# in connectors/.env; both are forwarded via `docker -e`.) +# PREFILL/DECODE_CUDAGRAPH_MODE, *_MORI_BACKEND, KV_CACHE_MEMORY_BYTES, +# VLLM_ROCM_USE_AITER_MLA, ...). GPU_MEMORY_UTILIZATION is optional here; +# if omitted, the launcher provides a topology-aware fallback (0.70 at +# EP>=32, 0.8 otherwise). The Docker image ships NONE of these — it is a +# clean binary/library artifact, so the same image serves any model/cluster. +# (Platform + MoRI fabric env lives in connectors/.env; both +# are forwarded via `docker -e`.) # # BOUNDARY — the launcher/connector owns the wideEP parallelism DEGREE and transport: # --data-parallel-size / --enable-expert-parallel / --all2all-backend / @@ -32,6 +34,14 @@ # today, so the dp: blocks below are intentionally empty — placeholders so per-model # dp: tuning can be added later without launcher edits (both connectors now append it). +# Prefix cache (agentic): the default xPyD perf sweeps run cache-free — the +# connectors hardcode "--no-enable-prefix-caching" on the emitted vllm serve argv +# (NOT a yaml flag). The agentic trace-replay path re-enables prefix caching by +# STRIPPING that flag in the connector (moriio.sh / rixl.sh) when BENCHMARK_SCRIPT=agentic +# (or ENABLE_PREFIX_CACHE=1). That strip is exercised only by the agentic-tagged +# models.json entries (BENCHMARK_SCRIPT=agentic), not by any perf-sweep entry; there is +# nothing to set here. Do not add a prefix-cache flag to these entries. + # ============================ Dense / TP-capable ============================ Llama-3.1-405B-Instruct-FP8-KV: @@ -143,7 +153,6 @@ _deepseek_recipe_env: &deepseek_recipe_env KV_BLOCK_SIZE: "16" KV_CACHE_DTYPE: "fp8" KV_CACHE_MEMORY_BYTES: "20000000000" - GPU_MEMORY_UTILIZATION: "0.80" VLLM_CUDAGRAPH_MODE: "PIECEWISE" PREFILL_CUDAGRAPH_MODE: "NONE" DECODE_CUDAGRAPH_MODE: "PIECEWISE" @@ -178,3 +187,97 @@ DeepSeek-R1: dp: "" decode: dp: "" + +# ============================ MoE + DSA (wideEP only) ============================ + +# GLM-5.1-FP8 (zai-org/GLM-5.1-FP8, arch GlmMoeDsaForCausalLM): MLA + DeepSeek +# Sparse Attention (DSA). 78 layers (3 dense + 75 MoE), 256 routed experts top-8 +# + 1 shared, FP8 block 128. wideEP-only (see WIDE_EP_ONLY_MODELS in the slurm). +# +# IMAGE CONTRACT: config only, NO vLLM source fixes. The DSA enablement (upstream +# #47766, DSA indexer KV transfer over MoRIIO) must be IN-SOURCE in the image's vLLM. +# Build from docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile, whose VLLM_REF +# pin IS the contract. DOCKER_IMAGE_NAME is submit-time, not set here. +# +# NOT SUPPORTED: images predating the in-source DSA fixes boot and serve but return +# silently wrong output (rocm/pytorch-private:glm-dockerimage-built-09072026 scored NIAH +# 2k 0/10). There is no runtime patcher and no fallback — rebuild at the pinned VLLM_REF. +# +# GLM vs DeepSeek (both MLA MoE; GLM is DSA-sparse): +# - KV_BLOCK_SIZE=1 DSA sparse indexer REQUIRES block-size 1 (DS uses 16) +# - VLLM_ROCM_USE_AITER_MLA=1 AITER sparse MLA path ON (DS sets 0) +# - prefill cudagraph NONE capture deadlocks on this stack; decode captures fine +# block=1 and AITER_MLA=1 are already the moriio.sh defaults; set explicitly so the +# recipe survives a connector default change. +# +# dp_flags carry the GLM tool/reasoning parsers and apply to both roles; the dp: blocks +# stay empty like the DeepSeek family because the recipe comes from env:. +# +# Long-context: the earlier ~18k collapse is fixed in-source (#47766 persistent sparse-MLA +# metadata), not by config. Left at native max_model_len; capping was tested, no benefit. +GLM-5.1-FP8: + env: + VLLM_USE_V1: "1" + # v0.27 wraps layer_name in a torch OpaqueBase for the unified MLA custom ops, and + # the boxing FAILS on this ROCm torch build -> "RuntimeError: unknown parameter type" + # on the first real MLA decode (the compile path returns early, so it passes + # boot+warmup then crashes on the first request). 0 restores the plain-string path. + VLLM_USE_LAYERNAME: "0" + VLLM_ROCM_USE_AITER: "1" + VLLM_ROCM_USE_AITER_RMSNORM: "1" + VLLM_ROCM_USE_AITER_MLA: "1" + KV_BLOCK_SIZE: "1" + KV_CACHE_DTYPE: "fp8" + VLLM_CUDAGRAPH_MODE: "NONE" + PREFILL_CUDAGRAPH_MODE: "NONE" + # FULL graph for the uniform-decode shapes, PIECEWISE for everything else. Chosen + # over plain PIECEWISE because it measured faster for decode on MI300X; the gain + # narrows as EP widens, since the cross-node all2all inside the captured step does + # not compress. Costs a longer boot (capture is slower and the captured graphs hold + # more memory than PIECEWISE) - notable for CI. Set PIECEWISE to fall back. + DECODE_CUDAGRAPH_MODE: "FULL_AND_PIECEWISE" + CUDAGRAPH_CAPTURE_SIZES: "1 2 4 8 16 32 64 128 256" + # Moves graph partitioning to inductor codegen (MLA KV-update becomes an eager + # boundary). The connector default is OFF, so GLM opts in here and no other model's + # --compilation-config changes. Not a correctness requirement - an A/B was accuracy- + # neutral and within noise on latency - but kept ON because it is the configuration + # this recipe was exercised with. + USE_INDUCTOR_GRAPH_PARTITION: "1" + VLLM_ALL2ALL_BACKEND: "mori_high_throughput" + PREFILL_MORI_BACKEND: "mori_high_throughput" + DECODE_MORI_BACKEND: "mori_low_latency" + # MoRI EP shared-memory heap (16 GiB). The dispatch/combine width itself is set by + # decode.dp below, not by this. + MORI_SHMEM_HEAP_SIZE: "17179869184" + # DSA sparse-indexer logits-buffer cap (crash fix). The indexer only sub-chunks the + # query dim when M*N*4 exceeds this budget; the 512MB default lets an 8192-token + # prefill launch fp8_mqa_logits at grid=(8192,), which HARD-FAULTS the worker on + # gfx942 (silent GPU fault -> DP group collapse -> 503 at >=8k prompts). 64MB forces + # M-dim sub-chunking (~2k tokens/chunk) so the buffer and launch stay bounded. + VLLM_SPARSE_INDEXER_MAX_LOGITS_MB: "64" + # NCCL heartbeat watchdog: at long context (>~8k) a DP rank's sparse-MLA/MoE all2all + # can exceed the default HeartbeatMonitor timeout, so the monitor declares the rank + # dead and tears down the whole process group -> prefill EngineCore crash -> 503. + # Disable the monitor-triggered teardown and extend the timeouts. + TORCH_NCCL_ENABLE_MONITORING: "0" + TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC: "1800" + TORCH_NCCL_DUMP_ON_TIMEOUT: "0" + TORCH_NCCL_BLOCKING_WAIT: "0" + TORCH_NCCL_ASYNC_ERROR_HANDLING: "1" + NCCL_IB_TIMEOUT: "22" + # Per-shape warmup before each measured isl/osl cell, which the sweep's global 32/32 + # con=1 warmup does not cover. Kept because the latency figures above were measured + # with it on, not because it was shown to help: an A/B at 1024/1024 con=8 was neutral. + # The script default is OFF, so no other model's sweep changes. + SHAPE_WARMUP: "1" + dp_flags: "--tool-call-parser glm47 --reasoning-parser glm45 --enable-auto-tool-choice --chat-template-content-format string" + prefill: + dp: "" + decode: + # PERF: the MoRI EP dispatch width comes from max_num_batched_tokens, so the decode + # role otherwise runs an 8192-token-wide all2all every step, which dominates decode + # latency. It cannot be shrunk via env: mori derives recv capacity from the send + # width, so the buffer must still cover vLLM's profiling dummy run. Keep >= typical + # prompt length: smaller values cut per-token latency but inflate TTFT sharply; + # 2048 was the balance point. + dp: "--max-num-batched-tokens 2048" diff --git a/scripts/vllm_dissag/parse_to_csv.py b/scripts/vllm_dissag/parse_to_csv.py index e772e394..c67425ed 100644 --- a/scripts/vllm_dissag/parse_to_csv.py +++ b/scripts/vllm_dissag/parse_to_csv.py @@ -143,6 +143,69 @@ def _get_run_metadata(pipeline: str = "vllm"): } +def parse_niah_log(log_file: str) -> Dict[int, Dict]: + """Parse NIAH benchmark log file and extract retrieval results per context length. + + Scans for summary lines emitted by benchmark_niah.py: + words= 2000 mean=9.7/10 min=9 max=10 (n=3) + Returns {n_words: {'mean': float, 'min': int, 'max': int, 'n': int}}. + """ + results = {} + with open(log_file, 'r') as f: + for line in f: + # Match: words= 2000 mean=9.7/10 min=9 max=10 (n=3) + m = re.search( + r'words=\s*(\d+)\s+mean=([\d.]+)/10\s+min=(\d+)\s+max=(\d+)\s+\(n=(\d+)\)', + line + ) + if m: + n_words = int(m.group(1)) + results[n_words] = { + 'mean': float(m.group(2)), + 'min': int(m.group(3)), + 'max': int(m.group(4)), + 'n': int(m.group(5)), + } + return results + + +def save_niah_perf_csv(results: Dict[int, Dict], output_file: str, + model_name: str = "", pipeline: str = "vllm"): + """Save NIAH results in madengine perf.csv format (one row per context length).""" + if not results: + print("No NIAH results to save to perf.csv.") + return + + meta = _get_run_metadata(pipeline) + + fieldnames = [ + 'model', 'n_gpus', 'nnodes', 'gpus_per_node', 'training_precision', + 'pipeline', 'args', 'tags', 'docker_file', 'base_docker', 'docker_sha', + 'docker_image', 'git_commit', 'machine_name', 'deployment_type', 'launcher', + 'gpu_architecture', 'performance', 'metric', 'relative_change', 'status', + 'build_duration', 'test_duration', 'dataname', 'data_provider_type', + 'data_size', 'data_download_duration', 'build_number', + 'additional_docker_run_options', + ] + + with open(output_file, 'w', newline='') as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + + for n_words in sorted(results.keys()): + data = results[n_words] + row = { + 'model': model_name, + 'performance': f"{data['mean']:.1f}", + 'metric': f"retrieval/10 (niah words={n_words} seeds={data['n']})", + 'status': 'SUCCESS', + } + row.update(meta) + writer.writerow(row) + + print(f"Saved {len(results)} NIAH rows to perf.csv: {output_file}") + + def save_perf_csv(results: Dict[Tuple[int, int, int], Dict], output_file: str, model_name: str = "", pipeline: str = "vllm"): """Save results in madengine perf.csv format.""" @@ -190,6 +253,8 @@ def main(): parser.add_argument('-o', '--output', type=str, help='Output CSV file name (default: _results.csv)') parser.add_argument('--perf-csv', type=str, help='Also generate madengine perf.csv at this path') parser.add_argument('--model-name', type=str, default='', help='Model name for perf.csv') + parser.add_argument('--niah', action='store_true', + help='Parse NIAH retrieval log instead of throughput sweep (requires --perf-csv)') args = parser.parse_args() @@ -201,6 +266,22 @@ def main(): print(f"Parsing log file: {log_file}") + # NIAH mode: parse retrieval scores, write perf.csv only + if args.niah: + if not args.perf_csv: + print("Error: --niah requires --perf-csv") + sys.exit(1) + results = parse_niah_log(log_file) + if not results: + print("No NIAH results found in log file.") + return + save_niah_perf_csv(results, args.perf_csv, args.model_name) + print(f"\nSummary (NIAH):") + print(f" Context lengths parsed: {len(results)}") + print(f" perf.csv: {args.perf_csv}") + return + + # Default: throughput sweep mode results = parse_benchmark_log(log_file) if not results: diff --git a/scripts/vllm_dissag/run_xPyD_models.slurm b/scripts/vllm_dissag/run_xPyD_models.slurm index c71fc7e8..8b29966d 100755 --- a/scripts/vllm_dissag/run_xPyD_models.slurm +++ b/scripts/vllm_dissag/run_xPyD_models.slurm @@ -87,6 +87,7 @@ VALID_MODELS=( \ "DeepSeek-R1" \ "Qwen3-32B" \ "Qwen3-30B-A3B" \ + "GLM-5.1-FP8" \ ) # Models allowed for CONNECTOR=moriio WIDE_EP=1 (MoRI-EP; legacy RUN_MORI=1) @@ -94,6 +95,7 @@ MORI_EP_VALID_MODELS=( \ "DeepSeek-V3" \ "DeepSeek-V3-5layer" \ "DeepSeek-R1" \ + "GLM-5.1-FP8" \ ) # Models allowed for CONNECTOR=rixl WIDE_EP=1 EP_BACKEND=deepep (legacy RUN_DEEPEP=1) @@ -179,7 +181,8 @@ WIDE_EP="${WIDE_EP:-0}" # the MoRI-EP / DeepEP recipe (block=16, MLA off, per-role cudagraph). Running them # in TP mode is unsupported — the TP argv would double the model's own # --compilation-config and drop the mandatory +quant_fp8 op. Reject early. -WIDE_EP_ONLY_MODELS=( "DeepSeek-V3" "DeepSeek-V3-5layer" "DeepSeek-R1" ) +# GLM-5.1-FP8 (MLA+DSA) is validated only under MoRI-EP wideEP; moriio+TP is untested. +WIDE_EP_ONLY_MODELS=( "DeepSeek-V3" "DeepSeek-V3-5layer" "DeepSeek-R1" "GLM-5.1-FP8" ) model_is_wide_ep_only() { local m="$1" for x in "${WIDE_EP_ONLY_MODELS[@]}"; do [[ "$m" == "$x" ]] && return 0; done @@ -427,11 +430,17 @@ BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-}" # Benchmark script selector: BENCHMARK_SCRIPT tag -> file run by the launcher. # sweep (default) -> benchmark_xPyD.sh (general concurrency sweep) # long_context -> benchmark_long_context.sh (per-shape warmup, c=1-first) +# agentic -> benchmark_agentic.sh (aiperf agentic trace replay) +# The launcher/connectors auto-enable prefix caching + server /metrics when +# BENCHMARK_SCRIPT=agentic. +# niah -> benchmark_niah.sh (needle-in-a-haystack; NIAH_WORDS) BENCHMARK_SCRIPT="${BENCHMARK_SCRIPT:-sweep}" case "$BENCHMARK_SCRIPT" in sweep) BENCHMARK_SCRIPT_FILE="benchmark_xPyD.sh" ;; long_context) BENCHMARK_SCRIPT_FILE="benchmark_long_context.sh" ;; - *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context)" >&2; exit 1 ;; + agentic) BENCHMARK_SCRIPT_FILE="benchmark_agentic.sh" ;; + niah) BENCHMARK_SCRIPT_FILE="benchmark_niah.sh" ;; + *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context, agentic, niah)" >&2; exit 1 ;; esac if [[ ! -f "$BENCHMARK_SCRIPT_FILE" ]]; then echo "Error: selected benchmark script '$BENCHMARK_SCRIPT_FILE' not found in $(pwd)." >&2 @@ -470,6 +479,7 @@ export IPADDRS="$(echo "${IPS[*]}" | sed 's/ /,/g')" export BENCHMARK_ITR=$BENCHMARK_ITR export BENCHMARK_CON="${BENCHMARK_CON}" export BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS}" +export BENCHMARK_SCRIPT="${BENCHMARK_SCRIPT}" export BENCHMARK_SCRIPT_FILE="${BENCHMARK_SCRIPT_FILE}" export DOCKER_CONT_NAME="container_${MODEL_NAME}_${SLURM_JOB_ID}" @@ -482,7 +492,7 @@ srun --nodelist="$SELECTED_NODELIST_SRUN" bash -c ' echo "Rank $SLURM_PROCID on $(hostname)"; docker ps -q | xargs --no-run-if-empty docker stop; docker rm -f $DOCKER_CONT_NAME 2>/dev/null || true; -fuser -k 5000/tcp 2>/dev/null || true; +fuser -k ${BARRIER_PORT:-5000}/tcp 2>/dev/null || true; fuser -k 2222/tcp 2>/dev/null || true; fuser -k 15000/tcp 2>/dev/null || true; sleep 2; @@ -525,13 +535,31 @@ done for _pattern in libmlx5.so* libionic*.so* libbnxt_re*.so* libefa.so* libhns.so*; do for _vlib in $_LIBDIR/${_pattern}; do - [ -e "$_vlib" ] && _RDMA_MOUNTS="$_RDMA_MOUNTS -v $_vlib:$_vlib:ro" + # Require a regular file AFTER symlink resolution: `-e` also accepts a directory + # matching one of these globs, and docker then bind-mounts a directory onto a + # file target -> "not a directory" -> container create exit 125. + [ -f "$_vlib" ] && _RDMA_MOUNTS="$_RDMA_MOUNTS -v $_vlib:$_vlib:ro" done done [ -d "$_LIBDIR/libibverbs" ] && _RDMA_MOUNTS="$_RDMA_MOUNTS -v $_LIBDIR/libibverbs:$_LIBDIR/libibverbs:ro" [ -d /etc/libibverbs.d ] && _RDMA_MOUNTS="$_RDMA_MOUNTS -v /etc/libibverbs.d:/etc/libibverbs.d:ro" echo "[host-rdma] mounts: $_RDMA_MOUNTS" +# AGENTIC_CONFIG may use $HOME (e.g. from models.json); expand at submit time. +[ -n "${AGENTIC_CONFIG:-}" ] && AGENTIC_CONFIG="${AGENTIC_CONFIG/#$HOME/$HOME}" + +# GPU_MEMORY_UTILIZATION precedence: +# 1. Submit-time `-e GPU_MEMORY_UTILIZATION=X` (explicit override) +# 2. models.yaml env: block (static per-model value) +# 3. Launcher fallback below (topology-aware: 0.70 at EP>=32, else 0.8) +# +# Models needing topology-dependent values (e.g., GLM) should OMIT the key from +# yaml and use the fallback. Models needing a static value can specify it in yaml. +# The 0.70 at EP>=32 accounts for larger MoRI dispatch/combine buffers at high EP +# width; measured on GLM but applies to any MoRI model as a safe default. +_EP_WIDTH=$(( yD * ${GPUS_PER_NODE:-8} )) +if [ "$_EP_WIDTH" -ge 32 ]; then _GPU_MEM_UTIL_LAUNCHER_DEFAULT=0.70; else _GPU_MEM_UTIL_LAUNCHER_DEFAULT=0.8; fi +echo "GPU_MEMORY_UTILIZATION launcher fallback: ${_GPU_MEM_UTIL_LAUNCHER_DEFAULT} (decode EP width ${_EP_WIDTH})" docker run --rm \ --device /dev/dri \ @@ -564,6 +592,27 @@ docker run --rm \ -e MASTER_PORT=$MASTER_PORT \ -e MODEL_PATH=$MODEL_PATH \ -e NIXL_COOKBOOK_PATH=$NIXL_COOKBOOK_PATH \ + -e NIXL_REPO_DIR=$NIXL_REPO_DIR \ + ${AGENTIC_CONC:+-e AGENTIC_CONC=$AGENTIC_CONC} \ + ${DURATION:+-e DURATION=$DURATION} \ + ${AGENTIC_PORT:+-e AGENTIC_PORT=$AGENTIC_PORT} \ + ${MODEL_PREFIX:+-e MODEL_PREFIX=$MODEL_PREFIX} \ + ${WEKA_LOADER_OVERRIDE:+-e WEKA_LOADER_OVERRIDE=$WEKA_LOADER_OVERRIDE} \ + ${AIPERF_UNSAFE_OVERRIDE:+-e AIPERF_UNSAFE_OVERRIDE=$AIPERF_UNSAFE_OVERRIDE} \ + ${AGENTIC_CACHE_WARMUP_DURATION:+-e AGENTIC_CACHE_WARMUP_DURATION=$AGENTIC_CACHE_WARMUP_DURATION} \ + ${MAX_MODEL_LEN:+-e MAX_MODEL_LEN=$MAX_MODEL_LEN} \ + ${AGENTIC_CONFIG:+-e AGENTIC_CONFIG=$AGENTIC_CONFIG} \ + ${AGENTIC_WORKLOAD:+-e AGENTIC_WORKLOAD=$AGENTIC_WORKLOAD} \ + ${SUITE_CORPUS_DIR:+-e SUITE_CORPUS_DIR=$SUITE_CORPUS_DIR} \ + ${SUITE_CORPUS_FORCE:+-e SUITE_CORPUS_FORCE=$SUITE_CORPUS_FORCE} \ + ${AGENTIC_STRICT_CONTEXT:+-e AGENTIC_STRICT_CONTEXT=$AGENTIC_STRICT_CONTEXT} \ + ${AGENTIC_HF_ISL_TAIL:+-e AGENTIC_HF_ISL_TAIL=$AGENTIC_HF_ISL_TAIL} \ + ${AIPERF_PIN:+-e AIPERF_PIN=$AIPERF_PIN} \ + ${AGENTIC_UTILS_PIN:+-e AGENTIC_UTILS_PIN=$AGENTIC_UTILS_PIN} \ + ${HF_HOME:+-e HF_HOME=$HF_HOME} \ + ${ENABLE_SERVER_METRICS:+-e ENABLE_SERVER_METRICS=$ENABLE_SERVER_METRICS} \ + ${ENABLE_PREFIX_CACHE:+-e ENABLE_PREFIX_CACHE=$ENABLE_PREFIX_CACHE} \ + ${DRY_RUN:+-e DRY_RUN=$DRY_RUN} \ -e xP=$xP \ -e yD=$yD \ -e USER_NAME=$USER_NAME \ @@ -578,6 +627,7 @@ docker run --rm \ ${CONNECTOR:+-e CONNECTOR=$CONNECTOR} \ ${WIDE_EP:+-e WIDE_EP=$WIDE_EP} \ ${EP_BACKEND:+-e EP_BACKEND=$EP_BACKEND} \ + ${BARRIER_PORT:+-e BARRIER_PORT=$BARRIER_PORT} \ ${RUN_MORI:+-e RUN_MORI=$RUN_MORI} \ ${RUN_DEEPEP:+-e RUN_DEEPEP=$RUN_DEEPEP} \ ${VLLM_ALL2ALL_BACKEND:+-e VLLM_ALL2ALL_BACKEND=$VLLM_ALL2ALL_BACKEND} \ @@ -589,7 +639,15 @@ docker run --rm \ ${ROUTER_BINARY:+-e ROUTER_BINARY=$ROUTER_BINARY} \ ${KV_CACHE_DTYPE:+-e KV_CACHE_DTYPE=$KV_CACHE_DTYPE} \ ${MORIIO_TOY_PROXY:+-e MORIIO_TOY_PROXY=$MORIIO_TOY_PROXY} \ + ${BENCHMARK_SCRIPT:+-e BENCHMARK_SCRIPT=$BENCHMARK_SCRIPT} \ ${BENCHMARK_SCRIPT_FILE:+-e BENCHMARK_SCRIPT_FILE=$BENCHMARK_SCRIPT_FILE} \ + ${NIAH_WORDS:+-e NIAH_WORDS=$NIAH_WORDS} \ + ${NIAH_SEEDS:+-e NIAH_SEEDS=$NIAH_SEEDS} \ + ${NIAH_WARMUP:+-e NIAH_WARMUP=$NIAH_WARMUP} \ + ${NIAH_MAXTOK:+-e NIAH_MAXTOK=$NIAH_MAXTOK} \ + ${NIAH_TIMEOUT:+-e NIAH_TIMEOUT=$NIAH_TIMEOUT} \ + ${SHAPE_WARMUP:+-e SHAPE_WARMUP=$SHAPE_WARMUP} \ + ${USE_INDUCTOR_GRAPH_PARTITION:+-e USE_INDUCTOR_GRAPH_PARTITION=$USE_INDUCTOR_GRAPH_PARTITION} \ ${PREFILL_CUDAGRAPH_MODE:+-e PREFILL_CUDAGRAPH_MODE=$PREFILL_CUDAGRAPH_MODE} \ ${DECODE_CUDAGRAPH_MODE:+-e DECODE_CUDAGRAPH_MODE=$DECODE_CUDAGRAPH_MODE} \ ${CUDAGRAPH_CAPTURE_SIZES:+-e CUDAGRAPH_CAPTURE_SIZES="$CUDAGRAPH_CAPTURE_SIZES"} \ @@ -613,7 +671,8 @@ docker run --rm \ ${MORI_NUM_QP_PER_PE:+-e MORI_NUM_QP_PER_PE=$MORI_NUM_QP_PER_PE} \ ${VLLM_MORIIO_QP_PER_TRANSFER:+-e VLLM_MORIIO_QP_PER_TRANSFER=$VLLM_MORIIO_QP_PER_TRANSFER} \ ${VLLM_MORIIO_NUM_WORKERS:+-e VLLM_MORIIO_NUM_WORKERS=$VLLM_MORIIO_NUM_WORKERS} \ - -e GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.8} \ + ${GPU_MEMORY_UTILIZATION:+-e GPU_MEMORY_UTILIZATION=$GPU_MEMORY_UTILIZATION} \ + -e _GPU_MEM_UTIL_LAUNCHER_DEFAULT=$_GPU_MEM_UTIL_LAUNCHER_DEFAULT \ -e GPUS_PER_NODE=${GPUS_PER_NODE:-8} \ ${GPU_MAX_HW_QUEUES:+-e GPU_MAX_HW_QUEUES=$GPU_MAX_HW_QUEUES} \ ${HIP_FORCE_DEV_KERNARG:+-e HIP_FORCE_DEV_KERNARG=$HIP_FORCE_DEV_KERNARG} \ diff --git a/scripts/vllm_dissag/tests/argv_assert.sh b/scripts/vllm_dissag/tests/argv_assert.sh index 9c85da89..2d3fd885 100755 --- a/scripts/vllm_dissag/tests/argv_assert.sh +++ b/scripts/vllm_dissag/tests/argv_assert.sh @@ -78,6 +78,27 @@ _parse() { # $1=connector ; reads its .env with same logic as the slurm _has "$(_parse moriio)" "-e PYTORCH_HIP_ALLOC_CONF=expandable_segments:False" "parse yields HIP_ALLOC -e arg" _has "$(PYTORCH_HIP_ALLOC_CONF=expandable_segments:True _parse moriio)" "-e PYTORCH_HIP_ALLOC_CONF=expandable_segments:True" "submit-time override wins" +# Per-shape warmup must stay opt-in: it is on the shared sweep path, so a default-on gate +# would change the measured TPOT of every already-validated recipe. GLM opts in via its +# models.yaml env:, and both GLM-only knobs need a docker -e line or they cannot be +# A/B-tested from the submit side (the recipe applies whenever the key is absent). +echo "" +echo "=== per-shape warmup is opt-in, not default-on ===" +B="$(cat "$DIR/benchmark_xPyD.sh")" +_has "$B" '${SHAPE_WARMUP:-0}' "benchmark_xPyD.sh: warmup gate defaults OFF" +_hasnot "$B" '${SHAPE_WARMUP:-1}' "benchmark_xPyD.sh: gate is not default-on" +_OPTIN="$(python3 - "$DIR/models.yaml" <<'PY' +import sys, yaml +y = yaml.safe_load(open(sys.argv[1])) or {} +optin = [m for m, c in y.items() + if isinstance(c, dict) and (c.get("env") or {}).get("SHAPE_WARMUP") == "1"] +print("[" + ",".join(sorted(optin)) + "]") +PY +)" +_has "$_OPTIN" "[GLM-5.1-FP8]" "models.yaml: GLM-5.1-FP8 is the ONLY warmup opt-in" +_has "$(cat "$SLURM")" '${SHAPE_WARMUP:+-e SHAPE_WARMUP=' "slurm forwards SHAPE_WARMUP override" +_has "$(cat "$SLURM")" '${USE_INDUCTOR_GRAPH_PARTITION:+-e USE_INDUCTOR_GRAPH_PARTITION=' "slurm forwards IGP override" + echo "" echo "======================================================" echo " argv_assert: ${pass} passed, ${fail} failed" diff --git a/scripts/vllm_dissag/tests/gate_check.sh b/scripts/vllm_dissag/tests/gate_check.sh index 66f61674..3700bf8d 100755 --- a/scripts/vllm_dissag/tests/gate_check.sh +++ b/scripts/vllm_dissag/tests/gate_check.sh @@ -3,9 +3,10 @@ # Verifies the (MODEL x CONNECTOR x WIDE_EP x EP_BACKEND) enablement gate in # run_xPyD_models.slurm accepts exactly the supported combos and rejects the rest. # -# Strategy: source the slurm's model lists + axis/validation logic in a harness that -# stops right after the gate (never reaches docker/srun), then assert exit status. -# This runs the REAL gate code path, so a future edit to the lists/gate is caught. +# Strategy: re-run a hand-kept MIRROR of the slurm's model lists + axis/validation logic +# in a harness that stops right after the gate (never reaches docker/srun), then assert +# the verdict. NOTE the mirror is a copy, not the real code path: editing the lists in +# run_xPyD_models.slurm without editing the mirror below will NOT be caught here. # # Usage: bash tests/gate_check.sh (exit 0 = all pass) set -u @@ -28,10 +29,11 @@ _run_gate() { set -u # --- mirror of run_xPyD_models.slurm gate (keep in sync) --- VALID_MODELS=( "Llama-3.1-405B-Instruct-FP8-KV" "amd-Llama-3.3-70B-Instruct-FP8-KV" \ - "DeepSeek-V3" "DeepSeek-V3-5layer" "gpt-oss-120b" "DeepSeek-R1" "Qwen3-32B" "Qwen3-30B-A3B" ) - MORI_EP_VALID_MODELS=( "DeepSeek-V3" "DeepSeek-V3-5layer" "DeepSeek-R1" ) + "DeepSeek-V3" "DeepSeek-V3-5layer" "gpt-oss-120b" "DeepSeek-R1" "Qwen3-32B" "Qwen3-30B-A3B" \ + "GLM-5.1-FP8" ) + MORI_EP_VALID_MODELS=( "DeepSeek-V3" "DeepSeek-V3-5layer" "DeepSeek-R1" "GLM-5.1-FP8" ) DEEPEP_VALID_MODELS=( "DeepSeek-V3" "DeepSeek-V3-5layer" "DeepSeek-R1" ) - WIDE_EP_ONLY_MODELS=( "DeepSeek-V3" "DeepSeek-V3-5layer" "DeepSeek-R1" ) + WIDE_EP_ONLY_MODELS=( "DeepSeek-V3" "DeepSeek-V3-5layer" "DeepSeek-R1" "GLM-5.1-FP8" ) MODEL_NAME="${MODEL_NAME:-None}" _in(){ local n="$1"; shift; for x in "$@"; do [[ "$n" == "$x" ]] && return 0; done; return 1; } _in "$MODEL_NAME" "${VALID_MODELS[@]}" || { echo REJECT; exit 0; } @@ -87,6 +89,11 @@ _case ALLOW "DSV3 rixl wideEP(deepep)" DeepSeek-V3 rixl _case ALLOW "R1 moriio wideEP" DeepSeek-R1 moriio 1 _case REJECT "DSV3 moriio TP" DeepSeek-V3 moriio 0 _case REJECT "DSV3 rixl TP" DeepSeek-V3 rixl 0 +# GLM-5.1-FP8 — moriio wideEP only (no TP, no DeepEP) +_case ALLOW "GLM moriio wideEP(mori)" GLM-5.1-FP8 moriio 1 +_case REJECT "GLM moriio TP" GLM-5.1-FP8 moriio 0 +_case REJECT "GLM rixl TP" GLM-5.1-FP8 rixl 0 +_case REJECT "GLM rixl wideEP(deepep)" GLM-5.1-FP8 rixl 1 # cross-pairs _case REJECT "DSV3 moriio+deepep xpair" DeepSeek-V3 moriio 1 deepep _case REJECT "DSV3 rixl+mori xpair" DeepSeek-V3 rixl 1 mori diff --git a/scripts/vllm_dissag/tests/run_interactive.sh b/scripts/vllm_dissag/tests/run_interactive.sh index d6e1e67a..802d5f2c 100755 --- a/scripts/vllm_dissag/tests/run_interactive.sh +++ b/scripts/vllm_dissag/tests/run_interactive.sh @@ -19,7 +19,7 @@ RUN_FILE_FULL="$NIXL_COOKBOOK_PATH/vllm_disagg.sh" # cleanup any stale container/ports on this node docker rm -f "$DOCKER_CONT_NAME" 2>/dev/null || true -fuser -k 5000/tcp 2>/dev/null || true +fuser -k "${BARRIER_PORT:-5000}/tcp" 2>/dev/null || true fuser -k 2222/tcp 2>/dev/null || true fuser -k 15000/tcp 2>/dev/null || true fuser -k 30000/tcp 2>/dev/null || true diff --git a/scripts/vllm_dissag/vllm_disagg.sh b/scripts/vllm_dissag/vllm_disagg.sh index 06fbf84f..7b794265 100755 --- a/scripts/vllm_dissag/vllm_disagg.sh +++ b/scripts/vllm_dissag/vllm_disagg.sh @@ -176,6 +176,10 @@ PY )" [[ -n "$_yaml_env" ]] && eval "$_yaml_env" + # Fallback: launcher-calculated default if not set by submit-time or yaml + : ${GPU_MEMORY_UTILIZATION:=$_GPU_MEM_UTIL_LAUNCHER_DEFAULT} + export GPU_MEMORY_UTILIZATION + # 2) Resolve per-role flag strings for the active PARALLEL_MODE. eval "$(python3 - <<'PY' import os, shlex, yaml @@ -247,8 +251,38 @@ if [ "$NODE_RANK" -eq 0 ]; then # connector_start_proxy sets BENCHMARK_PORT (router->ROUTER_PORT, toy->PROXY_PORT). # Fall back to PROXY_PORT only if the connector didn't set it. export BENCHMARK_PORT="${BENCHMARK_PORT:-${PROXY_PORT}}" + # Agentic replay driver reads the endpoint via AGENTIC_PORT; for the default + # sweep it is exactly BENCHMARK_PORT (the router/proxy port), so the non-agentic + # path is unaffected. + export AGENTIC_PORT="${BENCHMARK_PORT}" + + # Agentic path only: the vLLM PD router serves /v1/chat/completions but 503s on + # /v1/models under MoRIIO service discovery (empty HTTP worker registry). The + # shared harness (scripts/common/agentic_lib.sh) gates readiness + served-model + # resolution on GET /v1/models, so start a tiny side-port shim that answers + # /v1/models (gated on the router's /health) and stream-proxies everything else + # to the router, then point AGENTIC_PORT at the shim. Byte-for-byte no-op for the + # default sweep (guarded by BENCHMARK_SCRIPT_FILE). + _agentic_shim_pid="" + if [[ "${BENCHMARK_SCRIPT:-}" == "agentic" ]]; then + _shim_port="${AGENTIC_SHIM_PORT:-$((BENCHMARK_PORT + 1))}" + _shim_model="${MODEL:-${MODEL_PATH}}" + _shim_prefill="${AGENTIC_SERVER_METRICS%% *}" + AGENTIC_SHIM_PORT="${_shim_port}" \ + AGENTIC_SHIM_UPSTREAM="127.0.0.1:${BENCHMARK_PORT}" \ + AGENTIC_SHIM_MODEL="${_shim_model}" AGENTIC_SHIM_PREFILL="${_shim_prefill}" \ + python3 "$NIXL_COOKBOOK_PATH/agentic_models_shim.py" \ + > >(tee /run_logs/${SLURM_JOB_ID}/agentic_models_shim_NODE${NODE_RANK}.log >/dev/null) 2>&1 & + _agentic_shim_pid=$! + export AGENTIC_PORT="${_shim_port}" + echo "[agentic-shim] models shim on :${_shim_port} -> router :${BENCHMARK_PORT} (model=${_shim_model})" + sleep 3 + fi + bash "$NIXL_COOKBOOK_PATH/${BENCHMARK_SCRIPT_FILE:-benchmark_xPyD.sh}" + [[ -n "$_agentic_shim_pid" ]] && { kill "$_agentic_shim_pid" 2>/dev/null || true; } + echo "Killing the proxy server.." pkill -P $proxy_pid 2>/dev/null; kill $proxy_pid 2>/dev/null || true echo "Killing the prefill master server.." diff --git a/tools/fetch_primus.sh b/tools/fetch_primus.sh new file mode 100755 index 00000000..11062092 --- /dev/null +++ b/tools/fetch_primus.sh @@ -0,0 +1,76 @@ +#!/usr/bin/env bash +############################################################################### +# +# MIT License +# +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# + +# Check Primus out into scripts/Primus. The JAX backends need it before anything else +# happens: scripts/jax-{maxtext,maxdiffusion}/get_models_json.py glob its example configs +# to enumerate models, and both primus_* dockerfiles COPY the tree into the image. Without +# it, discovery reports zero models instead of a missing prerequisite. +# +# Run on the host, from anywhere: +# tools/fetch_primus.sh +# +# Idempotent, so it is safe in CI or a Makefile. Override PRIMUS_URL, PRIMUS_REF, or +# PRIMUS_DIR for a fork, another branch or commit, or a different location. +set -uo pipefail + +PRIMUS_URL="${PRIMUS_URL:-https://github.com/AMD-AGI/Primus}" +PRIMUS_REF="${PRIMUS_REF:-main}" + +MAD_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +PRIMUS_DIR="${PRIMUS_DIR:-$MAD_ROOT/scripts/Primus}" + +log() { echo "[fetch-primus] $*"; } +die() { echo "[fetch-primus] ERROR: $*" >&2; exit 1; } + +command -v git >/dev/null || die "git not found on PATH." + +if git -C "$PRIMUS_DIR" rev-parse --git-dir >/dev/null 2>&1; then + log "already checked out at $PRIMUS_DIR" +elif [[ -d "$PRIMUS_DIR" ]] && [[ -z "$(ls -A "$PRIMUS_DIR" 2>/dev/null)" ]]; then + # Empty dir left by an uninitialized git submodule; remove so clone succeeds. + rmdir "$PRIMUS_DIR" + log "cloning $PRIMUS_URL ($PRIMUS_REF) into $PRIMUS_DIR" + git clone --branch "$PRIMUS_REF" "$PRIMUS_URL" "$PRIMUS_DIR" \ + || die "clone failed. For a private repo, check that your git credentials can read $PRIMUS_URL." +elif [[ -e "$PRIMUS_DIR" ]]; then + die "$PRIMUS_DIR exists but is not a git checkout. Move it aside and re-run." +else + log "cloning $PRIMUS_URL ($PRIMUS_REF) into $PRIMUS_DIR" + # Deliberately not --recursive, and no submodules are initialized afterwards. Both + # primus_* images take their framework from the base image (/workspace/maxtext, + # /workspace/maxdiffusion) and pin MAXTEXT_PATH / MAXDIFFUSION_PATH to it. A submodule + # checkout would only add an unpatched second copy of the same commit to the build context. + git clone --branch "$PRIMUS_REF" "$PRIMUS_URL" "$PRIMUS_DIR" \ + || die "clone failed. For a private repo, check that your git credentials can read $PRIMUS_URL." +fi + +maxtext_configs=$(find "$PRIMUS_DIR/examples/maxtext/configs" -name '*.yaml' 2>/dev/null | wc -l) +maxdiff_configs=$(find "$PRIMUS_DIR/examples/maxdiffusion/configs" -name '*.yaml' 2>/dev/null | wc -l) +log "$(git -C "$PRIMUS_DIR" rev-parse --short HEAD) on $(git -C "$PRIMUS_DIR" rev-parse --abbrev-ref HEAD)" +# Config files, not discovered models: discovery filters the multi-node-only ones out. +log "config files found: $maxtext_configs maxtext, $maxdiff_configs maxdiffusion" +[[ "$maxtext_configs" -gt 0 ]] || die "no maxtext configs found; is $PRIMUS_REF the right ref?" diff --git a/tools/run_models.py b/tools/run_models.py index 6c885995..8ca1f31f 100644 --- a/tools/run_models.py +++ b/tools/run_models.py @@ -267,9 +267,9 @@ def run_model( update_dict(docker_build_args, mad_secrets) build_args = " ".join([f"--build-arg {key}='{value}'" for key, value in docker_build_args.items()]) - docker_context = "./docker" - model_docker_image = f"ci-{model_name}" - model_docker_container = f"container_ci-{model_name}" + docker_context = model.get("dockercontext", "./docker") + model_docker_image = f"ci-{model_name}".lower() + model_docker_container = f"container_ci-{model_name}".lower() # Store the data for the run details run_details.model = model_name @@ -435,7 +435,7 @@ def run_model( update_perf_csv(exception_result="perf_entry.json", perf_csv=output) # Clean up the instance of docker del docker - sys.exit(1) + return False test_duration = time.time() - test_start_time logger.info(f"Test duration: {test_duration} seconds") @@ -530,7 +530,8 @@ def main() -> bool: console = Console(live_output=args.live_output) # Load models.json file to list of dictionary. - models = load_models() + models_data = load_models() + models = models_data["models"] if isinstance(models_data, dict) and "models" in models_data else models_data user_tags = None if args.tags: