Add support for Intel Gaudi hpu accelerators (#566)

* Add support for Intel Gaudi hpu accelerators * Fixing the `find_spec` dep. * Fixing unused import. --------- Co-authored-by: Nicolas Patry <[email protected]>
huggingface · Feb 4, 2025 · ee109c6 · ee109c6 · github-actions · Feb 4, 2025
1 parent ea1a2d0
commit ee109c6
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 1 deletion.
diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
@@ -267,7 +267,8 @@ enum Device {
     Xpu(usize),
     Xla(usize),
     Mlu(usize),
-    /// User didn't specify acceletor, torch
+    Hpu,
+    /// User didn't specify accelerator, torch
     /// is responsible for choosing.
     Anonymous(usize),
 }
@@ -296,6 +297,7 @@ impl<'source> FromPyObject<'source> for Device {
                 "xpu" => Ok(Device::Xpu(0)),
                 "xla" => Ok(Device::Xla(0)),
                 "mlu" => Ok(Device::Mlu(0)),
+                "hpu" => Ok(Device::Hpu),
                 name if name.starts_with("cuda:") => parse_device(name).map(Device::Cuda),
                 name if name.starts_with("npu:") => parse_device(name).map(Device::Npu),
                 name if name.starts_with("xpu:") => parse_device(name).map(Device::Xpu),
@@ -327,6 +329,7 @@ impl<'py> IntoPyObject<'py> for Device {
             Device::Xpu(n) => format!("xpu:{n}").into_pyobject(py).map(|x| x.into_any()),
             Device::Xla(n) => format!("xla:{n}").into_pyobject(py).map(|x| x.into_any()),
             Device::Mlu(n) => format!("mlu:{n}").into_pyobject(py).map(|x| x.into_any()),
+            Device::Hpu => "hpu".into_pyobject(py).map(|x| x.into_any()),
             Device::Anonymous(n) => n.into_pyobject(py).map(|x| x.into_any()),
         }
     }

diff --git a/bindings/python/tests/test_pt_comparison.py b/bindings/python/tests/test_pt_comparison.py
@@ -170,6 +170,24 @@ def test_npu(self):
         for k, v in reloaded.items():
             self.assertTrue(torch.allclose(data[k], reloaded[k]))
 
+    def test_hpu(self):
+        # must be run to load torch with Intel Gaudi bindings
+        try:
+            import habana_frameworks.torch.core as htcore
+        except ImportError:
+            self.skipTest("HPU is not available")
+
+        data = {
+            "test1": torch.zeros((2, 2), dtype=torch.float32).to("hpu"),
+            "test2": torch.zeros((2, 2), dtype=torch.float16).to("hpu"),
+        }
+        local = "./tests/data/out_safe_pt_mmap_small_hpu.safetensors"
+        save_file(data, local)
+
+        reloaded = load_file(local, device="hpu")
+        for k, v in reloaded.items():
+            self.assertTrue(torch.allclose(data[k], reloaded[k]))
+
     @unittest.skipIf(not torch.cuda.is_available(), "Cuda is not available")
     def test_anonymous_accelerator(self):
         data = {