[perf]: stop holding the whole checkpoint during DiT load, unblocking MiniMax H3 on one GB10 - #1714
Merged
SolitaryThinker merged 2 commits intoAug 26, 2026
Conversation
mergify
Bot
added
type: perf
Contributor
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI
🔴 PR merge requirementsWaiting for
This rule is failing.
|
KyleNeverGivesUp
changed the title
KyleNeverGivesUp
added a commit
to KyleNeverGivesUp/FastVideo
that referenced
this pull request
Aug 18, 2026
…ader's Every offload flag moves weights to host memory so the device can drop them. That trade only pays when the two are separate pools. On a unified-memory device they are one, so the move frees nothing, the copy is a pure loss, and the peak becomes the sum instead of the max. hao-ai-lab#1710 added `Platform.has_unified_memory()` and used it at one site: the text encoder's FSDP offload during loading. That is not enough, because a single flag acts in more than one place. `text_encoder_cpu_offload` does three things: component_loader.py:353 picks the load-time target device, so the model is placed on the host component_loader.py:407 gates the FSDP offload path, which hao-ai-lab#1710 covers minimax_h3_conditioning.py:292-302 gates a `.to(device)` before the conditioning forward and a `.to("cpu")` after it Gating only the loader leaves the model on the host and moves the copy to inference time. On a DGX Spark loading MiniMax H3 that is a 48 GB move in the middle of generation with about 11 GiB free, and the worker is killed. Decide it once, in `check_fastvideo_args`, so every call site sees the same answer. `UNIFIED_MEMORY_OFFLOAD_FLAGS` names the five flags and a test asserts the tuple still matches the dataclass, because missing one is silent: the run works and quietly pays twice, which is exactly how the text encoder survived the first pass at this. `use_fsdp_inference` is deliberately left alone. Sharding across ranks is a separate decision from where the weights live, and a unified-memory host can have more than one of these devices. MPS keeps its own branch, which does disable it, and reaches this one through `elif`. hao-ai-lab#1710's loader guard is not made redundant. `TextEncoderLoader.load` takes an explicit `cpu_offload` argument that bypasses the args-level decision, and that path still needs it. Measured on a DGX Spark, GB10, 121 GiB unified memory, one GPU, with hao-ai-lab#1710, hao-ai-lab#1711 and hao-ai-lab#1714 also applied. Before this change the conditioning stage is terminated 1.5 s in, every time, with nothing reported. After it, the run gets through input preparation, conditioning, latent preparation and the full denoising loop, and is terminated in the video decode instead. Getting past that last step also needs the video VAE decoder in fp16 rather than the fp32 it is pinned to at `models/vaes/minimax_h3_video.py:565`. With that as a local patch on top, MiniMax H3 completes a text to video generation on one GB10: 320 by 192, 124 frames, video and audio, 11.69 s. It could not load at all before this series. The VAE change is not in this PR because the encode path has to stay fp32 for FL2VA and Ref2VA, so it needs its own design and its own measurements. fastvideo/tests/platforms/test_unified_memory_offload.py: 11 passed Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
4 tasks
SolitaryThinker
force-pushed
the
release-checkpoint-tensors-during-load
branch
2 times, most recently
from
August 23, 2026 22:55
281245b to
c57c3b1
Compare
KyleNeverGivesUp
force-pushed
the
release-checkpoint-tensors-during-load
branch
from
August 26, 2026 10:29
c57c3b1 to
93bcee6
Compare
`load_model_from_full_model_state_dict` holds the whole checkpoint alive while
it copies it, so peak memory during a load is checkpoint plus model rather than
just model.
`hf_to_custom_state_dict` drains the weight iterator into one dict before a
single parameter is placed, and every value in that dict is a zero-copy view
into a mmap'd safetensors shard: `weight_utils.py:177-180` opens each shard with
`safe_open` on device `cpu` and hands out `f.get_tensor` views. The dict is
therefore holding the file, not 639 small objects. As the loop touches each
tensor the pages it faults in stay referenced for the rest of the load and the
kernel cannot reclaim them, while the copy allocates the same bytes again on the
device.
Drain the dict instead of iterating it. The loop reads each entry once and
`custom_param_sd` is not used after it, so nothing needs it to stay whole. `pop`
rather than a trailing `del` because the loop continues in two places and the
release has to happen on every path.
Measured on a DGX Spark, GB10, 121 GiB unified memory, loading the transformer
of noctuashap/MiniMax-H3-pruned-r16: 20.17 B parameters, 639 tensors, 37.6 GiB.
Logging `torch.cuda.memory_allocated`, `VmRSS` and `MemAvailable` every ten
parameters:
before 20.8 GiB of parameters placed cost 41.9 GiB of MemAvailable, 2.01x,
and the worker was killed at about 380 of 639
after 37.1 GiB of parameters placed cost 39.5 GiB of MemAvailable, 1.06x,
and the load completed
`VmRSS` shows it directly. Before it climbs monotonically with the device
allocation and never comes down. After it stays between 61.4 and 64.5 GiB for
the whole load, dropping back each time a shard's mappings are released.
Load time is unchanged: 0.38 s per parameter before, 0.37 s after. The bytes are
still read off disk, only the retention goes away.
Not device specific and not model specific. On a discrete GPU the duplicate is
host RAM rather than one shared pool, so the win is a halved host-side peak
instead of the run completing.
fastvideo/tests/loader/test_fsdp_load_releases_checkpoint.py: 3 passed, and
reverting this change fails two of them.
SolitaryThinker
force-pushed
the
release-checkpoint-tensors-during-load
branch
from
August 26, 2026 19:11
93bcee6 to
54c7652
Compare
Collaborator
|
/merge |
Collaborator
|
/test full |
SolitaryThinker
approved these changes
Aug 26, 2026
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
load_model_from_full_model_state_dictretained every source tensor until theentire model copy completed. For MiniMax-H3 this kept the checkpoint storage and
the progressively materialized DiT resident together, exceeding the practical
working set of a single 121 GiB GB10.
This change drains the converted state dictionary as parameters are consumed.
It is generic loader behavior: discrete-GPU systems reduce host-memory pressure,
while unified-memory systems reduce pressure on the shared physical pool.
Changes
popeach source tensor beforecopying it into the model.
distribution, strict/non-strict handling, and unused-key accounting.
dictionary and prove that entries are released during the load, including
continue paths and mixed-dtype placement.
backed by safetensors mappings or ordinary allocated tensors; the optimization
is valid for both.
Dependency
This PR introduces
fastvideo/tests/loader/. Merged PR #1710 owns the sharedunit-lane collection entry for that directory, so this branch is rebased onto
its merge commit rather than duplicating the CI wiring. Final head:
54c7652b3245bb5add9ee950949f8dd9a57457b8.Verification
The original GB10 measurement remains representative of the motivating working
set: draining the dictionary allowed all 20.17B transformer parameters to load,
where retaining the source set repeatedly exhausted memory partway through the
copy. This PR treats that as memory-lifecycle evidence, not as a performance or
quality baseline.
Checklist