diff --git a/CMakeLists.txt b/CMakeLists.txt index a4b98e4..69a186f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -85,6 +85,9 @@ add_library(mlxforge_core STATIC src/tokenizer/bpe.cpp src/tokenizer/spm.cpp src/tokenizer/tokenizer.cpp + src/inspect/safetensors_header.cpp + src/inspect/model_schema.cpp + src/inspect/schematic_html.cpp ) target_include_directories(mlxforge_core PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/src) # stb_image is header-only and used only inside src/vision/image_decode.cpp, so it diff --git a/apps/mlxforge_cli.cpp b/apps/mlxforge_cli.cpp index 2fbdfca..9087653 100644 --- a/apps/mlxforge_cli.cpp +++ b/apps/mlxforge_cli.cpp @@ -24,6 +24,12 @@ // - Embeds text and prints the (by default unit-normalized) vector. With no flags the model // self-selects its convention (a Qwen3-Embedding checkpoint uses last-token pooling + a // trailing EOS). The embedding smoke/golden-reference harness for the library. +// mlxforge-cli schematic [--out file.html] [--open] +// - Writes a self-contained interactive HTML infographic of the model's architecture (block +// schematic with matrix dims, parameter distribution, per-layer tensor explorer), read from +// metadata only — safetensors headers or the GGUF tensor directory; no weights are loaded +// and no MLX arrays are created. Defaults to ./-schematic.html; --open opens it in +// the default browser. // // / is either a local model directory or a HuggingFace repo id (e.g. mlx-community/Llama-3.2-1B-Instruct-4bit), // which will be downloaded on first use. @@ -32,6 +38,8 @@ #include #include #include +#include +#include #include #include #include @@ -43,6 +51,8 @@ #include "core/logging.h" #include "core/model_source.h" #include "core/weights.h" +#include "inspect/model_schema.h" +#include "inspect/schematic_html.h" #include "model/model_factory.h" #include "model/qwen3_vl.h" #include "model/vision/vit.h" @@ -445,6 +455,57 @@ int run_embed(const std::string& spec, const std::string& text, return 0; } +// Architecture infographic: read tensor metadata only (safetensors headers / +// the GGUF tensor directory — no weight load, no MLX arrays) and write a +// self-contained interactive HTML schematic of the model. +int run_schematic(const std::string& spec, std::string out_path, bool open_after) { + const std::string resolved = mlxforge::resolve_model_dir(spec); + + // A repo-id spec reads better as the display name than the resolved snapshot + // dir's hash basename; the builders fall back to the basename when empty. + const std::string display = mlxforge::looks_like_repo_id(spec) ? spec : ""; + + mlxforge::inspect::ModelSchema schema; + if (mlxforge::is_gguf_path(resolved)) { + schema = mlxforge::inspect::build_schema_from_gguf(resolved, display); + } else { + const auto cfg = mlxforge::ModelConfig::from_file(resolved + "/config.json"); + schema = mlxforge::inspect::build_schema_from_safetensors(resolved, cfg, display); + } + + if (out_path.empty()) { + // -schematic.html from the display name's last path component. + std::string base = schema.model_name; + if (const size_t slash = base.find_last_of('/'); slash != std::string::npos) + base = base.substr(slash + 1); + if (const size_t dot = base.rfind(".gguf"); dot != std::string::npos) base.resize(dot); + out_path = base + "-schematic.html"; + } + + const std::string html = mlxforge::inspect::render_schematic_html(schema.to_json()); + std::ofstream out(out_path, std::ios::binary); + if (!out || !(out << html)) { + mlxforge::log::error("schematic: cannot write '{}'", out_path); + return 1; + } + out.close(); + + mlxforge::log::info("schematic: {} ({}) — {} tensors, {} params, {} layers", schema.model_name, + schema.family, schema.tensors.size(), schema.total_params, schema.cfg.n_layers); + // The output path is the command's primary output (like dump-weights). + std::printf("%s\n", out_path.c_str()); + + if (open_after) { + // Single-quote the path for the shell, escaping any embedded quote. + std::string quoted = "'"; + for (char c : out_path) quoted += (c == '\'') ? std::string("'\\''") : std::string(1, c); + quoted += "'"; + if (std::system(("open " + quoted).c_str()) != 0) + mlxforge::log::warn("schematic: failed to open '{}'", out_path); + } + return 0; +} + } // namespace int main(int argc, char** argv) { @@ -552,6 +613,25 @@ int main(int argc, char** argv) { } return run_embed(argv[2], argv[3], opts); } + if (cmd == "schematic") { + // Architecture infographic from tensor metadata (no weight load). + if (argc < 3) { + std::fprintf(stderr, "usage: mlxforge-cli schematic [--out file.html] [--open]\n"); + return 2; + } + std::string out_path; + bool open_after = false; + for (int i = 3; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--out" && i + 1 < argc) out_path = argv[++i]; + else if (a == "--open") open_after = true; + else { + std::fprintf(stderr, "schematic: unknown argument '%s'\n", a.c_str()); + return 2; + } + } + return run_schematic(argv[2], out_path, open_after); + } // No subcommand: run the smoke test by default return run_smoke(); diff --git a/src/core/gguf.cpp b/src/core/gguf.cpp index 31c8620..4d9d9ae 100644 --- a/src/core/gguf.cpp +++ b/src/core/gguf.cpp @@ -662,6 +662,73 @@ GgufModel load_gguf_config_and_tokenizer(const std::string& gguf_path) { return g; } +std::string ggml_type_name(uint32_t t) { + // llama.cpp's ggml_type enum (ids 4/5 were removed upstream and never ship). + static const std::unordered_map kNames = { + {0, "F32"}, {1, "F16"}, {2, "Q4_0"}, {3, "Q4_1"}, {6, "Q5_0"}, + {7, "Q5_1"}, {8, "Q8_0"}, {9, "Q8_1"}, {10, "Q2_K"}, {11, "Q3_K"}, + {12, "Q4_K"}, {13, "Q5_K"}, {14, "Q6_K"}, {15, "Q8_K"}, {16, "IQ2_XXS"}, + {17, "IQ2_XS"}, {18, "IQ3_XXS"}, {19, "IQ1_S"}, {20, "IQ4_NL"}, {21, "IQ3_S"}, + {22, "IQ2_S"}, {23, "IQ4_XS"}, {24, "I8"}, {25, "I16"}, {26, "I32"}, + {27, "I64"}, {28, "F64"}, {29, "IQ1_M"}, {30, "BF16"}, + }; + auto it = kNames.find(t); + return it != kNames.end() ? it->second : "type_" + std::to_string(t); +} + +double ggml_bits_per_weight(uint32_t t) { + // bytes-per-block / weights-per-block, so the per-block scale/min overhead is + // included (Q4_0: 18 B / 32 w = 4.5 bpw; K-quants are 256-weight super-blocks). + static const std::unordered_map kBpw = { + {0, 32.0}, {1, 16.0}, {2, 4.5}, {3, 5.0}, {6, 5.5}, + {7, 6.0}, {8, 8.5}, {9, 9.0}, {10, 2.625}, {11, 3.4375}, + {12, 4.5}, {13, 5.5}, {14, 6.5625}, {15, 9.125}, {16, 2.0625}, + {17, 2.3125}, {18, 3.0625}, {19, 1.5625}, {20, 4.5}, {21, 3.4375}, + {22, 2.5}, {23, 4.25}, {24, 8.0}, {25, 16.0}, {26, 32.0}, + {27, 64.0}, {28, 64.0}, {29, 1.75}, {30, 16.0}, + }; + auto it = kBpw.find(t); + return it != kBpw.end() ? it->second : 0.0; +} + +GgufInspection inspect_gguf(const std::string& gguf_path) { + const GgufMetadata meta = parse_gguf_metadata(gguf_path); + GgufInspection out; + out.head = model_head_from_metadata(meta); + + const uint64_t alignment = + std::max(1, static_cast(gm_int_or(meta, "general.alignment", 32))); + const GgufDirectory dir = parse_gguf_directory(gguf_path, alignment); + + std::ifstream f(gguf_path, std::ios::binary | std::ios::ate); + if (!f) throw std::runtime_error("gguf: cannot open '" + gguf_path + "'"); + out.file_bytes = static_cast(f.tellg()); + + // Per-tensor bytes come from the gaps between consecutive data offsets (the + // last tensor runs to end-of-file). Exact for every type, including ones this + // build cannot load — no bytes-per-type table involved. Offsets are + // alignment-padded, so each tensor absorbs its own trailing pad (< alignment). + std::vector order(dir.tensors.size()); + for (size_t i = 0; i < order.size(); ++i) order[i] = i; + std::sort(order.begin(), order.end(), + [&](size_t a, size_t b) { return dir.tensors[a].offset < dir.tensors[b].offset; }); + + out.tensors.resize(dir.tensors.size()); + for (size_t k = 0; k < order.size(); ++k) { + const GgufTensorInfo& t = dir.tensors[order[k]]; + GgufTensorMeta& m = out.tensors[order[k]]; + m.name = t.name; + m.canonical = remap_gguf_key(t.name).value_or(""); + m.ggml_type = t.ggml_type; + for (auto it = t.dims.rbegin(); it != t.dims.rend(); ++it) + m.shape.push_back(static_cast(*it)); + const uint64_t end = (k + 1 < order.size()) ? dir.tensors[order[k + 1]].offset + : out.file_bytes - dir.data_start; + m.bytes = end - t.offset; + } + return out; +} + GgufModel load_gguf_model(const std::string& gguf_path) { log::info("gguf: loading '{}'", gguf_path); const GgufMetadata meta = parse_gguf_metadata(gguf_path); diff --git a/src/core/gguf.h b/src/core/gguf.h index 8baa601..b5ec673 100644 --- a/src/core/gguf.h +++ b/src/core/gguf.h @@ -52,4 +52,38 @@ GgufModel load_gguf_model(const std::string& gguf_path); // server uses it on the main thread while the worker loads the weights. GgufModel load_gguf_config_and_tokenizer(const std::string& gguf_path); +// ----- Metadata-only tensor inspection (CLI `schematic`) -------------------- +// The tensor-info section of a GGUF file carries every tensor's name, ggml +// type, dims and data offset; reading it costs a header parse, not a weight +// load. inspect_gguf exposes that directory (plus the config/tokenizer head) +// for model-introspection tooling, creating no MLX arrays. + +struct GgufTensorMeta { + std::string name; // raw ggml name ("blk.0.attn_q.weight") + std::string canonical; // remapped HF key; "" if unrecognized/dropped + uint32_t ggml_type = 0; + std::vector shape; // MLX order (ggml dims reversed); LOGICAL dims + uint64_t bytes = 0; // on-disk bytes, derived from the data offsets +}; + +struct GgufInspection { + GgufModel head; // config + tokenizer material, no weights + std::vector tensors; + uint64_t file_bytes = 0; +}; + +// Parse the metadata + tensor directory of a GGUF file. Creates no MLX arrays, +// so it is safe to call on any thread. Throws like load_gguf_config_and_tokenizer +// on an unsupported architecture or a malformed file. +GgufInspection inspect_gguf(const std::string& gguf_path); + +// Display name for a ggml tensor type id ("Q4_K", "F16", ...); "type_" for +// an id this build does not know. +std::string ggml_type_name(uint32_t t); + +// Effective bits per weight for a ggml tensor type, including the per-block +// scale/min overhead (e.g. Q4_0 = 4.5: 18 bytes per 32 weights); 0 if unknown. +// Display/estimation only — exact byte counts come from the data offsets. +double ggml_bits_per_weight(uint32_t t); + } // namespace mlxforge diff --git a/src/inspect/model_schema.cpp b/src/inspect/model_schema.cpp new file mode 100644 index 0000000..2bfc580 --- /dev/null +++ b/src/inspect/model_schema.cpp @@ -0,0 +1,516 @@ +#include "inspect/model_schema.h" + +#include +#include +#include +#include +#include + +#include "core/gguf.h" +#include "core/logging.h" +#include "core/weights.h" +#include "inspect/safetensors_header.h" + +namespace mlxforge::inspect { + +namespace { + +bool ends_with(const std::string& s, const std::string& suffix) { + return s.size() >= suffix.size() && + s.compare(s.size() - suffix.size(), suffix.size(), suffix) == 0; +} +bool starts_with(const std::string& s, const std::string& prefix) { + return s.rfind(prefix, 0) == 0; +} +bool contains(const std::string& s, const std::string& sub) { + return s.find(sub) != std::string::npos; +} + +std::string path_basename(const std::string& path) { + std::string p = path; + while (!p.empty() && p.back() == '/') p.pop_back(); + const size_t slash = p.find_last_of('/'); + return slash == std::string::npos ? p : p.substr(slash + 1); +} + +uint64_t prod(const std::vector& shape) { + uint64_t n = 1; + for (int64_t d : shape) n *= static_cast(d); + return n; +} + +// Preferred in-block display order for the decode-matmul list (q before k +// before v..., not alphabetical). Unknown modules sort after, by name. +int module_rank(const std::string& module) { + static const std::pair kRanks[] = { + {"self_attn.q_proj", 0}, {"self_attn.k_proj", 1}, {"self_attn.v_proj", 2}, + {"self_attn.o_proj", 3}, {"mlp.gate.", 4}, {"mlp.gate_proj", 5}, + {"mlp.up_proj", 6}, {"mlp.down_proj", 7}, + {"mlp.switch_mlp.gate_proj", 5}, {"mlp.switch_mlp.up_proj", 6}, + {"mlp.switch_mlp.down_proj", 7}, + }; + for (const auto& [key, rank] : kRanks) + if (contains(module, key)) return rank; + return 8; +} + +} // namespace + +std::string component_of(const std::string& key) { + if (starts_with(key, "visual.")) return "vision"; + if (starts_with(key, "model.embed_tokens")) return "embed"; + if (starts_with(key, "lm_head")) return "lm_head"; + if (starts_with(key, "model.norm")) return "norm"; + if (starts_with(key, "model.layers.")) { + if (contains(key, ".self_attn.")) return "attn"; + if (contains(key, ".linear_attn.")) return "linear_attn"; + // MoE pieces: the stacked experts (switch_mlp), raw per-expert tensors, + // the router (".mlp.gate." — distinct from the dense ".mlp.gate_proj."), + // and any shared expert. + if (contains(key, ".mlp.switch_mlp.") || contains(key, ".mlp.experts.") || + contains(key, ".mlp.gate.") || contains(key, ".mlp.shared_expert")) { + return "moe"; + } + if (contains(key, "layernorm") || contains(key, ".norm")) return "norm"; + if (contains(key, ".mlp.")) return "mlp"; + return "other"; + } + return "other"; +} + +int layer_of(const std::string& key) { + static const std::string kPrefix = "model.layers."; + if (!starts_with(key, kPrefix)) return -1; + const size_t dot = key.find('.', kPrefix.size()); + if (dot == std::string::npos) return -1; + const std::string idx = key.substr(kPrefix.size(), dot - kPrefix.size()); + if (idx.empty() || !std::all_of(idx.begin(), idx.end(), + [](unsigned char c) { return std::isdigit(c); })) { + return -1; + } + return std::stoi(idx); +} + +namespace { + +// Collapse raw per-expert MoE tensors ("...mlp.experts..gate_proj.weight") +// into one row per (layer, proj): "...mlp.experts.*.gate_proj.weight" with the +// expert count prepended to the shape. Keeps the table readable on 128-expert +// models; pre-stacked switch_mlp checkpoints pass through untouched. +std::vector aggregate_per_expert(std::vector in) { + struct Agg { + TensorEntry entry; // first expert's entry, name rewritten + uint64_t count = 0; + }; + std::map aggs; + std::vector out; + out.reserve(in.size()); + + for (auto& e : in) { + const size_t exp_pos = e.name.find(".experts."); + if (exp_pos == std::string::npos) { + out.push_back(std::move(e)); + continue; + } + const size_t idx_begin = exp_pos + std::strlen(".experts."); + size_t idx_end = idx_begin; + while (idx_end < e.name.size() && std::isdigit(static_cast(e.name[idx_end]))) + ++idx_end; + if (idx_end == idx_begin || idx_end >= e.name.size() || e.name[idx_end] != '.') { + out.push_back(std::move(e)); // not the per-expert pattern + continue; + } + const std::string agg_name = + e.name.substr(0, idx_begin) + "*" + e.name.substr(idx_end); + auto [it, fresh] = aggs.try_emplace(agg_name); + if (fresh) { + it->second.entry = e; + it->second.entry.name = agg_name; + } else { + it->second.entry.params += e.params; + it->second.entry.bytes += e.bytes; + } + ++it->second.count; + } + + for (auto& [_, agg] : aggs) { + agg.entry.shape.insert(agg.entry.shape.begin(), static_cast(agg.count)); + if (!agg.entry.stored_shape.empty()) + agg.entry.stored_shape.insert(agg.entry.stored_shape.begin(), + static_cast(agg.count)); + out.push_back(std::move(agg.entry)); + } + return out; +} + +// Family label, mirroring create_model's dispatch (model/model_factory.cpp). +std::string detect_family(const ModelConfig& cfg, bool has_qk_norm) { + if (cfg.full_attention_interval > 0) return "qwen3.5-hybrid"; + if (cfg.num_experts > 0) return "qwen3-moe"; + if (cfg.has_vision_tower()) return "qwen3-vl"; + if (has_qk_norm) return "qwen3"; + return "llama"; +} + +// Aggregates, totals, sort order, derived math — shared by both builders. +void finalize(ModelSchema& s) { + const ModelConfig& cfg = s.cfg; + s.head_dim = cfg.head_dim > 0 ? cfg.head_dim : (cfg.n_heads > 0 ? cfg.hidden / cfg.n_heads : 0); + s.gqa_ratio = cfg.n_kv_heads > 0 ? cfg.n_heads / cfg.n_kv_heads : 1; + s.n_full_attn_layers = 0; + for (int i = 0; i < cfg.n_layers; ++i) + if (!cfg.is_linear_layer(i)) ++s.n_full_attn_layers; + // fp16 cache: K + V per full-attention layer, 2 bytes per element. + s.kv_bytes_per_token = + static_cast(s.n_full_attn_layers) * 2.0 * cfg.n_kv_heads * s.head_dim * 2.0; + + bool has_qk_norm = false; + for (const auto& e : s.tensors) + if (contains(e.name, ".self_attn.q_norm.")) { has_qk_norm = true; break; } + s.family = detect_family(cfg, has_qk_norm); + s.tied_embeddings = std::none_of(s.tensors.begin(), s.tensors.end(), [](const TensorEntry& e) { + return starts_with(e.name, "lm_head."); + }); + + std::sort(s.tensors.begin(), s.tensors.end(), [](const TensorEntry& a, const TensorEntry& b) { + if (a.layer != b.layer) return a.layer < b.layer; + return a.name < b.name; + }); + + s.total_params = 0; + s.total_bytes = s.dropped_bytes; + s.by_component.clear(); + s.by_layer.assign(std::max(0, cfg.n_layers), ComponentAgg{}); + for (const auto& e : s.tensors) { + s.total_params += e.params; + s.total_bytes += e.bytes; + auto& comp = s.by_component[e.component]; + comp.params += e.params; + comp.bytes += e.bytes; + if (e.layer >= 0 && e.layer < static_cast(s.by_layer.size())) { + s.by_layer[e.layer].params += e.params; + s.by_layer[e.layer].bytes += e.bytes; + } + } + + // Decode matmul shapes (M=1) from the actual tensors of representative + // layers — the first full-attention layer, plus the first linear-attention + // layer for hybrid models. Tensor-derived dims catch config/checkpoint drift + // (and the attn_output_gate 2x q_proj) for free. + s.decode_matmuls.clear(); + std::vector reps; + for (int i = 0; i < cfg.n_layers; ++i) + if (!cfg.is_linear_layer(i)) { reps.push_back(i); break; } + for (int i = 0; i < cfg.n_layers; ++i) + if (cfg.is_linear_layer(i)) { reps.push_back(i); break; } + for (int rep : reps) { + const std::string prefix = "model.layers." + std::to_string(rep) + "."; + std::vector block; + for (const auto& e : s.tensors) { + if (e.layer != rep || e.shape.size() < 2) continue; + if (e.component == "norm" || e.component == "other") continue; + MatmulShape m; + m.name = e.name.substr(prefix.size()); + if (ends_with(m.name, ".weight")) m.name.resize(m.name.size() - std::strlen(".weight")); + m.in = e.shape.back(); + m.out = e.shape[e.shape.size() - 2]; + m.quant = e.quant; + if (e.shape.size() == 3) { // stacked experts (E, out, in) + m.note = std::to_string(cfg.num_experts_per_tok) + " of " + + std::to_string(e.shape.front()) + " experts active"; + } else if (contains(m.name, "mlp.gate") && !contains(m.name, "gate_proj")) { + m.note = "router"; + } else if (cfg.is_linear_layer(rep)) { + m.note = "linear attention"; + } + block.push_back(std::move(m)); + } + std::stable_sort(block.begin(), block.end(), [](const MatmulShape& a, const MatmulShape& b) { + return module_rank(a.name) < module_rank(b.name); + }); + s.decode_matmuls.insert(s.decode_matmuls.end(), block.begin(), block.end()); + } +} + +// Majority vote over the labels of the big (2-D+) layer weights, so norms and +// odd buffers never skew the summary; `mixed` reports a heterogeneous file. +std::string majority_label(const std::vector& labels, bool& mixed) { + mixed = false; + if (labels.empty()) return ""; + std::map counts; + for (const auto& l : labels) ++counts[l]; + std::string best; + int best_n = -1; + for (const auto& [label, n] : counts) + if (n > best_n) { best = label; best_n = n; } + mixed = counts.size() > 1; + return best; +} + +} // namespace + +ModelSchema build_schema_from_safetensors(const std::string& model_dir, const ModelConfig& cfg, + const std::string& model_name) { + ModelSchema s; + s.cfg = cfg; + s.format = "safetensors (MLX)"; + s.model_name = model_name.empty() ? path_basename(model_dir) : model_name; + + // Canonicalize keys the same way the loader does, keeping the vision tower + // unconditionally — the schematic should show it even when the engine's + // text-only load would drop it. Dropped buffers still count toward bytes. + std::map canon; + for (auto& e : read_safetensors_dir(model_dir)) { + auto key = sanitize_key(e.name, /*keep_vision=*/true); + if (!key) { + s.dropped_bytes += e.nbytes; + continue; + } + e.name = *key; + canon.emplace(*key, std::move(e)); + } + + // A ".scales" sibling marks ".weight" as MLX-quantized; the + // triplet folds into one logical tensor with the packed shape unpacked. + std::set quant_bases; + for (const auto& [name, _] : canon) { + if (!ends_with(name, ".scales")) continue; + const std::string base = name.substr(0, name.size() - std::strlen(".scales")); + if (canon.count(base + ".weight")) quant_bases.insert(base); + } + + // A ".scales"/".biases" of a quantized base folds into the base's + // ".weight" row below, so it never becomes a row of its own. + auto is_folded_quant_sibling = [&](const std::string& name) { + for (const char* suf : {".scales", ".biases"}) { + if (ends_with(name, suf) && + quant_bases.count(name.substr(0, name.size() - std::strlen(suf)))) { + return true; + } + } + return false; + }; + + std::vector quant_labels, dense_labels; + for (const auto& [name, e] : canon) { + if (is_folded_quant_sibling(name)) continue; + + TensorEntry t; + t.name = name; + t.layer = layer_of(name); + t.component = component_of(name); + t.dtype = e.dtype; + t.shape = e.shape; + t.bytes = e.nbytes; + + const bool is_weight = ends_with(name, ".weight"); + const std::string base = + is_weight ? name.substr(0, name.size() - std::strlen(".weight")) : ""; + if (is_weight && quant_bases.count(base)) { + const QuantParams qp = cfg.quant_for(base); + t.stored_shape = e.shape; + // Packed uint32 columns: in * bits / 32 -> logical in = cols * 32 / bits. + t.shape.back() = e.shape.back() * 32 / qp.bits; + const auto& scales = canon.at(base + ".scales"); + const int64_t scales_in = scales.shape.empty() ? 0 : scales.shape.back(); + if (scales_in * qp.group_size != t.shape.back()) { + log::warn( + "inspect: quant shape mismatch for '{}': packed-derived in={} vs " + "scales-derived in={} (bits={} gs={}) — check quant config", + base, t.shape.back(), scales_in * qp.group_size, qp.bits, qp.group_size); + } + t.quant = std::to_string(qp.bits) + "b gs" + std::to_string(qp.group_size); + t.bytes += scales.nbytes; + if (auto it = canon.find(base + ".biases"); it != canon.end()) + t.bytes += it->second.nbytes; + if (t.shape.size() >= 2) + quant_labels.push_back(std::to_string(qp.bits) + "-bit gs" + + std::to_string(qp.group_size)); + } else if (t.shape.size() >= 2 && t.layer >= 0) { + std::string d = e.dtype; + std::transform(d.begin(), d.end(), d.begin(), + [](unsigned char c) { return std::tolower(c); }); + dense_labels.push_back(d == "f16" ? "fp16" : d == "bf16" ? "bf16" : d); + } + t.params = prod(t.shape); + s.tensors.push_back(std::move(t)); + } + + s.tensors = aggregate_per_expert(std::move(s.tensors)); + bool mixed = false; + s.quant_summary = !quant_labels.empty() ? majority_label(quant_labels, mixed) + " MLX" + : majority_label(dense_labels, mixed); + if (s.quant_summary.empty()) s.quant_summary = "fp16"; + if (mixed) s.quant_summary += " (mixed)"; + finalize(s); + return s; +} + +ModelSchema build_schema_from_gguf(const std::string& gguf_path, const std::string& model_name) { + GgufInspection insp = inspect_gguf(gguf_path); + + ModelSchema s; + s.cfg = insp.head.config; + s.format = "GGUF"; + s.model_name = model_name.empty() ? path_basename(gguf_path) : model_name; + + std::vector quant_labels; + for (const auto& tm : insp.tensors) { + TensorEntry t; + t.name = tm.canonical.empty() ? tm.name : tm.canonical; + t.layer = layer_of(t.name); + t.component = component_of(t.name); + t.shape = tm.shape; // GGUF dims are already logical + t.dtype = ggml_type_name(tm.ggml_type); + t.bytes = tm.bytes; + const bool dense = tm.ggml_type == 0 /*F32*/ || tm.ggml_type == 1 /*F16*/ || + tm.ggml_type == 30 /*BF16*/; + if (!dense) t.quant = t.dtype; + // rope_freqs is a baked frequency table, not model parameters. + t.params = (tm.name == "rope_freqs.weight") ? 0 : prod(t.shape); + if (t.shape.size() >= 2 && t.layer >= 0) quant_labels.push_back(t.dtype); + s.tensors.push_back(std::move(t)); + } + + bool mixed = false; + s.quant_summary = majority_label(quant_labels, mixed); + if (s.quant_summary.empty()) s.quant_summary = "F16"; + s.quant_summary += " GGUF"; + if (mixed) s.quant_summary += " (mixed)"; + s.cfg.tie_word_embeddings = std::none_of( + s.tensors.begin(), s.tensors.end(), + [](const TensorEntry& e) { return starts_with(e.name, "lm_head."); }); + finalize(s); + return s; +} + +nlohmann::json ModelSchema::to_json() const { + nlohmann::json j; + j["header"] = { + {"name", model_name}, + {"family", family}, + {"format", format}, + {"quant", quant_summary}, + {"model_type", cfg.model_type}, + {"params", total_params}, + {"bytes", total_bytes}, + {"dropped_bytes", dropped_bytes}, + {"context_length", cfg.max_position_embeddings}, + {"vocab", cfg.vocab}, + {"tied_embeddings", tied_embeddings}, + }; + + nlohmann::json arch = { + {"hidden", cfg.hidden}, + {"n_layers", cfg.n_layers}, + {"n_heads", cfg.n_heads}, + {"n_kv_heads", cfg.n_kv_heads}, + {"gqa_ratio", gqa_ratio}, + {"head_dim", head_dim}, + {"intermediate_size", cfg.intermediate_size}, + }; + bool has_qk_norm = false; + for (const auto& e : tensors) + if (e.name.find(".self_attn.q_norm.") != std::string::npos) { has_qk_norm = true; break; } + arch["qk_norm"] = has_qk_norm; + if (cfg.num_experts > 0) { + int n_moe = 0; + for (int i = 0; i < cfg.n_layers; ++i) + if (cfg.is_moe_layer(i)) ++n_moe; + arch["moe"] = {{"experts", cfg.num_experts}, + {"top_k", cfg.num_experts_per_tok}, + {"moe_intermediate", cfg.moe_intermediate_size}, + {"n_moe_layers", n_moe}}; + } else { + arch["moe"] = nullptr; + } + if (cfg.full_attention_interval > 0) { + arch["hybrid"] = {{"full_attention_interval", cfg.full_attention_interval}, + {"n_linear_layers", cfg.n_layers - n_full_attn_layers}, + {"linear_num_key_heads", cfg.linear_num_key_heads}, + {"linear_num_value_heads", cfg.linear_num_value_heads}, + {"linear_key_head_dim", cfg.linear_key_head_dim}, + {"linear_value_head_dim", cfg.linear_value_head_dim}, + {"conv_kernel", cfg.linear_conv_kernel_dim}}; + } else { + arch["hybrid"] = nullptr; + } + { + nlohmann::json rope = {{"theta", cfg.rope_theta}}; + if (cfg.rope_scaling.has_value()) { + rope["type"] = cfg.rope_scaling->rope_type; + rope["factor"] = cfg.rope_scaling->factor; + } else if (cfg.rope_freq_factors.has_value()) { + rope["type"] = "llama3 (baked)"; + } else { + rope["type"] = "none"; + } + arch["rope"] = std::move(rope); + } + if (cfg.vision.has_value()) { + const VisionConfig& v = *cfg.vision; + arch["vision"] = {{"depth", v.depth}, + {"hidden", v.hidden}, + {"intermediate_size", v.intermediate_size}, + {"num_heads", v.num_heads}, + {"patch_size", v.patch_size}, + {"spatial_merge_size", v.spatial_merge_size}, + {"out_hidden_size", v.out_hidden_size}, + {"deepstack_indexes", v.deepstack_visual_indexes}}; + } else { + arch["vision"] = nullptr; + } + j["arch"] = std::move(arch); + + // Quantized-KV variants mirror cache/kv_quant's layout: packed values plus + // fp16 scale+bias per group of 64 — approximate (group padding ignored). + auto kv_quant_bytes = [&](int bits) { + return static_cast(n_full_attn_layers) * 2.0 * cfg.n_kv_heads * head_dim * + (bits / 8.0 + 4.0 / 64.0); + }; + nlohmann::json matmuls = nlohmann::json::array(); + for (const auto& m : decode_matmuls) { + matmuls.push_back({{"name", m.name}, + {"in", m.in}, + {"out", m.out}, + {"quant", m.quant}, + {"note", m.note}}); + } + j["derived"] = {{"kv_bytes_per_token", kv_bytes_per_token}, + {"kv_bytes_per_token_kv8", kv_quant_bytes(8)}, + {"kv_bytes_per_token_kv4", kv_quant_bytes(4)}, + {"n_full_attn_layers", n_full_attn_layers}, + {"decode_matmuls", std::move(matmuls)}}; + + nlohmann::json components = nlohmann::json::object(); + for (const auto& [name, agg] : by_component) + components[name] = {{"params", agg.params}, {"bytes", agg.bytes}}; + j["components"] = std::move(components); + + nlohmann::json layers = nlohmann::json::array(); + for (size_t i = 0; i < by_layer.size(); ++i) { + const int idx = static_cast(i); + const char* kind = cfg.is_linear_layer(idx) ? "linear" + : cfg.is_moe_layer(idx) ? "moe" + : "attn"; + layers.push_back({{"idx", idx}, + {"kind", kind}, + {"params", by_layer[i].params}, + {"bytes", by_layer[i].bytes}}); + } + j["layers"] = std::move(layers); + + nlohmann::json rows = nlohmann::json::array(); + for (const auto& e : tensors) { + nlohmann::json row = {{"name", e.name}, {"layer", e.layer}, {"component", e.component}, + {"shape", e.shape}, {"dtype", e.dtype}, {"quant", e.quant}, + {"params", e.params}, {"bytes", e.bytes}}; + if (!e.stored_shape.empty() && e.stored_shape != e.shape) + row["stored_shape"] = e.stored_shape; + rows.push_back(std::move(row)); + } + j["tensors"] = std::move(rows); + return j; +} + +} // namespace mlxforge::inspect diff --git a/src/inspect/model_schema.h b/src/inspect/model_schema.h new file mode 100644 index 0000000..3461dd1 --- /dev/null +++ b/src/inspect/model_schema.h @@ -0,0 +1,92 @@ +// Architecture schema for the CLI `schematic` command. +// +// Builds a structured, JSON-serializable description of a model — every tensor +// (logical shape, dtype, quantization, bytes), per-component and per-layer +// parameter aggregates, and the derived numbers an inference engineer cares +// about (GQA ratio, KV-cache bytes/token, decode matmul shapes) — from +// metadata only: safetensors headers (inspect/safetensors_header) or the GGUF +// tensor directory (core/gguf's inspect_gguf). No MLX arrays are created. +#pragma once + +#include +#include +#include +#include + +#include + +#include "core/config.h" + +namespace mlxforge::inspect { + +struct TensorEntry { + std::string name; // canonical key (post sanitize/remap) + int layer = -1; // decoder layer index; -1 = global/vision + std::string component; // embed|attn|linear_attn|mlp|moe|norm|lm_head|vision|other + std::vector shape; // LOGICAL dims (quantized weights unpacked) + std::vector stored_shape; // on-disk dims (uint32-packed for MLX quant) + std::string dtype; // "F16", "BF16", "U32", "Q4_K", ... + std::string quant; // "" (dense) | "4b gs64" | "Q4_0" | "type_" + uint64_t params = 0; // logical element count (0 for non-param buffers) + uint64_t bytes = 0; // on-disk bytes (incl. folded scales/biases) +}; + +struct ComponentAgg { + uint64_t params = 0; + uint64_t bytes = 0; +}; + +// One decode-step matmul (M=1): activation [1, in] x weight [out, in]. +struct MatmulShape { + std::string name; // module path within the block ("self_attn.q_proj", ...) + int64_t in = 0; + int64_t out = 0; + std::string quant; // quant string of the weight ("" = dense) + std::string note; // e.g. "8 of 128 experts active" +}; + +struct ModelSchema { + // Header card. + std::string model_name; + std::string format; // "safetensors (MLX)" | "GGUF" + std::string family; // llama|qwen3|qwen3-moe|qwen3.5-hybrid|qwen3-vl + std::string quant_summary; // "4-bit gs64 MLX" | "Q4_K GGUF" | "fp16" | ... + uint64_t total_params = 0; + uint64_t total_bytes = 0; // on-disk weight bytes (incl. dropped buffers) + uint64_t dropped_bytes = 0; // bytes of buffers the engine never loads + bool tied_embeddings = false; + + ModelConfig cfg; // architecture hyperparameters (to_json picks fields) + + // Derived numbers. + int head_dim = 0; // resolved (cfg.head_dim, else hidden/n_heads) + int gqa_ratio = 1; // n_heads / n_kv_heads + int n_full_attn_layers = 0; // == n_layers unless hybrid + double kv_bytes_per_token = 0; // fp16 cache, full-attention layers only + std::vector decode_matmuls; // representative block(s), tensor-derived + + std::vector tensors; // sorted by (layer, name) + std::map by_component; + std::vector by_layer; // size n_layers + + // The JSON blob embedded into the HTML page (shape documented in + // schematic_html.cpp's template). + nlohmann::json to_json() const; +}; + +// Classify a canonical tensor key into its component bucket; exposed for the +// table-driven unit tests. +std::string component_of(const std::string& canonical_key); +// Decoder layer index of a canonical key ("model.layers.N." prefix), -1 if global. +int layer_of(const std::string& canonical_key); + +// Build the schema from a safetensors model directory (header parse only). +// `model_name` defaults to the directory basename when empty. +ModelSchema build_schema_from_safetensors(const std::string& model_dir, const ModelConfig& cfg, + const std::string& model_name = ""); + +// Build the schema from a GGUF file (metadata + tensor directory only). +ModelSchema build_schema_from_gguf(const std::string& gguf_path, + const std::string& model_name = ""); + +} // namespace mlxforge::inspect diff --git a/src/inspect/safetensors_header.cpp b/src/inspect/safetensors_header.cpp new file mode 100644 index 0000000..b495121 --- /dev/null +++ b/src/inspect/safetensors_header.cpp @@ -0,0 +1,98 @@ +#include "inspect/safetensors_header.h" + +#include +#include +#include +#include + +#include + +#include "core/logging.h" +#include "core/weights.h" + +namespace mlxforge::inspect { + +namespace { +// Headers are typically a few hundred KiB even for 100B+ checkpoints; anything +// past this is a corrupt length field, not a real header. +constexpr uint64_t kMaxHeaderBytes = 256ull * 1024 * 1024; +} // namespace + +std::vector read_safetensors_header(const std::string& file) { + std::ifstream f(file, std::ios::binary); + if (!f) throw std::runtime_error("inspect: cannot open '" + file + "'"); + + uint64_t header_len = 0; + f.read(reinterpret_cast(&header_len), sizeof(header_len)); + if (!f || header_len == 0 || header_len > kMaxHeaderBytes) { + throw std::runtime_error("inspect: corrupt safetensors header length in '" + file + "'"); + } + + std::string header(header_len, '\0'); + f.read(&header[0], static_cast(header_len)); + if (!f) throw std::runtime_error("inspect: truncated safetensors header in '" + file + "'"); + + nlohmann::json j; + try { + j = nlohmann::json::parse(header); + } catch (const nlohmann::json::exception& e) { + throw std::runtime_error("inspect: bad safetensors header JSON in '" + file + "': " + e.what()); + } + + std::vector out; + out.reserve(j.size()); + for (const auto& [key, v] : j.items()) { + if (key == "__metadata__") continue; + SafetensorsEntry e; + e.name = key; + e.dtype = v.value("dtype", ""); + if (v.contains("shape")) e.shape = v["shape"].get>(); + if (v.contains("data_offsets") && v["data_offsets"].is_array() && + v["data_offsets"].size() == 2) { + const uint64_t begin = v["data_offsets"][0].get(); + const uint64_t end = v["data_offsets"][1].get(); + if (end < begin) { + throw std::runtime_error("inspect: invalid data_offsets for '" + key + "' in '" + file + + "'"); + } + e.nbytes = end - begin; + } + out.push_back(std::move(e)); + } + return out; +} + +std::vector read_safetensors_dir(const std::string& model_dir) { + // Mirror load_weights' shard discovery (core/weights.cpp): prefer the sharded + // layout, but only when every shard the index names exists — some exports + // ship a consolidated model.safetensors alongside a stale index.json. + const std::string index_path = model_dir + "/model.safetensors.index.json"; + std::ifstream index_file(index_path); + if (index_file) { + nlohmann::json index_json; + index_file >> index_json; + const auto weight_map = parse_shard_index(index_json); + const auto files = shard_files(weight_map); + const bool all_present = std::all_of(files.begin(), files.end(), [&](const std::string& f) { + return std::ifstream(model_dir + "/" + f).good(); + }); + if (all_present) { + log::debug("inspect: sharded checkpoint, {} files", files.size()); + std::vector out; + for (const auto& file : files) { + auto shard = read_safetensors_header(model_dir + "/" + file); + out.insert(out.end(), std::make_move_iterator(shard.begin()), + std::make_move_iterator(shard.end())); + } + return out; + } + log::debug("inspect: index.json shards absent; falling back to single file"); + } + const std::string single = model_dir + "/model.safetensors"; + if (!std::ifstream(single)) { + throw std::runtime_error("inspect: no model.safetensors[.index.json] in '" + model_dir + "'"); + } + return read_safetensors_header(single); +} + +} // namespace mlxforge::inspect diff --git a/src/inspect/safetensors_header.h b/src/inspect/safetensors_header.h new file mode 100644 index 0000000..87573cb --- /dev/null +++ b/src/inspect/safetensors_header.h @@ -0,0 +1,34 @@ +// Metadata-only safetensors reader for model inspection. +// +// A .safetensors file starts with an 8-byte little-endian u64 header length +// followed by a JSON header mapping tensor name -> {dtype, shape, data_offsets}. +// Reading just that header recovers every tensor's name/shape/dtype/byte-size +// without touching the weight data — no MLX arrays, no GPU, fast even on the +// multi-GiB shards of a 70B model. Used by the CLI `schematic` command; the +// engine's real weight loading stays in core/weights. +#pragma once + +#include +#include +#include + +namespace mlxforge::inspect { + +struct SafetensorsEntry { + std::string name; // raw key as stored in the file (not canonicalized) + std::string dtype; // header string verbatim ("F16", "BF16", "F32", "U32", ...) + std::vector shape; // on-disk dims (packed for MLX-quantized weights) + uint64_t nbytes = 0; // data_offsets[1] - data_offsets[0] +}; + +// Parse one .safetensors file's JSON header. Throws on a missing/unreadable +// file, a corrupt header, or an implausibly large header length. +std::vector read_safetensors_header(const std::string& file); + +// Read the headers of every shard in a model directory, mirroring +// load_weights' shard discovery: prefer model.safetensors.index.json when all +// the shards it names are present (stale-index exports exist), else fall back +// to the single model.safetensors. Throws if neither layout is found. +std::vector read_safetensors_dir(const std::string& model_dir); + +} // namespace mlxforge::inspect diff --git a/src/inspect/schematic_html.cpp b/src/inspect/schematic_html.cpp new file mode 100644 index 0000000..046d1d1 --- /dev/null +++ b/src/inspect/schematic_html.cpp @@ -0,0 +1,704 @@ +#include "inspect/schematic_html.h" + +#include +#include + +namespace mlxforge::inspect { + +namespace { + +// Splice point for the schema JSON inside the data ...") can never terminate the data element. JSON.parse restores +// the original characters; the JS side renders all data via textContent. +std::string escape_json_for_html(const std::string& json) { + std::string out; + out.reserve(json.size()); + for (char c : json) { + if (c == '<') { + out += "\\u003c"; + } else { + out += c; + } + } + return out; +} + +// The page template: a pastel "drafting paper" single-file infographic. +// Inline CSS/JS only — no external fonts, scripts or styles — so the file +// renders offline. Adjacent raw literals concatenate. +const char* kTemplate = + R"HTML( + + + + +model schematic · mlxforge + + + +
+
mlxforge · model schematic
+

+
+
+
+
+
+

architecture

+
+
+
+

forward pass

+
+
+
+

parameter distribution

+
+
+

per-layer parameters

+
+
+
+
+

tensor explorer

+
+ + +
+
+
+
+
generated by mlxforge-cli schematic — metadata-only read of the checkpoint +(safetensors headers / GGUF tensor directory); no weights were loaded.
+ +)HTML" + R"HTML( +)HTML" + R"HTML( + + +)HTML"; + +} // namespace + +std::string render_schematic_html(const nlohmann::json& schema) { + std::string page = kTemplate; + const size_t pos = page.find(kMarker); + if (pos == std::string::npos) { + throw std::logic_error("schematic_html: template is missing the schema marker"); + } + page.replace(pos, std::strlen(kMarker), escape_json_for_html(schema.dump())); + return page; +} + +} // namespace mlxforge::inspect diff --git a/src/inspect/schematic_html.h b/src/inspect/schematic_html.h new file mode 100644 index 0000000..e603e2c --- /dev/null +++ b/src/inspect/schematic_html.h @@ -0,0 +1,18 @@ +// Self-contained HTML renderer for the CLI `schematic` command. +// +// Takes a ModelSchema JSON blob (model_schema.h's to_json) and splices it into +// an embedded HTML template — inline CSS/JS/SVG only, no external references — +// so the output is a single file that renders offline. The page draws the +// transformer-block schematic, stat cards, parameter-distribution bars and a +// collapsible per-layer tensor explorer with vanilla JS. +#pragma once + +#include + +#include + +namespace mlxforge::inspect { + +std::string render_schematic_html(const nlohmann::json& schema); + +} // namespace mlxforge::inspect diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 47843e2..df83d68 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -49,6 +49,9 @@ add_executable(mlxforge_tests tokenizer/bpe_test.cpp tokenizer/spm_test.cpp model/quantized_test.cpp + inspect/safetensors_header_test.cpp + inspect/model_schema_test.cpp + inspect/schematic_html_test.cpp # The C ABI is compiled directly into the test binary (it is otherwise only in # the shared library), so capi_test can drive the engine through mlxforge.h. ${CMAKE_SOURCE_DIR}/src/capi/mlxforge.cpp diff --git a/tests/core/gguf_test.cpp b/tests/core/gguf_test.cpp index b3fcca4..d5e8b09 100644 --- a/tests/core/gguf_test.cpp +++ b/tests/core/gguf_test.cpp @@ -228,6 +228,90 @@ TEST_CASE("GGUF loader rejects Qwen3.5/qwen3next as unsupported") { std::remove(path.c_str()); } +TEST_CASE("GGUF inspect_gguf exposes the tensor directory without loading weights") { + // Hand-craft a tiny llama GGUF with one F16 tensor and one tensor of an + // unknown ggml type. inspect_gguf must report both (the unknown type is shown, + // not rejected — only loading needs to dequantize) with exact offset-derived + // byte counts. + constexpr uint32_t kF16 = 1; + constexpr uint32_t kUnknownType = 99; + constexpr uint64_t kAlign = 32; + + std::string kv; + uint64_t nkv = 0; + auto kv_str = [&](const std::string& k, const std::string& v) { + put_str(kv, k); put(kv, 8); put_str(kv, v); ++nkv; + }; + auto kv_u32 = [&](const std::string& k, uint32_t v) { + put_str(kv, k); put(kv, 4); put(kv, v); ++nkv; + }; + auto kv_f32 = [&](const std::string& k, float v) { + put_str(kv, k); put(kv, 6); put(kv, v); ++nkv; + }; + kv_str("general.architecture", "llama"); + kv_u32("llama.block_count", 1); + kv_u32("llama.embedding_length", 8); + kv_u32("llama.attention.head_count", 2); + kv_f32("llama.attention.layer_norm_rms_epsilon", 1e-5f); + + // token_embd: F16 ne {8, 4} -> MLX shape (4, 8), 64 data bytes (aligned). + // blk.0.attn_q: unknown type, 50 data bytes padded to 64. + std::string tinfo; + put_str(tinfo, "token_embd.weight"); + put(tinfo, 2); + put(tinfo, 8); + put(tinfo, 4); + put(tinfo, kF16); + put(tinfo, 0); + put_str(tinfo, "blk.0.attn_q.weight"); + put(tinfo, 2); + put(tinfo, 8); + put(tinfo, 8); + put(tinfo, kUnknownType); + put(tinfo, 64); + + std::string b; + put(b, 0x46554747); // "GGUF" + put(b, 3); // version + put(b, 2); // tensor count + put(b, nkv); + b += kv; + b += tinfo; + b.append((kAlign - (b.size() % kAlign)) % kAlign, '\0'); // align to data_start + b.append(64, '\0'); // token_embd data + b.append(64, '\0'); // attn_q data (50 used, padded) + + const std::string path = + (std::filesystem::temp_directory_path() / "mlxforge_inspect.gguf").string(); + { std::ofstream(path, std::ios::binary).write(b.data(), b.size()); } + + mlxforge::GgufInspection insp = mlxforge::inspect_gguf(path); + std::remove(path.c_str()); + + CHECK(insp.head.config.model_type == "llama"); + CHECK(insp.head.config.n_layers == 1); + CHECK(insp.file_bytes == b.size()); + REQUIRE(insp.tensors.size() == 2); + + const auto& embd = insp.tensors[0]; + CHECK(embd.name == "token_embd.weight"); + CHECK(embd.canonical == "model.embed_tokens.weight"); + CHECK(embd.shape == std::vector{4, 8}); // ggml ne reversed + CHECK(embd.ggml_type == kF16); + CHECK(embd.bytes == 64); + + const auto& q = insp.tensors[1]; + CHECK(q.canonical == "model.layers.0.self_attn.q_proj.weight"); + CHECK(q.ggml_type == kUnknownType); + CHECK(q.bytes == 64); // runs to end-of-file + + CHECK(mlxforge::ggml_type_name(12) == "Q4_K"); + CHECK(mlxforge::ggml_type_name(1) == "F16"); + CHECK(mlxforge::ggml_type_name(kUnknownType) == "type_99"); + CHECK(mlxforge::ggml_bits_per_weight(2) == doctest::Approx(4.5)); + CHECK(mlxforge::ggml_bits_per_weight(kUnknownType) == 0.0); +} + TEST_CASE("GGUF forward pass produces the golden first token") { if (!gguf_available()) { MESSAGE("MLXFORGE_GGUF_MODEL not present; skipping"); diff --git a/tests/inspect/model_schema_test.cpp b/tests/inspect/model_schema_test.cpp new file mode 100644 index 0000000..a928113 --- /dev/null +++ b/tests/inspect/model_schema_test.cpp @@ -0,0 +1,363 @@ +// ModelSchema builders: quant-triplet folding + packed-shape math, component +// classification, derived KV/GQA math, MoE aggregation, JSON totals. Pure +// units run over synthetic files; the model-gated cases self-skip when the +// cached checkpoints are absent. +#include + +#include +#include +#include +#include + +#include + +#include "core/config.h" +#include "inspect/model_schema.h" + +namespace { + +namespace fs = std::filesystem; +using mlxforge::ModelConfig; +using mlxforge::inspect::ModelSchema; +using mlxforge::inspect::TensorEntry; + +void write_safetensors(const std::string& path, const nlohmann::json& header) { + const std::string h = header.dump(); + const uint64_t len = h.size(); + std::ofstream f(path, std::ios::binary); + f.write(reinterpret_cast(&len), sizeof(len)); + f.write(h.data(), static_cast(h.size())); +} + +struct TempDir { + fs::path dir; + explicit TempDir(const std::string& name) : dir(fs::temp_directory_path() / name) { + fs::remove_all(dir); + fs::create_directories(dir); + } + ~TempDir() { fs::remove_all(dir); } + std::string str() const { return dir.string(); } +}; + +const TensorEntry& find_tensor(const ModelSchema& s, const std::string& name) { + for (const auto& e : s.tensors) + if (e.name == name) return e; + REQUIRE_MESSAGE(false, "tensor not found: " << name); + static TensorEntry dummy; + return dummy; +} + +// Append helpers for hand-crafted GGUF files (mirrors tests/core/gguf_test.cpp). +template +void put(std::string& b, T v) { + b.append(reinterpret_cast(&v), sizeof(T)); +} +void put_str(std::string& b, const std::string& s) { + put(b, s.size()); + b += s; +} + +} // namespace + +TEST_CASE("component classification buckets canonical keys") { + using mlxforge::inspect::component_of; + using mlxforge::inspect::layer_of; + struct Case { + const char* key; + const char* component; + int layer; + }; + const Case cases[] = { + {"model.embed_tokens.weight", "embed", -1}, + {"lm_head.weight", "lm_head", -1}, + {"model.norm.weight", "norm", -1}, + {"model.layers.0.self_attn.q_proj.weight", "attn", 0}, + {"model.layers.3.self_attn.q_norm.weight", "attn", 3}, + {"model.layers.12.input_layernorm.weight", "norm", 12}, + {"model.layers.12.post_attention_layernorm.weight", "norm", 12}, + {"model.layers.5.mlp.gate_proj.weight", "mlp", 5}, + {"model.layers.5.mlp.down_proj.weight", "mlp", 5}, + {"model.layers.7.mlp.gate.weight", "moe", 7}, // router, not gate_proj + {"model.layers.7.mlp.switch_mlp.up_proj.weight", "moe", 7}, + {"model.layers.7.mlp.experts.31.down_proj.weight", "moe", 7}, + {"model.layers.2.linear_attn.in_proj_qkvz.weight", "linear_attn", 2}, + {"visual.blocks.4.attn.q_proj.weight", "vision", -1}, + {"visual.patch_embed.weight", "vision", -1}, + {"rope_freqs.weight", "other", -1}, + }; + for (const auto& c : cases) { + CAPTURE(c.key); + CHECK(component_of(c.key) == c.component); + CHECK(layer_of(c.key) == c.layer); + } +} + +TEST_CASE("safetensors schema folds quant triplets and unpacks packed shapes") { + TempDir td("mlxforge_schema_quant"); + // q_proj: 4-bit gs64, logical [256, 512] -> packed weight [256, 64] U32, + // scales/biases [256, 8] F16. o_proj: an 8-bit gs32 override, logical [4, 32] + // -> packed [4, 8] U32, scales [4, 1]. + const nlohmann::json header = { + {"model.embed_tokens.weight", + {{"dtype", "F16"}, {"shape", {10, 512}}, {"data_offsets", {0, 10240}}}}, + {"model.layers.0.self_attn.q_proj.weight", + {{"dtype", "U32"}, {"shape", {256, 64}}, {"data_offsets", {10240, 75776}}}}, + {"model.layers.0.self_attn.q_proj.scales", + {{"dtype", "F16"}, {"shape", {256, 8}}, {"data_offsets", {75776, 79872}}}}, + {"model.layers.0.self_attn.q_proj.biases", + {{"dtype", "F16"}, {"shape", {256, 8}}, {"data_offsets", {79872, 83968}}}}, + {"model.layers.0.self_attn.o_proj.weight", + {{"dtype", "U32"}, {"shape", {4, 8}}, {"data_offsets", {83968, 84096}}}}, + {"model.layers.0.self_attn.o_proj.scales", + {{"dtype", "F16"}, {"shape", {4, 1}}, {"data_offsets", {84096, 84104}}}}, + {"model.norm.weight", {{"dtype", "F16"}, {"shape", {512}}, {"data_offsets", {84104, 85128}}}}, + // A buffer sanitize_key drops; its bytes must still count toward disk size. + {"model.layers.0.self_attn.rotary_emb.inv_freq", + {{"dtype", "F32"}, {"shape", {32}}, {"data_offsets", {85128, 85256}}}}, + }; + write_safetensors(td.str() + "/model.safetensors", header); + + ModelConfig cfg; + cfg.n_layers = 1; + cfg.hidden = 512; + cfg.n_heads = 8; + cfg.n_kv_heads = 4; + cfg.vocab = 10; + cfg.quant_group_size = 64; + cfg.quant_bits = 4; + cfg.quant_overrides["model.layers.0.self_attn.o_proj"] = {32, 8}; + + const ModelSchema s = mlxforge::inspect::build_schema_from_safetensors(td.str(), cfg, "tiny"); + + // scales/biases never appear as rows. + CHECK(s.tensors.size() == 4); + + const auto& q = find_tensor(s, "model.layers.0.self_attn.q_proj.weight"); + CHECK(q.shape == std::vector{256, 512}); // 64 packed cols * 32 / 4 bits + CHECK(q.stored_shape == std::vector{256, 64}); + CHECK(q.params == 256 * 512); + CHECK(q.quant == "4b gs64"); + CHECK(q.bytes == 65536 + 4096 + 4096); // weight + scales + biases + CHECK(q.component == "attn"); + + const auto& o = find_tensor(s, "model.layers.0.self_attn.o_proj.weight"); + CHECK(o.shape == std::vector{4, 32}); // 8 packed cols * 32 / 8 bits + CHECK(o.quant == "8b gs32"); // override honored + + // No lm_head tensor -> tied embeddings; dropped buffer bytes still counted. + CHECK(s.tied_embeddings); + CHECK(s.dropped_bytes == 128); + CHECK(s.total_bytes == 85256); + CHECK(s.family == "llama"); + CHECK(s.quant_summary == "4-bit gs64 MLX (mixed)"); // the o_proj 8-bit override + + // Decode matmuls are tensor-derived ([out, in] -> in/out). + REQUIRE(!s.decode_matmuls.empty()); + CHECK(s.decode_matmuls[0].name == "self_attn.q_proj"); + CHECK(s.decode_matmuls[0].in == 512); + CHECK(s.decode_matmuls[0].out == 256); +} + +TEST_CASE("safetensors schema aggregates raw per-expert MoE tensors") { + TempDir td("mlxforge_schema_moe"); + const nlohmann::json header = { + {"model.layers.0.mlp.experts.0.gate_proj.weight", + {{"dtype", "F16"}, {"shape", {3, 4}}, {"data_offsets", {0, 24}}}}, + {"model.layers.0.mlp.experts.1.gate_proj.weight", + {{"dtype", "F16"}, {"shape", {3, 4}}, {"data_offsets", {24, 48}}}}, + {"model.layers.0.mlp.gate.weight", + {{"dtype", "F16"}, {"shape", {2, 4}}, {"data_offsets", {48, 64}}}}, + }; + write_safetensors(td.str() + "/model.safetensors", header); + + ModelConfig cfg; + cfg.n_layers = 1; + cfg.hidden = 4; + cfg.n_heads = 1; + cfg.n_kv_heads = 1; + cfg.num_experts = 2; + cfg.num_experts_per_tok = 2; + + const ModelSchema s = mlxforge::inspect::build_schema_from_safetensors(td.str(), cfg); + + const auto& experts = find_tensor(s, "model.layers.0.mlp.experts.*.gate_proj.weight"); + CHECK(experts.shape == std::vector{2, 3, 4}); // expert count prepended + CHECK(experts.params == 24); + CHECK(experts.bytes == 48); + CHECK(experts.component == "moe"); + CHECK(s.family == "qwen3-moe"); +} + +TEST_CASE("derived KV math is hybrid-aware") { + TempDir td("mlxforge_schema_kv"); + write_safetensors(td.str() + "/model.safetensors", + {{"model.embed_tokens.weight", + {{"dtype", "F16"}, {"shape", {4, 8}}, {"data_offsets", {0, 64}}}}}); + + ModelConfig dense; + dense.n_layers = 16; + dense.hidden = 2048; + dense.n_heads = 32; + dense.n_kv_heads = 8; + dense.head_dim = 64; + const ModelSchema sd = mlxforge::inspect::build_schema_from_safetensors(td.str(), dense); + CHECK(sd.n_full_attn_layers == 16); + CHECK(sd.gqa_ratio == 4); + CHECK(sd.kv_bytes_per_token == doctest::Approx(16 * 2 * 8 * 64 * 2)); + + ModelConfig hybrid = dense; + hybrid.n_layers = 48; + hybrid.full_attention_interval = 4; // every 4th layer full -> 12 of 48 + const ModelSchema sh = mlxforge::inspect::build_schema_from_safetensors(td.str(), hybrid); + CHECK(sh.n_full_attn_layers == 12); + CHECK(sh.kv_bytes_per_token == doctest::Approx(12 * 2 * 8 * 64 * 2)); + CHECK(sh.family == "qwen3.5-hybrid"); +} + +TEST_CASE("schema JSON totals equal the sum of tensor rows") { + TempDir td("mlxforge_schema_json"); + const nlohmann::json header = { + {"model.embed_tokens.weight", + {{"dtype", "F16"}, {"shape", {10, 8}}, {"data_offsets", {0, 160}}}}, + {"model.layers.0.self_attn.q_proj.weight", + {{"dtype", "F16"}, {"shape", {8, 8}}, {"data_offsets", {160, 288}}}}, + }; + write_safetensors(td.str() + "/model.safetensors", header); + + ModelConfig cfg; + cfg.n_layers = 1; + cfg.hidden = 8; + cfg.n_heads = 2; + cfg.n_kv_heads = 2; + const ModelSchema s = mlxforge::inspect::build_schema_from_safetensors(td.str(), cfg); + const nlohmann::json j = s.to_json(); + + for (const char* key : {"header", "arch", "derived", "components", "layers", "tensors"}) { + CAPTURE(key); + CHECK(j.contains(key)); + } + + uint64_t params = 0, bytes = 0; + for (const auto& row : j["tensors"]) { + params += row["params"].get(); + bytes += row["bytes"].get(); + } + CHECK(j["header"]["params"].get() == params); + CHECK(j["header"]["bytes"].get() == bytes); // no dropped buffers here + CHECK(params == 10 * 8 + 8 * 8); +} + +TEST_CASE("GGUF schema reports logical shapes and ggml quant types") { + // Tiny llama GGUF: one F16 embed + one Q4_0-typed projection. + constexpr uint32_t kF16 = 1, kQ4_0 = 2; + constexpr uint64_t kAlign = 32; + + std::string kv; + uint64_t nkv = 0; + auto kv_str = [&](const std::string& k, const std::string& v) { + put_str(kv, k); put(kv, 8); put_str(kv, v); ++nkv; + }; + auto kv_u32 = [&](const std::string& k, uint32_t v) { + put_str(kv, k); put(kv, 4); put(kv, v); ++nkv; + }; + auto kv_f32 = [&](const std::string& k, float v) { + put_str(kv, k); put(kv, 6); put(kv, v); ++nkv; + }; + kv_str("general.architecture", "llama"); + kv_u32("llama.block_count", 1); + kv_u32("llama.embedding_length", 64); + kv_u32("llama.attention.head_count", 2); + kv_f32("llama.attention.layer_norm_rms_epsilon", 1e-5f); + + std::string tinfo; + put_str(tinfo, "token_embd.weight"); + put(tinfo, 2); + put(tinfo, 64); // ggml ne: innermost first + put(tinfo, 4); + put(tinfo, kF16); + put(tinfo, 0); + put_str(tinfo, "blk.0.attn_q.weight"); + put(tinfo, 2); + put(tinfo, 64); + put(tinfo, 64); + put(tinfo, kQ4_0); + put(tinfo, 512); // 4*64 fp16 = 512 bytes (aligned) + + std::string b; + put(b, 0x46554747); + put(b, 3); + put(b, 2); + put(b, nkv); + b += kv; + b += tinfo; + b.append((kAlign - (b.size() % kAlign)) % kAlign, '\0'); + b.append(512, '\0'); // token_embd: 4*64 fp16 + b.append(64 * 64 / 32 * 18, '\0'); // attn_q: Q4_0, 18 bytes per 32 weights + + const std::string path = + (fs::temp_directory_path() / "mlxforge_schema_gguf.gguf").string(); + { std::ofstream(path, std::ios::binary).write(b.data(), b.size()); } + + const ModelSchema s = mlxforge::inspect::build_schema_from_gguf(path, "tiny-gguf"); + std::remove(path.c_str()); + + CHECK(s.format == "GGUF"); + CHECK(s.model_name == "tiny-gguf"); + CHECK(s.family == "llama"); + CHECK(s.quant_summary == "Q4_0 GGUF"); + CHECK(s.tied_embeddings); // no output.weight + + const auto& embd = find_tensor(s, "model.embed_tokens.weight"); + CHECK(embd.shape == std::vector{4, 64}); // logical, ne reversed + CHECK(embd.dtype == "F16"); + CHECK(embd.quant.empty()); + CHECK(embd.params == 256); + + const auto& q = find_tensor(s, "model.layers.0.self_attn.q_proj.weight"); + CHECK(q.shape == std::vector{64, 64}); + CHECK(q.dtype == "Q4_0"); + CHECK(q.quant == "Q4_0"); + CHECK(q.params == 4096); + CHECK(q.bytes == 64 * 64 / 32 * 18); // offset-derived: runs to end-of-file +} + +TEST_CASE("schematic schema matches the cached 4-bit Llama checkpoint") { + const std::string dir = MLXFORGE_MODEL_DIR_4BIT; + if (dir.empty() || !std::ifstream(dir + "/config.json").good()) { + MESSAGE("MLXFORGE_MODEL_DIR_4BIT not present; skipping"); + return; + } + const ModelConfig cfg = ModelConfig::from_file(dir + "/config.json"); + const ModelSchema s = mlxforge::inspect::build_schema_from_safetensors(dir, cfg); + + CHECK(s.family == "llama"); + CHECK(s.quant_summary.find("4-bit") != std::string::npos); + CHECK(s.tied_embeddings); + // Llama-3.2-1B is ~1.24B parameters; the schema's logical count must land + // within 2% (the packed-shape unpacking is what this gates). + CHECK(s.total_params > 1.21e9); + CHECK(s.total_params < 1.26e9); + // Every per-layer tensor classifies into a real bucket. + for (const auto& e : s.tensors) { + CAPTURE(e.name); + if (e.layer >= 0) CHECK(e.component != "other"); + } +} + +TEST_CASE("schematic schema reads the cached GGUF checkpoint") { + const std::string path = MLXFORGE_GGUF_MODEL; + if (path.empty() || !std::ifstream(path).good()) { + MESSAGE("MLXFORGE_GGUF_MODEL not present; skipping"); + return; + } + const ModelSchema s = mlxforge::inspect::build_schema_from_gguf(path); + CHECK(s.family == "llama"); + CHECK(s.cfg.n_layers == 16); + CHECK(s.by_layer.size() == 16); + CHECK(s.total_params > 1.21e9); + CHECK(s.total_params < 1.26e9); + for (const auto& e : s.tensors) { + CAPTURE(e.name); + if (e.layer >= 0) CHECK(e.component != "other"); + } +} diff --git a/tests/inspect/safetensors_header_test.cpp b/tests/inspect/safetensors_header_test.cpp new file mode 100644 index 0000000..c9a7d11 --- /dev/null +++ b/tests/inspect/safetensors_header_test.cpp @@ -0,0 +1,123 @@ +// Metadata-only safetensors header parsing: pure I/O + JSON units over tiny +// synthetic files written into the temp dir. No model or MLX arrays needed. +#include + +#include +#include +#include +#include +#include + +#include + +#include "inspect/safetensors_header.h" + +namespace { + +namespace fs = std::filesystem; +using mlxforge::inspect::SafetensorsEntry; + +// Write a synthetic .safetensors: 8-byte LE header length + JSON header. The +// data section is zero bytes — the reader never touches it. +void write_safetensors(const std::string& path, const nlohmann::json& header) { + const std::string h = header.dump(); + const uint64_t len = h.size(); + std::ofstream f(path, std::ios::binary); + f.write(reinterpret_cast(&len), sizeof(len)); + f.write(h.data(), static_cast(h.size())); +} + +// A scratch directory that cleans itself up. +struct TempDir { + fs::path dir; + explicit TempDir(const std::string& name) : dir(fs::temp_directory_path() / name) { + fs::remove_all(dir); + fs::create_directories(dir); + } + ~TempDir() { fs::remove_all(dir); } + std::string str() const { return dir.string(); } +}; + +std::map by_name(const std::vector& v) { + std::map m; + for (const auto& e : v) m.emplace(e.name, e); + return m; +} + +} // namespace + +TEST_CASE("safetensors header parses names, dtypes, shapes and byte sizes") { + TempDir td("mlxforge_st_header"); + const nlohmann::json header = { + {"__metadata__", {{"format", "pt"}}}, + {"model.embed_tokens.weight", + {{"dtype", "BF16"}, {"shape", {128, 64}}, {"data_offsets", {0, 16384}}}}, + {"model.layers.0.self_attn.q_proj.weight", + {{"dtype", "F16"}, {"shape", {64, 64}}, {"data_offsets", {16384, 24576}}}}, + {"model.norm.weight", {{"dtype", "F32"}, {"shape", {64}}, {"data_offsets", {24576, 24832}}}}, + }; + const std::string file = td.str() + "/model.safetensors"; + write_safetensors(file, header); + + const auto entries = mlxforge::inspect::read_safetensors_header(file); + CHECK(entries.size() == 3); // __metadata__ skipped + const auto m = by_name(entries); + + const auto& embed = m.at("model.embed_tokens.weight"); + CHECK(embed.dtype == "BF16"); + CHECK(embed.shape == std::vector{128, 64}); + CHECK(embed.nbytes == 16384); + + const auto& norm = m.at("model.norm.weight"); + CHECK(norm.dtype == "F32"); + CHECK(norm.shape == std::vector{64}); + CHECK(norm.nbytes == 256); +} + +TEST_CASE("safetensors dir merges sharded headers via the index") { + TempDir td("mlxforge_st_sharded"); + const nlohmann::json index = {{"weight_map", + {{"a.weight", "model-00001-of-00002.safetensors"}, + {"b.weight", "model-00002-of-00002.safetensors"}}}}; + std::ofstream(td.str() + "/model.safetensors.index.json") << index.dump(); + write_safetensors(td.str() + "/model-00001-of-00002.safetensors", + {{"a.weight", {{"dtype", "F16"}, {"shape", {4}}, {"data_offsets", {0, 8}}}}}); + write_safetensors(td.str() + "/model-00002-of-00002.safetensors", + {{"b.weight", {{"dtype", "F16"}, {"shape", {8}}, {"data_offsets", {0, 16}}}}}); + + const auto entries = mlxforge::inspect::read_safetensors_dir(td.str()); + CHECK(entries.size() == 2); + const auto m = by_name(entries); + CHECK(m.at("a.weight").nbytes == 8); + CHECK(m.at("b.weight").nbytes == 16); +} + +TEST_CASE("safetensors dir falls back to the single file on a stale index") { + TempDir td("mlxforge_st_stale"); + // The index references a shard that was never downloaded; the consolidated + // model.safetensors alongside it must win (mirrors load_weights' fallback). + const nlohmann::json index = {{"weight_map", {{"a.weight", "model-00001-of-00009.safetensors"}}}}; + std::ofstream(td.str() + "/model.safetensors.index.json") << index.dump(); + write_safetensors(td.str() + "/model.safetensors", + {{"a.weight", {{"dtype", "F16"}, {"shape", {4}}, {"data_offsets", {0, 8}}}}}); + + const auto entries = mlxforge::inspect::read_safetensors_dir(td.str()); + CHECK(entries.size() == 1); + CHECK(entries[0].name == "a.weight"); +} + +TEST_CASE("safetensors dir throws when neither layout exists") { + TempDir td("mlxforge_st_empty"); + CHECK_THROWS_AS(mlxforge::inspect::read_safetensors_dir(td.str()), std::runtime_error); +} + +TEST_CASE("safetensors header rejects a corrupt header length") { + TempDir td("mlxforge_st_corrupt"); + const std::string file = td.str() + "/model.safetensors"; + { + std::ofstream f(file, std::ios::binary); + const uint64_t bogus = ~0ull; // far past the sanity cap + f.write(reinterpret_cast(&bogus), sizeof(bogus)); + } + CHECK_THROWS_AS(mlxforge::inspect::read_safetensors_header(file), std::runtime_error); +} diff --git a/tests/inspect/schematic_html_test.cpp b/tests/inspect/schematic_html_test.cpp new file mode 100644 index 0000000..b5dcce4 --- /dev/null +++ b/tests/inspect/schematic_html_test.cpp @@ -0,0 +1,109 @@ +// The schematic HTML renderer: self-contained output (no external refs), the +// embedded schema JSON round-trips, and hostile tensor names cannot break out +// of the data ", start); + REQUIRE(end != std::string::npos); + return html.substr(start, end - start); +} + +} // namespace + +TEST_CASE("schematic HTML is a self-contained page with no external references") { + const std::string html = mlxforge::inspect::render_schematic_html(sample_schema()); + + CHECK(html.rfind("", 0) == 0); + CHECK(html.find("") != std::string::npos); + // Offline-only: no external scripts, styles, fonts or fetches. + CHECK(html.find("