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

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

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

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

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

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

***

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

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

```clojure theme={null}
(amp/step! scaler optimizer)
```

***

## float16 vs bfloat16

<Tabs>
  <Tab title="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!`

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

  <Tab title="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

    ```clojure theme={null}
    (let [loss (amp/autocast {:device :cuda :dtype :bfloat16}
                 (compute-loss model input target))]
      (.backward loss)
      (.step optimizer))
    ```
  </Tab>
</Tabs>

***

## Full AMP Training Loop

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

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

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

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


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