[Bug] Deferred init issue in `te.pytorch.ops`

Describe the bug

OperationFuser caches Parameter objects during construction at fuser.py:357. On the first forward, it calls reset_recipe_state() and then pre_first_fuser_forward() at fuser.py:461, but never refreshes the cached parameters before passing them into autograd at fuser.py:544.

The following existing ops replace meta Parameters with new CUDA Parameter objects during deferred initialization:

• BasicLinear
• Bias
• LayerNorm
• RMSNorm
• GroupedLinear

I reproduced the problem for all five on an H100 using BF16 deferred initialization. Forward successfully materializes the parameters on CUDA, but backward fails with:

RuntimeError: Function _OperationFuserAutogradFunctionBackward returned an
invalid gradient at index 1 - expected device meta but got cuda:0

Quantized deferred initialization has the same problem: BasicLinear correctly materializes its weight as a CUDA Float8Tensor, but backward still returns a CUDA gradient for the stale meta Parameter. Therefore, reset_recipe_state() must remain before pre_first_fuser_forward().

ops.Linear has an additional issue. It registers top-level weight and bias aliases at linear.py:143. After deferred initialization, the internal basic-op parameters are CUDA tensors, while op.weight and op.bias still reference the old meta tensors. Refreshing the fuser cache fixes the autograd mismatch, but these aliases must also be synchronized.

The current tests do not cover this path: test_linear defaults to device="cuda" at test_fusible_ops.py:1106, and there is no meta-device deferred-init case.

In short, this should be treated as a pre-existing, general OperationFuser deferred-initialization bug. The generic fix is to refresh the flattened basic-op parameter cache immediately after all first-forward materialization hooks, plus separately refresh the public parameter aliases in ops.Linear.

Steps/Code to reproduce bug

Please list minimal steps or code snippet for us to be able to reproduce the bug.

A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.

Expected behavior

A clear and concise description of what you expected to happen.

Environment overview (please complete the following information)

• Environment location: [Bare-metal, Docker, Cloud(specify cloud provider - AWS, Azure, GCP, Collab)]
• Method of Transformer Engine install: [pip install or from source]. Please specify exact commands you used to install.
• If method of install is [Docker], provide docker pull & docker run commands used

Environment details

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:

• OS version
• PyTorch version
• Python version
• Transformer Engine version
• CUDA version
• CUDNN version

Device details

• GPU model

Additional context

Add any other context about the problem here.