Hacktoberfest 2026: những issue maintainer đã đánh dấu cho tháng Mười, đang mở và phù hợp người mới. Xem issue Hacktoberfest

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

Đang mở
  • Có gì không ổn?
#3,294 4 bình luận 0 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 2 ngày

Chưa có ai nhận issue này.

Đánh giá

Độ khó
3/5
Thời gian dự kiến
1-2 ngày
Mức phù hợp với người mới
68/100
Loại issue
Lỗi
Độ rõ ràng
Đặc tả rõ ràng
Mức độ hoạt động
Ít trao đổi
Công nghệ
python, pytorch
Lĩnh vực
machine-learning

Hướng nghiên cứu

Bắt đầu bằng cách chạy bản tái hiện NVFP4Quantizer được cung cấp, sau đó theo dõi các đường dispatch và dequantization trong pytorch/quantized_tensor.py, pytorch/tensor/nvfp4_tensor.py và pytorch/tensor/storage/nvfp4_tensor_storage.py. Công việc được hoàn tất khi việc áp dụng xq_cpu + 0 sau khi chuyển tensor đã lượng tử hóa sang CPU không còn gây ra RecursionError.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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.

Ngôn ngữ chính
Python
Star
3.6k
Fork
862
Merge trung bình
4 ngày 21 giờ
Pull request đã merge (30 ngày)
62

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của NVIDIA/TransformerEngine

Tất cả issue của NVIDIA/TransformerEngine

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.