From d6f6a78a76c8f22ed58bf95d3d97a4e8095a353e Mon Sep 17 00:00:00 2001 From: "hzji210@gmail.com" Date: Wed, 26 Nov 2025 21:50:17 +0800 Subject: [PATCH] Workaround for torch-npu's lack of support for creating nested tensors from NPU tensors. --- verl/__init__.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/verl/__init__.py b/verl/__init__.py index 3fb53c2996d..65ab0854b2c 100644 --- a/verl/__init__.py +++ b/verl/__init__.py @@ -52,6 +52,26 @@ patch_hub() if is_npu_available: + # Workaround for torch-npu's lack of support for creating nested tensors from NPU tensors. + # + # ``` + # >>> a, b = torch.arange(3).npu(), torch.arange(5).npu() + 3 + # >>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged) + # ``` + # throws "not supported in npu" on Ascend NPU. + # See https://github.com/Ascend/pytorch/blob/294cdf5335439b359991cecc042957458a8d38ae/torch_npu/utils/npu_intercept.py#L109 + # for details. + + import torch + + try: + if hasattr(torch.nested.nested_tensor, "__wrapped__"): + torch.nested.nested_tensor = torch.nested.nested_tensor.__wrapped__ + if hasattr(torch.nested.as_nested_tensor, "__wrapped__"): + torch.nested.as_nested_tensor = torch.nested.as_nested_tensor.__wrapped__ + except AttributeError: + pass + from .models.transformers import npu_patch as npu_patch package_name = "transformers"