From d334f49755417a2c9c7b967ed3b8569bba7ce522 Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Sun, 9 Aug 2026 02:07:42 +0200 Subject: [PATCH 1/7] feat(cuda): add distributed training runtime Use direct NCCL bindings because current JavaCPP ProcessGroup constructors cannot safely model intrusive pointers. --- deps.edn | 6 +- examples/distributed_training.clj | 102 ++++++ src/clorch/amp.clj | 157 ++++++++++ src/clorch/cuda.clj | 28 +- src/clorch/data.clj | 155 +++++++++- src/clorch/distributed.clj | 429 ++++++++++++++++++++++++++ src/clorch/distributed/checkpoint.clj | 112 +++++++ src/clorch/distributed/launch.clj | 271 ++++++++++++++++ src/clorch/distributed/nccl.clj | 358 +++++++++++++++++++++ src/clorch/distributed/worker.clj | 45 +++ src/clorch/nn/functional.clj | 61 +++- src/clorch/nn/parallel.clj | 278 +++++++++++++++++ src/clorch/torch.clj | 52 +++- test/clorch/amp_test.clj | 39 +++ test/clorch/cuda_smoke_test.clj | 143 ++++++--- test/clorch/distributed_test.clj | 66 ++++ test/clorch/nn_functional_test.clj | 13 + test/clorch/torch_test.clj | 4 +- 18 files changed, 2249 insertions(+), 70 deletions(-) create mode 100644 examples/distributed_training.clj create mode 100644 src/clorch/amp.clj create mode 100644 src/clorch/distributed.clj create mode 100644 src/clorch/distributed/checkpoint.clj create mode 100644 src/clorch/distributed/launch.clj create mode 100644 src/clorch/distributed/nccl.clj create mode 100644 src/clorch/distributed/worker.clj create mode 100644 src/clorch/nn/parallel.clj create mode 100644 test/clorch/amp_test.clj create mode 100644 test/clorch/distributed_test.clj diff --git a/deps.edn b/deps.edn index 5c979c2..9fde3cd 100644 --- a/deps.edn +++ b/deps.edn @@ -1,8 +1,6 @@ {:paths ["src" "test" "examples"] - :deps {org.bytedeco/pytorch-platform {:mvn/version "2.1.2-1.5.10"} - org.bytedeco/pytorch-platform-gpu {:mvn/version "2.1.2-1.5.10"} - org.bytedeco/cuda-platform-redist {:mvn/version "12.3-8.9-1.5.10"} - org.bytedeco/openblas-platform {:mvn/version "0.3.26-1.5.10"} + :deps {org.bytedeco/pytorch-platform {:mvn/version "2.10.0-1.5.13"} + org.bytedeco/pytorch-platform-gpu {:mvn/version "2.10.0-1.5.13"} cnuernber/dtype-next {:mvn/version "11.007"} nrepl/nrepl {:mvn/version "1.1.1"} org.clojure/data.json {:mvn/version "2.5.0"} diff --git a/examples/distributed_training.clj b/examples/distributed_training.clj new file mode 100644 index 0000000..ba4dd52 --- /dev/null +++ b/examples/distributed_training.clj @@ -0,0 +1,102 @@ +(ns distributed-training + "Runnable local multi-GPU DDP training example." + (:require [clorch.amp :as amp] + [clorch.data :as data] + [clorch.distributed :as dist] + [clorch.nn :as nn] + [clorch.nn.functional :as F] + [clorch.nn.parallel :as ddp] + [clorch.optim :as optim] + [clorch.torch :as t]) + (:import [org.bytedeco.pytorch Scalar Tensor])) + +(defn- training-example [sample-index sample-count] + (let [x (/ (double sample-index) (double (max 1 (dec sample-count)))) + features [x (* x x) (Math/sin (* Math/PI x)) (Math/cos (* Math/PI x))] + target [(+ (* 0.3 x) (* -0.2 x x) (* 0.5 (Math/sin (* Math/PI x))))]] + [features target])) + +(defn- batch-tensors [indices sample-count] + (let [examples (mapv #(training-example % sample-count) indices)] + {:input (t/to (t/tensor (mapv first examples) {:dtype :float32}) :cuda) + :target (t/to (t/tensor (mapv second examples) {:dtype :float32}) :cuda)})) + +(defn- scaled-loss [loss divisor] + (with-open [value (Scalar. (double divisor))] + (.div ^Tensor loss value))) + +(defn train-worker + "Worker entrypoint used by `run-local!`; one invocation per GPU." + [{:keys [rank world-size process-group args]}] + (let [{:keys [epochs sample-count batch-size accumulation precision checkpoint-path] + :or {epochs 2 + sample-count 256 + batch-size 16 + accumulation 1 + precision :bfloat16}} args + _ (t/manual-seed 1337) + sampler (data/distributed-sampler + sample-count {:num-replicas world-size :rank rank :seed 1337}) + model (nn/to (nn/linear 4 1) :cuda) + optimizer (optim/adamw (nn/parameters model) :lr 0.01) + scaler (amp/grad-scaler {:enabled? (= precision :float16)}) + global-step (atom 0) + last-loss (atom nil)] + (with-open [parallel-model + (ddp/distributed-data-parallel + model {:process-group process-group :bucket-cap-mb 25.0})] + (dotimes [epoch epochs] + (data/set-epoch! sampler epoch) + (let [batches (partition-all batch-size (data/sample-indices sampler))] + (doseq [micro-batches (partition-all accumulation batches)] + (optim/zero-grad optimizer) + (doseq [[micro-index indices] (map-indexed vector micro-batches)] + (let [final-micro? (= micro-index (dec (count micro-batches))) + train-micro! + (fn [] + (t/with-torch + (let [{:keys [input target]} + (batch-tensors indices sample-count) + loss (amp/autocast + {:device :cuda :dtype precision} + (F/mse-loss (nn/forward parallel-model input) + target))] + (reset! last-loss (t/item-float loss)) + (amp/backward! scaler + (scaled-loss loss (count micro-batches))))))] + (if final-micro? + (train-micro!) + (ddp/no-sync (train-micro!))))) + (ddp/optimizer-step! parallel-model optimizer {:scaler scaler}) + (swap! global-step inc)))) + (when checkpoint-path + (dist/save-checkpoint! + process-group checkpoint-path + {:model model + :optimizer optimizer + :sampler sampler + :scaler scaler + :state {:epoch epochs :global-step @global-step}})) + (let [parameter (.get (nn/parameters model) 0) + rank-zero-parameter (.clone ^Tensor parameter)] + (dist/broadcast! process-group rank-zero-parameter {:root-rank 0}) + {:rank rank + :steps @global-step + :loss @last-loss + :precision precision + :parameters-synchronized? + (.allclose ^Tensor parameter ^Tensor rank-zero-parameter)})))) + +(defn run-local! + "Launches one training worker JVM per CUDA device and waits for completion." + [devices & [options]] + (let [job (dist/launch! + {:nproc-per-node (count devices) + :devices devices + :main 'distributed-training/train-worker + :args (or options {})})] + (dist/await-job! job) + {:status (dist/job-status job) + :logs (into {} + (map (fn [rank] [rank (dist/job-logs job rank)])) + (range (count devices)))})) diff --git a/src/clorch/amp.clj b/src/clorch/amp.clj new file mode 100644 index 0000000..10130c2 --- /dev/null +++ b/src/clorch/amp.clj @@ -0,0 +1,157 @@ +(ns clorch.amp + "CUDA/CPU autocast and dynamic gradient scaling." + (:require [clorch.torch :as t]) + (:import [org.bytedeco.pytorch Optimizer Scalar Tensor TensorVector] + [org.bytedeco.pytorch.global torch torch$DeviceType torch$ScalarType])) + +(defn- device-type [device] + (case device + :cuda torch$DeviceType/CUDA + :cpu torch$DeviceType/CPU + (throw (ex-info "Autocast device must be :cuda or :cpu" {:device device})))) + +(defn- scalar-type [dtype] + (case dtype + :float16 torch$ScalarType/Half + :bfloat16 torch$ScalarType/BFloat16 + (throw (ex-info "Autocast dtype must be :float16 or :bfloat16" {:dtype dtype})))) + +(defn call-with-autocast + "Runs `f` with thread-local autocast state, restoring prior state afterward." + [{:keys [device dtype enabled? cache?] + :or {device :cuda dtype :bfloat16 enabled? true cache? true}} + f] + (let [native-device-type (device-type device) + dtype (scalar-type dtype)] + (when (and enabled? (not (torch/is_autocast_available native-device-type))) + (throw (ex-info "Autocast is unavailable for this device" {:device device}))) + (let [old-enabled (torch/is_autocast_enabled native-device-type) + old-dtype (torch/get_autocast_dtype native-device-type) + old-cache (torch/is_autocast_cache_enabled)] + (try + (torch/set_autocast_dtype native-device-type dtype) + (torch/set_autocast_cache_enabled (boolean cache?)) + (torch/set_autocast_enabled native-device-type (boolean enabled?)) + (torch/increment_nesting) + (f) + (finally + (when (zero? (torch/decrement_nesting)) + (torch/clear_cache)) + (torch/set_autocast_enabled native-device-type old-enabled) + (torch/set_autocast_dtype native-device-type old-dtype) + (torch/set_autocast_cache_enabled old-cache)))))) + +(defmacro autocast + "Runs body under CUDA/CPU autocast. Defaults to CUDA bfloat16." + [options & body] + `(call-with-autocast ~options (fn [] ~@body))) + +(defrecord GradScaler + [scale growth-factor backoff-factor growth-interval growth-tracker enabled?]) + +(defn grad-scaler + "Creates a dynamic FP16 loss scaler. Disable it for bfloat16 training." + ([] (grad-scaler {})) + ([{:keys [initial-scale growth-factor backoff-factor growth-interval enabled?] + :or {initial-scale 65536.0 + growth-factor 2.0 + backoff-factor 0.5 + growth-interval 2000 + enabled? true}}] + (when-not (pos? initial-scale) + (throw (ex-info ":initial-scale must be positive" {:initial-scale initial-scale}))) + (when-not (> growth-factor 1.0) + (throw (ex-info ":growth-factor must exceed 1" {:growth-factor growth-factor}))) + (when-not (< 0.0 backoff-factor 1.0) + (throw (ex-info ":backoff-factor must be between 0 and 1" + {:backoff-factor backoff-factor}))) + (when-not (pos? growth-interval) + (throw (ex-info ":growth-interval must be positive" + {:growth-interval growth-interval}))) + (->GradScaler (atom (double initial-scale)) + (double growth-factor) + (double backoff-factor) + (long growth-interval) + (atom 0) + (boolean enabled?)))) + +(defn current-scale [scaler] + @(:scale scaler)) + +(defn scaler-state [scaler] + {:scale @(:scale scaler) + :growth-tracker @(:growth-tracker scaler)}) + +(defn load-scaler-state! [scaler {:keys [scale growth-tracker]}] + (when-not (and (number? scale) (pos? scale) + (integer? growth-tracker) (not (neg? growth-tracker))) + (throw (ex-info "Invalid gradient scaler checkpoint state" + {:scale scale :growth-tracker growth-tracker}))) + (reset! (:scale scaler) (double scale)) + (reset! (:growth-tracker scaler) (long growth-tracker)) + scaler) + +(defn scale-loss + "Multiplies loss by current scale when scaling is enabled." + [scaler loss] + (if (:enabled? scaler) + (with-open [factor (Scalar. (double @(:scale scaler)))] + (.mul ^Tensor loss factor)) + loss)) + +(defn backward! + "Scales loss and runs backward." + [scaler loss] + (.backward ^Tensor (scale-loss scaler loss)) + nil) + +(defn- optimizer-gradients [^Optimizer optimizer] + (let [^TensorVector parameters (.parameters optimizer)] + (keep (fn [index] + (let [parameter (.get parameters (long index)) + gradient (.grad ^Tensor parameter)] + (when (.defined gradient) gradient))) + (range (.size parameters))))) + +(defn- gradients-finite? [gradients] + (every? (fn [^Tensor gradient] + (t/with-torch + (.item_bool (.all (.isfinite gradient))))) + gradients)) + +(defn- globally-finite? [finite?] + (if ((requiring-resolve 'clorch.distributed/initialized?)) + (t/with-torch + (let [flag (t/tensor [(if finite? 0 1)] {:dtype :int32 :device :cuda})] + ((requiring-resolve 'clorch.distributed/all-reduce!) flag {:op :max}) + (zero? (.item_int ^Tensor flag)))) + finite?)) + +(defn- unscale-gradients! [gradients scale] + (with-open [divisor (Scalar. (double scale))] + (doseq [^Tensor gradient gradients] + (.dividePut gradient divisor)))) + +(defn- update-scale! [scaler finite?] + (if finite? + (let [steps (swap! (:growth-tracker scaler) inc)] + (when (>= steps (:growth-interval scaler)) + (swap! (:scale scaler) * (:growth-factor scaler)) + (reset! (:growth-tracker scaler) 0))) + (do + (swap! (:scale scaler) * (:backoff-factor scaler)) + (reset! (:growth-tracker scaler) 0)))) + +(defn step! + "Unscales gradients, synchronizes overflow state, conditionally steps optimizer, + and updates the dynamic scale. Returns true when the optimizer stepped." + [scaler ^Optimizer optimizer] + (if-not (:enabled? scaler) + (do (.step optimizer) true) + (let [gradients (vec (optimizer-gradients optimizer)) + finite? (globally-finite? (gradients-finite? gradients))] + (when finite? + (unscale-gradients! gradients @(:scale scaler)) + (.step optimizer)) + (update-scale! scaler finite?) + finite?))) diff --git a/src/clorch/cuda.clj b/src/clorch/cuda.clj index acbf6a9..80357d2 100644 --- a/src/clorch/cuda.clj +++ b/src/clorch/cuda.clj @@ -1,6 +1,7 @@ (ns clorch.cuda (:require [clorch.platform :as platform]) - (:import [org.bytedeco.pytorch.global torch] + (:import [org.bytedeco.cuda.global cudart] + [org.bytedeco.pytorch.global torch] [org.bytedeco.javacpp Loader])) ;; Force loading with GPU extension if possible @@ -30,6 +31,31 @@ [] (torch/cuda_device_count)) +(defn- check-runtime! [operation status] + (when-not (= cudart/cudaSuccess status) + (throw (ex-info (str operation " failed: " + (.getString (cudart/cudaGetErrorString status))) + {:operation operation + :status status})))) + +(defn set-device! + "Sets the CUDA device used by the current worker thread and returns its index." + [device-index] + (let [device-index (clojure.core/int device-index)] + (when-not (<= 0 device-index (dec (device-count))) + (throw (ex-info "CUDA device index is out of range" + {:device-index device-index + :device-count (device-count)}))) + (check-runtime! "cudaSetDevice" (cudart/cudaSetDevice device-index)) + device-index)) + +(defn current-device + "Returns the CUDA device index selected for the current worker thread." + [] + (let [device-index (int-array 1)] + (check-runtime! "cudaGetDevice" (cudart/cudaGetDevice device-index)) + (aget device-index 0))) + (defn synchronize "Waits for all kernels in all streams on a CUDA device to complete." ([] (torch/cuda_synchronize)) diff --git a/src/clorch/data.clj b/src/clorch/data.clj index be02818..bafe104 100644 --- a/src/clorch/data.clj +++ b/src/clorch/data.clj @@ -2,6 +2,7 @@ (:require [clojure.edn :as edn] [clojure.java.io :as io]) (:import [java.io BufferedReader BufferedWriter InputStreamReader OutputStreamWriter] + [java.util SplittableRandom] [java.util.concurrent Callable Future LinkedBlockingQueue ThreadFactory ThreadPoolExecutor TimeUnit TimeoutException] [java.util.concurrent.atomic AtomicInteger])) @@ -27,6 +28,11 @@ (defprotocol IProcessDataset (process-spec [this] "Returns EDN-serializable worker spec, e.g. {:factory 'my.ns/mk-ds :args [...]}")) +(defprotocol ISampler + (sample-indices [this] "Returns this rank's indices for the current epoch.") + (set-epoch! [this epoch] "Changes the deterministic shuffle epoch.") + (sampler-state [this] "Returns checkpointable sampler state.") + (load-sampler-state! [this state] "Restores checkpointable sampler state.")) (def ^:private javadataset-extension (delay (let [cls (Class/forName "org.bytedeco.pytorch.JavaDataset")] @@ -279,15 +285,130 @@ :get-item (fn [idx] {:data (select X-t 0 idx) :target (select y-t 0 idx)})))) +(defn- ceil-div [dividend divisor] + (quot (+ dividend (dec divisor)) divisor)) + +(defn- shuffled-range [n seed] + (let [values (long-array (range n)) + random (SplittableRandom. (long seed))] + (loop [i (dec n)] + (when (pos? i) + (let [j (.nextInt random (inc i)) + tmp (aget values i)] + (aset-long values i (aget values j)) + (aset-long values j tmp) + (recur (dec i))))) + (vec values))) + +(defn- sampler-size [dataset-size replicas drop-last?] + (if drop-last? + (if (zero? (mod dataset-size replicas)) + (quot dataset-size replicas) + (max 0 (ceil-div (- dataset-size replicas) replicas))) + (ceil-div dataset-size replicas))) + +(deftype DistributedSampler + [dataset-size replicas rank shuffle? seed drop-last? epoch] + ISampler + (sample-indices [_] + (let [base (if shuffle? + (shuffled-range dataset-size (+ seed @epoch)) + (vec (range dataset-size))) + samples-per-rank (sampler-size dataset-size replicas drop-last?) + total-size (* samples-per-rank replicas) + indices (cond + (zero? total-size) [] + drop-last? (subvec base 0 total-size) + (= total-size dataset-size) base + :else (into base (take (- total-size dataset-size) + (cycle base))))] + (->> indices + (drop rank) + (take-nth replicas) + (take samples-per-rank) + vec))) + (set-epoch! [this next-epoch] + (when (neg? next-epoch) + (throw (ex-info "Sampler epoch must be non-negative" + {:epoch next-epoch}))) + (reset! epoch (long next-epoch)) + this) + (sampler-state [_] + {:epoch @epoch + :seed seed + :rank rank + :replicas replicas + :dataset-size dataset-size + :drop-last? drop-last? + :shuffle? shuffle?}) + (load-sampler-state! [this state] + (doseq [[field expected] [[:seed seed] + [:rank rank] + [:replicas replicas] + [:dataset-size dataset-size] + [:drop-last? drop-last?] + [:shuffle? shuffle?]]] + (when-not (= expected (get state field)) + (throw (ex-info "Sampler checkpoint is incompatible" + {:field field + :expected expected + :actual (get state field)})))) + (set-epoch! this (:epoch state))) + clojure.lang.Seqable + (seq [this] + (seq (sample-indices this))) + clojure.lang.ILookup + (valAt [this lookup-key] + (.valAt this lookup-key nil)) + (valAt [_ lookup-key not-found] + (case lookup-key + :dataset-size dataset-size + :num-replicas replicas + :rank rank + :shuffle? shuffle? + :seed seed + :drop-last? drop-last? + :epoch @epoch + not-found))) + +(defn distributed-sampler + "Partitions a dataset deterministically across distributed ranks. + + Defaults :num-replicas and :rank from WORLD_SIZE and RANK. Call set-epoch! + before each epoch so every rank applies the same new shuffle." + [dataset-source & args] + (let [options (if (map? (first args)) (first args) (apply hash-map args)) + dataset-size (if (number? dataset-source) (long dataset-source) (long (get-size dataset-source))) + replicas (long (or (:num-replicas options) + (some-> (System/getenv "WORLD_SIZE") Long/parseLong) + 1)) + rank (long (or (:rank options) + (some-> (System/getenv "RANK") Long/parseLong) + 0)) + shuffle? (get options :shuffle? true) + seed (long (get options :seed 0)) + drop-last? (get options :drop-last? false)] + (when (neg? dataset-size) + (throw (ex-info "Dataset size must be non-negative" + {:dataset-size dataset-size}))) + (when-not (pos? replicas) + (throw (ex-info "Sampler replica count must be positive" + {:num-replicas replicas}))) + (when-not (<= 0 rank (dec replicas)) + (throw (ex-info "Sampler rank is outside replica range" + {:rank rank :num-replicas replicas}))) + (DistributedSampler. dataset-size replicas rank shuffle? seed drop-last? (atom 0)))) + (defn- resolve-worker-backend [num-workers worker-backend] (if (= worker-backend :auto) (if (pos? num-workers) :process :thread) worker-backend)) -(deftype ClorchDataloader [ds batch-size shuffle? idxs drop-last? num-workers prefetch-factor collate-fn worker-backend timeout-ms worker-init-fn] +(deftype ClorchDataloader [ds batch-size shuffle? sampler idxs drop-last? num-workers prefetch-factor collate-fn worker-backend timeout-ms worker-init-fn] clojure.lang.Seqable (seq [_] - (let [chunks (build-chunks idxs batch-size drop-last?)] + (let [epoch-indices (if sampler (sample-indices sampler) idxs) + chunks (build-chunks epoch-indices batch-size drop-last?)] (cond (and (pos? num-workers) (= worker-backend :process)) (seq (process-parallel-batches ds chunks collate-fn num-workers prefetch-factor timeout-ms worker-init-fn)) @@ -297,6 +418,9 @@ :else (map #(build-batch ds % collate-fn) chunks)))) + java.lang.Iterable + (iterator [this] + (clojure.lang.RT/iter (seq this))) clojure.lang.ILookup (valAt [this k] (.valAt this k nil)) (valAt [_ k not-found] @@ -304,7 +428,8 @@ :dataset ds :batch-size batch-size :shuffle? shuffle? - :indices idxs + :sampler sampler + :indices (if sampler (sample-indices sampler) idxs) :drop-last? drop-last? :num-workers num-workers :prefetch-factor prefetch-factor @@ -316,7 +441,7 @@ (defn dataloader [ds & args] (let [opts (if (map? (first args)) (first args) (apply hash-map args)) - {:keys [batch-size shuffle? drop-last? num-workers prefetch-factor collate-fn worker-backend timeout-ms worker-init-fn] + {:keys [batch-size shuffle? sampler drop-last? num-workers prefetch-factor collate-fn worker-backend timeout-ms worker-init-fn] :or {batch-size 32 shuffle? true drop-last? false @@ -326,15 +451,21 @@ worker-backend :auto timeout-ms nil worker-init-fn nil}} opts + _ (when (and sampler (contains? opts :shuffle?) shuffle?) + (throw (ex-info "A dataloader with a sampler cannot also shuffle" + {:shuffle? shuffle?}))) _ (when (= "org.bytedeco.pytorch.JavaDataset" (.getName (class ds))) (ensure-javadataset-extension!)) worker-backend (resolve-worker-backend num-workers worker-backend) n (get-size ds) - idxs (if shuffle? - (let [randperm (torch-fn :randperm) - item-float (torch-fn :item-float) - tseq (torch-fn :tseq) - perm (randperm n)] - (mapv #(long (item-float %)) (tseq perm))) - (vec (clojure.core/range n)))] - (->ClorchDataloader ds batch-size shuffle? idxs drop-last? num-workers prefetch-factor collate-fn worker-backend timeout-ms worker-init-fn))) + idxs (when-not sampler + (if shuffle? + (let [randperm (torch-fn :randperm) + item-float (torch-fn :item-float) + tseq (torch-fn :tseq) + perm (randperm n)] + (mapv #(long (item-float %)) (tseq perm))) + (vec (clojure.core/range n))))] + (->ClorchDataloader ds batch-size (and shuffle? (nil? sampler)) sampler idxs + drop-last? num-workers prefetch-factor collate-fn + worker-backend timeout-ms worker-init-fn))) diff --git a/src/clorch/distributed.clj b/src/clorch/distributed.clj new file mode 100644 index 0000000..a22fb43 --- /dev/null +++ b/src/clorch/distributed.clj @@ -0,0 +1,429 @@ +(ns clorch.distributed + "Multi-process distributed runtime. + + One JVM represents one rank. NCCL collectives use JavaCPP's public CUDA + bindings directly; this avoids a known ProcessGroup intrusive-pointer defect + in the current JavaCPP PyTorch preset." + (:refer-clojure :exclude [send]) + (:require [clojure.string :as str] + [clorch.platform :as platform]) + (:import [java.util.concurrent TimeoutException] + [org.bytedeco.javacpp SizeTPointer] + [org.bytedeco.javacpp.chrono Milliseconds] + [org.bytedeco.pytorch Store StringVector TCPStore TCPStoreOptions Tensor])) + +(platform/configure-platform-extension!) + +(declare destroy-process-group!) + +(defrecord ProcessGroupContext + [^Store store + backend-state + backend + rank + world-size + timeout-ms + master-address + master-port + local-rank + operation-counter] + java.lang.AutoCloseable + (close [this] + (destroy-process-group! this))) + +(defrecord StoreWork [task result]) + +(defonce ^:private default-context (atom nil)) + +(defn- environment [environment-name] + (System/getenv environment-name)) + +(defn- parse-long! [label value] + (try + (Long/parseLong (str value)) + (catch NumberFormatException cause + (throw (ex-info (str label " must be an integer") + {:label label :value value} + cause))))) + +(defn- option-or-env [options option-key environment-key default-value] + (if (contains? options option-key) + (get options option-key) + (or (environment environment-key) default-value))) + +(defn- validate-rendezvous! + [{:keys [backend rank world-size timeout-ms master-address master-port local-rank]}] + (when-not (= :nccl backend) + (throw (ex-info "The current JavaCPP release supports Clorch distributed execution through :nccl only" + {:backend backend + :reason :javacpp-process-group-construction-defect}))) + (when-not (pos? world-size) + (throw (ex-info "World size must be positive" {:world-size world-size}))) + (when-not (<= 0 rank (dec world-size)) + (throw (ex-info "Rank must be within the process group" + {:rank rank :world-size world-size}))) + (when-not (<= 0 local-rank) + (throw (ex-info "Local rank must be non-negative" {:local-rank local-rank}))) + (when-not (pos? timeout-ms) + (throw (ex-info "Process-group timeout must be positive" + {:timeout-ms timeout-ms}))) + (when (str/blank? master-address) + (throw (ex-info "Master address must not be blank" {}))) + (when-not (<= 1 master-port 65535) + (throw (ex-info "Master port must be between 1 and 65535" + {:master-port master-port})))) + +(defn- rendezvous-options [options] + (let [resolved {:backend (keyword (option-or-env options :backend "CLORCH_DIST_BACKEND" :nccl)) + :rank (parse-long! "rank" (option-or-env options :rank "RANK" 0)) + :world-size (parse-long! "world-size" (option-or-env options :world-size "WORLD_SIZE" 1)) + :local-rank (parse-long! "local-rank" (option-or-env options :local-rank "LOCAL_RANK" 0)) + :timeout-ms (parse-long! "timeout-ms" (option-or-env options :timeout-ms "CLORCH_DIST_TIMEOUT_MS" 300000)) + :master-address (str (option-or-env options :master-address "MASTER_ADDR" "127.0.0.1")) + :master-port (parse-long! "master-port" (option-or-env options :master-port "MASTER_PORT" 29500))}] + (validate-rendezvous! resolved) + resolved)) + +(defn- tcp-store + [{:keys [rank world-size timeout-ms master-address master-port]}] + (let [timeout (Milliseconds. timeout-ms) + workers (SizeTPointer. (long-array [world-size])) + options (doto (TCPStoreOptions.) + (.port (unchecked-short master-port)) + (.isServer (zero? rank)) + (.numWorkers workers) + (.waitWorkers true) + (.multiTenant false) + (.timeout timeout) + (.useLibUV true))] + (TCPStore. master-address options))) + +(defn init-process-group! + "Initializes this JVM's default NCCL process group. + + Options default from RANK, LOCAL_RANK, WORLD_SIZE, MASTER_ADDR, and + MASTER_PORT. One JVM must be launched per GPU." + ([] (init-process-group! {})) + ([options] + (locking default-context + (when @default-context + (throw (ex-info "A default process group is already initialized" + {:context @default-context}))) + (let [{:keys [rank world-size local-rank] :as resolved} + (rendezvous-options options)] + (when-not ((requiring-resolve 'clorch.cuda/available?)) + (throw (ex-info "NCCL requires an available CUDA runtime" + {:backend :nccl :rank rank}))) + (let [store (tcp-store resolved)] + (try + (let [backend-state + ((requiring-resolve 'clorch.distributed.nccl/create!) + store rank world-size local-rank) + context (map->ProcessGroupContext + (assoc resolved + :store store + :backend-state backend-state + :operation-counter (atom 0)))] + (reset! default-context context) + context) + (catch Throwable cause + (try (.close store) (catch Throwable _)) + (throw cause)))))))) + +(defn initialized? + "Returns true when this JVM owns a live default process group." + [] + (boolean @default-context)) + +(defn current-process-group + "Returns the default ProcessGroupContext, throwing when uninitialized." + [] + (or @default-context + (throw (ex-info "Distributed process group is not initialized" {})))) + +(defn- context-or-current [context] + (or context (current-process-group))) + +(defn rank + "Returns this worker's global rank." + ([] (:rank (current-process-group))) + ([context] (:rank (context-or-current context)))) + +(defn local-rank + "Returns this worker's node-local rank." + ([] (:local-rank (current-process-group))) + ([context] (:local-rank (context-or-current context)))) + +(defn world-size + "Returns the number of ranks in the process group." + ([] (:world-size (current-process-group))) + ([context] (:world-size (context-or-current context)))) + +(defn backend + "Returns the process group's backend keyword." + ([] (:backend (current-process-group))) + ([context] (:backend (context-or-current context)))) + +(defn rank-zero? + "Returns true for global rank zero." + ([] (zero? (rank))) + ([context] (zero? (rank context)))) + +(defn- submit-store-work [operation result async?] + (if async? + (->StoreWork (future (operation)) result) + (do (operation) result))) + +(defn await! + "Waits for asynchronous distributed work and returns its result." + ([work] (await! work (:timeout-ms (current-process-group)))) + ([work timeout-ms] + (cond + (instance? StoreWork work) + (let [outcome (deref (:task work) timeout-ms ::timeout)] + (when (= ::timeout outcome) + (throw (TimeoutException. + (str "Distributed operation exceeded " timeout-ms " ms")))) + (:result work)) + + (and (map? work) (contains? work :event)) + ((requiring-resolve 'clorch.distributed.nccl/await!) work timeout-ms) + + :else + (throw (ex-info "Unknown distributed work value" + {:value-type (some-> work class str)}))))) + +(defn- finish-work [work async? timeout-ms] + (if async? + work + (await! work timeout-ms))) + +(defn all-reduce! + "Reduces CUDA tensors across every rank in place." + ([tensors] (all-reduce! nil tensors {})) + ([tensors options] (all-reduce! nil tensors options)) + ([context tensors {:keys [op async? timeout-ms] + :or {op :sum async? false}}] + (let [context (context-or-current context) + timeout-ms (or timeout-ms (:timeout-ms context)) + work ((requiring-resolve 'clorch.distributed.nccl/all-reduce!) + (:backend-state context) tensors op)] + (finish-work work async? timeout-ms)))) + +(defn broadcast! + "Broadcasts CUDA tensors from root rank in place." + ([tensors] (broadcast! nil tensors {})) + ([tensors options] (broadcast! nil tensors options)) + ([context tensors {:keys [root-rank root-tensor async? timeout-ms] + :or {root-rank 0 root-tensor 0 async? false}}] + (when-not (zero? root-tensor) + (throw (ex-info "Direct NCCL broadcast supports root tensor index zero only" + {:root-tensor root-tensor}))) + (let [context (context-or-current context) + timeout-ms (or timeout-ms (:timeout-ms context)) + work ((requiring-resolve 'clorch.distributed.nccl/broadcast!) + (:backend-state context) tensors root-rank)] + (finish-work work async? timeout-ms)))) + +(defn reduce! + "Reduces CUDA tensors in place onto root rank." + ([tensors] (reduce! nil tensors {})) + ([tensors options] (reduce! nil tensors options)) + ([context tensors {:keys [op root-rank root-tensor async? timeout-ms] + :or {op :sum root-rank 0 root-tensor 0 async? false}}] + (when-not (zero? root-tensor) + (throw (ex-info "Direct NCCL reduce supports root tensor index zero only" + {:root-tensor root-tensor}))) + (let [context (context-or-current context) + timeout-ms (or timeout-ms (:timeout-ms context)) + work ((requiring-resolve 'clorch.distributed.nccl/reduce!) + (:backend-state context) tensors op root-rank)] + (finish-work work async? timeout-ms)))) + +(defn all-gather-into! + "Gathers equal-sized CUDA inputs into a preallocated output tensor." + ([output input] (all-gather-into! nil output input {})) + ([output input options] (all-gather-into! nil output input options)) + ([context ^Tensor output ^Tensor input {:keys [async? timeout-ms] + :or {async? false}}] + (let [context (context-or-current context) + timeout-ms (or timeout-ms (:timeout-ms context)) + work ((requiring-resolve 'clorch.distributed.nccl/all-gather-into!) + (:backend-state context) output input (:world-size context))] + (finish-work work async? timeout-ms)))) + +(defn reduce-scatter-into! + "Reduces an input CUDA tensor and scatters equal chunks into output." + ([output input] (reduce-scatter-into! nil output input {})) + ([output input options] (reduce-scatter-into! nil output input options)) + ([context ^Tensor output ^Tensor input {:keys [op async? timeout-ms] + :or {op :sum async? false}}] + (let [context (context-or-current context) + timeout-ms (or timeout-ms (:timeout-ms context)) + work ((requiring-resolve 'clorch.distributed.nccl/reduce-scatter-into!) + (:backend-state context) output input (:world-size context) op)] + (finish-work work async? timeout-ms)))) + +(defn all-to-all-single! + "Performs all-to-all into preallocated CUDA output with optional split sizes." + ([output input] (all-to-all-single! nil output input {})) + ([output input options] (all-to-all-single! nil output input options)) + ([context ^Tensor output ^Tensor input + {:keys [output-split-sizes input-split-sizes async? timeout-ms] + :or {output-split-sizes [] input-split-sizes [] async? false}}] + (let [context (context-or-current context) + timeout-ms (or timeout-ms (:timeout-ms context)) + work ((requiring-resolve 'clorch.distributed.nccl/all-to-all-single!) + (:backend-state context) output input (:world-size context) + output-split-sizes input-split-sizes)] + (finish-work work async? timeout-ms)))) + +(defn send + "Sends CUDA tensors to a peer rank. Returns input after completion or work." + ([tensors destination] (send nil tensors destination {})) + ([tensors destination options] (send nil tensors destination options)) + ([context tensors destination {:keys [tag async? timeout-ms] + :or {tag 0 async? false}}] + (let [context (context-or-current context) + timeout-ms (or timeout-ms (:timeout-ms context)) + work ((requiring-resolve 'clorch.distributed.nccl/send!) + (:backend-state context) tensors destination tag)] + (finish-work work async? timeout-ms)))) + +(defn receive! + "Receives CUDA tensors from a peer rank in place." + ([tensors source] (receive! nil tensors source {})) + ([tensors source options] (receive! nil tensors source options)) + ([context tensors source {:keys [tag async? timeout-ms] + :or {tag 0 async? false}}] + (let [context (context-or-current context) + timeout-ms (or timeout-ms (:timeout-ms context)) + work ((requiring-resolve 'clorch.distributed.nccl/receive!) + (:backend-state context) tensors source tag)] + (finish-work work async? timeout-ms)))) + +(defn- next-operation-id! [context] + (swap! (:operation-counter context) inc)) + +(defn- store-barrier! [context operation-id timeout-ms] + (let [^Store store (:store context) + prefix (str "clorch/barrier/" operation-id "/") + rank-key (str prefix (:rank context)) + store-keys (StringVector.)] + (.set store rank-key "ready") + (dotimes [peer (:world-size context)] + (.push_back store-keys (str prefix peer))) + (._wait store store-keys (Milliseconds. timeout-ms)))) + +(defn barrier! + "Blocks until every rank enters the barrier." + ([] (barrier! nil {})) + ([options] (barrier! nil options)) + ([context {:keys [async? timeout-ms] :or {async? false}}] + (let [context (context-or-current context) + timeout-ms (or timeout-ms (:timeout-ms context)) + operation-id (next-operation-id! context)] + (submit-store-work #(store-barrier! context operation-id timeout-ms) + nil async?)))) + +(defn rank-zero-call! + "Runs f only on rank zero and propagates its success or failure to all ranks." + ([f] (rank-zero-call! nil f)) + ([context f] + (let [context (context-or-current context) + ^Store store (:store context) + operation-id (next-operation-id! context) + result-key (str "clorch/rank-zero-result/" operation-id) + result (when (rank-zero? context) + (try + (let [value (f)] + (.set store result-key "ok") + {:value value}) + (catch Throwable cause + (.set store result-key + (str "error\n" (.getName (class cause)) "\n" + (.getMessage cause))) + {:error cause}))) + status (.get_to_str store result-key)] + (if (= status "ok") + (:value result) + (let [[_ class-name message] (str/split status #"\n" 3)] + (if-let [cause (:error result)] + (throw cause) + (throw (ex-info "Rank-zero operation failed" + {:remote-class class-name + :remote-message message})))))))) + +(defn destroy-process-group! + "Synchronizes and releases the native communicator. Repeated calls are safe." + ([] + (when-let [context @default-context] + (destroy-process-group! context))) + ([context] + (locking default-context + (when (and context + (or (identical? context @default-context) + (nil? @default-context))) + (try + ((requiring-resolve 'clorch.distributed.nccl/destroy!) + (:backend-state context)) + (finally + (try (.close ^Store (:store context)) (catch Throwable _)) + (when (identical? context @default-context) + (reset! default-context nil)))))) + nil)) + +(defmacro with-process-group + "Initializes a process group for body and always destroys it." + [options & body] + `(let [context# (init-process-group! ~options)] + (try + ~@body + (finally + (destroy-process-group! context#))))) + +(defn launch! + "Launches local worker JVMs. See clorch.distributed.launch/launch!." + [config] + ((requiring-resolve 'clorch.distributed.launch/launch!) config)) + +(defn job-status [job] + ((requiring-resolve 'clorch.distributed.launch/status) job)) + +(defn await-job! + ([job] + ((requiring-resolve 'clorch.distributed.launch/await!) job)) + ([job timeout-ms] + ((requiring-resolve 'clorch.distributed.launch/await!) job timeout-ms))) + +(defn stop-job! [job] + ((requiring-resolve 'clorch.distributed.launch/stop!) job)) + +(defn job-logs + ([job worker-rank] + ((requiring-resolve 'clorch.distributed.launch/logs) job worker-rank)) + ([job worker-rank max-bytes] + ((requiring-resolve 'clorch.distributed.launch/logs) job worker-rank max-bytes))) + +(defn checkpoint! + "Runs a checkpoint writer on rank zero and propagates failure to every rank." + ([write-checkpoint!] + (rank-zero-call! write-checkpoint!)) + ([context write-checkpoint!] + (rank-zero-call! context write-checkpoint!))) + +(defn save-checkpoint! + "Atomically saves distributed model, optimizer, sampler, and training state." + ([path values] + (save-checkpoint! nil path values)) + ([context path values] + ((requiring-resolve 'clorch.distributed.checkpoint/save!) + context path values))) + +(defn load-checkpoint! + "Restores a distributed checkpoint on every rank and returns training state." + ([path targets] + (load-checkpoint! nil path targets)) + ([context path targets] + ((requiring-resolve 'clorch.distributed.checkpoint/load!) + context path targets))) diff --git a/src/clorch/distributed/checkpoint.clj b/src/clorch/distributed/checkpoint.clj new file mode 100644 index 0000000..63e1b3b --- /dev/null +++ b/src/clorch/distributed/checkpoint.clj @@ -0,0 +1,112 @@ +(ns clorch.distributed.checkpoint + "Atomic rank-zero checkpoints for distributed training." + (:require [clojure.edn :as edn] + [clojure.java.io :as io] + [clorch.distributed :as dist] + [clorch.torch :as t]) + (:import [java.nio.file AtomicMoveNotSupportedException CopyOption Files + StandardCopyOption] + [java.util UUID] + [org.bytedeco.pytorch Generator] + [org.bytedeco.pytorch.global torch])) + +(def ^:private checkpoint-version 1) + +(defn- metadata-path [path] + (str path ".edn")) + +(defn- ensure-edn! [value] + (let [encoded (pr-str value)] + (when-not (= value (edn/read-string encoded)) + (throw (ex-info "Checkpoint training state must round-trip through EDN" + {:state value}))) + encoded)) + +(defn- atomic-move! [source destination] + (let [source-path (.toPath (io/file source)) + destination-path (.toPath (io/file destination))] + (try + (Files/move source-path destination-path + (into-array CopyOption + [StandardCopyOption/ATOMIC_MOVE + StandardCopyOption/REPLACE_EXISTING])) + (catch AtomicMoveNotSupportedException _ + (Files/move source-path destination-path + (into-array CopyOption + [StandardCopyOption/REPLACE_EXISTING])))))) + +(defn- sampler-state [sampler] + (when sampler + ((requiring-resolve 'clorch.data/sampler-state) sampler))) + +(defn- archive-values [model optimizer cpu-rng-state] + (cond-> {:model model + :cpu-rng-state cpu-rng-state} + optimizer (assoc :optimizer optimizer))) + +(defn save! + "Writes model, optimizer, CPU RNG, sampler, scaler, and EDN training state on rank zero. + + The tensor archive is `path`; metadata is committed last to `path.edn`, so a + missing metadata file always denotes an incomplete checkpoint." + [context path {:keys [model optimizer sampler scaler state]}] + (when-not model + (throw (ex-info "Checkpoint requires :model" {}))) + (let [context (or context (dist/current-process-group)) + metadata {:version checkpoint-version + :world-size (:world-size context) + :sampler (sampler-state sampler) + :scaler (when scaler + ((requiring-resolve 'clorch.amp/scaler-state) scaler)) + :state state} + metadata-content (ensure-edn! metadata)] + (dist/rank-zero-call! + context + (fn [] + (let [suffix (str ".tmp." (UUID/randomUUID)) + archive-temp (str path suffix) + metadata-temp (str (metadata-path path) suffix) + ^Generator cpu-generator (torch/getDefaultCPUGenerator) + cpu-rng-state (.get_state cpu-generator)] + (try + (t/save (archive-values model optimizer cpu-rng-state) + archive-temp) + (spit metadata-temp metadata-content) + (atomic-move! archive-temp path) + (atomic-move! metadata-temp (metadata-path path)) + (finally + (Files/deleteIfExists (.toPath (io/file archive-temp))) + (Files/deleteIfExists (.toPath (io/file metadata-temp)))))))) + (dist/barrier! context) + path)) + +(defn load! + "Restores model, optimizer, CPU RNG, sampler, scaler, and training state on every rank." + [context path {:keys [model optimizer sampler scaler]}] + (when-not model + (throw (ex-info "Checkpoint requires :model" {}))) + (let [context (or context (dist/current-process-group)) + metadata-file (io/file (metadata-path path))] + (when-not (.isFile metadata-file) + (throw (ex-info "Checkpoint metadata is missing or incomplete" + {:path (metadata-path path)}))) + (let [metadata (edn/read-string (slurp metadata-file))] + (when-not (= checkpoint-version (:version metadata)) + (throw (ex-info "Unsupported checkpoint version" + {:expected checkpoint-version + :actual (:version metadata)}))) + (let [^Generator cpu-generator (torch/getDefaultCPUGenerator) + cpu-rng-state (.get_state cpu-generator)] + (t/load (archive-values model optimizer cpu-rng-state) path) + (.set_state cpu-generator cpu-rng-state)) + (when (and sampler (:sampler metadata)) + ((requiring-resolve 'clorch.data/load-sampler-state!) + sampler + (assoc (:sampler metadata) + :rank (:rank context) + :replicas (:world-size context)))) + (when (and scaler (:scaler metadata)) + ((requiring-resolve 'clorch.amp/load-scaler-state!) + scaler (:scaler metadata))) + (dist/barrier! context) + (:state metadata)))) diff --git a/src/clorch/distributed/launch.clj b/src/clorch/distributed/launch.clj new file mode 100644 index 0000000..cff5685 --- /dev/null +++ b/src/clorch/distributed/launch.clj @@ -0,0 +1,271 @@ +(ns clorch.distributed.launch + "Local multi-process launcher used by REPLs and production entrypoints." + (:require [clojure.edn :as edn] + [clojure.java.io :as io]) + (:import [java.io File RandomAccessFile] + [java.net ServerSocket] + [java.nio.charset StandardCharsets] + [java.nio.file Files] + [java.time Instant] + [java.util UUID] + [java.util.concurrent LinkedBlockingQueue TimeUnit TimeoutException])) + +(declare stop!) + +(defrecord WorkerProcess [rank ^Process process ^File stdout ^File stderr]) + +(defrecord DistributedJob + [id config workers status started-at monitor config-file] + java.lang.AutoCloseable + (close [this] + (stop! this))) + +(defn- ensure-java-25! [] + (let [feature (.feature (Runtime/version))] + (when (< feature 25) + (throw (ex-info "Distributed workers require Java 25 or newer" + {:java-feature feature}))))) + +(defn- available-port [] + (with-open [socket (ServerSocket. 0)] + (.setReuseAddress socket true) + (.getLocalPort socket))) + +(defn- normalize-config [config] + (let [nproc (long (get config :nproc-per-node 1)) + backend (keyword (get config :backend :nccl)) + main (:main config) + devices (:devices config) + timeout-ms (long (get config :timeout-ms 300000))] + (when-not (pos? nproc) + (throw (ex-info ":nproc-per-node must be positive" + {:nproc-per-node nproc}))) + (when-not (= :nccl backend) + (throw (ex-info ":backend must be :nccl" + {:backend backend}))) + (when-not (qualified-symbol? main) + (throw (ex-info ":main must be a namespace-qualified symbol" + {:main main}))) + (when-not (pos? timeout-ms) + (throw (ex-info ":timeout-ms must be positive" + {:timeout-ms timeout-ms}))) + (when (and devices (not= nproc (count devices))) + (throw (ex-info "Device count must equal :nproc-per-node" + {:nproc-per-node nproc :devices devices}))) + (when (and devices (not= (count devices) (count (set devices)))) + (throw (ex-info "Each distributed worker requires a distinct CUDA device" + {:devices devices}))) + (when (nil? devices) + (throw (ex-info "NCCL launch requires explicit :devices" + {:nproc-per-node nproc}))) + (assoc config + :nproc-per-node nproc + :backend backend + :timeout-ms timeout-ms + :master-address (str (get config :master-address "127.0.0.1")) + :master-port (long (get config :master-port (available-port))) + :args (get config :args {})))) + +(defn- ensure-edn! [value] + (let [encoded (pr-str value) + decoded (edn/read-string encoded)] + (when-not (= value decoded) + (throw (ex-info "Worker configuration must round-trip through EDN" + {:value value}))) + encoded)) + +(defn- create-log-directory [config job-id] + (let [directory (or (:log-dir config) + (str (System/getProperty "java.io.tmpdir") + File/separator "clorch-distributed-" job-id)) + path (.toPath (io/file directory))] + (Files/createDirectories path (make-array java.nio.file.attribute.FileAttribute 0)) + (.toFile path))) + +(defn- write-worker-config! [config log-directory] + (let [path (.toPath (io/file log-directory "worker-config.edn")) + content (ensure-edn! (select-keys config [:main :args]))] + (Files/writeString path content StandardCharsets/UTF_8 + (into-array java.nio.file.OpenOption [])) + (.toFile path))) + +(defn- java-command [config config-file] + (let [java-home (or (:java-home config) (System/getProperty "java.home")) + java-executable (str java-home File/separator "bin" File/separator "java") + classpath (System/getProperty "java.class.path")] + (vec (concat [java-executable] + (:jvm-opts config) + ["-cp" classpath + "clojure.main" "-m" "clorch.distributed.worker" + (.getAbsolutePath ^File config-file)])))) + +(defn- configure-environment! + [^ProcessBuilder builder config rank] + (let [environment (.environment builder) + devices (:devices config) + assigned-device (when devices (nth devices rank)) + cuda-local-rank (if assigned-device 0 rank)] + (doseq [[environment-key value] (:environment config)] + (.put environment (name environment-key) (str value))) + (doseq [[environment-key value] {"RANK" rank + "LOCAL_RANK" cuda-local-rank + "WORLD_SIZE" (:nproc-per-node config) + "LOCAL_WORLD_SIZE" (:nproc-per-node config) + "MASTER_ADDR" (:master-address config) + "MASTER_PORT" (:master-port config) + "CLORCH_DIST_BACKEND" (name (:backend config)) + "CLORCH_DIST_TIMEOUT_MS" (:timeout-ms config)}] + (.put environment environment-key (str value))) + (when assigned-device + (.put environment "CUDA_VISIBLE_DEVICES" (str assigned-device))) + builder)) + +(defn- start-worker! [config config-file log-directory rank] + (let [stdout (io/file log-directory (str "rank-" rank ".out.log")) + stderr (io/file log-directory (str "rank-" rank ".err.log")) + builder (ProcessBuilder. ^java.util.List (java-command config config-file))] + (.directory builder (io/file (or (:working-directory config) + (System/getProperty "user.dir")))) + (.redirectOutput builder stdout) + (.redirectError builder stderr) + (configure-environment! builder config rank) + (->WorkerProcess rank (.start builder) stdout stderr))) + +(defn- terminate-process! [^Process process] + (when (.isAlive process) + (.destroy process) + (when-not (.waitFor process 10 TimeUnit/SECONDS) + (.destroyForcibly process) + (.waitFor process 10 TimeUnit/SECONDS)))) + +(defn- monitor-workers [workers status config-file] + (let [completions (LinkedBlockingQueue.)] + (doseq [{:keys [rank ^Process process]} workers] + (future + (.put completions [rank (.waitFor process)]))) + (loop [remaining (count workers) + exits (sorted-map) + failure nil] + (if (zero? remaining) + (let [final-status (if failure + {:state :failed + :failed-rank (first failure) + :exit-code (second failure) + :exit-codes exits + :finished-at (Instant/now)} + {:state :succeeded + :exit-codes exits + :finished-at (Instant/now)})] + (reset! status final-status) + (Files/deleteIfExists (.toPath ^File config-file)) + final-status) + (let [[rank exit-code] (.take completions) + first-failure (or failure (when-not (zero? exit-code) + [rank exit-code]))] + (when (and (nil? failure) first-failure) + (doseq [{:keys [^Process process]} workers] + (terminate-process! process))) + (recur (dec remaining) + (assoc exits rank exit-code) + first-failure)))))) + +(defn launch! + "Launches one local JVM per distributed rank and returns a job handle. + + Required: :main qualified worker function. NCCL also requires :devices, one + CUDA device identifier per rank. Worker stdout and stderr are file-backed and + never accumulated in coordinator memory." + [config] + (ensure-java-25!) + (let [config (normalize-config config) + job-id (str (UUID/randomUUID)) + log-directory (create-log-directory config job-id) + config-file (write-worker-config! config log-directory) + status (atom {:state :starting}) + workers (atom [])] + (try + (dotimes [rank (:nproc-per-node config)] + (swap! workers conj (start-worker! config config-file log-directory rank))) + (reset! status {:state :running}) + (let [monitor (future (monitor-workers @workers status config-file))] + (->DistributedJob job-id config @workers status (Instant/now) + monitor config-file)) + (catch Throwable cause + (doseq [{:keys [^Process process]} @workers] + (terminate-process! process)) + (Files/deleteIfExists (.toPath config-file)) + (throw cause))))) + +(defn status + "Returns immutable job and per-rank process status." + [job] + {:id (:id job) + :state @(:status job) + :started-at (:started-at job) + :workers (mapv (fn [{:keys [rank ^Process process stdout stderr]}] + {:rank rank + :pid (.pid process) + :alive? (.isAlive process) + :stdout (.getAbsolutePath ^File stdout) + :stderr (.getAbsolutePath ^File stderr)}) + (:workers job))}) + +(defn await! + "Waits for every rank. Throws when a rank fails or timeout elapses." + ([job] + (await! job nil)) + ([job timeout-ms] + (let [result (if timeout-ms + (deref (:monitor job) timeout-ms ::timeout) + @(:monitor job))] + (when (= ::timeout result) + (throw (TimeoutException. + (str "Distributed job exceeded " timeout-ms " ms")))) + (when (= :failed (:state result)) + (throw (ex-info "Distributed job failed" result))) + result))) + +(defn stop! + "Terminates all live ranks and waits for their process trees to exit." + [job] + (doseq [{:keys [^Process process]} (:workers job)] + (terminate-process! process)) + (when-not (realized? (:monitor job)) + (reset! (:status job) {:state :stopping})) + @(:monitor job)) + +(defn log-paths + "Returns rank-indexed stdout and stderr paths." + [job] + (into (sorted-map) + (map (fn [{:keys [rank stdout stderr]}] + [rank {:stdout (.getAbsolutePath ^File stdout) + :stderr (.getAbsolutePath ^File stderr)}]) + (:workers job)))) + +(defn- read-last-bytes [^File file max-bytes] + (if-not (.exists file) + "" + (with-open [input (RandomAccessFile. file "r")] + (let [length (.length input) + start (max 0 (- length max-bytes)) + size (int (- length start)) + payload (byte-array size)] + (.seek input start) + (.readFully input payload) + (String. payload StandardCharsets/UTF_8))))) + +(defn logs + "Returns a bounded tail of one rank's stdout and stderr." + ([job rank] + (logs job rank 65536)) + ([job rank max-bytes] + (when-not (pos? max-bytes) + (throw (ex-info "Log byte limit must be positive" + {:max-bytes max-bytes}))) + (let [{:keys [stdout stderr]} + (or (some #(when (= rank (:rank %)) %) (:workers job)) + (throw (ex-info "Unknown distributed rank" {:rank rank})))] + {:rank rank + :stdout (read-last-bytes stdout max-bytes) + :stderr (read-last-bytes stderr max-bytes)}))) diff --git a/src/clorch/distributed/nccl.clj b/src/clorch/distributed/nccl.clj new file mode 100644 index 0000000..c7b86af --- /dev/null +++ b/src/clorch/distributed/nccl.clj @@ -0,0 +1,358 @@ +(ns clorch.distributed.nccl + "Direct NCCL backend for Clorch distributed collectives. + + This namespace is loaded only for :nccl process groups so CPU-only JVMs never + initialize CUDA classes." + (:require [clorch.cuda :as cuda]) + (:import [org.bytedeco.cuda.cudart CUevent_st CUstream_st] + [org.bytedeco.cuda.global cudart nccl] + [org.bytedeco.cuda.nccl ncclComm ncclUniqueId] + [org.bytedeco.javacpp BytePointer IntPointer Pointer] + [org.bytedeco.pytorch ByteVector Store TensorBase TensorVector])) + +(defrecord NcclWork [^CUevent_st event ^ncclComm communicator result released?]) + +(defn- nccl-message [status] + (if-let [message (nccl/ncclGetErrorString status)] + (.getString message) + (str "NCCL status " status))) + +(defn- check-nccl! [operation status] + (when-not (= nccl/ncclSuccess status) + (throw (ex-info (str operation " failed: " (nccl-message status)) + {:operation operation :status status}))) + status) + +(defn- check-cuda! [operation status] + (when-not (= cudart/cudaSuccess status) + (throw (ex-info (str operation " failed with CUDA status " status) + {:operation operation :status status}))) + status) + +(defn- scalar-type->nccl [^TensorBase tensor] + (case (.toString (.scalar_type tensor)) + "Byte" nccl/ncclUint8 + "Char" nccl/ncclInt8 + "Int" nccl/ncclInt32 + "Long" nccl/ncclInt64 + "Half" nccl/ncclFloat16 + "Float" nccl/ncclFloat32 + "Double" nccl/ncclFloat64 + "BFloat16" nccl/ncclBfloat16 + "Bool" nccl/ncclUint8 + (throw (ex-info "NCCL does not support this tensor dtype" + {:dtype (.toString (.scalar_type tensor))})))) + +(defn- validate-tensor! [^TensorBase tensor] + (when-not (.is_cuda tensor) + (throw (ex-info "NCCL collectives require CUDA tensors" {}))) + (when-not (.is_contiguous tensor) + (throw (ex-info "NCCL collectives require contiguous tensors" {}))) + tensor) + +(defn- tensors [value] + (cond + (instance? TensorBase value) [value] + (instance? TensorVector value) + (mapv #(.get ^TensorVector value (long %)) (range (.size ^TensorVector value))) + (sequential? value) (mapv validate-tensor! value) + :else (throw (ex-info "Collectives require a tensor or tensor collection" + {:value-type (some-> value class str)})))) + +(defn- reduction [operation] + (case operation + :sum nccl/ncclSum + :avg nccl/ncclAvg + :product nccl/ncclProd + :min nccl/ncclMin + :max nccl/ncclMax + (throw (ex-info "Unsupported NCCL reduction operation" + {:operation operation})))) + +(defn- unique-id-bytes [^ncclUniqueId unique-id] + (byte-array + (map #(.internal unique-id (int %)) + (range nccl/NCCL_UNIQUE_ID_BYTES)))) + +(defn- install-unique-id! [^ncclUniqueId unique-id ^ByteVector byte-values] + (when-not (= nccl/NCCL_UNIQUE_ID_BYTES (.size byte-values)) + (throw (ex-info "Invalid NCCL rendezvous identifier" + {:expected nccl/NCCL_UNIQUE_ID_BYTES + :actual (.size byte-values)}))) + (dotimes [index nccl/NCCL_UNIQUE_ID_BYTES] + (.internal unique-id index (.get byte-values (long index)))) + unique-id) + +(defn create! + "Creates a rank-local NCCL communicator and non-blocking communication stream." + [^Store store rank world-size local-rank] + (cuda/set-device! local-rank) + (let [unique-id (ncclUniqueId.) + rendezvous-key "clorch/nccl/unique-id"] + (if (zero? rank) + (do + (check-nccl! "ncclGetUniqueId" (nccl/ncclGetUniqueId unique-id)) + (.set store rendezvous-key (ByteVector. (unique-id-bytes unique-id)))) + (install-unique-id! unique-id (.get store rendezvous-key))) + (let [communicator (ncclComm.) + stream (CUstream_st.)] + (try + (check-nccl! "ncclCommInitRank" + (nccl/ncclCommInitRank communicator (int world-size) + unique-id (int rank))) + (check-cuda! "cudaStreamCreateWithFlags" + (cudart/cudaStreamCreateWithFlags + stream cudart/cudaStreamNonBlocking)) + {:communicator communicator :stream stream} + (catch Throwable cause + (when-not (.isNull communicator) + (nccl/ncclCommAbort communicator)) + (throw cause)))))) + +(defn destroy! [{:keys [^ncclComm communicator ^CUstream_st stream]}] + (when (and stream (not (.isNull stream))) + (cudart/cudaStreamSynchronize stream)) + (when (and communicator (not (.isNull communicator))) + (let [finalize-status (nccl/ncclCommFinalize communicator)] + (when-not (= nccl/ncclSuccess finalize-status) + (nccl/ncclCommAbort communicator)) + (nccl/ncclCommDestroy communicator))) + (when (and stream (not (.isNull stream))) + (cudart/cudaStreamDestroy stream)) + nil) + +(defn abort! [{:keys [^ncclComm communicator]}] + (when (and communicator (not (.isNull communicator))) + (nccl/ncclCommAbort communicator)) + nil) + +(defn- enqueue! [{:keys [^ncclComm communicator ^CUstream_st stream]} + result operation] + (let [ready (CUevent_st.) + done (CUevent_st.)] + (try + (check-cuda! "cudaEventCreateWithFlags" + (cudart/cudaEventCreateWithFlags + ready cudart/cudaEventDisableTiming)) + (check-cuda! "cudaEventCreateWithFlags" + (cudart/cudaEventCreateWithFlags + done cudart/cudaEventDisableTiming)) + (check-cuda! "cudaEventRecord" (cudart/cudaEventRecord ready)) + (check-cuda! "cudaStreamWaitEvent" + (cudart/cudaStreamWaitEvent stream ready)) + (operation communicator stream) + (check-cuda! "cudaEventRecord" + (cudart/cudaEventRecord done stream)) + (check-cuda! "cudaStreamWaitEvent" (cudart/cudaStreamWaitEvent + (CUstream_st. (cast Pointer nil)) done)) + (cudart/cudaEventDestroy ready) + (->NcclWork done communicator result (atom false)) + (catch Throwable cause + (when-not (.isNull ready) (cudart/cudaEventDestroy ready)) + (when-not (.isNull done) (cudart/cudaEventDestroy done)) + (throw cause))))) + +(defn await! + "Waits for NCCL work with a bounded timeout and returns its result." + [^NcclWork work timeout-ms] + (when-not work + (throw (ex-info "Cannot await nil NCCL work" {}))) + (when-not @(:released? work) + (let [deadline (+ (System/nanoTime) (* (long timeout-ms) 1000000))] + (loop [] + (let [status (cudart/cudaEventQuery (:event work))] + (cond + (= status cudart/cudaSuccess) + (let [async-status (IntPointer. 1)] + (check-nccl! "ncclCommGetAsyncError" + (nccl/ncclCommGetAsyncError + (:communicator work) async-status)) + (check-nccl! "NCCL asynchronous operation" (.get async-status 0)) + (cudart/cudaEventDestroy (:event work)) + (reset! (:released? work) true)) + + (= status cudart/cudaErrorNotReady) + (if (< (System/nanoTime) deadline) + (do (Thread/sleep 1) (recur)) + (throw (java.util.concurrent.TimeoutException. + (str "NCCL operation exceeded " timeout-ms " ms")))) + + :else + (check-cuda! "cudaEventQuery" status)))))) + (:result work)) + +(defn all-reduce! [backend tensors-value operation] + (let [values (mapv validate-tensor! (tensors tensors-value))] + (enqueue! + backend tensors-value + (fn [communicator stream] + (check-nccl! "ncclGroupStart" (nccl/ncclGroupStart)) + (try + (doseq [^TensorBase tensor values] + (check-nccl! "ncclAllReduce" + (nccl/ncclAllReduce + (.data_ptr tensor) (.data_ptr tensor) (.numel tensor) + (scalar-type->nccl tensor) (reduction operation) + communicator stream))) + (finally + (check-nccl! "ncclGroupEnd" (nccl/ncclGroupEnd)))))))) + +(defn broadcast! [backend tensors-value root-rank] + (let [values (mapv validate-tensor! (tensors tensors-value))] + (enqueue! + backend tensors-value + (fn [communicator stream] + (check-nccl! "ncclGroupStart" (nccl/ncclGroupStart)) + (try + (doseq [^TensorBase tensor values] + (check-nccl! "ncclBroadcast" + (nccl/ncclBroadcast + (.data_ptr tensor) (.data_ptr tensor) (.numel tensor) + (scalar-type->nccl tensor) (int root-rank) + communicator stream))) + (finally + (check-nccl! "ncclGroupEnd" (nccl/ncclGroupEnd)))))))) + +(defn reduce! [backend tensors-value operation root-rank] + (let [values (mapv validate-tensor! (tensors tensors-value))] + (enqueue! + backend tensors-value + (fn [communicator stream] + (check-nccl! "ncclGroupStart" (nccl/ncclGroupStart)) + (try + (doseq [^TensorBase tensor values] + (check-nccl! "ncclReduce" + (nccl/ncclReduce + (.data_ptr tensor) (.data_ptr tensor) (.numel tensor) + (scalar-type->nccl tensor) (reduction operation) + (int root-rank) communicator stream))) + (finally + (check-nccl! "ncclGroupEnd" (nccl/ncclGroupEnd)))))))) + +(defn all-gather-into! [backend ^TensorBase output ^TensorBase input world-size] + (validate-tensor! output) + (validate-tensor! input) + (when-not (= (.numel output) (* world-size (.numel input))) + (throw (ex-info "All-gather output size must equal input size times world size" + {:output-elements (.numel output) + :input-elements (.numel input) + :world-size world-size}))) + (when-not (= (.toString (.scalar_type output)) + (.toString (.scalar_type input))) + (throw (ex-info "All-gather input and output dtypes must match" {}))) + (enqueue! + backend output + (fn [communicator stream] + (check-nccl! "ncclAllGather" + (nccl/ncclAllGather + (.data_ptr input) (.data_ptr output) (.numel input) + (scalar-type->nccl input) communicator stream))))) + +(defn reduce-scatter-into! + [backend ^TensorBase output ^TensorBase input world-size operation] + (validate-tensor! output) + (validate-tensor! input) + (when-not (= (.numel input) (* world-size (.numel output))) + (throw (ex-info "Reduce-scatter input size must equal output size times world size" + {:output-elements (.numel output) + :input-elements (.numel input) + :world-size world-size}))) + (when-not (= (.toString (.scalar_type output)) + (.toString (.scalar_type input))) + (throw (ex-info "Reduce-scatter input and output dtypes must match" {}))) + (enqueue! + backend output + (fn [communicator stream] + (check-nccl! "ncclReduceScatter" + (nccl/ncclReduceScatter + (.data_ptr input) (.data_ptr output) (.numel output) + (scalar-type->nccl input) (reduction operation) + communicator stream))))) + +(defn- byte-offset-pointer [^TensorBase tensor element-offset] + (doto (BytePointer. (.data_ptr tensor)) + (.position (* (long element-offset) (.element_size tensor))))) + +(defn all-to-all-single! + [backend ^TensorBase output ^TensorBase input world-size output-splits input-splits] + (validate-tensor! output) + (validate-tensor! input) + (when-not (= (.toString (.scalar_type output)) + (.toString (.scalar_type input))) + (throw (ex-info "All-to-all input and output dtypes must match" {}))) + (let [default-input-count (quot (.numel input) world-size) + default-output-count (quot (.numel output) world-size) + input-counts (if (seq input-splits) input-splits + (repeat world-size default-input-count)) + output-counts (if (seq output-splits) output-splits + (repeat world-size default-output-count)) + input-counts (mapv long input-counts) + output-counts (mapv long output-counts)] + (when-not (and (= world-size (count input-counts)) + (= world-size (count output-counts)) + (= (.numel input) (reduce + input-counts)) + (= (.numel output) (reduce + output-counts))) + (throw (ex-info "All-to-all split sizes must partition input and output" + {:world-size world-size + :input-splits input-counts + :output-splits output-counts}))) + (enqueue! + backend output + (fn [communicator stream] + (check-nccl! "ncclGroupStart" (nccl/ncclGroupStart)) + (try + (loop [peer 0 input-offset 0 output-offset 0] + (when (< peer world-size) + (check-nccl! "ncclSend" + (nccl/ncclSend + (byte-offset-pointer input input-offset) + (input-counts peer) (scalar-type->nccl input) peer + communicator stream)) + (check-nccl! "ncclRecv" + (nccl/ncclRecv + (byte-offset-pointer output output-offset) + (output-counts peer) (scalar-type->nccl output) peer + communicator stream)) + (recur (inc peer) + (+ input-offset (input-counts peer)) + (+ output-offset (output-counts peer))))) + (finally + (check-nccl! "ncclGroupEnd" (nccl/ncclGroupEnd)))))))) + +(defn send! [backend tensors-value destination tag] + (when-not (zero? tag) + (throw (ex-info "NCCL point-to-point operations do not support tags" + {:tag tag}))) + (let [values (mapv validate-tensor! (tensors tensors-value))] + (enqueue! + backend tensors-value + (fn [communicator stream] + (check-nccl! "ncclGroupStart" (nccl/ncclGroupStart)) + (try + (doseq [^TensorBase tensor values] + (check-nccl! "ncclSend" + (nccl/ncclSend + (.data_ptr tensor) (.numel tensor) + (scalar-type->nccl tensor) (int destination) + communicator stream))) + (finally + (check-nccl! "ncclGroupEnd" (nccl/ncclGroupEnd)))))))) + +(defn receive! [backend tensors-value source tag] + (when-not (zero? tag) + (throw (ex-info "NCCL point-to-point operations do not support tags" + {:tag tag}))) + (let [values (mapv validate-tensor! (tensors tensors-value))] + (enqueue! + backend tensors-value + (fn [communicator stream] + (check-nccl! "ncclGroupStart" (nccl/ncclGroupStart)) + (try + (doseq [^TensorBase tensor values] + (check-nccl! "ncclRecv" + (nccl/ncclRecv + (.data_ptr tensor) (.numel tensor) + (scalar-type->nccl tensor) (int source) + communicator stream))) + (finally + (check-nccl! "ncclGroupEnd" (nccl/ncclGroupEnd)))))))) diff --git a/src/clorch/distributed/worker.clj b/src/clorch/distributed/worker.clj new file mode 100644 index 0000000..a11759b --- /dev/null +++ b/src/clorch/distributed/worker.clj @@ -0,0 +1,45 @@ +(ns clorch.distributed.worker + (:require [clojure.edn :as edn] + [clojure.java.io :as io] + [clorch.cuda :as cuda] + [clorch.distributed :as dist])) + +(defn- ensure-java-25! [] + (let [feature (.feature (Runtime/version))] + (when (< feature 25) + (throw (ex-info "Distributed workers require Java 25 or newer" + {:java-feature feature}))))) + +(defn- resolve-entrypoint [entrypoint] + (let [namespace-symbol (symbol (namespace entrypoint))] + (require namespace-symbol) + (or (ns-resolve namespace-symbol (symbol (name entrypoint))) + (throw (ex-info "Distributed worker entrypoint does not exist" + {:main entrypoint}))))) + +(defn- worker-context [process-group args] + {:rank (dist/rank process-group) + :local-rank (dist/local-rank process-group) + :world-size (dist/world-size process-group) + :backend (dist/backend process-group) + :process-group process-group + :args args}) + +(defn -main [& [config-path]] + (ensure-java-25!) + (when-not config-path + (throw (ex-info "Distributed worker requires a configuration path" {}))) + (let [{:keys [main args]} (edn/read-string (slurp (io/file config-path))) + backend (keyword (or (System/getenv "CLORCH_DIST_BACKEND") "nccl")) + local-rank (Long/parseLong (or (System/getenv "LOCAL_RANK") "0"))] + (cuda/set-device! local-rank) + (try + (dist/with-process-group {:backend backend} + (let [process-group (dist/current-process-group) + entrypoint (resolve-entrypoint main) + result (entrypoint (worker-context process-group args))] + (println (pr-str {:rank (dist/rank process-group) + :status :completed + :result result})))) + (finally + (shutdown-agents))))) diff --git a/src/clorch/nn/functional.clj b/src/clorch/nn/functional.clj index 0b39874..bcd6a2e 100644 --- a/src/clorch/nn/functional.clj +++ b/src/clorch/nn/functional.clj @@ -1,6 +1,7 @@ (ns clorch.nn.functional (:require [clorch.torch :as t]) - (:import [org.bytedeco.pytorch.global torch])) + (:import [org.bytedeco.pytorch DoubleOptional TensorOptional] + [org.bytedeco.pytorch.global torch])) (defn relu [tensor] (torch/relu (t/->tensor tensor))) @@ -9,7 +10,7 @@ (torch/relu6 (t/->tensor tensor))) (defn leaky-relu [tensor & [negative-slope]] - (torch/leaky_relu (t/->tensor tensor) (double (or negative-slope 0.01)) false)) + (torch/leaky_relu (t/->tensor tensor) (org.bytedeco.pytorch.Scalar. (double (or negative-slope 0.01))))) (defn gelu [tensor] (torch/gelu (t/->tensor tensor))) @@ -41,13 +42,17 @@ (torch/hardsigmoid (t/->tensor tensor))) (defn elu [tensor & [alpha]] - (torch/elu (t/->tensor tensor) (double (or alpha 1.0)) false)) + (torch/elu (t/->tensor tensor) + (org.bytedeco.pytorch.Scalar. (double (or alpha 1.0))) + (org.bytedeco.pytorch.Scalar. 1.0) + (org.bytedeco.pytorch.Scalar. 1.0))) (defn selu [tensor] (torch/selu (t/->tensor tensor))) (defn celu [tensor & [alpha]] - (torch/celu (t/->tensor tensor) (double (or alpha 1.0)) false)) + (torch/celu (t/->tensor tensor) + (org.bytedeco.pytorch.Scalar. (double (or alpha 1.0))))) (defn softplus [tensor & [beta threshold]] (torch/softplus (t/->tensor tensor) @@ -58,7 +63,9 @@ (torch/softsign (t/->tensor tensor))) (defn softmin [tensor dim] - (torch/softmin (t/->tensor tensor) (long dim) (org.bytedeco.pytorch.ScalarTypeOptional.))) + (torch/softmax (torch/neg (t/->tensor tensor)) + (long dim) + (org.bytedeco.pytorch.ScalarTypeOptional.))) (defn log-softmax [tensor dim-or-opts] (let [dim (if (map? dim-or-opts) (:dim dim-or-opts) dim-or-opts)] @@ -74,13 +81,24 @@ (torch/tanhshrink (t/->tensor tensor))) (defn threshold [tensor threshold-val value & {:keys [inplace] :or {inplace false}}] - (torch/threshold (t/->tensor tensor) (double threshold-val) (double value) (boolean inplace))) + (let [tensor (t/->tensor tensor) + threshold-value (org.bytedeco.pytorch.Scalar. (double threshold-val)) + value (org.bytedeco.pytorch.Scalar. (double value))] + (if inplace + (torch/threshold_ tensor threshold-value value) + (torch/threshold tensor threshold-value value)))) (defn glu [tensor & [dim]] (torch/glu (t/->tensor tensor) (long (or dim -1)))) (defn rrelu [tensor & {:keys [lower upper training inplace] :or {lower 0.125 upper 0.3333333333333333 training false inplace false}}] - (torch/rrelu (t/->tensor tensor) (double lower) (double upper) (boolean training) (boolean inplace))) + (let [tensor (t/->tensor tensor) + lower (org.bytedeco.pytorch.Scalar. (double lower)) + upper (org.bytedeco.pytorch.Scalar. (double upper)) + generator (org.bytedeco.pytorch.GeneratorOptional.)] + (if inplace + (torch/rrelu_ tensor lower upper (boolean training) generator) + (torch/rrelu tensor lower upper (boolean training) generator)))) (defn prelu [tensor weight] (torch/prelu (t/->tensor tensor) (t/->tensor weight))) @@ -93,7 +111,34 @@ (torch/dropout (t/->tensor tensor) (double p) (boolean training?))) (defn dropout2d [tensor p & {:keys [training?] :or {training? true}}] - (torch/dropout2d (t/->tensor tensor) (double p) (boolean training?) false)) + (let [options (org.bytedeco.pytorch.DropoutFuncOptions.)] + (.put (.p options) (double p)) + (.put (.training options) (boolean training?)) + (.put (.inplace options) false) + (torch/dropout2d (t/->tensor tensor) options))) + +(defn scaled-dot-product-attention + "Runs LibTorch's fused scaled-dot-product attention dispatcher. + + CUDA automatically selects Flash Attention, memory-efficient attention, or + the math kernel according to dtype, shape, mask, and hardware support." + [query attention-key value & {:keys [attention-mask dropout-p causal? scale enable-gqa?] + :or {dropout-p 0.0 causal? false enable-gqa? false}}] + (with-open [mask-option (if attention-mask + (TensorOptional. (t/->tensor attention-mask)) + (TensorOptional.)) + scale-option (if scale + (DoubleOptional. (double scale)) + (DoubleOptional.))] + (torch/scaled_dot_product_attention + (t/->tensor query) + (t/->tensor attention-key) + (t/->tensor value) + mask-option + (double dropout-p) + (boolean causal?) + scale-option + (boolean enable-gqa?)))) (defn max-pool1d [tensor kernel & {:keys [stride padding dilation ceil-mode] :or {padding 0 dilation 1 ceil-mode false}}] (let [k (if (number? kernel) [kernel] kernel) diff --git a/src/clorch/nn/parallel.clj b/src/clorch/nn/parallel.clj new file mode 100644 index 0000000..3c02bec --- /dev/null +++ b/src/clorch/nn/parallel.clj @@ -0,0 +1,278 @@ +(ns clorch.nn.parallel + "Distributed neural-network wrappers backed by direct NCCL collectives." + (:require [clorch.distributed :as dist] + [clorch.nn :as nn]) + (:import [org.bytedeco.pytorch Module Scalar SizeTVector Tensor + TensorBase TensorTensorHook TensorVector] + [org.bytedeco.pytorch.global torch])) + +(def ^:dynamic *synchronize-gradients* + "False inside no-sync so DDP accumulates local gradients without collectives." + true) + +(defn- parameter-vector [model] + (let [parameters (nn/parameters model :trainable-only true)] + (when (zero? (.size parameters)) + (throw (ex-info "DistributedDataParallel requires trainable parameters" {}))) + parameters)) + +(defn- tensor-vector->vec [^TensorVector values] + (mapv #(Tensor. (.get values (long %))) (range (.size values)))) + +(defn- native-buffers [model] + (let [output (TensorVector.)] + (letfn [(walk [node] + (cond + (instance? Module node) + (let [buffers (.buffers ^Module node)] + (dotimes [index (.size buffers)] + (.push_back output (.get buffers (long index))))) + + (record? node) (run! walk (vals node)) + (map? node) (run! walk (vals node)) + (sequential? node) (run! walk node)))] + (walk model) + output))) + +(defn- validate-parameters! [context parameters] + (doseq [[index ^Tensor parameter] (map-indexed vector parameters)] + (when-not (.is_cuda parameter) + (throw (ex-info "DDP parameters must be on the rank-local CUDA device" + {:parameter-index index + :local-rank (:local-rank context)}))) + (when-not (= (:local-rank context) (long (.index (.device parameter)))) + (throw (ex-info "DDP parameter is on the wrong CUDA device" + {:parameter-index index + :parameter-device (long (.index (.device parameter))) + :local-rank (:local-rank context)})))) + parameters) + +(defn- parameter-signature [parameters] + (mapv (fn [^Tensor parameter] + {:shape (let [sizes (.sizes parameter)] + (mapv #(.get sizes (long %)) (range (.size sizes)))) + :dtype (.toString (.scalar_type parameter)) + :device (long (.index (.device parameter)))}) + parameters)) + +(defn- verify-parameters! [context parameters] + (let [operation-id (swap! (:operation-counter context) inc) + prefix (str "clorch/ddp/signature/" operation-id "/") + rank-key (str prefix (:rank context)) + signature (pr-str (parameter-signature parameters)) + store (:store context) + store-keys (org.bytedeco.pytorch.StringVector.)] + (.set store rank-key signature) + (dotimes [peer (:world-size context)] + (.push_back store-keys (str prefix peer))) + (._wait store store-keys + (org.bytedeco.javacpp.chrono.Milliseconds. (:timeout-ms context))) + (let [signatures + (mapv #(.get_to_str store (str prefix %)) + (range (:world-size context)))] + (when-not (apply = signatures) + (throw (ex-info "DDP parameter shapes, dtypes, and devices differ across ranks" + {:rank (:rank context)})))))) + +(defn- bucket-indices [^TensorVector parameters bucket-cap-bytes] + (let [limits (SizeTVector. + (long-array [(long (torch/kDefaultFirstBucketBytes)) + (long bucket-cap-bytes)])) + native-buckets (.get0 + (torch/compute_bucket_assignment_by_size + parameters limits))] + (mapv (fn [bucket-index] + (let [indices (.get native-buckets (long bucket-index))] + (mapv #(.get indices (long %)) (range (.size indices))))) + (range (.size native-buckets))))) + +(defn- make-bucket [parameters indices] + (let [native-parameters (TensorVector.)] + (doseq [index indices] + (.push_back native-parameters (parameters index))) + (let [buffer (torch/flatten_dense_tensors native-parameters) + views (torch/unflatten_dense_tensors buffer native-parameters)] + {:indices indices + :parameters native-parameters + :buffer buffer + :views views + :remaining (atom (count indices)) + :gradients (atom (vec (repeat (count indices) nil))) + :lock (Object.)}))) + +(defn- await-pending! [context pending] + (let [works (locking pending + (let [works @pending] + (reset! pending []) + works))] + (doseq [work works] + (dist/await! work (:timeout-ms context))))) + +(defn- reset-prior-gradient! [^Tensor parameter ^Tensor gradient] + (let [prior (.grad parameter)] + (when (.defined prior) + (.add_ gradient prior) + (.reset (.mutable_grad parameter))))) + +(defn- complete-bucket! + [context bucket pending] + (let [buffer (:buffer bucket) + work (dist/all-reduce! context buffer {:op :sum :async? true})] + (.div_ ^Tensor buffer (Scalar. (double (:world-size context)))) + (doseq [slot (range (count (:indices bucket)))] + (.copy_ ^Tensor (@(:gradients bucket) slot) + (.get ^TensorVector (:views bucket) (long slot)))) + (swap! pending conj work) + (reset! (:remaining bucket) (count (:indices bucket))) + (reset! (:gradients bucket) (vec (repeat (count (:indices bucket)) nil))))) + +(defn- gradient-hook + [context bucket slot ^Tensor parameter synchronize? pending classloader] + (proxy [TensorTensorHook] [] + (call [gradient-base] + ;; Native autograd workers do not inherit the application classloader. + (.setContextClassLoader (Thread/currentThread) classloader) + (when @synchronize? + (let [gradient (Tensor. ^TensorBase gradient-base) + bucket-lock (:lock bucket)] + (reset-prior-gradient! parameter gradient) + (locking bucket-lock + (.copy_ ^Tensor (.get ^TensorVector (:views bucket) (long slot)) + gradient) + (swap! (:gradients bucket) assoc slot gradient) + (when (zero? (swap! (:remaining bucket) dec)) + (complete-bucket! context bucket pending))))) + gradient-base))) + +(defn- install-gradient-hooks! + [context parameters buckets synchronize? pending] + (let [classloader (.getContextClassLoader (Thread/currentThread))] + (vec + (mapcat + (fn [bucket] + (mapv + (fn [slot parameter-index] + (let [parameter (parameters parameter-index) + hook (gradient-hook context bucket slot parameter + synchronize? pending classloader) + hook-id (.register_hook ^Tensor parameter hook)] + {:parameter parameter :hook hook :hook-id hook-id})) + (range (count (:indices bucket))) + (:indices bucket))) + buckets)))) + +(deftype DistributedDataParallel + [model context parameters buckets hooks buffers broadcast-buffers? + synchronize? pending closed?] + nn/IModule + (-forward [_ input] + (when @closed? + (throw (ex-info "DistributedDataParallel has been closed" {}))) + (await-pending! context pending) + (doseq [bucket buckets] + (when-not (= (count (:indices bucket)) @(:remaining bucket)) + (throw (ex-info "Previous backward pass did not produce every parameter gradient" + {:missing-gradient-count @(:remaining bucket) + :bucket-indices (:indices bucket)})))) + (when (and broadcast-buffers? (pos? (.size ^TensorVector buffers))) + (dist/broadcast! context buffers {:root-rank 0})) + (reset! synchronize? *synchronize-gradients*) + (nn/forward model input)) + (-train [this training?] + (nn/train model training?) + this) + (-to [_ _] + (throw (ex-info "Move a model to its rank-local device before wrapping it in DistributedDataParallel" + {}))) + clojure.lang.IFn + (invoke [this input] + (nn/forward this input)) + clojure.lang.ILookup + (valAt [this lookup-key] + (.valAt this lookup-key nil)) + (valAt [_ lookup-key not-found] + (case lookup-key + :model model + :context context + :parameters parameters + :buckets buckets + :buffers buffers + :broadcast-buffers? broadcast-buffers? + :closed? @closed? + not-found)) + java.lang.AutoCloseable + (close [_] + (when (compare-and-set! closed? false true) + (await-pending! context pending) + (doseq [{:keys [^Tensor parameter hook-id]} hooks] + (.remove_hook parameter hook-id))))) + +(defn distributed-data-parallel + "Wraps a rank-local model in synchronous NCCL data parallelism. + + Parameters are verified and broadcast from rank zero. Reusable gradient + buckets launch asynchronous NCCL all-reduce operations as backward produces + each bucket. Construct and move the model to its rank-local CUDA device first. + + Options: + - :process-group ProcessGroupContext; defaults to current group + - :bucket-cap-mb gradient bucket capacity, default 25 MiB + - :broadcast-buffers? synchronize module buffers before forward, default true" + [model & args] + (let [options (if (map? (first args)) (first args) (apply hash-map args)) + context (or (:process-group options) (dist/current-process-group)) + bucket-cap-mb (double (get options :bucket-cap-mb 25.0)) + broadcast-buffers? (get options :broadcast-buffers? true)] + (when-not (pos? bucket-cap-mb) + (throw (ex-info "DDP bucket capacity must be positive" + {:bucket-cap-mb bucket-cap-mb}))) + (when (get options :find-unused-parameters? false) + (throw (ex-info "Direct NCCL DDP requires every trainable parameter in each synchronized backward pass" + {:find-unused-parameters? true}))) + (when (get options :gradient-as-bucket-view? false) + (throw (ex-info "Direct NCCL DDP does not expose gradients as bucket views" + {:gradient-as-bucket-view? true}))) + (let [native-parameters (parameter-vector model) + parameters (validate-parameters! + context (tensor-vector->vec native-parameters)) + _ (verify-parameters! context parameters) + _ (dist/broadcast! context native-parameters {:root-rank 0}) + buckets (mapv #(make-bucket parameters %) + (bucket-indices native-parameters + (long (* bucket-cap-mb 1024.0 1024.0)))) + synchronize? (atom true) + pending (atom []) + hooks (install-gradient-hooks! context parameters buckets + synchronize? pending)] + (DistributedDataParallel. + model context parameters buckets hooks (native-buffers model) + broadcast-buffers? synchronize? pending (atom false))))) + +(defmacro no-sync + "Accumulates local gradients in body without NCCL synchronization." + [& body] + `(binding [*synchronize-gradients* false] + ~@body)) + +(defn synchronize! + "Waits for all gradient buckets launched by the most recent backward pass." + [ddp] + (await-pending! (:context ddp) (.-pending ^DistributedDataParallel ddp)) + ddp) + +(defn optimizer-step! + "Waits for DDP reductions, then steps an optimizer directly or through AMP." + ([ddp optimizer] + (optimizer-step! ddp optimizer {})) + ([ddp optimizer {:keys [scaler]}] + (synchronize! ddp) + (if scaler + ((requiring-resolve 'clorch.amp/step!) scaler optimizer) + (do + (.step optimizer) + true)))) + +(defn reducer-statistics + "Returns bucket element counts for capacity and profiling diagnostics." + [ddp] + (mapv #(long (.numel ^Tensor (:buffer %))) (:buckets ddp))) diff --git a/src/clorch/torch.clj b/src/clorch/torch.clj index d98f9e5..0eb4752 100644 --- a/src/clorch/torch.clj +++ b/src/clorch/torch.clj @@ -148,9 +148,8 @@ :Char {:kw :int8 :bytes 1} :Byte {:kw :uint8 :bytes 1} :Bool {:kw :uint8 :bytes 1} - ;; Fallbacks for other types if they are convertible or viewable - :Half {:kw :float32 :bytes 2} - :BFloat16 {:kw :float32 :bytes 2} + :Half {:kw :float16 :bytes 2} + :BFloat16 {:kw :bfloat16 :bytes 2} :ComplexFloat {:kw :float32 :bytes 8} :ComplexDouble {:kw :float64 :bytes 16} :QInt8 {:kw :int8 :bytes 1} @@ -1025,7 +1024,13 @@ ```clojure (add (ones [2]) (ones [2])) ```" - [a b] (torch/add (->tensor a) (->tensor b))) + [a b] + (cond + (number? b) (with-open [scalar (Scalar. (clojure.core/double b))] + (torch/add (->tensor a) scalar)) + (number? a) (with-open [scalar (Scalar. (clojure.core/double a))] + (torch/add scalar (->tensor b))) + :else (torch/add (->tensor a) (->tensor b)))) (defn sub "Element-wise subtraction of two tensors or a tensor and a scalar. @@ -1040,7 +1045,13 @@ ```clojure (sub (ones [2]) 0.5) ```" - [a b] (torch/sub (->tensor a) (->tensor b))) + [a b] + (cond + (number? b) (with-open [scalar (Scalar. (clojure.core/double b))] + (torch/sub (->tensor a) scalar)) + (number? a) (with-open [scalar (Scalar. (clojure.core/double a))] + (torch/subtract scalar (->tensor b))) + :else (torch/sub (->tensor a) (->tensor b)))) (defn mul "Element-wise multiplication of two tensors or a tensor and a scalar. @@ -1055,7 +1066,13 @@ ```clojure (mul (ones [2]) 2.0) ```" - [a b] (torch/mul (->tensor a) (->tensor b))) + [a b] + (cond + (number? b) (with-open [scalar (Scalar. (clojure.core/double b))] + (torch/mul (->tensor a) scalar)) + (number? a) (with-open [scalar (Scalar. (clojure.core/double a))] + (torch/multiply scalar (->tensor b))) + :else (torch/mul (->tensor a) (->tensor b)))) (defn div "Element-wise division of two tensors or a tensor and a scalar. @@ -1070,7 +1087,13 @@ ```clojure (div (ones [2]) 2.0) ```" - [a b] (torch/div (->tensor a) (->tensor b))) + [a b] + (cond + (number? b) (with-open [scalar (Scalar. (clojure.core/double b))] + (torch/div (->tensor a) scalar)) + (number? a) (with-open [scalar (Scalar. (clojure.core/double a))] + (torch/divide scalar (->tensor b))) + :else (torch/div (->tensor a) (->tensor b)))) (defn lerp "Does a linear interpolation of two tensors based on a weight. @@ -3600,6 +3623,13 @@ (.save ^org.bytedeco.pytorch.Module thing sub-archive) (.write archive (clojure.core/str prefix) sub-archive))) + (instance? org.bytedeco.pytorch.Optimizer thing) + (if (empty? prefix) + (.save ^org.bytedeco.pytorch.Optimizer thing archive) + (with-open [sub-archive (org.bytedeco.pytorch.OutputArchive.)] + (.save ^org.bytedeco.pytorch.Optimizer thing sub-archive) + (.write archive (clojure.core/str prefix) sub-archive))) + (map? thing) (doseq [[k v] thing] (let [k-str (if (keyword? k) (name k) (str k))] @@ -3622,6 +3652,14 @@ (.load ^org.bytedeco.pytorch.Module thing sub-archive) thing)) + (instance? org.bytedeco.pytorch.Optimizer thing) + (if (empty? prefix) + (do (.load ^org.bytedeco.pytorch.Optimizer thing archive) thing) + (with-open [sub-archive (org.bytedeco.pytorch.InputArchive.)] + (.read archive (clojure.core/str prefix) sub-archive) + (.load ^org.bytedeco.pytorch.Optimizer thing sub-archive) + thing)) + (map? thing) (do (doseq [[k v] thing] (let [k-str (if (keyword? k) (name k) (str k))] diff --git a/test/clorch/amp_test.clj b/test/clorch/amp_test.clj new file mode 100644 index 0000000..5f2742b --- /dev/null +++ b/test/clorch/amp_test.clj @@ -0,0 +1,39 @@ +(ns clorch.amp-test + (:require [clojure.test :refer [deftest is testing]] + [clorch.amp :as amp] + [clorch.torch :as t]) + (:import [org.bytedeco.pytorch Tensor] + [org.bytedeco.pytorch.global torch torch$DeviceType])) + +(deftest cpu-autocast-state-test + (let [before (torch/is_autocast_enabled torch$DeviceType/CPU) + output (amp/autocast {:device :cpu :dtype :bfloat16} + (t/matmul (t/randn [8 8]) (t/randn [8 8])))] + (is (= "BFloat16" (.toString (.scalar_type ^Tensor output)))) + (is (= before (torch/is_autocast_enabled torch$DeviceType/CPU))))) + +(deftest autocast-restores-state-after-failure-test + (let [before (torch/is_autocast_enabled torch$DeviceType/CPU)] + (is (thrown-with-msg? Exception #"body failed" + (amp/autocast {:device :cpu :dtype :bfloat16} + (throw (Exception. "body failed"))))) + (is (= before (torch/is_autocast_enabled torch$DeviceType/CPU))))) + +(deftest gradient-scaler-state-test + (let [scaler (amp/grad-scaler {:initial-scale 1024.0})] + (is (= 1024.0 (amp/current-scale scaler))) + (amp/load-scaler-state! scaler {:scale 64.0 :growth-tracker 7}) + (is (= {:scale 64.0 :growth-tracker 7} + (amp/scaler-state scaler))) + (is (thrown? clojure.lang.ExceptionInfo + (amp/load-scaler-state! scaler + {:scale -1.0 :growth-tracker 0}))))) + +(deftest gradient-scaler-validation-test + (testing "invalid dynamic scaling factors fail at construction" + (is (thrown? clojure.lang.ExceptionInfo + (amp/grad-scaler {:initial-scale 0.0}))) + (is (thrown? clojure.lang.ExceptionInfo + (amp/grad-scaler {:growth-factor 1.0}))) + (is (thrown? clojure.lang.ExceptionInfo + (amp/grad-scaler {:backoff-factor 1.0}))))) diff --git a/test/clorch/cuda_smoke_test.clj b/test/clorch/cuda_smoke_test.clj index 10c834b..9c90973 100644 --- a/test/clorch/cuda_smoke_test.clj +++ b/test/clorch/cuda_smoke_test.clj @@ -1,48 +1,117 @@ (ns clorch.cuda-smoke-test - (:require [clorch.torch :as torch] - [clorch.cuda :as cuda]) - (:import [org.bytedeco.javacpp Loader])) + (:require [clojure.java.io :as io] + [clorch.amp :as amp] + [clorch.autograd :as autograd] + [clorch.cuda :as cuda] + [clorch.data :as data] + [clorch.distributed :as dist] + [clorch.nn :as nn] + [clorch.nn.functional :as F] + [clorch.nn.parallel :as ddp] + [clorch.optim :as optim] + [clorch.torch :as torch]) + (:import [java.net ServerSocket] + [org.bytedeco.javacpp Loader] + [org.bytedeco.pytorch Tensor])) (defn preload! [] - (println "Preloading native libraries...") - (doseq [cls ["org.bytedeco.openblas.global.openblas_nolapack" - "org.bytedeco.cuda.global.nvrtc" - "org.bytedeco.pytorch.global.torch"]] - (try - (println " Loading" cls) - (Loader/load (Class/forName cls)) - (catch Throwable e - (println " Failed to load" cls ":" (.getMessage e)))))) + (doseq [class-name ["org.bytedeco.openblas.global.openblas_nolapack" + "org.bytedeco.cuda.global.nvrtc" + "org.bytedeco.pytorch.global.torch"]] + (Loader/load (Class/forName class-name)))) + +(defn- check! [condition message data] + (when-not condition + (throw (ex-info message data)))) + +(defn- available-port [] + (with-open [socket (ServerSocket. 0)] + (.getLocalPort socket))) + +(defn- parameters [model] + (let [values (nn/parameters model)] + (mapv #(.get values (long %)) (range (.size values))))) (defn run-test [] (preload!) - (println "\n=== Clorch CUDA Smoke Test ===") - (try - (let [available? (cuda/available?)] - (println "CUDA Available:" available?) - (if available? - (try - (println "Device Count:" (cuda/device-count)) - (let [t (torch/tensor [1.0 2.0 3.0] {:device "cuda:0"})] - (println "Tensor created on GPU:" t) - (println "Tensor device:" (.device (torch/->tensor t))) - (let [t-cpu (torch/to t :cpu)] - (println "Moved back to CPU:" t-cpu)) - (println "SUCCESS: CUDA backend is working.")) - (catch Exception e - (println "FAILURE: CUDA reported available but operation failed:") - (println e) - (when (instance? Throwable e) - (.printStackTrace e)))) - (println "SKIPPED: No CUDA hardware detected. Logic is verified, but cannot run on this machine."))) - (catch Exception e - (println "FATAL ERROR during CUDA check:") - (println e) - (.printStackTrace e)))) + (if-not (cuda/available?) + (do + (println "SKIPPED: CUDA hardware is unavailable.") + :skipped) + (do + (cuda/set-device! 0) + (let [x (torch/ones [2 4] {:device :cuda})] + (check! (.is_cuda ^Tensor (torch/add x 2.0)) + "CUDA scalar arithmetic returned a CPU tensor" {}) + (check! (= :float16 + (amp/autocast {:device :cuda :dtype :float16} + (torch/dtype + (torch/matmul x (torch/transpose x 0 1))))) + "CUDA autocast did not select float16" {})) + (let [q (torch/randn [1 2 8 16] {:device :cuda :dtype :float16}) + output (F/scaled-dot-product-attention q q q :causal? true)] + (check! (= [1 2 8 16] (torch/size output)) + "Fused attention returned the wrong shape" + {:shape (torch/size output)})) + (dist/with-process-group + {:backend :nccl + :rank 0 + :local-rank 0 + :world-size 1 + :master-port (available-port) + :timeout-ms 30000} + (let [model (nn/to (nn/linear 4 2) :cuda) + optimizer (optim/adam (nn/parameters model)) + scaler (amp/grad-scaler {:initial-scale 128.0}) + sampler (data/distributed-sampler 16 {:num-replicas 1 :rank 0}) + checkpoint (str (System/getProperty "java.io.tmpdir") + "/clorch-cuda-smoke.pt")] + (with-open [parallel-model + (ddp/distributed-data-parallel + model {:process-group (dist/current-process-group)})] + ;; DDP must own independent tensor handles after constructor locals collect. + (System/gc) + (optim/zero-grad optimizer) + (let [input (torch/randn [8 4] {:device :cuda}) + target (torch/zeros [8 2] {:device :cuda}) + loss (amp/autocast {:device :cuda :dtype :float16} + (F/mse-loss (nn/forward parallel-model input) target))] + (amp/backward! scaler loss) + (check! (ddp/optimizer-step! + parallel-model optimizer {:scaler scaler}) + "Finite AMP step was skipped" {}))) + (data/set-epoch! sampler 3) + (let [before (mapv #(.clone ^Tensor %) (parameters model))] + (dist/save-checkpoint! + checkpoint {:model model + :optimizer optimizer + :sampler sampler + :scaler scaler + :state {:step 1}}) + (autograd/no-grad + (doseq [^Tensor parameter (parameters model)] + (.zero_ parameter))) + (data/set-epoch! sampler 0) + (check! (= {:step 1} + (dist/load-checkpoint! + checkpoint {:model model + :optimizer optimizer + :sampler sampler + :scaler scaler})) + "Checkpoint training state did not round-trip" {}) + (check! (= 3 (:epoch sampler)) + "Sampler epoch did not restore" {:epoch (:epoch sampler)}) + (check! (every? true? + (map #(.equal ^Tensor %1 ^Tensor %2) + before (parameters model))) + "Model weights did not restore" {}) + (io/delete-file checkpoint true) + (io/delete-file (str checkpoint ".edn") true)))) + (println "CUDA, NCCL, DDP, AMP, fused attention, and checkpoint smoke passed.") + true))) -(defn -main [& args] +(defn -main [& _] (run-test)) -;; Allow running as script too (when (= (System/getProperty "clojure.main.filename") *file*) (run-test)) diff --git a/test/clorch/distributed_test.clj b/test/clorch/distributed_test.clj new file mode 100644 index 0000000..83b1361 --- /dev/null +++ b/test/clorch/distributed_test.clj @@ -0,0 +1,66 @@ +(ns clorch.distributed-test + (:require [clojure.set :as set] + [clojure.test :refer [deftest is]] + [clorch.data :as data])) + +(defn process-dataset [n] + (data/dataset + :size (constantly n) + :get-item (fn [index] + {:data (long index) + :target (long (* index index))}) + :process-spec {:factory 'clorch.distributed-test/process-dataset + :args [n]})) + +(deftest distributed-sampler-test + (let [rank-zero (data/distributed-sampler + 10 {:num-replicas 2 :rank 0 :seed 41}) + rank-one (data/distributed-sampler + 10 {:num-replicas 2 :rank 1 :seed 41}) + epoch-zero-a (data/sample-indices rank-zero) + epoch-zero-b (data/sample-indices rank-one)] + (is (= 5 (count epoch-zero-a))) + (is (= 5 (count epoch-zero-b))) + (is (= (set (range 10)) (set (concat epoch-zero-a epoch-zero-b)))) + (is (empty? (set/intersection + (set epoch-zero-a) (set epoch-zero-b)))) + (data/set-epoch! rank-zero 1) + (data/set-epoch! rank-one 1) + (is (= (set (range 10)) + (set (concat (data/sample-indices rank-zero) + (data/sample-indices rank-one))))) + (is (not= epoch-zero-a (data/sample-indices rank-zero))))) + +(deftest sampler-checkpoint-test + (let [sampler (data/distributed-sampler + 12 {:num-replicas 3 :rank 1 :seed 99})] + (data/set-epoch! sampler 8) + (let [state (data/sampler-state sampler)] + (data/set-epoch! sampler 0) + (data/load-sampler-state! sampler state) + (is (= 8 (:epoch sampler))) + (is (= state (data/sampler-state sampler)))))) + +(deftest sampler-and-shuffle-conflict-test + (let [sampler (data/distributed-sampler + 8 {:num-replicas 1 :rank 0})] + (is (thrown-with-msg? + clojure.lang.ExceptionInfo + #"cannot also shuffle" + (data/dataloader (process-dataset 8) + {:sampler sampler :shuffle? true}))))) + +(deftest process-dataloader-preserves-order-test + (let [loader (data/dataloader + (process-dataset 10) + {:batch-size 3 + :shuffle? false + :num-workers 2 + :prefetch-factor 2 + :worker-backend :process + :timeout-ms 30000}) + batches (vec loader)] + (is (= [[0 1 2] [3 4 5] [6 7 8] [9]] + (mapv :data batches))) + (is (= [[0 1 4] [9 16 25] [36 49 64] [81]] + (mapv :target batches))))) diff --git a/test/clorch/nn_functional_test.clj b/test/clorch/nn_functional_test.clj index 9db61be..892a627 100644 --- a/test/clorch/nn_functional_test.clj +++ b/test/clorch/nn_functional_test.clj @@ -88,3 +88,16 @@ (is (near? (t/item-float (F/cosine-similarity x1 x2 :dim 0)) 0.0))) (testing "Pairwise Distance" (is (near? (t/item-float (F/pairwise-distance x1 x2)) (Math/sqrt 2.0))))))) + +(deftest scaled-dot-product-attention-test + (t/with-torch + (let [query (t/randn [1 2 4 8]) + attention-key (t/randn [1 2 4 8]) + value (t/randn [1 2 4 8]) + actual (F/scaled-dot-product-attention query attention-key value) + scores (t/mul + (t/matmul query (t/transpose attention-key -2 -1)) + (/ 1.0 (Math/sqrt 8.0))) + expected (t/matmul (F/softmax scores -1) value)] + (is (= [1 2 4 8] (t/size actual))) + (is (t/allclose actual expected {:rtol 1e-4 :atol 1e-5}))))) diff --git a/test/clorch/torch_test.clj b/test/clorch/torch_test.clj index 50db9a1..96aaafd 100644 --- a/test/clorch/torch_test.clj +++ b/test/clorch/torch_test.clj @@ -96,7 +96,9 @@ [:float16 "Half"] [:bfloat16 "BFloat16"] [:complex64 "ComplexFloat"] [:complex128 "ComplexDouble"]]] (let [t (torch/ones [2] {:dtype kw})] - (is (= expected (.toString (.scalar_type t))))))))) + (is (= expected (.toString (.scalar_type t)))))) + (is (= :float16 (torch/dtype (torch/ones [2] {:dtype :float16})))) + (is (= :bfloat16 (torch/dtype (torch/ones [2] {:dtype :bfloat16}))))))) (deftest float64-tensor-construction-test (testing "Float64 tensors preserve dtype and values" From 439cb4e047b3e11e459fb0b94e35424df8f6ded0 Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Sun, 9 Aug 2026 02:28:54 +0200 Subject: [PATCH 2/7] fix(pytorch): adapt scalar shrink bindings --- src/clorch/nn/functional.clj | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/clorch/nn/functional.clj b/src/clorch/nn/functional.clj index bcd6a2e..e6dd1cf 100644 --- a/src/clorch/nn/functional.clj +++ b/src/clorch/nn/functional.clj @@ -72,10 +72,12 @@ (torch/log_softmax (t/->tensor tensor) (long dim) (org.bytedeco.pytorch.ScalarTypeOptional.)))) (defn hardshrink [tensor & [lambd]] - (torch/hardshrink (t/->tensor tensor) (double (or lambd 0.5)))) + (with-open [threshold (org.bytedeco.pytorch.Scalar. (double (or lambd 0.5)))] + (torch/hardshrink (t/->tensor tensor) threshold))) (defn softshrink [tensor & [lambd]] - (torch/softshrink (t/->tensor tensor) (double (or lambd 0.5)))) + (with-open [threshold (org.bytedeco.pytorch.Scalar. (double (or lambd 0.5)))] + (torch/softshrink (t/->tensor tensor) threshold))) (defn tanhshrink [tensor] (torch/tanhshrink (t/->tensor tensor))) From 70412cb5c138cd55e83da0c188ca3228dbd15441 Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Sun, 9 Aug 2026 02:28:59 +0200 Subject: [PATCH 3/7] test(distributed): cover worker failure cleanup --- test/clorch/cuda_smoke_test.clj | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/test/clorch/cuda_smoke_test.clj b/test/clorch/cuda_smoke_test.clj index 9c90973..9e7c18f 100644 --- a/test/clorch/cuda_smoke_test.clj +++ b/test/clorch/cuda_smoke_test.clj @@ -1,5 +1,6 @@ (ns clorch.cuda-smoke-test (:require [clojure.java.io :as io] + [clojure.string :as str] [clorch.amp :as amp] [clorch.autograd :as autograd] [clorch.cuda :as cuda] @@ -32,6 +33,28 @@ (let [values (nn/parameters model)] (mapv #(.get values (long %)) (range (.size values))))) +(defn- verify-launcher-failure! [] + (let [job (dist/launch! {:nproc-per-node 1 + :devices [0] + :main 'clorch.cuda-smoke-test/missing-worker + :args {}}) + failure (try + (dist/await-job! job 120000) + nil + (catch clojure.lang.ExceptionInfo exception + (ex-data exception))) + status (dist/job-status job) + logs (dist/job-logs job 0)] + (check! (= :failed (:state failure)) + "Launcher did not propagate worker failure" + {:failure failure}) + (check! (every? (complement :alive?) (:workers status)) + "Launcher left a failed worker alive" + {:status status}) + (check! (str/includes? (:stderr logs) "does not exist") + "Worker logs omitted the failure diagnostic" + {:logs logs}))) + (defn run-test [] (preload!) (if-not (cuda/available?) @@ -107,6 +130,7 @@ "Model weights did not restore" {}) (io/delete-file checkpoint true) (io/delete-file (str checkpoint ".edn") true)))) + (verify-launcher-failure!) (println "CUDA, NCCL, DDP, AMP, fused attention, and checkpoint smoke passed.") true))) From 099de04abaf78901b51582c77d490d10abd21f94 Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Sun, 9 Aug 2026 02:29:05 +0200 Subject: [PATCH 4/7] docs(distributed): add CUDA training guide --- PORTING_STATUS.md | 41 +++++++-- README.md | 22 +++-- docs/distributed.md | 188 +++++++++++++++++++++++++++++++++++++++++ docs/index.md | 5 +- docs/llms.txt | 8 +- docs/mkdocs.yml | 1 + docs/pytorch-parity.md | 28 +++--- 7 files changed, 265 insertions(+), 28 deletions(-) create mode 100644 docs/distributed.md diff --git a/PORTING_STATUS.md b/PORTING_STATUS.md index 6eb1f06..e1e84be 100644 --- a/PORTING_STATUS.md +++ b/PORTING_STATUS.md @@ -174,6 +174,8 @@ This document tracks what's been ported from PyTorch to Clorch, and what still n | - `bce-loss`, `bce-with-logits-loss` | [x] | | **Similarity** | [x] | | - `cosine-similarity`, `pairwise-distance` | [x] | +| **Fused Attention** | [x] | +| - `scaled-dot-product-attention` | [x] | --- @@ -203,6 +205,8 @@ This document tracks what's been ported from PyTorch to Clorch, and what still n | `IDataset` protocol, `dataset`, `defdataset` | [x] | | `dataloader` (Thread & Process workers) | [x] | | `tensor-dataset` | [x] | +| `distributed-sampler`, `sample-indices`, `set-epoch!` | [x] | +| `sampler-state`, `load-sampler-state!` | [x] | --- @@ -210,10 +214,35 @@ This document tracks what's been ported from PyTorch to Clorch, and what still n | Feature | Status | |---------|--------| -| `available?` | [x] | +| `available?`, `device-count` | [x] | +| `set-device!`, `current-device` | [x] | +| `synchronize` | [x] | +| `manual-seed`, `manual-seed-all` | [x] | --- +### Automatic Mixed Precision (`clorch.amp`) + +| Feature | Status | +|---------|--------| +| `autocast`, `call-with-autocast` | [x] | +| `grad-scaler`, `scale-loss`, `backward!`, `step!` | [x] | +| `scaler-state`, `load-scaler-state!` | [x] | + +--- + +### Distributed Training (`clorch.distributed`, `clorch.nn.parallel`) + +| Feature | Status | +|---------|--------| +| NCCL process-group lifecycle and rank metadata | [x] | +| All-reduce, broadcast, reduce, gather, reduce-scatter, all-to-all | [x] | +| Point-to-point send/receive and barriers | [x] | +| Managed local worker launch, status, logs, cancellation, failure propagation | [x] | +| `DistributedDataParallel`, bucketed gradient synchronization, `no-sync` | [x] | +| Rank-zero model, optimizer, sampler, scaler, CPU RNG, and training checkpoints | [x] | +--- + ## What's Not Yet Implemented ### Specialized Tensors & Ops @@ -262,10 +291,12 @@ This document tracks what's been ported from PyTorch to Clorch, and what still n | Autograd | 5 | ~5 | | NN Layers | ~80 | ~20 | | Optimizers | 7 | ~10 | -| Data | 5 | ~5 | -| CUDA | 1 | ~10 | -| **Total** | **~268** | **~60** | +| Data | ~8 | ~4 | +| CUDA | ~7 | ~8 | +| AMP | ~6 | ~4 | +| Distributed | ~20 | ~10 | +| **Total** | **Pending version-pinned recount** | **Pending version-pinned recount** | --- -*Last updated: 2026-03-13* +*Last updated: 2026-08-08* diff --git a/README.md b/README.md index 9eb59c4..d20b0ce 100644 --- a/README.md +++ b/README.md @@ -10,13 +10,14 @@ MIT license

-Clorch is a Clojure deep-learning library backed by LibTorch. It provides PyTorch-style tensors, autograd, neural-network modules, optimizers, data loading, and explicit CPU/CUDA device placement through a REPL-friendly API. +Clorch is a Clojure deep-learning library backed by LibTorch. It provides PyTorch-style tensors, autograd, neural-network modules, optimizers, data loading, explicit CPU/CUDA placement, and NCCL distributed training through a REPL-friendly API. ## Highlights -- **Tensor and training APIs:** tensor operations, autograd, losses, optimizers, data loaders, state dictionaries, and native-memory scopes +- **Tensor and training APIs:** tensor operations, autograd, losses, optimizers, data loaders, state dictionaries, AMP, and native-memory scopes - **Model building:** standard layers, custom `nn/defmodel` modules, architecture summaries, and checkpoint loading -- **LLM components:** RMSNorm, RoPE, grouped-query attention, SwiGLU, causal masks, KV caches, and autoregressive generation +- **Distributed CUDA:** NCCL collectives, managed rank processes, distributed sampling, synchronous DDP, gradient accumulation, and rank-zero checkpoints +- **LLM components:** RMSNorm, RoPE, grouped-query attention, SwiGLU, fused scaled-dot-product attention, causal masks, KV caches, and autoregressive generation - **NanoChat-inspired example:** a compact, single-device Llama-style training and chat demo with checkpointing and optional KV caching, inspired by [Karpathy's NanoChat](https://github.com/karpathy/nanochat) - **CPU and CUDA:** native LibTorch execution with explicit tensor and model placement @@ -151,6 +152,7 @@ Full documentation: **[antlobach.github.io/clorch](https://antlobach.github.io/c | [Neural networks](docs/nn.md) | Modules, custom models, layers, and summaries | | [Functional API](docs/functional.md) | Stateless layers, activations, pooling, and losses | | [Optimizers](docs/optimizers.md) | SGD, Adam, AdamW, RMSprop, and Adagrad | +| [Distributed training](docs/distributed.md) | NCCL collectives, worker launch, DDP, AMP, sampling, and checkpoints | | [Memory management](docs/memory.md) | Native allocation scopes and pointer lifecycles | | [PyTorch parity](docs/pytorch-parity.md) | Measured coverage, current gaps, and roadmap | @@ -164,6 +166,9 @@ Clorch follows PyTorch concepts, but it does not expose every PyTorch symbol. Ch |---|---| | `torch` | `clorch.torch` | | `torch.cuda` | `clorch.cuda` | +| `torch.amp` | `clorch.amp` | +| `torch.distributed` | `clorch.distributed` | +| `DistributedDataParallel` | `clorch.nn.parallel` | | `torch.autograd` | `clorch.autograd` | | `torch.nn` | `clorch.nn` | | `torch.nn.functional` | `clorch.nn.functional` | @@ -178,7 +183,7 @@ When generating Clorch code: - Keep long-lived models and optimizers outside `with-torch`; use one scope per allocating batch or generation step. - Return a JVM scalar or `nil` from `with-torch` unless a tensor must escape. - Reuse patterns from the [examples](examples/) instead of inventing wrapper APIs. -- Check [PyTorch Parity](docs/pytorch-parity.md) before assuming support for distributed training, AMP, Flash Attention, or `torch.compile`. +- Read the [Distributed CUDA Training](docs/distributed.md) guide before generating multi-process training code. ## Examples @@ -187,10 +192,11 @@ When generating Clorch code: - [Synthetic training](examples/synthetic.clj): end-to-end model training - [Modern Llama](examples/modern_llama.clj): RoPE, GQA, SwiGLU, and KV caching - [Compact Llama chat](examples/nanochat.clj): small-corpus training, checkpointing, and token generation +- [Distributed CUDA training](examples/distributed_training.clj): NCCL workers, DDP, AMP, accumulation, sampling, and checkpoints ## Project status -The cross-language comparison suite currently passes 40 of 40 numerical scenarios. The repository's tracked feature catalog marks 268 of approximately 328 capabilities as implemented, or about 81.7%. That catalog does not enumerate the entire upstream PyTorch API. +The cross-language comparison suite currently passes 40 of 40 numerical scenarios. The tracked feature catalog predates the PyTorch 2.10 distributed-training milestone and needs a version-pinned recount before Clorch publishes another breadth percentage. Read [PyTorch Parity](docs/pytorch-parity.md) for the capability table and roadmap to 100% of a version-pinned target. @@ -208,6 +214,12 @@ Run the CPU release suite: clojure -Sthreads 1 -M -m clorch.release-check --mode cpu ``` +Run the GPU release suite on a Linux NVIDIA host: + +```bash +clojure -Sthreads 1 -M -m clorch.release-check --mode gpu +``` + Run all shipped examples: ```bash diff --git a/docs/distributed.md b/docs/distributed.md new file mode 100644 index 0000000..36af83b --- /dev/null +++ b/docs/distributed.md @@ -0,0 +1,188 @@ +# Distributed CUDA Training + +Clorch runs one JVM per CUDA device and coordinates ranks with NCCL. The public API covers process-group lifecycle, collectives, local worker launch, distributed sampling, synchronous data parallelism, AMP, and rank-zero checkpoints. + +## Requirements + +- Linux with NVIDIA GPUs and a working driver +- Java 25 or newer for the managed launcher +- One distinct CUDA device per local rank +- GPU-enabled JavaCPP PyTorch and CUDA native libraries on the classpath + +Clorch currently supports the `:nccl` backend. It does not provide Gloo, RPC, FSDP, tensor parallelism, or elastic membership. + +## Run the training example + +The shipped example launches one worker JVM for each device: + +```clojure +(require '[distributed-training :as training]) + +(training/run-local! + [0 1] + {:epochs 4 + :sample-count 1024 + :batch-size 32 + :accumulation 2 + :precision :bfloat16 + :checkpoint-path "/tmp/clorch-ddp.pt"}) +``` + +`run-local!` waits for every rank and returns job status plus file-backed stdout and stderr. Use `:float16` to enable dynamic loss scaling. Bfloat16 uses autocast without a scaler. + +## Worker entrypoints + +`dist/launch!` accepts a namespace-qualified worker function and EDN arguments: + +```clojure +(require '[clorch.distributed :as dist]) + +(def job + (dist/launch! + {:nproc-per-node 2 + :devices [0 1] + :main 'my.training/train-worker + :args {:epochs 10} + :timeout-ms 300000})) + +(dist/job-status job) +(dist/await-job! job) +(dist/job-logs job 0) +``` + +Clorch invokes the worker with this map: + +```clojure +{:rank 0 + :local-rank 0 + :world-size 2 + :backend :nccl + :process-group context + :args {:epochs 10}} +``` + +The launcher sets `RANK`, `LOCAL_RANK`, `WORLD_SIZE`, `LOCAL_WORLD_SIZE`, `MASTER_ADDR`, `MASTER_PORT`, and `CUDA_VISIBLE_DEVICES` for each child. If one rank exits with an error, the launcher terminates the remaining ranks and preserves each rank's logs on disk. Call `dist/stop-job!` to cancel a running job. + +## Process groups and collectives + +Workers launched by Clorch receive an initialized process group. For a custom launcher, initialize from the standard environment variables: + +```clojure +(require '[clorch.distributed :as dist]) + +(dist/with-process-group {:backend :nccl} + (dist/all-reduce! tensor {:op :sum}) + (dist/barrier!)) +``` + +Available collective operations: + +- `all-reduce!`, `broadcast!`, and `reduce!` +- `all-gather-into!` and `reduce-scatter-into!` +- `all-to-all-single!` +- `send` and `receive!` +- `barrier!` + +Collectives mutate CUDA tensors in place. Every rank must call collectives in the same order with compatible shapes, dtypes, and split sizes. Pass `:async? true` to receive a work handle, then call `dist/await!` before reading the result. + +## DistributedDataParallel + +Create the model on the rank-local CUDA device before wrapping it: + +```clojure +(require '[clorch.nn :as nn] + '[clorch.nn.parallel :as ddp] + '[clorch.optim :as optim]) + +(def model (nn/to (nn/linear 128 32) :cuda)) +(def optimizer (optim/adamw (nn/parameters model) :lr 3e-4)) + +(with-open [parallel-model + (ddp/distributed-data-parallel + model {:bucket-cap-mb 25.0 + :broadcast-buffers? true})] + ;; forward, backward, then: + (ddp/optimizer-step! parallel-model optimizer)) +``` + +The constructor verifies parameter signatures across ranks and broadcasts parameters from rank zero. Backward hooks copy gradients into reusable buckets and start asynchronous all-reduce operations. `ddp/optimizer-step!` waits for pending reductions before stepping the optimizer. + +Each synchronized backward pass must produce a gradient for every trainable parameter. `:find-unused-parameters? true` and `:gradient-as-bucket-view? true` are unsupported and fail during construction. + +### Gradient accumulation + +Use `ddp/no-sync` around every micro-batch except the last one: + +```clojure +(doseq [[index batch] (map-indexed vector micro-batches)] + (let [train! #(train-micro-batch! parallel-model batch)] + (if (= index (dec (count micro-batches))) + (train!) + (ddp/no-sync (train!))))) + +(ddp/optimizer-step! parallel-model optimizer) +``` + +## Distributed sampling + +Each rank needs a disjoint, deterministic sample stream: + +```clojure +(require '[clorch.data :as data]) + +(def sampler + (data/distributed-sampler + dataset-size + {:num-replicas (dist/world-size) + :rank (dist/rank) + :seed 1337 + :shuffle? true + :drop-last? false})) + +(data/set-epoch! sampler epoch) +(doseq [indices (partition-all batch-size (data/sample-indices sampler))] + (train-batch! indices)) +``` + +Call `set-epoch!` before each epoch. Otherwise each epoch reuses the same permutation. Sampler state includes epoch, seed, rank, replica count, dataset size, shuffle mode, and drop-last mode. + +## Mixed precision + +```clojure +(require '[clorch.amp :as amp]) + +(def scaler (amp/grad-scaler {:initial-scale 65536.0})) + +(let [loss (amp/autocast {:device :cuda :dtype :float16} + (compute-loss parallel-model batch))] + (amp/backward! scaler loss) + (ddp/optimizer-step! parallel-model optimizer {:scaler scaler})) +``` + +`amp/step!` unscales gradients, checks finite values across all initialized ranks, skips the optimizer step on overflow, and updates the scale. Autocast restores the thread's previous dtype, enabled flag, and cache setting after the body exits. + +## Checkpoints + +```clojure +(dist/save-checkpoint! + "/checkpoints/model.pt" + {:model model + :optimizer optimizer + :sampler sampler + :scaler scaler + :state {:epoch epoch :global-step step}}) + +(def training-state + (dist/load-checkpoint! + "/checkpoints/model.pt" + {:model model + :optimizer optimizer + :sampler sampler + :scaler scaler})) +``` + +Rank zero writes the tensor archive and EDN metadata through temporary files, then commits them with atomic moves. Every rank restores model weights, optimizer state, CPU RNG state, sampler state, scaler state, and the supplied EDN training state. Clorch does not capture CUDA generator state; store and reapply your CUDA seed when exact CUDA random-stream replay matters. + +## Current verification scope + +The release suite covers CPU behavior and a single-GPU CUDA path, including NCCL world size one, DDP backward, AMP overflow handling, fused scaled-dot-product attention, checkpoints, worker failures, and process cleanup. Two-rank execution requires a machine with at least two visible NVIDIA GPUs and belongs in release validation for changes to collectives or DDP. diff --git a/docs/index.md b/docs/index.md index 129abe5..42680f9 100644 --- a/docs/index.md +++ b/docs/index.md @@ -5,7 +5,7 @@ title: Home # Clorch -**Clorch** is a high-performance deep learning library for Clojure with CPU and NVIDIA CUDA GPU execution. Its tensor, autograd, neural-network, optimizer, data, and device APIs follow PyTorch's names and behavior closely while remaining idiomatic and REPL-friendly. +**Clorch** is a high-performance deep learning library for Clojure with CPU and NVIDIA CUDA GPU execution. Its tensor, autograd, neural-network, optimizer, data, AMP, and NCCL distributed APIs follow PyTorch's names and behavior closely while remaining idiomatic and REPL-friendly. --- @@ -16,6 +16,7 @@ title: Home * **Broad PyTorch Coverage**: Supports more than 100 tensor operations, activations, and layers. * **Transparent Memory**: Uses deterministic `with-torch` scopes to prevent native memory bloat. * **LLM Architecture**: Includes embeddings, RMSNorm, RoPE, grouped-query attention, SwiGLU, KV caching, and autoregressive generation. +* **Distributed CUDA**: Provides NCCL collectives, local rank launch, DDP gradient synchronization, distributed sampling, AMP, and rank-zero checkpoints. --- @@ -63,6 +64,7 @@ title: Home - [**Functional API**](functional.md): Stateless operations and low-level control. - [**Activation Functions**](activations.md): Complete list of 25+ supported activations. - [**Loss Functions**](losses.md): Criteria for regression and classification. +- [**Distributed CUDA Training**](distributed.md): Worker launch, collectives, DDP, AMP, sampling, and checkpoints. ### Operations & Stability - [**Memory Management**](memory.md): How to use `with-torch` effectively. @@ -92,3 +94,4 @@ Check out the `examples/` directory for full implementations: * `autograd_tutorial.clj`: Basic math and gradients. * `modern_llama.clj`: Llama-3 style components (RoPE, GQA, SwiGLU). * `synthetic.clj`: A complete training loop from data to optimization. +* `distributed_training.clj`: Multi-process NCCL DDP training with AMP and checkpointing. diff --git a/docs/llms.txt b/docs/llms.txt index dabf107..d45c9ce 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -1,14 +1,14 @@ # Clorch -> Clorch is a Clojure deep-learning library backed by LibTorch. It provides PyTorch-style tensors, autograd, neural-network modules, optimizers, data loading, and explicit CPU/CUDA device placement through a REPL-friendly API. +> Clorch is a Clojure deep-learning library backed by LibTorch. It provides PyTorch-style tensors, autograd, neural-network modules, optimizers, data loading, explicit CPU/CUDA placement, AMP, and NCCL distributed training through a REPL-friendly API. Clorch follows PyTorch concepts but does not expose every PyTorch symbol. Verify public vars in the linked documentation or source instead of translating Python names mechanically. -Namespace map: `torch` → `clorch.torch`; `torch.cuda` → `clorch.cuda`; `torch.autograd` → `clorch.autograd`; `torch.nn` → `clorch.nn`; `torch.nn.functional` → `clorch.nn.functional`; `torch.optim` → `clorch.optim`; `torch.distributions` → `clorch.distributions`; `torch.linalg` → `clorch.linalg`; `Dataset` and `DataLoader` → `clorch.data`. +Namespace map: `torch` → `clorch.torch`; `torch.cuda` → `clorch.cuda`; `torch.amp` → `clorch.amp`; `torch.distributed` → `clorch.distributed`; `DistributedDataParallel` → `clorch.nn.parallel`; `torch.autograd` → `clorch.autograd`; `torch.nn` → `clorch.nn`; `torch.nn.functional` → `clorch.nn.functional`; `torch.optim` → `clorch.optim`; `torch.distributions` → `clorch.distributions`; `torch.linalg` → `clorch.linalg`; `Dataset` and `DataLoader` → `clorch.data`. Runtime rules: resolve one application-level device and place models and tensors explicitly. Keep models and optimizers outside repeated `with-torch` scopes. Use one scope per allocating batch or generation step, returning a JVM scalar or `nil` unless a tensor must escape. -Current limits: single-device training only. Distributed training, AMP, dedicated Flash Attention, and a `torch.compile` equivalent are not implemented. Consult the parity guide before making coverage claims. +Current limits: distributed execution is NCCL-only and the managed launcher is local to one host. Gloo, RPC, FSDP, tensor parallelism, elastic membership, CUDA RNG checkpointing, and a `torch.compile` equivalent are not implemented. The release suite requires a multi-GPU NVIDIA host to validate more than one rank. ## Start Here @@ -33,6 +33,7 @@ Current limits: single-device training only. Distributed training, AMP, dedicate - [Optimizers](https://raw.githubusercontent.com/antlobach/clorch/main/docs/optimizers.md): SGD, Adam, AdamW, RMSprop, and Adagrad. - [Loss functions](https://raw.githubusercontent.com/antlobach/clorch/main/docs/losses.md): Regression and classification losses. - [Initialization](https://raw.githubusercontent.com/antlobach/clorch/main/docs/init.md): Parameter initialization helpers. +- [Distributed CUDA training](https://raw.githubusercontent.com/antlobach/clorch/main/docs/distributed.md): NCCL process groups, collectives, worker lifecycle, DDP, gradient accumulation, AMP, distributed sampling, and checkpoints. - [Memory management](https://raw.githubusercontent.com/antlobach/clorch/main/docs/memory.md): Native-memory scopes, retention, release, REPL sessions, and repeated workloads. ## Examples @@ -43,6 +44,7 @@ Current limits: single-device training only. Distributed training, AMP, dedicate - [Synthetic training](https://raw.githubusercontent.com/antlobach/clorch/main/examples/synthetic.clj): Compact end-to-end supervised training. - [Modern Llama](https://raw.githubusercontent.com/antlobach/clorch/main/examples/modern_llama.clj): RMSNorm, RoPE, GQA, SwiGLU, causal masking, and KV caching. - [NanoChat-inspired example](https://raw.githubusercontent.com/antlobach/clorch/main/examples/nanochat.clj): Compact single-device Llama-style training, checkpointing, generation, and terminal chat inspired by Karpathy's NanoChat. +- [Distributed CUDA training](https://raw.githubusercontent.com/antlobach/clorch/main/examples/distributed_training.clj): Local multi-process DDP with AMP, gradient accumulation, distributed sampling, and checkpoints. - [LLMs from scratch](https://raw.githubusercontent.com/antlobach/clorch/main/examples/llms_from_scratch.clj): GPT-style transformer construction with custom Clorch modules. ## Project Status diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 4eec1b4..7aa4ff6 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -25,6 +25,7 @@ nav: - Optimizers: optimizers.md - Initialization: init.md - Autograd: autograd.md + - Distributed CUDA Training: distributed.md - PyTorch Parity: pytorch-parity.md - Memory Management: memory.md - Performance & Profiling: profiling.md diff --git a/docs/pytorch-parity.md b/docs/pytorch-parity.md index 9bc44bd..d6aaaf9 100644 --- a/docs/pytorch-parity.md +++ b/docs/pytorch-parity.md @@ -7,12 +7,12 @@ Clorch follows PyTorch's tensor, autograd, module, optimizer, data, and device c | Measurement | Current result | Meaning | |---|---:|---| | Cross-language numerical parity suite | **100% (40/40 scenarios)** | Every currently paired Clorch/Python scenario passes within its configured numerical tolerance. | -| Tracked feature-catalog coverage | **approximately 81.7% (268 of approximately 328 capabilities)** | Coverage of the explicit capability catalog maintained in this repository. | +| Tracked feature-catalog coverage | **Pending recount after the PyTorch 2.10 upgrade** | `PORTING_STATUS.md` records the implemented surfaces, but the prior 268/328 estimate predates AMP and distributed training. | | Entire upstream PyTorch public API | **Not yet enumerated** | A whole-project percentage requires a version-pinned inventory of every upstream public symbol and behavior. | -The 81.7% figure is the best current breadth estimate: $268 / (268 + 60) \times 100 \approx 81.7\%$. It is not a claim that 81.7% of every symbol in every PyTorch package exists. The 100% figure applies only to the 40 cross-language scenarios currently exercised by `tests_comparison/compare_torch.py`. +The 100% figure applies only to the 40 cross-language scenarios exercised by `tests_comparison/compare_torch.py`. Clorch will publish another breadth percentage after generating a version-pinned PyTorch 2.10 inventory. -The release suite also verifies CUDA discovery, GPU tensor allocation, device reporting, CUDA computation, synchronization, and transfer back to CPU on NVIDIA hardware. +The release suite verifies CPU behavior and a single-GPU CUDA path, including CUDA discovery, NCCL world size one, DDP backward, AMP overflow handling, fused scaled-dot-product attention, checkpoints, worker failures, and process cleanup. Multi-rank validation requires a host with at least two visible NVIDIA GPUs. ## Capability comparison @@ -25,12 +25,12 @@ The release suite also verifies CUDA discovery, GPU tensor allocation, device re | Autograd | `requires-grad`, `backward`, gradient access, detach, and `no-grad` | Custom autograd functions, `gradcheck`, forward-mode AD, inference mode, hooks, anomaly detection, and full graph-control APIs | | Neural-network modules | Linear, convolution, transposed convolution, normalization, pooling, padding, recurrent layers, embeddings, activations, dropout, and custom `defmodel` modules | Native transformer modules, complete container modules, hooks, parametrizations, pruning, lazy modules, and remaining specialized layers | | Functional API and losses | Core linear/convolutional functions, activations, normalization, pooling, interpolation, padding, and common classification/regression losses | Remaining functional operators and full loss/reduction/weighting option parity | -| Optimizers | SGD, Adam, AdamW, RMSprop, Adagrad, zeroing, and stepping | Remaining optimizers, learning-rate schedulers, optimizer state dictionaries, parameter-group mutation, and hooks | -| Data loading | Dataset protocol, tensor datasets, batching, shuffling, thread workers, and process workers | Iterable datasets, distributed samplers, DataPipes, collator breadth, pin-memory semantics, and complete worker lifecycle parity | -| Serialization and JIT | Tensor/model save-load, state dictionaries, weight save-load, and JIT load-save-forward | `torch.export`, compilation, tracing breadth, package APIs, ONNX export, and complete archive compatibility | -| CUDA | Automatic native-backend selection, availability, device count, synchronization, seeding, and explicit CUDA tensor/model placement | Streams, events, graphs, AMP/autocast, memory statistics, allocator controls, peer access, and complete multi-GPU APIs | +| Optimizers | SGD, Adam, AdamW, RMSprop, Adagrad, zeroing, stepping, and native checkpoint restore | Remaining optimizers, learning-rate schedulers, public optimizer state dictionaries, parameter-group mutation, and hooks | +| Data loading | Dataset protocol, tensor datasets, batching, shuffling, thread/process workers, and deterministic distributed samplers | Iterable datasets, DataPipes, collator breadth, pin-memory semantics, and complete worker lifecycle parity | +| Serialization and JIT | Tensor/model save-load, state dictionaries, weight save-load, JIT load-save-forward, and atomic rank-zero training checkpoints | `torch.export`, compilation, tracing breadth, package APIs, ONNX export, CUDA RNG checkpointing, and complete archive compatibility | +| CUDA | Automatic backend selection, availability, device count and selection, synchronization, seeding, explicit placement, autocast, dynamic gradient scaling, and fused scaled-dot-product attention | Streams, events, graphs, memory/allocator controls, peer access, and the remaining multi-GPU APIs | | Sparse and quantized computation | Basic quantized and complex dtype exposure | Sparse layouts and operators, complete quantization workflows, observers, prepared/converted modules, and quantized kernels | -| Distributed training | Not implemented | DistributedDataParallel, RPC, collectives, NCCL/Gloo integration, FSDP, tensor parallelism, and distributed checkpointing | +| Distributed training | Direct NCCL collectives, local rank launcher, distributed samplers, synchronous DDP, gradient accumulation, failure propagation, and rank-zero checkpoints | Gloo, RPC, FSDP, tensor parallelism, elastic membership, unused-parameter discovery, multi-node launcher support, and broader multi-GPU validation | | Compiler and performance stack | Native LibTorch execution and explicit native-memory scopes | `torch.compile`-equivalent graph capture, Dynamo/Inductor-style compilation, profiler breadth, and benchmark tooling | | Domain libraries | Core deep-learning library only | TorchVision, TorchAudio, TorchText, TorchData, and ecosystem-specific model/data APIs | @@ -48,10 +48,10 @@ The release suite also verifies CUDA discovery, GPU tensor allocation, device re | Autoregressive generation | Implemented | `nn/generate` and the NanoChat generation loop | | GPT-style transformer construction | Implemented with custom modules | `examples/llms_from_scratch.clj` | | Llama-style blocks | Implemented with custom modules | `examples/modern_llama.clj`, `examples/nanochat.clj` | -| Flash/scaled-dot-product attention kernels | Not implemented as dedicated fused APIs | Add fused attention bindings and kernel selection | -| Mixed-precision training | Partial dtype support; no complete AMP workflow | Add autocast, gradient scaling, and AMP state management | +| Flash/scaled-dot-product attention kernels | Implemented | `F/scaled-dot-product-attention` uses LibTorch's fused dispatcher | +| Mixed-precision training | Implemented for CUDA/CPU autocast and dynamic gradient scaling | `clorch.amp` | | Quantized LLM inference | Not complete | Add end-to-end weight/activation quantization and quantized kernels | -| Distributed LLM training/inference | Not implemented | Add collectives, DDP/FSDP, tensor parallelism, and distributed checkpoints | +| Distributed LLM training/inference | Partial | NCCL DDP training exists; add FSDP, tensor parallelism, RPC, and multi-node orchestration | ## Roadmap to 100% @@ -59,9 +59,9 @@ The release suite also verifies CUDA discovery, GPU tensor allocation, device re 2. **Turn inventory into contracts.** Map every upstream symbol to implemented, partial, intentionally different, or missing; attach behavior tests to every implemented mapping. 3. **Close core tensor gaps.** Finish dense overloads and edge cases, then sparse, nested, meta, quantized, and named-tensor behavior. 4. **Complete autograd and modules.** Add custom functions, grad checking, hooks, inference controls, native transformer/container modules, and remaining functional/loss APIs. -5. **Complete optimization and data.** Add optimizer state, schedulers, remaining algorithms, iterable/distributed data loading, and worker/pinning parity. -6. **Complete accelerator support.** Add CUDA streams/events/graphs, AMP, memory controls, multi-GPU execution, and distributed collectives. +5. **Complete optimization and data.** Add public optimizer state, schedulers, remaining algorithms, iterable data loading, DataPipes, and worker/pinning parity. +6. **Complete accelerator support.** Add CUDA streams, events, graphs, allocator controls, Gloo, FSDP, tensor parallelism, elastic execution, and broader multi-GPU conformance tests. 7. **Complete export and ecosystem surfaces.** Add compile/export/ONNX/package behavior and explicitly scope domain libraries. 8. **Prove the result.** Run generated cross-language conformance tests across Linux, macOS, Windows, CPU, CUDA, supported dtypes, and error/edge-case behavior. -Until that version-pinned inventory exists, report both current numbers together: **100% of 40 tested numerical scenarios and approximately 81.7% of the repository's tracked feature catalog**. +Until that inventory exists, report the measured result without inventing breadth precision: **100% of 40 tested numerical scenarios; tracked catalog percentage pending recount**. From 015a401340dbc0d1b8b87ee91f966f7824f0c2c8 Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Sun, 9 Aug 2026 03:52:35 +0200 Subject: [PATCH 5/7] fix(ddp): commit reduced gradients before step Gradient hooks could return before NCCL results reached every parameter, leaving later gradients rank-local. Commit reduced buckets after waiting for collective completion. --- examples/distributed_training.clj | 10 +++++--- src/clorch/nn/parallel.clj | 41 +++++++++++++++++------------- test/clorch/cuda_smoke_test.clj | 42 +++++++++++++++++++++++++++++++ 3 files changed, 72 insertions(+), 21 deletions(-) diff --git a/examples/distributed_training.clj b/examples/distributed_training.clj index ba4dd52..6da9361 100644 --- a/examples/distributed_training.clj +++ b/examples/distributed_training.clj @@ -77,15 +77,17 @@ :sampler sampler :scaler scaler :state {:epoch epochs :global-step @global-step}})) - (let [parameter (.get (nn/parameters model) 0) - rank-zero-parameter (.clone ^Tensor parameter)] - (dist/broadcast! process-group rank-zero-parameter {:root-rank 0}) + (let [model-parameters (t/->vector (nn/parameters model)) + rank-zero-parameters (mapv #(.clone ^Tensor %) model-parameters)] + (dist/broadcast! process-group rank-zero-parameters {:root-rank 0}) {:rank rank :steps @global-step :loss @last-loss :precision precision :parameters-synchronized? - (.allclose ^Tensor parameter ^Tensor rank-zero-parameter)})))) + (every? true? + (map #(.allclose ^Tensor %1 ^Tensor %2) + model-parameters rank-zero-parameters))})))) (defn run-local! "Launches one training worker JVM per CUDA device and waits for completion." diff --git a/src/clorch/nn/parallel.clj b/src/clorch/nn/parallel.clj index 3c02bec..06e2e42 100644 --- a/src/clorch/nn/parallel.clj +++ b/src/clorch/nn/parallel.clj @@ -97,16 +97,29 @@ :buffer buffer :views views :remaining (atom (count indices)) - :gradients (atom (vec (repeat (count indices) nil))) :lock (Object.)}))) +(defn- commit-bucket-gradients! [context bucket] + (with-open [divisor (Scalar. (double (:world-size context)))] + (.div_ ^Tensor (:buffer bucket) divisor)) + (doseq [slot (range (count (:indices bucket)))] + (let [^Tensor parameter + (.get ^TensorVector (:parameters bucket) (long slot)) + ^Tensor gradient (.grad parameter)] + (when-not (.defined gradient) + (throw (ex-info "Backward pass did not materialize a parameter gradient" + {:parameter-index ((:indices bucket) slot)}))) + (.copy_ gradient + (.get ^TensorVector (:views bucket) (long slot)))))) + (defn- await-pending! [context pending] - (let [works (locking pending - (let [works @pending] - (reset! pending []) - works))] - (doseq [work works] - (dist/await! work (:timeout-ms context))))) + (let [entries (locking pending + (let [entries @pending] + (reset! pending []) + entries))] + (doseq [{:keys [work bucket]} entries] + (dist/await! work (:timeout-ms context)) + (commit-bucket-gradients! context bucket)))) (defn- reset-prior-gradient! [^Tensor parameter ^Tensor gradient] (let [prior (.grad parameter)] @@ -116,15 +129,10 @@ (defn- complete-bucket! [context bucket pending] - (let [buffer (:buffer bucket) - work (dist/all-reduce! context buffer {:op :sum :async? true})] - (.div_ ^Tensor buffer (Scalar. (double (:world-size context)))) - (doseq [slot (range (count (:indices bucket)))] - (.copy_ ^Tensor (@(:gradients bucket) slot) - (.get ^TensorVector (:views bucket) (long slot)))) - (swap! pending conj work) - (reset! (:remaining bucket) (count (:indices bucket))) - (reset! (:gradients bucket) (vec (repeat (count (:indices bucket)) nil))))) + (let [work (dist/all-reduce! context (:buffer bucket) + {:op :sum :async? true})] + (swap! pending conj {:work work :bucket bucket}) + (reset! (:remaining bucket) (count (:indices bucket))))) (defn- gradient-hook [context bucket slot ^Tensor parameter synchronize? pending classloader] @@ -139,7 +147,6 @@ (locking bucket-lock (.copy_ ^Tensor (.get ^TensorVector (:views bucket) (long slot)) gradient) - (swap! (:gradients bucket) assoc slot gradient) (when (zero? (swap! (:remaining bucket) dec)) (complete-bucket! context bucket pending))))) gradient-base))) diff --git a/test/clorch/cuda_smoke_test.clj b/test/clorch/cuda_smoke_test.clj index 9e7c18f..da7e1d2 100644 --- a/test/clorch/cuda_smoke_test.clj +++ b/test/clorch/cuda_smoke_test.clj @@ -33,6 +33,47 @@ (let [values (nn/parameters model)] (mapv #(.get values (long %)) (range (.size values))))) +(defn ddp-all-parameter-worker [{:keys [rank process-group]}] + (preload!) + (let [model (nn/to (nn/linear 2 1) :cuda) + optimizer (optim/sgd (nn/parameters model) :lr 0.1) + features (if (zero? rank) + [[1.0 2.0] [2.0 -1.0]] + [[-1.0 1.0] [0.5 3.0]]) + targets (if (zero? rank) + [[0.5] [1.0]] + [[-0.5] [2.0]])] + (with-open [parallel-model + (ddp/distributed-data-parallel + model {:process-group process-group :bucket-cap-mb 1.0})] + (optim/zero-grad optimizer) + (let [input (torch/tensor features {:dtype :float32 :device :cuda}) + target (torch/tensor targets {:dtype :float32 :device :cuda}) + loss (F/mse-loss (nn/forward parallel-model input) target)] + (autograd/backward loss)) + (ddp/optimizer-step! parallel-model optimizer) + (let [model-parameters (parameters model) + rank-zero-parameters (mapv #(.clone ^Tensor %) model-parameters)] + (dist/broadcast! process-group rank-zero-parameters {:root-rank 0}) + (check! (every? true? + (map #(.allclose ^Tensor %1 ^Tensor %2) + model-parameters rank-zero-parameters)) + "DDP did not synchronize every model parameter" + {:rank rank}) + {:rank rank :parameters-synchronized? true})))) + +(defn- verify-two-device-ddp! [] + (when (<= 2 (cuda/device-count)) + (let [job (dist/launch! {:nproc-per-node 2 + :devices [0 1] + :main 'clorch.cuda-smoke-test/ddp-all-parameter-worker + :timeout-ms 120000})] + (dist/await-job! job 180000) + (check! (= {0 0, 1 0} (get-in (dist/job-status job) + [:state :exit-codes])) + "Two-rank DDP worker failed" + {:status (dist/job-status job)})))) + (defn- verify-launcher-failure! [] (let [job (dist/launch! {:nproc-per-node 1 :devices [0] @@ -131,6 +172,7 @@ (io/delete-file checkpoint true) (io/delete-file (str checkpoint ".edn") true)))) (verify-launcher-failure!) + (verify-two-device-ddp!) (println "CUDA, NCCL, DDP, AMP, fused attention, and checkpoint smoke passed.") true))) From dc3fe789c2a0089342ae2c0b2c8ceb23e7dabb51 Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Sun, 9 Aug 2026 03:52:42 +0200 Subject: [PATCH 6/7] docs: add multi-GPU setup guide --- README.md | 143 ++++++++++++++++++++++++++++++++++++++++++-- docs/distributed.md | 16 +++-- 2 files changed, 150 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index d20b0ce..4ea53a0 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,18 @@ Start a REPL from the project directory: clj ``` +### Supported runtime + +| Component | Supported version | +|---|---| +| Java | OpenJDK/Temurin 21 through 25; CI runs the full 21–25 matrix | +| Clojure | Clojure 1.12.x | +| Clojure CLI | A current 1.12.x CLI; CI uses `1.12.5.1664` | +| Native PyTorch | JavaCPP PyTorch `2.10.0-1.5.13` | + +Use a JDK, not a JRE. Java 24 and 25 should be started with +`--enable-native-access=ALL-UNNAMED` for JavaCPP native loading. + ## Quick start ### Tensors and autograd @@ -104,11 +116,14 @@ Non-trainable params: 0 The constructor arguments configure the model, the binding vector registers modules or parameters, and the `forward` form defines execution. Registered fields participate in `nn/parameters`, `nn/to`, and state dictionaries. Read [Custom Models with `defmodel`](docs/nn.md#custom-models-with-defmodel) or the [minimal source example](examples/simple.clj). -## CPU and CUDA +## CPU, CUDA, and multi-GPU training -Clorch selects the native backend when `clorch.torch` loads. It uses CUDA when GPU natives are available and NVIDIA hardware is detected; otherwise it loads the CPU backend. +Clorch selects the native backend when `clorch.torch` loads. It uses CUDA when +GPU natives are available and NVIDIA hardware is detected; otherwise it loads +the CPU backend. -Backend selection does not move tensors or models. Choose a device once, then place models and data on it: +Backend selection does not move tensors or models. Choose a device once, then +place models and data on it: ```clojure (require '[clorch.cuda :as cuda]) @@ -121,9 +136,127 @@ Backend selection does not move tensors or models. Choose a device once, then pl Environment overrides: - `CLORCH_FORCE_CPU=1` forces the CPU backend. -- `CLORCH_FORCE_GPU=1` requests the CUDA backend. It still requires compatible natives, hardware, and an NVIDIA driver. +- `CLORCH_FORCE_GPU=1` requests the CUDA backend. It still requires compatible + native libraries, hardware, and an NVIDIA driver. + +### Multi-GPU requirements + +Distributed training uses one worker JVM per GPU and NCCL for communication. +The managed launcher covers multiple GPUs on one Linux host. + +- Two or more NVIDIA GPUs visible to the same process +- A working NVIDIA driver with CUDA 13 support +- CUDA 13.1 user-space libraries, including cuBLAS +- cuDNN 9 and NCCL 2; the validated stack uses cuDNN 9.19 and NCCL 2.29.2 +- Java 21 through 25 and Clojure 1.12.x +- One distinct CUDA device per rank +- Enough host RAM and CUDA VRAM for one model replica per rank +- A writable checkpoint directory and an available local TCP port + +For an Ubuntu host configured with NVIDIA's CUDA package repository, the +required runtime packages are: + +```bash +sudo apt-get update +sudo apt-get install cuda-libraries-13-1 libcudnn9-cuda-13 libnccl2 +``` + +Confirm that the host exposes each GPU before starting Clojure: + +```bash +nvidia-smi -L +java -version +clojure -Sdescribe +``` + +Set native-loading variables before the JVM starts. `JAVA_TOOL_OPTIONS` is +inherited by launcher-created worker JVMs: + +```bash +export CLORCH_FORCE_GPU=1 +export LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH:-}" +export JAVA_TOOL_OPTIONS="--enable-native-access=ALL-UNNAMED" + +CLOJURE_DISABLE_RLWRAP=1 clojure -M:dev +``` + +Check CUDA from the REPL: + +```clojure +(require '[clorch.cuda :as cuda] + '[clorch.torch :as t]) + +{:available (cuda/available?) + :devices (cuda/device-count)} +;; => {:available true, :devices 2} +``` + +### Launch the distributed training example + +Clone this repository and start its nREPL so the `examples` path is available. +The example below launches ranks for physical devices 0 and 1, waits for both +workers, and returns status plus per-rank logs: + +```clojure +(require '[distributed-training :as training]) + +(def result + (training/run-local! + [0 1] + {:epochs 4 + :sample-count 1024 + :batch-size 32 + :accumulation 2 + :precision :bfloat16 + :checkpoint-path "/checkpoints/clorch-ddp.pt"})) +``` + +Use `:float16` for autocast with dynamic gradient scaling, or `:bfloat16` for +autocast without a scaler. `:accumulation 2` performs two micro-batches per +optimizer step and suppresses DDP synchronization until the final +micro-batch. + +The launcher sets `RANK`, `LOCAL_RANK`, `WORLD_SIZE`, `LOCAL_WORLD_SIZE`, +`MASTER_ADDR`, `MASTER_PORT`, and `CUDA_VISIBLE_DEVICES` for every child. Worker +code should place its model and batches on `:cuda`; the launcher maps that +logical device to the assigned physical GPU. + +For an application worker, launch a namespace-qualified function: + +```clojure +(require '[clorch.distributed :as dist]) + +(def job + (dist/launch! + {:nproc-per-node 2 + :devices [0 1] + :main 'my.training/train-worker + :args {:epochs 10} + :timeout-ms 300000})) + +(dist/job-status job) +(dist/await-job! job) +(dist/job-logs job 0) +``` + +Every rank must execute collectives in the same order and train from a +disjoint `distributed-sampler` partition. Create the model on `:cuda` before +wrapping it with `distributed-data-parallel`, call `data/set-epoch!` each +epoch, and use `ddp/optimizer-step!` instead of stepping the optimizer +directly. Only rank zero writes checkpoints; all ranks participate in save and +restore barriers. + +Run the GPU release check after configuring the host: + +```bash +CLORCH_FORCE_GPU=1 clojure -Sthreads 1 -M -m clorch.release-check --mode gpu +``` -Clorch currently supports single-device training. Distributed training, DDP, FSDP, collectives, and distributed checkpointing are not implemented. +On a host with at least two GPUs, the CUDA smoke includes a two-rank DDP check +that verifies every model parameter remains synchronized. Read +[Distributed CUDA Training](docs/distributed.md) for worker code, collectives, +sampling, AMP, gradient accumulation, failure handling, and checkpoint +restore. ## Native memory diff --git a/docs/distributed.md b/docs/distributed.md index 36af83b..b178a17 100644 --- a/docs/distributed.md +++ b/docs/distributed.md @@ -5,9 +5,9 @@ Clorch runs one JVM per CUDA device and coordinates ranks with NCCL. The public ## Requirements - Linux with NVIDIA GPUs and a working driver -- Java 25 or newer for the managed launcher +- Java 21 through 25 and Clojure 1.12.x - One distinct CUDA device per local rank -- GPU-enabled JavaCPP PyTorch and CUDA native libraries on the classpath +- CUDA 13.1 user-space libraries, cuDNN 9, and NCCL 2 on the classpath Clorch currently supports the `:nccl` backend. It does not provide Gloo, RPC, FSDP, tensor parallelism, or elastic membership. @@ -105,7 +105,11 @@ Create the model on the rank-local CUDA device before wrapping it: (ddp/optimizer-step! parallel-model optimizer)) ``` -The constructor verifies parameter signatures across ranks and broadcasts parameters from rank zero. Backward hooks copy gradients into reusable buckets and start asynchronous all-reduce operations. `ddp/optimizer-step!` waits for pending reductions before stepping the optimizer. +The constructor verifies parameter signatures across ranks and broadcasts +parameters from rank zero. Backward hooks copy local gradients into reusable +buckets and start asynchronous all-reduce operations. `ddp/optimizer-step!` +waits for pending reductions, commits every averaged bucket back to the model +gradients, and then steps the optimizer. Each synchronized backward pass must produce a gradient for every trainable parameter. `:find-unused-parameters? true` and `:gradient-as-bucket-view? true` are unsupported and fail during construction. @@ -185,4 +189,8 @@ Rank zero writes the tensor archive and EDN metadata through temporary files, th ## Current verification scope -The release suite covers CPU behavior and a single-GPU CUDA path, including NCCL world size one, DDP backward, AMP overflow handling, fused scaled-dot-product attention, checkpoints, worker failures, and process cleanup. Two-rank execution requires a machine with at least two visible NVIDIA GPUs and belongs in release validation for changes to collectives or DDP. +The release suite covers CPU behavior and CUDA execution, including NCCL, DDP +backward, AMP overflow handling, fused scaled-dot-product attention, +checkpoints, worker failures, and process cleanup. Two-rank validation on two +RTX A5000 GPUs covers NCCL gradient reduction, parameter synchronization, +bfloat16 AMP, gradient accumulation, and rank-zero checkpoint creation. From 933d890bd05d4784bf1abe1b0e539985c253107d Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Sun, 9 Aug 2026 03:52:47 +0200 Subject: [PATCH 7/7] ci: test Java 21 through 25 --- .github/workflows/ci.yml | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 663f9c3..d654f90 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -34,21 +34,27 @@ concurrency: jobs: test: - name: Clojure tests + name: Clojure tests (Java ${{ matrix.java }}) runs-on: ubuntu-latest timeout-minutes: 45 + strategy: + fail-fast: false + matrix: + java: ["21", "22", "23", "24", "25"] + env: CLORCH_FORCE_CPU: "1" + JAVA_TOOL_OPTIONS: "--enable-native-access=ALL-UNNAMED" steps: - name: Check out source uses: actions/checkout@v4 - - name: Set up Java 17 + - name: Set up Java ${{ matrix.java }} uses: actions/setup-java@v4 with: distribution: temurin - java-version: "17" + java-version: ${{ matrix.java }} - name: Set up Clojure CLI uses: DeLaGuardo/setup-clojure@13.4