Repository: huggingface/transformers
Pull request: fix error for _register_pytree_node in torch2.1.0 and  fix bf16 assertion in xpu and npu

# What does this PR do?

fix issue #37838

## fix error for _register_pytree_node

The PR [Apply torchfix to replace deprecated functions: `_pytree._register_pytree_node` and `torch.cpu.amp.autocast`](https://github.com/huggingface/transformers/commit/71b35387fd6d71487bd29e694ed10d925203e031) introduced a compatibility issue when using **Torch 2.1.0**.

Specifically, in **Torch 2.1.0**, the correct import should still be:

```python
from torch.utils._pytree import _register_pytree_node
```

However, the PR incorrectly tries to import:

```python
from torch.utils._pytree import register_pytree_node
```

This leads to the following error when running with **Torch 2.1.0**:

```
ImportError: cannot import name 'register_pytree_node' from 'torch.utils._pytree'
```

## fix bf16 assertion in xpu and npu

The PR [Add XPU case to `is_torch_bf16_gpu_available`](https://github.com/huggingface/transformers/commit/954f31cd818c431312f452c4e10bcbc0bdde42a2) introduced a compatibility issue:

- It **only handled XPU** and **ignored other devices** like NPU.
- It **also triggered an error** when XPU is **not available** in the current PyTorch build.
  - Directly accessing `torch.xpu.is_available()` without checking if `torch.xpu` exists causes an `AttributeError`.

When `torch.xpu` is not present (which is common in standard PyTorch installations), the following code pattern:

```python
if torch.xpu.is_available():
    ...
```

causes:

```
AttributeError: module 'torch' has no attribute 'xpu'
```

We use a safer access pattern:

```python
if hasattr(torch, "xpu") and torch.xpu.is_available():
    ...
```

This ensures compatibility even when `torch.xpu` is not defined.

In addition, I have also added support for NPU devices by checking:

```python
if is_torch_npu_available():
    ...
```

## Before submitting
- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request), Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

## Who can review?

- trainer: @zach-huggingface and @SunMarc

Changed files:
filename: src/transformers/utils/generic.py
patch: @@ -352,9 +352,9 @@ def __init_subclass__(cls) -> None:
                     serialized_type_name=f"{cls.__module__}.{cls.__name__}",
                 )
             else:
-                from torch.utils._pytree import register_pytree_node
+                from torch.utils._pytree import _register_pytree_node
 
-                register_pytree_node(
+                _register_pytree_node(
                     cls,
                     _model_output_flatten,
                     partial(_model_output_unflatten, output_type=cls),

filename: src/transformers/utils/import_utils.py
patch: @@ -533,10 +533,12 @@ def is_torch_bf16_gpu_available() -> bool:
 
     if torch.cuda.is_available():
         return torch.cuda.is_bf16_supported()
-    if torch.xpu.is_available():
+    if is_torch_xpu_available():
         return torch.xpu.is_bf16_supported()
     if is_torch_hpu_available():
         return True
+    if is_torch_npu_available():
+        return torch.npu.is_bf16_supported()
     return False