From 388034a380f759194c46c1bb5de333019ba4ad60 Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Wed, 14 Oct 2026 23:21:08 +0200 Subject: [PATCH 1/7] feat: add checkpoint loading primitives --- README.md | 2 +- src/clorch/nn.clj | 73 +++++++++++++++++++++++++++++++++++ src/clorch/nn/init.clj | 34 +++++++++++++++- src/clorch/torch.clj | 5 ++- test/clorch/nn_init_test.clj | 20 +++++++--- test/clorch/nn_model_test.clj | 13 ++++++- test/clorch/torch_test.clj | 6 +++ 7 files changed, 142 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index aa07450..419c9a0 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ CI Documentation Mintlify documentation - Ask DeepWiki + Ask DeepWiki Context7 documentation MIT license

diff --git a/src/clorch/nn.clj b/src/clorch/nn.clj index 9f422b5..4b7ceb5 100644 --- a/src/clorch/nn.clj +++ b/src/clorch/nn.clj @@ -408,6 +408,79 @@ (walk thing prefix) @res)))) +(defn load-flat-state-dict! + "Copies a flat, dotted-path state dictionary into a model. + + Native parameter views are consumed while their owning parameter list is alive, + avoiding borrowed JavaCPP pointers escaping the traversal." + [thing values] + (let [loaded (atom #{})] + (letfn [(copy-value! [path destination] + (when-some [source (get values path)] + (torch/copy- destination source) + (swap! loaded conj path))) + (walk [node path] + (cond + (instance? Module node) + (let [named-p (.named_parameters ^Module node) + pairs (.pairs named-p)] + (dotimes [i (.size named-p)] + (let [parameter-name (.getString (.first pairs (clojure.core/long i))) + full-name (if (empty? path) + parameter-name + (str path "." parameter-name))] + (copy-value! full-name (.second pairs (clojure.core/long i)))))) + + (instance? Parameter node) + (copy-value! path (:tensor node)) + + (instance? Tensor node) + (copy-value! path node) + + (map? node) + (doseq [[k v] node] + (walk v (if (empty? path) + (clojure.core/name k) + (str path "." (clojure.core/name k))))) + + (vector? node) + (doseq [[i v] (map-indexed vector node)] + (walk v (str path "[" i "]")))))] + (walk thing "") + @loaded))) + +(defn flat-parameter-names + "Returns the flat paths accepted by load-flat-state-dict! without retaining + borrowed native parameter views." + [thing] + (let [names (atom #{})] + (letfn [(walk [node path] + (cond + (instance? Module node) + (let [named-p (.named_parameters ^Module node) + pairs (.pairs named-p)] + (dotimes [i (.size named-p)] + (let [parameter-name (.getString (.first pairs (clojure.core/long i)))] + (swap! names conj + (if (empty? path) + parameter-name + (str path "." parameter-name)))))) + + (or (instance? Parameter node) (instance? Tensor node)) + (swap! names conj path) + + (map? node) + (doseq [[k v] node] + (walk v (if (empty? path) + (clojure.core/name k) + (str path "." (clojure.core/name k))))) + + (vector? node) + (doseq [[i v] (map-indexed vector node)] + (walk v (str path "[" i "]")))))] + (walk thing "") + @names))) + ;; --- Layer Implementations --- (defrecord LayerNormRecord [gamma beta eps] diff --git a/src/clorch/nn/init.clj b/src/clorch/nn/init.clj index 50a49bf..83db59e 100644 --- a/src/clorch/nn/init.clj +++ b/src/clorch/nn/init.clj @@ -33,12 +33,12 @@ (torch/xavier_normal_ (t/->tensor tensor) (double (or gain 1.0))))) (defn kaiming-uniform! [tensor & {:keys [a mode non-linearity] - :or {a 0 mode :fan-in non-linearity :leaky-relu}}] + :or {a 0 mode :fan-in non-linearity :leaky-relu}}] (no-grad (torch/kaiming_uniform_ (t/->tensor tensor) (double a) (->fan-mode mode) (->non-linearity non-linearity)))) (defn kaiming-normal! [tensor & {:keys [a mode non-linearity] - :or {a 0 mode :fan-in non-linearity :leaky-relu}}] + :or {a 0 mode :fan-in non-linearity :leaky-relu}}] (no-grad (torch/kaiming_normal_ (t/->tensor tensor) (double a) (->fan-mode mode) (->non-linearity non-linearity)))) @@ -58,6 +58,36 @@ (no-grad (torch/zeros_ (t/->tensor tensor)))) +(defn- normal-cdf [value] + (t/with-torch + (* 0.5 + (+ 1.0 + (double + (t/item-float + (t/erf (t/tensor (/ (double value) (Math/sqrt 2.0)))))))))) + +(defn trunc-normal! + "Fills a tensor from N(mean, std²), redrawn outside [a, b]. + + Uses the inverse-CDF algorithm from torch.nn.init.trunc_normal_." + [tensor & {:keys [mean std a b] + :or {mean 0.0 std 1.0 a -2.0 b 2.0}}] + (when-not (and (pos? std) (< a b)) + (throw (IllegalArgumentException. + "trunc-normal! requires a positive std and a < b"))) + (let [lower (normal-cdf (/ (- a mean) std)) + upper (normal-cdf (/ (- b mean) std)) + native (t/->tensor tensor)] + (no-grad + (torch/uniform_ native (- (* 2.0 lower) 1.0) (- (* 2.0 upper) 1.0)) + (t/copy- native + (-> native + t/erfinv + (t/mul (* std (Math/sqrt 2.0))) + (t/add mean) + (t/clamp a b))))) + tensor) + (defn ones! [tensor] (no-grad (torch/ones_ (t/->tensor tensor)))) diff --git a/src/clorch/torch.clj b/src/clorch/torch.clj index 3ce93e9..25ac038 100644 --- a/src/clorch/torch.clj +++ b/src/clorch/torch.clj @@ -29,6 +29,7 @@ ;; --- Memory Management --- (def ^:private ^ThreadLocal session-scope (ThreadLocal.)) +(def ^:private float-array-class (class (float-array 0))) (defn release! "Manually deletes the native C++ memory associated with a tensor or collection of tensors. @@ -193,6 +194,8 @@ (instance? Tensor data) data (and (record? data) (:tensor data)) (:tensor data) (number? data) (torch/tensor (clojure.core/float data)) + (instance? float-array-class data) + (torch/tensor ^floats data (org.bytedeco.pytorch.TensorOptions.)) (or (vector? data) (seq? data)) (let [shape (loop [curr data s []] (if (coll? curr) (recur (first curr) (conj s (count curr))) s)) flat-data (clojure.core/flatten data)] @@ -1613,7 +1616,7 @@ (let [dst-t (->tensor dst) src-t (->tensor src)] (with-open [_ (NoGradGuard.)] - (.set_ dst-t src-t)) + (.copy_ dst-t src-t)) dst-t)) (defn multinomial diff --git a/test/clorch/nn_init_test.clj b/test/clorch/nn_init_test.clj index 34fd08b..9cc1a26 100644 --- a/test/clorch/nn_init_test.clj +++ b/test/clorch/nn_init_test.clj @@ -13,28 +13,36 @@ (testing "ones!" (init/ones! x) (is (every? #(= 1.0 %) (mapv t/item-float (t/tseq (t/reshape x [-1])))))) - + (testing "zeros!" (init/zeros! x) (is (every? #(= 0.0 %) (mapv t/item-float (t/tseq (t/reshape x [-1])))))) - + (testing "constant!" (init/constant! x 3.14) (is (every? #(near? % 3.14) (mapv t/item-float (t/tseq (t/reshape x [-1])))))) - + (testing "uniform!" (init/uniform! x -1 1) (let [v (mapv t/item-float (t/tseq (t/reshape x [-1])))] (is (every? #(and (>= % -1) (<= % 1)) v)))) - + (testing "normal!" (init/normal! x 0 1) (is (some? x))) - + (testing "xavier-uniform!" (init/xavier-uniform! x) (is (some? x))) - + (testing "kaiming-normal!" (init/kaiming-normal! x :non-linearity :relu) (is (some? x)))))) + +(deftest truncated-normal-test + (t/with-torch + (let [x (t/empty [10000])] + (init/trunc-normal! x :mean 0.5 :std 0.25 :a 0.0 :b 1.0) + (is (>= (t/item-float (t/min x)) 0.0)) + (is (<= (t/item-float (t/max x)) 1.0)) + (is (< (Math/abs (- 0.5 (t/item-float (t/mean x)))) 0.02))))) diff --git a/test/clorch/nn_model_test.clj b/test/clorch/nn_model_test.clj index a20d732..cd75d2e 100644 --- a/test/clorch/nn_model_test.clj +++ b/test/clorch/nn_model_test.clj @@ -2,7 +2,6 @@ (:require [clorch.torch :as t] [clorch.nn :as nn] [clorch.nn.functional :as F] - [clorch.autograd :as autograd] [clojure.test :refer [deftest is testing]])) (deftest model-lifecycle-test @@ -27,6 +26,7 @@ (deftest custom-model-test (t/with-torch + #_{:clj-kondo/ignore [:unresolved-symbol]} (nn/defmodel MyModel [in out] [l1 (nn/linear in 16) l2 (nn/linear 16 out)] @@ -40,3 +40,14 @@ out (nn/forward m x)] (is (= [2 5] (t/size out))) (is (= 4 (.size (nn/parameters m))))))) + +(deftest flat-state-dict-loading-test + (t/with-torch + (let [model (nn/sequential (nn/linear 3 2)) + names (nn/flat-parameter-names model) + state {"[0].weight" (t/full [2 3] 2.0) + "[0].bias" (t/full [2] -1.0)} + loaded (nn/load-flat-state-dict! model state) + output (nn/forward model (t/ones [1 3]))] + (is (= #{"[0].weight" "[0].bias"} names loaded)) + (is (= 10.0 (double (t/item-float (t/sum output)))))))) diff --git a/test/clorch/torch_test.clj b/test/clorch/torch_test.clj index 96aaafd..b8e66d4 100644 --- a/test/clorch/torch_test.clj +++ b/test/clorch/torch_test.clj @@ -108,6 +108,12 @@ (is (= [2] (torch/size t))) (is (near? 3.75 (torch/item-float (torch/sum t)))))))) +(deftest primitive-float-array-construction-test + (torch/with-torch + (let [tensor (torch/tensor (float-array [1.25 -2.5 3.75]))] + (is (= [3] (torch/size tensor))) + (is (near? 2.5 (torch/item-float (torch/sum tensor))))))) + (deftest math-ops-test (testing "Basic arithmetic value check" (torch/with-torch From cf9312fa1fe588b0f7e30a5ee565a62132706857 Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Thu, 15 Oct 2026 00:15:59 +0200 Subject: [PATCH 2/7] feat: add reusable transformer modules --- benchmarks/transformer_clorch.clj | 41 ++++ benchmarks/transformer_pytorch.py | 32 +++ run_transformer_benchmark.sh | 13 ++ run_transformer_comparison.sh | 13 ++ src/clorch/nn.clj | 243 ++++++++++++++++++++++ src/clorch/nn/functional.clj | 42 +++- src/clorch/torch.clj | 46 ++-- test/clorch/nn_transformer_test.clj | 91 ++++++++ tests_comparison/compare_transformer.clj | 87 ++++++++ tests_comparison/transformer_reference.py | 97 +++++++++ 10 files changed, 689 insertions(+), 16 deletions(-) create mode 100644 benchmarks/transformer_clorch.clj create mode 100644 benchmarks/transformer_pytorch.py create mode 100755 run_transformer_benchmark.sh create mode 100755 run_transformer_comparison.sh create mode 100644 test/clorch/nn_transformer_test.clj create mode 100644 tests_comparison/compare_transformer.clj create mode 100644 tests_comparison/transformer_reference.py diff --git a/benchmarks/transformer_clorch.clj b/benchmarks/transformer_clorch.clj new file mode 100644 index 0000000..8128c23 --- /dev/null +++ b/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)})))) diff --git a/benchmarks/transformer_pytorch.py b/benchmarks/transformer_pytorch.py new file mode 100644 index 0000000..3c84251 --- /dev/null +++ b/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)})) diff --git a/run_transformer_benchmark.sh b/run_transformer_benchmark.sh new file mode 100755 index 0000000..cb1ca75 --- /dev/null +++ b/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 ["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))' diff --git a/run_transformer_comparison.sh b/run_transformer_comparison.sh new file mode 100755 index 0000000..374a83e --- /dev/null +++ b/run_transformer_comparison.sh @@ -0,0 +1,13 @@ +#!/bin/bash +set -euo pipefail + +reference_file="$(mktemp -t clorch-transformer-reference.XXXXXX.json)" +trap 'rm -f "$reference_file"' EXIT + +uv run --python 3.11 --no-project \ + --with "torch==2.1.2" \ + python tests_comparison/transformer_reference.py "$reference_file" + +clojure -Sthreads 1 \ + -Sdeps '{:paths ["tests_comparison"] :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 -m compare-transformer "$reference_file" diff --git a/src/clorch/nn.clj b/src/clorch/nn.clj index 4b7ceb5..9b363a5 100644 --- a/src/clorch/nn.clj +++ b/src/clorch/nn.clj @@ -3,6 +3,7 @@ (:require [clojure.string :as str] [clorch.torch :as torch] [clorch.nn.functional :as F] + [clorch.nn.init :as init] [clorch.autograd :as autograd]) (:import [org.bytedeco.pytorch LinearImpl LinearOptions ReLUImpl SigmoidImpl SoftmaxImpl DropoutImpl DropoutOptions FlattenImpl Conv1dImpl Conv1dOptions Conv2dImpl Conv2dOptions Conv3dImpl Conv3dOptions ConvTranspose1dImpl ConvTranspose1dOptions ConvTranspose2dImpl ConvTranspose2dOptions ConvTranspose3dImpl ConvTranspose3dOptions BatchNorm1dImpl BatchNorm2dImpl BatchNorm3dImpl BatchNormOptions GroupNormImpl GroupNormOptions InstanceNorm1dImpl InstanceNorm2dImpl InstanceNorm3dImpl InstanceNormOptions RNNImpl RNNOptions LSTMImpl LSTMOptions GRUImpl GRUOptions EmbeddingImpl EmbeddingOptions Module Tensor TensorVector TanhImpl ELUImpl SELUImpl CELUImpl CELUOptions ELUOptions LeakyReLUImpl LeakyReLUOptions ReLU6Impl SoftplusImpl SoftplusOptions SoftsignImpl SoftminImpl MishImpl HardtanhImpl HardtanhOptions LogSigmoidImpl LogSoftmaxImpl HardshrinkImpl HardshrinkOptions SoftshrinkImpl SoftshrinkOptions TanhshrinkImpl ThresholdImpl ThresholdOptions GLUImpl GLUOptions RReLUImpl RReLUOptions AlphaDropoutImpl FeatureAlphaDropoutImpl BilinearImpl BilinearOptions] [org.bytedeco.pytorch.global torch$ScalarType] @@ -885,6 +886,248 @@ (defn ~model-name ~args (let [~@init-bindings] (~(symbol (str "->" record-name)) ~@field-syms)))))) +(defn- apply-attention-rope + [tensor rope] + (if-not rope + tensor + (let [{:keys [cos sin interleaved?]} + (if (map? rope) + rope + {:cos (first rope) :sin (second rope)})] + (torch/apply-rope tensor cos sin :interleaved? (boolean interleaved?))))) + +(defn- project-attention-inputs + [^LinearImpl in-projection bias? dimension query attention-key value] + (if (and (identical? query attention-key) + (identical? query value)) + (torch/chunk (forward in-projection query) 3 -1) + (let [weight (.weight in-projection) + bias (when bias? (.bias in-projection)) + q-weight (torch/ix weight [0 dimension] :_) + q-bias (when bias (torch/ix bias [0 dimension])) + q (F/linear query q-weight q-bias)] + (if (identical? attention-key value) + (let [kv-weight (torch/ix weight [dimension (* 3 dimension)] :_) + kv-bias (when bias (torch/ix bias [dimension (* 3 dimension)])) + [k v] (torch/chunk + (F/linear attention-key kv-weight kv-bias) + 2 + -1)] + [q k v]) + (let [k-weight (torch/ix weight [dimension (* 2 dimension)] :_) + v-weight (torch/ix weight [(* 2 dimension) (* 3 dimension)] :_) + k-bias (when bias (torch/ix bias [dimension (* 2 dimension)])) + v-bias (when bias (torch/ix bias [(* 2 dimension) (* 3 dimension)]))] + [q + (F/linear attention-key k-weight k-bias) + (F/linear value v-weight v-bias)]))))) + +(defrecord MultiHeadAttentionRecord + [in-projection out-projection dimension head-count head-dim + dropout-p bias? training?] + IModule + (-forward [_ input] + (let [input (if (map? input) input {:query input}) + query (torch/->tensor (:query input)) + attention-key (torch/->tensor (or (:key input) query)) + value (torch/->tensor (or (:value input) attention-key)) + cached-key (:cached-key input) + cached-value (:cached-value input) + _ (when (not= (boolean cached-key) (boolean cached-value)) + (throw (IllegalArgumentException. + "cached-key and cached-value must be supplied together"))) + [batch-size query-length _] (torch/size query) + [q projected-key projected-value] + (if cached-key + [(let [[q _ _] (project-attention-inputs + in-projection bias? dimension query query query)] + q) + nil + nil] + (project-attention-inputs + in-projection bias? dimension query attention-key value)) + key-length (if cached-key + (nth (torch/size cached-key) 2) + (second (torch/size projected-key))) + q (torch/reshape q + [batch-size query-length head-count head-dim]) + k (when projected-key + (torch/reshape projected-key + [batch-size key-length head-count head-dim])) + v (when projected-value + (torch/reshape projected-value + [batch-size key-length head-count head-dim])) + rope (:rope input) + q (apply-attention-rope q rope) + k (when k (apply-attention-rope k rope)) + q (torch/transpose q 1 2) + k (or cached-key (torch/transpose k 1 2)) + v (or cached-value (torch/transpose v 1 2)) + attention-mask + (F/combine-attention-masks (:attention-mask input) + (:key-padding-mask input)) + attended + (F/scaled-dot-product-attention + q k v + :attention-mask attention-mask + :dropout-p (if @training? dropout-p 0.0) + :causal? (boolean (:causal? input))) + merged (torch/reshape + (torch/transpose attended 1 2) + [batch-size query-length dimension]) + output (forward out-projection merged)] + (if (:return-kv? input) + {:output output :key k :value v} + output))) + (-train [this training] + (train in-projection training) + (train out-projection training) + (reset! training? (boolean training)) + this) + (-to [this dtype-or-device] + (assoc this + :in-projection (to in-projection dtype-or-device) + :out-projection (to out-projection dtype-or-device))) + clojure.lang.IFn + (invoke [this input] (forward this input))) + +(defn multihead-attention + "Creates batch-first multi-head attention with packed QKV projection. + + Forward accepts a tensor for self-attention or a map with `:query`, optional + `:key`/`:value`, `:attention-mask`, `:key-padding-mask`, `:causal?`, `:rope`, + `:cached-key`/`:cached-value`, and `:return-kv?`." + [dimension head-count & {dropout-probability :dropout + bias :bias + :or {dropout-probability 0.0 bias true}}] + (when-not (zero? (mod dimension head-count)) + (throw (IllegalArgumentException. + "dimension must be divisible by head-count"))) + (let [in-projection (linear dimension (* 3 dimension) :bias bias) + out-projection (linear dimension dimension :bias bias)] + (init/xavier-uniform! (.weight ^LinearImpl in-projection)) + (when bias + (init/zeros! (.bias ^LinearImpl in-projection))) + (->MultiHeadAttentionRecord + in-projection out-projection dimension head-count + (quot dimension head-count) (double dropout-probability) (boolean bias) (atom true)))) + +(defn- transformer-activation [activation] + (case activation + :relu (relu) + :gelu (gelu) + (if (satisfies? IModule activation) + activation + (throw (IllegalArgumentException. + "activation must be :relu, :gelu, or an IModule"))))) + +#_{:clj-kondo/ignore [:unresolved-symbol]} +(defmodel TransformerEncoderLayer + [dimension head-count & {:keys [feedforward-dim dropout activation + norm-first? bias] + :or {dropout 0.0 activation :relu + norm-first? false bias true}}] + [attention (multihead-attention dimension head-count + :dropout dropout :bias bias) + norm1 (layernorm dimension) + norm2 (layernorm dimension) + feedforward [(linear dimension (or feedforward-dim (* 4 dimension)) + :bias bias) + (transformer-activation activation) + (clorch.nn/dropout dropout) + (linear (or feedforward-dim (* 4 dimension)) dimension + :bias bias)] + residual-dropout1 (clorch.nn/dropout dropout) + residual-dropout2 (clorch.nn/dropout dropout) + pre-norm? norm-first?] + (forward [input] + (let [input (if (map? input) input {:x input}) + x (:x input) + cached? (boolean (:cached-key input)) + normalized-query (if pre-norm? (forward norm1 x) x) + normalized-key + (when-not cached? + (when-let [attention-key (:key input)] + (if pre-norm? + (forward norm1 attention-key) + attention-key))) + normalized-value + (when-not cached? + (when-let [value (:value input)] + (if pre-norm? + (forward norm1 value) + value))) + attention-result + (forward attention + (assoc input + :query normalized-query + :key normalized-key + :value normalized-value)) + attended (if (:return-kv? input) + (:output attention-result) + attention-result) + attention-residual + (torch/add x (forward residual-dropout1 attended)) + after-attention + (if pre-norm? + attention-residual + (forward norm1 attention-residual)) + feedforward-input + (if pre-norm? + (forward norm2 after-attention) + after-attention) + output + (torch/add after-attention + (forward residual-dropout2 + (forward feedforward feedforward-input))) + output (if pre-norm? output (forward norm2 output))] + (if (:return-kv? input) + (assoc attention-result :output output) + output)))) + +(defn transformer-encoder-layer + "Creates a batch-first transformer encoder layer. + + The forward map accepts `:x` plus the attention options accepted by + `multihead-attention`. Optional `:key` and `:value` enable cross-attention." + [dimension head-count & options] + (clojure.core/apply TransformerEncoderLayer dimension head-count options)) + +#_{:clj-kondo/ignore [:unresolved-symbol]} +(defmodel TransformerEncoder [blocks] + [layers (vec blocks)] + (forward [input] + (let [input (if (map? input) input {:x input})] + (if (:return-kv? input) + (reduce (fn [{:keys [output cache]} [index layer]] + (let [cached (nth (:cache input) index nil) + result + (forward layer + (cond-> (assoc input :x output) + cached + (assoc :cached-key (:key cached) + :cached-value (:value cached))))] + {:output (:output result) + :cache (conj cache + {:key (:key result) + :value (:value result)})})) + {:output (:x input) :cache []} + (map-indexed vector layers)) + (reduce (fn [state layer] + (forward layer (assoc input :x state))) + (:x input) + layers))))) + +(defn transformer-encoder + "Creates an encoder from a layer collection or a zero-argument layer factory." + ([layers] + (TransformerEncoder layers)) + ([layer-factory layer-count] + (when-not (fn? layer-factory) + (throw (IllegalArgumentException. + "two-argument transformer-encoder requires a layer factory"))) + (TransformerEncoder (repeatedly layer-count layer-factory)))) + (defn generate "Generates tokens autoregressively." [model idx max-new-tokens context-size] diff --git a/src/clorch/nn/functional.clj b/src/clorch/nn/functional.clj index 0be72d8..2088bae 100644 --- a/src/clorch/nn/functional.clj +++ b/src/clorch/nn/functional.clj @@ -121,6 +121,44 @@ (.put (.inplace options) false) (torch/dropout2d (t/->tensor tensor) options))) +(defn combine-attention-masks + "Combines an SDPA attention mask with a key-padding mask. + + `attention-mask` follows `scaled-dot-product-attention`: true boolean values + participate in attention, while floating values are additive logits biases. + `key-padding-mask` follows PyTorch module convention: true values are ignored. + Its shape must be `[batch, key-length]`." + [attention-mask key-padding-mask] + (if-not key-padding-mask + attention-mask + (let [[batch-size key-length :as shape] (t/size key-padding-mask)] + (when-not (= 2 (count shape)) + (throw (IllegalArgumentException. + "key-padding-mask must have shape [batch, key-length]"))) + (let [padding-mask (t/reshape (t/->tensor key-padding-mask) + [batch-size 1 1 key-length]) + allowed (t/logical-not padding-mask)] + (if-not attention-mask + allowed + (let [attention-mask (t/->tensor attention-mask) + scalar-type (.toString (.scalar_type attention-mask))] + (cond + (= "Bool" scalar-type) + (t/logical-and attention-mask allowed) + + (.is_floating_point attention-mask) + (let [padding-bias (t/to + (t/masked-fill + (t/zeros [batch-size 1 1 key-length]) + padding-mask + Double/NEGATIVE_INFINITY) + (t/dtype attention-mask))] + (t/add attention-mask padding-bias)) + + :else + (throw (IllegalArgumentException. + "attention-mask must be boolean or floating point"))))))))) + (defn scaled-dot-product-attention "Runs LibTorch's fused scaled-dot-product attention dispatcher. @@ -271,8 +309,8 @@ (defn pad [input padding & {:keys [mode value] :or {mode :constant value 0.0}}] (let [^longs p-arr (if (number? padding) - (->pad-array padding (- (count (t/size input)) 2)) - (long-array (mapv clojure.core/long padding)))] + (->pad-array padding (- (count (t/size input)) 2)) + (long-array (mapv clojure.core/long padding)))] (torch/pad (t/->tensor input) p-arr (clojure.core/name mode) (org.bytedeco.pytorch.DoubleOptional. (double value))))) (defn interpolate [input & {:keys [size scale-factor mode] :or {mode :nearest}}] diff --git a/src/clorch/torch.clj b/src/clorch/torch.clj index 25ac038..cc489e7 100644 --- a/src/clorch/torch.clj +++ b/src/clorch/torch.clj @@ -2015,24 +2015,42 @@ :else (throw (IllegalArgumentException. (str "Invalid indexer: " spec)))))) (.index ^Tensor t-final idx-vec)))) -(defn precompute-rope-freqs [dim seq-len & {:keys [theta] :or {theta 10000.0}}] +(defn precompute-rope-freqs + "Precomputes rotary-embedding cosine and sine tables. + + `interleaved?` selects adjacent pairs `(0,1), (2,3), ...`; the default uses + the half-split layout `(0,D/2), (1,D/2+1), ...`." + [dim seq-len & {:keys [theta interleaved?] + :or {theta 10000.0 interleaved? false}}] (let [freqs (div (arange 0 dim 2) (clojure.core/double dim)) inv-freq (pow (clojure.core/float theta) (mul freqs -1.0)) - t (arange seq-len) - freqs-matrix (outer t inv-freq) - emb (cat [freqs-matrix freqs-matrix] -1)] + freqs-matrix (outer (arange seq-len) inv-freq) + emb (if interleaved? + (repeat-interleave freqs-matrix 2 -1) + (cat [freqs-matrix freqs-matrix] -1))] [(cos emb) (sin emb)])) -(defn apply-rope [x cos-emb sin-emb] - (let [[_B _T _H D] (size x) - cos-emb (reshape cos-emb [1 _T 1 D]) - sin-emb (reshape sin-emb [1 _T 1 D]) - half-d (quot D 2) - x1 (ix x :_ :_ :_ [0 half-d]) - x2 (ix x :_ :_ :_ [half-d D]) - ;; standard rope rotation: [x1, x2] -> [x1*cos - x2*sin, x1*sin + x2*cos] - neg-x2 (mul x2 -1.0) - rotated-x (cat [neg-x2 x1] -1)] +(defn apply-rope + "Applies rotary embeddings to `[batch, sequence, heads, head-dim]`. + + Set `interleaved?` to rotate adjacent pairs instead of half-split pairs." + [x cos-emb sin-emb & {:keys [interleaved?] :or {interleaved? false}}] + (let [[batch-size sequence-length heads head-dim] (size x) + cos-emb (reshape (ix cos-emb [0 sequence-length] :_) + [1 sequence-length 1 head-dim]) + sin-emb (reshape (ix sin-emb [0 sequence-length] :_) + [1 sequence-length 1 head-dim]) + rotated-x + (if interleaved? + (let [pairs (reshape x [batch-size sequence-length heads + (quot head-dim 2) 2]) + [even odd] (unbind pairs -1)] + (reshape (stack [(mul odd -1.0) even] -1) + [batch-size sequence-length heads head-dim])) + (let [half-dim (quot head-dim 2) + first-half (ix x :_ :_ :_ [0 half-dim]) + second-half (ix x :_ :_ :_ [half-dim head-dim])] + (cat [(mul second-half -1.0) first-half] -1)))] (add (mul x cos-emb) (mul rotated-x sin-emb)))) ;; --- Extended Math --- diff --git a/test/clorch/nn_transformer_test.clj b/test/clorch/nn_transformer_test.clj new file mode 100644 index 0000000..82b3767 --- /dev/null +++ b/test/clorch/nn_transformer_test.clj @@ -0,0 +1,91 @@ +(ns clorch.nn-transformer-test + (:require [clojure.test :refer [deftest is testing]] + [clorch.nn :as nn] + [clorch.torch :as t])) + +(set! *warn-on-reflection* true) + +(defn- max-error [left right] + (t/item-float (t/max (t/abs (t/sub left right))))) + +(deftest interleaved-rope-rotates-adjacent-pairs + (t/with-torch + (let [x (t/tensor [[[[1.0 2.0 3.0 4.0]]]]) + cosines (t/zeros [1 4]) + sines (t/ones [1 4]) + half-split (t/apply-rope x cosines sines) + interleaved (t/apply-rope x cosines sines :interleaved? true)] + (is (< (max-error half-split + (t/tensor [[[[-3.0 -4.0 1.0 2.0]]]])) + 1.0e-6)) + (is (< (max-error interleaved + (t/tensor [[[[-2.0 1.0 -4.0 3.0]]]])) + 1.0e-6))))) + +(deftest cached-cross-attention-matches-full-attention + (t/manual-seed 11) + (let [attention (nn/multihead-attention 8 2 :dropout 0.0) + query (t/randn [2 3 8]) + context (t/randn [2 5 8]) + full (nn/forward attention {:query query + :key context + :value context + :return-kv? true})] + (t/with-torch + (let [cached (nn/forward attention + {:query query + :cached-key (:key full) + :cached-value (:value full)})] + (is (= [2 3 8] (t/size cached))) + (is (< (max-error (:output full) cached) 1.0e-6)))))) + +(deftest key-padding-mask-excludes-padded-values + (t/manual-seed 13) + (let [attention (nn/multihead-attention 8 2 :dropout 0.0) + query (t/randn [2 3 8]) + context (t/randn [2 5 8]) + changed (t/clone context) + _ (t/copy- (t/ix changed 1 [3 5] :_) (t/full [2 8] 1000.0)) + padding-mask (t/tensor [[false false false false false] + [false false false true true]] + {:dtype :bool})] + (t/with-torch + (let [baseline (nn/forward attention {:query query + :key context + :value context + :key-padding-mask padding-mask}) + modified (nn/forward attention {:query query + :key changed + :value changed + :key-padding-mask padding-mask})] + (testing "unmasked first row remains ordinary cross-attention" + (is (= [2 3 8] (t/size baseline)))) + (testing "masked values cannot affect the second row" + (is (< (max-error (t/ix baseline 1 :_ :_) + (t/ix modified 1 :_ :_)) + 1.0e-5))))))) + +(deftest transformer-encoder-runs-pre-and-post-norm + (t/manual-seed 17) + (t/with-torch + (let [x (t/randn [2 4 8]) + pre-layer (nn/transformer-encoder-layer + 8 2 :feedforward-dim 16 :dropout 0.0 + :activation :gelu :norm-first? true) + post-layer (nn/transformer-encoder-layer + 8 2 :feedforward-dim 16 :dropout 0.0 + :activation :relu :norm-first? false) + encoder (nn/transformer-encoder [pre-layer post-layer])] + (is (= [2 4 8] (t/size (nn/forward pre-layer x)))) + (is (= [2 4 8] (t/size (nn/forward post-layer x)))) + (is (= [2 4 8] (t/size (nn/forward encoder x))))))) + +(deftest eval-mode-disables-attention-dropout + (t/manual-seed 19) + (let [attention (nn/multihead-attention 8 2 :dropout 0.75) + x (t/randn [2 5 8])] + (nn/train attention false) + (t/with-torch + (let [first-output (nn/forward attention x) + second-output (nn/forward attention x)] + (is (< (max-error first-output second-output) 1.0e-7)))))) diff --git a/tests_comparison/compare_transformer.clj b/tests_comparison/compare_transformer.clj new file mode 100644 index 0000000..6bbb8d6 --- /dev/null +++ b/tests_comparison/compare_transformer.clj @@ -0,0 +1,87 @@ +(ns compare-transformer + (:require [clojure.data.json :as json] + [clorch.nn :as nn] + [clorch.torch :as t])) + +(set! *warn-on-reflection* true) + +(def specs + {"attention.in-projection.weight" [0.002 -0.10] + "attention.in-projection.bias" [0.003 -0.05] + "attention.out-projection.weight" [0.004 -0.10] + "attention.out-projection.bias" [0.005 0.02] + "feedforward[0].weight" [0.001 -0.08] + "feedforward[0].bias" [0.002 0.01] + "feedforward[3].weight" [0.0015 -0.06] + "feedforward[3].bias" [0.003 -0.02] + "norm1.gamma" [0.001 1.0] + "norm1.beta" [0.002 -0.01] + "norm2.gamma" [0.0015 0.9] + "norm2.beta" [0.001 0.02]}) + +(def attention-specs + {"in-projection.weight" [0.002 -0.10] + "in-projection.bias" [0.003 -0.05] + "out-projection.weight" [0.004 -0.10] + "out-projection.bias" [0.005 0.02]}) + +(defn- initialize! [model parameter-specs] + (doseq [[parameter-name tensor] (nn/named-parameters model)] + (let [[scale offset] (get parameter-specs parameter-name) + shape (t/size tensor) + values (t/to (t/add (t/mul (t/arange (reduce * 1 shape)) scale) + offset) + :float32)] + (when-not scale + (throw (ex-info "Missing deterministic parameter specification" + {:parameter parameter-name}))) + (t/copy- tensor (t/reshape values shape)))) + model) + +(defn- max-error [actual expected] + (t/item-float (t/max (t/abs (t/sub actual (t/tensor expected)))))) + +(defn -main [reference-path] + (let [reference (json/read-str (slurp reference-path)) + attention (initialize! + (nn/multihead-attention 8 2 :dropout 0.0) + attention-specs) + layer (initialize! + (nn/transformer-encoder-layer + 8 2 :feedforward-dim 16 :dropout 0.0 + :activation :gelu :norm-first? true) + specs)] + (nn/train attention false) + (nn/train layer false) + (t/with-torch + (let [query (t/sub (t/mul (t/reshape (t/arange 48) [2 3 8]) 0.01) 0.2) + context (t/sub (t/mul (t/reshape (t/arange 80) [2 5 8]) 0.008) 0.15) + padding-mask (t/tensor [[false false false false false] + [false false false true true]] + {:dtype :bool}) + rope-input (t/tensor [[[[1.0 2.0 3.0 4.0]] + [[-1.0 0.5 2.0 -3.0]]]]) + angles (t/tensor [[0.2 0.2 0.6 0.6] + [0.4 0.4 0.8 0.8]]) + results + {"self_attention" (nn/forward attention query) + "cross_attention" (nn/forward attention + {:query query + :key context + :value context + :key-padding-mask padding-mask}) + "transformer" (nn/forward layer query) + "interleaved_rope" (t/apply-rope + rope-input + (t/cos angles) + (t/sin angles) + :interleaved? true)} + errors (into {} + (map (fn [[result-name actual]] + [result-name (max-error actual (get reference result-name))])) + results) + tolerance 2.0e-5] + (println (pr-str errors)) + (when (some #(>= (double %) tolerance) (vals errors)) + (throw (ex-info "Clorch transformer parity check failed" + {:errors errors :tolerance tolerance}))))))) diff --git a/tests_comparison/transformer_reference.py b/tests_comparison/transformer_reference.py new file mode 100644 index 0000000..0443132 --- /dev/null +++ b/tests_comparison/transformer_reference.py @@ -0,0 +1,97 @@ +import json +import sys + +import torch + + +torch.set_num_threads(1) +torch.manual_seed(0) + +SPECS = { + "in_proj_weight": (0.002, -0.10), + "in_proj_bias": (0.003, -0.05), + "out_proj.weight": (0.004, -0.10), + "out_proj.bias": (0.005, 0.02), + "linear1.weight": (0.001, -0.08), + "linear1.bias": (0.002, 0.01), + "linear2.weight": (0.0015, -0.06), + "linear2.bias": (0.003, -0.02), + "norm1.weight": (0.001, 1.0), + "norm1.bias": (0.002, -0.01), + "norm2.weight": (0.0015, 0.9), + "norm2.bias": (0.001, 0.02), +} + + +def values_like(tensor, scale, offset): + return (torch.arange(tensor.numel(), dtype=torch.float32) * scale + offset).reshape(tensor.shape) + + +def initialize(module): + with torch.no_grad(): + for name, parameter in module.named_parameters(): + suffix = name.removeprefix("self_attn.") + scale, offset = SPECS[suffix] + parameter.copy_(values_like(parameter, scale, offset)) + + +def interleaved_rope(x, cosines, sines): + pairs = x.reshape(*x.shape[:-1], x.shape[-1] // 2, 2) + even, odd = pairs.unbind(-1) + rotated = torch.stack((-odd, even), dim=-1).reshape_as(x) + return x * cosines.reshape(1, cosines.shape[0], 1, -1) + rotated * sines.reshape(1, sines.shape[0], 1, -1) + + +def main(output_path): + attention = torch.nn.MultiheadAttention(8, 2, dropout=0.0, batch_first=True) + layer = torch.nn.TransformerEncoderLayer( + 8, + 2, + dim_feedforward=16, + dropout=0.0, + activation="gelu", + batch_first=True, + norm_first=True, + ) + initialize(attention) + initialize(layer) + attention.eval() + layer.eval() + + query = (torch.arange(48, dtype=torch.float32).reshape(2, 3, 8) * 0.01) - 0.2 + context = (torch.arange(80, dtype=torch.float32).reshape(2, 5, 8) * 0.008) - 0.15 + padding_mask = torch.tensor( + [[False, False, False, False, False], + [False, False, False, True, True]] + ) + rope_input = torch.tensor([[[[1.0, 2.0, 3.0, 4.0]], + [[-1.0, 0.5, 2.0, -3.0]]]]) + angles = torch.tensor([[0.2, 0.2, 0.6, 0.6], + [0.4, 0.4, 0.8, 0.8]]) + + with torch.inference_mode(): + self_attention = attention(query, query, query, need_weights=False)[0] + cross_attention = attention( + query, + context, + context, + key_padding_mask=padding_mask, + need_weights=False, + )[0] + transformer = layer(query) + rope = interleaved_rope(rope_input, torch.cos(angles), torch.sin(angles)) + + with open(output_path, "w", encoding="utf-8") as output: + json.dump( + { + "self_attention": self_attention.tolist(), + "cross_attention": cross_attention.tolist(), + "transformer": transformer.tolist(), + "interleaved_rope": rope.tolist(), + }, + output, + ) + + +if __name__ == "__main__": + main(sys.argv[1]) From 2e2e2034557115e5cda5af9212cc6547d63e7131 Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Thu, 15 Oct 2026 00:15:59 +0200 Subject: [PATCH 3/7] feat: add SafeTensors state dict IO --- src/clorch/safetensors.clj | 318 +++++++++++++++++++++++++++++++ test/clorch/safetensors_test.clj | 71 +++++++ 2 files changed, 389 insertions(+) create mode 100644 src/clorch/safetensors.clj create mode 100644 test/clorch/safetensors_test.clj diff --git a/src/clorch/safetensors.clj b/src/clorch/safetensors.clj new file mode 100644 index 0000000..2252c76 --- /dev/null +++ b/src/clorch/safetensors.clj @@ -0,0 +1,318 @@ +(ns clorch.safetensors + "SafeTensors state-dictionary I/O without Python or pickle execution." + (:refer-clojure :exclude [load-file]) + (:require [clojure.data.json :as json] + [clorch.nn :as nn] + [clorch.torch :as t]) + (:import [java.io RandomAccessFile] + [java.nio ByteBuffer ByteOrder] + [java.nio.channels FileChannel] + [java.nio.charset StandardCharsets] + [java.nio.file AtomicMoveNotSupportedException Files Path StandardCopyOption StandardOpenOption] + [org.bytedeco.javacpp BytePointer Pointer] + [org.bytedeco.pytorch ScalarTypeOptional Tensor TensorOptions] + [org.bytedeco.pytorch.global torch torch$ScalarType])) + +(set! *warn-on-reflection* true) + +(def ^:private dtype-info + {"BOOL" {:bytes 1 :native "Bool"} + "U8" {:bytes 1 :native "Byte"} + "I8" {:bytes 1 :native "Char"} + "I16" {:bytes 2 :native "Short"} + "I32" {:bytes 4 :native "Int"} + "I64" {:bytes 8 :native "Long"} + "F16" {:bytes 2 :native "Half"} + "BF16" {:bytes 2 :native "BFloat16"} + "F32" {:bytes 4 :native "Float"} + "F64" {:bytes 8 :native "Double"} + "C64" {:bytes 8 :native "ComplexFloat"} + "C128" {:bytes 16 :native "ComplexDouble"}}) + +(def ^:private native->dtype + (into {} (map (fn [[safe {:keys [native]}]] [native safe]) dtype-info))) + +(defn- scalar-type + ^torch$ScalarType [dtype] + (case dtype + "BOOL" (torch/kBool) + "U8" (torch/kByte) + "I8" (torch/kChar) + "I16" (torch/kShort) + "I32" (torch/kInt) + "I64" (torch/kLong) + "F16" (torch/kHalf) + "BF16" (torch/kBFloat16) + "F32" (torch/kFloat) + "F64" (torch/kDouble) + "C64" (torch/kComplexFloat) + "C128" (torch/kComplexDouble) + (throw (ex-info "Unsupported SafeTensors dtype" {:dtype dtype})))) + +(defn- path + ^Path [value] + (if (instance? Path value) + value + (Path/of (str value) (make-array String 0)))) + +(defn- read-fully! + [^FileChannel channel ^ByteBuffer buffer position] + (loop [offset (long position)] + (when (.hasRemaining buffer) + (let [read-count (.read channel buffer offset)] + (when (neg? read-count) + (throw (ex-info "Unexpected end of SafeTensors file" + {:position offset}))) + (recur (+ offset read-count))))) + (.flip buffer)) + +(defn- write-fully! + [^FileChannel channel ^ByteBuffer buffer] + (while (.hasRemaining buffer) + (.write channel buffer))) + +(defn- read-header + [^FileChannel channel] + (let [file-size (.size channel)] + (when (< file-size 8) + (throw (ex-info "SafeTensors file is shorter than its header length" {}))) + (let [length-buffer (doto (ByteBuffer/allocate 8) + (.order ByteOrder/LITTLE_ENDIAN))] + (read-fully! channel length-buffer 0) + (let [header-length (.getLong length-buffer)] + (when (or (neg? header-length) + (> header-length (- file-size 8)) + (> header-length (* 100 1024 1024))) + (throw (ex-info "Invalid SafeTensors header length" + {:header-length header-length + :file-size file-size}))) + (let [header-buffer (ByteBuffer/allocate + (Math/toIntExact header-length))] + (read-fully! channel header-buffer 8) + {:data-start (+ 8 header-length) + :data-length (- file-size 8 header-length) + :header (json/read-str + (String. (.array header-buffer) + StandardCharsets/UTF_8))}))))) + +(defn- element-count + [shape] + (reduce * 1 (map long shape))) + +(defn- validate-entry! + [tensor-name data-length {:strs [dtype shape data_offsets]}] + (let [{bytes-per-element :bytes} (get dtype-info dtype) + [start end] data_offsets + expected-bytes (when bytes-per-element + (* bytes-per-element (element-count shape)))] + (when-not (and (string? tensor-name) + bytes-per-element + (vector? shape) + (every? #(and (integer? %) (not (neg? (long %)))) shape) + (= 2 (count data_offsets)) + (integer? start) + (integer? end) + (<= 0 start end data-length) + (= expected-bytes (- end start))) + (throw (ex-info "Invalid SafeTensors tensor entry" + {:name tensor-name + :dtype dtype + :shape shape + :data-offsets data_offsets + :data-length data-length}))) + {:name tensor-name + :dtype dtype + :shape (mapv long shape) + :start (long start) + :end (long end)})) + +(defn- validate-entries! + [header data-length] + (let [validated-entries (->> (dissoc header "__metadata__") + (map (fn [[tensor-name entry]] + (validate-entry! tensor-name data-length entry))) + (sort-by :start) + vec)] + (loop [previous-end (long 0) + [entry & remaining] validated-entries] + (when entry + (when-not (= previous-end (:start entry)) + (throw (ex-info "SafeTensors data offsets overlap or contain gaps" + {:previous-end previous-end + :entry entry}))) + (recur (long (:end entry)) remaining))) + (when-not (= data-length (or (:end (peek validated-entries)) 0)) + (throw (ex-info "SafeTensors data does not cover the file" + {:data-length data-length + :last-offset (or (:end (peek validated-entries)) 0)}))) + validated-entries)) + +(defn- tensor-options + ^TensorOptions [dtype] + (with-open [options (TensorOptions.) + scalar (ScalarTypeOptional. (scalar-type dtype))] + (.dtype options scalar))) + +(defn- read-tensor + [^FileChannel channel data-start {:keys [dtype shape start end]}] + (let [byte-count (- end start) + options (tensor-options dtype) + buffer (ByteBuffer/allocate (Math/toIntExact byte-count))] + (try + (read-fully! channel buffer (+ data-start start)) + (with-open [pointer (BytePointer. (.array buffer)) + view (torch/from_blob pointer (long-array shape) options)] + (.clone ^Tensor view)) + (finally + (.close options))))) + +(defn load-file + "Loads a SafeTensors file as a map of string tensor names to CPU tensors. + + The caller owns the returned tensors and must release them with their model or + enclosing native-memory scope. Pickle and executable metadata are never read." + [file] + (when-not (= ByteOrder/LITTLE_ENDIAN (ByteOrder/nativeOrder)) + (throw (UnsupportedOperationException. + "SafeTensors requires little-endian native tensor storage"))) + (with-open [random-access (RandomAccessFile. (str (path file)) "r") + channel (.getChannel random-access)] + (let [{:keys [data-start data-length header]} (read-header channel) + validated-entries (validate-entries! header data-length)] + (into {} + (map (fn [{tensor-name :name :as entry}] + [tensor-name (read-tensor channel data-start entry)])) + validated-entries)))) + +(defn metadata + "Returns the string metadata map from a SafeTensors file." + [file] + (with-open [random-access (RandomAccessFile. (str (path file)) "r") + channel (.getChannel random-access)] + (or (get-in (read-header channel) [:header "__metadata__"]) {}))) + +(defn- tensor-entry + [offset [tensor-name value]] + (when-not (and (string? tensor-name) (not= tensor-name "__metadata__")) + (throw (ex-info "SafeTensors tensor names must be strings other than __metadata__" + {:name tensor-name}))) + (let [tensor (t/->tensor value) + native-name (.toString (.scalar_type tensor)) + dtype (get native->dtype native-name) + byte-count (Math/multiplyExact (.numel tensor) (.element_size tensor))] + (when-not dtype + (throw (ex-info "Tensor dtype is unsupported by SafeTensors" + {:name tensor-name :dtype native-name}))) + [{:name tensor-name + :tensor tensor + :dtype dtype + :shape (t/size tensor) + :start offset + :end (Math/addExact (long offset) byte-count)} + (Math/addExact (long offset) byte-count)])) + +(defn- entries + [tensors] + (first + (reduce (fn [[result offset] pair] + (let [[entry next-offset] (tensor-entry offset pair)] + [(conj result entry) next-offset])) + [[] 0] + (sort-by key tensors)))) + +(defn- padded-header + [tensor-entries file-metadata] + (let [tensor-header + (into (sorted-map) + (map (fn [{tensor-name :name :keys [dtype shape start end]}] + [tensor-name {"dtype" dtype + "shape" shape + "data_offsets" [start end]}])) + tensor-entries) + header (cond-> tensor-header + (seq file-metadata) (assoc "__metadata__" file-metadata)) + header-bytes (.getBytes (json/write-str header) StandardCharsets/UTF_8) + padding (mod (- 8 (mod (alength header-bytes) 8)) 8) + result (byte-array (+ (alength header-bytes) padding))] + (System/arraycopy header-bytes 0 result 0 (alength header-bytes)) + (java.util.Arrays/fill result (alength header-bytes) (alength result) (byte 32)) + result)) + +(defn- write-tensor! + [^FileChannel channel tensor] + (t/with-torch + (let [^Tensor contiguous (.contiguous (.cpu (.detach (t/->tensor tensor)))) + byte-count (* (.numel contiguous) (.element_size contiguous))] + (when (pos? byte-count) + (let [^Pointer pointer (.data_ptr contiguous)] + (.capacity pointer byte-count) + (write-fully! channel (.asByteBuffer pointer)))) + nil))) + +(defn save-file! + "Atomically writes a map of string tensor names to a SafeTensors file. + + Optional `:metadata` must be a string-to-string map. Tensors are detached, + made contiguous, and copied to CPU before their raw bytes are written." + [tensors file & {:as options}] + (let [file-metadata (get options :metadata {})] + (when-not (map? tensors) + (throw (IllegalArgumentException. "tensors must be a map"))) + (when-not (and (map? file-metadata) + (every? (fn [[metadata-key metadata-value]] + (and (string? metadata-key) (string? metadata-value))) + file-metadata)) + (throw (IllegalArgumentException. + "metadata must contain only string keys and values"))) + (when-not (= ByteOrder/LITTLE_ENDIAN (ByteOrder/nativeOrder)) + (throw (UnsupportedOperationException. + "SafeTensors requires little-endian native tensor storage"))) + (let [destination (path file) + parent (or (.getParent destination) + (Path/of "." (make-array String 0))) + tensor-entries (entries tensors) + ^bytes header (padded-header tensor-entries file-metadata)] + (Files/createDirectories parent + (make-array java.nio.file.attribute.FileAttribute 0)) + (let [temporary + (Files/createTempFile + parent + (str "." (.getFileName destination) "-") + ".tmp" + (make-array java.nio.file.attribute.FileAttribute 0))] + (try + (with-open [channel (FileChannel/open + temporary + (into-array java.nio.file.OpenOption + [StandardOpenOption/WRITE + StandardOpenOption/TRUNCATE_EXISTING]))] + (let [length-buffer (doto (ByteBuffer/allocate 8) + (.order ByteOrder/LITTLE_ENDIAN) + (.putLong (long (alength header))) + (.flip))] + (write-fully! channel length-buffer) + (write-fully! channel (ByteBuffer/wrap header)) + (doseq [{:keys [tensor]} tensor-entries] + (write-tensor! channel tensor)) + (.force channel true))) + (try + (Files/move temporary destination + (into-array java.nio.file.CopyOption + [StandardCopyOption/ATOMIC_MOVE + StandardCopyOption/REPLACE_EXISTING])) + (catch AtomicMoveNotSupportedException _ + (Files/move temporary destination + (into-array java.nio.file.CopyOption + [StandardCopyOption/REPLACE_EXISTING])))) + (finally + (Files/deleteIfExists temporary)))) + destination))) + +(defn load-state-dict! + "Loads a SafeTensors file into a model using flat dotted parameter names." + [model file] + (t/with-torch + (let [state (load-file file)] + (nn/load-flat-state-dict! model state) + nil)) + model) diff --git a/test/clorch/safetensors_test.clj b/test/clorch/safetensors_test.clj new file mode 100644 index 0000000..f81c553 --- /dev/null +++ b/test/clorch/safetensors_test.clj @@ -0,0 +1,71 @@ +(ns clorch.safetensors-test + (:require [clojure.test :refer [deftest is testing]] + [clorch.safetensors :as safe] + [clorch.torch :as t]) + (:import [java.io FileOutputStream] + [java.nio.file Files Path] + [org.bytedeco.pytorch Tensor])) + +(set! *warn-on-reflection* true) + +(defn- temporary-file + ^Path [] + (Files/createTempFile "clorch-safetensors-test-" ".safetensors" + (make-array java.nio.file.attribute.FileAttribute 0))) + +(defn- max-error + [^Tensor actual ^Tensor expected] + (t/item-float + (t/max + (t/abs + (t/sub (t/to actual :float32) + (t/to expected :float32)))))) + +(deftest small-round-trip-preserves-values-dtypes-and-metadata + (let [file (temporary-file)] + (try + (t/with-torch + (let [values [1.25 -2.5 3.0] + tensors {"bool" (t/tensor [true false true] {:dtype :bool}) + "bfloat16" (t/tensor values {:dtype :bfloat16}) + "float16" (t/tensor values {:dtype :float16}) + "float32" (t/tensor values) + "float64" (t/tensor values {:dtype :float64}) + "int8" (t/tensor [1 -2 3] {:dtype :int8}) + "int32" (t/tensor [1 -2 3] {:dtype :int32}) + "int64" (t/tensor [1 -2 3] {:dtype :int64}) + "uint8" (t/tensor [1 2 3] {:dtype :uint8})}] + (safe/save-file! tensors file :metadata {"source" "small-test"}) + (let [loaded (safe/load-file file)] + (is (= (set (keys tensors)) (set (keys loaded)))) + (is (= {"source" "small-test"} (safe/metadata file))) + (doseq [[tensor-name expected] tensors] + (testing tensor-name + (is (= (.toString (.scalar_type ^Tensor expected)) + (.toString (.scalar_type ^Tensor (get loaded tensor-name))))) + (is (zero? (max-error (get loaded tensor-name) expected)))))))) + (finally + (Files/deleteIfExists file))))) + +(deftest malformed-data-offset-is-rejected + (let [file (temporary-file) + ^bytes header-bytes + (.getBytes + "{\"x\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[0,8]}}") + padding (mod (- 8 (mod (alength header-bytes) 8)) 8) + padded (byte-array (+ (alength header-bytes) padding)) + contents (java.nio.ByteBuffer/allocate (+ 8 (alength padded) 4))] + (try + (System/arraycopy header-bytes 0 padded 0 (alength header-bytes)) + (java.util.Arrays/fill padded (alength header-bytes) (alength padded) (byte 32)) + (.order contents java.nio.ByteOrder/LITTLE_ENDIAN) + (.putLong contents (long (alength padded))) + (.put contents padded) + (.putFloat contents 1.0) + (with-open [output (FileOutputStream. (.toFile file))] + (.write output ^bytes (.array contents))) + (is (thrown-with-msg? clojure.lang.ExceptionInfo + #"Invalid SafeTensors tensor entry" + (safe/load-file file))) + (finally + (Files/deleteIfExists file))))) From e26c893f063a4fb32b492e97c938715128798bf4 Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Sun, 18 Oct 2026 19:10:30 +0200 Subject: [PATCH 4/7] docs: document transformer, SafeTensors, and rope APIs --- PORTING_STATUS.md | 23 +++++++++++-- docs/functional.md | 22 +++++++++++++ docs/init.md | 12 +++++++ docs/mkdocs.yml | 1 + docs/nn.md | 62 +++++++++++++++++++++++++++++++++++ docs/pytorch-parity.md | 8 +++-- docs/safetensors.md | 46 ++++++++++++++++++++++++++ docs/tensors.md | 21 ++++++++++++ run_transformer_benchmark.sh | 2 +- run_transformer_comparison.sh | 2 +- 10 files changed, 191 insertions(+), 8 deletions(-) create mode 100644 docs/safetensors.md diff --git a/PORTING_STATUS.md b/PORTING_STATUS.md index e1e84be..81b3f3c 100644 --- a/PORTING_STATUS.md +++ b/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* diff --git a/docs/functional.md b/docs/functional.md index e604f78..0b64138 100644 --- a/docs/functional.md +++ b/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. diff --git a/docs/init.md b/docs/init.md index 80925e2..02a8fca 100644 --- a/docs/init.md +++ b/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`). diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 7aa4ff6..515352d 100644 --- a/docs/mkdocs.yml +++ b/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 diff --git a/docs/nn.md b/docs/nn.md index 88e7447..fdf679d 100644 --- a/docs/nn.md +++ b/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 diff --git a/docs/pytorch-parity.md b/docs/pytorch-parity.md index d6aaaf9..37cb758 100644 --- a/docs/pytorch-parity.md +++ b/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` | diff --git a/docs/safetensors.md b/docs/safetensors.md new file mode 100644 index 0000000..dfa43aa --- /dev/null +++ b/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. diff --git a/docs/tensors.md b/docs/tensors.md index a6694ce..d19f297 100644 --- a/docs/tensors.md +++ b/docs/tensors.md @@ -207,3 +207,24 @@ 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) + +;; 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`. diff --git a/run_transformer_benchmark.sh b/run_transformer_benchmark.sh index cb1ca75..1e33760 100755 --- a/run_transformer_benchmark.sh +++ b/run_transformer_benchmark.sh @@ -9,5 +9,5 @@ uv run --python 3.11 --no-project \ python benchmarks/transformer_pytorch.py clojure -Sthreads 1 \ - -Sdeps '{:paths ["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"}}}' \ + -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))' diff --git a/run_transformer_comparison.sh b/run_transformer_comparison.sh index 374a83e..342f3be 100755 --- a/run_transformer_comparison.sh +++ b/run_transformer_comparison.sh @@ -9,5 +9,5 @@ uv run --python 3.11 --no-project \ python tests_comparison/transformer_reference.py "$reference_file" clojure -Sthreads 1 \ - -Sdeps '{:paths ["tests_comparison"] :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"}}}' \ + -Sdeps '{:paths ["src" "tests_comparison"] :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 -m compare-transformer "$reference_file" From 49c687f53b0ee2b01481fd2dfee8b03a1aa04d65 Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Sun, 18 Oct 2026 20:10:24 +0200 Subject: [PATCH 5/7] fix: move defmodel fields across devices --- src/clorch/nn.clj | 5 +++-- test/clorch/nn_model_test.clj | 6 +++++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/clorch/nn.clj b/src/clorch/nn.clj index 9b363a5..ecc9102 100644 --- a/src/clorch/nn.clj +++ b/src/clorch/nn.clj @@ -859,6 +859,7 @@ (defmacro defmodel [model-name args init-bindings & forward-specs] (let [field-syms (take-nth 2 init-bindings) + field-keys (mapv keyword field-syms) forward-fn (first (filter #(= (first %) 'forward) forward-specs)) forward-args (second forward-fn) forward-body (nnext forward-fn) @@ -869,13 +870,13 @@ (-forward [this# ~(first forward-args)] (let [{:keys [~@field-syms]} this#] ~@forward-body)) (-train [this# bool#] - (doseq [f# [~@field-syms]] + (doseq [f# [~@field-keys]] (let [v# (get this# f#)] (when (or (satisfies? IModule v#) (coll? v#)) (train v# bool#)))) this#) (-to [this# dtype-or-device#] - (let [updates# (into {} (for [f# [~@field-syms]] + (let [updates# (into {} (for [f# [~@field-keys]] (let [v# (get this# f#)] [f# (if (or (satisfies? IModule v#) (coll? v#) (instance? org.bytedeco.pytorch.Tensor v#)) (to v# dtype-or-device#) diff --git a/test/clorch/nn_model_test.clj b/test/clorch/nn_model_test.clj index cd75d2e..cf39617 100644 --- a/test/clorch/nn_model_test.clj +++ b/test/clorch/nn_model_test.clj @@ -39,7 +39,11 @@ x (t/randn [2 10]) out (nn/forward m x)] (is (= [2 5] (t/size out))) - (is (= 4 (.size (nn/parameters m))))))) + (is (= 4 (.size (nn/parameters m)))) + (testing "Device and dtype conversion reaches defmodel fields" + (let [m-double (nn/to m :float64)] + (is (every? #(= :float64 (t/dtype %)) + (t/->vector (nn/parameters m-double))))))))) (deftest flat-state-dict-loading-test (t/with-torch From 656b2e6e07a02f1451e34ee0280fe5ffb909f18f Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Sun, 18 Oct 2026 20:12:32 +0200 Subject: [PATCH 6/7] fix: preserve leaf parameters during model transfer --- src/clorch/nn.clj | 8 ++++---- test/clorch/nn_model_test.clj | 14 ++++++++------ 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/src/clorch/nn.clj b/src/clorch/nn.clj index ecc9102..2f64d5c 100644 --- a/src/clorch/nn.clj +++ b/src/clorch/nn.clj @@ -71,7 +71,9 @@ (dotimes [i (.size thing)] (.push_back out ^Tensor (torch/to (.get thing (clojure.core/long i)) dtype-or-device))) out) - (instance? Parameter thing) (assoc thing :tensor (torch/to (:tensor thing) dtype-or-device)) + (instance? Parameter thing) + (parameter (torch/to (:tensor thing) dtype-or-device) + :requires-grad (.requires_grad ^Tensor (:tensor thing))) (map? thing) (into (empty thing) (for [[k v] thing] [k (to v dtype-or-device)])) (sequential? thing) (mapv #(to % dtype-or-device) thing) (vector? thing) (mapv #(to % dtype-or-device) thing) @@ -878,9 +880,7 @@ (-to [this# dtype-or-device#] (let [updates# (into {} (for [f# [~@field-keys]] (let [v# (get this# f#)] - [f# (if (or (satisfies? IModule v#) (coll? v#) (instance? org.bytedeco.pytorch.Tensor v#)) - (to v# dtype-or-device#) - v#)])))] + [f# (to v# dtype-or-device#)])))] (clojure.core/merge this# updates#))) clojure.lang.IFn (invoke [this# x#] (forward this# x#))) diff --git a/test/clorch/nn_model_test.clj b/test/clorch/nn_model_test.clj index cf39617..8cebbff 100644 --- a/test/clorch/nn_model_test.clj +++ b/test/clorch/nn_model_test.clj @@ -29,21 +29,23 @@ #_{:clj-kondo/ignore [:unresolved-symbol]} (nn/defmodel MyModel [in out] [l1 (nn/linear in 16) - l2 (nn/linear 16 out)] + l2 (nn/linear 16 out) + norm (nn/layernorm out)] (forward [x] (let [x1 (nn/forward l1 x) x2 (F/relu x1)] - (nn/forward l2 x2)))) + (nn/forward norm (nn/forward l2 x2))))) (let [m (MyModel 10 5) x (t/randn [2 10]) out (nn/forward m x)] (is (= [2 5] (t/size out))) - (is (= 4 (.size (nn/parameters m)))) + (is (= 6 (.size (nn/parameters m)))) (testing "Device and dtype conversion reaches defmodel fields" - (let [m-double (nn/to m :float64)] - (is (every? #(= :float64 (t/dtype %)) - (t/->vector (nn/parameters m-double))))))))) + (let [m-double (nn/to m :float64) + params (t/->vector (nn/parameters m-double))] + (is (every? #(= :float64 (t/dtype %)) params)) + (is (every? #(.is_leaf ^org.bytedeco.pytorch.Tensor %) params))))))) (deftest flat-state-dict-loading-test (t/with-torch From f18a62dd24c415b408bf822c5dcf65bac0fd794f Mon Sep 17 00:00:00 2001 From: antlobach <246185237+antlobach@users.noreply.github.com> Date: Sun, 18 Oct 2026 20:21:11 +0200 Subject: [PATCH 7/7] feat: place RoPE tables on a device --- docs/tensors.md | 3 +++ src/clorch/torch.clj | 8 +++++--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/docs/tensors.md b/docs/tensors.md index d19f297..90a57ff 100644 --- a/docs/tensors.md +++ b/docs/tensors.md @@ -221,6 +221,9 @@ A built-in utility for high-quality text generation. It filters the probability ;; (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) diff --git a/src/clorch/torch.clj b/src/clorch/torch.clj index cc489e7..5fcf185 100644 --- a/src/clorch/torch.clj +++ b/src/clorch/torch.clj @@ -2019,15 +2019,17 @@ "Precomputes rotary-embedding cosine and sine tables. `interleaved?` selects adjacent pairs `(0,1), (2,3), ...`; the default uses - the half-split layout `(0,D/2), (1,D/2+1), ...`." - [dim seq-len & {:keys [theta interleaved?] + the half-split layout `(0,D/2), (1,D/2+1), ...`. Pass `device` to create + tables for inputs on a non-default device." + [dim seq-len & {:keys [theta interleaved? device] :or {theta 10000.0 interleaved? false}}] (let [freqs (div (arange 0 dim 2) (clojure.core/double dim)) inv-freq (pow (clojure.core/float theta) (mul freqs -1.0)) freqs-matrix (outer (arange seq-len) inv-freq) emb (if interleaved? (repeat-interleave freqs-matrix 2 -1) - (cat [freqs-matrix freqs-matrix] -1))] + (cat [freqs-matrix freqs-matrix] -1)) + emb (if device (to emb device) emb)] [(cos emb) (sin emb)])) (defn apply-rope