[perf]: MiniMax H3 on GB10 - skip text encoder CPU offload on unified memory (5m49s to 30ms) - #1710
Conversation
mergify
Bot
added
type: perf
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI
🔴 PR merge requirementsWaiting for
This rule is failing.
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 42bca784f8
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
Verified on the GB10 with the change in place and 30 ms from weights read to module ready, against 5 min 49 s on main with the same flag set. For completeness, the three cases side by side, same machine, same checkpoint:
The middle row is what the table in the description was measured from. The last row is the one that matters here, since it keeps the user's request intact and lets the platform decide whether it can be honoured. |
KyleNeverGivesUp
changed the title
KyleNeverGivesUp
changed the title
…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>
42bca78 to
1d874b6
Compare
fb1c107 to
0fa8b1f
Compare
|
/merge |
On integrated CUDA and MPS devices, host and accelerator allocations share one physical pool. Text-encoder CPU offload therefore frees no memory and can briefly retain both copies; on DGX Spark this made the 63 GB MiniMax H3 encoder spend 5m49s sharding before the process was killed. Apply the decision to the worker-local FastVideoArgs only after that worker binds its selected device. This avoids creating a CUDA context in the parent process and avoids classifying every worker from device 0. Clearing the text_encoder_cpu_offload flag keeps load placement, FSDP, generic text encoding, and MiniMax H3 conditioning consistent. Keep a role-aware loader fallback for direct text-encoder callers. ImageEncoderLoader shares this implementation, so its role remains unchanged here and can opt into the same policy separately. CUDA classification follows cudaDeviceProp integrated, while discrete and unknown platforms preserve existing behavior.
5345aa5 to
450aad5
Compare
|
/test full |
1 similar comment
|
/test full |
|
The full suite run from Why I cannot reproduce them. My box is a single GB10. Both lanes need more, and the training one also refuses to run on this device at all:
What I could check. I ran the unit lane's pytest command on a GB10 for this branch and for its merge base. The failures there are identical between the two, so nothing in the unit lane is caused by these changes. I could not find a second pull request that has run either of the two failing lanes, so there is no comparison point for them. One thing that may matter for the SSIM lane. The build came from What would help. The failing output from those two lanes. If they were already failing before this branch, or if the SSIM lane failed on a test unrelated to the loader and platform changes here, that closes it. If they are real, the most likely candidate in this diff is the |
…fails The Trigger Merge Gate workflow cancels stale Buildkite builds in step 2 and starts the merge-gate build in step 7. A step failure ends the job, so when the cancel step dies the gate build is never started at all. The check then goes red for a reason unrelated to the pull request, and nothing in the log distinguishes "the tests failed" from "the tests never ran". That is what happened on hao-ai-lab#1710: jq: error (at <stdin>:1): Cannot index string with string "env" Process completed with exit code 5 `curl` is called without `--fail-with-body`, so an HTTP error is treated as success and its body is piped onward. Buildkite answers a rate limit or an unauthorized read with an object such as {"message": "Not Found"}, and `.[]` over an object yields its values, so `.env` then runs against a string and jq exits non-zero. Three changes, all to the same step: `continue-on-error: true`, because cancelling stale builds only saves agent time. Failing to cancel wastes an agent; failing to trigger means untested code, and step 7 keeps its hard failure. `--fail-with-body` on the lookup, with the response echoed as a warning, so a lookup that fails says why instead of feeding an error body to the parser. `if type == "array" then .[] else empty end` in the filter, so a response that is not a build list yields no matches rather than aborting. Verified against six response shapes: a matching build, a build for another PR, the error object from hao-ai-lab#1710, an empty array, a build with no env, and a bare string. The first returns the build number and the rest return nothing, where the error object previously exited 5.
2 participants
Purpose
CPU offload assumes host memory and device memory are two different pools. On a
device whose GPU reads host RAM they are one, so offloading frees nothing, and
the way it is implemented makes peak memory worse rather than better.
Part of #1709.
The measurement
DGX Spark (GB10, 121 GB usable), loading MiniMax H3's text encoder, a 63 GB
bf16 checkpoint, memory sampled every 5 s with
free. Baseline 8 GB, machineotherwise idle.
free'susedcolumn excludesbuff/cache, so this is processmemory.
Loading weights tooktoLoaded module text_encoderThe weights themselves arrive in about 17 s in both cases, and 63 GB on disk
becomes roughly 67 GB resident, which is expected. With offload on, memory then
climbs steadily for the next five minutes while
fully_shardwalks the moduletree, and the run is OOM-killed before the DiT is reached.
Nothing is upcast. The checkpoint is uniformly bf16 (
{'BF16': 43}across thefirst shard's header,
"dtype": "bfloat16"in its config), so the extra 34 GBis not dtype conversion.
Why it happens
shard_modelcallsfully_shard(..., offload_policy=CPUOffloadPolicy(...))permodule, bottom up. Each call allocates that module's parameters on the host and
copies them over before the device copy goes away. On a discrete card the two
allocations come from different pools, so the device side empties as the host
side fills. On unified memory both come from the same RAM, so during the walk
the machine holds part of the model twice, and the walk is long enough that the
overlap is substantial.
Changes
Platform.has_unified_memory, defaultFalse, so discrete accelerators keepthe existing path unchanged.
MpsPlatformreturnsTrue.CudaPlatformBasereadscudaDeviceProp::integrated, which torch surfaces asis_integrated. Confirmed on the hardware in question:The probe is guarded: a torch build without the field, or a machine with no
visible device, returns
Falserather than changing offload behaviour onhardware that cannot be classified.
The text encoder loader skips sharding on such a device and logs why, in a
branch next to the existing MPS one.
The
text_encoderoffload flag itself is deliberately left alone. It stateswhat the user wants, and the same script should keep working on both an A100,
where offload is real, and a Spark, where it cannot be. The platform decides
whether the request is achievable, and says so in the log rather than silently
doing nothing.
MPS already skipped sharding, but for an unrelated reason ("not compatible").
That branch stays as it is; this adds a second reason that applies to CUDA
devices with the same memory topology.
Test Plan
pytest fastvideo/tests/platforms/test_unified_memory.py -q pre-commit run --files fastvideo/platforms/interface.py fastvideo/platforms/cuda.py \ fastvideo/platforms/mps.py fastvideo/models/loader/component_loader.py \ fastvideo/tests/platforms/test_unified_memory.pyCPU only, no accelerator needed. The CUDA cases monkeypatch
torch.cuda.get_device_properties, since what needs guarding is the decision,not the driver.
Covered: the base default stays
False; MPS isTrue; CUDA followsis_integratedin both directions; a torch build without the attribute fallsback to
False; a probe that raises does not take the run down.Test Results
Test output
mypy could not run in my checkout, whose directory name contains a hyphen and
so is not a valid package name. It passes in CI.
End to end on the GB10, the numbers in the table above. Reproduce with:
and compare the gap between the two log lines with and without this change.
Scope
This does not make H3 fit on a single GB10. Its four components are 111 GB on
disk against 121 GB of memory, so it still runs out, just without the extra
34 GB and five minutes on top. What this fixes is a cost that every unified
memory device pays on every model that requests text encoder offload, which is
the default in several of the example scripts.
Checklist