diff --git a/examples/rlix/run_miles_dual.py b/examples/rlix/run_miles_dual.py index f637de0944..4a3b802e4d 100644 --- a/examples/rlix/run_miles_dual.py +++ b/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", diff --git a/examples/rlix/run_miles_rlix.py b/examples/rlix/run_miles_rlix.py index 681c11f24f..9fce1e9118 100644 --- a/examples/rlix/run_miles_rlix.py +++ b/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", diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 93e6e330f5..58636ec301 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/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") diff --git a/miles/backends/megatron_utils/tms_utils.py b/miles/backends/megatron_utils/tms_utils.py index 84c0703de5..872878d4b4 100644 --- a/miles/backends/megatron_utils/tms_utils.py +++ b/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." ) diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index c262ea308d..7b7e654033 100644 --- a/miles/ray/actor_group.py +++ b/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": diff --git a/tests/fast/backends/test_tms_utils.py b/tests/fast/backends/test_tms_utils.py index 3993744024..e38e94aa74 100644 --- a/tests/fast/backends/test_tms_utils.py +++ b/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 diff --git a/tests/fast/utils/test_rlix_env_forwarding.py b/tests/fast/utils/test_rlix_env_forwarding.py new file mode 100644 index 0000000000..c17742bff5 --- /dev/null +++ b/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}"