Skip to content

fix(launch): default TRITON_CACHE_DIR to a node-local path - #3877

Merged
faresobeid merged 3 commits into
mainfrom
fix/triton-cache-node-local
Oct 8, 2026
Merged

faresobeid merged 3 commits into
mainfrom
fix/triton-cache-node-local

Conversation

@faresobeid

@faresobeid faresobeid commented Oct 5, 2026 •

Copy link
Copy Markdown
Contributor

Processes that share ~/.triton on a shared filesystem can hang or crash (seen with Nemotron-H, where many ranks compile the Mamba Triton kernels at once; reproduced below with Qwen3.5 and GLM-4.5-Air).

This PR adds TRITON_CACHE_DIR = /tmp/triton-<uid> to DEFAULT_COMMON_ENV_VARS. It is per-user because /tmp is shared by all users on a node. Like the other launcher defaults, it reaches every launched component: the local rl/sft/inference launchers (also used by the single-node sbatch scripts) and the trainer, orchestrator, inference (incl. prefill/decode), and online-eval env blocks of the multi-node and inference sbatch scripts. [env_vars] (or a component's env_vars) overrides it; a value exported in the shell or .env does not. It also removes the TRITON_CACHE_DIR overrides from six example configs (glm-4.5-air swe/swe-budget/search/terminal, glm-5.3 and qwen3-30b-a3b sft h200 base): they pointed every user at the same /tmp dir and are now covered by the default. docs/scaling.md says so in one paragraph.

vLLM 0.31's own default: when TRITON_CACHE_DIR is unset, its Inductor adaptor points it at $VLLM_CACHE_ROOT/torch_compile_cache/<hash>/rank_<i>_<j>/triton_cache (~/.cache/vllm by default, so usually on the shared home). That only covers Triton kernels compiled after torch.compile initializes, and it is still on the shared filesystem. vLLM keeps a TRITON_CACHE_DIR that is already set, so with this PR vLLM uses /tmp/triton-<uid> too.

Measured cost (one H200 node, vLLM 0.31 venv, Triton 3.7.1)

Each run starts from a fresh process with the same config. "Ready" is launch → first trainer step logged (Qwen) or launch → /health OK (GLM). For GLM, the first generated token came 1.3–2 s after ready in every run that came up. Inductor's cache dir was emptied before each run except the warm node-local one.

Triton cache Qwen3.5-4B SFT, 8 GPUs, compiled blocks (step 1 time) GLM-4.5-Air vLLM TP8
shared, cold (today on a new cluster / after a Triton or torch update) failed 2/2: FileNotFoundError on a cache file that another rank was writing (causal_conv1d_fwd_kernel.llir, l2norm_fwd_kernel.ptx) failed 2/2: 4 of 8 ranks hit FileNotFoundError in the cache during Inductor autotune, the other 4 hung in a collective until killed (1st run). The 2nd run crashed instead of hanging.
shared, warm not measured cleanly (see below) 344 s
node-local, cold (this PR, first job on a node) 233 s (step 1: 91 s), 164 kernels compiled 375 s with vLLM's compile cache warm (106 kernels); 465 s if vLLM's cache is also cold (87 kernels)
node-local, warm (this PR, later jobs on the same node) 138 s (step 1: 8.3 s) 344 s
  • A cold node-local cache costs about 95 s per node for Qwen3.5 (FLA + Inductor kernels) and about 31 s for GLM-4.5-Air inference when vLLM's own compile cache is warm. Warm node-local and warm shared take the same time.
  • The rest of the GLM startup does not depend on the Triton cache: about 3.5 min of Python imports and process spawn from the shared venv, 13 s to load weights.
  • vLLM: VLLM_CACHE_ROOT (~/.cache/vllm, shared home) holds vLLM's AOT/Inductor artifacts. With this PR it stays shared. On a new node, vLLM loads those artifacts (torch.compile 0.6 s instead of 30 s), and the Triton binaries are compiled into /tmp/triton-<uid> (106 kernels, +31 s). TORCHINDUCTOR_CACHE_DIR (default /tmp/torchinductor_<user>, already node-local) stayed empty in these runs. Inductor's Triton kernels go to TRITON_CACHE_DIR.
  • A crash on the shared cache can also break later starts. After the first failed GLM run, the shared vLLM cache only had AOT artifacts from the 4 ranks that got through. On the next start, those 4 ranks loaded the artifacts and the other 4 recompiled, so the ranks diverged and hung (Triton cache warm-ish, 60 kernels). For the same reason the Qwen "shared, warm" run started from a partial cache and failed again. The warm-shared GLM number comes from a shared dir seeded with a complete cache.

Possible middle ground (not implemented)

Because the cold cost is above 1 min per node for FLA-heavy models, a run could seed the node-local cache from a shared read-only copy at job start. For example, on each node before the components start:
[ -n "$TRITON_CACHE_SEED" ] && cp -rn "$TRITON_CACHE_SEED"/. "$TRITON_CACHE_DIR"/
The seed would be a warm /tmp/triton-<uid> copied once to the shared FS. Triton keys entries by source, compiler, and GPU, so stale entries are never reused, only skipped. Only reads hit the shared FS, so the concurrent-write race does not occur. I left it out to keep this PR minimal.


Note

Low Risk
Launch-time env default only; behavior is overridable via config and fixes known multi-rank Triton cache races, with a one-time cold-compile cost on each node.

Overview
Sets a global launcher default so every trainer, orchestrator, and inference process gets TRITON_CACHE_DIR=/tmp/triton-<uid> via DEFAULT_COMMON_ENV_VARS, avoiding concurrent Triton kernel writes under ~/.triton on a shared filesystem (hangs / FileNotFoundError under multi-rank compile).

Removes redundant [env_vars] overrides from six GLM-4.5-Air RL and GLM-5.3 / Qwen3-30B SFT example configs that previously pinned Triton cache paths under /tmp; those runs now inherit the per-user node-local default (still overridable through config env_vars).

Reviewed by Cursor Bugbot for commit f216420. Bugbot is set up for automated code reviews on this repo. Configure here.

@faresobeid
faresobeid force-pushed the fix/triton-cache-node-local branch from 70e3001 to ff9d3ea Compare October 6, 2026 17:19
@faresobeid
faresobeid marked this pull request as ready for review October 7, 2026 13:34
Processes sharing ~/.triton on a shared FS hang or crash. Add TRITON_CACHE_DIR=/tmp/triton-<uid> to DEFAULT_COMMON_ENV_VARS so every launched component gets it; env_vars in the config override it.
@faresobeid
faresobeid force-pushed the fix/triton-cache-node-local branch from ff9d3ea to f435bc5 Compare October 7, 2026 19:23
Comment thread docs/scaling.md Resolved
Signed-off-by: faresobeid <111092724+faresobeid@users.noreply.github.com>
@faresobeid
faresobeid merged commit 765bbeb into main Oct 8, 2026
19 of 20 checks passed
@faresobeid
faresobeid deleted the fix/triton-cache-node-local branch October 8, 2026 10:59
faresobeid added a commit that referenced this pull request Oct 11, 2026
## What
- `VLLM_CACHE_ROOT=/tmp/.vllm-cache-<uid>` and
`FLASHINFER_WORKSPACE_BASE=/tmp/.flashinfer-cache-<uid>` are added to
`DEFAULT_INFERENCE_ENV_VARS`. `TRITON_CACHE_DIR` got the same treatment
in #3877.
- The by-hand copies (`/tmp/.vllm-cache`, `/tmp/.flashinfer-cache`) are
removed from the four GLM-4.5-Air recipes. Those paths were shared by
every user on a node.
- docs/scaling.md gets one paragraph that lists the node-local cache
defaults.

The paths are dot-prefixed so the SLURM start-up cleanup (`rm -rf
/dev/shm/vllm-* /dev/shm/vllm_* /tmp/vllm-* /tmp/vllm_* /tmp/torch-*
/tmp/torchelastic_*` in the multi_node_rl, multi_node_sft and inference
sbatch templates) does not wipe them. These are the same names the
recipes used before.

Precedence is the same as for the other launcher defaults: `[env_vars]`
or `[inference.env_vars]` override them; a value exported in the shell
does not.

## Why
- Without these, vLLM's AOT/Inductor compile cache lives in
`~/.cache/vllm` and FlashInfer's JIT cache under `$HOME`, which is
usually the shared filesystem.
- Measured while testing #3877 (GLM-4.5-Air TP8, vLLM 0.31): a crashed
start left a partial shared vLLM cache. On the next start, 4 TP ranks
loaded those artifacts and the other 4 recompiled. The ranks diverged
and hung in a collective. So one crash breaks every later start that
uses that cache.
- The original scaleswe recipe (#2862) moved this cache to `/tmp`
because the NFS cache hung inference startup.
- Cost: the first start on each new node has a cold vLLM cache.
GLM-4.5-Air: 465 s cold vs 375 s with a warm vLLM cache, so about 1.5
min once per node.

## Overlap
- #3627 (garrett361, draft) adds
`VLLM_CACHE_ROOT`/`FLASHINFER_WORKSPACE_BASE` under `/tmp/garrett/` to
new DeepSeek-V4 recipes. With this PR those lines can go.

## Test plan
- `pytest tests/unit/orchestrator tests/unit/test_configs.py
tests/unit/utils -m "not gpu"`: 365 passed (rebased on origin/main
41a8ed0).
- Ran the cleanup globs in a scratch dir that held `.vllm-cache-<uid>`,
`.flashinfer-cache-<uid>` and, as a control, `vllm-cache-<uid>`. Only
the control matched.
- `ruff check` and `ruff format --check` on `process.py`: clean.
- `rl @ examples/advanced/glm-4.5-air/swe.toml --dry-run`:
- The rendered `rl.sbatch` exports both vars only in the inference block
(`/tmp/.vllm-cache-<uid>`, `/tmp/.flashinfer-cache-<uid>`).
- With `--inference.env-vars '{"VLLM_CACHE_ROOT": "/scratch/vc"}'`, the
sbatch exports `/scratch/vc`.

<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> **Low Risk**
> Launcher default env vars for inference only; configs can still
override, with no auth or training logic changes.
> 
> **Overview**
> Inference launches now default **vLLM** and **FlashInfer** compile/JIT
caches to per-user paths under `/tmp` (`VLLM_CACHE_ROOT`,
`FLASHINFER_WORKSPACE_BASE`), matching the existing Triton cache
behavior and avoiding shared `$HOME` or NFS caches that can leave
tensor-parallel ranks on mismatched artifacts and hang.
> 
> The four GLM-4.5-Air example recipes drop redundant
`[inference.env_vars]` blocks that duplicated those settings.
**docs/scaling.md** documents the launcher defaults and notes overrides
via `[env_vars]` or component `env_vars`.
> 
> <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit
fe8a34d. Bugbot is set up for automated
code reviews on this repo. Configure
[here](https://www.cursor.com/dashboard/bugbot).</sup>
<!-- /CURSOR_SUMMARY -->
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants