Model lifecycle and interop
ABI version and layout arithmetic, safetensors/GGUF/ONNX/HuggingFace interop math, data/tensor/pipeline parallelism formulas, and checkpoint/quantization sizing, across four independent Float64 kernel
Four small modules for the arithmetic around a model artifact rather than the arithmetic inside it: the layout and version-compatibility contract a checkpoint's bytes have to satisfy (model_abi), the format-specific math for reading a model someone else produced (model_interop, covering safetensors, GGUF, ONNX and HuggingFace configs), the communication and memory formulas that describe splitting one model across many devices (model_parallel), and the sizing/compression math around saving and shrinking a checkpoint (model_serialization). Every function in all four modules takes and returns plain Float64 — there is no tensor, buffer, or struct type anywhere in this family, and no shared status-code vocabulary either: a degenerate input gets a clamped or defaulted number back, never an error code.
from algenta import MojoRuntime
runtime = MojoRuntime(mode="local")
runtime.execute("model_abi", "attention_head_dim", [4096.0, 32.0])["result"]
# 128.0 -- Llama-2-7B: 4096 hidden dim / 32 heads = 128 per head
runtime.execute("model_abi", "ffn_dim_from_model_dim", [4096.0, 8.0 / 3.0, 256.0])["result"]
# 11008.0 -- matches Llama-2-7B's real intermediate_size (4096 * 8/3, rounded up to a 256 multiple)The modules
model_abi
tensor-layout strides, dtype byte widths, NCHW/NHWC activation indexing, weight-shard byte offsets, and ABI version compatibility between a saved model and the running kernel
model_interop
safetensors tensor alignment, GGUF metadata sizing, ONNX per-op FLOPs, HuggingFace config/rope fields, checkpoint-version and tokenizer-vocab compatibility scores, and a perplexity helper for model-card quality proxies
model_parallel
data/tensor/pipeline parallelism math — gradient averaging, ring all-reduce volume, pipeline bubble fraction, layer-to-stage assignment, ZeRO-3 shard sizing, and comm/compute overlap ratios
model_serialization
checkpoint byte-size estimates, quantization bit widths, LoRA adapter parameter counts, pruning speedup, weight-sharing/sparsity ratios, and a rolling version-hash accumulator
Same division, two provenances
model_abi.attention_head_dim(model_dim, n_heads) and model_interop.hf_config_head_dim(hidden_size, num_attention_heads) are the identical formula — model_dim / n_heads, falling back to model_dim itself when n_heads <= 0 — duplicated across the two modules rather than shared. That is deliberate: model_abi speaks the runtime's own config fields, model_interop speaks a HuggingFace config.json's field names. The module boundary tracks which document produced the number, not a different piece of arithmetic. model_abi.model_config_head_count(model_dim, head_dim) is the same division run the other way (model_dim / head_dim, defaulting to 1.0 when head_dim <= 0), so attention_head_dim and model_config_head_count invert each other: attention_head_dim(4096, 32) == 128 and model_config_head_count(4096, 128) == 32.
Two RoPE-scaling formulas that are not interchangeable
Both modules have a function for stretching RoPE's base theta to a longer context, and they compute genuinely different things. model_abi.rope_theta_scale(base_theta, original_context, target_context) squares the context ratio — base_theta * (target/original)^2 — despite its own docstring citing the full NTK-aware exponent dim / (dim - 2); the function never receives a dim argument at all, so what actually ships is the fixed-quadratic approximation, not the per-dimension NTK formula. model_interop.hf_rope_theta_scale(base_theta, rope_scaling_factor) is a flat multiply by whatever factor a HuggingFace rope_scaling config supplies — no context-length ratio involved. Doubling context from 4096 to 8192 and passing an equivalent factor of 2.0 to the other function gives different answers:
A caller who assumes these two are the same operation under different names will scale theta by 2x too much or too little, depending on which one they meant to call.
Worked examples — model_abi
ffn_dim_from_model_dim rounds the raw FFN size up to the next multiple of round_to — always the next one, even when the raw value already lands exactly on a multiple:
activation_layout_index gives a different flat offset for the identical logical coordinate depending on the layout flag — the entire reason the function exists:
Indexing a NCHW-trained activation buffer with the NHWC formula (or vice versa) does not error — every index it produces is still in range, it just silently points at the wrong element. abi_version_compatible turns a version-diff comparison into the same yes/no shape every guard in this family uses:
Worked examples — model_interop
safetensors_tensor_offset pads a running byte offset up to the next alignment boundary, matching the real safetensors format's 8-byte tensor alignment — and is a no-op when the offset is already aligned:
onnx_node_flops reproduces the standard 2·M·N·K FLOP count for a matmul node:
gguf_kv_metadata_size's own docstring documents the GGUF string encoding as 8 + len (an 8-byte length prefix plus the character bytes), but the code it ships computes base + 8.0 * value_count — multiplying, not adding the fixed 8-byte prefix once. For a 21-byte key ("general.architecture") holding a 5-character string value, the true GGUF wire size is 25 + 8 + 5 = 38 bytes; the function as written returns:
Anything estimating real GGUF file sizes from this function's string branch will overcount; the int32/float32/bool branches (value_type 0, 1, 3) match their documented per-value byte costs exactly. checkpoint_version_compat decays smoothly rather than being a hard yes/no like abi_version_compatible:
tokenizer_vocab_overlap is a plain Jaccard index over vocab sizes and a shared-token count — a Llama tokenizer (32,000 tokens) against a GPT-2 tokenizer (50,257 tokens) sharing a rough 5,000 tokens comes back low, quantifying why weights rarely port across tokenizer families without a remap:
model_card_perplexity inverts the standard cross-entropy-to-perplexity relationship (exp(-mean log-likelihood)); a mean log-likelihood of exactly -ln(10) per token comes back as a perplexity of exactly 10:
Worked examples — model_parallel
tensor_parallel_split_dim and zero_stage3_param_bytes are both a plain division by a device count — one for a weight dimension, the other for parameter bytes:
pipeline_bubble_fraction shows why microbatch count matters more than stage count alone — the same 4 pipeline stages produce a very different bubble depending on how many microbatches fill them:
gradient_all_reduce_bytes and communication_computation_ratio chain together: the first gives a ring all-reduce's communication volume, the second checks whether that volume is small enough to hide behind compute.
pipeline_stage_assignment splits 32 layers evenly across 4 stages, and the boundary layers land exactly where floor(layer_id * n_stages / total_layers) says they should:
ddp_sync_overhead clamps at 1.0 rather than reporting an overhead above 100%, even when the inputs describe an impossible measurement (gradient sync time exceeding the whole step it is part of):
The clamp hides a measurement error rather than surfacing it — a 600ms sync inside a 500ms step is a bad input, not a real 120% overhead, and the function has no way to say so.
Worked examples — model_serialization
weight_quantization_bits is a staircase onto standard hardware widths, not a continuous function — several input precisions collapse onto the same bucket, and there is no 3-bit or 5-through-7-bit bucket at all:
checkpoint_size_estimate matches the well-known "weights, or weights times three for Adam" rule of thumb:
lora_parameter_count is where the efficiency argument for LoRA becomes a concrete number: a rank-8 adapter over 4 target modules on a 4096-wide model is a tiny fraction of the 7B base model it adapts.
pruning_sparsity_speedup scales the ideal 1/(1-sparsity) memory-bound speedup by how much of the workload is actually memory-bound — the same sparsity level gives a much smaller realized speedup than the ideal figure alone would suggest:
version_hash_component truncates a weight value to three decimal places (Int64(param_value * 1000.0)) before folding it into a rolling-hash accumulator step — two weights differing only past the third decimal place contribute an identical component:
No status vocabulary, no structs
Unlike attnk, tensorx, or llmk, none of these four modules define a struct or a shared status-code range — the runtime contract lists zero structs for all four (model_abi, model_interop, model_parallel, model_serialization each report "structs": []). Every function is a bare Float64-in/Float64-out arithmetic helper, and every guard against a degenerate input (world_size <= 0, total_layers <= 0, weight_norm <= 0, an unrecognized op_type or value_type code) returns a plain defaulted or clamped number rather than an error sentinel. A caller that wants to detect "this input was nonsensical" has to notice the fallback value itself — there is no TX_ERR_*-style code anywhere in this family to check instead.
A unit mismatch worth watching
model_abi.dtype_byte_width returns bytes per element (fp16 → 2.0); model_serialization .checkpoint_size_estimate expects dtype_bits in bits (fp16 → 16.0). Feeding one function's output straight into the other without multiplying by 8 silently understates a checkpoint size by 8x — there is nothing in either signature that stops a Float64 meant as bytes from being accepted as bits.
Distinct from serialization_io
Bit-exact weight serialization documents serialization_io, a completely different module: it round-trips named Float64 arrays to a versioned string with bit-exact fidelity, for the classical-stats namespace's own weight persistence. It shares zero functions with model_serialization — weights_to_string, weights_from_string, and string_checksum are the whole of its public API. model_serialization never serializes an actual weight value; it estimates byte sizes, parameter counts, and compression ratios about a model's weights (checkpoint_size_estimate, lora_parameter_count, weight_sharing_ratio) without ever writing one out. The "serialization" in each name refers to a different problem: one is a wire format, the other is a sizing calculator.
Verification
All four mojo/tests/libs/test_model_*.mojo files behind this page (test_model_abi.mojo, test_model_interop.mojo, test_model_parallel.mojo, test_model_serialization.mojo) are auto-generated coverage smoke tests: each calls every function in its module exactly once, with every argument fixed at 0.5, and checks only that the result is finite — none pin a specific numeric output, and none exercise the multi-branch behavior (dtype codes, op_type/value_type dispatch, layout flags, clamping) that the worked examples above depend on. Every number in this page was hand-traced against the real formulas in mojo/engine/model_abi.mojo, model_interop.mojo, model_parallel.mojo, and model_serialization.mojo, and is reproducible by re-running the same arithmetic — including the gguf_kv_metadata_size and ffn_dim_from_model_dim results that diverge from what each function's own docstring describes.
Next: see how a checkpoint's derived facts (including this family's own head-dim and RoPE fields) gate real model support in the model support matrix, or how a loaded checkpoint's tensors actually run a decode step in LLM inference kernels.
Last updated
Was this helpful?