Add support for Cambricon mlu devices

[huismiling]

What does this PR do?

Add support for Cambricon mlu devices.

Transformers and Accelerate have supported cambricon mlu (huggingface/transformers#29627, huggingface/accelerate#2552).

This PR enables users to leverage the cambricon mlu for training and inference with safetensors.

Test code:

import torch
import torch_mlu
from safetensors.torch import save_file
from safetensors import safe_open

test = {"a": torch.tensor([1, 2, 3])}

save_file(test, 'test.safetensors')

with safe_open("test.safetensors", framework="pt", device="mlu") as f:
   for key in f.keys():
      print(f.get_tensor(key).device)

[huismiling]

@Narsil @ArthurZucker
Hi, could you help to merge this PR ?

[kalvdans]

I don't see that the data is transferred into the device itself with this PR, they stay as numpy arrays?

[huismiling]

@kalvdans
Hi, Cambricon MLU is a AI processor that support AI frameworks like PyTorch, TensorFlow, etc. So, It is possible to run Transformers/Accelerate on MLUs to train foundation model. Website: https://www.cambricon.com .

I ran the test code on MLUs. mlu:0 is printed.

[kalvdans]

I ran the code withuot "import torch_mlu" and got the error
"""
RuntimeError: Expected one of cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip, ve, fpga, maia, xla, lazy, vulkan, mps, meta, hpu, mtia, privateuseone device type at start of device string: mlu
"""
on the get_tensor() call.

Where can I get the torch_mlu module from to test it, and does it have side-effects on module load?

[huismiling]

@kalvdans
Hi, torch_mlu is a module designed to work with the MLU cards, making the MLU card as easy to use as a GPU. Unfortunately, due to certain reasons, only some customers are able to obtain the MLU module and cards at present. It is not available for individual developer.

Test code is ran as below.

Python 3.10.14 (main, May  6 2024, 19:42:50) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> import torch_mlu
>>> from safetensors.torch import save_file
>>> from safetensors import safe_open
>>> test = {"a": torch.tensor([1, 2, 3], device="mlu")}
>>> save_file(test, 'test.safetensors')
>>> f = safe_open("test.safetensors", framework="pt", device="mlu")
>>> f.get_tensor("a").device
device(type='mlu', index=0)
>>> torch.mlu.is_available()
True

[kalvdans]

Thanks @huismiling for explaining. I'll leave it to others to decide if they want proprietary untestable code in a public library.

off-topic, but I recommend Cambricon to make the registering with torch explicit, as tools such as "uv check --fix" will remove seemingly unused imports.

[huismiling]

@kalvdans Thanks.

[ArthurZucker]

LGTM not sure if docs need an update as well @Narsil