diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index c3947adf..7be0ca7a 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp @@ -95,9 +95,14 @@ int64_t host_device_ptr(int64_t host_ptr) { } void host_register(int64_t addr, int64_t nbytes) { - const cudaError_t err = - cudaHostRegister(reinterpret_cast(addr), static_cast(nbytes), - cudaHostRegisterPortable | cudaHostRegisterMapped); + // Zero-size registrations are rejected by the driver with cudaErrorInvalidValue. + // Callers pin-after-fill small scale banks whose padded size can round to zero + // on dense layers; clamp so they register as a 1-byte mapped region instead. + const size_t reg_nbytes = + static_cast(nbytes > 0 ? nbytes : 1); + const cudaError_t err = cudaHostRegister( + reinterpret_cast(addr), reg_nbytes, + cudaHostRegisterPortable | cudaHostRegisterMapped); TORCH_CHECK(err == cudaSuccess, "cudaHostRegister failed: ", cudaGetErrorString(err)); } diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py index 5b50b25e..6c00718f 100644 --- a/python/freetoken/moe/host_banks.py +++ b/python/freetoken/moe/host_banks.py @@ -22,6 +22,7 @@ import mmap import os import queue +import sys import threading from concurrent.futures import ThreadPoolExecutor from enum import Enum @@ -88,8 +89,41 @@ def pin(self) -> None: try: host_register(self.addr, len(self._buf)) except RuntimeError as exc: + msg = str(exc).lower() + lock_words = {"lock", "locked", "locking", "memlock"} + if "out of memory" in msg or bool(lock_words & set(msg.split())) or "memlock" in msg: + if sys.platform == "win32": + hint = ( + "the driver refused to lock more host memory (on Windows/WDDM " + "the pageable-locking pool is roughly half of system RAM); " + "close other processes, add RAM, or serve some layers with " + "--moe-backend cpu" + ) + else: + try: + import resource + + soft = resource.getrlimit(resource.RLIMIT_MEMLOCK)[0] + if soft == resource.RLIM_INFINITY: + rlimit = "unlimited" + elif soft >= 1073741824: + rlimit = f"{soft / 1073741824:.0f} GB" + else: + rlimit = f"{soft // 1048576} MB" + except Exception: + rlimit = "?" + hint = ( + "the kernel refused to lock more host memory: RLIMIT_MEMLOCK " + f"is {rlimit} (often RAM/8 by default). Check `ulimit -l`, " + "raise it via /etc/security/limits.d/*.conf and re-login -- " + "an existing SSH ControlMaster session keeps serving the old " + "limit until `ssh -O exit `" + ) + else: + hint = "see the driver error above" raise RuntimeError( - f"cudaHostRegister failed for {len(self._buf) / 2**30:.1f} GiB" + f"cudaHostRegister failed for a {len(self._buf) / 2**30:.2f} GiB bank " + f"({msg}): {hint}" ) from exc self._pinned = True