Multi-GPU Distributed Training with NCCL, DDP, and AMP

NCCL workers, collective ops, DDP gradient sync, gradient accumulation, distributed sampling, and atomic rank-zero checkpoints in one Clorch workflow.

Clorch runs one JVM per CUDA device and coordinates ranks with NCCL. The distributed API covers process-group lifecycle, collective operations, local worker launch, distributed sampling, synchronous data-parallel training with DDP, automatic mixed precision, and rank-zero checkpoints. This page covers every aspect of that system in depth.

---

Requirements

Clorch currently supports only the :nccl backend. Gloo, RPC, FSDP, tensor parallelism, and elastic membership are not yet implemented.

Component	Required version
Linux	Any modern distribution with NVIDIA driver
NVIDIA GPUs	One distinct device per rank
Java	OpenJDK/Temurin 25+ (distributed workers require Java 25 or newer)
Clojure	1.12.x
CUDA user-space	13.1
cuDNN	9 (tested with 9.19)
NCCL	2 (tested with 2.29.2)

1. Install CUDA runtime packages

   On an Ubuntu host configured with NVIDIA’s CUDA package repository:

   sudo apt-get update
   sudo apt-get install cuda-libraries-13-1 libcudnn9-cuda-13 libnccl2
   
2. Verify GPU visibility

   nvidia-smi -L
   java -version
   clojure -Sdescribe
   
3. Set environment variables

   These variables must be present before the Clojure process starts. JAVA_TOOL_OPTIONS is inherited by every worker JVM spawned by the launcher.

   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
   
4. Confirm CUDA from the REPL

   (require '[clorch.cuda :as cuda])
   
   {:available (cuda/available?)
    :devices   (cuda/device-count)}
   ;; => {:available true, :devices 2}
   

---

Running the Training Example

The shipped example in examples/distributed_training.clj demonstrates a complete DDP training loop with AMP and checkpointing.

(require '[distributed-training :as training])

;; Launch one worker per GPU, wait for completion, return logs
(def result
  (training/run-local!
   [0 1]
   {:epochs            4
    :sample-count      1024
    :batch-size        32
    :accumulation      2
    :precision         :bfloat16
    :checkpoint-path   "/tmp/clorch-ddp.pt"}))

Use :float16 to enable dynamic loss scaling. Use :bfloat16 for autocast without a scaler. :accumulation 2 performs two micro-batches per optimizer step and suppresses DDP synchronization until the final micro-batch.

---

Launching Workers with dist/launch!

For your own training namespace, call dist/launch! directly with a config map:

(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}))

Job Management

(dist/job-status job)       ;; current state map
(dist/await-job! job)       ;; block until all ranks complete
(dist/job-logs job 0)       ;; tail of rank 0's stdout and stderr
(dist/stop-job! job)        ;; terminate all live ranks

Worker Entrypoint Map

Clorch invokes your :main function with a single map argument:

{:rank          0         ;; global rank
 :local-rank    0         ;; rank within this node
 :world-size    2         ;; total number of ranks
 :backend       :nccl
 :process-group context   ;; initialized ProcessGroupContext
 :args          {:epochs 10}}

Environment Variables Set by the Launcher

The launcher sets the following for each child JVM:

Variable	Value
RANK	Global rank index
LOCAL_RANK	Rank within the node
WORLD_SIZE	Total number of ranks
LOCAL_WORLD_SIZE	Ranks on this node
MASTER_ADDR	Rendezvous host (default 127.0.0.1)
MASTER_PORT	Rendezvous port (auto-selected)
CUDA_VISIBLE_DEVICES	Assigned physical GPU index
CLORCH_DIST_BACKEND	nccl
CLORCH_DIST_TIMEOUT_MS	Timeout in milliseconds

---

Process Groups and Collectives

Workers launched by Clorch receive an initialized process group. For a custom launcher, initialize from environment variables:

(dist/with-process-group {:backend :nccl}
  (dist/all-reduce! tensor {:op :sum})
  (dist/barrier!))

dist/with-process-group initializes the group, runs the body, and always destroys the group, even on error.

Available Collective Operations

Every rank must call collectives in the same order with compatible shapes, dtypes, and split sizes. Collectives operate on CUDA tensors in place.

all-reduce!

Reduces tensors across every rank in place.

;; Sum gradients across all ranks
(dist/all-reduce! gradient-tensor {:op :sum})

;; Async variant — returns a work handle
(def work (dist/all-reduce! gradient-tensor {:op :sum :async? true}))
(dist/await! work)

Supported op values: :sum, :avg, :min, :max, :band, :bor, :bxor.

broadcast!

Broadcasts tensors from the root rank to all other ranks in place.

(dist/broadcast! model-parameters {:root-rank 0})

reduce!

Reduces tensors in place onto the root rank only.

(dist/reduce! tensor {:op :sum :root-rank 0})

gather / scatter

;; Gather equal-sized inputs into a pre-allocated output
(dist/all-gather-into! output-tensor input-tensor)

;; Reduce input and scatter chunks into output
(dist/reduce-scatter-into! output-tensor input-tensor {:op :sum})

;; All-to-all with optional split sizes
(dist/all-to-all-single! output input
  {:output-split-sizes [8 8]
   :input-split-sizes  [8 8]})

send / receive!

Point-to-point communication between specific ranks.

;; Send from rank 0 to rank 1
(when (= rank 0)
  (dist/send tensors 1))

;; Receive on rank 1 from rank 0
(when (= rank 1)
  (dist/receive! tensors 0))

barrier!

Block until every rank has entered the barrier.

(dist/barrier!)

;; Async variant
(def work (dist/barrier! {:async? true}))
(dist/await! work)

---

DistributedDataParallel (DDP)

DDP replicates a model across all ranks, synchronizes gradients via all-reduce after each backward pass, and averages them before the optimizer step.

(require '[clorch.nn :as nn]
         '[clorch.nn.parallel :as ddp]
         '[clorch.optim :as optim])

;; 1. Create model on the rank-local CUDA device
(def model     (nn/to (nn/linear 128 32) :cuda))
(def optimizer (optim/adamw (nn/parameters model) :lr 3e-4))

;; 2. Wrap with DDP (must happen after nn/to)
(with-open [parallel-model
            (ddp/distributed-data-parallel
             model {:bucket-cap-mb     25.0
                    :broadcast-buffers? true})]

  ;; 3. Forward, backward, then step
  (let [loss (nn/forward parallel-model input)]
    (nn/zero-grad optimizer)
    (.backward loss)
    (ddp/optimizer-step! parallel-model optimizer)))

The constructor broadcasts parameters from rank zero, verifying signature consistency across ranks. Backward hooks bucket local gradients and initiate asynchronous all-reduce operations. ddp/optimizer-step! waits for pending reductions, commits averaged gradients to the model, and steps the optimizer.

:find-unused-parameters? true and :gradient-as-bucket-view? true are unsupported and will throw during construction. Every synchronized backward pass must produce a gradient for every trainable parameter.

Gradient Accumulation

Use ddp/no-sync around every micro-batch except the last to suppress expensive all-reduce on intermediate steps:

(doseq [[micro-index batch] (map-indexed vector micro-batches)]
  (let [train! #(train-micro-batch! parallel-model batch)]
    (if (= micro-index (dec (count micro-batches)))
      (train!)                      ;; last micro-batch — sync gradients
      (ddp/no-sync (train!)))))     ;; intermediate — skip all-reduce

(ddp/optimizer-step! parallel-model optimizer)

---

Distributed Sampling

Each rank needs a disjoint, deterministic subset of the dataset. data/distributed-sampler partitions indices across replicas using a seeded shuffle.

(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}))

;; Call set-epoch! before every epoch to reshuffle
(data/set-epoch! sampler epoch)

;; Iterate over this rank's indices in batches
(doseq [indices (partition-all batch-size (data/sample-indices sampler))]
  (train-batch! indices))

Always call data/set-epoch! before each epoch. Omitting it causes every epoch to reuse the same permutation, which can cause training to overfit to the same mini-batches.

data/distributed-sampler defaults :num-replicas and :rank from the WORLD_SIZE and RANK environment variables when omitted, so worker entrypoints do not need to thread these values explicitly.

---

Mixed Precision

For AMP within distributed training, see the AMP page for full details. A brief integration example:

(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)
  ;; Pass scaler to ddp/optimizer-step! for synchronized overflow detection
  (ddp/optimizer-step! parallel-model optimizer {:scaler scaler}))

---

Checkpoints

Only rank zero writes the checkpoint. All ranks participate in save and restore barriers to keep execution synchronized.

Saving

(dist/save-checkpoint!
 "/checkpoints/model.pt"
 {:model     model
  :optimizer optimizer
  :sampler   sampler
  :scaler    scaler
  :state     {:epoch epoch :global-step step}})

Restoring

(def training-state
  (dist/load-checkpoint!
   "/checkpoints/model.pt"
   {:model     model
    :optimizer optimizer
    :sampler   sampler
    :scaler    scaler}))

;; Returns the :state map passed at save time
(:epoch training-state)
(:global-step training-state)

What Gets Saved

Payload	Description
Model weights	Tensor archive via LibTorch’s save
Optimizer state	Optimizer moments and step counts
CPU RNG state	Enables reproducible data augmentation on resume
Sampler state	Epoch, seed, rank, replicas, shuffle, and drop-last flags
Scaler state	Current scale and growth tracker
EDN metadata	Training state map (epoch, global step, anything serializable)

The tensor archive and EDN metadata file are written to temporary files first and then moved atomically, so a crash during writing never leaves a corrupt checkpoint.

Clorch does not capture CUDA generator state. If your data pipeline performs CUDA random operations (e.g., CUDA augmentations), save and reapply your CUDA seed manually when exact replay matters.

---

Current Verification Scope

The release suite covers CPU behavior and CUDA execution paths 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.

Run the GPU release check after configuring the host to confirm your environment:

CLORCH_FORCE_GPU=1 clojure -Sthreads 1 -M -m clorch.release-check --mode gpu