> ## Documentation Index
> Fetch the complete documentation index at: https://antlobach-clorch-182a83cb.mintlify.site/llms.txt
> Use this file to discover all available pages before exploring further.

# 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

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

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

<Steps>
  <Step title="Install CUDA runtime packages">
    On an Ubuntu host configured with NVIDIA's CUDA package repository:

    ```bash theme={null}
    sudo apt-get update
    sudo apt-get install cuda-libraries-13-1 libcudnn9-cuda-13 libnccl2
    ```
  </Step>

  <Step title="Verify GPU visibility">
    ```bash theme={null}
    nvidia-smi -L
    java -version
    clojure -Sdescribe
    ```
  </Step>

  <Step title="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.

    ```bash theme={null}
    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
    ```
  </Step>

  <Step title="Confirm CUDA from the REPL">
    ```clojure theme={null}
    (require '[clorch.cuda :as cuda])

    {:available (cuda/available?)
     :devices   (cuda/device-count)}
    ;; => {:available true, :devices 2}
    ```
  </Step>
</Steps>

***

## Running the Training Example

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

```clojure theme={null}
(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:

```clojure theme={null}
(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

```clojure theme={null}
(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:

```clojure theme={null}
{: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:

```clojure theme={null}
(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.

<Tabs>
  <Tab title="all-reduce!">
    Reduces tensors across every rank in place.

    ```clojure theme={null}
    ;; 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`.
  </Tab>

  <Tab title="broadcast!">
    Broadcasts tensors from the root rank to all other ranks in place.

    ```clojure theme={null}
    (dist/broadcast! model-parameters {:root-rank 0})
    ```
  </Tab>

  <Tab title="reduce!">
    Reduces tensors in place onto the root rank only.

    ```clojure theme={null}
    (dist/reduce! tensor {:op :sum :root-rank 0})
    ```
  </Tab>

  <Tab title="gather / scatter">
    ```clojure theme={null}
    ;; 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]})
    ```
  </Tab>

  <Tab title="send / receive!">
    Point-to-point communication between specific ranks.

    ```clojure theme={null}
    ;; 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))
    ```
  </Tab>

  <Tab title="barrier!">
    Block until every rank has entered the barrier.

    ```clojure theme={null}
    (dist/barrier!)

    ;; Async variant
    (def work (dist/barrier! {:async? true}))
    (dist/await! work)
    ```
  </Tab>
</Tabs>

***

## 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.

```clojure theme={null}
(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.

<Warning>
  `: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.
</Warning>

### Gradient Accumulation

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

```clojure theme={null}
(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.

```clojure theme={null}
(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))
```

<Warning>
  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.
</Warning>

`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](/advanced/amp) for full details. A brief integration example:

```clojure theme={null}
(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

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

### Restoring

```clojure theme={null}
(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.

<Warning>
  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.
</Warning>

***

## 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.

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

  ```bash theme={null}
  CLORCH_FORCE_GPU=1 clojure -Sthreads 1 -M -m clorch.release-check --mode gpu
  ```
</Info>


This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.