test(distributed): adapt to new torch_neuronx initialization

huggingface · Jan 2, 2025 · fa0d0b9 · fa0d0b9
1 parent a2039d4
commit fa0d0b9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/distributed_utils.py b/tests/distributed_utils.py
@@ -207,8 +207,8 @@ def _dist_run(self, local_rank, num_procs, master_port, tp_size, pp_size):
                     raise RuntimeError("self.torchelastic_run_id was not set, it is needed to run a distributed test.")
                 os.environ["TORCHELASTIC_RUN_ID"] = self.torchelastic_run_id
 
-                # Now that the environment has been set, we can configure the PJRT environment.
-                torch_neuronx.xla.configure_pjrt_environment()
+                # Now that the environment has been set, we can initialize the XLA environment.
+                torch_neuronx.initialization.initialize()
 
             if self.init_distributed:
                 dist.init_process_group(backend=self.backend, rank=local_rank, world_size=num_procs)