Automatic Mixed Precision: float16, bfloat16, and AMP

clorch.amp provides autocast for float16 or bfloat16, dynamic gradient scaling via grad-scaler, and DDP-aware optimizer stepping with overflow detection.

Automatic Mixed Precision (AMP) reduces memory consumption and increases throughput by running forward passes in a lower-precision floating-point format while keeping weights and critical operations in float32. Modern NVIDIA GPUs have dedicated tensor cores that execute half-precision matrix multiplications several times faster than full-precision equivalents. Clorch’s clorch.amp namespace provides the two building blocks you need: the autocast macro for precision narrowing and the grad-scaler record for dynamic loss scaling.

(require '[clorch.amp :as amp])

---

amp/autocast

autocast is a macro that enables LibTorch’s thread-local autocast state for the duration of its body. Eligible operations in the body run in the specified dtype; unsupported operations fall back to float32 automatically. Autocast restores the previous dtype, enabled flag, and cache setting when the body exits, even on error.

;; Default: CUDA, bfloat16
(amp/autocast {:device :cuda :dtype :bfloat16}
  (nn/forward model input))

;; Explicit float16 for older hardware
(amp/autocast {:device :cuda :dtype :float16}
  (nn/forward model input))

;; CPU autocast (useful for inference profiling)
(amp/autocast {:device :cpu :dtype :bfloat16}
  (nn/forward model input))

Options:

Key	Default	Description
:device	:cuda	:cuda or :cpu
:dtype	:bfloat16	:float16 or :bfloat16
:enabled?	true	Set false to disable without removing the call site
:cache?	true	Whether to cache autocast type promotions

---

amp/grad-scaler

Float16 has a limited dynamic range; gradients near the minimum representable value underflow to zero. The grad scaler multiplies the loss by a large scale factor before backward, and then divides the gradients back before the optimizer step. If gradients overflow to Inf or NaN, the scaler discards that step and reduces the scale.

;; Default scale: 65536.0, growth factor 2.0, backoff 0.5
(def scaler (amp/grad-scaler))

;; Custom initial scale
(def scaler (amp/grad-scaler {:initial-scale 32768.0}))

Constructor options:

Key	Default	Description
:initial-scale	65536.0	Starting loss scale
:growth-factor	2.0	Multiplier after growth-interval clean steps
:backoff-factor	0.5	Multiplier after an overflow step
:growth-interval	2000	Steps between scale increases
:enabled?	true	Pass false for bfloat16 (no scaling needed)

For bfloat16 training, create the scaler with {:enabled? false} or skip it entirely. Bfloat16 has the same dynamic range as float32, so loss scaling is not required.

---

Scaling and Backward

amp/backward! multiplies the loss by the current scale and calls .backward. It does not unscale gradients; that happens inside amp/step!.

(amp/backward! scaler loss)

---

amp/step!

amp/step! performs the following atomically:

1. Collects all gradients from the optimizer’s parameter list.
2. Checks whether every gradient is finite (no Inf or NaN).
3. In distributed training, performs an all-reduce of the finite flag so the decision is consistent across all ranks.
4. If all gradients are finite, unscales them by dividing by the current scale and calls .step on the optimizer.
5. Updates the dynamic scale: increases it after growth-interval clean steps, decreases it after an overflow.
6. Returns true if the optimizer stepped, false if the step was skipped.

(amp/step! scaler optimizer)

---

float16 vs bfloat16

float16

• Smaller dynamic range (1e-4 to 65504)
• Requires dynamic loss scaling to avoid underflow
• Faster on older Pascal/Volta hardware
• Use amp/grad-scaler and amp/backward!

(def scaler (amp/grad-scaler {:initial-scale 65536.0}))

(let [loss (amp/autocast {:device :cuda :dtype :float16}
             (compute-loss model input target))]
  (amp/backward! scaler loss)
  (amp/step! scaler optimizer))

bfloat16

• Same dynamic range as float32 (no underflow risk)
• Simpler to use — no scaler required
• Recommended for modern Ampere/Hopper hardware
• Use autocast alone without a scaler

(let [loss (amp/autocast {:device :cuda :dtype :bfloat16}
             (compute-loss model input target))]
  (.backward loss)
  (.step optimizer))

---

Full AMP Training Loop

The following example shows a complete float16 AMP loop including gradient accumulation.

(require '[clorch.amp :as amp]
         '[clorch.nn :as nn]
         '[clorch.nn.functional :as F]
         '[clorch.optim :as optim]
         '[clorch.torch :as t])

(def model     (nn/to my-model :cuda))
(def optimizer (optim/adamw (nn/parameters model) :lr 3e-4))
(def scaler    (amp/grad-scaler {:initial-scale 65536.0}))

(dotimes [epoch epochs]
  (doseq [batch dataloader]
    (t/with-torch
      (optim/zero-grad optimizer)
      (let [loss (amp/autocast {:device :cuda :dtype :float16}
                   (F/cross-entropy
                    (nn/forward model (:data batch))
                    (:target batch)))]
        (amp/backward! scaler loss)
        (amp/step! scaler optimizer)))))

---

Integration with DDP

When combining AMP and DDP, pass the scaler to ddp/optimizer-step! instead of calling amp/step! directly. This ensures the overflow flag is synchronized across all distributed ranks before any rank steps the optimizer.

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

(with-open [parallel-model (ddp/distributed-data-parallel model {})]
  (doseq [micro-batches (partition-all accumulation batches)]
    (optim/zero-grad optimizer)
    (doseq [[idx batch] (map-indexed vector micro-batches)]
      (let [final? (= idx (dec (count micro-batches)))
            train! (fn []
                     (t/with-torch
                       (let [loss (amp/autocast {:device :cuda :dtype :float16}
                                    (F/mse-loss (nn/forward parallel-model (:input batch))
                                                (:target batch)))]
                         (amp/backward! scaler loss))))]
        (if final?
          (train!)
          (ddp/no-sync (train!)))))
    ;; scaler is passed here so DDP can synchronize the finite flag
    (ddp/optimizer-step! parallel-model optimizer {:scaler scaler})))

---

Inspecting Scaler State

;; Current loss scale
(amp/current-scale scaler)
;; → 65536.0

;; Full state (for logging or checkpointing)
(amp/scaler-state scaler)
;; → {:scale 65536.0 :growth-tracker 42}

The scaler state is automatically saved and restored by dist/save-checkpoint! and dist/load-checkpoint!.