SPB Git
spboucher.ai

spb/forge

Public MIT

Forge — LLM training from scratch in pure C++20 + Metal on Apple Silicon.

llmtrainingcppmetalapple-silicon

Add .forge — Apple-native, git-style weight format with zero-copy loading

A .forge is a model repository: tiny JSON manifests (one commit per save,
with parent links) over content-addressed shards. Tensors are 16KB-page-
aligned inside shards padded to page multiples, so loading is mmap +
newBuffer(bytesNoCopy) — on unified memory the file-cache pages ARE the
GPU memory. Saves are deltas: only tensors whose FNV-1a hash changed since
the parent manifest are written. Shards cap at 95MB (GitHub-pushable).
Store f32 (zero-copy alias at load) or f16/bf16 (half size).

- src/core/fmodel.{h,cpp}: save() + Snapshot zero-copy reader
- Tensor::from_buffer: views over externally-owned MTLBuffers
- forge export CLI; generate/eval accept .forge repos directly
- trainer commits weights natively to <out>/model.forge at each checkpoint
  (forge_save/forge_dtype config keys); .bin keeps optimizer state for resume
- tools/fmodel.py: inspect, log (history), to-safetensors (pure numpy)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simon-pierre boucher committed (Aug 5, 2026) parent 4054156 Browse files

7 changed files +702 −13

7 changed files
Unified Split
modified CMakeLists.txt +1 −0
@@ -66,6 +66,7 @@ add_library(forge_core STATIC
66 src/core/allocator.cpp 66 src/core/allocator.cpp
67 src/core/device.cpp 67 src/core/device.cpp
68 src/core/autograd.cpp 68 src/core/autograd.cpp
69+ src/core/fmodel.cpp
69 src/ops/cpu/cpu_ops.cpp 70 src/ops/cpu/cpu_ops.cpp
70 src/ops/metal/metal_ops.cpp 71 src/ops/metal/metal_ops.cpp
71 src/ops/ops.cpp 72 src/ops/ops.cpp
added src/core/fmodel.cpp +359 −0
@@ -0,0 +1,359 @@
1+// Author: Simon-Pierre Boucher — contact@spboucher.ai
2+#include "core/fmodel.h"
3+
4+#include "core/device.h"
5+
6+#include <Metal/Metal.hpp>
7+
8+#include <nlohmann/json.hpp>
9+
10+#include <fcntl.h>
11+#include <sys/mman.h>
12+#include <sys/stat.h>
13+#include <unistd.h>
14+
15+#include <cinttypes>
16+#include <cstdio>
17+#include <cstdlib>
18+#include <cstring>
19+#include <filesystem>
20+#include <fstream>
21+#include <unordered_map>
22+#include <unordered_set>
23+
24+namespace forge::fmodel {
25+
26+namespace {
27+
28+namespace fs = std::filesystem;
29+using nlohmann::json;
30+
31+// Apple Silicon VM page — the alignment contract of the whole format.
32+constexpr size_t kPage = 16384;
33+constexpr uint32_t kShardMagic = 0x53475246; // "FRGS"
34+constexpr uint32_t kVersion = 1;
35+
36+[[noreturn]] void die(const std::string& msg) {
37+ std::fprintf(stderr, "forge/fmodel: %s\n", msg.c_str());
38+ std::abort();
39+}
40+
41+size_t page_align(size_t n) { return (n + kPage - 1) & ~(kPage - 1); }
42+
43+uint64_t fnv1a64(const void* data, size_t n) {
44+ const uint8_t* p = static_cast<const uint8_t*>(data);
45+ uint64_t h = 1469598103934665603ull;
46+ for (size_t i = 0; i < n; ++i) {
47+ h ^= p[i];
48+ h *= 1099511628211ull;
49+ }
50+ return h;
51+}
52+
53+std::string hex64(uint64_t v) {
54+ char buf[17];
55+ std::snprintf(buf, sizeof(buf), "%016" PRIx64, v);
56+ return buf;
57+}
58+
59+DType dtype_from_name(const std::string& s) {
60+ if (s == "f32") return DType::F32;
61+ if (s == "f16") return DType::F16;
62+ if (s == "bf16") return DType::BF16;
63+ if (s == "u16") return DType::U16;
64+ if (s == "i32") return DType::I32;
65+ die("unknown dtype in manifest: " + s);
66+}
67+
68+// Convert one f32 tensor into `dtype` bytes (identity for f32).
69+std::vector<uint8_t> convert_out(const Tensor& t, DType dtype) {
70+ const int64_t n = t.numel();
71+ const float* src = t.data<float>();
72+ std::vector<uint8_t> out(size_t(n) * dtype_size(dtype));
73+ if (dtype == DType::F32) {
74+ std::memcpy(out.data(), src, out.size());
75+ } else if (dtype == DType::F16) {
76+ f16_t* d = reinterpret_cast<f16_t*>(out.data());
77+ for (int64_t i = 0; i < n; ++i) d[i] = f16_t(src[i]);
78+ } else if (dtype == DType::BF16) {
79+ uint16_t* d = reinterpret_cast<uint16_t*>(out.data());
80+ for (int64_t i = 0; i < n; ++i) d[i] = float_to_bf16(src[i]);
81+ } else {
82+ die("save: unsupported storage dtype");
83+ }
84+ return out;
85+}
86+
87+std::vector<std::pair<std::string, Var>> dedupe(
88+ const std::vector<std::pair<std::string, Var>>& named) {
89+ std::vector<std::pair<std::string, Var>> out;
90+ std::unordered_set<const void*> seen;
91+ for (const auto& [name, p] : named) {
92+ if (!p.defined()) continue;
93+ if (seen.insert(p.id()).second) out.emplace_back(name, p);
94+ }
95+ return out;
96+}
97+
98+json load_json(const fs::path& p) {
99+ std::ifstream in(p);
100+ if (!in) die("cannot open " + p.string());
101+ return json::parse(in);
102+}
103+
104+} // namespace
105+
106+std::string save(const std::string& repo_dir, const std::string& config_json,
107+ const std::vector<std::pair<std::string, Var>>& named_params,
108+ const SaveOptions& opts) {
109+ const fs::path repo(repo_dir);
110+ fs::create_directories(repo / "manifests");
111+ fs::create_directories(repo / "objects");
112+
113+ // Parent manifest (if any): hash -> existing location, for delta reuse.
114+ json parent;
115+ std::string parent_name;
116+ const fs::path head = repo / "manifest-latest.json";
117+ if (fs::exists(head)) {
118+ parent = load_json(head);
119+ parent_name = parent.value("self", "");
120+ }
121+ std::unordered_map<std::string, json> known; // content hash -> tensor entry
122+ if (parent.contains("tensors"))
123+ for (const auto& [name, e] : parent.at("tensors").items())
124+ known.emplace(e.at("hash").get<std::string>(), e);
125+
126+ // Hash every (converted) tensor; split into reused vs new.
127+ struct Pending {
128+ std::string name;
129+ std::vector<uint8_t> bytes;
130+ std::string hash;
131+ std::vector<int64_t> shape;
132+ };
133+ json tensors = json::object();
134+ std::vector<Pending> fresh;
135+ size_t reused = 0, reused_bytes = 0;
136+ for (const auto& [name, p] : dedupe(named_params)) {
137+ Pending pd;
138+ pd.name = name;
139+ pd.shape = p.value().shape();
140+ pd.bytes = convert_out(p.value(), opts.dtype);
141+ pd.hash = hex64(fnv1a64(pd.bytes.data(), pd.bytes.size()));
142+ auto it = known.find(pd.hash);
143+ if (it != known.end()) {
144+ tensors[name] = it->second; // unchanged since parent: reference it
145+ ++reused;
146+ reused_bytes += pd.bytes.size();
147+ } else {
148+ fresh.push_back(std::move(pd));
149+ }
150+ }
151+
152+ // Pack new tensors into page-aligned shards capped at shard_mb.
153+ const size_t cap = size_t(opts.shard_mb) * 1024 * 1024;
154+ size_t si = 0;
155+ size_t written_bytes = 0;
156+ while (si < fresh.size()) {
157+ std::vector<uint8_t> blob(kPage, 0); // shard header page
158+ {
159+ uint32_t* h = reinterpret_cast<uint32_t*>(blob.data());
160+ h[0] = kShardMagic;
161+ h[1] = kVersion;
162+ }
163+ json members = json::array();
164+ uint64_t shard_hash = 1469598103934665603ull;
165+ size_t end = si;
166+ while (end < fresh.size()) {
167+ const size_t start = page_align(blob.size());
168+ const size_t next = page_align(start + fresh[end].bytes.size());
169+ if (end > si && next > cap + kPage) break; // shard full
170+ blob.resize(start, 0);
171+ blob.insert(blob.end(), fresh[end].bytes.begin(), fresh[end].bytes.end());
172+ members.push_back(json{{"i", end}, {"offset", start}});
173+ shard_hash ^= fnv1a64(fresh[end].hash.data(), fresh[end].hash.size());
174+ shard_hash *= 1099511628211ull;
175+ ++end;
176+ }
177+ blob.resize(page_align(blob.size()), 0);
178+
179+ const std::string shard_rel = "objects/" + hex64(shard_hash) + ".fshard";
180+ const fs::path shard_path = repo / shard_rel;
181+ if (!fs::exists(shard_path)) { // content-addressed: identical = skip
182+ const fs::path tmp = shard_path.string() + ".tmp";
183+ std::ofstream out(tmp, std::ios::binary);
184+ out.write(reinterpret_cast<const char*>(blob.data()),
185+ std::streamsize(blob.size()));
186+ out.close();
187+ fs::rename(tmp, shard_path);
188+ }
189+ written_bytes += blob.size();
190+
191+ for (const auto& m : members) {
192+ const Pending& pd = fresh[m.at("i").get<size_t>()];
193+ tensors[pd.name] = json{
194+ {"dtype", dtype_name(opts.dtype)}, {"shape", pd.shape},
195+ {"hash", pd.hash}, {"shard", shard_rel},
196+ {"offset", m.at("offset")}, {"nbytes", pd.bytes.size()}};
197+ }
198+ si = end;
199+ }
200+
201+ // Commit: numbered manifest + refreshed HEAD.
202+ int64_t seq = 0;
203+ for (const auto& e : fs::directory_iterator(repo / "manifests")) {
204+ (void)e;
205+ ++seq;
206+ }
207+ char mname[64];
208+ std::snprintf(mname, sizeof(mname), "manifest-%06lld.json",
209+ static_cast<long long>(seq));
210+ json manifest{{"fmodel", kVersion},
211+ {"self", mname},
212+ {"parent", parent_name},
213+ {"step", opts.step},
214+ {"tag", opts.tag},
215+ {"dtype", dtype_name(opts.dtype)},
216+ {"config", json::parse(config_json.empty() ? "{}" : config_json)},
217+ {"tensors", tensors}};
218+ const fs::path mpath = repo / "manifests" / mname;
219+ {
220+ std::ofstream out(mpath);
221+ out << manifest.dump(1);
222+ }
223+ fs::copy_file(mpath, head, fs::copy_options::overwrite_existing);
224+
225+ std::printf("fmodel: %s — %zu tensors reused (%.1f MB), %zu written (%.1f MB)\n",
226+ (repo / mname).c_str(), reused, double(reused_bytes) / 1e6,
227+ fresh.size(), double(written_bytes) / 1e6);
228+ return mpath.string();
229+}
230+
231+// ---- Snapshot ---------------------------------------------------------------
232+
233+struct Snapshot::Impl {
234+ struct Mapped {
235+ void* base = nullptr;
236+ size_t len = 0;
237+ MTL::Buffer* buffer = nullptr;
238+ };
239+ struct Entry {
240+ DType dtype;
241+ std::vector<int64_t> shape;
242+ std::string shard;
243+ size_t offset, nbytes;
244+ };
245+
246+ fs::path repo;
247+ std::string config;
248+ int64_t step = -1;
249+ std::vector<std::string> order;
250+ std::unordered_map<std::string, Entry> index;
251+ std::unordered_map<std::string, Mapped> maps; // shard rel-path -> mapping
252+
253+ ~Impl() {
254+ for (auto& [rel, m] : maps) {
255+ if (m.buffer) m.buffer->release();
256+ if (m.base) munmap(m.base, m.len);
257+ }
258+ }
259+
260+ const Mapped& mapped(const std::string& rel) {
261+ auto it = maps.find(rel);
262+ if (it != maps.end()) return it->second;
263+
264+ const fs::path p = repo / rel;
265+ const int fd = ::open(p.c_str(), O_RDONLY);
266+ if (fd < 0) die("cannot open shard " + p.string());
267+ struct stat st;
268+ if (fstat(fd, &st) != 0) die("fstat failed: " + p.string());
269+ const size_t len = size_t(st.st_size);
270+ if (len % kPage != 0) die("shard not page-padded: " + p.string());
271+
272+ // MAP_PRIVATE + RW: pages stay clean (backed by the file cache)
273+ // unless someone writes — and the read-only contract means nobody
274+ // does. newBuffer(bytesNoCopy) needs page alignment + page-multiple
275+ // length, which mmap and the writer guarantee.
276+ void* base = mmap(nullptr, len, PROT_READ | PROT_WRITE, MAP_PRIVATE, fd, 0);
277+ ::close(fd);
278+ if (base == MAP_FAILED) die("mmap failed: " + p.string());
279+ if (reinterpret_cast<const uint32_t*>(base)[0] != kShardMagic)
280+ die("bad shard magic: " + p.string());
281+
282+ MTL::Buffer* buf = Device::get().mtl()->newBuffer(
283+ base, len, MTL::ResourceStorageModeShared, nullptr);
284+ if (!buf) die("newBuffer(bytesNoCopy) failed: " + p.string());
285+ return maps.emplace(rel, Mapped{base, len, buf}).first->second;
286+ }
287+};
288+
289+Snapshot Snapshot::open(const std::string& path) {
290+ fs::path manifest(path);
291+ fs::path repo;
292+ if (fs::is_directory(manifest)) {
293+ repo = manifest;
294+ manifest = repo / "manifest-latest.json";
295+ } else {
296+ repo = manifest.parent_path();
297+ if (repo.filename() == "manifests") repo = repo.parent_path();
298+ }
299+ if (!fs::exists(manifest)) die("no manifest at " + manifest.string());
300+
301+ auto impl = std::make_shared<Impl>();
302+ impl->repo = repo;
303+ json j = load_json(manifest);
304+ impl->config = j.value("config", json::object()).dump();
305+ impl->step = j.value("step", int64_t(-1));
306+ for (const auto& [name, e] : j.at("tensors").items()) {
307+ Impl::Entry ent;
308+ ent.dtype = dtype_from_name(e.at("dtype").get<std::string>());
309+ ent.shape = e.at("shape").get<std::vector<int64_t>>();
310+ ent.shard = e.at("shard").get<std::string>();
311+ ent.offset = e.at("offset").get<size_t>();
312+ ent.nbytes = e.at("nbytes").get<size_t>();
313+ impl->order.push_back(name);
314+ impl->index.emplace(name, std::move(ent));
315+ }
316+
317+ Snapshot s;
318+ s.impl_ = std::move(impl);
319+ return s;
320+}
321+
322+const std::string& Snapshot::config_json() const { return impl_->config; }
323+int64_t Snapshot::step() const { return impl_->step; }
324+std::vector<std::string> Snapshot::names() const { return impl_->order; }
325+bool Snapshot::has(const std::string& name) const { return impl_->index.count(name); }
326+
327+const std::vector<int64_t>& Snapshot::shape(const std::string& name) const {
328+ auto it = impl_->index.find(name);
329+ if (it == impl_->index.end()) die("no tensor " + name);
330+ return it->second.shape;
331+}
332+
333+Tensor Snapshot::tensor_f32(const std::string& name) const {
334+ auto it = impl_->index.find(name);
335+ if (it == impl_->index.end()) die("no tensor " + name);
336+ const Impl::Entry& e = it->second;
337+ const Impl::Mapped& m = impl_->mapped(e.shard);
338+
339+ if (e.dtype == DType::F32)
340+ return Tensor::from_buffer(m.buffer, e.offset, e.shape); // zero-copy
341+
342+ // f16/bf16 storage: convert into a fresh f32 tensor.
343+ Tensor out = Tensor::empty(e.shape);
344+ const int64_t n = out.numel();
345+ float* dst = out.data<float>();
346+ const uint8_t* src = static_cast<const uint8_t*>(m.base) + e.offset;
347+ if (e.dtype == DType::F16) {
348+ const f16_t* s = reinterpret_cast<const f16_t*>(src);
349+ for (int64_t i = 0; i < n; ++i) dst[i] = float(s[i]);
350+ } else if (e.dtype == DType::BF16) {
351+ const uint16_t* s = reinterpret_cast<const uint16_t*>(src);
352+ for (int64_t i = 0; i < n; ++i) dst[i] = bf16_to_float(s[i]);
353+ } else {
354+ die("tensor_f32: unsupported stored dtype for " + name);
355+ }
356+ return out;
357+}
358+
359+} // namespace forge::fmodel
added src/core/fmodel.h +80 −0
@@ -0,0 +1,80 @@
1+// Author: Simon-Pierre Boucher — contact@spboucher.ai
2+#pragma once
3+
4+#include "core/autograd.h"
5+#include "core/tensor.h"
6+
7+#include <cstdint>
8+#include <memory>
9+#include <string>
10+#include <utility>
11+#include <vector>
12+
13+// .forge — Forge's home-made, Apple-native, git-style weight format.
14+//
15+// A .forge is a DIRECTORY (a model repository):
16+//
17+// model.forge/
18+// manifest-latest.json <- HEAD: copy of the newest manifest
19+// manifests/manifest-000042.json <- one "commit" per save(): step, tag,
20+// parent, config, tensor index
21+// objects/<16-hex>.fshard <- immutable content-addressed shards
22+//
23+// Three properties make it worth existing next to .pt/safetensors/gguf:
24+//
25+// 1. APPLE-NATIVE ZERO-COPY. Every tensor is aligned to the Apple Silicon
26+// 16 KB page inside its shard, and shards are padded to page multiples,
27+// so loading is mmap + newBuffer(bytesNoCopy): the file-cache pages ARE
28+// the GPU memory (MTLStorageModeShared unified memory). No parse, no
29+// memcpy — a multi-GB model "loads" in milliseconds.
30+//
31+// 2. GIT-STYLE DELTA SAVES. Tensors are content-addressed (FNV-1a 64).
32+// save() rewrites only tensors whose bytes changed since the parent
33+// manifest; unchanged ones are referenced in place. Manifests are tiny
34+// JSON commits with a parent link, so a repo carries its whole history.
35+//
36+// 3. GITHUB-FRIENDLY SHARDING. New tensors are packed into shard files
37+// capped at shard_mb (default 95 MB — under GitHub's 100 MB limit), so
38+// a repo can be pushed as-is.
39+//
40+// dtypes: f32 (zero-copy alias at load) · f16 / bf16 (half size on disk,
41+// converted to f32 at load until the mixed-precision kernels land).
42+namespace forge::fmodel {
43+
44+struct SaveOptions {
45+ DType dtype = DType::F32; // f32 | f16 | bf16 storage
46+ int64_t shard_mb = 95; // shard cap; a bigger single tensor gets its own
47+ std::string tag; // optional human label for the manifest
48+ int64_t step = -1; // training step recorded in the manifest
49+};
50+
51+// Snapshot the (deduped) parameters into repo_dir, creating it if needed.
52+// Returns the manifest path. Repeated saves write only changed tensors.
53+std::string save(const std::string& repo_dir, const std::string& config_json,
54+ const std::vector<std::pair<std::string, Var>>& named_params,
55+ const SaveOptions& opts);
56+
57+// Zero-copy reader for one manifest.
58+class Snapshot {
59+public:
60+ // path: a .forge directory (opens manifest-latest.json) or an explicit
61+ // manifest .json inside one. Dies with a message on malformed repos.
62+ static Snapshot open(const std::string& path);
63+
64+ const std::string& config_json() const;
65+ int64_t step() const;
66+ std::vector<std::string> names() const;
67+ bool has(const std::string& name) const;
68+ const std::vector<int64_t>& shape(const std::string& name) const;
69+
70+ // f32 tensors alias the mmapped shard (READ-ONLY by contract: they feed
71+ // forward passes, never optimizer updates). Other dtypes convert into a
72+ // fresh f32 tensor.
73+ Tensor tensor_f32(const std::string& name) const;
74+
75+private:
76+ struct Impl;
77+ std::shared_ptr<Impl> impl_;
78+};
79+
80+} // namespace forge::fmodel
modified src/core/tensor.cpp +21 −2
@@ -20,15 +20,21 @@ namespace {
20 } // namespace 20 } // namespace
21 21
22 // Owns one pooled MTLBuffer; returns it to the pool when the last Tensor 22 // Owns one pooled MTLBuffer; returns it to the pool when the last Tensor
23−// view drops it. 23+// view drops it. External storages (fmodel mmaps) retain a caller-owned
24+// buffer instead and release it directly — the allocator never sees it.
24 struct Tensor::Storage { 25 struct Tensor::Storage {
25 MTL::Buffer* buffer = nullptr; 26 MTL::Buffer* buffer = nullptr;
27+ bool external = false;
26 28
27 explicit Storage(size_t nbytes) { 29 explicit Storage(size_t nbytes) {
28 buffer = Device::get().allocator().acquire(nbytes); 30 buffer = Device::get().allocator().acquire(nbytes);
29 } 31 }
32+ Storage(MTL::Buffer* ext, bool) : buffer(ext), external(true) {
33+ buffer->retain();
34+ }
30 ~Storage() { 35 ~Storage() {
31− Device::get().allocator().release(buffer); 36+ if (external) buffer->release();
37+ else Device::get().allocator().release(buffer);
32 } 38 }
33 Storage(const Storage&) = delete; 39 Storage(const Storage&) = delete;
34 Storage& operator=(const Storage&) = delete; 40 Storage& operator=(const Storage&) = delete;
@@ -58,6 +64,19 @@ Tensor Tensor::empty(std::vector<int64_t> shape, DType dtype) {
58 return t; 64 return t;
59 } 65 }
60 66
67+Tensor Tensor::from_buffer(MTL::Buffer* buffer, size_t byte_offset,
68+ std::vector<int64_t> shape, DType dtype) {
69+ Tensor t;
70+ t.dtype_ = dtype;
71+ t.shape_ = std::move(shape);
72+ t.strides_ = contiguous_strides(t.shape_);
73+ if (byte_offset % dtype_size(dtype) != 0)
74+ die("Tensor::from_buffer: offset not aligned to dtype size");
75+ t.offset_ = int64_t(byte_offset / dtype_size(dtype));
76+ t.storage_ = std::make_shared<Storage>(buffer, true);
77+ return t;
78+}
79+
61 Tensor Tensor::zeros(std::vector<int64_t> shape, DType dtype) { 80 Tensor Tensor::zeros(std::vector<int64_t> shape, DType dtype) {
62 Tensor t = empty(std::move(shape), dtype); 81 Tensor t = empty(std::move(shape), dtype);
63 std::memset(t.raw(), 0, t.nbytes()); 82 std::memset(t.raw(), 0, t.nbytes());
modified src/core/tensor.h +8 −0
@@ -25,6 +25,14 @@ public:
25 static Tensor zeros(std::vector<int64_t> shape, DType dtype = DType::F32); 25 static Tensor zeros(std::vector<int64_t> shape, DType dtype = DType::F32);
26 static Tensor full(std::vector<int64_t> shape, float value, DType dtype = DType::F32); 26 static Tensor full(std::vector<int64_t> shape, float value, DType dtype = DType::F32);
27 27
28+ // Wrap an externally-owned MTLBuffer region (e.g. an mmapped .fmodel
29+ // shard bridged with newBuffer(bytesNoCopy)). The buffer is retained for
30+ // the storage's lifetime and released — NOT returned to the allocator
31+ // pool — when the last view drops. byte_offset must be a multiple of the
32+ // dtype size (fmodel aligns to the 16 KB page, far stricter).
33+ static Tensor from_buffer(MTL::Buffer* buffer, size_t byte_offset,
34+ std::vector<int64_t> shape, DType dtype = DType::F32);
35+
28 bool defined() const { return storage_ != nullptr; } 36 bool defined() const { return storage_ != nullptr; }
29 DType dtype() const { return dtype_; } 37 DType dtype() const { return dtype_; }
30 int64_t ndim() const { return int64_t(shape_.size()); } 38 int64_t ndim() const { return int64_t(shape_.size()); }
modified src/main.cpp +104 −11
@@ -12,6 +12,7 @@
12 #include <Metal/Metal.hpp> 12 #include <Metal/Metal.hpp>
13 13
14 #include "core/device.h" 14 #include "core/device.h"
15+#include "core/fmodel.h"
15 #include "nn/config.h" 16 #include "nn/config.h"
16 #include "nn/transformer.h" 17 #include "nn/transformer.h"
17 #include "ops/metal/metal_ops.h" 18 #include "ops/metal/metal_ops.h"
@@ -23,11 +24,14 @@
23 24
24 #include <cmath> 25 #include <cmath>
25 #include <cstdio> 26 #include <cstdio>
27+#include <filesystem>
26 #include <fstream> 28 #include <fstream>
27 #include <map> 29 #include <map>
30+#include <memory>
28 #include <random> 31 #include <random>
29 #include <sstream> 32 #include <sstream>
30 #include <string> 33 #include <string>
34+#include <unordered_set>
31 #include <vector> 35 #include <vector>
32 36
33 namespace { 37 namespace {
@@ -41,7 +45,11 @@ void print_usage() {
41 " generate --checkpoint <ckpt> --tokenizer <model> --prompt <text>\n" 45 " generate --checkpoint <ckpt> --tokenizer <model> --prompt <text>\n"
42 " [--temp t] [--top-k k] [--max-tokens n] [--seed s]\n" 46 " [--temp t] [--top-k k] [--max-tokens n] [--seed s]\n"
43 " eval --checkpoint <ckpt> --data <val.bin> [--batches n]\n" 47 " eval --checkpoint <ckpt> --data <val.bin> [--batches n]\n"
44− " info [--config <json>]\n"); 48+ " export --checkpoint <ckpt.bin> --out <model.forge> [--dtype f32|f16|bf16]\n"
49+ " [--shard-mb n] [--tag label] (git-style delta repo, zero-copy load)\n"
50+ " info [--config <json>]\n"
51+ "\n"
52+ "generate/eval also accept a .forge repo (or one of its manifests) as --checkpoint.\n");
45 } 53 }
46 54
47 std::map<std::string, std::string> parse_flags(int argc, char** argv, int start) { 55 std::map<std::string, std::string> parse_flags(int argc, char** argv, int start) {
@@ -115,6 +123,74 @@ int cmd_train(const std::map<std::string, std::string>& flags) {
115 return 0; 123 return 0;
116 } 124 }
117 125
126+// True when the path is a .forge repo directory or a manifest json inside one.
127+bool is_forge_repo(const std::string& path) {
128+ namespace fs = std::filesystem;
129+ if (fs::is_directory(path)) return fs::exists(fs::path(path) / "manifest-latest.json");
130+ return path.size() > 5 && path.rfind(".json") == path.size() - 5;
131+}
132+
133+// Build the model described by a .forge snapshot and alias its weights.
134+// f32 repos are ZERO-COPY: parameters point straight into the mmapped,
135+// page-aligned shards (unified memory — the file cache feeds the GPU).
136+std::unique_ptr<forge::nn::Transformer> load_forge_model(const std::string& path,
137+ forge::Config* out_cfg) {
138+ forge::fmodel::Snapshot snap = forge::fmodel::Snapshot::open(path);
139+ nlohmann::json j = nlohmann::json::parse(snap.config_json());
140+ forge::Config cfg;
141+ if (j.contains("model")) j.at("model").get_to(cfg.model);
142+ if (j.contains("train")) j.at("train").get_to(cfg.train);
143+ auto model = std::make_unique<forge::nn::Transformer>(cfg.model, 0);
144+
145+ std::unordered_set<const void*> seen;
146+ for (const auto& [name, p] : model->named_parameters()) {
147+ if (!seen.insert(p.id()).second) continue; // tied param: already aliased
148+ if (!snap.has(name)) {
149+ std::fprintf(stderr, "forge: %s missing from %s\n", name.c_str(),
150+ path.c_str());
151+ std::exit(1);
152+ }
153+ if (snap.shape(name) != p.value().shape()) {
154+ std::fprintf(stderr, "forge: shape mismatch for %s\n", name.c_str());
155+ std::exit(1);
156+ }
157+ p.value() = snap.tensor_f32(name);
158+ }
159+ if (out_cfg) *out_cfg = cfg;
160+ std::printf("loaded %s (step %lld)\n", path.c_str(),
161+ static_cast<long long>(snap.step()));
162+ return model;
163+}
164+
165+int cmd_export(const std::map<std::string, std::string>& flags) {
166+ const std::string ckpt = flag(flags, "checkpoint");
167+ const std::string out = flag(flags, "out");
168+ if (ckpt.empty() || out.empty()) {
169+ print_usage();
170+ return 1;
171+ }
172+ select_backend(flag(flags, "backend", "metal"));
173+
174+ const std::string config_json = forge::train::read_checkpoint_config(ckpt);
175+ nlohmann::json j = nlohmann::json::parse(config_json);
176+ forge::ModelConfig mc;
177+ j.at("model").get_to(mc);
178+ forge::nn::Transformer model(mc, 0);
179+ forge::train::CheckpointData meta =
180+ forge::train::load_checkpoint(ckpt, model.named_parameters(), nullptr);
181+
182+ forge::fmodel::SaveOptions opts;
183+ const std::string dt = flag(flags, "dtype", "f32");
184+ opts.dtype = dt == "f16" ? forge::DType::F16
185+ : dt == "bf16" ? forge::DType::BF16
186+ : forge::DType::F32;
187+ opts.shard_mb = std::stoll(flag(flags, "shard-mb", "95"));
188+ opts.tag = flag(flags, "tag");
189+ opts.step = meta.step;
190+ forge::fmodel::save(out, config_json, model.named_parameters(), opts);
191+ return 0;
192+}
193+
118 int cmd_generate(const std::map<std::string, std::string>& flags) { 194 int cmd_generate(const std::map<std::string, std::string>& flags) {
119 const std::string ckpt = flag(flags, "checkpoint"); 195 const std::string ckpt = flag(flags, "checkpoint");
120 const std::string tok_path = flag(flags, "tokenizer"); 196 const std::string tok_path = flag(flags, "tokenizer");
@@ -128,12 +204,21 @@ int cmd_generate(const std::map<std::string, std::string>& flags) {
128 const int64_t max_tokens = std::stoll(flag(flags, "max-tokens", "256")); 204 const int64_t max_tokens = std::stoll(flag(flags, "max-tokens", "256"));
129 const uint64_t seed = std::stoull(flag(flags, "seed", "1234")); 205 const uint64_t seed = std::stoull(flag(flags, "seed", "1234"));
130 206
131− // model config comes from the checkpoint itself 207+ // model config comes from the checkpoint / .forge repo itself
132− nlohmann::json j = nlohmann::json::parse(forge::train::read_checkpoint_config(ckpt)); 208+ std::unique_ptr<forge::nn::Transformer> model_ptr;
133 forge::ModelConfig mc; 209 forge::ModelConfig mc;
134− j.at("model").get_to(mc); 210+ if (is_forge_repo(ckpt)) {
135− forge::nn::Transformer model(mc, 0); 211+ forge::Config cfg;
136− forge::train::load_checkpoint(ckpt, model.named_parameters(), nullptr); 212+ model_ptr = load_forge_model(ckpt, &cfg);
213+ mc = cfg.model;
214+ } else {
215+ nlohmann::json j =
216+ nlohmann::json::parse(forge::train::read_checkpoint_config(ckpt));
217+ j.at("model").get_to(mc);
218+ model_ptr = std::make_unique<forge::nn::Transformer>(mc, 0);
219+ forge::train::load_checkpoint(ckpt, model_ptr->named_parameters(), nullptr);
220+ }
221+ forge::nn::Transformer& model = *model_ptr;
137 222
138 forge::tok::BPETokenizer tokenizer; 223 forge::tok::BPETokenizer tokenizer;
139 tokenizer.load(tok_path); 224 tokenizer.load(tok_path);
@@ -200,12 +285,19 @@ int cmd_eval(const std::map<std::string, std::string>& flags) {
200 select_backend(flag(flags, "backend", "metal")); 285 select_backend(flag(flags, "backend", "metal"));
201 const int64_t batches = std::stoll(flag(flags, "batches", "50")); 286 const int64_t batches = std::stoll(flag(flags, "batches", "50"));
202 287
203− nlohmann::json j = nlohmann::json::parse(forge::train::read_checkpoint_config(ckpt)); 288+ std::unique_ptr<forge::nn::Transformer> model_ptr;
204 forge::Config cfg; 289 forge::Config cfg;
205− if (j.contains("model")) j.at("model").get_to(cfg.model); 290+ if (is_forge_repo(ckpt)) {
206− if (j.contains("train")) j.at("train").get_to(cfg.train); 291+ model_ptr = load_forge_model(ckpt, &cfg);
207− forge::nn::Transformer model(cfg.model, 0); 292+ } else {
208− forge::train::load_checkpoint(ckpt, model.named_parameters(), nullptr); 293+ nlohmann::json j =
294+ nlohmann::json::parse(forge::train::read_checkpoint_config(ckpt));
295+ if (j.contains("model")) j.at("model").get_to(cfg.model);
296+ if (j.contains("train")) j.at("train").get_to(cfg.train);
297+ model_ptr = std::make_unique<forge::nn::Transformer>(cfg.model, 0);
298+ forge::train::load_checkpoint(ckpt, model_ptr->named_parameters(), nullptr);
299+ }
300+ forge::nn::Transformer& model = *model_ptr;
209 301
210 forge::train::DataLoader loader(data, cfg.model.context_length, 0); 302 forge::train::DataLoader loader(data, cfg.model.context_length, 0);
211 const int64_t B = cfg.train.batch_size, T = cfg.model.context_length; 303 const int64_t B = cfg.train.batch_size, T = cfg.model.context_length;
@@ -243,6 +335,7 @@ int main(int argc, char** argv) {
243 else if (command == "train") rc = cmd_train(flags); 335 else if (command == "train") rc = cmd_train(flags);
244 else if (command == "generate") rc = cmd_generate(flags); 336 else if (command == "generate") rc = cmd_generate(flags);
245 else if (command == "eval") rc = cmd_eval(flags); 337 else if (command == "eval") rc = cmd_eval(flags);
338+ else if (command == "export") rc = cmd_export(flags);
246 else { 339 else {
247 print_usage(); 340 print_usage();
248 rc = 1; 341 rc = 1;
added tools/fmodel.py +129 −0
@@ -0,0 +1,129 @@
1+# Author: Simon-Pierre Boucher — contact@spboucher.ai
2+"""Inspect and convert .forge weight repositories (see src/core/fmodel.h for
3+the format: git-style manifests + content-addressed page-aligned shards).
4+
5+Usage:
6+ python3 tools/fmodel.py inspect model.forge [--manifest manifest-000002.json]
7+ python3 tools/fmodel.py log model.forge # manifest history
8+ python3 tools/fmodel.py to-safetensors model.forge out.safetensors
9+ [--manifest m.json] # pure-numpy writer, no torch needed
10+
11+Only numpy is required. The safetensors export makes .forge weights loadable
12+from PyTorch/HF (`safetensors.torch.load_file`) for interop.
13+"""
14+import argparse
15+import json
16+import os
17+import struct
18+
19+import numpy as np
20+
21+PAGE = 16384
22+NP_DTYPE = {"f32": np.float32, "f16": np.float16, "u16": np.uint16, "i32": np.int32}
23+ST_DTYPE = {"f32": "F32", "f16": "F16", "bf16": "BF16", "u16": "U16", "i32": "I32"}
24+
25+
26+def resolve_manifest(repo, manifest=None):
27+ if os.path.isfile(repo): # given a manifest path directly
28+ return os.path.dirname(os.path.dirname(os.path.abspath(repo))) \
29+ if os.path.basename(os.path.dirname(repo)) == "manifests" \
30+ else os.path.dirname(os.path.abspath(repo)), repo
31+ path = os.path.join(repo, "manifests", manifest) if manifest else \
32+ os.path.join(repo, "manifest-latest.json")
33+ return repo, path
34+
35+
36+def load(repo, manifest=None):
37+ repo, mpath = resolve_manifest(repo, manifest)
38+ with open(mpath) as f:
39+ return repo, json.load(f)
40+
41+
42+def tensor_bytes(repo, entry):
43+ with open(os.path.join(repo, entry["shard"]), "rb") as f:
44+ f.seek(entry["offset"])
45+ return f.read(entry["nbytes"])
46+
47+
48+def tensor_array(repo, entry):
49+ raw = tensor_bytes(repo, entry)
50+ dt = entry["dtype"]
51+ if dt == "bf16": # numpy has no bf16: widen via bit tricks
52+ u16 = np.frombuffer(raw, dtype=np.uint16).astype(np.uint32) << 16
53+ return u16.view(np.float32).reshape(entry["shape"])
54+ return np.frombuffer(raw, dtype=NP_DTYPE[dt]).reshape(entry["shape"])
55+
56+
57+def cmd_inspect(args):
58+ repo, m = load(args.repo, args.manifest)
59+ cfg = m.get("config", {}).get("model", {})
60+ print(f"{m.get('self', '?')} step={m.get('step')} tag='{m.get('tag', '')}'"
61+ f" dtype={m.get('dtype')} parent='{m.get('parent', '')}'")
62+ if cfg:
63+ print(f"model: {cfg.get('name')} layers={cfg.get('n_layers')}"
64+ f" d_model={cfg.get('d_model')} vocab={cfg.get('vocab_size')}")
65+ total = 0
66+ shards = {}
67+ for name, e in sorted(m["tensors"].items()):
68+ total += e["nbytes"]
69+ shards.setdefault(e["shard"], 0)
70+ shards[e["shard"]] += e["nbytes"]
71+ print(f" {name:<40} {e['dtype']:<5} {str(e['shape']):<20}"
72+ f" {e['nbytes'] / 1e6:8.2f} MB {e['shard'].split('/')[-1][:8]}…")
73+ print(f"{len(m['tensors'])} tensors, {total / 1e6:.1f} MB across "
74+ f"{len(shards)} shard(s)")
75+
76+
77+def cmd_log(args):
78+ repo, _ = resolve_manifest(args.repo)
79+ mdir = os.path.join(repo, "manifests")
80+ for name in sorted(os.listdir(mdir)):
81+ with open(os.path.join(mdir, name)) as f:
82+ m = json.load(f)
83+ shards = {e["shard"] for e in m["tensors"].values()}
84+ print(f"{name} step={m.get('step'):>8} tag='{m.get('tag', '')}'"
85+ f" dtype={m.get('dtype')} shards={len(shards)}"
86+ f" parent='{m.get('parent', '')}'")
87+
88+
89+def cmd_to_safetensors(args):
90+ repo, m = load(args.repo, args.manifest)
91+ header = {"__metadata__": {"format": "forge",
92+ "step": str(m.get("step", -1)),
93+ "config": json.dumps(m.get("config", {}))}}
94+ offset = 0
95+ order = sorted(m["tensors"].items())
96+ for name, e in order:
97+ header[name] = {"dtype": ST_DTYPE[e["dtype"]], "shape": e["shape"],
98+ "data_offsets": [offset, offset + e["nbytes"]]}
99+ offset += e["nbytes"]
100+ hjson = json.dumps(header, separators=(",", ":")).encode()
101+ pad = (8 - len(hjson) % 8) % 8 # safetensors aligns the header to 8
102+ hjson += b" " * pad
103+ with open(args.out, "wb") as f:
104+ f.write(struct.pack("<Q", len(hjson)))
105+ f.write(hjson)
106+ for name, e in order:
107+ f.write(tensor_bytes(repo, e))
108+ print(f"wrote {args.out} ({(8 + len(hjson) + offset) / 1e6:.1f} MB, "
109+ f"{len(order)} tensors)")
110+
111+
112+def main():
113+ ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
114+ sub = ap.add_subparsers(dest="cmd", required=True)
115+ for name, fn in [("inspect", cmd_inspect), ("log", cmd_log),
116+ ("to-safetensors", cmd_to_safetensors)]:
117+ p = sub.add_parser(name)
118+ p.add_argument("repo", help=".forge directory or manifest json")
119+ p.add_argument("--manifest", default=None,
120+ help="specific manifest inside manifests/")
121+ if name == "to-safetensors":
122+ p.add_argument("out")
123+ p.set_defaults(fn=fn)
124+ args = ap.parse_args()
125+ args.fn(args)
126+
127+
128+if __name__ == "__main__":
129+ main()
130
Menu

Theme
↑↓ navigate ↵ open esc close / or ⌘K anywhere