fix: keep PyTorch-executed weights on the device when the mutable module offloads

[tp5uiuc]

What — Gates deallocate_module() on the user-facing offload_module_to_cpu setting at
the three call sites on ToT that ignores it. Also puts PyTorch-executed submodules back on the target
device after every offload. Adds a test pinning what the setting guarantees for a partially
converted model.

Why — A hybrid graph module does not own its weights. torch.export.export keeps the source
module's Parameter objects, ExportedProgram.module() re-registers those same objects, and both
partitioners share attributes by setattr — so the compiled module and the module it was compiled
from hold identical Parameter instances. A nn.Module.to() rebinds .data on them and so moving one
moves the other. deallocate_module() is a .to('cpu'), so any graph the partitioner did not hand
to TensorRT in full had its PyTorch-executed weights dragged off the device while activations stayed
on it, leading to errors like this

    RuntimeError: Expected all tensors to be on the same device, but got mat1 is on
    cuda:0, different from other tensors on cpu (wrapper_CUDA_addmm)

A fully converted graph does not see this error, because TensorRT submodules keep their weights inside the
engine. Related PRs #3418 added the unconditional deallocate at all four call sites and
#3736 later gated it behind offload_module_to_cpu in compile() only. The other three call sites of update_refit_condition(), refit_gm() and the static load() kept offloading regardless.

How — Do both halves together: gate the three offload sites that ignored the setting, and put
PyTorch-executed submodules back on the target device afterwards. Doing only one of them means the bug is still open. The offload_module_to_cpu test now asserts what the contract actually guarantees for a
partially converted model rather than that the whole source model lands on CPU.

Testing — Reproduced on an L40S (SM 8.9) with torch_executed_ops forcing the fallback and no capability guards involved, so this is neither card- nor backend-specific. Confirmation sweep (T4 SM 7.5 + L40S, driver 595.58.03): 5 runtime/ and 1 models/ failures closed, no status change on the L40S.

Cost / Gotchas — Behaviour change: offload_module_to_cpu=True on a partially converted model
now leaves PyTorch-executed weights on the device. That is what makes the model runnable, but a
caller relying on "everything ends up on CPU" will see a difference.

🤖 Generated with Claude Code

[tp5uiuc]

This is maybe a behavior change, but it is needed to keep the code from breaking with the offload_module_to_cpu flag. It is needed because export share Parameter identity between the original model and the graph module. Having a .to("cpu") on the original model yanks the weights that _run_on_gpu_* needs to execute in PyTorch. Pinning these submodules back is actually the same placement compile_module already does here

TensorRT/py/torch_tensorrt/dynamo/_compiler.py

Lines 1428 to 1442 in 0fad50d

	for name, _ in partitioned_module.named_children():
	submodule = getattr(partitioned_module, name)
	# filter on the GraphModule
	if not isinstance(submodule, torch.fx.graph_module.GraphModule):
	continue
	# Criteria for a module to be convertible to TRT
	if settings.use_fast_partitioner and "_run_on_acc" not in name:
	dryrun_tracker.to_run_in_torch.extend(parse_non_trt_nodes(submodule))
	logger.debug(
	"Submodule in PyTorch: %s\n %s",
	str(name),
	str(submodule.graph),
	)
	submodule.to(to_torch_device(settings.device))
	continue

and so this restores an invariant. Without it we get errors in the MR description.

[tp5uiuc]

CI summary

failing check	cause
standard / test (dynamo-models, standard, l2)	test_cross_runtime_serde::test_save_cpp_load_python and ::test_save_python_load_cpp — Windows only
standard / test (dynamo-models-critical, standard, l1)	test_hf_gqa_model::test_dynamic_head_dim_with_hf_model[False-dtype0] — Tensor-likes are not close
executorch-runtime-test	OSError: libcurand.so.10
build / build-windows-arm64-cross, build / prepare-upload	The operation was canceled — superseded run; the retry of both passed
gate ×3	aggregators

Why these are not from this PR: the same test_cross_runtime_serde pair fails identically on #4643, #4644, #4645, #4646, #4647 and #4648, whose change sets are disjoint — e.g. #4648 changes _exporter.py (the save path) and #4647 changes neither that nor anything else in common, yet both fail the same two tests with byte-identical error text. executorch-runtime-test reproduces on main nightlies and on unrelated PRs (#4638, #4642). The gate entries are aggregators of the above.

This PR changes only _MutableTorchTensorRTModule.py plus two test files; test_hf_gqa_model.py contains no reference to Mutable, offload or deallocate.