From 3e86d04d6fcb5d80a7257cc1c9e063a83e8abd10 Mon Sep 17 00:00:00 2001 From: littlemex Date: Fri, 3 Jul 2026 09:47:07 +0000 Subject: [PATCH] fix(actor_group): select torch_memory_saver LD_PRELOAD lib by CUDA runtime, not filename With --offload-train on the megatron backend, actor_group.py picked the torch_memory_saver preload .so from a hard-coded [_cu12, unsuffixed] list by os.path.exists (first that exists). On a CUDA 13 image the cu12 build exists but links libcudart.so.12 (absent), so LD_PRELOAD makes every child -- incl. the train worker -- die with 'libcudart.so.12: cannot open shared object file'. The correct cu13 build ships in the same wheel but is never enumerated. Delegate .so selection to torch_memory_saver's own CUDA-aware resolver (get_binary_path_from_package) when present; a resolver that raises is propagated rather than masked. For older builds without the resolver, fall back keyed on the CUDA major torch was built for (torch.version.cuda) -- which is compile-time metadata needing no GPU/driver in this process. This matters because slime resolves on the Ray driver, which may be a CPU-only head node with no libcuda.so.1: selecting the cu build by existence there is correct (LD_PRELOAD takes effect on the GPU workers). Only when the major cannot be determined do we probe ctypes.CDLL loadability as a last resort, and if even that is impossible (no libcuda), fall back to the first existing candidate with a warning rather than failing a config the workers could run. --- slime/ray/actor_group.py | 94 +++++++++++++++++++---- tests/test_tms_preload_resolution.py | 108 +++++++++++++++++++++++++++ 2 files changed, 188 insertions(+), 14 deletions(-) create mode 100644 tests/test_tms_preload_resolution.py diff --git a/slime/ray/actor_group.py b/slime/ray/actor_group.py index 4237b80329..413c4b9724 100644 --- a/slime/ray/actor_group.py +++ b/slime/ray/actor_group.py @@ -10,6 +10,85 @@ from slime.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, add_default_ray_env_vars +def _resolve_tms_preload_lib(torch_memory_saver): + """Path to the torch_memory_saver preload .so matching the CUDA runtime. + + Selecting by filename existence is wrong on CUDA 13 (the cu12 .so exists but + links the absent libcudart.so.12, crashing every LD_PRELOAD child). Prefer + the library's own CUDA-aware resolver; otherwise key on the CUDA major torch + was built for (compile-time metadata, so it works on a GPU-less Ray driver), + and only dlopen-probe as a last resort since that needs the host libcuda. + """ + stem = "torch_memory_saver_hook_mode_preload" + + # Prefer the library resolver; only a *missing* one (older build) falls back. + try: + from torch_memory_saver.utils import get_binary_path_from_package + except ImportError: + get_binary_path_from_package = None + if get_binary_path_from_package is not None: + return str(get_binary_path_from_package(stem)) + + import ctypes + + base = os.path.dirname(os.path.dirname(torch_memory_saver.__file__)) + + def _first_existing(names): + for name in names: + path = os.path.join(base, name) + if os.path.exists(path): + return path + return None + + # Key on torch's build-time CUDA major (no GPU/driver needed here). + major = None + try: + import torch + + cuda = getattr(torch.version, "cuda", None) + major = cuda.split(".", 1)[0] if cuda else None + except ImportError: + pass + + if major is not None: + path = _first_existing([f"{stem}_cu{major}.abi3.so"]) + if path is not None: + return path + + # Major unknown: fall back to a dlopen loadability probe (needs libcuda). + candidates = [f"{stem}.abi3.so", f"{stem}_cu12.abi3.so", f"{stem}_cu13.abi3.so"] + missing, failed = [], [] + for name in candidates: + path = os.path.join(base, name) + if not os.path.exists(path): + missing.append(name) + continue + try: + ctypes.CDLL(path) + except OSError as e: + failed.append(f"{name}: {e}") + continue + return path + + # dlopen impossible too (e.g. GPU-less driver without libcuda): pick by + # existence with a warning rather than fail a config the workers could run. + fallback_path = _first_existing(candidates) + if fallback_path is not None: + import warnings + + warnings.warn( + "No torch_memory_saver preload lib dlopen'd here (missing libcuda.so.1?); " + f"using {os.path.basename(fallback_path)} by existence. " + f"dlopen failures: {failed or 'none'}.", + stacklevel=2, + ) + return fallback_path + + raise FileNotFoundError( + "Could not find a torch_memory_saver preload library under " f"{base}. Not found: {missing or 'none'}." + ) + + class RayTrainGroup: """ A group of ray actors @@ -73,20 +152,7 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): if self.args.offload_train and self.args.train_backend == "megatron": import torch_memory_saver - for path in [ - "torch_memory_saver_hook_mode_preload_cu12.abi3.so", - "torch_memory_saver_hook_mode_preload.abi3.so", - ]: - dynlib_path = os.path.join( - os.path.dirname(os.path.dirname(torch_memory_saver.__file__)), - path, - ) - if os.path.exists(dynlib_path): - break - else: - raise FileNotFoundError( - "Cannot find torch_memory_saver dynamic library. Please make sure torch_memory_saver is properly installed." - ) + dynlib_path = _resolve_tms_preload_lib(torch_memory_saver) env_vars["LD_PRELOAD"] = dynlib_path env_vars["TMS_INIT_ENABLE"] = "1" diff --git a/tests/test_tms_preload_resolution.py b/tests/test_tms_preload_resolution.py new file mode 100644 index 0000000000..0860fbd933 --- /dev/null +++ b/tests/test_tms_preload_resolution.py @@ -0,0 +1,108 @@ +"""Unit tests for actor_group._resolve_tms_preload_lib. + +Loads slime.ray.actor_group with heavy deps (ray) stubbed, then exercises the +torch_memory_saver preload .so selection: resolver-first, then a CUDA-major +keyed existence fallback, and finally a dlopen probe. Mirrors the module-stub +pattern in test_megatron_argument_validation.py. +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + + +def load_actor_group(monkeypatch): + ray_mod = types.ModuleType("ray") + pg_mod = types.ModuleType("ray.util.placement_group") + sched_mod = types.ModuleType("ray.util.scheduling_strategies") + pg_mod.PlacementGroup = object + sched_mod.PlacementGroupSchedulingStrategy = object + utils_mod = types.ModuleType("slime.ray.utils") + utils_mod.NOSET_VISIBLE_DEVICES_ENV_VARS_LIST = [] + utils_mod.add_default_ray_env_vars = lambda *a, **k: None + + monkeypatch.setitem(sys.modules, "ray", ray_mod) + monkeypatch.setitem(sys.modules, "ray.util.placement_group", pg_mod) + monkeypatch.setitem(sys.modules, "ray.util.scheduling_strategies", sched_mod) + monkeypatch.setitem(sys.modules, "slime.ray.utils", utils_mod) + + module_path = Path(__file__).resolve().parents[1] / "slime" / "ray" / "actor_group.py" + module_name = "test_actor_group_module" + sys.modules.pop(module_name, None) + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def _fake_tms(tmp_path, *, with_resolver, resolver_ret=None, resolver_raises=None): + """A fake torch_memory_saver module rooted so its .so dir is tmp_path.""" + pkg = tmp_path / "torch_memory_saver" + pkg.mkdir(exist_ok=True) + tms = types.ModuleType("torch_memory_saver") + tms.__file__ = str(pkg / "__init__.py") + utils = types.ModuleType("torch_memory_saver.utils") + if with_resolver: + + def get_binary_path_from_package(stem): + if resolver_raises is not None: + raise resolver_raises + return resolver_ret + + utils.get_binary_path_from_package = get_binary_path_from_package + return tms, utils + + +def _touch(tmp_path, name): + p = tmp_path / name + p.write_bytes(b"") + return str(p) + + +def test_prefers_library_resolver(monkeypatch, tmp_path): + ag = load_actor_group(monkeypatch) + want = str(tmp_path / "torch_memory_saver_hook_mode_preload_cu13.abi3.so") + tms, utils = _fake_tms(tmp_path, with_resolver=True, resolver_ret=want) + monkeypatch.setitem(sys.modules, "torch_memory_saver.utils", utils) + assert ag._resolve_tms_preload_lib(tms) == want + + +def test_resolver_error_propagates(monkeypatch, tmp_path): + # A resolver that exists but raises must NOT be masked by the fallback. + ag = load_actor_group(monkeypatch) + tms, utils = _fake_tms(tmp_path, with_resolver=True, resolver_raises=RuntimeError("boom")) + monkeypatch.setitem(sys.modules, "torch_memory_saver.utils", utils) + with pytest.raises(RuntimeError, match="boom"): + ag._resolve_tms_preload_lib(tms) + + +def test_fallback_keys_on_cuda_major_by_existence(monkeypatch, tmp_path): + # No resolver: pick cu by existence (no dlopen), so it works on a + # GPU-less driver. torch.version.cuda drives the major. + ag = load_actor_group(monkeypatch) + tms, _ = _fake_tms(tmp_path, with_resolver=False) + monkeypatch.delitem(sys.modules, "torch_memory_saver.utils", raising=False) + torch_mod = types.ModuleType("torch") + torch_mod.version = types.SimpleNamespace(cuda="13.0") + monkeypatch.setitem(sys.modules, "torch", torch_mod) + # Only the cu13 build exists; the fallback must select it by name. + _touch(tmp_path, "torch_memory_saver_hook_mode_preload_cu12.abi3.so") + want = _touch(tmp_path, "torch_memory_saver_hook_mode_preload_cu13.abi3.so") + assert ag._resolve_tms_preload_lib(tms) == want + + +def test_fallback_cuda12_selects_cu12(monkeypatch, tmp_path): + # CUDA 12 build: the cu12 variant is selected by existence, unchanged. + ag = load_actor_group(monkeypatch) + tms, _ = _fake_tms(tmp_path, with_resolver=False) + monkeypatch.delitem(sys.modules, "torch_memory_saver.utils", raising=False) + torch_mod = types.ModuleType("torch") + torch_mod.version = types.SimpleNamespace(cuda="12.4") + monkeypatch.setitem(sys.modules, "torch", torch_mod) + want = _touch(tmp_path, "torch_memory_saver_hook_mode_preload_cu12.abi3.so") + _touch(tmp_path, "torch_memory_saver_hook_mode_preload_cu13.abi3.so") + assert ag._resolve_tms_preload_lib(tms) == want