SPB Git

spb/forge Public MIT

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

llmtrainingcppmetalapple-silicon
C++ 61.2% C 23% Python 7.6% TeX 7.2% CMake 1.1%

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 5 days ago (Aug 5, 2026) parent 4054156

Showing 7 changed files with +702 and −13

modified CMakeLists.txt +1 −0
@@ -66,6 +66,7 @@ add_library(forge_core STATIC
66 66 src/core/allocator.cpp
67 67 src/core/device.cpp
68 68 src/core/autograd.cpp
69 + src/core/fmodel.cpp
69 70 src/ops/cpu/cpu_ops.cpp
70 71 src/ops/metal/metal_ops.cpp
71 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 20 } // namespace
21 21
22 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 25 struct Tensor::Storage {
25 26 MTL::Buffer* buffer = nullptr;
27 + bool external = false;
26 28
27 29 explicit Storage(size_t nbytes) {
28 30 buffer = Device::get().allocator().acquire(nbytes);
29 31 }
32 + Storage(MTL::Buffer* ext, bool) : buffer(ext), external(true) {
33 + buffer->retain();
34 + }
30 35 ~Storage() {
31 − Device::get().allocator().release(buffer);
36 + if (external) buffer->release();
37 + else Device::get().allocator().release(buffer);
32 38 }
33 39 Storage(const Storage&) = delete;
34 40 Storage& operator=(const Storage&) = delete;
@@ -58,6 +64,19 @@ Tensor Tensor::empty(std::vector<int64_t> shape, DType dtype) {
58 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 80 Tensor Tensor::zeros(std::vector<int64_t> shape, DType dtype) {
62 81 Tensor t = empty(std::move(shape), dtype);
63 82 std::memset(t.raw(), 0, t.nbytes());
modified src/core/tensor.h +8 −0
@@ -25,6 +25,14 @@ public:
25 25 static Tensor zeros(std::vector<int64_t> shape, DType dtype = DType::F32);
26 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 36 bool defined() const { return storage_ != nullptr; }
29 37 DType dtype() const { return dtype_; }
30 38 int64_t ndim() const { return int64_t(shape_.size()); }
modified src/main.cpp +104 −11
@@ -12,6 +12,7 @@
12 12 #include <Metal/Metal.hpp>
13 13
14 14 #include "core/device.h"
15 +#include "core/fmodel.h"
15 16 #include "nn/config.h"
16 17 #include "nn/transformer.h"
17 18 #include "ops/metal/metal_ops.h"
@@ -23,11 +24,14 @@
23 24
24 25 #include <cmath>
25 26 #include <cstdio>
27 +#include <filesystem>
26 28 #include <fstream>
27 29 #include <map>
30 +#include <memory>
28 31 #include <random>
29 32 #include <sstream>
30 33 #include <string>
34 +#include <unordered_set>
31 35 #include <vector>
32 36
33 37 namespace {
@@ -41,7 +45,11 @@ void print_usage() {
41 45 " generate --checkpoint <ckpt> --tokenizer <model> --prompt <text>\n"
42 46 " [--temp t] [--top-k k] [--max-tokens n] [--seed s]\n"
43 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 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 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 194 int cmd_generate(const std::map<std::string, std::string>& flags) {
119 195 const std::string ckpt = flag(flags, "checkpoint");
120 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 204 const int64_t max_tokens = std::stoll(flag(flags, "max-tokens", "256"));
129 205 const uint64_t seed = std::stoull(flag(flags, "seed", "1234"));
130 206
131 − // model config comes from the checkpoint itself
132 − nlohmann::json j = nlohmann::json::parse(forge::train::read_checkpoint_config(ckpt));
207 + // model config comes from the checkpoint / .forge repo itself
208 + std::unique_ptr<forge::nn::Transformer> model_ptr;
133 209 forge::ModelConfig mc;
134 − j.at("model").get_to(mc);
135 − forge::nn::Transformer model(mc, 0);
136 − forge::train::load_checkpoint(ckpt, model.named_parameters(), nullptr);
210 + if (is_forge_repo(ckpt)) {
211 + forge::Config cfg;
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 223 forge::tok::BPETokenizer tokenizer;
139 224 tokenizer.load(tok_path);
@@ -200,12 +285,19 @@ int cmd_eval(const std::map<std::string, std::string>& flags) {
200 285 select_backend(flag(flags, "backend", "metal"));
201 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 289 forge::Config cfg;
205 − if (j.contains("model")) j.at("model").get_to(cfg.model);
206 − if (j.contains("train")) j.at("train").get_to(cfg.train);
207 − forge::nn::Transformer model(cfg.model, 0);
208 − forge::train::load_checkpoint(ckpt, model.named_parameters(), nullptr);
290 + if (is_forge_repo(ckpt)) {
291 + model_ptr = load_forge_model(ckpt, &cfg);
292 + } else {
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 302 forge::train::DataLoader loader(data, cfg.model.context_length, 0);
211 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 335 else if (command == "train") rc = cmd_train(flags);
244 336 else if (command == "generate") rc = cmd_generate(flags);
245 337 else if (command == "eval") rc = cmd_eval(flags);
338 + else if (command == "export") rc = cmd_export(flags);
246 339 else {
247 340 print_usage();
248 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