[Bug]: Qwen3.5-397B-A17B (bf16) OOMs host RAM on the PyTorch backend; AutoDeploy loads fine on the same node

[KyleShao1016]

System Info

• GPU: 8× NVIDIA H200 (SM90), 2,015 GB host RAM
• TensorRT-LLM: 1.3.0rc16 (verified against the source tree)
• Container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc16
• Model: Qwen3.5-397B-A17B (bf16; hybrid full-attention + Gated DeltaNet)
• Serving: trtllm-serve, TP8

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

trtllm-serve Qwen3.5-397B-A17B --backend pytorch --tp_size 8

Expected behavior

The model loads within available host RAM (it fits in GPU across 8×H200). The AutoDeploy backend loads the same checkpoint on the same node using ~69 GB committed host RAM (mmap-shared).

actual behavior

The PyTorch backend exhausts host RAM during weight load and OOMs the node. Measured peak committed host RAM (the OOM-relevant metric):

Deployment	Modeling path	Loader behavior	Peak committed host	Result
Qwen3.5 AutoDeploy (accelerate, mmap)	—	mmap/shared	69 GB	served (~10 min)
Qwen3.5 PyTorch (qwen3_5/qwen3_next)	accumulates, never frees	~2,013 GB	OOM	never served (54 min)

additional notes

Root cause (read against the 1.3.0rc16 tree):

1. The shared reader HfWeightLoader._load_safetensors_file (tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py) uses safetensors.torch.load_file (no mmap) → each of the 8 TP ranks materializes a private host copy.
2. Per-model load_weights diverges:
   • deepseek_v3 (modeling_deepseekv3.py) drains the host dict as it streams to GPU via .cuda() + weights.mark_consumed(...) → host stays bounded and releases.
   • qwen3_5 / qwen3_next (modeling_qwen3_next.py) does not call .cuda()/mark_consumed
     during load → the host weight dict is never drained → it accumulates → ~8× checkpoint →
     ~2 TB → OOM.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

[nv-guomingz]

@KyleShao1016 I'll take a look on PyTorch flow com issue ASAP.

[KyleShao1016]

@nv-guomingz Thanks! I'd be happy to help investigate this further and contribute to a fix if that would be helpful.

[CodersAcademy006]

Hey @KyleShao1016 looking at the source tree, the fix is a four-line change in tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py.

Root cause: HfWeightLoader._load_weights_in_parallel correctly returns a ConsumableWeightsDict. But Qwen3NextHfWeightMapper.preprocess_weights builds a plain new_weights = {} and returns it, silently dropping the wrapper. _load_weights_impl_v2 guards every mark_consumed call with hasattr(weights, 'mark_consumed'), so with a plain dict that check is always False, nothing is ever freed, and all 8 TP ranks hold the full checkpoint in host RAM simultaneously (~8 x ~250 GB = ~2 TB on this config).

DeepSeek avoids this because it uses a dedicated DeepseekV3WeightLoader that takes a ConsumableWeightsDict directly and calls .cuda() + mark_consumed() inline. The qwen3_next path goes through the generic _load_weights_impl_v2 which has the right plumbing but never gets a consumable dict to work with.

The fix mirrors the identical pattern already used in rename_by_params_map in BaseWeightMapper:

--- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py
+++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py
@@ -1,6 +1,8 @@
 import torch
 from torch import nn
 
+from tensorrt_llm._torch.models.checkpoints.base_weight_loader import \
+    ConsumableWeightsDict
 from tensorrt_llm._torch.models.checkpoints.hf.qwen2_moe_weight_mapper import \
     Qwen2MoeHfWeightMapper
 from tensorrt_llm._torch.models.modeling_utils import register_mapper
@@ -107,4 +109,7 @@ class Qwen3NextHfWeightMapper(Qwen2MoeHfWeightMapper):
             else:
                 new_weights[key] = weights[name]
 
-        return new_weights
+        # Preserve ConsumableWeightsDict so _load_weights_impl_v2 can call
+        # mark_consumed() after each module load and free host tensors
+        # incrementally, matching the behaviour of rename_by_params_map.
+        if isinstance(weights, ConsumableWeightsDict):
+            return ConsumableWeightsDict(new_weights)
+        return new_weights

Happy to open a PR with this if useful to @nv-guomingz.

Please correct me if I am wrong. Thank You

[KyleShao1016]

Hi @nv-guomingz — I opened PR #15435 with a proposed fix for this.

Root cause (a correction to the original diagnosis): it's not the non-mmap load_file / missing mark_consumed path. The weights actually reach the GPU fine (meta-init, ~108 GB/rank) and the source dict is mmap-backed (~0.9 GB committed), so nothing accumulates on the host. The real issue is CPU thread oversubscription during weight load:

• run_concurrently uses an unbounded thread pool (CPython default min(32, cpu+4)), and each .contiguous() on the fused MoE expert tensors runs with full-core OpenMP intra-op parallelism.
• On a 192-core node with TP8 that's roughly 8 ranks × 32 workers × ~192 threads fighting over 192 cores → it deterministically wedges at the first large MoE layer (module 887/1749).
• The ~1.8 TB host-RAM growth is a symptom of the stalled load (in-flight scratch buffers piling up), not retained weights.

py-spy on the hung ranks showed every worker spinning in qwen3_5_weight_mapper.handle_special_instance_module → transpose(-1,-2).contiguous(), with top at ~2000%+ CPU/rank.

The fix bounds both the pool size and the per-op intra-op thread count to each rank's fair share of cores (cores // local_ranks), so local_ranks × pool × intra_op ≈ total cores. It honors an explicit OMP_NUM_THREADS and restores the thread count afterward. With it, Qwen3.5-397B-A17B (bf16) loads and serves on 8×H200 with no env-var workarounds (it previously wedged at 51%).

Where I'd value your input: the parameter choice is intentionally conservative — a fixed cap of 4 workers/rank plus the cores // local_ranks split. I'm not sure that's optimal across models/hardware: it ignores per-op scratch size, and cores // local_ranks can under-count when a launcher already pins per-rank affinity. If you have a preferred convention for sizing this (or think it should be memory-aware, or a config knob), I'm glad to rework it.

Happy to iterate however is most useful — thanks!