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:
BasicLinearBiasLayerNormRMSNormGroupedLinear
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
Additional context
Add any other context about the problem here.
Describe the bug
OperationFusercaches Parameter objects during construction at fuser.py:357. On the first forward, it callsreset_recipe_state()and thenpre_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:
BasicLinearBiasLayerNormRMSNormGroupedLinearI reproduced the problem for all five on an H100 using BF16 deferred initialization. Forward successfully materializes the parameters on CUDA, but backward fails with:
Quantized deferred initialization has the same problem:
BasicLinearcorrectly materializes its weight as a CUDAFloat8Tensor, but backward still returns a CUDA gradient for the stale meta Parameter. Therefore,reset_recipe_state()must remain beforepre_first_fuser_forward().ops.Linearhas an additional issue. It registers top-levelweightandbiasaliases at linear.py:143. After deferred initialization, the internal basic-op parameters are CUDA tensors, whileop.weightandop.biasstill 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_lineardefaults todevice="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
OperationFuserdeferred-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 inops.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)
docker pull&docker runcommands usedEnvironment details
If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
Device details
Additional context
Add any other context about the problem here.