trainstream 0.1.1__py3-none-any.whl

trainstream/__init__.py

@@ -0,0 +1,62 @@
+"""TrainStream: PCIe offloading engine for PyTorch training.
+
+    import trainstream
+    model, optimizer = trainstream.wrap(model, optimizer)
+
+Optimizer state, gradients and (a profiled fraction of) activations are moved
+to pinned host RAM over dedicated CUDA copy streams, so models that would not
+fit in VRAM with a stock optimizer can still be trained on a single GPU.
+"""
+
+import sys
+
+try:
+    import torch  # noqa: F401  (loads libtorch / c10 before the extension)
+except ImportError as _exc:  # pragma: no cover - depends on the environment
+    raise ImportError(
+        "TrainStream needs PyTorch with CUDA. Install it first, e.g.\n"
+        "    pip install torch --index-url https://download.pytorch.org/whl/cu130") from _exc
+
+try:
+    from . import _C  # prebuilt (build.bat / setup.py build_ext / --no-build-isolation install)
+except ImportError:
+    from . import _build
+
+    _C = _build.load()  # compiled on first import, then loaded from the cache
+    sys.modules[__name__ + "._C"] = _C
+
+from .hooks import ActivationOffloader, GradientOffloader, ParamReadyHooks, find_blocks
+from .optimizer import OffloadAdamW
+from .profiler import HardwareProfile, profile_hardware
+from .wrapper import TrainStreamSession, memory_report, unwrap, wrap
+
+__version__ = "0.1.1"
+
+__all__ = [
+    "wrap",
+    "unwrap",
+    "memory_report",
+    "OffloadAdamW",
+    "ActivationOffloader",
+    "GradientOffloader",
+    "ParamReadyHooks",
+    "TrainStreamSession",
+    "HardwareProfile",
+    "profile_hardware",
+    "find_blocks",
+    "pinned_stats",
+    "transfer_stats",
+]
+
+
+def pinned_stats():
+    """Pin-logging allocator counters plus live bytes per allocation tag."""
+    stats = dict(_C.pinned_stats())
+    stats["by_tag"] = dict(_C.pinned_tag_bytes())
+    return stats
+
+
+def transfer_stats(device=None):
+    """Cumulative bytes / copies issued on TrainStream's D2H and H2D streams."""
+    idx = torch.cuda.current_device() if device is None else (torch.device(device).index or 0)
+    return dict(_C.transfer_stats(idx))

trainstream/_build.py

@@ -0,0 +1,201 @@
+"""Compile-on-first-import for trainstream._C.
+
+`pip install trainstream` installs a pure-Python wheel that carries the C++/CUDA
+sources in trainstream/csrc. The first `import trainstream` compiles them
+against the PyTorch that is installed (a few minutes) and caches the result;
+later imports load the cached module directly, without a compiler.
+
+The cache lives in %LOCALAPPDATA%\\trainstream (Windows) or
+~/.cache/trainstream (Linux); TRAINSTREAM_CACHE_DIR overrides it. Builds are
+keyed by the sources, Python, PyTorch, CUDA and GPU architectures, so upgrading
+any of them triggers a fresh build instead of loading a stale one.
+"""
+
+import glob
+import hashlib
+import importlib.machinery
+import importlib.util
+import os
+import re
+import shutil
+import subprocess
+import sys
+
+import torch
+
+NAME = "trainstream_C"
+CSRC = os.path.join(os.path.dirname(os.path.abspath(__file__)), "csrc")
+SOURCES = ["bindings.cpp", "memory_pinner.cpp", "offload_adamw.cpp", "offload_adamw_kernel.cu"]
+HEADERS = ["kernels.h", "offload_adamw.h", "trainstream.h"]
+
+if sys.platform == "win32":
+    CXX_FLAGS = ["/O2", "/EHsc", "/W3", "/wd4244", "/wd4267", "/wd4251", "/wd4275", "/wd4996"]
+    NVCC_FLAGS = ["-O3", "--use_fast_math", "-lineinfo", "--expt-relaxed-constexpr", "-Xcompiler", "/wd4819"]
+    LIB_EXT = ".pyd"
+else:
+    CXX_FLAGS = ["-O3", "-Wall", "-Wno-unused-function", "-fvisibility=hidden"]
+    NVCC_FLAGS = ["-O3", "--use_fast_math", "-lineinfo", "--expt-relaxed-constexpr"]
+    LIB_EXT = ".so"
+
+
+class BuildError(ImportError):
+    pass
+
+
+def _cache_root():
+    if os.environ.get("TRAINSTREAM_CACHE_DIR"):
+        return os.environ["TRAINSTREAM_CACHE_DIR"]
+    if sys.platform == "win32":
+        base = os.environ.get("LOCALAPPDATA") or os.path.expanduser("~\\AppData\\Local")
+        return os.path.join(base, "trainstream", "build")
+    base = os.environ.get("XDG_CACHE_HOME") or os.path.expanduser("~/.cache")
+    return os.path.join(base, "trainstream")
+
+
+def _arch_list():
+    if os.environ.get("TORCH_CUDA_ARCH_LIST"):
+        return os.environ["TORCH_CUDA_ARCH_LIST"]
+    caps = sorted({torch.cuda.get_device_capability(i) for i in range(torch.cuda.device_count())})
+    return ";".join(f"{major}.{minor}" for major, minor in caps)
+
+
+def _build_key(arch):
+    h = hashlib.sha256()
+    for f in SOURCES + HEADERS:
+        with open(os.path.join(CSRC, f), "rb") as fh:
+            h.update(f.encode() + b"\0" + fh.read())
+    h.update(repr((sys.version_info[:2], torch.__version__, torch.version.cuda, arch,
+                   CXX_FLAGS, NVCC_FLAGS)).encode())
+    tag = f"py{sys.version_info[0]}{sys.version_info[1]}-torch{torch.__version__}"
+    return re.sub(r"[^A-Za-z0-9_.-]", "_", tag) + "-" + h.hexdigest()[:12]
+
+
+def _import_library(path):
+    loader = importlib.machinery.ExtensionFileLoader(NAME, path)
+    spec = importlib.util.spec_from_file_location(NAME, path, loader=loader)
+    module = importlib.util.module_from_spec(spec)
+    loader.exec_module(module)
+    return module
+
+
+def _version_key(path):
+    return tuple(int(p) for p in re.findall(r"\d+", os.path.basename(path)))
+
+
+def _setup_cuda_home():
+    if os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") or shutil.which("nvcc"):
+        return
+    if sys.platform == "win32":
+        pattern = os.path.join(os.environ.get("ProgramFiles", "C:\\Program Files"),
+                               "NVIDIA GPU Computing Toolkit", "CUDA", "v*")
+    else:
+        pattern = "/usr/local/cuda*"
+    found = [d for d in glob.glob(pattern) if os.path.isfile(os.path.join(
+        d, "bin", "nvcc.exe" if sys.platform == "win32" else "nvcc"))]
+    if not found:
+        return
+    major = (torch.version.cuda or "").split(".")[0]
+    same_major = [d for d in found if _version_key(d)[:1] == (int(major),)] if major.isdigit() else []
+    os.environ["CUDA_HOME"] = max(same_major or found, key=_version_key)
+
+
+def _setup_msvc():
+    """Import the x64 MSVC environment (vcvars64.bat) when cl.exe is not on PATH."""
+    if shutil.which("cl"):
+        return
+    vswhere = os.path.join(os.environ.get("ProgramFiles(x86)", "C:\\Program Files (x86)"),
+                           "Microsoft Visual Studio", "Installer", "vswhere.exe")
+    if not os.path.isfile(vswhere):
+        return
+    try:
+        out = subprocess.check_output(
+            [vswhere, "-latest", "-products", "*", "-requires",
+             "Microsoft.VisualStudio.Component.VC.Tools.x86.x64", "-property", "installationPath"],
+            text=True, stderr=subprocess.DEVNULL)
+    except (OSError, subprocess.CalledProcessError):
+        return
+    roots = [line.strip() for line in out.splitlines() if line.strip()]
+    if not roots:
+        return
+    vcvars = os.path.join(roots[0], "VC", "Auxiliary", "Build", "vcvars64.bat")
+    if not os.path.isfile(vcvars):
+        return
+    try:
+        env = subprocess.check_output(f'cmd /s /c ""{vcvars}" >nul && set"',
+                                      text=True, stderr=subprocess.DEVNULL)
+    except (OSError, subprocess.CalledProcessError):
+        return
+    for line in env.splitlines():
+        key, sep, value = line.partition("=")
+        if sep and key:
+            os.environ[key] = value
+    os.environ.setdefault("DISTUTILS_USE_SDK", "1")
+
+
+def _setup_ninja():
+    if shutil.which("ninja"):
+        return
+    try:
+        import ninja
+    except ImportError:
+        return
+    os.environ["PATH"] = ninja.BIN_DIR + os.pathsep + os.environ.get("PATH", "")
+
+
+def _hint():
+    if sys.platform == "win32":
+        return ("TrainStream needs the CUDA toolkit (same major version as PyTorch's CUDA, "
+                "https://developer.nvidia.com/cuda-downloads) and Visual Studio 2022+ Build Tools "
+                "with the 'Desktop development with C++' workload.")
+    return ("TrainStream needs the CUDA toolkit (same major version as PyTorch's CUDA) "
+            "and g++ 11+ (e.g. apt install build-essential).")
+
+
+def load():
+    """Return the compiled extension, building it first if no cached build matches."""
+    if torch.version.cuda is None:
+        raise BuildError(
+            f"TrainStream needs a CUDA build of PyTorch, but torch {torch.__version__} is CPU-only. "
+            "Install one from https://pytorch.org/get-started/locally/ , e.g.\n"
+            "    pip install torch --index-url https://download.pytorch.org/whl/cu130")
+    if not torch.cuda.is_available() and not os.environ.get("TORCH_CUDA_ARCH_LIST"):
+        raise BuildError("TrainStream needs an NVIDIA GPU, but torch.cuda.is_available() is False "
+                         "(set TORCH_CUDA_ARCH_LIST to compile without one).")
+
+    arch = _arch_list()
+    build_dir = os.path.join(_cache_root(), _build_key(arch))
+    library = os.path.join(build_dir, NAME + LIB_EXT)
+    if os.path.isfile(library):
+        return _import_library(library)
+
+    os.makedirs(build_dir, exist_ok=True)
+    _setup_ninja()
+    _setup_cuda_home()
+    if sys.platform == "win32":
+        _setup_msvc()
+    os.environ["TORCH_CUDA_ARCH_LIST"] = arch
+
+    from torch.utils import cpp_extension
+    if os.environ.get("CUDA_HOME"):
+        cpp_extension.CUDA_HOME = os.environ["CUDA_HOME"]
+    if cpp_extension.CUDA_HOME is None:
+        raise BuildError("TrainStream could not find the CUDA toolkit (nvcc). " + _hint()
+                         + " Set CUDA_HOME if it is installed somewhere unusual.")
+
+    print(f"TrainStream: compiling its CUDA extension for torch {torch.__version__} "
+          f"(first import only, takes a few minutes; cached in {build_dir})", file=sys.stderr, flush=True)
+    try:
+        module = cpp_extension.load(
+            name=NAME,
+            sources=[os.path.join(CSRC, f) for f in SOURCES],
+            extra_include_paths=[CSRC],
+            extra_cflags=CXX_FLAGS,
+            extra_cuda_cflags=NVCC_FLAGS,
+            build_directory=build_dir,
+            verbose=bool(os.environ.get("TRAINSTREAM_VERBOSE_BUILD")),
+        )
+    except Exception as exc:
+        raise BuildError(f"TrainStream failed to compile its CUDA extension. {_hint()}\n"
+                         f"Set TRAINSTREAM_VERBOSE_BUILD=1 to see the full compiler output.\n{exc}") from exc
+    print("TrainStream: build finished.", file=sys.stderr, flush=True)
+    return module

trainstream/csrc/bindings.cpp

@@ -0,0 +1,139 @@
+// TrainStream — Python bindings (module: trainstream._C).
+//
+// Every call that can block on the GPU or on worker threads releases the GIL,
+// so autograd's backward thread and the training thread never deadlock on it.
+
+#include <torch/extension.h>
+
+#include <cstdio>
+#include <exception>
+
+#include "offload_adamw.h"
+#include "trainstream.h"
+
+namespace py = pybind11;
+using namespace trainstream;
+
+namespace {
+using nogil = py::call_guard<py::gil_scoped_release>;
+
+// An exception escaping a destructor (e.g. a CUDA error surfacing inside the
+// caching allocator's free path) calls std::terminate, and the default abort
+// handler prints only a stack. Print the message first, then chain.
+std::terminate_handler g_prev_terminate = nullptr;
+void print_then_terminate() {
+  if (auto ex = std::current_exception()) {
+    try {
+      std::rethrow_exception(ex);
+    } catch (const std::exception& e) {
+      std::fprintf(stderr, "[trainstream] terminate: %s\n", e.what());
+    } catch (...) {
+      std::fprintf(stderr, "[trainstream] terminate: unknown exception\n");
+    }
+    std::fflush(stderr);
+  }
+  if (g_prev_terminate) g_prev_terminate();
+  std::abort();
+}
+
+std::map<std::string, std::string> build_info() {
+  int runtime_version = 0, driver_version = 0;
+  cudaRuntimeGetVersion(&runtime_version);
+  cudaDriverGetVersion(&driver_version);
+  return {{"cuda_runtime", std::to_string(runtime_version)},
+          {"cuda_driver", std::to_string(driver_version)},
+          {"compiled_with_cuda", std::to_string(CUDART_VERSION)},
+#ifdef _MSVC_LANG
+          {"cxx_standard", std::to_string(_MSVC_LANG)}};  // MSVC pins __cplusplus to 199711L
+#else
+          {"cxx_standard", std::to_string(__cplusplus)}};
+#endif
+}
+}  // namespace
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.doc() = "TrainStream native runtime: pinned memory, DMA streams, offloaded AdamW";
+  g_prev_terminate = std::set_terminate(print_then_terminate);
+
+  m.def("build_info", &build_info);
+  m.def("now_ms", &now_ms);
+  m.def("init", [](int device) { runtime(device); }, py::arg("device") = -1);
+  m.def("is_dense", [](const at::Tensor& t) { return t.is_non_overlapping_and_dense(); },
+        "True if the tensor occupies exactly numel*itemsize contiguous bytes (any dimension order).");
+
+  // ---- events -------------------------------------------------------------
+  py::class_<Event, std::shared_ptr<Event>>(m, "Event")
+      .def(py::init([](int device, bool timing) { return make_event(resolve_device(device), timing); }),
+           py::arg("device") = -1, py::arg("timing") = false)
+      .def("record", [](Event& e) { e.record(compute_stream(e.device())); },
+           "Record on the current PyTorch compute stream.")
+      .def("query", &Event::query)
+      .def("synchronize", &Event::synchronize, nogil())
+      .def("wait", [](Event& e) { e.block(compute_stream(e.device())); },
+           "Make the current compute stream wait for this event (GPU-side, non-blocking for the host).")
+      .def("elapsed_ms", &Event::elapsed_ms, nogil())
+      .def_property_readonly("recorded", &Event::recorded)
+      .def_property_readonly("device", &Event::device);
+  m.def("record_event", &record_on_current, py::arg("device") = -1, py::arg("timing") = false);
+
+  // ---- pinned memory --------------------------------------------------------
+  m.def("pinned_empty", &pinned_empty, py::arg("sizes"), py::arg("dtype"), py::arg("tag") = "user",
+        py::arg("mapped") = false, py::arg("cached") = true);
+  m.def("pinned_stats", &pinned_stats);
+  m.def("pinned_tag_bytes", &pinned_tag_bytes);
+  m.def("pinned_log", &pinned_log, py::arg("limit") = 256);
+  m.def("pinned_empty_cache", &pinned_empty_cache, nogil());
+  m.def("pinned_set_cache_limit", &pinned_set_cache_limit);
+
+  // ---- transfers ------------------------------------------------------------
+  m.def("d2h_async", &d2h_async, py::arg("dst"), py::arg("src"));
+  m.def("h2d_async", &h2d_async, py::arg("dst"), py::arg("src"), py::arg("after") = EventPtr());
+  m.def("h2d_alloc_async", &h2d_alloc_async, py::arg("src"), py::arg("after") = EventPtr(),
+        py::arg("device") = -1, py::arg("sizes") = std::vector<int64_t>{},
+        py::arg("strides") = std::vector<int64_t>{});
+  m.def("transfer_stats", &transfer_stats, py::arg("device") = -1);
+  m.def("synchronize_streams", [](int device) {
+        Runtime& rt = runtime(device);
+        TS_CUDA_CHECK(cudaStreamSynchronize(rt.d2h));
+        TS_CUDA_CHECK(cudaStreamSynchronize(rt.h2d));
+        TS_CUDA_CHECK(cudaStreamSynchronize(rt.upd));
+      }, py::arg("device") = -1, nogil());
+
+  // ---- benchmarks -----------------------------------------------------------
+  m.def("benchmark_pcie", &benchmark_pcie, py::arg("device") = -1, py::arg("bytes") = 256LL << 20,
+        py::arg("iters") = 10, nogil());
+  m.def("benchmark_matmul", &benchmark_matmul, py::arg("device") = -1, py::arg("n") = 4096,
+        py::arg("dtype") = at::kBFloat16, py::arg("iters") = 20, nogil());
+  m.def("benchmark_cpu_adamw", &benchmark_cpu_adamw, py::arg("numel") = 1 << 24, py::arg("threads") = 0,
+        nogil());
+  m.def("benchmark_zero_copy_adamw", &benchmark_zero_copy_adamw, py::arg("device") = -1,
+        py::arg("numel") = 1 << 24, nogil());
+
+  // ---- optimizer engine -----------------------------------------------------
+  py::class_<OffloadAdamW>(m, "OffloadAdamWEngine")
+      .def(py::init<int, const std::string&, int, int64_t>(), py::arg("device"), py::arg("mode") = "cpu",
+           py::arg("threads") = 0, py::arg("chunk_elems") = 1 << 20)
+      .def("add_param", &OffloadAdamW::add_param, py::arg("param"), py::arg("pin_moments") = true)
+      .def("offload_grad", &OffloadAdamW::offload_grad, py::arg("id"), py::arg("grad"),
+           py::arg("apply_update") = false, py::arg("lr") = 0.0, py::arg("beta1") = 0.9, py::arg("beta2") = 0.999,
+           py::arg("eps") = 1e-8, py::arg("weight_decay") = 0.0, nogil())
+      .def("step", &OffloadAdamW::step, py::arg("ids"), py::arg("lrs"), py::arg("beta1s"), py::arg("beta2s"),
+           py::arg("epss"), py::arg("wds"), py::arg("grads"), py::arg("max_grad_norm") = 0.0,
+           py::arg("compute_norm") = false, nogil())
+      .def("wait_ready", &OffloadAdamW::wait_ready, nogil())
+      .def("synchronize", &OffloadAdamW::synchronize, nogil())
+      .def("last_grad_norm", &OffloadAdamW::last_grad_norm, nogil())
+      .def("state", &OffloadAdamW::state, nogil())
+      .def("load_state",
+           [](OffloadAdamW& e, int64_t id, const at::Tensor& m, const at::Tensor& v,
+              const std::optional<at::Tensor>& master, int64_t step) {
+             e.load_state(id, m, v, master.value_or(at::Tensor()), step);
+           },
+           py::arg("id"), py::arg("exp_avg"), py::arg("exp_avg_sq"), py::arg("master") = py::none(),
+           py::arg("step") = 0, nogil())
+      .def("get_step", &OffloadAdamW::get_step)
+      .def("stats", &OffloadAdamW::stats)
+      .def_property_readonly("mode", &OffloadAdamW::mode)
+      .def_property_readonly("num_params", &OffloadAdamW::num_params)
+      .def_property_readonly("threads", &OffloadAdamW::threads);
+}

trainstream/csrc/kernels.h

@@ -0,0 +1,49 @@
+// TrainStream — device kernel launch interface.
+//
+// This header is deliberately free of PyTorch includes so that nvcc only has to
+// parse CUDA runtime headers when compiling offload_adamw_kernel.cu. The C++
+// translation units translate at::ScalarType into kernels::DType before calling.
+#pragma once
+
+#include <cstdint>
+#include <cuda_runtime.h>
+
+namespace trainstream {
+namespace kernels {
+
+enum class DType : int { kFloat32 = 0, kFloat16 = 1, kBFloat16 = 2 };
+
+// Hyper-parameters for one AdamW step on one tensor. Bias corrections are folded
+// on the host so the kernel does no pow() per element:
+//   step_size              = lr / (1 - beta1^t)
+//   bias_correction2_sqrt  = sqrt(1 - beta2^t)
+struct AdamWParams {
+  float lr;
+  float beta1;
+  float beta2;
+  float eps;
+  float weight_decay;
+  float step_size;
+  float bias_correction2_sqrt;
+  float max_grad_norm;  // only read when a grad_sq_norm pointer is supplied
+};
+
+// Fused AdamW where the optimizer state lives in *mapped pinned host memory*.
+// `exp_avg`, `exp_avg_sq` and (optionally) `master` are device-visible aliases of
+// cudaHostAlloc(cudaHostAllocMapped) buffers, so every state read/write is a PCIe
+// transaction and the moments never occupy VRAM. `param` and `grad` live in VRAM.
+// If `master` is non-null the fp32 master copy is updated and `param` receives a
+// rounded copy (mixed-precision training with fp16/bf16 weights).
+// If `grad_sq_norm` is non-null it points at a device float holding sum(g^2) over
+// all gradients; the kernel derives the global-norm clip coefficient from it.
+cudaError_t adamw_zero_copy(DType dtype, void* param, const void* grad, float* exp_avg,
+                            float* exp_avg_sq, float* master, int64_t n,
+                            const AdamWParams& hp, const float* grad_sq_norm, int sm_count,
+                            cudaStream_t stream);
+
+// out[0] += sum(grad[i]^2). Used for global gradient-norm clipping on the GPU path.
+cudaError_t grad_sq_norm_accumulate(DType dtype, const void* grad, int64_t n, float* out,
+                                    int sm_count, cudaStream_t stream);
+
+}  // namespace kernels
+}  // namespace trainstream