trainstream 0.1.1__py3-none-any.whl

trainstream/csrc/offload_adamw.h

@@ -0,0 +1,154 @@
+// TrainStream — offloaded AdamW engine.
+//
+// Two execution modes share one bookkeeping model ("slots", one per parameter):
+//
+//  * "cpu": gradients are streamed D2H as soon as autograd produces them; host
+//    worker threads accumulate them into fp32 buffers, run AdamW on fp32 master
+//    weights held in pinned RAM, and stream the updated weights back H2D on a
+//    dedicated copy stream. step() returns immediately; the next forward pass
+//    waits per-module (wait_ready) so layer 0 can start computing while the
+//    CPU is still updating layer N.
+//
+//  * "gpu": gradients stay in VRAM; a CUDA kernel updates the parameter in
+//    place while reading/writing the moments through mapped pinned memory
+//    (zero-copy over PCIe). Moments never occupy VRAM.
+#pragma once
+
+#include "trainstream.h"
+
+#include <condition_variable>
+#include <functional>
+#include <mutex>
+#include <thread>
+#include <deque>
+
+namespace trainstream {
+
+class WorkerPool {
+ public:
+  WorkerPool(int threads, int device);
+  ~WorkerPool();
+  int size() const { return static_cast<int>(workers_.size()); }
+  // Jobs submitted to the same worker run in FIFO order.
+  void submit(size_t worker, std::function<void()> fn);
+
+ private:
+  struct Worker {
+    std::thread thread;
+    std::mutex mu;
+    std::condition_variable cv;
+    std::deque<std::function<void()>> queue;
+    bool stop = false;
+  };
+  std::vector<std::unique_ptr<Worker>> workers_;
+};
+
+class OffloadAdamW {
+ public:
+  OffloadAdamW(int device, const std::string& mode, int threads, int64_t chunk_elems);
+  ~OffloadAdamW();
+
+  // pin_moments=false keeps exp_avg/exp_avg_sq in ordinary RAM (cpu mode only).
+  int64_t add_param(const at::Tensor& param, bool pin_moments = true);
+
+  // cpu mode: stream one gradient to host and queue its accumulation. With
+  // apply_update the AdamW update for this tensor is queued right behind it
+  // ("eager" / optimizer-in-backward), overlapping host math with the rest of
+  // backward; the H2D upload is ordered after every kernel queued so far.
+  void offload_grad(int64_t id, const at::Tensor& grad, bool apply_update = false, double lr = 0.0,
+                    double beta1 = 0.9, double beta2 = 0.999, double eps = 1e-8, double weight_decay = 0.0);
+
+  // Queue one optimizer step for `ids` (in the order the next forward needs
+  // them). `grads` is used by gpu mode only. Returns the global grad norm when
+  // it was computed on the host (cpu mode with clipping / compute_norm), else -1.
+  double step(const std::vector<int64_t>& ids, const std::vector<double>& lrs,
+              const std::vector<double>& beta1s, const std::vector<double>& beta2s,
+              const std::vector<double>& epss, const std::vector<double>& wds,
+              const std::vector<at::Tensor>& grads, double max_grad_norm, bool compute_norm);
+
+  // Host-waits until each slot's update has been *issued*, then makes the
+  // current compute stream wait (on the GPU) for the H2D / kernel to land.
+  void wait_ready(const std::vector<int64_t>& ids);
+  void synchronize();
+
+  double last_grad_norm();
+  std::tuple<at::Tensor, at::Tensor, at::Tensor> state(int64_t id);  // master|undef, m, v
+  void load_state(int64_t id, const at::Tensor& exp_avg, const at::Tensor& exp_avg_sq,
+                  const at::Tensor& master, int64_t step);
+  int64_t get_step(int64_t id);
+  std::map<std::string, double> stats();
+  std::string mode() const { return gpu_mode_ ? "gpu" : "cpu"; }
+  int64_t num_params() const { return static_cast<int64_t>(slots_.size()); }
+  int threads() const { return pool_ ? pool_->size() : 0; }
+
+ private:
+  struct Slot {
+    int64_t id = 0;
+    at::Tensor param;  // flat alias of the user's parameter (shares storage)
+    int64_t numel = 0;
+    at::ScalarType pdtype = at::kFloat;
+    at::Tensor master, exp_avg, exp_avg_sq;  // fp32, pinned host
+    float* d_master = nullptr;               // device aliases (gpu mode)
+    float* d_m = nullptr;
+    float* d_v = nullptr;
+    at::Tensor grad_stage, grad_acc, param_stage;  // cpu mode staging
+    at::Tensor grad_keep;  // VRAM gradient kept alive until its DMA / kernel is done
+    int64_t step = 0;
+    bool has_grad = false;
+    bool acc_fresh = true;
+    EventPtr d2h_done;
+    EventPtr ready;
+    std::mutex mu;
+    std::condition_variable cv;
+    int inflight = 0;
+    int64_t nchunks = 1;
+  };
+  using ChunkFn = std::function<void(int64_t chunk, int64_t begin, int64_t end)>;
+
+  Slot& slot(int64_t id);
+  size_t worker_for(int64_t id, int64_t chunk) const;
+  void run_batch(Slot& s, std::shared_ptr<ChunkFn> body, std::shared_ptr<std::function<void()>> finalize);
+  void wait_idle(Slot& s);
+  // Queue AdamW + H2D for one slot. Reads the fp32 accumulator, or the raw
+  // staging buffer once `from_stage` (its D2H event) has completed.
+  void enqueue_update(Slot& s, double lr, double beta1, double beta2, double eps, double wd, double clip,
+                      EventPtr after, EventPtr from_stage);
+  void record_error(const std::string& msg);
+  void check_error();
+  double step_cpu(const std::vector<int64_t>& ids, const std::vector<double>& lrs,
+                  const std::vector<double>& beta1s, const std::vector<double>& beta2s,
+                  const std::vector<double>& epss, const std::vector<double>& wds,
+                  double max_grad_norm, bool compute_norm);
+  void step_gpu(const std::vector<int64_t>& ids, const std::vector<double>& lrs,
+                const std::vector<double>& beta1s, const std::vector<double>& beta2s,
+                const std::vector<double>& epss, const std::vector<double>& wds,
+                const std::vector<at::Tensor>& grads, double max_grad_norm, bool compute_norm);
+
+  int device_;
+  bool gpu_mode_;
+  int64_t chunk_;
+  std::unique_ptr<WorkerPool> pool_;
+  std::vector<std::unique_ptr<Slot>> slots_;
+
+  std::mutex err_mu_;
+  std::string error_;
+
+  // gpu-mode global-norm scratch
+  float* norm_dev_ = nullptr;
+  at::Tensor norm_host_;
+  EventPtr norm_ev_;
+  double last_norm_ = -1.0;
+
+  // telemetry (microseconds / bytes)
+  std::atomic<int64_t> update_us_{0}, accumulate_us_{0}, ready_stall_us_{0}, grad_stall_us_{0};
+  std::atomic<int64_t> grad_bytes_{0}, param_bytes_{0}, steps_{0}, steps_eager_{0};
+  std::atomic<int64_t> host_state_bytes_{0}, pinned_state_bytes_{0}, pageable_state_bytes_{0};
+  std::atomic<double> step_start_ms_{0.0}, last_finalize_ms_{0.0};
+};
+
+// Throughput of the host AdamW loop (elements / second) across `threads`.
+std::map<std::string, double> benchmark_cpu_adamw(int64_t numel, int threads);
+// Throughput of the zero-copy kernel with moments in mapped pinned memory.
+std::map<std::string, double> benchmark_zero_copy_adamw(int device, int64_t numel);
+
+}  // namespace trainstream

trainstream/csrc/offload_adamw_kernel.cu

@@ -0,0 +1,183 @@
+// TrainStream — CUDA kernels for the zero-copy offloaded AdamW path.
+//
+// The optimizer moments are never resident in VRAM: the kernel dereferences
+// mapped pinned host pointers directly, so each warp-wide load becomes a PCIe
+// read request and each store a posted PCIe write. The kernel is therefore
+// PCIe-bound by design (~16-24 bytes of host traffic per parameter), and the
+// main tuning lever is keeping enough independent requests in flight to cover
+// the ~1-2 us PCIe round-trip latency. We do that with a 4-way unrolled
+// "load everything first, then compute" body and a grid sized to fill every SM.
+
+#include "kernels.h"
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+namespace trainstream {
+namespace kernels {
+namespace {
+
+constexpr int kThreads = 256;
+constexpr int kUnroll = 4;
+constexpr int kBlocksPerSm = 16;
+
+template <typename T>
+struct Cvt;
+
+template <>
+struct Cvt<float> {
+  __device__ __forceinline__ static float load(float x) { return x; }
+  __device__ __forceinline__ static float store(float x) { return x; }
+};
+
+template <>
+struct Cvt<__half> {
+  __device__ __forceinline__ static float load(__half x) { return __half2float(x); }
+  __device__ __forceinline__ static __half store(float x) { return __float2half_rn(x); }
+};
+
+template <>
+struct Cvt<__nv_bfloat16> {
+  __device__ __forceinline__ static float load(__nv_bfloat16 x) { return __bfloat162float(x); }
+  __device__ __forceinline__ static __nv_bfloat16 store(float x) { return __float2bfloat16_rn(x); }
+};
+
+inline int grid_for(int64_t n, int per_block, int sm_count) {
+  int64_t blocks = (n + per_block - 1) / per_block;
+  const int64_t cap = static_cast<int64_t>(sm_count > 0 ? sm_count : 80) * kBlocksPerSm;
+  if (blocks > cap) blocks = cap;
+  if (blocks < 1) blocks = 1;
+  return static_cast<int>(blocks);
+}
+
+template <typename T>
+__global__ void __launch_bounds__(kThreads)
+    adamw_zero_copy_kernel(T* __restrict__ param, const T* __restrict__ grad,
+                           float* __restrict__ exp_avg, float* __restrict__ exp_avg_sq,
+                           float* __restrict__ master, int64_t n, AdamWParams hp,
+                           const float* __restrict__ grad_sq_norm) {
+  float clip = 1.0f;
+  if (grad_sq_norm != nullptr) {
+    const float norm = sqrtf(*grad_sq_norm);
+    clip = fminf(1.0f, hp.max_grad_norm / (norm + 1e-6f));
+  }
+  const float decay = 1.0f - hp.lr * hp.weight_decay;
+  const float one_minus_b1 = 1.0f - hp.beta1;
+  const float one_minus_b2 = 1.0f - hp.beta2;
+
+  const int64_t tile = static_cast<int64_t>(kThreads) * kUnroll;
+  const int64_t stride = static_cast<int64_t>(gridDim.x) * tile;
+
+  for (int64_t base = static_cast<int64_t>(blockIdx.x) * tile + threadIdx.x; base < n;
+       base += stride) {
+    float g[kUnroll], m[kUnroll], v[kUnroll], p[kUnroll];
+
+    // Phase 1: issue every load for the unrolled lanes before consuming any of
+    // them, so up to 3*kUnroll host reads per thread are in flight together.
+#pragma unroll
+    for (int u = 0; u < kUnroll; ++u) {
+      const int64_t i = base + static_cast<int64_t>(u) * kThreads;
+      if (i < n) {
+        g[u] = Cvt<T>::load(grad[i]) * clip;
+        m[u] = exp_avg[i];
+        v[u] = exp_avg_sq[i];
+        p[u] = master != nullptr ? master[i] : Cvt<T>::load(param[i]);
+      }
+    }
+
+    // Phase 2: AdamW math (matches torch.optim.AdamW, foreach=False semantics).
+#pragma unroll
+    for (int u = 0; u < kUnroll; ++u) {
+      const int64_t i = base + static_cast<int64_t>(u) * kThreads;
+      if (i < n) {
+        float pu = p[u] * decay;
+        const float mu = m[u] + one_minus_b1 * (g[u] - m[u]);  // lerp
+        const float vu = hp.beta2 * v[u] + one_minus_b2 * g[u] * g[u];
+        const float denom = sqrtf(vu) / hp.bias_correction2_sqrt + hp.eps;
+        pu -= hp.step_size * (mu / denom);
+        exp_avg[i] = mu;
+        exp_avg_sq[i] = vu;
+        if (master != nullptr) master[i] = pu;
+        param[i] = Cvt<T>::store(pu);
+      }
+    }
+  }
+}
+
+template <typename T>
+__global__ void __launch_bounds__(kThreads)
+    sq_norm_kernel(const T* __restrict__ grad, int64_t n, float* __restrict__ out) {
+  __shared__ float warp_sums[kThreads / 32];
+  float acc = 0.0f;
+  const int64_t stride = static_cast<int64_t>(gridDim.x) * kThreads;
+  for (int64_t i = static_cast<int64_t>(blockIdx.x) * kThreads + threadIdx.x; i < n; i += stride) {
+    const float x = Cvt<T>::load(grad[i]);
+    acc += x * x;
+  }
+#pragma unroll
+  for (int off = 16; off > 0; off >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, off);
+  if ((threadIdx.x & 31) == 0) warp_sums[threadIdx.x >> 5] = acc;
+  __syncthreads();
+  if (threadIdx.x < 32) {
+    acc = threadIdx.x < (kThreads / 32) ? warp_sums[threadIdx.x] : 0.0f;
+#pragma unroll
+    for (int off = 16; off > 0; off >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, off);
+    if (threadIdx.x == 0) atomicAdd(out, acc);
+  }
+}
+
+template <typename T>
+cudaError_t launch_adamw(void* param, const void* grad, float* m, float* v, float* master,
+                         int64_t n, const AdamWParams& hp, const float* norm, int sm_count,
+                         cudaStream_t stream) {
+  if (n == 0) return cudaSuccess;
+  const int grid = grid_for(n, kThreads * kUnroll, sm_count);
+  adamw_zero_copy_kernel<T><<<grid, kThreads, 0, stream>>>(
+      static_cast<T*>(param), static_cast<const T*>(grad), m, v, master, n, hp, norm);
+  return cudaGetLastError();
+}
+
+template <typename T>
+cudaError_t launch_sq_norm(const void* grad, int64_t n, float* out, int sm_count,
+                           cudaStream_t stream) {
+  if (n == 0) return cudaSuccess;
+  const int grid = grid_for(n, kThreads, sm_count);
+  sq_norm_kernel<T><<<grid, kThreads, 0, stream>>>(static_cast<const T*>(grad), n, out);
+  return cudaGetLastError();
+}
+
+}  // namespace
+
+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) {
+  switch (dtype) {
+    case DType::kFloat32:
+      return launch_adamw<float>(param, grad, exp_avg, exp_avg_sq, master, n, hp, grad_sq_norm,
+                                 sm_count, stream);
+    case DType::kFloat16:
+      return launch_adamw<__half>(param, grad, exp_avg, exp_avg_sq, master, n, hp, grad_sq_norm,
+                                  sm_count, stream);
+    case DType::kBFloat16:
+      return launch_adamw<__nv_bfloat16>(param, grad, exp_avg, exp_avg_sq, master, n, hp,
+                                         grad_sq_norm, sm_count, stream);
+  }
+  return cudaErrorInvalidValue;
+}
+
+cudaError_t grad_sq_norm_accumulate(DType dtype, const void* grad, int64_t n, float* out,
+                                    int sm_count, cudaStream_t stream) {
+  switch (dtype) {
+    case DType::kFloat32:
+      return launch_sq_norm<float>(grad, n, out, sm_count, stream);
+    case DType::kFloat16:
+      return launch_sq_norm<__half>(grad, n, out, sm_count, stream);
+    case DType::kBFloat16:
+      return launch_sq_norm<__nv_bfloat16>(grad, n, out, sm_count, stream);
+  }
+  return cudaErrorInvalidValue;
+}
+
+}  // namespace kernels
+}  // namespace trainstream

trainstream/csrc/trainstream.h

@@ -0,0 +1,132 @@
+// TrainStream — shared runtime declarations (events, streams, pinned memory,
+// asynchronous transfers). Implemented in memory_pinner.cpp.
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_runtime.h>
+
+#include <atomic>
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <string>
+#include <tuple>
+#include <vector>
+
+#define TS_CUDA_CHECK(expr)                                                                   \
+  do {                                                                                       \
+    const cudaError_t _ts_err = (expr);                                                      \
+    TORCH_CHECK(_ts_err == cudaSuccess, "TrainStream CUDA error: ", cudaGetErrorString(_ts_err), \
+                " (", #expr, ") at ", __FILE__, ":", __LINE__);                              \
+  } while (0)
+
+namespace trainstream {
+
+// ---------------------------------------------------------------------------
+// Events
+// ---------------------------------------------------------------------------
+// RAII wrapper around cudaEvent_t. Non-timing events are recycled through a
+// per-device pool because transfer-heavy workloads create thousands per step.
+// Re-recording a recycled event is safe: cudaStreamWaitEvent snapshots the
+// event's state at call time, so earlier waiters are unaffected.
+class Event {
+ public:
+  Event(int device, bool timing);
+  ~Event();
+  Event(const Event&) = delete;
+  Event& operator=(const Event&) = delete;
+
+  void record(cudaStream_t stream);
+  bool query() const;                    // true if completed (or never recorded)
+  void synchronize() const;              // host blocks until completion
+  void block(cudaStream_t stream) const; // GPU-side: stream waits for this event
+  float elapsed_ms(const Event& end) const;
+  bool recorded() const { return recorded_; }
+  int device() const { return device_; }
+
+ private:
+  cudaEvent_t ev_ = nullptr;
+  int device_;
+  bool timing_;
+  bool recorded_ = false;
+};
+using EventPtr = std::shared_ptr<Event>;
+
+EventPtr make_event(int device, bool timing = false);
+EventPtr record_on_current(int device, bool timing = false);
+
+// ---------------------------------------------------------------------------
+// Runtime: dedicated streams + traffic counters, one instance per device.
+// ---------------------------------------------------------------------------
+struct Runtime {
+  explicit Runtime(int device);
+
+  int device;
+  int sm_count = 0;
+  cudaStream_t d2h = nullptr;  // device -> host (gradients, activations out)
+  cudaStream_t h2d = nullptr;  // host -> device (parameters, activation prefetch)
+  cudaStream_t upd = nullptr;  // zero-copy optimizer kernels
+
+  std::atomic<uint64_t> bytes_d2h{0};
+  std::atomic<uint64_t> bytes_h2d{0};
+  std::atomic<uint64_t> copies_d2h{0};
+  std::atomic<uint64_t> copies_h2d{0};
+};
+
+Runtime& runtime(int device);
+int resolve_device(int device);  // -1 -> current device
+cudaStream_t compute_stream(int device);
+std::vector<int> initialized_devices();
+
+// Tell the CUDA caching allocator that `t`'s memory is in use on `stream`, so
+// the block is not handed out again until that stream's queued work finishes.
+void record_stream(const at::Tensor& t, cudaStream_t stream, int device);
+
+// ---------------------------------------------------------------------------
+// Pinned host memory (pin-logging caching allocator)
+// ---------------------------------------------------------------------------
+// cached=true  -> returned to a size-bucketed free list on release (cheap reuse;
+//                 used for activation stashes that churn every step).
+// cached=false -> cudaFreeHost on release (persistent optimizer state).
+// mapped=true  -> cudaHostAllocMapped: device-addressable for zero-copy kernels.
+at::Tensor pinned_empty(const std::vector<int64_t>& sizes, at::ScalarType dtype,
+                        const std::string& tag, bool mapped, bool cached);
+std::map<std::string, int64_t> pinned_stats();
+std::map<std::string, int64_t> pinned_tag_bytes();
+// (timestamp_ms, op, bytes, tag) — op is one of "alloc", "free", "hit", "release".
+std::vector<std::tuple<double, std::string, int64_t, std::string>> pinned_log(int64_t limit);
+void pinned_empty_cache();
+void pinned_set_cache_limit(int64_t bytes);
+void* pinned_device_pointer(const at::Tensor& t);  // for mapped allocations
+
+// ---------------------------------------------------------------------------
+// Asynchronous transfers on the dedicated streams.
+// ---------------------------------------------------------------------------
+// D2H: the copy stream first waits for the *current* compute stream so it
+// reads fully produced data. The caller must keep `src` alive until the
+// returned event completes (no recordStream; see memory_pinner.cpp).
+EventPtr d2h_async(const at::Tensor& dst, const at::Tensor& src);
+// H2D into an existing device tensor; waits for `after` (if any) first.
+EventPtr h2d_async(const at::Tensor& dst, const at::Tensor& src, const EventPtr& after);
+// H2D into a freshly allocated device tensor (from the compute stream's pool;
+// the copy is ordered after work already queued there). The consumer must
+// wait on the returned event before using or releasing the tensor.
+// With `sizes`/`strides` the device tensor is created with that (dense) layout,
+// restoring the exact strides a tensor had before it was paged out.
+std::tuple<at::Tensor, EventPtr> h2d_alloc_async(const at::Tensor& src, const EventPtr& after,
+                                                 int device, const std::vector<int64_t>& sizes = {},
+                                                 const std::vector<int64_t>& strides = {});
+std::map<std::string, int64_t> transfer_stats(int device);
+
+// ---------------------------------------------------------------------------
+// Micro-benchmarks (CUDA-event timed)
+// ---------------------------------------------------------------------------
+std::map<std::string, double> benchmark_pcie(int device, int64_t bytes, int iters);
+std::map<std::string, double> benchmark_matmul(int device, int64_t n, at::ScalarType dtype,
+                                               int iters);
+
+double now_ms();
+
+}  // namespace trainstream