Skip to content

Navigation Menu

Sign in
Sign up

Idea: smarter approaches to loading model weights in MFlux #717

ianscrivener started this conversation in General
Discussion options

Model Weight Loading Suggestions for MFlux


By Claude Opus 5 - discussion document — preliminary. Based on a source read of MFlux at
commit 89d2aff.

Why this is worth looking at

MFlux currently loads every weight of every component into memory before it
constructs a single model object. All four components — VAE, transformer, and
both text encoders — are read from disk and held simultaneously, then copied
into the model, and only released later by a callback that runs after the
prompt has been encoded.

For a Flux.1-dev generation that means peak memory during loading is roughly
the whole model, before any generation has started. On a 128 GB Mac Studio that
is invisible. On a 16 GB MacBook Air it is the difference between working and
not working.

This document sketches what could be done about it, roughly in order of
value-per-unit-of-effort. Nothing here is a criticism of the current design —
load-everything is the right first implementation, and it is what every
framework starts with.

Note on scope: MFlux supports several model families, not just Flux. Krea 2,
Z-Image, Qwen-Image, Ideogram4 and others all have different component counts
and sizes. Anything proposed here should work at the level of "components in a
model definition" rather than being written against Flux's particular four.


Two different processes, two different problems

MFlux does two broadly different things with model weights, and they have
almost nothing in common from a memory point of view.

Saving - mflux-save — reading a model in the standard Hugging Face safetensors layout
and writing it back out in MFlux's own format, optionally quantized. This is
what MFlux-save does. It is a transform: read, convert, write, done. Nothing
needs to remain in memory afterwards.

Inference - mflux-generate — loading a model in order to generate images. Weights must be
resident while they are being used, and the question is how much can be freed
between phases.

The save path is by far the easier of the two to fix, and it is where the
current behaviour is most clearly wrong: it appears to hold roughly twice the
model in memory
, when it needs to hold approximately one component.

TLDR, a lazy-loaded mflux-save can save a substantial amount of memory with negligible save speed loss.

mflux-save has an unambiguous solutions. Everything in the
inference path depends on how much memory the user has; the save path does not.
See Saving a model below, and
A9.

There is a third consideration that cuts across both: MFlux is one-and-done.
Each invocation loads everything, generates, and exits, discarding all of it. A
user producing ten images pays the full load cost ten times. See
Process lifetime below, and
A10.

worth considering: mflux-generate-server...


There is no single right answer — there are two regimes

This section is about the inference path only. The save path has a single
correct answer and is covered separately below.

For inference: memory and speed pull in opposite directions, and which one
matters depends entirely on the user's hardware.

If the model comfortably fits, eager loading is optimal and nothing below
should be applied. Load everything once, keep it resident, and never touch the
disk again. Freeing a component you will need in thirty seconds is pure waste,
and streaming weights from disk when you have RAM to spare is strictly slower
for no benefit. A user with 128 GB and a 24 GB model wants the current
behaviour, possibly more aggressive — keep components loaded between
generations too.

If the model does not comfortably fit, the calculus inverts. Sequential
loading, freeing between phases and streaming blocks are the difference between
running slowly and not running at all. Some added load time is an easy trade
against a generation that would otherwise fail.

Most of this document describes the second regime, because that is where the
current design leaves value on the table. But the first regime is the one many
users are in, and any change should be measured against both. A change that
halves peak memory while adding thirty seconds to every generation is a clear
win for one group and a clear regression for the other.

The practical conclusion is that the strategy should be chosen at runtime
from available memory and model size
, not fixed at build time. That framing
runs through everything below and is developed in
A0.


The shape of the problem

Three distinct costs get conflated when people say "memory":

  1. Peak during load — how much is resident at the worst moment while
    reading from disk. Currently: everything, twice over, briefly.
  2. Steady state during generation — how much stays resident while
    denoising. Currently: the transformer and VAE, after text encoders are
    dropped.
  3. Time to first image — how long the user waits before anything happens.
    Currently dominated by reading every byte of every component.

These pull in different directions. Streaming weights from disk lowers (1) and
(2) but can hurt (3) if done naively. Worth being explicit about which one any
change targets. A1

There is a fourth cost that only appears in the memory-constrained regime:
time per denoising step, if weights are being read from disk during the
forward pass rather than sitting in memory. In the comfortable regime this is
zero. In the constrained regime it can dominate everything else.


The save path

Converting a model does not need the model in memory. It needs one component
in memory at a time
, and arguably one tensor at a time.

The sequence should be:

for each component:
 read it
 convert (rename, reshape, quantize)
 write it
 release it

Peak memory is then the size of the largest single component — for most models
the transformer, and for a quantized save, the transformer at its quantized
size if conversion happens per tensor as it is read.

What appears to happen instead: the loading machinery is shared with the
inference path, so the save path inherits its behaviour — every component is
read and accumulated before anything is written, and the model objects may be
constructed as well even though nothing is going to run a forward pass.

That gives roughly two copies of a model that never needed one.

Why this is worth doing first:

  • No regime dependency. A user with 128 GB gains nothing from holding a
    component they have already written to disk, so there is no policy to
    design and no override to expose.
  • No risk of regression. Nothing about generation speed is affected.
  • It is the operation most likely to fail outright on a constrained machine,
    because a user converting a model to fit their hardware currently needs
    enough memory to hold it unconverted first.

That last point is worth stating plainly: someone quantizing a model because
it does not fit currently needs it to fit in order to quantize it.

A9

One exception worth preserving

MFlux-generate --quantize quantizes on the fly and then immediately generates
with the result. There, holding the converted weights is correct — they are
about to be used.

The distinction is whether there is a consumer waiting:

Command After conversion Correct behaviour
MFlux-save writes and exits release each component after writing
MFlux-generate --quantize generates immediately retain, subject to the inference regime

Sharing one code path between these is fine; sharing one policy is not.


Process lifetime: one-and-done

Every MFlux invocation starts from nothing. It reads the weights, builds the
models, quantizes if asked, generates one image, and exits — throwing all of it
away. The next invocation does the same work again.

For a user generating ten images this is ten full load cycles, and only one of
them was unavoidable.

The operating system's page cache does not solve this. It makes the second
read of the same file fast, so the disk cost largely disappears. But everything
after the read still happens every time: deserializing into arrays,
constructing the model objects, applying the weights, and — if quantizing on
the fly — running the entire quantization pass over every tensor.

That compute is repeated in full on every run, and on a large model it is not
small.

Three levels of fix

Batch within one invocation. Accept multiple prompts or multiple seeds in a
single run and generate them all with the models loaded once. No architectural
change — just a loop around the generation call instead of around the process.
This covers the most common case, which is generating variations, and it is by
some distance the cheapest thing on this list.

Interactive session. A mode that loads the models, then accepts prompts
until told to stop. The same mechanism as batch, with better ergonomics for
exploration where the user does not know the next prompt in advance.

Persistent server. A background process holding models in memory, with
short-lived clients sending requests. The largest change, and the one with the
most operational surface — lifetime management, eviction, concurrency.

Worth noting before pursuing the third: ComfyUI already is this, and it is
where most users doing high-volume work already are. Duplicating it is a
significant investment to reach a workflow that exists elsewhere.

Why this belongs in a memory document

Because batching changes the right answer to every other question here.

The current behaviour frees text encoders after prompt encoding. For a single
image that is correct. For ten images it means loading and freeing the text
encoder ten times to save memory during a denoise loop that had room for it.

Any batching work should therefore land alongside the strategy selection in
A0 — the loader needs to know whether more
generations are coming before it decides what to
discard. A10


Suggestions

These concern the inference path. For the save path, see above.

1. Load components in the order they are used, and free as you go

The cheapest structural win. Text encoders are needed once, at the start, and
never again. The transformer is needed for the whole denoise loop. The VAE is
needed once, at the end.

Loading all four upfront means the VAE sits in memory for the entire generation
doing nothing, and the text encoders sit there long after their work is done.

A sequential pattern — load encoder, encode, free, load transformer, denoise,
free, load VAE, decode — bounds peak memory at roughly the largest single
component rather than the sum of all of them.

MFlux already does part of this via MemorySaver, but only after loading
everything first. Moving the load itself into the sequence is the change.
A2

Regime note: this should be a policy, not a default. A user with ample
memory generating several images wants components to stay loaded — freeing the
text encoder after each prompt means re-reading it for the next one. Sequential
load-and-free is the right behaviour when memory is tight and the wrong
behaviour when it is not.

Value: large when constrained, negative when not. Effort: moderate —
needs the weight loader to become lazy per component rather than eager for all.


2. Don't hold the raw weights and the model copy at the same time

Currently the loaded weight dictionary and the constructed model both reference
the weights for a period, and the raw dictionary is not explicitly dropped.
Whether this doubles memory depends on whether the framework copies or aliases
the arrays, which is worth measuring rather than assuming.

If it does copy, releasing each component's raw weights immediately after they
are applied is a small change with a large effect on peak.
A3

Value: potentially large, unknown until measured. Effort: small.


3. Read only the tensors a component actually needs

Safetensors files carry a header listing every tensor's name, shape and byte
offset. There is no requirement to read all of them — you can seek to the ones
you want.

MFlux's own safetensors reader already parses the header and memory-maps
individual tensors at their offsets. It reads everything only because the loop
has no filter.

The reader change is trivial. The harder part is knowing which keys a component
needs before loading, since today that is derived from the loaded dictionary
rather than declared in advance. A4

Value: moderate on its own, large as an enabler for everything below.
Effort: small for the reader, moderate for the plumbing.


4. Stream transformer blocks rather than holding all of them

The most aggressive option, and the one that changes the memory ceiling
fundamentally.

A transformer's blocks are used in sequence, once per denoising step. In
principle you only need the current block's weights resident — load block n,
compute, release, load block n+1.

In practice this trades memory for I/O, and the trade is only good if the
weights can be re-read fast enough to keep the GPU busy. On a machine with fast
local storage and a model that does not fit in memory, it is the difference
between running slowly and not running at all.

ComfyUI does exactly this on supported hardware, faulting weights in at forward
time from a memory-mapped file. It is not exotic; it is becoming the norm.
A5

Regime note: the clearest example of a change that must never be on by
default. On a machine with headroom, streaming blocks adds I/O to every step of
every generation in exchange for memory nobody needed freed.

Value: very large for constrained machines, actively harmful for
unconstrained ones. Effort: significant.


5. Keep a component loaded across generations

A user generating ten images in a row currently pays the full load cost each
time if each run is a separate process. Even within a session, dropping and
reloading a component that will be needed again in thirty seconds is wasteful.

A cache keyed on component identity, with a memory budget and a
least-recently-used eviction policy, would let repeated generations skip most
of the loading work. This matters more for interactive use than for batch.
A6

Value: large for interactive use. Effort: moderate.


6. Quantize during load, not after

If a model is being quantized on the fly, the current order — load full
precision, then quantize — means peak memory is set by the full-precision
size even though the end state is much smaller.

Quantizing tensor by tensor as they are read would keep peak close to the
quantized size. This only applies to on-the-fly quantization, not to
pre-quantized models. A7

Value: moderate, situational. Effort: moderate.


7. Let the model definition declare its memory behaviour

Different model families have very different shapes. Some have one text
encoder, some have two; some have a vision tower that is unused for
text-to-image; component sizes vary by an order of magnitude.

Rather than hard-coding a strategy, the weight definition for each model could
declare what its components are, when they are needed, and what can be freed
after each phase. The loader then implements one policy driven by that
declaration.

This is the change that makes everything above work across all supported models
rather than just Flux. A8

Value: structural. Effort: moderate, and best done before the others.


Where the user-visible value is

Split by regime, because the same change helps one group and hurts the other:

Change Ample memory Tight memory
Sequential load + free Harmful — re-reads what it just freed Large — makes models run that currently don't
Block streaming Harmful — disk I/O on every step Very large — the difference between slow and impossible
Quantize during load Neutral Moderate — lowers the headroom needed to quantize
Cross-generation caching Large — cuts waiting between runs Situational — competes with freeing
Selective tensor reads Positive — never read what isn't used Positive — same
Skip unused components entirely Positive Positive
Declarative lifecycles Enabler Enabler
Sequential save path Positive — lower peak, no downside Large — makes conversion possible at all

Two of these help everyone regardless of hardware: not reading tensors you
never use
, and not loading components the task doesn't need — a vision
tower is dead weight in text-to-image on any machine. Those are the safest
things to do first.

Everything else needs a policy that knows which regime the user is in.


Suggested order

Unconditional wins first — these help every user regardless of hardware:

  1. Sequential save path (A9) — no regime
    dependency, no regression risk, and it is where the current behaviour is
    most clearly wrong
  2. Declarative component lifecycles (Performance report - 2021 M1 Pro 16GB #7 ) — the foundation, and immediately
    enables skipping components a task doesn't need
  3. Selective tensor reads (Added .gitignore #3 ) — never read what is never used
  4. Drop raw weights promptly (Automatically download models from Huggingface or used cache ones #2 ) — cheap, but measure first

Then the regime-dependent work, behind a strategy chosen at runtime:

  1. Runtime strategy selection (A0) —
    detect available memory, pick eager or sequential
  2. Sequential load and free (How much memory is needed? #1 ) — the constrained-regime win
  3. Cross-generation cache (Add CLI args for easier execution #5 ) — the ample-regime win
  4. Quantize during load (Performance report - 2023 M2 Max 96GB #6 )
  5. Block streaming (Generate image at different resolutions #4 ) — the ambitious one, constrained regime only

Items 1–4 carry no risk of regressing anyone's performance and could plausibly
be done together. Item 5 is what makes 6 and 9 safe to ship.

Item 1 is the smallest and highest-value of the lot: it is a self-contained
change to one code path, it needs no policy or detection, and it removes a
failure mode users hit today.



Appendix

A0: Choosing a strategy at runtime

The core observation: the optimal loading strategy is a function of available
memory divided by model size
, and both are knowable before loading starts.

The regimes

Ratio (available memory ÷ model size) Regime Strategy
> ×ばつ Comfortable Eager load everything; keep resident across generations; free nothing
×ばつ×ばつ Adequate Eager load; free components after their phase; do not cache across runs
×ばつ×ばつ Tight Sequential load and free; peak bounded by largest component
< ×ばつ Constrained Sequential plus block streaming; accept slower steps to run at all

The boundaries are illustrative — they would need calibrating against real
measurements. The structure is the point: a small number of named strategies,
selected automatically, with a manual override.

What needs to be known

Available memory. Queryable on all supported platforms. On unified-memory
Macs this is the whole system budget shared with everything else running, which
argues for a conservative margin rather than assuming the full amount is
available.

Model size. Known before loading from the file sizes on disk, adjusted for
whatever quantization is being applied on the fly. This does not require reading
the weights — just the file metadata.

Task shape. Whether components can be skipped entirely (a vision tower for
text-to-image), and whether the user is generating one image or many. The latter
strongly affects whether caching or freeing is correct.

Why automatic, with an override

Users should not have to understand any of this. The current --low-ram flag
is a blunt version of the same idea — it exists precisely because the right
behaviour differs by machine — but it puts the decision on the user and offers
only two settings.

An automatic choice with a manual override gets the common case right without
removing control:

(default) detect and choose
--memory-strategy eager
--memory-strategy sequential
--memory-strategy stream

The batch case deserves special handling

A user generating twenty images with the same model and different prompts has a
different optimum from a user generating one. In the batch case:

  • text encoders are needed once per prompt, so freeing them between images is
    wasteful if there is room to keep them
  • the transformer is needed continuously and should never be freed
  • load cost is amortised across many generations, so a slower load in exchange
    for faster steps is a good trade

This is the clearest case where the current behaviour — free the text encoders
after every prompt encoding — is measurably wrong for a user with memory to
spare.

A note on measuring

Any change proposed in this document should be benchmarked in both regimes
before being adopted. A useful minimum matrix:

  • a machine where the model fits with headroom
  • a machine where it barely fits
  • a machine where it does not fit without streaming

Reporting peak memory and wall-clock time for each. A change that improves one
column and regresses another is not necessarily wrong — but it needs to be a
policy rather than a default.

↩ back


A1: The three costs in detail

Peak during load is currently the sum of all components' raw weights, plus
whatever the model objects hold after update. Because all components are
accumulated into one structure before any model is built, there is no point at
which only part of the model is resident.

Steady state is better handled. A callback frees the text encoders after
the prompt is encoded and before the first denoise step, and frees the
transformer after the loop and before VAE decode. So during denoising the
resident set is transformer plus VAE, which is close to minimal.

The gap is that this callback runs after everything has already been loaded.
The peak has already happened.

Time to first image is dominated by disk reads on a cold start. Any change
that reads less will improve it; any change that reads the same bytes in
smaller pieces will not, and may make it worse through I/O overhead.

Worth measuring these separately — a change that halves peak memory while
adding two seconds to load time is a good trade for some users and a bad one
for others.

↩ back


A2: Sequential component loading

The current order is: load all components → construct all model objects → apply
weights → generate.

The proposed order is: for each component, in use order — load, construct,
apply, use, free.

Two complications:

Shared sources. Some components may live in the same file or share weights.
MFlux already has a raw-weights cache for this reason, which currently only
adds retention. A sequential loader would need to know which sources are shared
and either keep them or accept re-reading.

LoRA application. LoRA weights are applied after the base model is
constructed. If the transformer is loaded later in the sequence, LoRA
application has to move with it.

Expected effect. For Flux.1-dev the components are roughly: T5 encoder
~9.5 GB at bf16, transformer ~24 GB, CLIP ~250 MB, VAE ~170 MB. Sequential
loading bounds peak at the transformer alone rather than the sum — a saving of
roughly 10 GB at bf16, proportionally less when quantized.

For models with larger text encoders — Qwen-Image uses a 7B VLM — the saving is
correspondingly larger.

When this is the wrong thing to do. On a machine with ample memory the same
sequence costs time and saves nothing anyone needed. Worse, in a multi-image
session it means re-reading the text encoder for every prompt, turning a
one-time cost into a per-image one.

The correct behaviour there is the opposite: load everything once, keep it all,
and never touch the disk again for the rest of the session.

This is why sequential loading belongs behind a strategy selection rather than
becoming the default. The change worth making unconditionally is making the
loader capable of loading per component; whether it then frees them is a
policy question answered at runtime.

↩ back


A3: Raw weights and the model copy

After weights are read from disk they exist as a dictionary. They are then
applied to the constructed model objects. For a period both reference the same
data.

Whether this is a real doubling depends on whether the update operation copies
the arrays or aliases them. If it aliases, there is no doubling and this
suggestion is moot. If it copies, peak is briefly twice the component size.

The source read did not settle this — it is marked unclear and flagged as
needing a look at whether model.update copies or aliases.

This should be measured before any work is done. A single memory trace
during load would settle it, and the answer determines whether this is a
significant win or a non-issue.

If it does copy, the fix is to drop each component's entry from the raw
dictionary immediately after applying it, rather than holding the whole
structure until the initialisation function returns.

↩ back


A4: Selective tensor reads

A safetensors file begins with a JSON header mapping each tensor name to its
data type, shape, and byte range within the file. Reading one tensor means
parsing the header and reading that byte range. Nothing else needs to be
touched.

MFlux has a reader that does precisely this — header parse, then per-tensor
memory-map at the offset. It reads every tensor because the iteration has no
filter. Adding an optional set of names to read is a one-line change to that
loop.

The obstacle is upstream. Today the set of keys a component needs is derived
from the dictionary after loading, by a mapping step that renames and
restructures. To filter before loading, that mapping would need to be
invertible — able to say "for component X, I need these source keys" without
having read them.

For many models this is straightforward, since components live in their own
subdirectories and the mapping is a prefix rule. For models where the mapping is
more involved, it may need a declared key list per component rather than a
derived one.

Worth noting: the framework's own load function already memory-maps
safetensors, so a tensor that is never touched is arguably not resident. But the
subsequent steps — restructuring the dictionary, applying it to the model, and
quantizing — touch everything, so in practice the whole file is realised. A
genuine saving requires filtering before those steps, not relying on laziness
after them.

↩ back


A5: Block-level streaming

A diffusion transformer applies its blocks in sequence. Block n does not need
block n+1's weights to be resident. In principle the resident set could be one
block plus activations, rather than the whole transformer.

The obvious problem is that the blocks are re-used every denoising step. A
Flux generation at 28 steps would read the entire transformer 28 times. At
25 GB that is 700 GB of reads, which no storage system will do quickly.

What makes it viable anyway:

  • The operating system's page cache means re-reads hit RAM, not disk, when
    there is spare RAM. Streaming then costs almost nothing on a machine with
    headroom.
  • Fast local storage narrows the gap considerably. A modern NVMe drive reads at
    several GB/s.
  • For a user whose alternative is "does not run at all," slow is better than
    impossible.

What ComfyUI does, as a reference point: with its dynamic VRAM mode
enabled, weights are memory-mapped and faulted in at forward time, with each
tensor tagged with its file offset so it can be re-read on demand. It allocates
virtual address ranges rather than moving most weights eagerly, and only
force-loads very small modules. This is on by default on supported hardware.

Implementation shape for MFlux would be: keep the block modules constructed
but unweighted, and populate each block's parameters immediately before its
forward call, releasing after. The forward loop structure already iterates
blocks explicitly, so the insertion point is clear.

Interaction with quantization. A streamed block still has to be
dequantized. If quantization state is per-tensor and self-contained — as it is
for group-wise affine schemes — this composes cleanly. Formats that carry
model-wide state would not.

This must never be a default. For a user whose model fits, streaming adds
I/O to every block of every step of every image and frees memory that was not
needed. It is the single clearest example in this document of an optimisation
that is transformative for one group and a straight regression for another.

Gate it on the constrained regime, or behind an explicit flag, and never turn
it on speculatively.

↩ back


A6: Cross-generation component cache

Two cases, with different solutions.

Within one process. A user generating several images in a session. Here the
fix is not to free a component that will be needed again — a cache with a
memory budget, holding the most recently used components and evicting when
under pressure. The interaction with sequential loading needs care: the whole
point of sequential loading is to free things, and the whole point of caching is
not to.

A reasonable policy: free when memory pressure demands it, keep otherwise. That
requires knowing the memory ceiling, which is queryable on most platforms.

Across processes. A user running the CLI repeatedly. Nothing survives
process exit, so each run pays the full cost. Options include a persistent
service holding models in memory, or relying on the operating system's page
cache to make the second read cheap. The latter is free and already happens;
the former is a much larger architectural change.

The page cache means the second run of the same model is often substantially
faster than the first, which may make the in-process case the only one worth
solving.

This is the ample-memory regime's headline feature, and the mirror image of
sequential loading. Where a constrained user wants everything freed as early as
possible, a user with headroom wants nothing freed at all — including between
generations.

Both are the same underlying capability: knowing when a component is needed
again. One acts on it by freeing, the other by retaining. That is another
argument for the declarative lifecycle description in A8 — it answers the
question both policies depend on.

↩ back


A7: Quantize during load

When quantizing on the fly, the current sequence is: read full-precision
weights → construct model → apply weights → quantize.

Peak memory is therefore set by the full-precision size, even though the model
being generated is much smaller. A user quantizing a 24 GB transformer to 4 bits
needs 24 GB of headroom to produce a 7 GB model.

Quantizing per tensor as it is read would keep peak near the quantized size plus
one tensor. The largest single tensor is typically an embedding table, which
bounds the overhead.

Complications:

  • Any quantization scheme needing statistics across tensors — activation-aware
    approaches, for instance — cannot be done in a single streaming pass without
    a prior calibration pass.
  • The mapping and restructuring steps would need to tolerate a mix of
    quantized and unquantized entries mid-load.

For pre-quantized models this suggestion does not apply — the weights arrive
already small.

↩ back


A8: Declarative component lifecycles

Today the loading strategy is implicit in the initialisation code for each
model family. That means any improvement has to be made once per family, and
divergence between families is easy.

A declarative alternative: each model's weight definition already lists its
components. Extend that to describe, per component:

  • when it is needed — prompt encoding, denoising, decoding
  • whether it can be freed after that phase
  • whether it is shared with another component
  • its approximate size, for budgeting
  • whether it is required at all for a given task (a vision tower is unused
    for text-to-image)

The loader then implements a single policy against that description, and every
model family benefits from any improvement to it.

Why this comes first. Sequential loading, selective reads and caching all
need to know when a component is needed and when it can go. Building each of
them against per-family logic means building each of them several times.

A concrete payoff available immediately: models with a vision-language text
encoder carry a vision tower that does nothing during text-to-image generation.
Declaring it as not-required for that task means never reading it — a saving of
gigabytes for zero effort beyond the declaration. This one helps every user on
every machine, which makes it the safest change in the document.

It also supplies what runtime strategy selection needs. Choosing between
eager and sequential requires knowing the total size before loading and which
components a task actually uses. A declaration provides both without reading a
byte of weights, so the strategy can be picked before the first file is opened.

Without it, every strategy has to guess or measure as it goes.

↩ back


A9: The save path in detail

What the operation actually requires

Converting a model from the Hugging Face layout to MFlux's own is a streaming
transform. For each tensor: read it, possibly rename it, possibly reshape it,
possibly quantize it, write it out. Nothing read at step n is needed at step
n+1.

The theoretical minimum resident set is therefore one tensor. The practical
minimum is one component, since components are the natural unit for the
renaming and mapping logic.

Why it currently costs more

Two mechanisms, either or both of which may apply — worth confirming against
the source before acting.

Shared loading machinery. The save path uses the same weight loader as
inference, which reads every component into one accumulated structure before
returning. If save calls that loader and then writes, it has already paid the
full-model peak before writing its first byte.

Unnecessary model construction. If the save path constructs the nn.Module
objects and applies weights to them — because that is how the mapping and
quantization logic is reached — then the raw weights and the module parameters
coexist. Nothing is going to run a forward pass, so the modules exist purely as
a vehicle for the transform.

Together these plausibly account for the observed doubling.

The fix

Keep the mapping and quantization logic, change what surrounds it:

for component in definition.components:
 raw = read_component(component) # one component's tensors
 converted = map_and_quantize(raw) # existing logic
 write_component(converted) # to the output path
 del raw, converted # explicit release

If the quantization step can be applied per tensor rather than per component,
the loop can go one level finer and peak drops to roughly the largest single
tensor — typically an embedding table.

Why the constrained case matters most

A user converting a model to a quantized format is, very often, doing so
because the full-precision model does not fit their machine. The current
behaviour requires the unconverted model to fit in memory before it can be
converted to something that fits.

That is a circular requirement, and it is the kind of thing that makes people
conclude the tool does not work on their hardware.

Sequential saving removes it: peak becomes the largest component rather than
the whole model, and with per-tensor quantization, close to the largest tensor.

Interaction with MFlux-generate --quantize

The one case where retention is correct. That command quantizes and then
immediately generates, so the converted weights have a consumer.

The distinction is not the conversion logic — it is what happens afterwards:

  • write and exit → release each component as soon as it is written
  • write and generate → hand the converted component to the inference path,
    which applies its own regime-dependent policy

Implementing this as a callback or a sink the conversion loop writes to keeps
one code path with two behaviours, rather than two code paths.

Suggested measurement

Before and after, on one model:

  • peak resident memory during MFlux-save
  • wall-clock time for the save
  • the same for MFlux-save --quantize

Expect peak to fall substantially and time to be roughly unchanged, since the
same bytes are read and written either way. If time rises noticeably, the
per-component reads are likely fragmenting sequential I/O — worth checking
before concluding the approach is wrong.

↩ back

You must be logged in to vote

Replies: 0 comments

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
1 participant

AltStyle によって変換されたページ (->オリジナル) /