[Bug]: 142 GiB/rank host memory leak (leading to host OOM if KV cache offloading is enabled) when serving Kimi-K2.6-NVFP4 TP4 on GB300 (aarch64), due to MetaInitException on aten.fill_.Scalar (nn.LayerNorm init) silently falling back to whole-model CPU construction

[Weili-0234]

System Info

• Hardware: GB300 (Grace-Blackwell) node, aarch64, 4× GPU, TP4
• Container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc4 (also reproduces on current main — the relevant code is unchanged)
• torch: 2.10 NGC build (CPU allocator = mimalloc v2.2.4 on aarch64, enabled by Enable mimalloc on non-Windows platforms and make default for AArch64 builds pytorch/pytorch#164741 — note that PR shows as Closed because pytorch merges via pytorchmergebot, and it was briefly reverted then relanded the same day 2025-10-09; current pytorch main still sets USE_MIMALLOC ON for CPU_AARCH64 in CMakeLists.txt, and the NGC build logs mimalloc: v2.2.4 under MIMALLOC_VERBOSE=1)
• Model: nvidia/Kimi-K2.6-NVFP4 (≈554 GB checkpoint; any model whose modules construct stock nn.LayerNorm is affected)
• PyTorch backend (tensorrt_llm._torch), KV cache host offloading enabled (KvCacheConfig.host_cache_size)

Who can help?

Anyone familiar with the PyTorch-backend model loader / MetaInitMode (tensorrt_llm/_torch/models/modeling_utils.py).

Information

• The official example scripts
• My own modified scripts

Tasks

• My own task or dataset (give details below)

Reproduction

Trigger (any platform, CPU-only, a few seconds):

import torch.nn as nn
from tensorrt_llm._torch.models.modeling_utils import MetaInitMode

with MetaInitMode():
    nn.LayerNorm(64)   # raises MetaInitException: Meta tensor used in unsupported function: aten.fill_.Scalar

Deployment-scale consequence (aarch64):

1. Serve a large checkpoint whose model constructs nn.LayerNorm in __init__ (e.g. Kimi-K2.6's vision tower) with the PyTorch backend on an aarch64 node.
2. Engine log shows, once per rank per start (at info level only): Fallback to regular model init: ... MetaInitException: Meta tensor used in unsupported function: aten.fill_.Scalar
3. After weights land on GPU, check grep RssAnon /proc/<rank_pid>/status (RssAnon = anonymous resident set size, the process's actual physical-RAM footprint from heap/allocator arenas): each rank retains a weight-shard-sized private-anonymous region for process lifetime (137-142 GiB/rank for the ≈554 GB TP4 checkpoint; ≈568 GiB/node, 54% of host RAM), invisible to allocator statistics.
4. Allocator discriminator proving the pages are freed-but-retained (no rebuild):

MIMALLOC_VERBOSE=1 MIMALLOC_PURGE_DELAY=0 MIMALLOC_ABANDONED_PAGE_PURGE=1 <launch engine>
# log shows "mimalloc: v2.2.4" (env read); post-init RssAnon collapses to 5-6 GiB/rank

Expected behavior

Meta init survives value-irrelevant in-place initializers (fill_/zero_) that stock torch modules run in __init__; the model constructs on the meta device; no CPU materialization of the weights; host memory after load is a few GiB per rank.

actual behavior

MetaInitMode raises on aten.fill_.Scalar / aten.zero_.default, the loader's broad except Exception silently constructs the ENTIRE model on CPU memory, and on aarch64 the freed CPU copy is retained by mimalloc for process lifetime:

• 142 GiB/rank of dead (freed-but-retained) host memory, ≈568 GiB per node, for the ≈554 GB TP4 checkpoint;
• KvCacheConfig.host_cache_size cannot be raised beyond ≈30 GB/rank without host OOM, on a node that should support 150 GB+/rank;
• load time also regresses on every platform (full CPU init instead of meta skeleton).

Root cause of the retention half is upstream of TensorRT-LLM (mimalloc v2.2.4 arena purge gate; same class as the open pytorch/pytorch#178726, proposed fix pytorch/pytorch#178851) — but TensorRT-LLM turns a benign in-place init into that trap by aborting meta init on ops whose results are never consumed.

additional notes

• A fix PR is filed right after this issue (branch Weili-0234/TensorRT-LLM:fix-metainit-inplace-init-ops, will be auto-linked here): add a deterministic_init_ops pass-through (fill_.Scalar, fill_.Tensor, zero_.default) to MetaInitMode, plus regression tests (tests/unittest/_torch/models/test_meta_init_mode.py). Red/green verified on a GB300 node with the rc4 image: 3 failed + 1 passed before, 4 passed after.
• Searchable log signatures for others hitting this: Fallback to regular model init, MetaInitException: Meta tensor used in unsupported function: aten.fill_.Scalar, unexplained per-rank physical memory (RssAnon) ≈ weight-shard size on Grace (GB200/GB300).
• Suggested follow-up (not in the PR): log the whole-model CPU fallback at warning instead of info — at production log levels the fallback is currently invisible, which is why this took a full allocation-stack investigation to find.