fix(modeling): allow GPU-less meta-load of remote models that hard-im…

[khaiwang]

Allow GPU-less meta-load of remote models that hard-import CUDA kernels

Problem

Some HuggingFace models ship custom remote modeling code (trust_remote_code) that imports CUDA-only kernel packages at module-import time. For example, nvidia/NVIDIA-Nemotron-3-* (Nemotron-H) does, at the top of modeling_nemotron_h.py:

from mamba_ssm.ops.triton.layernorm_gated import rmsnorm_fn  # unguarded -> hard ImportError

mamba_ssm / causal_conv1d can't be installed or imported without a CUDA toolchain. Since nnsight builds a meta model client-side via AutoModelForCausalLM.from_config(..., trust_remote_code=True), a GPU-less nnsight/NDIF client hit ImportError: mamba-ssm is required ... before any meta tensor was created — so it couldn't construct the module tree, build an intervention graph, or run remote=True at all. This blocks the core "weak client, remote GPU" use case.

Fix

A meta model never runs a forward on the client (the forward happens on the GPU host), so these kernels are only imported, never called.

• meta_kernel_shim() (new modeling/_kernel_shim.py): wraps the meta from_config(...) and registers inert stand-ins for the kernel packages in sys.modules — only when they aren't really installed — then removes them. Each stub member raises if called, so a dispatched/real run can never silently use a stub instead of the real kernel. No-op on a GPU host.
• Consistent trust_remote_code=True across _load_config, _load_meta, and _load. Previously only the meta from_config() forced it, so the client tree could be built from a different implementation than the one dispatched/served, and _load_config used a native config that can't parse newer remote-only fields (the cause of KeyError '-' on newer Nemotron-H hybrid_override_pattern block types).

 src/nnsight/modeling/_kernel_shim.py | 136 +++++++  (new)
 src/nnsight/modeling/transformers.py |  26 ++-

Verification

• GPU-less meta-load now succeeds for Nemotron-3 4B and 120B (was: hard ImportError); gpt2 and other models unchanged.
• Tree equivalence: the meta module tree built without kernels is byte-identical (same module/param/buffer name→shape→dtype signature) to the tree built with real kernels — so a client's interventions map exactly onto the served model.
• End-to-end (real weights, Nemotron-3 4B on GPU): an nnsight trace produces logits byte-identical to plain transformers (argmax " Paris"); the fixed path is byte-identical to the prior CUDA path on baseline and on two interventions (ablation, ×0.5 scaling).
• No regression: gpt2 load is byte-identical with and without the change.

Notes

• A complementary upstream fix would guard the rmsnorm_fn import in NVIDIA's remote code (the adjacent kernel imports are already guarded by is_mamba_2_ssm_available()); this PR makes nnsight resilient without depending on that.
• The 120B forward was not tested (247 GB); the 4B uses the identical load/dispatch code path. The 4B forward ran the model's reference SSM path because the causal_conv1d CUDA kernel was unavailable on the test host (prebuilt wheel needs glibc 2.32, host has 2.31) — unrelated to this change, which does not touch the forward.

🤖 Generated with Claude Code