Persist trainer Triton/TorchInductor kernel caches in a Modal Volume

[micahtyong]

Summary

Lilo port of modal-projects/stitch#251. Every trainer definition runs with single_use_containers=True, so each trainer instance recompiled all Triton/Inductor kernels from scratch. This mounts one shared v2 Volume and points both compilers at it so later containers reuse earlier compiles.

New lilo/providers/modal/kernel_cache.py:

KERNEL_CACHE_VOLUME_NAME = "lilo-kernel-cache"
KERNEL_CACHE_ROOT = "/root/.cache/kernel-cache"
kernel_cache_volume = modal.Volume.from_name(KERNEL_CACHE_VOLUME_NAME, create_if_missing=True, version=2)
KERNEL_CACHE_ENV = {
    "TRITON_CACHE_DIR": f"{KERNEL_CACHE_ROOT}/triton",
    "TORCHINDUCTOR_CACHE_DIR": f"{KERNEL_CACHE_ROOT}/inductor",
}

Wired into all 13 trainer definitions (the _single variant inherits via the shared module) and the scoped trainer:

 TRAINER_VOLUMES = {
     ...,
     CHECKPOINT_ROOT: checkpoint_volume,
+    KERNEL_CACHE_ROOT: kernel_cache_volume,
 }
 ...
         backend_env={
+            **KERNEL_CACHE_ENV,      # first, so per-definition entries can override
             "LILO_BACKEND_CONFIG": ...,

run_engine_with_backend merges backend_env into the backend subprocess env, and a locally started Ray cluster inherits it. For the pre-existing-cluster path (256k multi-node), miles_runtime._WORKER_ENV_VARS now also forwards TRITON_CACHE_DIR/TORCHINDUCTOR_CACHE_DIR via runtime_env so actors on worker nodes write to the same volume.

No explicit commit(); v2 Volumes flush in the background and on container shutdown, same as #251.

Measurement

qwen3_8_27b_miles_lora_64k (H200:8, TP4/CP2, LoRA rank 32) in lilo-deploy, ephemeral app built from this branch. Cold = lilo-kernel-cache wiped (triton/ and inductor/ deleted) before the run; warm = a fresh trainer container ~10 min later against the populated volume. Each run: fwd/bwd+optim at 128 tokens, then 65,536 tokens, then both again.

phase	cold	warm	saved
first fwd/bwd, 128 tok (compile)	209.2 s	140.2 s	69 s
first fwd/bwd, 65,536 tok (compile)	110.5 s	88.7 s	22 s
first-step compile total	319.8 s	228.9 s	91 s (-28%)
steady fwd/bwd, 128 tok	3.0 s	1.8 s	—
steady fwd/bwd, 65,536 tok	11.1 s	7.6 s	—

Container start → ready for first step (Ray boot, actor init, HF→Megatron weight load; ~3–4 min) is not affected by this change — no Triton/Inductor compile happens in that phase.

Cache contents after the cold run: 574 Triton kernel dirs, Inductor fxgraph/ (231) + aotautograd/ (232) entries. The Triton dir count did not change after the warm run, i.e. every Triton kernel was a cache hit and no new kernels were compiled.

Takeaway: ~1.5 min saved per trainer start on this definition. The remaining ~140 s of the warm first step is not Triton/Inductor compile — it is dynamo tracing, autotune benchmarking, CUDA-graph capture and NCCL/TE warm-up that run every time regardless of cache. The 23 min Stitch saw on the 35B-A3B NVFP4 recipe does not transfer to this dense LoRA recipe. Also, trainer provisioning time was dominated by H200:8 queue wait (13 min on the cold run, ~2.5 min on the warm run) and is unaffected by this change.

Link to Devin session: https://modal.devinenterprise.com/sessions/21fe6974a320455ead92283cec40e59a
Open in Devin Desktop: https://modal.devinenterprise.com/desktop/session/21fe6974a320455ead92283cec40e59a?variant=devin
Requested by: @micahtyong

[devin-ai-integration]

I'll fix CI failures and address comments from users with write access that start with 'DevinAI' or '@devin'.

• Disable automatic comment, CI, and merge conflict monitoring

[micahtyong]

/devin review

[devin-ai-integration]

Starting Devin Review.

[devin-ai-integration]

✅ Devin Review: No Issues Found

Devin Review analyzed this PR and found no bugs or issues to report.

[kailash109]

Nice!! wondering if there's any more caching we can do to bring down the cold start times even more

ref: https://research.nvidia.com/labs/nemotron/files/NVIDIA-Nemotron-3-Ultra-Technical-Report.pdf table 9

in particular, looks like SGLANG_CACHE_DIR is not set, and setting FLASHINFER_WORKSPACE_BASE might help too