[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:
committed by
Satyam Srivastava
parent
a3d9725d15
commit
cc9bfd069d
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user