docs(installation-jax): add bare-metal JAX/MaxText install guide and scripts

docs/01-getting-started/bare-metal-installation.md

@@ -6,6 +6,8 @@ It is derived from the official training `Dockerfile` and installs the same comp
 
 > **Important**: This is a long, build-heavy process. A full from-source build (Flash Attention, TransformerEngine, aiter, Primus-Turbo, FBGEMM, rocSHMEM, etc.) can take **several hours** and needs a machine with many CPU cores, plenty of RAM, and tens of GB of free disk. The Docker image remains the recommended and best-supported path. Use this guide only when containers are not an option.
 
+> **Using the JAX MaxText backend instead?** This guide builds the PyTorch stack (Megatron-LM / TorchTitan). For the JAX / MaxText backend, follow the leaner [JAX bare-metal installation guide](bare-metal-installation-jax.md) instead.
+
 ---
 
 ## Quick path: automated install scripts

docs/01-getting-started/installation.md

@@ -117,7 +117,8 @@ Reproducing the training environment on the host means building the **same stack
 
 Follow the full, step-by-step guide here, which includes the exact pinned versions, environment variables, and automated install scripts:
 
-- **[Bare-metal installation: build the Primus training stack from source (no Docker)](./bare-metal-installation.md)**
+- **[Bare-metal installation (PyTorch: Megatron-LM / TorchTitan): build the Primus training stack from source (no Docker)](./bare-metal-installation.md)**
+- **[Bare-metal installation (JAX / MaxText): build the Primus JAX training stack from source (no Docker)](./bare-metal-installation-jax.md)**
 
 ### Verify
 

tools/installation-jax/README.md

@@ -0,0 +1,163 @@
+# Primus JAX / MaxText environment in a venv (no docker, no sudo) — v26.5
+
+Reproduces the Primus **v26.5 JAX training Dockerfile** in a Python virtual
+environment. Same package pins as the Dockerfile, adapted for a bare-metal host
+with no root and no containers.
+
+This is the JAX/MaxText counterpart to `tools/installation/` (the PyTorch /
+Megatron / TorchTitan stack). Use this one if you want to run the **JAX MaxText**
+training backend of Primus.
+
+## Why it differs from the Dockerfile
+
+| Constraint | Dockerfile | Here |
+|---|---|---|
+| ROCm | release **tarball** → `/opt/rocm` | same tarball → **`$ROCM_DIR`** (`$PRIMUS_JAX_BASE/rocm`, user-writable, **no sudo**) |
+| GPU arch | gfx942 + gfx950 | **auto-detected** from the host (`rocminfo`/KFD sysfs) for source builds/TE |
+| Build dir | container FS | venv on **`$PRIMUS_JAX_BASE`** (persistent); transient sources on local `/tmp` |
+| System deps | `apt install ...` | **skipped** (no sudo); documented in the guide's Section 2 |
+| MaxText setup | `setup.sh` (runs `apt`, prompts for a venv) | Python steps only (apt is a one-time root action; venv already made) |
+
+The key reason no sudo is required: the ROCm tarball is extracted into a
+user-writable dir (`$ROCM_DIR`) instead of `/opt/rocm`, so we don't depend on
+system ROCm or apt for the ROCm toolchain itself.
+
+> **Python 3.12+ required (3.12 preferred).** MaxText requires Python ≥ 3.12 (the
+> PyTorch recipe works on 3.10; this one does not). 3.12 is preferred because the
+> prebuilt ROCm wheels are `cp312`. You do **not** need `sudo` or a PPA:
+> - The scripts use `python3.12` (preferred) or `python3.13` if on PATH.
+> - Otherwise, if [`uv`](https://docs.astral.sh/uv/) is installed, `env.sh`
+>   auto-detects a uv-managed `>= 3.12` interpreter and `setup.sh` runs
+>   `uv python install 3.12` for you when none exists yet.
+> - No `uv`? Install it once (no root) and re-run:
+>   `python3 -m pip install --user uv` (or `curl -LsSf https://astral.sh/uv/install.sh | sh`).
+> - To force a specific interpreter: `export PRIMUS_PYTHON=/path/to/python3.12`.
+>
+> On Ubuntu 22.04, `apt install python3.12` fails (jammy has no such package) —
+> use the `uv` path above instead.
+
+> **Host OS: Ubuntu 24.04 / glibc ≥ 2.38 recommended, but 22.04 works.** The
+> prebuilt `transformer_engine_rocm_jax` wheel is built against the Dockerfile's
+> `ubuntu:24.04` base and needs `glibc ≥ 2.38` + `GLIBCXX_3.4.32` (GCC 13/14).
+> On Ubuntu 22.04 (glibc 2.35) the prebuilt TE wheel can't load
+> (`version 'GLIBC_2.38' not found`). **`setup.sh` handles this automatically:**
+> the `te` stage detects the host glibc and builds TransformerEngine from source
+> when it is < 2.38, so `bash setup.sh` works on both 22.04 and 24.04. (Only the
+> prebuilt TE wheel needs glibc ≥ 2.38; the ROCm JAX/PJRT wheels load on glibc
+> 2.35.) Check with `ldd --version`.
+
+## Run it
+
+```bash
+cd tools/installation-jax
+# PRIMUS_JAX_BASE is REQUIRED (no default). Point it at a directory you can write
+# to with tens of GB free — the venv, ROCm tarball, and checkouts all live here:
+export PRIMUS_JAX_BASE=/some/big/disk/primus-jax-env
+bash setup.sh                 # all default stages
+```
+
+If any stage fails the script stops immediately and prints which stage failed;
+fix the cause and re-run just that stage. Stages are idempotent:
+
+```bash
+bash setup.sh --list          # show stages
+bash setup.sh te              # reinstall just TransformerEngine
+bash setup.sh venv rocm jax   # venv + ROCm + JAX only
+```
+
+## Use the environment afterward
+
+```bash
+# Set the SAME PRIMUS_JAX_BASE you built with (required — env.sh errors without it)
+export PRIMUS_JAX_BASE=/some/big/disk/primus-jax-env
+source tools/installation-jax/env.sh   # activates venv + sets ROCm / NVTE / XLA env vars
+python -c "import jax; print(jax.devices())"
+
+# Primus is checked out at $WORKSPACE_DIR/Primus; MaxText at $MAXTEXT_DIR.
+cd "$WORKSPACE_DIR/Primus"
+./primus-cli direct -- train pretrain \
+  --config examples/maxtext/configs/MI300X/llama2_7B-pretrain.yaml
+```
+
+`env.sh` exports `MAXTEXT_PATH=$MAXTEXT_DIR`, so Primus runs the same MaxText
+checkout we installed the dependencies for.
+
+## Stages (default order, v26.5)
+
+`venv` → `rocm` → `maxtext` → `tf_source` → `jax` → `te` → `primus`
+→ `jaxreqs` → `rccl` → `manifest`
+
+- **venv** — create the venv (Python ≥ 3.12) and bootstrap `cmake`/`ninja`/`uv`.
+- **rocm** — download + extract the TheRock ROCm release tarball into `$ROCM_DIR`
+  (+ `amdsmi`).
+- **maxtext** — clone ROCm/MaxText (`release/v26.5`) and install its deps (the
+  Python part of MaxText's `setup.sh`) + the editable MaxText package.
+- **tf_source** — build **tensorflow-cpu 2.21 from source** (bazel, ~30–60 min);
+  fixes the ROCm-vs-TF LLVM symbol clash (SIGSEGV) and drops bundled NCCL.
+- **jax** — `jax`/`jaxlib` 0.10.0 + the ROCm `jax_rocm7_pjrt` / `jax_rocm7_plugin`
+  (installed after MaxText to override its stock jax).
+- **te** — prebuilt `transformer_engine_rocm_jax` wheel (+ `flax`, `pydantic`, ...);
+  auto-falls-back to a from-source build (`te_source`) on glibc < 2.38 hosts.
+- **primus** — clone Primus, init the `third_party/maxtext` submodule, drop the
+  stale `dataclasses` backports.
+- **jaxreqs** — install Primus' `requirements-jax.txt` (loguru, wandb, ...).
+- **rccl** — build **RCCL from source** and install it into `$ROCM_PATH/lib`.
+- **manifest** — dump `pip list` / `env` for reproducibility.
+
+Optional / alternative stages:
+
+- **te_source** — force the from-source TransformerEngine build regardless of
+  glibc. Normally unnecessary: the default `te` stage **auto-detects the host
+  glibc** and builds from source when it is < 2.38 (e.g. Ubuntu 22.04), where the
+  prebuilt wheel fails to load with `version 'GLIBC_2.38' not found`. Heavy build
+  (~30–60 min, compiles CK fused-attention kernels). Only the prebuilt TE wheel
+  needs glibc ≥ 2.38 — the ROCm JAX/PJRT wheels load fine on glibc 2.35.
+- **tf_cpu_fix** — lighter alternative to `tf_source`: `pip install tensorflow-cpu`
+  instead of the bazel build (avoids the bundled-NCCL clash; may still hit the
+  LLVM-symbol SIGSEGV on some ROCm 7.14 configs).
+
+```bash
+# Ubuntu 22.04 (glibc < 2.38): `bash setup.sh` already builds TE from source
+# automatically. To also skip the heavy tf_source bazel build, swap in tf_cpu_fix:
+bash setup.sh venv rocm maxtext tf_cpu_fix jax te primus jaxreqs rccl manifest
+```
+
+> **MaxText v26.5 and Primus.** MaxText v26.5 uses a **2-value**
+> `initialize()`/`run()` API. Primus `main` handles it — `MaxTextPretrainTrainer`
+> forwards `initialize()`'s tuple verbatim to `run()` (fix #912) — so this
+> recipe's pinned `release/v26.5` trains out of the box. Override
+> `MAXTEXT_BRANCH` only if you deliberately need a different MaxText release.
+
+## What is SKIPPED (needs sudo / apt — not reproducible here)
+
+- **System packages** (build toolchain, `numactl`, `gcsfuse`, RDMA/verbs libs):
+  a one-time root action, documented in Section 2 of the guide.
+- **AINIC** (`add-apt-repository`, `libionic-dev`): apt-only. Skipped.
+- **UCX + OpenMPI**: autotools source builds needed only for multi-node RDMA.
+  Single-node training works without them; see Section 4 of the guide.
+- **gcsfuse**: only needed to mount GCS buckets for data; not required for
+  synthetic-data or local-data runs.
+
+## Notes / gotchas vs. the Dockerfile
+
+- **Stage order matters (v26.5):** `maxtext` → `tf_source` → `jax` → `te`. MaxText's
+  `setup.sh` pulls in a stock `jax`/`tensorflow`; TF is then rebuilt from source and
+  the ROCm JAX/plugin is installed after (overriding MaxText's), and TE must come
+  after JAX or `jaxlib` gets clobbered. The stage order enforces this.
+- **TensorFlow + RCCL are built from source** in the default flow (matching the
+  v26.5 image). These are the two long builds; use `tf_cpu_fix` if you want to skip
+  the TF bazel build.
+- **No FlashAttention/aiter/torch stack.** The JAX MaxText path does not build the
+  PyTorch kernel libraries; attention fusion comes from `transformer_engine_rocm_jax`
+  + XLA.
+- **Two MaxText checkouts in the Docker image** (`/workspace/maxtext` and
+  `Primus/third_party/maxtext`) are collapsed here: we install deps from one
+  checkout and point `MAXTEXT_PATH` at it.
+- **Runtime-validated (single-node).** The default flow — with `te_source`
+  substituted for the prebuilt `te` — has been run end-to-end for single-node
+  MaxText pretraining on **gfx942 / Ubuntu 22.04 (glibc 2.35) / Python 3.12**.
+  On Ubuntu 24.04 (glibc ≥ 2.38) the prebuilt `te` wheel works directly. The
+  multi-node networking stack (UCX/OpenMPI/AINIC) is still not exercised here.
+
+Treat the reference `Dockerfile` as the authoritative, tested version
+combination; if you bump one pin you may need to bump the others.

tools/installation-jax/env.sh

@@ -0,0 +1,193 @@
+#!/usr/bin/env bash
+# env.sh — Primus JAX/MaxText venv environment (v26.5).
+# Source this both during the build (setup.sh does it) and every time you
+# want to USE the environment:   source env.sh
+#
+# This mirrors the v26.5 JAX training Dockerfile, adapted for a no-docker /
+# no-sudo bare-metal install. In v26.5 ROCm is delivered as a RELEASE TARBALL
+# (TheRock dist) extracted to a user-writable dir ($ROCM_DIR) — NOT the pip
+# `rocm-sdk-devel` wheels used by the v26.4 recipe. No system ROCm is required.
+#
+# NOTE: PRIMUS_JAX_BASE is REQUIRED (no default) — export it to a big-disk dir
+#       before sourcing this file. MaxText requires Python >= 3.12; see
+#       PRIMUS_PYTHON below.
+
+# ---- Install location (persistent) + transient build sources ----
+# PRIMUS_JAX_BASE is REQUIRED (there is intentionally no default): it is where
+# the venv, the extracted ROCm tarball, and the kept checkouts live. Point it at
+# a directory you can write to with tens of GB free, then (re-)source this file:
+#     export PRIMUS_JAX_BASE=/big/disk/primus-jax-env
+if [ -z "${PRIMUS_JAX_BASE:-}" ]; then
+    echo "[env] ERROR: PRIMUS_JAX_BASE is not set (no default). Set it to a directory" >&2
+    echo "[env]        you can write to with tens of GB free, then re-run:" >&2
+    echo "[env]            export PRIMUS_JAX_BASE=/big/disk/primus-jax-env" >&2
+    # env.sh is always sourced (by setup.sh and to use the env); stop sourcing
+    # without killing an interactive shell. Fall back to exit if run directly.
+    # shellcheck disable=SC2317  # 'exit 1' is reachable when this file is executed, not sourced
+    return 1 2>/dev/null || exit 1
+fi
+export PRIMUS_JAX_BASE
+export VENV_DIR="${VENV_DIR:-$PRIMUS_JAX_BASE/venv}"
+export WORKSPACE_DIR="${WORKSPACE_DIR:-$PRIMUS_JAX_BASE/workspace}"  # kept checkouts (maxtext, Primus)
+# Transient build sources: put on fast LOCAL /tmp (NFS is slow for compile I/O;
+# these dirs are deleted after each build anyway).
+export SRC_DIR="${SRC_DIR:-/tmp/primus-jax-build}"
+# ROCm release tarball (TheRock dist) is extracted here. This replaces the
+# Dockerfile's /opt/rocm and keeps the whole install no-sudo / user-writable.
+export ROCM_DIR="${ROCM_DIR:-$PRIMUS_JAX_BASE/rocm}"
+
+# MaxText clone (deps + editable install live here). Primus resolves the
+# MaxText backend from MAXTEXT_PATH (else its own third_party/maxtext), so we
+# point it at the checkout we actually installed the deps for.
+export MAXTEXT_DIR="${MAXTEXT_DIR:-$WORKSPACE_DIR/maxtext}"
+export MAXTEXT_PATH="${MAXTEXT_PATH:-$MAXTEXT_DIR}"
+
+# Python interpreter used to CREATE the venv. MaxText needs >= 3.12.
+# 3.12 is PREFERRED: it is the version the reference image (ubuntu:24.04) uses,
+# and the prebuilt ROCm wheels (transformer_engine_rocm_jax, jax_rocm7_*) are
+# built for cp312 — a 3.13 venv would need those built from source instead.
+# Resolution order:
+#   1. An explicit python3.12 / python3.13 on PATH (3.12 first).
+#   2. A uv-managed interpreter (>= 3.12), if `uv` is installed. This is the
+#      no-sudo path on distros whose apt has no python3.12 (e.g. Ubuntu 22.04):
+#      `uv python install 3.12` provides one without root. setup.sh will run
+#      that install automatically if none is found yet.
+#   3. Fall back to python3 (setup.sh errors clearly if that is < 3.12).
+if [ -z "${PRIMUS_PYTHON:-}" ]; then
+    for _py in python3.12 python3.13; do
+        if command -v "$_py" >/dev/null 2>&1; then export PRIMUS_PYTHON="$_py"; break; fi
+    done
+    if [ -z "${PRIMUS_PYTHON:-}" ] && command -v uv >/dev/null 2>&1; then
+        _uv_py="$(uv python find '>=3.12' 2>/dev/null || true)"
+        [ -n "$_uv_py" ] && [ -x "$_uv_py" ] && export PRIMUS_PYTHON="$_uv_py"
+    fi
+    export PRIMUS_PYTHON="${PRIMUS_PYTHON:-python3}"
+fi
+
+# Build parallelism.
+export MAX_JOBS="${MAX_JOBS:-128}"
+
+# ---- Target GPU architecture (auto-detected) ----
+# PYTORCH_ROCM_ARCH controls the gfx targets used for the source builds (TE,
+# RCCL) and TE's runtime NVTE_ROCM_ARCH. The name is kept for parity with the
+# Dockerfile / TE which read it. If you set it yourself it is respected as-is:
+#     export PYTORCH_ROCM_ARCH="gfx942;gfx950"
+# Otherwise it is auto-detected from the GPUs on this host. Detection needs no
+# sudo and no system ROCm: it uses `rocminfo` when available (e.g. after the
+# ROCm tarball is extracted) and otherwise falls back to the kernel KFD sysfs
+# topology, so it works even on a fresh machine before any install.
+
+# Decode a KFD gfx_target_version integer (e.g. 90402) into a gfx name (gfx942).
+_primus_gfxver_to_arch() {
+    local v="$1"
+    [ -n "$v" ] && [ "$v" != "0" ] || return 0
+    printf 'gfx%d%x%x\n' "$(( v / 10000 ))" "$(( (v / 100) % 100 ))" "$(( v % 100 ))"
+}
+
+# Echo a ";"-separated, de-duplicated list of detected gfx architectures.
+_primus_detect_gpu_arch() {
+    local arches=""
+    if command -v rocminfo >/dev/null 2>&1; then
+        arches="$(rocminfo 2>/dev/null \
+            | awk '/^[[:space:]]*Name:[[:space:]]*gfx/{print $2}' \
+            | sort -u | paste -sd';' -)"
+    fi
+    if [ -z "$arches" ] && [ -d /sys/class/kfd/kfd/topology/nodes ]; then
+        local f v a list=""
+        for f in /sys/class/kfd/kfd/topology/nodes/*/properties; do
+            [ -f "$f" ] || continue
+            v="$(awk '/^gfx_target_version/{print $2}' "$f" 2>/dev/null)"
+            a="$(_primus_gfxver_to_arch "$v")"
+            [ -n "$a" ] && list="${list}${a}"$'\n'
+        done
+        arches="$(printf '%s' "$list" | sort -u | sed '/^$/d' | paste -sd';' -)"
+    fi
+    printf '%s' "$arches"
+}
+
+if [ -z "${PYTORCH_ROCM_ARCH:-}" ]; then
+    _DETECTED_ARCH="$(_primus_detect_gpu_arch)"
+    if [ -n "$_DETECTED_ARCH" ]; then
+        export PYTORCH_ROCM_ARCH="$_DETECTED_ARCH"
+        echo "[env] detected GPU arch: PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH" >&2
+    else
+        export PYTORCH_ROCM_ARCH="gfx942;gfx950"
+        echo "[env] WARNING: no GPU detected; defaulting PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH (override by exporting it)" >&2
+    fi
+fi
+
+# Comma-separated variant for the tools/vars that expect commas.
+_ARCH_CSV="${PYTORCH_ROCM_ARCH//;/,}"
+export ROCM_AMDGPU_TARGETS="${ROCM_AMDGPU_TARGETS:-$_ARCH_CSV}"
+export GPU_ARCHS="${GPU_ARCHS:-$PYTORCH_ROCM_ARCH}"
+export HCC_AMDGPU_TARGET="${HCC_AMDGPU_TARGET:-$_ARCH_CSV}"
+export HIP_ARCHITECTURES="${HIP_ARCHITECTURES:-$_ARCH_CSV}"
+
+# Keep caches OFF the tiny home quota. Use local /tmp.
+export PIP_CACHE_DIR="${PIP_CACHE_DIR:-/tmp/primus-jax-cache/pip}"
+export XDG_CACHE_HOME="${XDG_CACHE_HOME:-/tmp/primus-jax-cache/xdg}"
+export UV_CACHE_DIR="${UV_CACHE_DIR:-/tmp/primus-jax-cache/uv}"
+mkdir -p "$PIP_CACHE_DIR" "$XDG_CACHE_HOME" "$UV_CACHE_DIR" 2>/dev/null || true
+
+# Activate the venv if it exists
+if [ -f "$VENV_DIR/bin/activate" ]; then
+    # shellcheck disable=SC1091
+    source "$VENV_DIR/bin/activate"
+fi
+
+# NOTE: the two HSA_* vars below are NOT in the v26.5 Dockerfile — they are a
+# carryover workaround for HSA_STATUS_ERROR_OUT_OF_RESOURCES. Harmless, but
+# remove them if you want an exact match to the image.
+export HSA_ENABLE_SCRATCH_ASYNC_RECLAIM=0
+export HSA_NO_SCRATCH_RECLAIM=1
+
+# Flags to fix the ROCm profiler hang issue (from the JAX Dockerfile base stage)
+export ROCPROFILER_QUEUE_INTERPOSITION=0
+export DEBUG_HIP_DYNAMIC_QUEUES=0
+
+# ---- ROCm path: from the extracted TheRock release tarball (v26.5) ----
+# setup.sh's `rocm` stage downloads therock-dist-*.tar.gz and extracts it into
+# $ROCM_DIR. Once present, wire the HIP/ROCm env at it (mirrors the Dockerfile's
+# /opt/rocm layout, whose LD path is lib + lib/rocm_sysdeps/lib).
+if [ -d "$ROCM_DIR/lib" ] || [ -x "$ROCM_DIR/bin/hipcc" ]; then
+    export ROCM_PATH="$ROCM_DIR"
+    export ROCM_HOME="$ROCM_DIR"   # Primus uses ROCM_HOME; set both
+    export HIP_PLATFORM=amd
+    export HIP_PATH="$ROCM_PATH"
+    export HIP_CLANG_PATH="$ROCM_PATH/llvm/bin"
+    export HIP_INCLUDE_PATH="$ROCM_PATH/include"
+    export HIP_LIB_PATH="$ROCM_PATH/lib"
+    export HIP_DEVICE_LIB_PATH="$ROCM_PATH/lib/llvm/amdgcn/bitcode"
+    # PATH: the Dockerfile puts /opt/rocm/lib on PATH too; mirror that.
+    export PATH="$ROCM_PATH/lib:$ROCM_PATH/bin:$HIP_CLANG_PATH:$PATH"
+    export LD_LIBRARY_PATH="$ROCM_PATH/lib:$ROCM_PATH/lib/rocm_sysdeps/lib:$ROCM_PATH/lib64:$ROCM_PATH/llvm/lib:${LD_LIBRARY_PATH:-}"
+    export LIBRARY_PATH="$ROCM_PATH/lib:$ROCM_PATH/lib64"
+    export CPATH="$HIP_INCLUDE_PATH"
+    export PKG_CONFIG_PATH="$ROCM_PATH/lib/pkgconfig"
+fi
+
+# ---- TransformerEngine (ROCm) runtime settings for JAX ----
+# The transformer_engine_rocm_jax wheel is prebuilt, so these are runtime knobs.
+export NVTE_ROCM_ARCH="$PYTORCH_ROCM_ARCH"
+export NVTE_USE_ROCM=1
+export NVTE_USE_HIPBLASLT=1
+export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1
+export NVTE_FUSED_ATTN=1
+export NVTE_CK_USES_BWD_V3=1
+export NVTE_CK_USES_FWD_V3=1
+export NVTE_CK_IS_V3_ATOMIC_FP32=1
+export NVTE_CK_HOW_V3_BF16_CVT=2
+
+# ---- AMD GPU runtime knobs (from the JAX Dockerfile) ----
+export GPU_MAX_HW_QUEUES=2
+export HIP_FORCE_DEV_KERNARG=1
+export HSA_FORCE_FINE_GRAIN_PCIE=1
+export NCCL_DEBUG=VERSION
+
+# NOT in the v26.5 Dockerfile: avoids a NaN-loss issue when training on gfx950
+# (no-op on gfx942). Keep it if you train on MI350/MI355; drop for exact parity.
+export RCCL_WARP_SPEED_AUTO=0
+
+# ---- XLA / JAX runtime settings ----
+export XLA_PYTHON_CLIENT_MEM_FRACTION=.9
+export XLA_FLAGS="${XLA_FLAGS:---xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --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 --xla_gpu_enable_command_buffer=''}"