Hacktoberfest 2026 : les issues que les mainteneurs ont marquées pour octobre, ouvertes et accessibles aux débutants. Parcourir les issues Hacktoberfest

[Bug] Using `QuantizedTensor` on CPU causes a `RecursionError`

Ouverte
  • Quelque chose ne va pas ?
#3,294 4 commentaires 0 réactions 0 personnes assignées Voir sur GitHub

Les mainteneurs répondent en général sous 2 jours

Personne n'a encore pris cette issue.

Évaluation

Difficulté
3/5
Temps estimé
1-2 jours
Accessibilité débutants
68/100
Type d'issue
Bug
Clarté
Clairement spécifiée
Activité
Calme
Stack technique
python, pytorch
Domaine
machine-learning

Piste de recherche

Commencez par exécuter la reproduction fournie de NVFP4Quantizer, puis suivez les chemins de dispatch et de déquantification dans pytorch/quantized_tensor.py, pytorch/tensor/nvfp4_tensor.py et pytorch/tensor/storage/nvfp4_tensor_storage.py. Le travail est terminé lorsque l’application de xq_cpu + 0 après le déplacement du tenseur quantifié vers le CPU ne lève plus de RecursionError.

Rédigé par le modèle d'indexation à partir du texte de l'issue.

Description

bug

Describe the bug

Trying to use an QuantizedTensor with certain Pytorch ops, after having moved it to cpu causes a RecursionError.

Steps/Code to reproduce bug

import torch
from transformer_engine.pytorch import NVFP4Quantizer

x = torch.randn(128, 128, device="cuda")
q = NVFP4Quantizer()
xq = q(x)
xq_cpu = xq.cpu()

xq_cpu + 0

will cause:

Traceback (most recent call last):
  File "/workspace/repro.py", line 9, in <module>
    xq_cpu + 0
    ~~~~~~~^~~
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 806, in __torch_dispatch__
    return super().__torch_dispatch__(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 804, in __torch_dispatch__
    args = tree_map(maybe_unwrap, args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1578, in tree_map
    return treespec.unflatten(map(func, *flat_args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1313, in unflatten
    leaves = list(leaves)
             ^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 772, in maybe_unwrap
    return arg.dequantize()
           ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 444, in dequantize
    return _FromNVFP4Func.apply(self, dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/autograd/function.py", line 625, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py", line 58, in forward
    cuda_tensor = tensor.to(device=torch.device("cuda"))
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 806, in __torch_dispatch__
    return super().__torch_dispatch__(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 804, in __torch_dispatch__
    args = tree_map(maybe_unwrap, args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1578, in tree_map
    return treespec.unflatten(map(func, *flat_args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1313, in unflatten
    leaves = list(leaves)
             ^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 772, in maybe_unwrap
    return arg.dequantize()
           ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 445, in dequantize
    return _FromNVFP4Func.forward(None, self, dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py", line 58, in forward
    cuda_tensor = tensor.to(device=torch.device("cuda"))
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 806, in __torch_dispatch__
    return super().__torch_dispatch__(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 804, in __torch_dispatch__
    args = tree_map(maybe_unwrap, args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1578, in tree_map
    return treespec.unflatten(map(func, *flat_args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1313, in unflatten
    leaves = list(leaves)
             ^^^^^^^^^^^^
[...]
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 772, in maybe_unwrap
    return arg.dequantize()
           ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 445, in dequantize
    return _FromNVFP4Func.forward(None, self, dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py", line 58, in forward
    cuda_tensor = tensor.to(device=torch.device("cuda"))
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 806, in __torch_dispatch__
    return super().__torch_dispatch__(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 804, in __torch_dispatch__
    args = tree_map(maybe_unwrap, args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1578, in tree_map
    return treespec.unflatten(map(func, *flat_args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1313, in unflatten
    leaves = list(leaves)
             ^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 772, in maybe_unwrap
    return arg.dequantize()
           ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 445, in dequantize
    return _FromNVFP4Func.forward(None, self, dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py", line 58, in forward
    cuda_tensor = tensor.to(device=torch.device("cuda"))
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/tensor/nvfp4_tensor.py", line 806, in __torch_dispatch__
    return super().__torch_dispatch__(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/quantized_tensor.py", line 804, in __torch_dispatch__
    args = tree_map(maybe_unwrap, args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1576, in tree_map
    leaves, treespec = tree_flatten(tree, is_leaf=is_leaf)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_pytree.py", line 1485, in tree_flatten
    treespec = helper(tree, leaves)
               ^^^^^^^^^^^^^^^^^^^^
RecursionError: maximum recursion depth exceeded

I noticed this when trying to use torch.testing.assert_close on a tensor quantized with NVFP4Quantizer and a custom implementation of NVFP4 quantization.

Expected behavior

xq_cpu + 0 should not cause RecursionError.

Environment overview (please complete the following information)

  • Environment location: [Docker]
  • Method of Transformer Engine install: [build from source].

Environment details

Should be irrelevant.

Device details

Should be irrelevant.

Additional context

n/a.

Langage dominant
Python
Étoiles
3.6k
Forks
862
Merge moyen
4 j 21 h
PR mergées (30 j)
62

Préparer son environnement

Par où commencer

  1. Lisez l'issue en entier, puis le guide de contribution du projet.
  2. Signalez en commentaire que vous la prenez — cela évite que deux personnes fassent le même travail.
  3. Forkez le dépôt et travaillez sur une branche.
  4. Ouvrez une pull request qui référence le numéro de l'issue.

Autres issues de NVIDIA/TransformerEngine

Toutes les issues de NVIDIA/TransformerEngine

Issues similaires

Plus d'issues Python

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.