Skip to content

nnef: lazy weight loading, so a pruned model never reads what it discards - #2524

Draft
czoli1976 wants to merge 1 commit into
sonos:mainfrom
czoli1976:feat/lazy-nnef-weights
Draft

nnef: lazy weight loading, so a pruned model never reads what it discards#2524
czoli1976 wants to merge 1 commit into
sonos:mainfrom
czoli1976:feat/lazy-nnef-weights

Conversation

@czoli1976

@czoli1976 czoli1976 commented Jul 26, 2026

Copy link
Copy Markdown
Contributor

Read a .dat's 128-byte header instead of its payload, wire the tensor as a LazyConst
carrying only its fact, and read the payload later for the constants still in the graph —
so a caller that prunes a model first never pays for the parts it discards.

This is a second take on #1386, which I found while looking for the "stale PR" you
mentioned on #2482. Its LazyConstProvider seam is the right one and I have kept it; what
changed is when the value appears.

Why not lazy TensorStorage

You suggested the exotics TensorStorage framework, and I tried to take it that way first.
It does not work for the case that motivates this: reaching concrete storage is a
type-based downcast (storage_as::<BlockQuantStorage>() in matmul/pack.rs), so a lazy
wrapper is simply not that type, and a q40 model would silently lose its packed matmul path
— block_quant_einsum_weights and the pack_a selection would decline to fire rather than
error. A file-backed TensorStorage still looks right for plain weights and for the
CUDA-direct-load case, but it cannot carry block-quant, so laziness has to sit one level up
in an op.

What changed against #1386

  • LazyConst is stateless and never evaluated; eval errors telling you to materialize
    first. Lazy values #1386 marked it stateful and re-read from disk on every eval, which for a weight
    touched each decode step is fatal, and which also stops it being a Const so declutter
    and the block-quant fusion no longer see it.
  • It must be excluded from eager evaluation explicitly: a node with no inputs satisfies
    "all inputs are constant" vacuously, so PropConst and compute_const_facts both tried
    to evaluate it. (wire_node's eager fold is already guarded by input_facts.len() > 0.)
  • LazyDatLoader in Lazy values #1386 still called read_tensor in full just to keep dt_shape, so
    the first pass streamed the whole model through memory anyway. read_tensor_header stops
    at 128 bytes; read_tensor is refactored to share it, so there is one implementation.
  • Directory-only is now deliberate and documented rather than incidental — a .nnef.tgz is
    a gzip stream with no random access, matching your "so file system" caveat. The eager
    path is untouched, and ResourceLoader needed no signature change.

Shape of it

let mut model = tract::nnef()?.load_lazy(dir)?;   // facts only
// ... prune: extract a subgraph, drop layers, whatever ...
materialize_lazy_consts(&mut model)?;             // reads only what is left

load_lazy returns an un-decluttered model and is an inherent method on api/rs's Nnef,
so no trait change and no impact on the proxy implementations.

Measured

On Qwen2.5-7B-Instruct-q40ef16, unpacked — 4085 MiB across 397 tensors:

  • lazy load types all 397 from headers, no payload read;
  • decluttering with lazy weights leaves ~226 nodes that could not be folded into constants
    (AddAxis, Cast, Reshape), and after materialize_lazy_consts plus a second
    declutter the op histogram is identical to an eager load;
  • cutting it in two and materializing each half reads 2051 MiB and 2034 MiB, each
    tensor exactly once.

Four tests in nnef/tests/lazy_weights.rs cover the small cases: typing without reading,
parity with an eager load, pruning before materializing, and declutter leaving lazy
constants alone. Header parity across every writable dtype, and for a block-quant tensor,
is covered in nnef/src/tensors.rs.

Caveats

  • Peak memory only improves if you actually prune; materializing everything is equivalent
    to an eager load plus a second declutter pass.
  • I have not measured wall-clock either way. Load does less I/O, declutter does more work,
    and I would rather not claim a number I have not taken.
  • graph.nnef already carries each variable's shape, so strictly only the dtype needs the
    header. I read both and cross-check them.

Happy to reshape any of this — in particular whether materialize_lazy_consts should be a
registered ModelTransform instead of a free function, and whether load_lazy belongs on
NnefInterface rather than as an inherent method. 🍍

@czoli1976 @claude
nnef: load a model's weights only when they are needed
44cf1c8 
Loading was all or nothing: every .dat was read in full before the graph was
even parsed, so a caller that only wants part of a model — one shard of it,
or a subgraph — still paid for the whole thing, and a model too big for the
machine could not be opened at all to be cut down. The .dat header is a fixed
128 bytes and fully describes the tensor that follows, so read that alone and
wire a LazyConst carrying the fact; materialize_lazy_consts then reads only
the constants still in the graph. LazyConst is excluded from eager evaluation
because a node with no inputs satisfies "all inputs are constant" vacuously.
Lazy loading needs to seek, so it is offered for unpacked directories only; a
.nnef.tgz is a gzip stream and keeps the existing eager path.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@czoli1976

czoli1976 commented Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

@kali lazy w loading

@kali

kali commented Jul 28, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Ha. Ok, indeed, I can see the problem with TensorStorage. Frustrating.

I'd like to avoid the load_lazy() pattern. I think we need a new method, not an offline knob, but I'd like to make it generic enough so it becomes a shared extension point. Same spirit as the Transforms. Maybe a load_with_options() with a JSON/RON serialized dictionary with a tract_nnef_lazy_weights: true to support your case ? WDYT ?

Agree to add it to the exclusion list in PropConst. At least for now. The statefull() / stateless() flag system need some attention, it has been abused so much... and throwing errors (like TooEarly) is ugly. eval() should work, even if it is stupidly expensive (re-reading from FS), to allow evaluation of declutter-but-not-optimised graph in debug/audit context.

For the materialization, I can see two options: a ModelTransform could actually do it, or it could be integrated in the codegen() and prepare phase.

@github-actions

github-actions Bot commented Jul 28, 2026

Copy link
Copy Markdown

⚠️ Bench vs main — no speed regressions · 3 secondary regression(s)

Reference: 2026-07-28 morning nightly run (0d old) · full report → run

Speed — evaltime · prefill · decode

no inference-speed regressions

⚠️ 3 secondary regression(s)
Δ metric device main → PR
⚠️ +38.9% en_tdnn_pyt_15M
load · pulse_120ms		 apple-m1-max 90 ms → 125 ms
⚠️ +33.0% en_tdnn_pyt_15M
load+optimize · pulse_120ms		 apple-m1-max 106 ms → 141 ms
⚠️ +5.5% llama_3_2_3B_instruct_q40ef16_541
load · cuda		 jetson-orin-nx 3.13 s → 3.3 s

@czoli1976

czoli1976 commented Aug 1, 2026

Copy link
Copy Markdown
Contributor Author

Could you think a bit about it and come up which one of the two sounds like the best option?

@kali

kali commented Aug 2, 2026

Copy link
Copy Markdown
Collaborator

I'm leaning towards codegen() time, so it's implicit and does not require extra attention from the user. Unless you hit a snag of course.

@czoli1976

czoli1976 commented Aug 2, 2026 via email

Copy link
Copy Markdown
Contributor Author

@kali

kali commented Aug 17, 2026

Copy link
Copy Markdown
Collaborator

@czoli1976 this one is yours to look at, now, right ?

@czoli1976

czoli1976 commented Aug 17, 2026 via email

Copy link
Copy Markdown
Contributor Author

@kali

kali commented Aug 17, 2026

Copy link
Copy Markdown
Collaborator

i'm putting it back in draft, the green tick is going to confuse me.

@kali
kali marked this pull request as draft August 17, 2026 08:28
@czoli1976

czoli1976 commented Aug 17, 2026 via email

Copy link
Copy Markdown
Contributor Author

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@czoli1976 @kali