Introduce multi-node SPMD initialization for Neuron

pytorch · Sep 19, 2024 · 92be233 · 92be233
1 parent d0ea5cc
commit 92be233
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/torch_xla/runtime.py b/torch_xla/runtime.py
@@ -252,6 +252,14 @@ def use_spmd(auto: Optional[bool] = False):
   if auto:
     torch_xla._XLAC._xla_set_auto_sharding()
     os.environ["XLA_AUTO_SPMD"] = "1"
+
+  if os.environ[xenv.PJRT_DEVICE] == 'NEURON':
+    # In case of Neuron, retrigger the PJRT initialization if possible.
+    try:
+      from torch_neuronx.initializer import initialize
+      initialize()
+    except ImportError:
+      pass
 
 
 def is_spmd():