feat: add reusable transformer, checkpoint, and SafeTensors support

PORTING_STATUS.md

@@ -152,6 +152,12 @@ This document tracks what's been ported from PyTorch to Clorch, and what still n
 | **Custom Models** | [x] |
 | - `SwiGLU` | [x] |
 | - `GroupedQueryAttention` | [x] |
+| **Transformer Modules** | [x] |
+| - `multihead-attention` | [x] |
+| - `transformer-encoder-layer` | [x] |
+| - `transformer-encoder` | [x] |
+| **State Dictionaries** | [x] |
+| - `load-flat-state-dict!`, `flat-parameter-names` | [x] |
 
 ---
 
@@ -176,6 +182,7 @@ This document tracks what's been ported from PyTorch to Clorch, and what still n
 | - `cosine-similarity`, `pairwise-distance` | [x] |
 | **Fused Attention** | [x] |
 | - `scaled-dot-product-attention` | [x] |
+| - `combine-attention-masks` | [x] |
 
 ---
 
@@ -194,6 +201,7 @@ This document tracks what's been ported from PyTorch to Clorch, and what still n
 |---------|--------|
 | `xavier-uniform!`, `xavier-normal!` | [x] |
 | `kaiming-uniform!`, `kaiming-normal!` | [x] |
+| `trunc-normal!` | [x] |
 | `normal!`, `uniform!`, `constant!`, `zeros!`, `ones!` | [x] |
 
 ---
@@ -243,6 +251,17 @@ This document tracks what's been ported from PyTorch to Clorch, and what still n
 | Rank-zero model, optimizer, sampler, scaler, CPU RNG, and training checkpoints | [x] |
 ---
 
+### SafeTensors (`clorch.safetensors`)
+
+| Feature | Status |
+|---------|--------|
+| `load-file`, `save-file!`, `metadata` | [x] |
+| `load-state-dict!` via flat dotted parameter names | [x] |
+| All specification dtypes (bool through complex double) | [x] |
+| Validated headers, non-overlapping offsets, atomic writes | [x] |
+
+---
+
 ## What's Not Yet Implemented
 
 ### Specialized Tensors & Ops
@@ -266,8 +285,6 @@ This document tracks what's been ported from PyTorch to Clorch, and what still n
 
 | Feature | Status |
 |---------|--------|
-| **Native Transformer Layers** | ❌ |
-| (Clojure implementations exist in examples) | |
 | **Container Modules** | ❌ |
 | - `ModuleList`, `ModuleDict`, `ParameterList` | ❌ |
 | **Utilities** | ❌ |
@@ -299,4 +316,4 @@ This document tracks what's been ported from PyTorch to Clorch, and what still n
 
 ---
 
-*Last updated: 2026-08-08*
+*Last updated: 2026-10-18*

README.md

@@ -8,7 +8,7 @@
   <a href="https://github.com/antlobach/clorch/actions/workflows/ci.yml"><img src="https://github.com/antlobach/clorch/actions/workflows/ci.yml/badge.svg" alt="CI"></a>
   <a href="https://antlobach.github.io/clorch/"><img src="https://img.shields.io/badge/docs-online-4051B5?logo=materialformkdocs&logoColor=white" alt="Documentation"></a>
   <a href="https://antlobach-clorch-182a83cb.mintlify.app"><img src="https://img.shields.io/badge/mintlify-docs-18E299?logo=mintlify&logoColor=white" alt="Mintlify documentation"></a>
-  <a href="https://deepwiki.com/antlobach/clorch"><img src="https://deepwiki.com/badge.svg" alt="Ask DeepWiki"></a>
+  <a href="https://deepwiki.com/antlobach/clorch"><img src="https://img.shields.io/badge/Ask-DeepWiki-2B2B2B" alt="Ask DeepWiki"></a>
   <a href="https://context7.com/antlobach/clorch"><img src="https://img.shields.io/badge/context7-docs-4051B5" alt="Context7 documentation"></a>
   <a href="LICENSE"><img src="https://img.shields.io/badge/license-MIT-4c1.svg" alt="MIT license"></a>
 </p>

benchmarks/transformer_clorch.clj

@@ -0,0 +1,41 @@
+(ns transformer-clorch
+  (:require [clorch.autograd :as autograd]
+            [clorch.nn :as nn]
+            [clorch.torch :as t])
+  (:import [org.bytedeco.pytorch.global torch]))
+
+(set! *warn-on-reflection* true)
+
+(def batch-size 4)
+(def sequence-length 256)
+(def dimension 512)
+(def heads 8)
+(def warmup 5)
+(def iterations 20)
+
+(defn- median [samples]
+  (let [ordered (vec (sort samples))]
+    (nth ordered (quot (count ordered) 2))))
+
+(defn benchmark []
+  (torch/set_num_threads 1)
+  (t/manual-seed 23)
+  (let [model (nn/multihead-attention dimension heads :dropout 0.0)
+        x (t/randn [batch-size sequence-length dimension])]
+    (nn/train model false)
+    (autograd/no-grad
+     (dotimes [_ warmup]
+       (t/with-torch
+         (nn/forward model x)
+         nil))
+     (let [samples
+           (repeatedly
+            iterations
+            (fn []
+              (let [started (System/nanoTime)]
+                (t/with-torch
+                  (nn/forward model x)
+                  nil)
+                (/ (double (- (System/nanoTime) started)) 1000000.0))))]
+       {:runtime "clorch"
+        :median-ms (median samples)}))))

benchmarks/transformer_pytorch.py

@@ -0,0 +1,32 @@
+import json
+import statistics
+import time
+
+import torch
+
+
+torch.set_num_threads(1)
+torch.manual_seed(23)
+BATCH = 4
+SEQUENCE = 256
+DIMENSION = 512
+HEADS = 8
+WARMUP = 5
+ITERATIONS = 20
+
+model = torch.nn.MultiheadAttention(
+    DIMENSION, HEADS, dropout=0.0, batch_first=True
+).eval()
+x = torch.randn(BATCH, SEQUENCE, DIMENSION)
+
+with torch.inference_mode():
+    for _ in range(WARMUP):
+        model(x, x, x, need_weights=False)
+
+    samples = []
+    for _ in range(ITERATIONS):
+        started = time.perf_counter_ns()
+        model(x, x, x, need_weights=False)
+        samples.append((time.perf_counter_ns() - started) / 1_000_000)
+
+print(json.dumps({"runtime": "pytorch", "median_ms": statistics.median(samples)}))

docs/functional.md

@@ -49,6 +49,28 @@ The `clorch.nn.functional` namespace (aliased as `F`) contains functional versio
 (F/pad input [1 1 1 1] :mode :reflect)
 ```
 
+### Attention
+
+```clojure
+;; LibTorch's fused scaled-dot-product attention
+;; query/key/value: [batch, heads, seq, head-dim]
+(F/scaled-dot-product-attention q k v
+                                :attention-mask mask     ; boolean or additive float bias
+                                :dropout-p 0.0
+                                :causal? false
+                                :scale scale)            ; optional custom scale
+
+;; Combines an SDPA mask with a module-style key-padding mask.
+;; attention-mask: true participates / float bias (SDPA convention);
+;; key-padding-mask: [batch, key], true is ignored (module convention).
+(F/combine-attention-masks attention-mask key-padding-mask)
+```
+
+`combine-attention-masks` broadcasts the padding mask to
+`[batch, 1, 1, key-length]` and returns a boolean mask when both inputs are
+boolean, or an additive float bias when the attention mask is floating point.
+It returns the attention mask unchanged when no padding mask is given.
+
 ## Loss Functions
 See the [Loss Functions](losses.md) guide for a full list.
 

docs/init.md

@@ -21,6 +21,18 @@ These functions modify the input tensor in-place.
 
 ## Research-Standard Initializers
 
+### Truncated Normal
+
+Draws from `N(mean, std²)` and redraws values that fall outside `[mean - a·std, mean + a·std]`,
+matching `torch.nn.init.trunc_normal_`.
+
+```clojure
+(init/trunc-normal! x :mean 0.0 :std 0.02 :a -2.0 :b 2.0)
+
+;; Defaults: mean 0.0, std 1.0, bounds ±2 standard deviations
+(init/trunc-normal! x)
+```
+
 ### Xavier (Glorot) Initialization
 Designed for layers with symmetric activation functions (like `tanh` or `sigmoid`).
 

docs/mkdocs.yml

@@ -24,6 +24,7 @@ nav:
       - Loss Functions: losses.md
       - Optimizers: optimizers.md
       - Initialization: init.md
+  - SafeTensors I/O: safetensors.md
   - Autograd: autograd.md
   - Distributed CUDA Training: distributed.md
   - PyTorch Parity: pytorch-parity.md

docs/nn.md

@@ -61,6 +61,68 @@ Directly wraps LibTorch's optimized C++ RNN implementations.
 (nn/constant-pad2d 1 3.14)
 ```
 
+## Attention & Transformers
+
+### Multi-Head Attention
+
+`multihead-attention` mirrors `torch.nn.MultiheadAttention` with a packed QKV
+projection and batch-first inputs.
+
+```clojure
+;; MultiheadAttention(embed_dim=8, num_heads=2)
+(def attention (nn/multihead-attention 8 2 :dropout 0.0 :bias true))
+
+;; Self-attention
+(nn/forward attention x)
+
+;; Cross-attention with masks and cached key/value
+(nn/forward attention
+            {:query query
+             :key context
+             :value context
+             :attention-mask additive-or-bool-mask   ; [batch, 1, query, key] or broadcastable
+             :key-padding-mask padding-mask          ; [batch, key]; true = ignored
+             :causal? false
+             :rope {:cos cosines :sin sines :interleaved? true}
+             :cached-key cached-key
+             :cached-value cached-value
+             :return-kv? true})
+```
+
+Calling it with a plain tensor is self-attention. With `:return-kv?` the forward
+returns `{:output ... :key ... :value ...}` where `key`/`value` have shape
+`[batch, heads, key-length, head-dim]` and can be fed back as
+`:cached-key`/`:cached-value` on the next step. `:rope` accepts the pair
+`[cosines sines]` or a map with `:interleaved?`.
+
+### Transformer Encoder Layers and Encoders
+
+`transformer-encoder-layer` mirrors `torch.nn.TransformerEncoderLayer`
+(pre-norm with `:norm-first? true`, the default is post-norm) and
+`transformer-encoder` stacks layers, mirroring `torch.nn.TransformerEncoder`.
+
+```clojure
+(def layer (nn/transformer-encoder-layer 512 8
+                                        :feedforward-dim 2048
+                                        :dropout 0.1
+                                        :activation :gelu
+                                        :norm-first? true))
+(def encoder (nn/transformer-encoder layer 6))
+;; or from an explicit collection:
+(def encoder (nn/transformer-encoder [layer-a layer-b]))
+
+(nn/forward encoder x)
+```
+
+The layer forward accepts `:x` plus the attention options above, including
+`:key`/`:value` for cross-attention and `:cached-key`/`:cached-value`. When
+`:return-kv?` is passed, the encoder also returns the per-layer key/value
+cache used for incremental decoding:
+
+```clojure
+(let [{:keys [output cache]} (nn/forward encoder x :return-kv? true)])
+```
+
 ## Containers & Utilities
 
 ### Sequential

docs/pytorch-parity.md

@@ -10,7 +10,7 @@ Clorch follows PyTorch's tensor, autograd, module, optimizer, data, and device c
 | 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 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 100% figure applies only to the 40 cross-language scenarios exercised by `tests_comparison/compare_torch.py`. The transformer modules and interleaved RoPE are additionally verified against a PyTorch reference by `run_transformer_comparison.sh` (`tests_comparison/compare_transformer.clj`). Clorch will publish another breadth percentage after generating a version-pinned PyTorch 2.10 inventory.
 
 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.
 
@@ -23,11 +23,11 @@ The release suite verifies CPU behavior and a single-GPU CUDA path, including CU
 | Indexing and slicing | Python-style `ix`, negative indices, ellipsis, stepped slices, tensor indices, masks, gather, select, and scatter-reduce | Full advanced-indexing mutation parity, sparse/nested indexing, and every upstream indexing edge case |
 | Linear algebra | Matrix products, decompositions, solving, inverse, determinant, eigen, SVD, QR, Cholesky, and norms | Remaining `torch.linalg` routines, batched edge cases, driver options, and complete complex-number coverage |
 | 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 |
+| Neural-network modules | Linear, convolution, transposed convolution, normalization, pooling, padding, recurrent layers, embeddings, activations, dropout, multi-head attention, transformer encoder layers/encoders, and custom `defmodel` modules | Complete container modules, remaining attention/transformer options, 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, 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 |
+| Serialization and JIT | Tensor/model save-load, state dictionaries, weight save-load, JIT load-save-forward, atomic rank-zero training checkpoints, and SafeTensors archives (`clorch.safetensors`) | `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 | 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 |
@@ -42,6 +42,8 @@ The release suite verifies CPU behavior and a single-GPU CUDA path, including CU
 | RMSNorm | Implemented | `nn/rmsnorm` |
 | Rotary position embeddings | Implemented | `torch/precompute-rope-freqs`, `torch/apply-rope` |
 | Grouped-query attention | Implemented | `nn/GroupedQueryAttention` |
+| Multi-head attention with KV caching | Implemented | `nn/multihead-attention` |
+| Transformer encoder blocks | Implemented | `nn/transformer-encoder-layer`, `nn/transformer-encoder` |
 | SwiGLU feed-forward block | Implemented | `nn/SwiGLU` |
 | Causal attention masks | Implemented | Tensor masking operations and Llama/GPT examples |
 | KV cache | Implemented in Llama-style model flow | `examples/modern_llama.clj`, `examples/nanochat.clj` |

docs/safetensors.md

@@ -0,0 +1,46 @@
+# SafeTensors I/O
+
+The `clorch.safetensors` namespace reads and writes [SafeTensors](https://huggingface.co/docs/safetensors) archives without Python, pickle execution, or Hugging Face tooling. Every dtype in the SafeTensors specification is supported.
+
+## Reading
+
+```clojure
+(require '[clorch.safetensors :as safe])
+
+(def tensors (safe/load-file "model.safetensors"))
+;; => {"layer.weight" #object[Tensor ...], ...}
+```
+
+- Returns a map of string tensor names to CPU tensors with the shapes and dtypes stored in the file.
+- Tensor data is validated before reading: dtype, shape, byte counts, and non-overlapping offsets covering the whole payload.
+- The returned tensors are caller-owned; release them with their enclosing native-memory scope (see [Memory Management](memory.md)).
+- Pickle payloads are never executed; only the JSON header and raw tensor bytes are read.
+
+```clojure
+(safe/metadata "model.safetensors")
+;; => {"format" "pt"} — the file's __metadata__ map, {} when absent
+```
+
+## Writing
+
+```clojure
+(safe/save-file! tensors "out.safetensors"
+                 :metadata {"format" "pt" "source" "clorch"})
+```
+
+- Tensors are detached, made contiguous, and copied to CPU before their bytes are written.
+- Optional `:metadata` must map strings to strings.
+- The file is written to a temporary file in the destination directory and moved into place atomically, so readers never observe a partial archive.
+
+## Loading Model Weights
+
+`load-state-dict!` loads an archive whose tensor names are the model's flat dotted parameter paths (see `nn/flat-parameter-names`):
+
+```clojure
+(nn/flat-parameter-names model)
+;; => #{"l1.weight" "l1.bias" "l2.weight" "l2.bias"}
+
+(safe/load-state-dict! model "model.safetensors")
+```
+
+Every supplied name is copied into the matching parameter; missing or unexpected names raise an error instead of being skipped. To adapt foreign checkpoint names to your model's names, load with `safe/load-file`, transform the map, and call `nn/load-flat-state-dict!` directly.

docs/tensors.md

@@ -207,3 +207,27 @@ A built-in utility for high-quality text generation. It filters the probability
 ```clojure
 (t/top-p (t/tensor [[0.1 0.8 0.1]]) 0.9)
 ```
+
+## Rotary Position Embeddings
+
+```clojure
+;; Cosine/sine tables of shape [1, sequence-length, head-dim].
+;; theta is the RoPE base (default 10000.0).
+;;
+;; Default layout rotates half-split pairs (0, D/2), (1, D/2+1), ...
+(t/precompute-rope-freqs head-dim seq-len)
+
+;; :interleaved? rotates adjacent pairs (0, 1), (2, 3), ...
+;; (the layout used by Llama-family and TabICL models)
+(t/precompute-rope-freqs head-dim seq-len :theta 100000.0 :interleaved? true)
+
+;; Build tables on the same device as CUDA inputs.
+(t/precompute-rope-freqs head-dim seq-len :device :cuda)
+
+;; Applies rope to [batch, sequence, heads, head-dim]; pass :interleaved? true
+;; when the tables were precomputed with :interleaved? true.
+(t/apply-rope x cosines sines :interleaved? true)
+```
+
+The `:interleaved?` flag must match between `precompute-rope-freqs` and
+`apply-rope`.

run_transformer_benchmark.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+set -euo pipefail
+
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+
+uv run --python 3.11 --no-project \
+  --with "torch==2.1.2" \
+  python benchmarks/transformer_pytorch.py
+
+clojure -Sthreads 1 \
+  -Sdeps '{:paths ["src" "benchmarks"] :deps {org.bytedeco/pytorch-platform {:mvn/version "2.7.1-1.5.12"} org.bytedeco/pytorch-platform-gpu {:mvn/version "2.7.1-1.5.12"}}}' \
+  -M -e '(require (quote transformer-clorch)) (prn (transformer-clorch/benchmark))'