feat(miles): allow tms preload on Blackwell with cu13+ torch wheels

examples/rlix/run_miles_dual.py

@@ -289,6 +289,7 @@ def _build_pipeline(
         pipeline_runtime_env_vars["PYTHONPATH"] = pythonpath
     for _k in (
         "MILES_TMS_HOOK_MODE",
+        "MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL",
         "MILES_MAX_RESIDUAL_GPU_MEM_GB",
         "MILES_SKIP_TMS_PAUSE",
         "MILES_SKIP_NODE_PG_PIN",

examples/rlix/run_miles_rlix.py

@@ -178,6 +178,7 @@ class MilesPipelineConfig:
     # the parent driver's env by default).
     for _k in (
         "MILES_TMS_HOOK_MODE",
+        "MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL",
         "MILES_MAX_RESIDUAL_GPU_MEM_GB",
         "MILES_SKIP_TMS_PAUSE",
         "MILES_SKIP_NODE_PG_PIN",

miles/backends/megatron_utils/actor.py

@@ -98,8 +98,10 @@ def init(
             # MILES_TMS_HOOK_MODE=torch switches torch_memory_saver into
             # PyTorch's CUDAPluggableAllocator path, avoiding the
             # LD_PRELOAD libc malloc hook that segfaults during
-            # build_cpu_bucket_cache on CUDA 12.9 / Blackwell. Must be
-            # set BEFORE any tms call that triggers _ensure_initialized.
+            # build_cpu_bucket_cache on Blackwell with pre-CUDA-13 wheels.
+            # The guard below is CUDA-version-aware: preload is allowed on
+            # cu13+ Blackwell. Must be set BEFORE any tms call that
+            # triggers _ensure_initialized.
             import os as _os
 
             mode = _os.environ.get("MILES_TMS_HOOK_MODE")

miles/backends/megatron_utils/tms_utils.py

@@ -4,8 +4,9 @@
 
 - ``preload`` (its default): an ``LD_PRELOAD`` libc-malloc interposer. Broadest
   catchment (every allocation, incl. NCCL / raw ``cudaMalloc``), but segfaults on
-  Blackwell-class GPUs under CUDA 12.9 — a raw SIGSEGV with no Python traceback,
-  typically on the first allocation (e.g. ``build_cpu_bucket_cache``).
+  Blackwell-class GPUs under pre-CUDA-13 wheels (observed on 12.9) — a raw
+  SIGSEGV with no Python traceback, typically on the first allocation (e.g.
+  ``build_cpu_bucket_cache``). Verified fixed on cu130 wheels.
 - ``torch``: PyTorch's ``CUDAPluggableAllocator``. Narrower (only torch
   allocations), but stable across architectures.
 
@@ -20,14 +21,21 @@
 import os
 
 import torch
+from packaging.version import parse
 
 logger = logging.getLogger(__name__)
 
 # Blackwell-class GPUs (datacenter B100/B200 == sm_100, consumer RTX 50xx ==
 # sm_120) have compute-capability major >= 10. ``preload`` segfaults there under
-# CUDA 12.9; ``torch`` mode is required.
+# pre-CUDA-13 wheels (observed on 12.9; verified fixed on cu130); ``torch`` mode
+# is required on those older stacks.
 TMS_PRELOAD_UNSAFE_CC_MAJOR = 10
 
+# preload verified safe on Blackwell from CUDA 13 wheels (RTX 5090 +
+# torch 2.11.0+cu130, 2026-07-05 audit: mock pause/resume + dual E2E both
+# pass); cu12.x wheels keep the historical tms 0.0.9 segfault.
+TMS_PRELOAD_SAFE_CUDA_MAJOR = 13
+
 # Escape hatch: set to "1" to proceed with preload on Blackwell anyway (e.g. once
 # a fixed tms/CUDA build is confirmed). Mirrors the repo's MILES_SKIP_* knobs.
 TMS_ALLOW_PRELOAD_ON_BLACKWELL_ENV = "MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL"
@@ -42,14 +50,32 @@ def resolve_tms_hook_mode(env_mode: str | None) -> str:
     return env_mode if env_mode in ("torch", "preload") else "preload"
 
 
+def _torch_cuda_major() -> int | None:
+    """Major version of the CUDA runtime the torch wheel was built against.
+
+    ``torch.version.cuda`` is ``None`` on CPU-only builds; treat that (or an
+    unparseable value) as "unknown" and return ``None`` so callers stay
+    conservative.
+    """
+    cuda_version = getattr(torch.version, "cuda", None)
+    if cuda_version is None:
+        return None
+    try:
+        return parse(str(cuda_version)).major
+    except Exception:  # unparseable (e.g. vendor-patched string) -> unknown
+        return None
+
+
 def assert_tms_hook_mode_matches_arch(env_mode: str | None) -> None:
     """Fail fast when the resolved tms hook mode will crash on this GPU.
 
     On Blackwell-class GPUs the ``"preload"`` hook segfaults inside libc on the
-    first allocation under CUDA 12.9. Because the crash is a tracebackless
-    SIGSEGV, we refuse to proceed and tell the operator exactly which knob to
-    set. Set ``MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL=1`` to bypass once a fixed
-    tms/CUDA combination is confirmed.
+    first allocation under pre-CUDA-13 wheels (observed on 12.9; verified fixed
+    on cu130). Because the crash is a tracebackless SIGSEGV, we refuse to
+    proceed on cu12.x/unknown stacks and tell the operator exactly which knob
+    to set. Torch wheels built against CUDA >= 13 are allowed through. Set
+    ``MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL=1`` to bypass on an older stack once
+    a fixed tms/CUDA combination is confirmed.
 
     Raises ``RuntimeError`` (not bare ``assert``) so it remains active under
     ``python -O``.
@@ -61,23 +87,39 @@ def assert_tms_hook_mode_matches_arch(env_mode: str | None) -> None:
     major, minor = torch.cuda.get_device_capability()
     if major < TMS_PRELOAD_UNSAFE_CC_MAJOR:
         return
+    cuda_major = _torch_cuda_major()
+    if cuda_major is not None and cuda_major >= TMS_PRELOAD_SAFE_CUDA_MAJOR:
+        logger.info(
+            "torch_memory_saver hook_mode 'preload' allowed on Blackwell (sm_%d%d) "
+            "because torch wheel CUDA %s >= %d (segfault verified fixed on cu130).",
+            major,
+            minor,
+            torch.version.cuda,
+            TMS_PRELOAD_SAFE_CUDA_MAJOR,
+        )
+        return
     if os.environ.get(TMS_ALLOW_PRELOAD_ON_BLACKWELL_ENV) == "1":
         logger.warning(
             "torch_memory_saver hook_mode 'preload' on Blackwell (sm_%d%d) is "
-            "known to segfault under CUDA 12.9; proceeding anyway because %s=1.",
+            "known to segfault under pre-CUDA-13 wheels (torch wheel CUDA: %s); "
+            "proceeding anyway because %s=1.",
             major,
             minor,
+            getattr(torch.version, "cuda", None),
             TMS_ALLOW_PRELOAD_ON_BLACKWELL_ENV,
         )
         return
     raise RuntimeError(
         f"torch_memory_saver hook_mode resolved to 'preload' on a Blackwell-class "
-        f"GPU ({torch.cuda.get_device_name()}, compute capability sm_{major}{minor}). "
-        f"'preload' uses an LD_PRELOAD libc-malloc hook that segfaults on this "
-        f"architecture under CUDA 12.9 — a raw SIGSEGV with no Python traceback, "
+        f"GPU ({torch.cuda.get_device_name()}, compute capability sm_{major}{minor}) "
+        f"with a pre-CUDA-13 torch wheel (torch.version.cuda="
+        f"{getattr(torch.version, 'cuda', None)!r}). "
+        f"'preload' uses an LD_PRELOAD libc-malloc hook that segfaults on Blackwell "
+        f"with pre-CUDA-13 wheels — a raw SIGSEGV with no Python traceback, "
         f"typically during build_cpu_bucket_cache. "
-        f"Fix: export MILES_TMS_HOOK_MODE=torch before launch "
-        f"(MILES_TMS_HOOK_MODE was {'unset' if env_mode is None else repr(env_mode)}). "
-        f"To override after confirming a fixed tms/CUDA build, set "
+        f"Fixes: (a) export MILES_TMS_HOOK_MODE=torch before launch "
+        f"(MILES_TMS_HOOK_MODE was {'unset' if env_mode is None else repr(env_mode)}); "
+        f"(b) upgrade to a cu13+ torch wheel (preload is verified safe there); or "
+        f"(c) after confirming a fixed tms/CUDA build, set "
         f"{TMS_ALLOW_PRELOAD_ON_BLACKWELL_ENV}=1."
     )

miles/ray/actor_group.py

@@ -125,9 +125,15 @@ def _allocate_gpus_via_placements(self, worker_placements, num_gpus_per_actor) -
             # Per-actor switch into torch_memory_saver "torch" hook mode
             # (CUDAPluggableAllocator) which avoids the LD_PRELOAD libc
             # malloc hook that segfaults during build_cpu_bucket_cache on
-            # CUDA 12.9 / Blackwell. The actor reads this env at init.
+            # Blackwell with pre-CUDA-13 wheels (the guard is CUDA-version
+            # aware: preload is allowed on cu13+ Blackwell). The actor
+            # reads this env at init.
             if (mode := _os.environ.get("MILES_TMS_HOOK_MODE")):
                 env_vars_base["MILES_TMS_HOOK_MODE"] = mode
+            # The guard's escape hatch is read inside the actor process;
+            # forward it so it works under Ray runtime_env isolation.
+            if (allow := _os.environ.get("MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL")):
+                env_vars_base["MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL"] = allow
 
         backend = self.args.train_backend
         if backend == "megatron":
@@ -219,6 +225,12 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor):
             env_vars["LD_PRELOAD"] = dynlib_path
             env_vars["TMS_INIT_ENABLE"] = "1"
             env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1"
+            # Forward MILES_TMS_HOOK_MODE for consistency with the
+            # placement path above (_allocate_gpus_via_placements).
+            if (mode := os.environ.get("MILES_TMS_HOOK_MODE")):
+                env_vars["MILES_TMS_HOOK_MODE"] = mode
+            if (allow := os.environ.get("MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL")):
+                env_vars["MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL"] = allow
 
         backend = self.args.train_backend
         if backend == "megatron":

tests/fast/backends/test_tms_utils.py

@@ -27,27 +27,50 @@ def test_resolve_hook_mode(env_mode, expected):
     assert resolve_tms_hook_mode(env_mode) == expected
 
 
-def _patch_gpu(monkeypatch, *, available=True, cc=(8, 9), name="NVIDIA L4"):
+def _patch_gpu(monkeypatch, *, available=True, cc=(8, 9), name="NVIDIA L4", cuda_version="12.9"):
+    # cuda_version default "12.9" preserves the historical (pre-cu13-guard) semantics
+    # and keeps the raise-tests deterministic on any host.
     monkeypatch.setattr(torch.cuda, "is_available", lambda: available)
     monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: cc)
     monkeypatch.setattr(torch.cuda, "get_device_name", lambda *a, **k: name)
+    monkeypatch.setattr(torch.version, "cuda", cuda_version, raising=False)
 
 
 # (major, minor) for Blackwell-class parts: B100/B200 == sm_100, RTX 50xx == sm_120
 @pytest.mark.parametrize("cc", [(10, 0), (12, 0)])
-# all of these resolve to preload, which is the unsafe mode on Blackwell
+# all of these resolve to preload, which is the unsafe mode on Blackwell + cu12.x
 @pytest.mark.parametrize("env_mode", [None, "preload", "bogus"])
 def test_raises_on_preload_blackwell(monkeypatch, cc, env_mode):
     monkeypatch.delenv("MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL", raising=False)
-    _patch_gpu(monkeypatch, cc=cc, name="NVIDIA B200")
+    _patch_gpu(monkeypatch, cc=cc, name="NVIDIA B200", cuda_version="12.9")
     with pytest.raises(RuntimeError, match="MILES_TMS_HOOK_MODE=torch"):
         assert_tms_hook_mode_matches_arch(env_mode)
 
 
+# preload verified safe on Blackwell from cu13 wheels (RTX 5090 + torch cu130)
+@pytest.mark.parametrize("cc", [(10, 0), (12, 0)])
+@pytest.mark.parametrize("cuda_version", ["13.0", "13.1"])
+# None resolves to preload too — the guard must allow both spellings on cu13+
+@pytest.mark.parametrize("env_mode", [None, "preload"])
+def test_no_raise_preload_blackwell_cu13(monkeypatch, cc, cuda_version, env_mode):
+    monkeypatch.delenv("MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL", raising=False)
+    _patch_gpu(monkeypatch, cc=cc, name="NVIDIA B200", cuda_version=cuda_version)
+    assert_tms_hook_mode_matches_arch(env_mode)  # cu13+ wheel: must not raise
+
+
+# unknown wheel CUDA (CPU build or unparseable string): stay conservative -> raise
+@pytest.mark.parametrize("cuda_version", [None, "not-a-version"])
+def test_raises_on_preload_blackwell_unknown_cuda(monkeypatch, cuda_version):
+    monkeypatch.delenv("MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL", raising=False)
+    _patch_gpu(monkeypatch, cc=(12, 0), name="NVIDIA RTX 5090", cuda_version=cuda_version)
+    with pytest.raises(RuntimeError, match="MILES_TMS_HOOK_MODE=torch"):
+        assert_tms_hook_mode_matches_arch("preload")
+
+
 # V100 (sm_70), A100 (sm_80), L4/Ada (sm_89), H100 (sm_90): preload is fine
 @pytest.mark.parametrize("cc", [(7, 0), (8, 0), (8, 9), (9, 0)])
 def test_no_raise_preload_pre_blackwell(monkeypatch, cc):
-    _patch_gpu(monkeypatch, cc=cc)
+    _patch_gpu(monkeypatch, cc=cc, cuda_version="12.9")
     assert_tms_hook_mode_matches_arch("preload")  # must not raise
     assert_tms_hook_mode_matches_arch(None)  # unset -> preload, still pre-Blackwell
 
@@ -59,8 +82,9 @@ def test_torch_mode_always_safe(monkeypatch, cc):
 
 
 def test_escape_hatch_allows_preload_on_blackwell(monkeypatch):
+    # escape hatch still matters on cu12.x/unknown stacks
     monkeypatch.setenv("MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL", "1")
-    _patch_gpu(monkeypatch, cc=(12, 0), name="NVIDIA RTX 5090")
+    _patch_gpu(monkeypatch, cc=(12, 0), name="NVIDIA RTX 5090", cuda_version="12.9")
     assert_tms_hook_mode_matches_arch("preload")  # bypassed -> no raise
 
 

tests/fast/utils/test_rlix_env_forwarding.py

@@ -0,0 +1,82 @@
+"""Source-level checks that operator env knobs survive the Ray runtime_env boundary.
+
+The tms preload-on-Blackwell guard relaxation (cu13+ wheels allowed) added two
+new knobs — MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL and MILES_MAX_RESIDUAL_GPU_MEM_GB —
+that only work if the rlix example drivers forward them into each pipeline's
+runtime_env, and the legacy actor_group path forwards MILES_TMS_HOOK_MODE like
+the placement path does (codex review requirement on the guard-relaxation PR).
+
+These tests read/AST-parse the source files instead of importing them: the
+example drivers pull in ray/miles heavyweight deps that fast tests must avoid.
+"""
+
+import ast
+from pathlib import Path
+
+# tests/fast/utils/test_rlix_env_forwarding.py -> parents[3] == repo root
+REPO_ROOT = Path(__file__).resolve().parents[3]
+
+RUN_MILES_DUAL = REPO_ROOT / "examples" / "rlix" / "run_miles_dual.py"
+RUN_MILES_RLIX = REPO_ROOT / "examples" / "rlix" / "run_miles_rlix.py"
+ACTOR_GROUP = REPO_ROOT / "miles" / "ray" / "actor_group.py"
+
+REQUIRED_FORWARDED_KEYS = (
+    "MILES_TMS_HOOK_MODE",
+    "MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL",
+    "MILES_MAX_RESIDUAL_GPU_MEM_GB",
+)
+
+
+def _forwarding_tuples(path: Path) -> list[set[str]]:
+    """All-string-constant tuples in ``path`` that mention MILES_TMS_HOOK_MODE."""
+    tree = ast.parse(path.read_text(), filename=str(path))
+    tuples = []
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Tuple):
+            continue
+        values = {el.value for el in node.elts if isinstance(el, ast.Constant) and isinstance(el.value, str)}
+        if len(values) == len(node.elts) and "MILES_TMS_HOOK_MODE" in values:
+            tuples.append(values)
+    return tuples
+
+
+def _function_sources(path: Path) -> dict[str, str]:
+    """Map function/method name -> source segment for every def in ``path``."""
+    source = path.read_text()
+    tree = ast.parse(source, filename=str(path))
+    return {
+        node.name: ast.get_source_segment(source, node)
+        for node in ast.walk(tree)
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
+    }
+
+
+def test_run_miles_dual_forwards_tms_knobs():
+    tuples = _forwarding_tuples(RUN_MILES_DUAL)
+    assert tuples, f"no env-forwarding tuple found in {RUN_MILES_DUAL}"
+    for keys in tuples:
+        for required in REQUIRED_FORWARDED_KEYS:
+            assert required in keys, f"{required} missing from forwarding tuple in {RUN_MILES_DUAL}"
+
+
+def test_run_miles_rlix_forwards_tms_knobs():
+    tuples = _forwarding_tuples(RUN_MILES_RLIX)
+    assert tuples, f"no env-forwarding tuple found in {RUN_MILES_RLIX}"
+    for keys in tuples:
+        for required in REQUIRED_FORWARDED_KEYS:
+            assert required in keys, f"{required} missing from forwarding tuple in {RUN_MILES_RLIX}"
+
+
+def test_actor_group_forwards_tms_knobs_on_both_paths():
+    # Both LD_PRELOAD injection blocks (placement + legacy) must forward
+    # MILES_TMS_HOOK_MODE (or actors on the legacy path silently fall back to
+    # preload and trip the Blackwell guard) AND the guard's escape hatch
+    # MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL (read inside the actor process, so
+    # it is inert under Ray runtime_env isolation unless forwarded).
+    functions = _function_sources(ACTOR_GROUP)
+    for func_name in ("_allocate_gpus_via_placements", "_allocate_gpus_for_actor"):
+        assert func_name in functions, f"{func_name} not found in {ACTOR_GROUP}"
+        for env_key in ("MILES_TMS_HOOK_MODE", "MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL"):
+            assert (
+                env_key in functions[func_name]
+            ), f"{func_name} in {ACTOR_GROUP} does not forward {env_key}"