[perf]: Load H3 DiT weights onto CUDA, not CPU

GB10 unified memory made the CPU-then-copy path take minutes for
the 35B FastH3 transformer. When host offload is off, read
safetensors onto the accelerator.
This commit is contained in:
Aryan Kumar
2026-08-30 23:27:24 -07:00
committed by Satyam Srivastava
parent a3d9725d15
commit cc9bfd069d
3 changed files with 21 additions and 1 deletions
@@ -161,6 +161,12 @@ is power-cycled. To avoid it:
on: "CPU" offload uses the same unified RAM. Multi-GPU FSDP sharding remains
available because it partitions weights without parking them in a separate
host pool.
- **MiniMax H3 / FastH3** still needs sequential loading on one GB10. The Qwen3-VL
conditioner is tens of gigabytes of BF16. If the DiT and VAEs load while that
encoder is still resident, the process is a typical `earlyoom` kill (Python is
preferred). The CUDA pipeline now encodes first, releases the encoder, then
loads DiT and VAEs onto the accelerator (`to_cpu` follows `cpu_offload`, which
is off here). See [Offloading](../../inference/offloading.md).
## Gotchas specific to the GB10
+10
View File
@@ -22,6 +22,16 @@ pool there, so offload adds transfers and duplicate residency instead of freeing
memory. CUDA FSDP sharding remains enabled when requested; MPS continues to
disable FSDP. `pin_cpu_memory` is not an offload mode and is left unchanged.
MiniMax H3 CUDA inference uses a second lever that does not copy weights to a
host pool. The pipeline loads the Qwen3-VL text encoder, runs conditioning, then
releases that encoder before it loads the DiT and video/audio VAEs. The MLX FastH3
runtime uses the same phase order. When host offload is off, DiT safetensors are
read onto the accelerator instead of CPU-then-copy. Input-preparation geometry
(spatial ratio, latent channels, audio sample rate) comes from the VAE arch
configs until those weights load. A later `generate()` on the same worker
currently re-enters conditioning after the encoder has been released; start a
new generator for a new prompt until prompt-cache reload exists.
## Behavior Explanation
!!! note
+5 -1
View File
@@ -239,7 +239,11 @@ def maybe_load_fsdp_model(
fsdp_shard_conditions=model._fsdp_shard_conditions,
pin_cpu_memory=pin_cpu_memory)
weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=True)
# Host offload is already disabled on unified memory (GB10). Staging the
# 35B FastH3 DiT on CPU and then copying to CUDA doubled that working set
# and took minutes. Follow cpu_offload: read onto the accelerator.
weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=cpu_offload)
logger.info("Loading transformer weights with to_cpu=%s", cpu_offload)
param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
dense_lora_patch = DenseLoRAPatch.from_adapter(
lora_path,