From a33b62272389fce0ef394ae95870234347682aaf Mon Sep 17 00:00:00 2001 From: Mir Mustafa Ali Date: Fri, 28 Aug 2026 16:51:47 -0500 Subject: [PATCH 1/5] vllm_dissag: GLM-5.1-FP8 (MLA+DSA) MoRI-EP WideEP disaggregated enablement (#206) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * vllm_dissag: GLM-5.1-FP8 (MLA+DSA) MoRI-EP WideEP disaggregated enablement Adds GLM-5.1-FP8 (GlmMoeDsaForCausalLM = MLA + DeepSeek Sparse Attention) to the MoRI-EP WideEP disaggregated serving path, stacked on the #171 unified launcher. Fully isolated from DeepSeek-V3/R1: GLM gets its own image + a MODEL_NAME-gated runtime path, so existing models are byte-identical to develop. Defects fixed (validated 1P/1D EP8 + 2P/2D EP16, NIAH 2k-35k = 10/10, no crash): - Long-context accuracy collapse: vLLM #47766 cache-key fix keeps the persistent sparse-MLA kernel ON (keys metadata on per-request context+query len). - 8k disagg prefill crash: DSA adds a 2nd (indexer) KV cache per layer that the single-geometry MoRIIO connector never transferred; paired + shipped prefill-> decode. Plus DSA invalid-token kernel fix (#45324) and shik-latest DP-notify. Changes: - docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile: NEW per-model image (raviguptaamd/vllm glm5.1-dsa-wideEP_on_shik_latest + aiter e03fa6040 + mori 42e895472b08 + router). The base vllm_disagg_inference Dockerfile (DSV3/R1) is left untouched. Future models add their own Dockerfile the same way. - models.json: card pyt_vllm_disagg_mori_glm-5.1-fp8 (GLM_SKIP_PATCHERS=1: image carries the DSA fixes in-source). - models.yaml: GLM-5.1-FP8 recipe (block=1, AITER MLA on, eager, mori backends). DeepSeek-V3 dp: caps (--max-num-seqs 64 --max-model-len 32768) to bound the newer base's decode logits workspace (isolated to the DSV3 entry). - connectors/moriio.sh: MODEL_NAME-gated GLM DSA runtime patchers (pure no-op for other models); GLM_SKIP_PATCHERS switch for baked-fix images. - 9 idempotent, anchor-based, self-skipping GLM DSA patcher scripts. KNOWN OPEN DEFECT (future work): 4P/4D EP32 emits corrupted tokens at all context lengths (suspect moriep all-to-all combine at scale); use 1P/1D and 2P/2D. Co-Authored-By: Claude * vllm_dissag: make NIAH harness thinking-model-aware (GLM-5.1) benchmark_niah.py mis-scored thinking models: it never disabled thinking and read only content + reasoning_content. GLM-5.1 emits chain-of-thought into the `reasoning` field and leaves `content` empty until the final answer, so with a small max_tokens the answer never lands in content -> a false 0/10 even when generation is correct. - Add chat_template_kwargs.enable_thinking=false so the answer goes to content. - Also score the `reasoning` field as a fallback. Verified against GLM-5.1-FP8: correct 9-10/10 retrieval across 2k-35k on all tested topologies (EP8/EP16/EP32) after the fix. Co-Authored-By: Claude * vllm_dissag: NIAH multi-seed support (NIAH_SEEDS) for variance-aware accuracy Needle layout is seeded, so a single run is deterministic (bit-exact on the same stack) but can't tell a real accuracy dip from single-needle variance. Add NIAH_SEEDS (default 0,1,2) to run each context length across multiple needle layouts; the summary now reports mean/min/max across seeds. Backward compatible: NIAH_SEEDS=0 reproduces the prior single-seed behavior. Co-Authored-By: Claude * vllm_dissag: NIAH gate robust to cold-start JIT (warmup + readiness probe) On a freshly-booted node the first request of each context shape pays the full JIT/kernel-autotune compile (minutes). The NIAH harness scored the FIRST request, so cold compile landed on a scored/gated request -> false 0/10 or timeout, failing the accuracy gate and skipping the perf sweep. Root-caused by reproducing on a cold boot (0 results) vs a warm server (all pass) on the same image. Fixes: - benchmark_niah.py: add a warmup pass (NIAH_WARMUP=1 default) — one throwaway request per context length before scoring, with a generous timeout, failures tolerated. Scored requests are then always warm. - benchmark_niah.py: distinguish TIMEOUT/ERROR from a wrong answer. Timeouts return a sentinel (excluded from mean, never counted as 0/10); summary flags NO-RESULT with guidance instead of silently reporting 0. - benchmark_niah.sh: replace the blind `sleep 10` with a /v1/models readiness poll (up to 5 min), and forward NIAH_WARMUP. Verified: patched harness on the warm server passes 10/10; cold-boot repro no longer produces false 0/10 because compile happens in the warmup pass. Co-Authored-By: Claude * [GLM-5.1] v0.27 level-set: 3.4x decode speedup, 200K context, RDMA + launcher fixes Validated on MI300X, 8 nodes, image rocmshared/pytorch-private:glm5.1-vllm027-b8 (base ci_base-dedbf6be8b + vLLM raviguptaamd/vllm@glm5.1-dsa-wideEP_on_vllm-v0.27 + aiter e03fa6040 + MoRI 42e895472b08 + router #181). PERF FIX (models.yaml decode.dp) -- the headline change: --max-num-batched-tokens 2048 on the DECODE role only. max_num_batched_tokens is a chunked-prefill SCHEDULER knob, but it also sizes the MoRI EP dispatch buffer (fused_moe/layer.py -> all2all_utils.py -> all2all.py max_num_inp_token_per_rank). At the 8192 default a decode instance ran an 8192-token-wide all2all every step, per layer, x78 layers, while decoding a handful of tokens: a fixed ~302ms/step floor, ~320x this model's HBM-bandwidth bound. Prefill keeps 8192 (it genuinely dispatches wide batches). 1024/64 con=8, warm: TPOT TTFT out tok/s 1P/1D 302 -> 88.0 ms 2431 -> 906 ms 24.9 -> 78.8 2P/2D 302 -> 94.1 ms 1633 ms 66.7 Published reference: 1P/1D ~89ms, 2P/2D ~91ms -> matched within 3%. Accuracy unaffected: NIAH 2k-200k clean on both topologies, no length collapse, memfault=0, latencies equal-or-better at every length. 200K validated (5.7x beyond the previously published 35K ceiling). Dockerfile: base -> ci_base-dedbf6be8b (matches the fork's upstream base), VLLM_REF -> the v0.27 branch, and WITH_MORI_BUILD/WITH_AITER_BUILD now default to 1 so a plain `docker build` reproduces the validated stack. Previously they defaulted to 0, which silently used the base's bundled aiter 0.1.19 -- that GPU-faults on the GLM DSA decode kernel. The pinned aiter e03fa6040 / MoRI 42e895472b08 must not be bumped without re-running long-context NIAH. connectors/moriio.sh: per-role env split (PREFILL_*/DECODE_* -> VLLM_MORI_*), mirroring the existing PREFILL/DECODE_MORI_BACKEND pattern -- models.yaml env: applies to BOTH roles, but prefill and decode need opposite values here. Also injects use_inductor_graph_partition (pairs with the vLLM splitting_ops fix). connectors/moriio.env: RDMA fabric -- MORI_IB_GID_INDEX=3 (RoCEv2 IPv4), MORI_RDMA_DEVICES/NCCL_IB_HCA restricted to the 8 GPU-local NICs (leaving the mgmt NICs in makes QPs form over a non-routable fabric -> ibverbs.cpp:189 timeouts), NCCL/GLOO control sockets on eth0. run_xPyD_models.slurm: libionic bind-mount requires a regular file after symlink resolve (a dangling symlink gave "OCI runtime create ... not a directory", container exit 125); prefer FABRIC_SUBNET over `hostname -I` first IP (nodes list a 10.224 overlay first, which made the socket_barrier advertise an unreachable NIC -> "Waiting for nodes" hang); GLM_KERNEL_PATCH/GLM_BACKEND_PATCH bind-mount hooks to test .py fixes without a rebuild; forward the new per-role env keys. vllm_disagg.sh: same FABRIC_SUBNET IP-selection fix for host_ip. benchmark_xPyD.sh: per-shape warmup at the REAL isl/osl before each shape's cells. The global warmup is isl=osl=32/con=1, which never exercises a 1024/8192/28672 prefill path or the decode cudagraph batch sizes, so the first measured cell absorbed residual JIT (observed 302ms vs ~88ms steady-state). Warmup output goes to a separate _SHAPEWARMUP.log so it cannot pollute the CSV. models.yaml (GLM-5.1-FP8): decode.dp perf fix above; recipe = prefill eager + mori_high_throughput, decode PIECEWISE cudagraph + mori_low_latency; VLLM_USE_LAYERNAME=0; VLLM_SPARSE_INDEXER_MAX_LOGITS_MB=64; NCCL heartbeat/timeout knobs for long-context collectives. Full operational playbook (including the dead ends) in skills_vllm_disagg.md. Co-Authored-By: Claude * [GLM-5.1] Default WITH_NIXL=0 (lean MoRI-EP-only image) GLM-5.1 is served over MoRI-EP (all2all) + MoRI-IO (KV transfer); the rixl connector's UCX/RIXL/rocSHMEM/DeepEP stack is never selected by this recipe, so building it only lengthens the build and ships unused transports. 0 is also how the validated image was actually built: glm5.1-vllm027-b8 used --build-arg WITH_NIXL=0, so the previous default of 1 meant a plain `docker build` did NOT reproduce the tested artifact. With this change, and with WITH_MORI_BUILD/WITH_AITER_BUILD already defaulting to 1, a no-flag build now matches the validated stack exactly. Set --build-arg WITH_NIXL=1 if you need the rixl connector from this same Dockerfile. Co-Authored-By: Claude * [GLM-5.1] Add long-context NIAH harness + vllm-disagg operational playbook niah_200k.py: needle-in-a-haystack sweep that validated GLM-5.1-FP8 to 200,049 tokens on both 1P/1D (EP8) and 2P/2D (EP16). Reports found/10, latency, and the server-reported prompt_tokens per length, and writes JSON. Model id is overridable via NIAH_MODEL so it is not GLM-specific. The existing benchmark_niah.* stop well short of this range; this covers the 64k-200k band. skills_vllm_disagg.md: operational playbook for vLLM PD-disaggregated WideEP on MI300X (MoRI-EP + MoRI-IO), written from this enablement. Documents, with measurements: - benchmarking method: ALWAYS discard the first post-boot run (cold Triton JIT made TTFT read 13.4s vs 906ms warm; with prefill eager the JIT cost lands in TTFT, not TPOT), and sanity-check against the HBM-bandwidth bound before blaming a kernel - the max_num_batched_tokens trap: a chunked-prefill SCHEDULER knob also sizes the MoRI EP dispatch buffer, so decode ran an 8192-token-wide all2all every step (302ms -> 88ms TPOT once sized for the real batch) - the DSA sentinel landmine: the invalid sparse-index sentinel must be 0, not -1, because aiter's mla_decode_fwd dereferences it (only bites at disagg long context) - three documented DEAD ENDS so they are not retried, including why mori's max_total_recv_tokens cannot decouple recv from send capacity (the clamp is a min()) - cache/boot behaviour (three caches with different rules, the aiter baton lock, measured boot times), readiness signals for multi-node topologies, per-role env plumbing, and RDMA fabric verification (a node can be SLURM-'alloc' with a dead fabric - verify with ping/ib_write_bw before blaming code) Co-Authored-By: Claude * [GLM-5.1] Remove dead-end MoRI EP env plumbing and its misleading comments Review cleanup. The VLLM_MORI_MAX_TOKENS_PER_RANK / VLLM_MORI_MAX_TOTAL_RECV_TOKENS per-role plumbing was written while chasing the decode-TPOT bug and is NOT what fixed it (the fix is `--max-num-batched-tokens 2048` on decode.dp). Worse, the comments asserted that max_total_recv_tokens keeps recv capacity large enough for vLLM's profiling dummy run -- which is false and was disproved by measurement: mori's MaxNumTokensToRecvPerRank() = min(ceil(maxTotalRecvTokens / worldSize), maxNumInpTokenPerRank) is a min(), so maxTotalRecvTokens can only LOWER recv capacity, never raise it above the send width. Anyone following those comments and setting the knobs would hit "Total recv token overflow" at boot (observed at 512, 2048, and with recv=65536). Removed: the per-role export block in moriio.sh, the stale models.yaml comment block, and the six dead keys from _RECIPE_ENV_KEYS. Replaced with a short NOTE in moriio.sh pointing at the real fix and at skills_vllm_disagg.md for the measurements and dead ends. No functional change to the validated configuration: the knobs defaulted to 0/unset, so the b8 runs never exercised them. Co-Authored-By: Claude * [GLM-5.1] Decode CUDA-graph FULL_AND_PIECEWISE: 2.1x (EP8) / 1.7x (EP16) faster decode Second perf lever, on top of the decode --max-num-batched-tokens fix. Switching the decode role from PIECEWISE to FULL_AND_PIECEWISE captures a FULL graph for the uniform-decode shapes and falls back to PIECEWISE for the rest. Measured on MI300X, ISL/OSL 1024/64, warm (first post-boot run discarded), image rocmshared/pytorch-private:glm5.1-pr176-verify: PIECEWISE FULL_AND_PIECEWISE gain 1P/1D EP8 TPOT c8 88.0 ms 41.8 ms 2.11x 1P/1D EP8 TPOT c16 91.5 ms 45.8 ms 2.00x 2P/2D EP16 TPOT c8 94.1 ms 55.4 ms 1.70x 2P/2D EP16 TPOT c16 96.2 ms 60.3 ms 1.60x 1P/1D outTP c8 78.8 tok/s 143.7 tok/s 1.82x 2P/2D outTP c8 66.7 tok/s 113.0 tok/s 1.69x 2P/2D TTFT c16 2221 ms 1400 ms 1.59x Both topologies now sit well below the published ~89 ms (EP8) / ~91 ms (EP16) reference. Accuracy is unaffected. NIAH 2k-35k: 1P/1D 51/60, 2P/2D 55/60 - both inside the PIECEWISE band (52-53 and 53 respectively; 2P/2D is in fact the best NIAH run of the campaign). No length collapse, no '!!!', memfault=0, latencies equal-or-better at every length. Caveats, stated rather than buried: - The gain SHRINKS as EP widens (2.11x at EP8 -> 1.70x at EP16). FULL captures the whole decode step including the cross-node all2all, which does not compress. This is validated at EP8 and EP16 ONLY; do not assume it holds at EP32 (which has a separate, unrelated correctness regression). - Capture costs ~92-94 s and ~3.0-3.5 GiB vs ~5 s / 7.2 GiB for PIECEWISE (fewer but larger graphs), so boot is longer. Irrelevant for a long-lived server, notable for CI. - Set DECODE_CUDAGRAPH_MODE=PIECEWISE to fall back. Two code paths branch on CUDAGraphMode.FULL (sparse_attn_indexer.py:411 and the MoRIIO READ-mode barrier); both are guards that skip host-side work under FULL, so enabling it is safe here - the MoRIIO one is already a no-op in WRITE mode. skills_vllm_disagg.md gains a section 2b with the table, the EP-width caveat, the capture cost, and the observation that FULL is far less cold-JIT sensitive (its warmup run already reads steady-state TPOT, where PIECEWISE's first run showed 13.7 s TTFT). Co-Authored-By: Claude * [GLM-5.1] Merge hygiene: drop lab leftovers, safe patcher defaults. Rebase onto develop after #197 so the GLM card lives in scripts/vllm_dissag/models.json. Remove the NFS silly-rename and keepalive/NIAH-200k/skills lab files, revert DeepSeek-V3 context caps and cluster-wide moriio.env NIC pins, turn persist-gate and DSA sentinel patchers opt-in (they crash the v0.27 image), gate DSA on GLM-5.*, and wire BENCHMARK_SCRIPT=niah for the 1P/1D smoke. * [GLM-5.1] Point the recipe at the v0.27 image it was validated on. The recipe named a pre-v0.27 image as "validated" and described patchers that no longer apply. Serving GLM-5.1-FP8 on that older image returns empty completions (NIAH 2k = 0/10), so name the image the Dockerfile actually builds and state that the patchers are an old-image fallback. * [GLM-5.1] Minimal MAD: drop runtime DSA patchers, make the image pin the contract PR 206 shipped GLM-5.1-FP8 with nine apply_glm_*.py runtime monkeypatch scripts (~1100 lines) invoked from the moriio connector. None of them run on the production path: the validated model card sets GLM_SKIP_PATCHERS=1 because the pinned vLLM (VLLM_REF=glm5.1-dsa-wideEP_on_vllm-v0.27) already carries every DSA fix in-source. They were dead weight that MAD had to keep in sync with a fork. This branch makes MAD carry catalog content only, and makes the image solely responsible for vLLM source fixes: - Delete all nine scripts/vllm_dissag/apply_glm_*.py. - connector_runtime_patch() in connectors/moriio.sh is now a documented no-op for every model; _glm_dsa_runtime_patch() and the GLM MODEL_NAME branch are gone. Non-GLM behavior is unchanged (it already returned 0 before reaching any GLM code), so DeepSeek/Llama/gpt-oss are untouched. - Remove the plumbing that only existed to feed the patchers: the GLM_SKIP_PATCHERS / GLM_PERSIST_GATE / GLM_DSA_SENTINEL_FIX docker -e forwards in run_xPyD_models.slurm, and GLM_SKIP_PATCHERS from the model card in models.json (with no patchers, "skip the patchers" is meaningless). - Rewrite the models.yaml and Dockerfile comments that promised runtime patchers. The docs now state the actual contract: the image must carry the DSA fixes in-source and VLLM_REF is the requirement, not a preference. Trade-off: this removes the fallback path for older images. On an image built from a vLLM ref that predates the DSA fixes, GLM-5.1 now boots and then produces garbage or stalls the disagg KV transfer, with nothing to recover it -- the fix is to rebuild the image, not to patch at runtime. In exchange MAD stops carrying a duplicate of upstream fork code that no supported configuration executes. The Dockerfile, model card, models.yaml recipe, NIAH harness and slurm allowlist are otherwise unchanged. Verified without GPUs: tests/argv_assert.sh 25 passed / 0 failed; models.json parses and the GLM card is intact; DRY_RUN argv for DeepSeek-V3, DeepSeek-V3-5layer, DeepSeek-R1, amd-Llama-3.3-70B-Instruct-FP8-KV and gpt-oss-120b is byte-identical to glm5.1-wideep-mad (so the only delta vs origin/develop is the pre-existing one from that branch); GLM-5.1-FP8 still emits --enable-expert-parallel, --block-size 1, and on decode --max-num-batched-tokens 2048 + FULL_AND_PIECEWISE; bash -n clean. * [GLM-5.1] Scope USE_INDUCTOR_GRAPH_PARTITION to GLM, restore other models' argv The knob was introduced by the GLM v0.27 level-set (it does not exist on origin/develop) with a default of 1, so every model served over the moriio connector silently got "use_inductor_graph_partition":true appended to its --compilation-config: DeepSeek-V3, DeepSeek-V3-5layer, DeepSeek-R1, amd-Llama-3.3-70B-Instruct-FP8-KV and gpt-oss-120b. The flag is a fix for the MLA custom-op boxing failure that only GLM-5.1 hits, so it belongs with the model, not with the connector. Flip the connector default to 0 and set USE_INDUCTOR_GRAPH_PARTITION=1 in the GLM-5.1-FP8 models.yaml env: block, the same mechanism the recipe already uses for VLLM_USE_LAYERNAME / KV_BLOCK_SIZE / DECODE_CUDAGRAPH_MODE. Add the key to _RECIPE_ENV_KEYS so the image-baked < models.yaml < submit-time -e precedence holds for it too (without it a submit-time override would lose to the recipe). Verified offline (DRY_RUN argv, no GPUs): the five models above are now byte-identical to origin/develop across moriio+rixl x wideEP+TP x prefill+decode (40 cells, zero diff, no use_inductor_graph_partition anywhere), GLM-5.1-FP8 prefill+decode argv is unchanged vs glm5.1-wideep-mad, and argv_assert is 25 passed / 0 failed. * [GLM-5.1] Fix false MORI_REF claim and stale image text in the GLM docs Comment-only. Three claims in the GLM catalog text did not match reality: 1. connectors/moriio.sh said MORI_REF=42e895472b08 had the large-transfer notify/mapping fixes #424/#436/#432 "baked in". It does not: the pin is ROCm/mori#366 (2026-06-05), while #424 merged 2026-06-26 and #432/#436 on 2026-06-29 — `gh api repos/ROCm/mori/compare/42e895472b08...` reports status "ahead" for all three. The pin is left alone (the image was validated with exactly this MoRI); only the comment is corrected. 2. models.yaml said the unsupported rocmshared lab image "is only usable on a branch that still ships patchers" — a dangling reference from a branch with no patchers. It now names the one supported image outright. 3. models.yaml's long-context caveat described the old mori v1.2.1 lab image ("garbage beyond ~16-18k tokens") as if it applied to the pinned image, which contradicts both the glmv5.1 Dockerfile STATUS and the NIAH 2k-35k numbers in this same recipe. Re-scoped to history and pointed at the in-source fix (vLLM #47766). #40018 is also labelled as an issue, which it is. Two audited citations needed no change: vLLM #47042 is only cited in benchmark_niah.{py,sh} / README.MD / ARCHITECTURE.md, all of which are origin/develop text that already calls it an issue (never a PR); and aiter#3921 (a gfx950 A8W8 change) was only cited by apply_glm_dsa_persistent_kernel_gate_fix.py, which this branch already deleted. * [GLM-5.1] Fix cross-model fallout: protect-list, fabric-IP default, patch mounts Three defects that reached beyond GLM-5.1-FP8: 1. _RECIPE_ENV_KEYS listed only 18 of the 39 keys a models.yaml `env:` block can set. Once MODELS_YAML_PROTECT exists the driver protects ONLY listed keys, so a submit-time `-e KEY=...` silently lost to the recipe for the 14 unlisted non-GLM keys (Llama-405B/70B, DeepSeek family, gpt-oss-120b) and 7 GLM ones. That contradicts the precedence documented here, in models.yaml and in ARCHITECTURE.md. The list is now the exact union of models.yaml env: keys. 2. FABRIC_SUBNET defaulted to the lab prefix 10.158., changing which NIC every model advertises on any cluster that happens to have a 10.158.x address, and the slurm never forwarded it, so an operator override applied on the host but not inside the container (split-brain MASTER_ADDR vs host_ip). Now opt-in: unset = hostname -I's first entry, i.e. develop's behaviour, and the var is forwarded. 3. GLM_KERNEL_PATCH / GLM_BACKEND_PATCH bind-mounts overlaid vLLM sparse-MLA source from a host path. They do not exist on develop, nothing else in the tree reads them, and a host-path source overlay is precisely the runtime patching this branch removed in favour of the image pin. Dropped. Also correct the `-e` -> `-f` rationale on the RDMA lib mounts: a dangling symlink already fails -e; what -f excludes is a path that resolves to a non-regular file. * [GLM-5.1] Correct false and stale claims in the GLM docs Verified against ROCm/mori, vllm-project/vllm and the pinned vLLM fork with gh api: - Dockerfile stage 1 claimed the MoRI pin is "v1.2.1 (the version for the 06_29 mori121 image)". MORI_REF=42e895472b08 is ROCm/mori#366 (2026-06-05) and is 32 commits BEHIND tag v1.2.1 (e31d426a, 2026-06-25). The text was copied from the base Dockerfile, which really does pin v1.2.1, and it also inverted models.yaml, which lists the mori121 image as NOT SUPPORTED. - Dockerfile stage 3 claimed the vLLM pin is "the 06_29 validated Wide-EP WRITE-mode branch (matches the published mori121 image)". It is the GLM DSA v0.27 branch; the 06_29 vLLM predates the in-source DSA fixes and cannot serve GLM-5.1-FP8 at all. - VLLM_REF is documented as "(HEAD cda3648602) = upstream v0.27 tip + 7 commits". The branch tip is d723eb305 (+9); it advanced on 2026-08-15, before the MAD commit that measured the 3.4x decode speedup. Since VLLM_REF is a branch name, `docker build` resolves it to that day's tip. Record the real state and tell builders to pass an explicit --build-arg VLLM_REF= for an auditable rebuild. - STATUS said "NIAH 2k-35k = 10/10" where the recipe records 51/60 and 55/60, and the vLLM pin note said "1P/1D + 2P/1D + 1P/2D ... decode PIECEWISE" where the validated topologies are 1P/1D EP8 + 2P/2D EP16 and decode is FULL_AND_PIECEWISE. - The `docker build -f` example inside the glmv5.1 Dockerfile named the BASE Dockerfile. - models.yaml's "3 recipe-defining ways" bullet still described decode as PIECEWISE with a ~3.7x ITL win after 638a033 moved it to FULL_AND_PIECEWISE. Reframed: the ~3.7x is PIECEWISE over eager, FULL_AND_PIECEWISE is the further 2.11x/1.70x already measured below. - benchmark_xPyD.sh attributed "302ms vs ~89ms" to residual JIT absorbed by the first measured cell. models.yaml and moriio.sh attribute the same 302->88ms to the 8192-token-wide MoRI EP all2all buffer, fixed by --max-num-batched-tokens. Note the distinction rather than leaving two root causes for one number. - vLLM #40018 is filed against gfx950 while this stack is gfx942; say so. - README told users to build "the one fullsource Dockerfile" and never mentioned that GLM-5.1-FP8 needs the per-model glmv5.1 image the same README's table cites. - benchmark_niah.sh passes every NIAH_* through explicitly except NIAH_SEEDS, which only worked by inheritance. * [GLM-5.1] Sync the gate_check allowlist mirror with the slurm The branch added GLM-5.1-FP8 to VALID_MODELS, MORI_EP_VALID_MODELS and WIDE_EP_ONLY_MODELS in run_xPyD_models.slurm but not to gate_check.sh's copy, so the new model's gate had no coverage. Add it, plus the four cases that pin the intended envelope (moriio wideEP only; TP and DeepEP rejected). The header also claimed the test "runs the REAL gate code path, so a future edit to the lists/gate is caught" — it does not, it re-runs a hand-kept mirror, which is exactly why this drift was possible. Say what it actually does. * [GLM-5.1] Name the validated v0.27 image tag and the 0/10 evidence. The IMAGE CONTRACT block said which ref the supported image is built from but never named the published artifact, and gave no evidence for why older images are unsupported. Name the tag (rocmshared/pytorch-private:glm5.1-vllm027-b8) and record the measurement: the pre-v0.27 image rocm/pytorch-private:glm-dockerimage-built-09072026 scored NIAH 2k 0/10 on the 1P/1D EP8 smoke, slurm job 216847. This branch ships no runtime patchers, so an unsupported image has no fallback and raises no error -- it serves and returns wrong output. Say that plainly. The glmv5.1 Dockerfile gets the same tag and the same evidence so the two files name one artifact; its vLLM stage already states that the mori121 image's vLLM cannot serve GLM-5.1-FP8, which is what models.yaml says too. Comment-only: no flags, no env, no pins move. * [GLM-5.1] Pin VLLM_REF to the validated commit instead of a branch name VLLM_REF defaulted to glm5.1-dsa-wideEP_on_vllm-v0.27, a mutable branch name, so `docker build` resolved to whatever the tip was on the day of the build. The branch has already moved once since validation (cda3648602 -> e8c186f71/d723eb305), which means "the image is the contract" could not hold: two builds of the same Dockerfile could ship different engines. Pin the tip that every number in models.yaml was measured on, d723eb305e, and update the places that described the ref as a branch (models.yaml, connectors/moriio.sh, the Dockerfile header). Moving the pin now requires an explicit --build-arg and a re-validation, which is the intent. * [GLM-5.1] Correct the described lineage of the pinned vLLM commit The header called dedbf6be8b the "upstream v0.27 tip". It is not: against vllm-project/vllm, dedbf6be8b is a main-branch commit of 2026-08-09, 270 commits ahead of the v0.27.0 tag and 17 behind releases/v0.27.0. Verified with the GitHub compare API, which also confirms the pinned d723eb305e is that commit plus exactly 9 ROCm commits, 0 behind. State what the commit actually is so a reader does not go looking for these fixes on the v0.27 release branch. * [GLM-5.1] Point VLLM_REF back at the branch, pending the mutable-ref fix Reverts the SHA pin from the previous commit. The ref goes back to glm5.1-dsa-wideEP_on_vllm-v0.27 so this Dockerfile stays consistent with how the other vllm_dissag images name their vLLM ref; the mutable-branch problem is being handled separately rather than solved one Dockerfile at a time. The REPRODUCIBILITY note stays and now names the validated commit outright, so an auditable rebuild is still one build arg away: --build-arg VLLM_REF=d723eb305eb78d1bda0ed357b2b54cc29487221f The lineage correction from the previous commit is kept and reworded for the branch form: the tip is dedbf6be8b plus 9 ROCm commits, and dedbf6be8b is upstream main, not the v0.27 release. * [GLM-5.1] Trim the connector_runtime_patch comment back to develop's connector_runtime_patch is a return-0 no-op shared by every moriio model, and this branch had grown its comment from 7 lines to 24. The MoRI paragraph was wrong for the model that dominates this connector: it told the reader their pin is 42e895472b08 and lacks ROCm/mori #424, #432 and #436, but DeepSeek-V3 builds from vllm_disagg_inference.ubuntu.amd.Dockerfile with MORI_REF=v1.2.1, 32 commits ahead of that sha and containing those fixes. The pin story is already told, correctly scoped to the GLM image, in the glmv5.1 Dockerfile. The GLM paragraph restated a rule the existing text already makes generic and duplicated detail that belongs in the Dockerfile header and the models.yaml recipe. Restores develop's wording plus one line pointing GLM readers at models.yaml, so the hunk is a one-line diff against develop. * [GLM-5.1] Drop the GLM line from connector_runtime_patch The function is unedited by this branch, so it needs no note from it. Leaves connector_runtime_patch byte-identical to develop; the GLM image contract is stated in models.yaml and the glmv5.1 Dockerfile, which is where it belongs. * [GLM-5.1] Correct the necessity claim for use_inductor_graph_partition The comments in moriio.sh and models.yaml stated the flag was required for GLM, citing "RuntimeError: unknown parameter type" on the first real MLA decode from the stable-ABI concat_and_cache_mla boxed kernel failing to compose in the Dynamo-FX partitioned graph. That does not reproduce on the pinned image. A/B on 1P/1D EP8 (slurm 217352.91 with the flag off vs .108 with it on) scored NIAH 2k 10/10 in both arms, with the emitted --compilation-config verified per arm from the live engine cmdline, and zero occurrences of that RuntimeError in any log. A matched 1024/1024 concurrency-8 sweep put mean TPOT at 42.57 ms off vs 42.11 ms on, inside run-to-run noise. Comments only; no behavior change. GLM keeps the flag ON because that is the configuration every published number was measured with, which the comments now say instead of claiming a correctness requirement. * [GLM-5.1] Drop the models.yaml protect-list, restore develop's env precedence _RECIPE_ENV_KEYS existed to add a third precedence tier so models.yaml could override image-baked ENV while a submit-time -e still won. That tier only matters if a supported image bakes a recipe key, and none does: docker inspect of the glmv5.1 image shows no KV_BLOCK_SIZE, no VLLM_ROCM_USE_AITER_MLA and no cudagraph settings -- its only ENV lines are build and toolchain vars. The case it defended is a lab image baking DeepSeek tuning, already unsupported for GLM. Against that it required a 35-key list in the slurm to be kept in sync with every models.yaml env: block, in a different file, or the documented precedence silently breaks for whichever key was missed. vllm_disagg.sh's yaml-env block is now byte-identical to develop. Diff versus develop drops from 79 to 44 lines in the slurm and 38 to 9 in the driver, the remainder being the unrelated FABRIC_SUBNET fix. Verified: DRY_RUN argv for GLM-5.1-FP8, DeepSeek-V3 and DeepSeek-R1 identical before and after; GLM still emits use_inductor_graph_partition. Submit-time overrides still win (USE_INDUCTOR_GRAPH_PARTITION=0 drops the key, KV_BLOCK_SIZE=16 beats the recipe's 1) because an exported value is in os.environ under develop's rule. tests/run_all.sh 25/25 plus the gate suite. * [GLM-5.1] Make per-shape benchmark warmup opt-in, not default-on benchmark_xPyD.sh is develop's default sweep path, shared by every model, so gating per-shape warmup on ${SHAPE_WARMUP:-1} would have changed the measured TPOT of recipes already validated without it (DeepSeek-V3/-R1, Llama-70B, gpt-oss-120b). Default to off and let GLM-5.1-FP8 opt in from its models.yaml env: block, which is the configuration its published latency numbers were measured under. Other models' sweeps are now byte-identical to develop. Also forward SHAPE_WARMUP and USE_INDUCTOR_GRAPH_PARTITION into the container conditionally. docker run only passes variables named in an explicit -e, so without this a submit-time export of either key never crosses the container boundary and models.yaml applies unopposed -- meaning neither knob could be A/B-tested from a wrapper script. Unset means not forwarded, so neither can shadow a recipe on a normal run. Covered by five new assertions in tests/argv_assert.sh, including that GLM-5.1-FP8 is the only warmup opt-in. * [GLM-5.1] Trim the comment bloat out of the GLM enablement diff The GLM changes added 409 comment lines against 271 lines of code, a 60% comment ratio that buried the actual recipe. models.yaml carried 130 comment lines for one model entry where the DeepSeek entries use roughly 17, and moriio.sh -- a file every model shares -- carried 25 comment lines for 6 lines of code. This removes 131 of them without touching a single line of code. Three kinds of comment are gone. Duplication: the MoRI EP buffer-width rationale was written out in full in both moriio.sh and models.yaml, so it now lives only at decode.dp where the setting actually is, and the same applies to the inductor-partition and FABRIC_SUBNET explanations. Lab narrative: the long-context collapse history, the unsupported-image inventory and the branch archaeology described a superseded image rather than the shipped one. Internal evidence: slurm step IDs, build labels and GitHub compare arithmetic that no reader outside this lab can act on. What stays is anything a future editor would break by removing the code it sits on -- KV_BLOCK_SIZE=1 for the DSA indexer, the prefill cudagraph deadlock, the sparse-indexer crash fix, the decode.dp TTFT floor -- plus the two safety warnings that describe SILENT failures: unsupported images return confident garbage rather than erroring, and long-context was fixed in-source so it must not be chased with config flags. Measurements are kept only as one-liners where they justify a non-obvious setting; the full tables moved to the work log. The Dockerfile is trimmed against its sibling vllm_disagg_inference Dockerfile rather than an arbitrary target: blocks copied from it are left byte-identical, and only the GLM-specific additions are compressed. Comment-only, proven three ways: the non-comment diff is empty, DRY_RUN argv is byte-identical for GLM, DeepSeek-V3, DeepSeek-R1, Llama-70B and DSV3+DeepEP, and the offline suites stay green at 24 gate + 30 argv assertions. * [GLM-5.1] Correct the per-shape warmup rationale to match measurement Both comments claimed the first measured cell of a shape absorbs residual JIT and reports an inflated TPOT. An A/B on allocation 217352 at 1024/1024 con=8 does not support that: with the warmup on TPOT was 42.13 ms and with it off 42.05 ms, and the ~21 s P99 TTFT in the con=1 cell that follows the warmup is present either way. State the real reason the recipe keeps it -- GLM's published figures were measured with it on -- rather than a mechanism that was not observed. Comments only; no code changes. * [GLM-5.1] Keep the FULL_AND_PIECEWISE measurement table in models.yaml The previous commit reduced this block to the two TPOT figures on the grounds that the full table was PR-description material. That was premature: the throughput column, the speedup multipliers and the NIAH accuracy row are the baseline the upstream DSA PRs get verified against, and they need to be readable next to the setting until those PRs merge. Restores the two measured rows and the accuracy line, keeps the tightened prose around them, and says in the comment that the rows are a temporary verification baseline so a later reader knows they can go once upstream lands. Comment-only: non-comment diff empty, DRY_RUN argv byte-identical for GLM, DeepSeek-V3, DeepSeek-R1, Llama-70B and DSV3+DeepEP, offline suites green at 24 gate + 30 argv. * [GLM-5.1] Drop the FABRIC_SUBNET opt-in; it fixes nothing by default The launcher advertises hostname -I's first entry on both sides of the container boundary, and socket_barrier binds --local-ip host_ip while peers dial --node-ips IPADDRS. When that first entry is not routable between nodes the barrier never completes and the job hangs at "Waiting for nodes" with no error. That is a real defect on this cluster: job 211823 came up with IPADDRS=10.224.2.61,10.224.0.194 and hung, and 211879 had one node on the overlay and one on the fabric and hung with a single engine up. Both predate GLM, so it is a cross-model launcher defect. The mitigation carried here was an opt-in FABRIC_SUBNET prefix filter applied to both lookups. Defaulting to unset, it fixes nothing for anyone: the next time enumeration flips, the run hangs exactly as before because nobody exports the variable. It was documented in no README, mentioned in no PR text and covered by no test, which makes it dead code in an enablement PR. A real fix means automatic routable-NIC selection, which changes NIC choice for every model on every cluster and belongs in its own launcher PR with its own testing. Reverting also removes a latent defect: _pick_fabric_ip had dropped develop's NR==1 guard, so a multi-line srun stdout would have yielded a multi-line MASTER_ADDR. vllm_disagg.sh is now byte-identical to develop, so this PR no longer touches it, and the slurm diff drops from 46 lines to 21. DRY_RUN argv is unchanged for GLM-5.1-FP8, DeepSeek-V3 and Llama-70B; suites stay green at 24 gate + 30 argv; no FABRIC_SUBNET reference remains in scripts/ or docker/. * Trim the glmv5.1 Dockerfile header to the reviewed comment budget Removes 14 comment lines of within-file duplication from the header, bringing the Dockerfile to the 180-line comment target from review: - the docker build invocation appeared twice (header and "build context" block); keep the first, which also shows DOCKER_IMAGE_NAME - WITH_NIXL=0/=1 was explained in the header and again at the ARG; keep the ARG-site copy, leave the one-line mention in the connector summary - BASE_IMAGE override was stated in both STATUS and the closing block; fold the ~30-60 min compile time into STATUS and drop the block - the "add a vllm_disagg_inference..Dockerfile per future model" guidance duplicates README.MD Also corrects a false claim inherited from the sibling Dockerfile: the pin bullet said "validated recipe knobs baked as ENV", which contradicts section 5 ("The image ships NO runtime recipe / tuning / platform ENV"). The recipe lives in models.yaml and is applied at launch. Comment-only: the non-comment diff is empty, all 34 Dockerfile directives are unchanged, and tests/run_all.sh is green (30 argv assertions + gate checks). * [GLM-5.1] Bump the component pins to the validated fresh stack Moves the recipe off the older pin set onto the stack that was re-validated end-to-end, and records the EP32 fix. MoRI 42e895472b08 -> 624002c897a3 AITER ROCm/aiter e03fa6040 -> raviguptaamd/aiter 624e43586b flydsl >=0.1.7,<0.1.9 -> ==0.3.1 vLLM glm5.1-dsa-wideEP_on_vllm-v0.27 (branch) -> 094820b5d (pinned) router ravgupta/discovery-dp-rank-roundrobin -> ravgupta/dp-roundrobin-on-tip Why each: * vLLM 094820b5d is branch glm5.1-dsa-wideEP_on_d626108b: the 10 GLM DSA commits rebased onto upstream d626108b (2026-08-20). It includes the MoRI combine() original-topk fix (623fdc946b), which resolves the EP32 token corruption the header previously called a KNOWN OPEN DEFECT. Measured at 4P/4D EP32 on the pre-bump stack: NIAH 8k 0/10 -> 10/10, perf-neutral. Pinned to a sha rather than a branch so builds are reproducible. * AITER: newer aiter claims native gfx942 support for gqa64 fp8 decode and routes GLM-5.1 to mla_a8w8_qh64_qseqlen1_gqaratio64_v3_ps, which GPU-faults. The 7-line fix lets gqa64 fall through to the capture-safe persistent view-fold, preserving cudagraph decode (FULL_AND_PIECEWISE) instead of forcing eager. Filed upstream as ROCm/aiter#4957. The fork ref is TEMPORARY -- revert AITER_REPO to ROCm/aiter once #4957 merges. This is what unblocks moving off the old e03fa6040 pin at all. * MoRI 624002c897a3 and flydsl 0.3.1 are the co-requisites of that aiter bump (flydsl_fp8_mqa_logits is missing on 0.1.x). Validation on the bumped pins: 1P/1D EP8 and 2P/2D EP16, prefill MoRI-HT eager + decode MoRI-LL FULL_AND_PIECEWISE -- NIAH 2k 10/10, 8k 10/10; 8192/1024 con32 TPOT ~59 ms. 4P/4D EP32 has NOT been re-validated on these pins (needs an 8-node allocation); the header says so explicitly rather than implying coverage. Also refreshes the header comments that described the old pins, so the rationale text matches the ARGs. Co-Authored-By: Claude * [GLM-5.1] Fix two claims the pin bump left behind Follow-up to 094d108, which bumped the pins but left three statements that no longer match the file. VLLM_REF is now a 40-char sha, so the REPRODUCIBILITY note contradicted itself: it still warned that "this default is a BRANCH NAME, so it is mutable" and then told the reader to pass --build-arg VLLM_REF= for an auditable rebuild. The mutable-ref hazard is what the bump removed, so say that instead. That note also claimed the pin "is the tip every number in models.yaml was measured on". It is not: the decode-cudagraph table in models.yaml was measured on d723eb305e and has not been re-measured on 094820b5d. Left as written it would credit the new stack with the old stack's numbers, which is the same class of false-provenance claim this PR has already corrected twice (e8a633d4e, 59f5fd8f8). README.MD's model table still called EP32 "known-broken" while the Dockerfile header two files away now says it is fixed. Matches the table to the header's wording, including that EP32 has not been re-validated on the bumped pins. Comment- and doc-only: no Dockerfile directive changes, all ARG values untouched. * [GLM-5.1] Make 4P/4D EP32 launchable, pin the router, record the validation Three fixes from validating an image built off this recipe unmodified (glm5.1-pr206-9f16198, built from PR head 9f16198). 1. EP32 could not boot. GPU_MEMORY_UTILIZATION 0.8 OOMs during EngineCore init at 4P/4D ("Tried to allocate 6.00 GiB ... 157.95 GiB allocated"), which SIGTERMs every rank; the socket barrier then prints "Waiting for nodes" for ~30 min, so it reads as a hang rather than a crash. The MoRI dispatch/combine buffers scale with EP width but the fraction did not. The default is now derived from decode EP width: 0.70 at EP>=32, 0.8 below. Measured: 0.70 boots (66.79 GiB KV cache) and holds through con=64; 2P/2D EP16 at 0.8 is unchanged (validated at stock in the same session). Submit-time -e still overrides. This had to go in the launcher, not models.yaml: the -e on the docker run line is unconditional, and vllm_disagg.sh skips any yaml key already in the environment, so the models.yaml GPU_MEMORY_UTILIZATION key never takes effect. Left in place (it documents intent) but the comment now says where the live value is. Note this is the only change here outside the glmv5.1 Dockerfile. 2. ROUTER_REF was a branch name, so two builds on different days could ship different routers - the same mutable-ref hazard the VLLM_REF sha pin removed. Pinned to 82dc9811. The comment above it still described the old PR #181 branch and its two commits; upstream has since merged the round-robin fix (prefill_dp_round_robin), so only the KV-notify fix is carried locally now. 3. STATUS said 4P/4D was fixed but not re-validated on these pins. It is now validated, so the claim is replaced with the measurements: NIAH, 3 seeds: 2P/2D 2k 9.7/10, 8k 9.7/10 | 4P/4D 2k 9.7/10, 8k 9.3/10 8192/1024 con32 TPOT median: 2P/2D 60.0 ms | 4P/4D 70.4 ms 0 failed requests across all 8 cells, 0 GPU faults 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 - so the text says ~93-97% retrieval rather than claiming a clean 10/10. Also records that EP32 costs ~17% TPOT over EP16, so nobody reads 4P/4D as a latency win. README.MD's model table matched to the same wording. Co-Authored-By: Claude * [GLM-5.1] Drop the quoted performance figures from the recipe Comment-only. No ARG, no env value, no code line changes. Reading this the way an outside reviewer would: the recipe quoted a lot of specific latency and throughput numbers and none of them said which stack produced them. That is a problem in both directions. Some were measured on a stack this recipe no longer pins - the decode-cudagraph table in models.yaml is from d723eb305e, as 9f16198 already noted - so a reader would credit the current pins with older results. The rest were real but were one configuration on one cluster, and a figure like "TPOT 41.8 ms" sitting in a recipe file reads as a spec that CI or a customer can hold you to. Neither belongs in a file whose job is to pin components. Removed the numbers, kept the reasoning that made each setting the right choice: FULL_AND_PIECEWISE is still justified (faster for decode, gain narrows as EP widens, costs boot time), max-num-batched-tokens 2048 still explains the TTFT-vs-latency tradeoff, and the inductor-partition A/B is still recorded as neutral. Anyone can re-measure; nobody has to trust a constant that may have gone stale. The STATUS block keeps the NIAH results - those are a correctness claim and are what the EP32 fix is actually about - and now states that a latency sweep was run without quoting it, with a note to re-measure on your own hardware. Also drops "1P/1D EP8" from the README model table. 1P/1D was validated on an earlier stack, not on these pins; only 2P/2D EP16 and 4P/4D EP32 were. I added that claim in error - it is the same false-provenance mistake 9f16198 corrected. The EP32 OOM comment in the launcher keeps its error string, since that is what someone chasing the "Waiting for nodes" hang will grep for, but loses the KV-cache figure. Co-Authored-By: Claude * [GLM-5.1] Address Cemberk's 3 MAD-integration design issues Issue 1 — NIAH is unschedulable and reports nothing: - Add pyt_vllm_disagg_mori_glm-5.1-fp8_niah card with BENCHMARK_SCRIPT=niah and niah tag so it's CI-schedulable and tag-filterable. - Add --niah mode to parse_to_csv.py: parses NIAH summary lines and emits madengine perf.csv rows (retrieval score per context length). - benchmark_niah.sh now calls parse_to_csv.py after the harness run. Issues 2 & 3 — GPU_MEMORY_UTILIZATION precedence fix: - Fix the broken yaml path: slurm now conditionally forwards GPU_MEMORY_UTILIZATION only if set by submit-time, and always passes _GPU_MEM_UTIL_LAUNCHER_DEFAULT as a separate var. - vllm_disagg.sh applies the launcher fallback after yaml loads, so: submit-time -e > models.yaml > launcher fallback (0.70@EP>=32, 0.8) - Remove inert GPU_MEMORY_UTILIZATION from GLM and Qwen yaml (they now use the topology-aware fallback). - Update header comment and README to document the precedence. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: raviguptaamd Co-authored-by: Claude Co-authored-by: Ravi Gupta --- ...gg_inference.glmv5.1.ubuntu.amd.Dockerfile | 351 ++++++++++++++++++ scripts/vllm_dissag/README.MD | 36 +- scripts/vllm_dissag/benchmark_niah.py | 85 ++++- scripts/vllm_dissag/benchmark_niah.sh | 24 +- scripts/vllm_dissag/benchmark_xPyD.sh | 27 ++ scripts/vllm_dissag/connectors/moriio.sh | 13 +- scripts/vllm_dissag/models.json | 63 ++++ scripts/vllm_dissag/models.yaml | 107 +++++- scripts/vllm_dissag/parse_to_csv.py | 81 ++++ scripts/vllm_dissag/run_xPyD_models.slurm | 37 +- scripts/vllm_dissag/tests/argv_assert.sh | 21 ++ scripts/vllm_dissag/tests/gate_check.sh | 19 +- scripts/vllm_dissag/vllm_disagg.sh | 4 + 13 files changed, 826 insertions(+), 42 deletions(-) create mode 100644 docker/vllm_disagg_inference.glmv5.1.ubuntu.amd.Dockerfile 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/scripts/vllm_dissag/README.MD b/scripts/vllm_dissag/README.MD index 17660abe..52679149 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,7 +256,7 @@ 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 ``` 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..ce6eb1aa 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). @@ -219,7 +225,7 @@ connector_launch_worker() { --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 @@ -242,6 +248,7 @@ connector_launch_worker() { "${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 diff --git a/scripts/vllm_dissag/models.json b/scripts/vllm_dissag/models.json index 6fc3849a..de50a465 100644 --- a/scripts/vllm_dissag/models.json +++ b/scripts/vllm_dissag/models.json @@ -278,6 +278,69 @@ }, "args": "-N 2 -n 2" }, + { + "name": "pyt_vllm_disagg_mori_glm-5.1-fp8", + "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" + ], + "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_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", diff --git a/scripts/vllm_dissag/models.yaml b/scripts/vllm_dissag/models.yaml index 23d66059..d3511624 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 / @@ -143,7 +145,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 +179,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..86a8366c 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,13 @@ 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) +# 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 ;; + niah) BENCHMARK_SCRIPT_FILE="benchmark_niah.sh" ;; + *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context, niah)" >&2; exit 1 ;; esac if [[ ! -f "$BENCHMARK_SCRIPT_FILE" ]]; then echo "Error: selected benchmark script '$BENCHMARK_SCRIPT_FILE' not found in $(pwd)." >&2 @@ -525,7 +530,10 @@ 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 @@ -533,6 +541,19 @@ done [ -d /etc/libibverbs.d ] && _RDMA_MOUNTS="$_RDMA_MOUNTS -v /etc/libibverbs.d:/etc/libibverbs.d:ro" echo "[host-rdma] mounts: $_RDMA_MOUNTS" +# 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 \ --device /dev/kfd \ @@ -590,6 +611,13 @@ docker run --rm \ ${KV_CACHE_DTYPE:+-e KV_CACHE_DTYPE=$KV_CACHE_DTYPE} \ ${MORIIO_TOY_PROXY:+-e MORIIO_TOY_PROXY=$MORIIO_TOY_PROXY} \ ${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 +641,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/vllm_disagg.sh b/scripts/vllm_dissag/vllm_disagg.sh index 06fbf84f..f13b3262 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 From 615c7a1c7fccee8bba8c2a314f4a225dde06071f Mon Sep 17 00:00:00 2001 From: afriedri Date: Fri, 28 Aug 2026 16:56:54 -0500 Subject: [PATCH 2/5] add ep back to qwen 235b + 397b (#232) --- scripts/vllm/configs/extended.yaml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 From 8346e20082851c0ae2a899cf10e637b2a6e01928 Mon Sep 17 00:00:00 2001 From: afriedri Date: Fri, 28 Aug 2026 16:58:27 -0500 Subject: [PATCH 3/5] DSv4 max-num-batched-tokens 8192 (#235) * DSv4 max-num-batched-tokens 8192 VLLM now defaults to 16384; too much at 1m context * Remove space --------- Co-authored-by: Rahul Garg --- scripts/vllm/configs/default.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 From b33923b3ce43293b4d53f49b2efba866ddaeae89 Mon Sep 17 00:00:00 2001 From: Mir Mustafa Ali Date: Mon, 31 Aug 2026 09:49:58 -0500 Subject: [PATCH 4/5] Add AgentX suite in SGLang and vLLM disaggregated launchers (#213) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * sglang_disagg: add agentic replay benchmark as a recipe-native option Adds aiperf-driven agentic trace-replay (inferencex-agentx-mvp) as a first-class benchmark for the sglang_disagg recipe, selectable via BENCHMARK_SCRIPT_FILE=benchmark_agentic.sh (or the AGENTIC=1 shorthand) and launched through the standard `sbatch run_xPyD_models.slurm` entrypoint. Recipe-only: no GitHub Actions / CI workflow dependency. Harness (new): - scripts/common/agentic_lib.sh (aiperf driver; passes --server-metrics) - scripts/common/validate_agentic_result.sh - scripts/common/preflight_weights.sh - scripts/common/verify_agentic_image.sh - scripts/sglang_disagg/benchmark_agentic.sh Launcher (gated behind agentic selection; default sweep unchanged): - sglang_disagg_mori_io_ep.sh: --enable-metrics --enable-metrics-for-all-schedulers, strip --disable-radix-cache, export AGENTIC_SERVER_METRICS + curl probe - run_xPyD_models.slurm: 10.158.x IP preference, agentic/metrics env passthrough, AGENTIC=1 shorthand - mori_ep_env.sh: MORI_IB_GID_INDEX=3, NCCL_IB_GID_INDEX=3 Docs: - README.MD: Agentic Replay Benchmark section (command, env knobs, supported models/modes, outputs, validation) - models.yaml: note on launcher-gated radix-cache strip - salloc_launch.sh: agentic sample commands Verified via cluster run 207554 (1P/1D DeepSeek-V3): server_metrics.present=true, gpu_cache_hit_rate=32.4%, error_rate ~2.9% (PASS). * vllm_dissag: add agentic replay benchmark as a recipe-native option Port the sglang agentic trace-replay feature to the vLLM disaggregated P/D launcher, reusing the endpoint-agnostic shared harness (scripts/common/*): - benchmark_agentic.sh: aiperf agentic hook, AGENTIC_PORT resolved from the vLLM router port (BENCHMARK_PORT, default 30000); finds agentic_lib.sh via NIXL_COOKBOOK_PATH/NIXL_REPO_DIR/_here ../common candidates. - run_xPyD_models.slurm: AGENTIC=1 shorthand + agentic BENCHMARK_SCRIPT case; forward agentic env (AGENTIC_CONC, DURATION, AGENTIC_PORT, MODEL_PREFIX, ENABLE_PREFIX_CACHE, ENABLE_SERVER_METRICS, DRY_RUN, ...) and NIXL_REPO_DIR so the in-container hook resolves common/. - vllm_disagg.sh: export AGENTIC_PORT=BENCHMARK_PORT before the hook exec. - connectors/moriio.sh, rixl.sh: on the agentic path strip the hardcoded --no-enable-prefix-caching (prefix cache ON) and export AGENTIC_SERVER_METRICS at the backend serve ports; gated so the default sweep argv is unchanged. - models.yaml: note prefix caching is re-enabled via the connector strip. - common/verify_agentic_image.sh: add a vLLM branch (AGENTIC_ENGINE=vllm) that checks the KV transfer backend import + vLLM OpenAI API server module. * vllm_dissag: add agentic_models_shim for /v1/models under MoRIIO discovery The vLLM PD router returns 503 on /v1/models in MoRIIO discovery mode (empty HTTP worker registry), but the shared agentic_lib.sh gates router readiness + served-model resolution on /v1/models. Add a lightweight shim that answers /v1/models (gated on the router's /health) and stream-proxies all other routes to the PD router; AGENTIC_PORT points at the shim. Wired into vllm_disagg.sh on the agentic path only (default sweep unchanged). * agentx: add config-driven multi-workload agentic benchmarking suite Unify regenerated (self-contained) and HF-downloaded trace sources behind one config-driven harness that replays both through aiperf's inferencex-agentx-mvp scenario against a single disaggregated endpoint. - scripts/common/agentx/: generic seed-deterministic weka_trace corpus generator (gen_agentx_profile.py), profile conformance verifier (verify_agentx_profile.py, 13-axis), and a PyYAML-optional config/profile loader (agentx_config.py). Ships agentic.example.yaml + caseA/caseB/custom/ small profiles. - scripts/common/benchmark_agentic_suite.sh: iterate a workloads list against one served endpoint; per workload materialize+verify (source=profile) or resolve --public-dataset (source=hf), run a context-compat check, sweep concurrency, and write per-workload result dirs + a combined suite_summary.json. - agentic_lib.sh: add materialize_corpus() and context_compat_check(); branch resolve_trace_loader() for the profile source; per-workload knobs in build_replay_cmd() (trajectory ratios, --max-context-length incl. Case-B 524288, num-dataset-entries); honor MODEL=auto and guard it under set -u. - sglang_disagg/vllm_dissag benchmark_agentic.sh route through the suite driver (AGENTIC_CONFIG); run_xPyD_models.slurm forward AGENTIC_CONFIG/AGENTIC_WORKLOAD/ SUITE_CORPUS_DIR/MAX_MODEL_LEN into the container. Validated end-to-end on MI300X with Qwen3-0.6B: verifier 13/13 (PyYAML and forced-fallback, byte-identical corpus) and a suite run at 0.0% error rate. * agentx: rename conformance presets caseA/caseB -> conformance_256k/conformance_512k Rename the two ExplainX/#173 conformance workload presets to role-based names: profiles/caseA.yaml -> profiles/conformance_256k.yaml profiles/caseB.yaml -> profiles/conformance_512k.yaml Updates the presets' name: field and all live preset references (agentic.example.yaml workload name:/preset:, agentx_config.py and gen/verify docstrings). The id_prefix session-id salt is intentionally preserved as "caseA" in both profiles (and DEFAULT_ID_PREFIX in gen_agentx_profile.py) to reproduce ROCm/MAD #173 byte-for-byte. Genuine #173 historical filename references (gen_caseA_conformance.py, verify_caseB.py, etc.) are kept as provenance. Also bundles the related, in-progress agentic-suite work: preset-chain merge + Tier 1/ Tier 2 corpus filtering (filter_weka_corpus.py, inferencex_256k.yaml), HF loader context gating, and suite driver/lib support. * agentic-suite: remove scripts/vllm_dissag/ (moved to local branch miali/vllm-disagg-agentic) * agentx: add core benchmarking docs (README, profiles authoring guide, scenarios cookbook) documents the config schema, profile/preset authoring + verify pre-gate, and a scenarios cookbook; docs-only, no code changes. * agentx docs: add run-timing notes, verifier output sample, and concurrency-sweep clarification * docs(sglang_disagg): add agentic-suite invocation examples and forward suite env vars Add 5 copy-paste agentic-suite invocation examples, an env-var table, and cross-links to the agentx core docs in the sglang_disagg README. Forward SUITE_CORPUS_FORCE, AGENTIC_STRICT_CONTEXT, and AGENTIC_HF_ISL_TAIL into the container from run_xPyD_models.slurm so the documented knobs take effect. * Restore scripts/vllm_dissag/ to match develop Revert the removal from 99faef3; the vLLM disaggregated launcher must stay in sync with develop on this branch. * vllm_dissag: add agentic-suite replay benchmark as recipe-native option Mirror the sglang_disagg agentic integration into the vLLM disaggregated launcher so AGENTIC=1 runs the multi-workload AgentX suite. Adds the benchmark_agentic.sh hook (AGENTIC_PORT resolves to the vLLM proxy BENCHMARK_PORT, with suite routing when AGENTIC_CONFIG/AGENTIC_WORKLOAD is set and DRY_RUN support), the agentic_models_shim.py for /v1/models under MoRIIO discovery, connector metrics export + prefix-cache handling for moriio/rixl, suite env-var forwarding in run_xPyD_models.slurm, and README invocation examples cross-linked to the agentx docs. * vllm_dissag: SSE-safe agentic shim + configurable rixl+TP barrier port 1. agentic_models_shim.py _proxy: force Content-Type: text/event-stream for streaming relays (request "stream": true or SSE upstream) so aiperf strips the SSE framing instead of marking every request invalid. 2. rixl.sh + run_xPyD_models.slurm + tests/run_interactive.sh: make the rixl+TP container barrier port env-overridable (BARRIER_PORT, default 5000), forwarded into the container and used by the host-side fuser cleanup. * docs(agentx): document all suite env vars and minimal required set Add a grouped environment variable reference covering previously undocumented user-facing knobs (context, warmup, HF/corpus, timing), add explicit minimal-required commands for the AGENTIC_CONFIG suite and config-less AGENTIC_WORKLOAD entry points, and split the misleading AGENTIC_WORKLOAD example into config-less vs config-filter variants. * docs(agentx): add integration template hook and close onboarding gaps Add scripts/common/agentx/templates/ with a copy-paste benchmark_agentic hook template and an integration guide so any backend (single-node, disaggregated, or other framework) can wire up the agentic suite by filling a few marked fields. Also make the core README Quick start runnable and close remaining onboarding gaps (serve-endpoint prerequisite, accurate RESULT_DIR default, qualified slurm path). * disagg: auto-resolve served max_model_len from real P/D worker In disaggregated serving the front-end routers (sglang_router :2322, vLLM agentic_models_shim) do not advertise max_model_len, so the agentic benchmark's context auto-detect returned 0 and ran with an unvalidated window (operators had to pin MAX_MODEL_LEN by hand). Each backend hook now queries the real prefill worker (first host:port in AGENTIC_SERVER_METRICS): - vLLM: worker /v1/models ModelCard.max_model_len - SGLang: worker /v1/models, falling back to /get_server_info On success it exports MAX_MODEL_LEN so it flows into both the suite and single-workload paths; if a worker is set but unresolvable it fails fast (no silent default). Non-disagg (no AGENTIC_SERVER_METRICS) and DRY_RUN fall through to the existing behavior unchanged. Adds node-free offline tests (stdlib HTTP stub) covering resolve, fail-fast, non-disagg fallthrough, and propagation into the suite plan. * docs(agentx): document disagg context resolution and sync template Close disagg doc/template drift from the served-context-window work: add the disaggregated worker-probe block and AGENTIC_RESOLVE_ONLY to the integration template, document AGENTIC_RESOLVE_ONLY and the dual role of AGENTIC_SERVER_METRICS (also probed for max_model_len on disagg), and correct the front-end /v1/models auto-detect wording in the sglang and vLLM launcher READMEs. * agentx: add source:corpus replay and configurable aiperf scenario Two suite configurability additions (both default to the existing behavior when unused, so the legacy single-workload path is byte-identical): - source: corpus replays an existing on-disk weka_trace corpus directory as-is via --input-file (no generate, no download). Verification is optional: it runs the profile pre-gate only when a preset/inline profile is supplied. ISL-tail resolution: explicit isl_tail: > profile-derived > 1048576 (the gate caps to the served window, so it only over-WARNs). Trajectory/num-dataset-entries follow the real-trace (hf) path. - run.scenario makes the aiperf --scenario configurable (default inferencex-agentx-mvp), overridable by AGENTIC_SCENARIO env or a per-workload scenario: key (re-derived each iteration so it can't leak). The workloads: list is unchanged. Docs (agentic.example.yaml, SCENARIOS.md) updated accordingly. * docs(agentx): add HOW_IT_WORKS explainer and onboarding section Add HOW_IT_WORKS.md: a toy-example walkthrough of the trace-replay mechanism (content-independence, seed-deterministic corpus generation, multi-turn prefix reuse via hash_ids, theoretical vs GPU cache-hit) with mermaid diagrams, mapped back to the conformance_256k profile. Also adds an "Onboarding new users" section to the agentx README (reading path, DRY_RUN preview, AGENTIC_WORKLOAD=small smoke run). Committing the doc also resolves the README's existing See-also link to it. * docs(agentx): cross-link launchers to onboarding; fix content-independence Add reciprocal "start here" links from the sglang and vLLM launcher READMEs into the AgentX core README/SCENARIOS/HOW_IT_WORKS, and extend the onboarding path to name both launchers (not just sglang). Correct HOW_IT_WORKS.md Section 2: input/prompt text is synthesized client-side by aiperf's PromptGenerator keyed by hash_ids (shared ids -> identical blocks -> cache hits); only output text is server-produced (measured, then discarded). Diagram updated to match. * fix(agentx): honor config concurrency/duration; log WARN verdict as WARN The backend hooks (sglang/vllm benchmark_agentic.sh) and the suite driver injected AGENTIC_CONC=16 / DURATION=120 (and 900) into the already-exported env before agentx_config.py runs, so the loader's env-override layer always clobbered the YAML run.concurrency/duration. Remove those default injections so precedence is: explicit env > config file > loader defaults (16/900). The legacy single-workload path is unchanged: build_replay_cmd still self-defaults via ${AGENTIC_CONC:-16}/${DURATION:-120}. Also add agentic_warn() and use it for the context clamp-and-run notice (CONTEXT_VERDICT=WARN), which previously printed an [agentic][ERROR] prefix despite being a warn+clamp+run, not a failure. Includes a co-located fix in resolve_served_max_model_len: the /server_info fallback now reads context_length only (not max_total_num_tokens, which is the KV-cache token-pool budget, not the per-request context window). * fix(agentx): close file handles, guard required fields + CLI arg values Address 3 validated CRITICAL review findings in the AgentX helper scripts: - Resource leaks: wrap json.dump/json.load in `with open(...)` in gen_agentx_profile.py (per-session write loop) and verify_agentx_profile.py (per-corpus read loop). - Missing required-field handling: emit a clear error (instead of a raw KeyError traceback) when a profile omits required distribution fields in gen_agentx_profile.py and verify_agentx_profile.py, and when agentx_config._isl_tail() has neither clamps.isl nor isl_p. - Hand-rolled arg loops: add a shared _argval() guard so a flag given without its value exits with a usage message instead of StopIteration, in gen_agentx_profile.py, verify_agentx_profile.py, filter_weka_corpus.py, and agentx_config.py. Issue #4 (walrus `:=` in vllm_dissag/benchmark_parser.py) is a FALSE POSITIVE: the codebase targets Python 3.8+ (runtime 3.10; pandas), so no version guard is needed. Happy path unchanged; offline suite 30/30. * fix(agentic): harden launcher IP resolution, MoRI preflight, weight check, uv install - A: run_xPyD_models.slurm derives MASTER_ADDR/node IP by preferring the 10.158.x fabric IP (unchanged on OCI), falling back to the first hostname -I address, and failing fast with a clear error when none resolve (was: empty value -> barrier hang). Applied at both sites. - C: verify_agentic_image.sh skips the Mooncake import check when RUN_MORI=1 (MoRI IO backend does not require Mooncake), keeping it for the default path. - F: preflight_weights.sh captures srun exit status and treats a nonzero exit or empty output as a hard failure instead of silently passing. - D: agentic_lib.sh ensure_agentic_uv prefers an on-PATH uv, pins the uv version, and gates the remote installer behind AGENTIC_ALLOW_UV_INSTALL=1. - B: gen_agentx_profile.py documents that floor division for block count is intentional for byte-for-byte reproducibility (no arithmetic change). * fix(agentic): forward AGENTIC_ALLOW_UV_INSTALL/AGENTIC_UV_VERSION into container The ISSUE D uv-install opt-in gate lives in agentic_lib.sh, which runs inside the disagg docker container. Env vars only reach that container through the launcher's explicit -e allowlist, so add the two new uv knobs (defaulting empty -> gate stays safe-by-default) to keep cluster runs able to install/use uv when AGENTIC_ALLOW_UV_INSTALL=1 is set. * test(agentx): add offline test + cluster submit/harvest harness run_offline.sh: node-free suite covering config resolution, corpus gen/verify, DRY_RUN plans, source:corpus, and scenario overrides. submit_matrix.sh / harvest.sh: sbatch matrix launcher and results harvester for the disagg agentic benchmark cells. * fix(agentic): default AGENTIC_ALLOW_UV_INSTALL=1 on the direct agentic path The agentic benchmark hooks are an explicit opt-in entrypoint, so default the pinned uv install to on (overridable with AGENTIC_ALLOW_UV_INSTALL=0) instead of making developers set it manually. Non-agentic launcher paths keep the gate off. * fix(agentic): forward HF_HOME conditionally in sglang launcher Use ${HF_HOME:+-e HF_HOME=$HF_HOME} so HF_HOME is only passed into the container when set on the host, matching the vLLM sibling launcher. The previous unconditional -e HF_HOME="${HF_HOME:-}" injected an empty HF_HOME on the default (non-agentic) sweep, which huggingface_hub reads as present-but-empty and resolves HF_HUB_CACHE to a relative hub/ under CWD. Restores byte-for-byte behavior for non-agentic sweeps. * fix(agentic): forward optional env conditionally in sglang launcher Convert the empty-default env-forwarding lines (-e VAR="${VAR:-}") in the per-node docker run block to the conditional idiom ${VAR:+-e VAR=$VAR}, so an unset host var no longer injects an empty string into the container (the same class of bug already fixed for HF_HOME in 6a8f32b). This matches the vLLM sibling launcher scripts/vllm_dissag/run_xPyD_models.slurm. Lines with a meaningful default (DOCKER_IMAGE_NAME, USE_CX7_NICS, BENCHMARK_COMBINATIONS) and always-set required vars are left unchanged. Addresses Cemberk's review item #3 on ROCm/MAD PR #213. * fix(agentic): make fabric subnet prefix configurable in launcher The sglang disagg launcher hardcoded a 10.158.x fabric-subnet preference when resolving MASTER_ADDR / node IPs from multi-homed `hostname -I` output. That embeds one cluster's (OCI useocpm2m) addressing repo-wide and silently changes node coordination for every sweep on other clusters. Introduce FABRIC_SUBNET_PREFIX (default "10.158.") near the node-info section and reuse an escaped, start-anchored regex (FABRIC_SUBNET_RE) at both the MASTER_ADDR and per-node IP selection sites. The default keeps OCI behavior byte-for-byte unchanged when unset; an empty prefix skips the prefix match and falls straight through to the first-address fallback. The existing first-address fallback and fail-fast-on-empty logic are preserved. Addresses Cemberk's review item #4 on ROCm/MAD PR #213. * feat(agentic): declare 4 agentic models.json entries; gate on mode; remove AGENTIC=1 Replace the implicit "AGENTIC=1 layered over every model" mode with declared, agentic-tagged models.json entries and mode-based gating, addressing Cemberk's PR #213 (ROCm/MAD) headline feedback. The default sweep path is byte-for-byte unchanged. - Add 4 validated, `agentic`-tagged models.json entries (BENCHMARK_SCRIPT=agentic): * pyt_sglang_disagg_mori_io_agentic_qwen3-32b (AGENTIC_WORKLOAD=conformance_256k) * pyt_sglang_disagg_mori_dp_agentic_deepseek-v3 (DP_MODE=1, AGENTIC_CONFIG suite) * pyt_vllm_disagg_mori_agentic_deepseek-v3 (AGENTIC_CONFIG suite) * pyt_vllm_disagg_mori_agentic_deepseek-v3-5layer (AGENTIC_CONFIG suite) - SGLang launcher: add a BENCHMARK_SCRIPT case selector (sweep|agentic) mirroring vLLM; sweep -> benchmark_xPyD.sh (default preserved), agentic -> benchmark_agentic.sh. - Hard-remove the redundant AGENTIC=1 shim from both launchers (breaking change). - Re-gate shared code on the mode (BENCHMARK_SCRIPT==agentic) instead of the resolved filename (sglang_disagg_mori_io_ep.sh, connectors/rixl.sh, connectors/moriio.sh, vllm_disagg.sh); execution paths keep BENCHMARK_SCRIPT_FILE. - Forward BENCHMARK_SCRIPT into both containers so in-container gates resolve. - models.yaml + READMEs: note the cache-flag strip is driven by BENCHMARK_SCRIPT=agentic and exercised only by agentic-tagged entries; add a validated-configs table, a `madengine run --tags agentic` discovery note, and a Breaking Change note. * chore(agentic): finish AGENTIC=1 removal in submit_matrix + stale docs Follow-up to the mode-based agentic gating. AGENTIC=1 was hard-removed from the launchers, so complete the removal in the remaining references: - submit_matrix.sh: functional fix — the campaign driver's per-cell env array used AGENTIC=1, which would now submit in SWEEP mode; switch to BENCHMARK_SCRIPT=agentic. - salloc_launch.sh: commented example recipes AGENTIC=1 -> BENCHMARK_SCRIPT=agentic. - benchmark_agentic.sh (vllm) + AgentX templates: reword header comments off the removed shorthand. - READMEs: reword the Breaking Change note off the literal AGENTIC=1 token. No AGENTIC=1 references remain (grep -rn 'AGENTIC=1' scripts/ is empty); default sweep path unchanged. * refactor(agentx): unify agentic hooks into shared --backend script Collapse the two byte-identical-except-three-axes agentic benchmark hooks (scripts/sglang_disagg/benchmark_agentic.sh, scripts/vllm_dissag/benchmark_agentic.sh) and the stale copy-paste template (scripts/common/agentx/templates/benchmark_agentic.template.sh) into ONE shared, --backend-parameterized source of truth at scripts/common/benchmark_agentic.sh plus two ~12-line shims that exec it with --backend sglang|vllm. The three per-backend axes now live in a single `case "$backend"`: 1. AGENTIC_PORT default (sglang 2322; vllm the BENCHMARK_PORT/PROXY_PORT/8000 cascade). 2. ctx-window resolver endpoint list (`ctx_endpoints=("path|kind" ...)`): sglang probes /v1/models then /get_server_info; vllm probes only /v1/models. The shared script reuses the sglang two-endpoint python parser (kind-driven) for both, looping the list. 3. lib-locator repo-dir candidates: a superset covering MOONCAKE_REPO_DIR, NIXL_COOKBOOK_PATH and NIXL_REPO_DIR (harmless when unset). Drift-proofing: the resolver, suite/DRY_RUN dispatch, and the 5-call replay tail now exist once, so future edits can no longer diverge between backends. Adding a backend is a case arm + a shim instead of copying a template (docs updated in templates/README.md). Net ~410 deletions vs ~206 insertions. Verified offline (no GPU): bash -n on all three; run_offline 31/31 PASS; both resolve_ctx_offline suites ALL PASS (sglang exercises the two-endpoint arm through the shim); per-backend port/endpoint parity confirmed. Cemberk PR #213 review item #7. * docs(agentx): consolidate SCENARIOS into README; trim launcher env duplication (Cemberk #9) - Fold all 10 SCENARIOS.md recipes into README.md as a new ## Examples section (Option A full fold); de-duplicate Scenario 1 and Scenario 9 against the existing Quick start / Preview sections via cross-links. - Delete scripts/common/agentx/SCENARIOS.md. - Retarget all cross-references (11 sites) to README.md#examples anchors. - Trim the duplicated AgentX env tables in the sglang/vllm launcher READMEs to a pointer at the central env reference plus launcher-specific rows only. * refactor(agentx): dedup per-backend ctx-resolve stub + offline test into common (Cemberk #8) * docs(agentx): add verify_agentx_profile.py example, container-path troubleshooting, and error reference Add a cookbook example showing the agentx_config.py --emit-json -> verify_agentx_profile.py flow (success + failure output), a troubleshooting entry for AGENTIC_CONFIG paths not visible inside the container bind mounts, and a consolidated symptom -> cause -> fix error reference table. * fix(agentx): resolve N1 parser + robustness guards - N1: Parse same-indent block sequences in fallback YAML loader - Gen guards: block_size>=1, turns non-empty/length-match/positive-weights - Verify guard: missing 'requests' in session JSON - Preset guard: file-not-found check before load Tests: add offline-gate assertions (fallback parser + guards) * fix(agentx): resolve shell driver failure masking - Suite driver: replace \`|| true\` with _suite_failed propagation - SGLang launcher: add set -o pipefail before tee - SGLang launcher: expand \$HOME in AGENTIC_CONFIG before docker forward Tests: add static guards (|| true removal, pipefail, \$HOME expansion) * fix(agentx): resolve blocking issues from Cemberk PR #213 2nd review Address two blocking issues preventing merge: 1. vLLM launcher /home/miali expansion - Add AGENTIC_CONFIG expansion before container forward (matches SGLang) - Use pattern (safer than eval) - Extend offline test to cover vLLM launcher 2. Example config references in models.json - Change 4 entries to use AGENTIC_WORKLOAD shorthand - Prevents 4.75h runs pointing at demo config - Comment out my_corpus/my_case in agentic.example.yaml - Add test_offline.yaml for offline tests (separate from user example) - Fix continue-on-failure for missing input_dir Testing: - Offline tests: 34/34 pass - JSON/YAML syntax: valid - DRY_RUN: verified with conformance_256k Resolves: Cemberk blocking issues #1 and #2 Ready for re-review Co-Authored-By: Claude Opus 4.6 * fix(vllm_dissag): add niah to BENCHMARK_SCRIPT valid-list error message Merge resolution kept only the agentic variant of the invalid-value message; niah is a working case arm and belongs in the valid list. --------- Co-authored-by: Mir Ali Co-authored-by: MIR-AMD Co-authored-by: Claude Opus 4.6 Co-authored-by: MIR-AMD --- scripts/common/agentic_lib.sh | 491 ++++++++++ scripts/common/agentx/HOW_IT_WORKS.md | 358 ++++++++ scripts/common/agentx/README.md | 846 ++++++++++++++++++ scripts/common/agentx/agentic.example.yaml | 92 ++ scripts/common/agentx/agentx_config.py | 580 ++++++++++++ scripts/common/agentx/filter_weka_corpus.py | 132 +++ scripts/common/agentx/gen_agentx_profile.py | 228 +++++ scripts/common/agentx/profiles/README.md | 181 ++++ .../agentx/profiles/conformance_256k.yaml | 34 + .../agentx/profiles/conformance_512k.yaml | 37 + .../agentx/profiles/custom.example.yaml | 50 ++ .../agentx/profiles/inferencex_256k.yaml | 20 + scripts/common/agentx/profiles/small.yaml | 15 + scripts/common/agentx/templates/README.md | 76 ++ scripts/common/agentx/tests/_stub_server.py | 46 + scripts/common/agentx/tests/harvest.sh | 107 +++ .../agentx/tests/resolve_ctx_offline.sh | 122 +++ scripts/common/agentx/tests/run_offline.sh | 329 +++++++ scripts/common/agentx/tests/submit_matrix.sh | 130 +++ scripts/common/agentx/tests/test_offline.yaml | 17 + .../common/agentx/verify_agentx_profile.py | 177 ++++ scripts/common/benchmark_agentic.sh | 152 ++++ scripts/common/benchmark_agentic_suite.sh | 261 ++++++ scripts/common/preflight_weights.sh | 48 + scripts/common/validate_agentic_result.sh | 75 ++ scripts/common/verify_agentic_image.sh | 133 +++ scripts/sglang_disagg/README.MD | 197 ++++ scripts/sglang_disagg/benchmark_agentic.sh | 12 + scripts/sglang_disagg/models.json | 68 ++ scripts/sglang_disagg/models.yaml | 11 +- scripts/sglang_disagg/run_xPyD_models.slurm | 84 +- scripts/sglang_disagg/salloc_launch.sh | 15 + .../sglang_disagg/sglang_disagg_mori_io_ep.sh | 58 +- scripts/vllm_dissag/README.MD | 171 ++++ scripts/vllm_dissag/agentic_models_shim.py | 183 ++++ scripts/vllm_dissag/benchmark_agentic.sh | 12 + scripts/vllm_dissag/connectors/moriio.sh | 22 +- scripts/vllm_dissag/connectors/rixl.sh | 28 +- scripts/vllm_dissag/models.json | 66 ++ scripts/vllm_dissag/models.yaml | 8 + scripts/vllm_dissag/run_xPyD_models.slurm | 34 +- scripts/vllm_dissag/tests/run_interactive.sh | 2 +- scripts/vllm_dissag/vllm_disagg.sh | 30 + 43 files changed, 5720 insertions(+), 18 deletions(-) create mode 100755 scripts/common/agentic_lib.sh create mode 100644 scripts/common/agentx/HOW_IT_WORKS.md create mode 100644 scripts/common/agentx/README.md create mode 100644 scripts/common/agentx/agentic.example.yaml create mode 100644 scripts/common/agentx/agentx_config.py create mode 100644 scripts/common/agentx/filter_weka_corpus.py create mode 100644 scripts/common/agentx/gen_agentx_profile.py create mode 100644 scripts/common/agentx/profiles/README.md create mode 100644 scripts/common/agentx/profiles/conformance_256k.yaml create mode 100644 scripts/common/agentx/profiles/conformance_512k.yaml create mode 100644 scripts/common/agentx/profiles/custom.example.yaml create mode 100644 scripts/common/agentx/profiles/inferencex_256k.yaml create mode 100644 scripts/common/agentx/profiles/small.yaml create mode 100644 scripts/common/agentx/templates/README.md create mode 100644 scripts/common/agentx/tests/_stub_server.py create mode 100644 scripts/common/agentx/tests/harvest.sh create mode 100755 scripts/common/agentx/tests/resolve_ctx_offline.sh create mode 100644 scripts/common/agentx/tests/run_offline.sh create mode 100644 scripts/common/agentx/tests/submit_matrix.sh create mode 100644 scripts/common/agentx/tests/test_offline.yaml create mode 100644 scripts/common/agentx/verify_agentx_profile.py create mode 100755 scripts/common/benchmark_agentic.sh create mode 100644 scripts/common/benchmark_agentic_suite.sh create mode 100755 scripts/common/preflight_weights.sh create mode 100755 scripts/common/validate_agentic_result.sh create mode 100755 scripts/common/verify_agentic_image.sh create mode 100755 scripts/sglang_disagg/benchmark_agentic.sh create mode 100755 scripts/vllm_dissag/agentic_models_shim.py create mode 100755 scripts/vllm_dissag/benchmark_agentic.sh 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 - < **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 be224de4..80f3735d 100644 --- a/scripts/sglang_disagg/models.json +++ b/scripts/sglang_disagg/models.json @@ -510,5 +510,73 @@ "BENCHMARK_COMBINATIONS": "1024/1024" }, "args": "-N 2 -n 2" + }, + { + "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 88f22e46..a87bbbcc 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 34949ddd..29c8499e 100755 --- a/scripts/sglang_disagg/run_xPyD_models.slurm +++ b/scripts/sglang_disagg/run_xPyD_models.slurm @@ -114,6 +114,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/"}" @@ -300,14 +314,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 @@ -353,6 +402,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,) @@ -394,14 +446,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 f2244936..633cb0c2 100755 --- a/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh +++ b/scripts/sglang_disagg/sglang_disagg_mori_io_ep.sh @@ -179,7 +179,26 @@ PREFILL_MODEL_CONFIG="${MODEL_BASE_FLAGS} ${MODEL_MODE_FLAGS} ${MODEL_PREFILL_FL DECODE_MODEL_CONFIG="${MODEL_BASE_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" @@ -534,15 +553,35 @@ PY echo "" 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 @@ -552,6 +591,11 @@ PY echo "Killing the co-located prefill server (pid=${_node0_prefill_pid})" kill "${_node0_prefill_pid}" + if [[ "${benchmark_status}" -ne 0 ]]; then + echo "ERROR: benchmark failed with status ${benchmark_status}" >&2 + exit "${benchmark_status}" + fi + elif [[ "$NODE_RANK" -ge 1 && "$NODE_RANK" -lt "$xP" ]]; then echo "${host_name}:${host_ip} is Prefill Node (Model: ${MODEL_NAME:-default})" # NODE_RANK 0..xP-1 map directly to PREFILL_NODE_RANK 0..xP-1 (proxy co-located on NODE_RANK=0). @@ -588,6 +632,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 @@ -675,6 +720,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_dissag/README.MD b/scripts/vllm_dissag/README.MD index 52679149..7b10df30 100644 --- a/scripts/vllm_dissag/README.MD +++ b/scripts/vllm_dissag/README.MD @@ -263,6 +263,177 @@ 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/connectors/moriio.sh b/scripts/vllm_dissag/connectors/moriio.sh index ce6eb1aa..3eb1094f 100644 --- a/scripts/vllm_dissag/connectors/moriio.sh +++ b/scripts/vllm_dissag/connectors/moriio.sh @@ -181,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 @@ -221,7 +231,7 @@ 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}" \ @@ -241,7 +251,7 @@ 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} \ @@ -320,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": "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 d3511624..2321fe13 100644 --- a/scripts/vllm_dissag/models.yaml +++ b/scripts/vllm_dissag/models.yaml @@ -34,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: diff --git a/scripts/vllm_dissag/run_xPyD_models.slurm b/scripts/vllm_dissag/run_xPyD_models.slurm index 86a8366c..8b29966d 100755 --- a/scripts/vllm_dissag/run_xPyD_models.slurm +++ b/scripts/vllm_dissag/run_xPyD_models.slurm @@ -430,13 +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" ;; + 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, niah)" >&2; exit 1 ;; + *) 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 @@ -475,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}" @@ -487,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; @@ -540,6 +545,8 @@ 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) @@ -585,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 \ @@ -599,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} \ @@ -610,6 +639,7 @@ 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} \ 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 f13b3262..7b794265 100755 --- a/scripts/vllm_dissag/vllm_disagg.sh +++ b/scripts/vllm_dissag/vllm_disagg.sh @@ -251,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.." From bfb0477ef8eab4c9594c234be8954e0362563033 Mon Sep 17 00:00:00 2001 From: Rahul Garg Date: Mon, 31 Aug 2026 10:50:54 -0400 Subject: [PATCH 5/5] Jax maxtext v26.6 (#238) --- .dockerignore | 22 + .gitignore | 16 + benchmark/jax_maxtext/README.md | 681 ++++++++---------- docker/jax_maxtext.ubuntu.amd.Dockerfile | 33 - .../primus_maxdiffusion.ubuntu.amd.Dockerfile | 84 +++ docker/primus_maxtext.ubuntu.amd.Dockerfile | 75 ++ models.json | 24 +- scripts/Primus | 2 +- .../extract_maxdiffusion_perf.py | 232 ++++++ scripts/jax-maxdiffusion/get_models_json.py | 139 ++++ scripts/jax-maxdiffusion/run.sh | 141 ++++ .../jax-maxtext/env_scripts/deepseek2_16b.yml | 68 -- .../env_scripts/deepseek2_env_16b.sh | 43 -- .../env_scripts/gfx950_deepseek2_16b.yml | 68 -- .../env_scripts/gfx950_deepseek2_env_16b.sh | 44 -- .../env_scripts/gfx950_llama2_70b.yml | 76 -- .../env_scripts/gfx950_llama2_70b_env.sh | 44 -- .../env_scripts/gfx950_llama2_7b.yml | 65 -- .../env_scripts/gfx950_llama2_7b_env.sh | 44 -- .../env_scripts/gfx950_llama3.1_405b.yml | 46 -- .../env_scripts/gfx950_llama3.3_70b.yml | 75 -- .../env_scripts/gfx950_llama3.3_70b_env.sh | 44 -- .../env_scripts/gfx950_llama3_70b.yml | 75 -- .../env_scripts/gfx950_llama3_70b_env.sh | 44 -- .../env_scripts/gfx950_llama3_8b.yml | 65 -- .../env_scripts/gfx950_llama3_8b_env.sh | 44 -- .../env_scripts/gfx950_mixtral_8x7b.yml | 71 -- .../env_scripts/gfx950_mixtral_8x7b_env.sh | 44 -- .../env_scripts/gfx950_qwen3_14b.yml | 62 -- .../env_scripts/gfx950_qwen3_14b_env.sh | 44 -- .../env_scripts/gfx950_qwen3_30b_a3b.yml | 65 -- .../env_scripts/gfx950_qwen3_30b_a3b_env.sh | 44 -- .../jax-maxtext/env_scripts/llama2_70b.yml | 75 -- .../jax-maxtext/env_scripts/llama2_70b_env.sh | 43 -- scripts/jax-maxtext/env_scripts/llama2_7b.yml | 64 -- .../jax-maxtext/env_scripts/llama2_7b_env.sh | 43 -- .../jax-maxtext/env_scripts/llama3.3_70b.yml | 74 -- .../env_scripts/llama3.3_70b_env.sh | 43 -- .../jax-maxtext/env_scripts/llama3_70b.yml | 74 -- .../jax-maxtext/env_scripts/llama3_70b_env.sh | 43 -- scripts/jax-maxtext/env_scripts/llama3_8b.yml | 64 -- .../jax-maxtext/env_scripts/llama3_8b_env.sh | 43 -- .../jax-maxtext/env_scripts/mixtral_8x7b.yml | 71 -- .../env_scripts/mixtral_8x7b_env.sh | 43 -- scripts/jax-maxtext/env_scripts/qwen3_14b.yml | 62 -- .../jax-maxtext/env_scripts/qwen3_14b_env.sh | 43 -- .../jax-maxtext/env_scripts/qwen3_30b_a3b.yml | 65 -- .../env_scripts/qwen3_30b_a3b_env.sh | 43 -- scripts/jax-maxtext/extract_maxtext_perf.py | 95 +++ scripts/jax-maxtext/get_models_json.py | 144 ++++ .../jax-maxtext/gpu-rocm/deepseek_v2_16b.sh | 75 -- scripts/jax-maxtext/gpu-rocm/llama2_70b.sh | 87 --- .../gpu-rocm/llama2_70b_multinode.sh | 177 ----- scripts/jax-maxtext/gpu-rocm/llama2_7b.sh | 76 -- .../gpu-rocm/llama2_7b_multinode.sh | 165 ----- scripts/jax-maxtext/gpu-rocm/llama3.3_70b.sh | 86 --- scripts/jax-maxtext/gpu-rocm/llama3_70b.sh | 86 --- .../gpu-rocm/llama3_70b_multinode.sh | 177 ----- .../gpu-rocm/llama3_70b_multinode_metal.sh | 204 ------ scripts/jax-maxtext/gpu-rocm/llama3_8b.sh | 77 -- .../gpu-rocm/llama3_8b_multinode.sh | 165 ----- scripts/jax-maxtext/gpu-rocm/readme.md | 39 - .../jax-maxtext_benchmark_report.py | 100 --- .../jax-maxtext_benchmark_report.sh | 167 ----- .../jax-maxtext_benchmark_setup.sh | 90 --- .../jax_maxtext_multinode_benchmark.sh | 389 ---------- scripts/jax-maxtext/models.json | 468 ------------ scripts/jax-maxtext/run.sh | 134 +++- scripts/primus_train/get_models_json.py | 25 +- scripts/primus_train/run.sh | 2 + tools/fetch_primus.sh | 76 ++ tools/run_models.py | 11 +- 72 files changed, 1444 insertions(+), 5188 deletions(-) create mode 100644 .dockerignore delete mode 100644 docker/jax_maxtext.ubuntu.amd.Dockerfile create mode 100644 docker/primus_maxdiffusion.ubuntu.amd.Dockerfile create mode 100644 docker/primus_maxtext.ubuntu.amd.Dockerfile create mode 100644 scripts/jax-maxdiffusion/extract_maxdiffusion_perf.py create mode 100644 scripts/jax-maxdiffusion/get_models_json.py create mode 100755 scripts/jax-maxdiffusion/run.sh delete mode 100644 scripts/jax-maxtext/env_scripts/deepseek2_16b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/deepseek2_env_16b.sh delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_deepseek2_16b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_deepseek2_env_16b.sh delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_llama2_70b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_llama2_70b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_llama2_7b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_llama2_7b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_llama3.1_405b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_llama3.3_70b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_llama3.3_70b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_llama3_70b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_llama3_70b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_llama3_8b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_llama3_8b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_mixtral_8x7b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_mixtral_8x7b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_qwen3_14b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_qwen3_14b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_qwen3_30b_a3b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/gfx950_qwen3_30b_a3b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/llama2_70b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/llama2_70b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/llama2_7b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/llama2_7b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/llama3.3_70b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/llama3.3_70b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/llama3_70b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/llama3_70b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/llama3_8b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/llama3_8b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/mixtral_8x7b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/mixtral_8x7b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/qwen3_14b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/qwen3_14b_env.sh delete mode 100644 scripts/jax-maxtext/env_scripts/qwen3_30b_a3b.yml delete mode 100644 scripts/jax-maxtext/env_scripts/qwen3_30b_a3b_env.sh create mode 100644 scripts/jax-maxtext/extract_maxtext_perf.py create mode 100644 scripts/jax-maxtext/get_models_json.py delete mode 100644 scripts/jax-maxtext/gpu-rocm/deepseek_v2_16b.sh delete mode 100644 scripts/jax-maxtext/gpu-rocm/llama2_70b.sh delete mode 100644 scripts/jax-maxtext/gpu-rocm/llama2_70b_multinode.sh delete mode 100644 scripts/jax-maxtext/gpu-rocm/llama2_7b.sh delete mode 100644 scripts/jax-maxtext/gpu-rocm/llama2_7b_multinode.sh delete mode 100644 scripts/jax-maxtext/gpu-rocm/llama3.3_70b.sh delete mode 100644 scripts/jax-maxtext/gpu-rocm/llama3_70b.sh delete mode 100644 scripts/jax-maxtext/gpu-rocm/llama3_70b_multinode.sh delete mode 100644 scripts/jax-maxtext/gpu-rocm/llama3_70b_multinode_metal.sh delete mode 100644 scripts/jax-maxtext/gpu-rocm/llama3_8b.sh delete mode 100644 scripts/jax-maxtext/gpu-rocm/llama3_8b_multinode.sh delete mode 100644 scripts/jax-maxtext/gpu-rocm/readme.md delete mode 100644 scripts/jax-maxtext/jax-maxtext_benchmark_report.py delete mode 100755 scripts/jax-maxtext/jax-maxtext_benchmark_report.sh delete mode 100755 scripts/jax-maxtext/jax-maxtext_benchmark_setup.sh delete mode 100644 scripts/jax-maxtext/jax_maxtext_multinode_benchmark.sh delete mode 100644 scripts/jax-maxtext/models.json mode change 100644 => 100755 scripts/jax-maxtext/run.sh create mode 100755 tools/fetch_primus.sh 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 e2d872c7..8eb205f6 100644 --- a/.gitignore +++ b/.gitignore @@ -76,3 +76,19 @@ venv.bak/ *.log *.out *.html + +# 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/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/jax-maxdiffusion/extract_maxdiffusion_perf.py b/scripts/jax-maxdiffusion/extract_maxdiffusion_perf.py new file mode 100644 index 00000000..389cd2fe --- /dev/null +++ b/scripts/jax-maxdiffusion/extract_maxdiffusion_perf.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python3 +""" +Extract JAX/MaxDiffusion performance metrics and write a madengine +multiple_results CSV (one row per metric). + +MaxDiffusion writes per-step metrics two ways: + + 1. A per-step stdout line (parsed as a fallback):: + + completed step: 12, seconds: 0.83, TFLOP/s/device: 210.4, loss: 0.123 + + Under the Primus launcher this line does NOT reliably reach the captured + training log (raw ``print`` from maxdiffusion.max_logging is dropped while + the trainer runs), so it cannot be relied on. + + 2. A JSON-lines metrics file, written directly by the trainer when + ``config.metrics_file`` is set (``max_utils.write_metrics_locally`` -> + ``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/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: