Jean-Zay specific imports

LLL-Orleans · Sep 26, 2024 · 4f06e98 · 4f06e98
1 parent 3aa2bf9
commit 4f06e98
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/fairseq/distributed/utils.py b/fairseq/distributed/utils.py
@@ -38,6 +38,15 @@
 
 logger = logging.getLogger(__name__)
 
+try:
+    import idr_torch
+    logger.info("Loaded idr_torch sucessfully")
+    if idr_torch.rank == 0:
+        print(">>> Training on ", len(idr_torch.hostname), " nodes and ", idr_torch.size, " processes")
+
+except ImportError:
+    logger.warning("Couldn't load idr_torch")
+
 
 def is_master(cfg: DistributedTrainingConfig):
     return cfg.distributed_rank == 0