fix(actor_group): select the torch_memory_saver LD_PRELOAD lib by CUDA runtime, not by filename (fixes libcudart.so.12 crash on CUDA 13)

[littlemex]

What / Why

With --offload-train on the Megatron backend, actor_group.py chooses the torch_memory_saver preload .so from a hard-coded list (_cu12, then unsuffixed) and picks the first that exists on disk. On a CUDA 13 image the cu12 build exists but links libcudart.so.12 (absent on CUDA 13), so LD_PRELOAD makes every child — including 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 considered.

This PR stops choosing by filename existence and instead uses the CUDA runtime to pick the right build, by delegating to torch_memory_saver's own resolver — the mechanism the library already provides for exactly this.

Root cause (concise)

• slime/ray/actor_group.py L64-L84: candidate list is ["..._preload_cu12.abi3.so", "..._preload.abi3.so"], selected via os.path.exists. No cu13 candidate; existence ≠ loadability.
• On CUDA 13 (torch 2.11.0+cu130, torch_memory_saver 0.0.9.post1): both hard-coded candidates are byte-identical cu12 builds that os.path.exists-pass but ctypes.CDLL-fail with libcudart.so.12: cannot open shared object file; only ..._preload_cu13.abi3.so loads.
• torch_memory_saver already ships a CUDA-aware resolver — utils.get_binary_path_from_package("torch_memory_saver_hook_mode_preload") — that detects the CUDA major (via torch.version.cuda, then a libcudart.so.<major> probe over (13, 12)) and returns the matching build. It is what the library uses internally in hooks/mode_preload.py.

The change

Replace the hard-coded filename loop with the library resolver, with a robust fallback for older torch_memory_saver versions that predate it. The fallback keys on the CUDA major torch was built for (torch.version.cuda) — compile-time metadata that needs no GPU/driver, so it works on the Ray driver, which may be a GPU-less (CPU-only) head node. A ctypes.CDLL loadability probe is only a last resort for when the major is unknown, since that probe needs the host libcuda.so.1 and would otherwise fail a config the GPU workers could run. (The version slime's Dockerfile pins today, @a193d9dd, already has the resolver, so the resolver path is what runs in practice; the fallback keeps the fix correct against older/other pins.) Only a missing resolver falls back — a resolver that exists but raises is propagated, not masked.

slime/ray/actor_group.py:

if self.args.offload_train and self.args.train_backend == "megatron":
    import torch_memory_saver

    dynlib_path = _resolve_tms_preload_lib(torch_memory_saver)

    env_vars["LD_PRELOAD"] = dynlib_path
    env_vars["TMS_INIT_ENABLE"] = "1"
    env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1"

with a module-level helper:

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(
        f"No torch_memory_saver preload library found under {base}. "
        f"Not found: {missing or 'none'}."
    )

Before / after on CUDA 13:

# before
os.path.exists("..._preload_cu12.abi3.so") -> True -> LD_PRELOAD=cu12
-> children: libcudart.so.12: cannot open shared object file  (train worker dies)

# after
get_binary_path_from_package(...) -> "..._preload_cu13.abi3.so"
-> LD_PRELOAD=cu13, loads cleanly, train worker starts

Alternatives considered (and why rejected)

I checked each candidate approach on a live CUDA 13 worker before choosing.

1. Just add cu13 to the hard-coded list (keep os.path.exists, put cu13 first).
   On this box it does pick a loadable cu13 (verified). Rejected: it still selects by existence, not loadability, and it re-hard-codes the CUDA set — the next CUDA bump (cu14) reintroduces the identical bug. It duplicates a decision the library already owns.

2. Just switch os.path.exists → ctypes.CDLL (keep the current cu12/unsuffixed list).
   Verified result: selects nothing (picked=None) on CUDA 13, because cu13 is not in the list — loadability alone cannot save a list that omits the right file. Rejected: incomplete.

3. Delegate to the resolver only, no fallback.
   Verified: returns the correct cu13. The version slime pins today (@a193d9dd) does have the resolver, so resolver-only would work against the current pin — but keying the fix solely on the presence of a newer API makes it silently regress to the old crash if the pin is ever bumped down or another environment ships an older torch_memory_saver (the resolver was added at some point; anything before it raises ImportError). Hence: resolver first, deterministic fallback second, so the fix is correct regardless of the pinned torch_memory_saver version.

4. A ctypes.CDLL loadability probe as the fallback's primary key.
   Tempting (loadability is the true test), but it runs in the process that calls this — the Ray driver, which may be a CPU-only head node with no libcuda.so.1. There every candidate (including the correct cu13) fails to dlopen, so the fallback would raise even though the GPU workers (where LD_PRELOAD actually takes effect) could load it. Rejected as the primary key. The fallback instead keys on torch.version.cuda (build-time metadata, no driver needed) and uses CDLL only as a last resort when the major is unknown, then degrades to existence-with-warning rather than failing a runnable config.

The chosen fix (resolver-first, then a CUDA-major existence fallback with a dlopen last resort) is the only option that is correct on CUDA 13, unchanged on CUDA 12, forward-compatible with future CUDA majors, safe on the older pinned torch_memory_saver, and safe on a GPU-less Ray driver. No new module, class, or config flag is introduced.

Blast radius / risk

• Low. On CUDA 12 the resolver returns the cu12/generic build exactly as before; the fallback's first loadable candidate is equivalent to today's behavior.
• The only behavioral change is on runtimes where the previously-selected build could not load (e.g. CUDA 13): those now select the loadable build, or fail with an explicit, actionable error instead of an opaque libcudart.so.12 crash in a child process.
• No CLI/API surface changes; offload_train=False is entirely unaffected.
• Does not conflict with fix:TorchMemorySaver observes invalid LD_PRELOAD. when add --disable-weights-backuper #1937 (which edits the guard condition of this same block, not the .so selection); the two changes are orthogonal and compose cleanly.

Test plan (with commands)

Deterministic check (no full job, runs in any pod on the image):

python3 - <<'PY'
import os, ctypes, torch, torch_memory_saver
# the helper this PR adds (import it, or paste it inline):
from slime.ray.actor_group import _resolve_tms_preload_lib
p = _resolve_tms_preload_lib(torch_memory_saver)
ctypes.CDLL(p)                                  # must load without OSError
print("torch.version.cuda =", torch.version.cuda)
print("resolved           =", os.path.basename(p))   # -> ..._preload_cu13.abi3.so on CUDA 13
print("CDLL_load          = OK")
PY

Expected on CUDA 13:

torch.version.cuda = 13.0
resolved           = torch_memory_saver_hook_mode_preload_cu13.abi3.so
CDLL_load          = OK

End-to-end (a Megatron --offload-train run): launch any offload_train Megatron recipe (e.g. Qwen3-4B), then confirm the train worker no longer hits the cu12 crash. The check below resolves the driver log by content, so it needs no hand-pasted Ray job id (finished jobs age out of ray job list):

LOG=$(ls -t /tmp/ray/session_latest/logs/job-driver-*.log 2>/dev/null | head -1)
echo "driver_log=$LOG"
grep -c "libcudart.so.12" "$LOG"
# Expected: 0   (before the fix: the train worker dies on this right after the
#                rollout servers come up)

Verified on hardware (H200, 2× p5en.48xlarge, CUDA 13.0)

• Patched helper invoked directly on the H200 workers and, crucially, on the GPU-less head:
  • resolver path (the one SLIME ships): resolves ..._preload_cu13.abi3.so, CDLL loads.
  • fallback path (resolver import forced unavailable): keys on torch.version.cuda=13.0 and selects ..._preload_cu13.abi3.so by existence — verified on the GPU-less head (no libcuda.so.1), where the old dlopen-gated approach failed every candidate. This is the CPU-only-driver case.
  • resolver-error propagation (resolver present but made to raise): the error propagates instead of being silently masked by the fallback.
• Unit tests (tests/test_tms_preload_resolution.py, 4 passed): resolver-first, resolver-error propagation, CUDA-major existence fallback, and CUDA 12 (cu12 selected, unchanged).
• Full --offload-train Qwen3-4B run: libcudart.so.12 errors = 0 (previously the train worker died on it), and the job advanced into the training phase.
• Contrast: before the change the two hard-coded candidates are byte-identical cu12 builds (md5 9e724677…) and both fail to load on CUDA 13; only the cu13 build (md5 1bd3044c…, never enumerated by the old code) loads.

Checklist

• .so selection delegates to torch_memory_saver's resolver when available; a resolver that raises is propagated, not masked.
• Fallback keys on the detected CUDA major (torch.version.cuda) by existence — safe on a GPU-less Ray driver — and only dlopen-probes when the major is unknown.
• Verified on CUDA 13 (H200): LD_PRELOAD resolves to the loadable ..._preload_cu13.abi3.so and the --offload-train run no longer hits libcudart.so.12.
• CUDA 12 unchanged: covered by test_fallback_cuda12_selects_cu12 (the cu12 build is selected by existence, as before).
• Unit tests added: tests/test_tms_preload_resolution.py (resolver-first, resolver-error propagation, CUDA-major existence fallback, CUDA 12) — 4 passed.