Fix the strong-reference fallback in make_weakref

[huthvincent]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Makes make_weakref's except RuntimeError handler do what its comment says — keep a strong reference — instead of raising.

Two defects in that handler, both on lines the refactor in #5451 left behind:

1. The warning interpolates arg.dtype, but the enclosing function is def make_weakref(ten, inplace=True) and no arg exists in that scope or at module scope, so the handler raises NameError on rank 0.
2. wr is bound only by wr = make_weak_ref(ten) inside the try. The handler assigns nothing and falls through to return wr, so on any rank other than 0 — where the get_rank() == 0 guard skips the warning — it raises UnboundLocalError.

Either way the fallback never happens: instead of accepting a memory overhead, the capture path dies. Before #5451 hoisted this body out of replace_with_weak_ref(arg), the same handler ended in return arg.

This is the suggestion posted during review of #5451, applied verbatim: #5451 (comment). It was not picked up before that PR merged, and both names are still broken on main and on dev.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required.
• Small updates (bug fixes, minor improvements): a linked issue is recommended.

Linked issue: none — 2-line bug fix. The diagnosis is already on record in #5451 (comment linked above).

Contribution process

Pre-checks

• I have added relevant unit tests — nothing in tests/ references make_weakref, and the handler needs transformer_engine's make_weak_ref to raise, which is awkward to arrange in a unit test. Happy to add one if you would like it.
• I have added relevant functional tests
• I have added proper typing to my code — the touched lines carry no annotations.
• I have added relevant documentation
• I have run the autoformatter.sh on my PR — black --skip-magic-trailing-comma --skip-string-normalization, isort and ruff check all clean on the changed file.

How to reproduce, and how we checked it

make_weakref returns early unless HAVE_TE_GRAPHS and torch.is_tensor(ten) and ten.is_from_global_mempool, so reproducing needs TE's make_weak_ref to raise for a tensor that already passed that gate. The handler's own comment names the trigger: "There is a known bug where some dtypes (e.g. torch.float64) are not mapped to a representation in transformer_engine/pytorch/utils.py." One supported configuration puts a float64 tensor on that boundary — --cuda-graph-scope moe_router with moe_router_dtype=fp64, whose router tensor is stamped is_from_global_mempool = True at cuda_graphs.py:1436 and then passed to make_weakref.

We checked the handler itself on 8xL4 / torch 2.12.0+cu130 by setting HAVE_TE_GRAPHS = True and replacing make_weak_ref with a function that raises RuntimeError:

• rank 0: NameError: name 'arg' is not defined
• rank 1: UnboundLocalError: cannot access local variable 'wr' where it is not associated with a value
• with this patch: returns the input tensor, no exception

What we did not do: we did not observe TE's real make_weak_ref raising — transformer_engine is not installed on the machine we tested on, so the gate was simulated in memory and only the handler was exercised. No performance claim is made anywhere in this PR and nothing was benchmarked.

The strongest objection to this change

This branch only executes if TE's make_weak_ref raises for a tensor already marked is_from_global_mempool, and no CI job exercises it — so this is a latent crash inside an error handler, and we cannot show anyone has hit it. Two things we deliberately did not fold in, both real and both separate arguments: the bare torch.distributed.get_rank() on line 480 itself raises when no default process group exists, and this module already imports log_single_rank; and except RuntimeError also catches a failure of ten.data = wr on line 473, which is not the dtype bug the comment names.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[huthvincent]

@NVIDIA/mcore-oncall — same ask as on #6859: this is Ready but copy-pr-bot is still waiting on validation, so it has not been built. Could someone run /ok to test 7cf6cad when convenient?

One file, +2 −1, restoring the strong-reference fallback in make_weakref that the comment above it already promises. Happy to rebase first if you would rather have it on a newer main.