returnn 1.20250220.200053__py3-none-any.whl → 1.20250221.114352__py3-none-any.whl

returnn/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250220.200053
+Version: 1.20250221.114352
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn/_setup_info_generated.py

@@ -1,2 +1,2 @@
-version = '1.20250220.200053'
-long_version = '1.20250220.200053+git.bb5c0aa'
+version = '1.20250221.114352'
+long_version = '1.20250221.114352+git.650b638'

returnn/torch/distributed.py

@@ -3,10 +3,11 @@ torch.distributed utils
 """
 
 from __future__ import annotations
-from typing import Optional, Any, Dict
+import ast
+import logging
 import os
 import socket
-import logging
+from typing import Optional, Any, Dict
 
 import torch
 from torch.nn.parallel import DistributedDataParallel
@@ -23,19 +24,31 @@ class DistributedContext:
     """
 
     def __init__(self, options: Dict[str, Any]):
-        import torch.distributed as dist
-
         self._opts = CollectionReadCheckCovered(options)
 
-        # when no backend is specified, both gloo and nccl backends will be created
-        # the gloo backend will be used for collectives with CPU tensors and
-        # the nccl backend will be used for collectives with CUDA tensors
-        dist.init_process_group(backend=self._opts.get("backend", None))
+        # Subprocesses have issues initializing torch.distributed process groups.
+        #
+        # We therefore pass rank/size information of the process group via an env
+        # variable that is automatically inherited in any created subprocess.
+        env_var_name = "_RETURNN_TORCH_DISTRIBUTED_INIT_INFO"
+        prev_init_info = os.environ.get(env_var_name)
+        if prev_init_info:
+            self.prev_init_info = ast.literal_eval(prev_init_info)
+            self._rank = self.prev_init_info["rank"]
+            self._size = self.prev_init_info["size"]
+        else:
+            import torch.distributed as dist
+
+            # when no backend is specified, both gloo and nccl backends will be created
+            # the gloo backend will be used for collectives with CPU tensors and
+            # the nccl backend will be used for collectives with CUDA tensors
+            dist.init_process_group(backend=self._opts.get("backend", None))
+            self._rank = dist.get_rank()
+            self._size = dist.get_world_size()
+            os.environ[env_var_name] = repr({"rank": self._rank, "size": self._size})
 
         self._local_rank = int(os.environ["LOCAL_RANK"])
         self._local_size = int(os.environ["LOCAL_WORLD_SIZE"])
-        self._rank = dist.get_rank()
-        self._size = dist.get_world_size()
 
         _logger.info(
             "Torch distributed initialized. Hostname %s, pid %i, rank %i / size %i, local rank %s / local size %s."
@@ -123,9 +136,9 @@ _is_set_up = False
 _ctx = None  # type: Optional[DistributedContext]
 
 
-def get_ctx(config=None) -> Optional[DistributedContext]:
+def get_ctx(config: Optional[Config] = None) -> Optional[DistributedContext]:
     """
-    :param Config|None config:
+    :param config:
     :returns: the global context if Torch distributed is enabled, or None otherwise.
       If we did not setup the context yet, it will automatically create it.
     """

returnn/util/basic.py

@@ -3773,9 +3773,9 @@ def should_write_to_disk(config):
     if config.typed_value("torch_distributed") is not None:
         assert BackendEngine.is_torch_selected(), "torch_distributed assumes PyTorch"
 
-        import torch.distributed
+        import returnn.torch.distributed as torch_distributed
 
-        if torch.distributed.get_rank() != 0:
+        if torch_distributed.get_ctx(config).rank() != 0:
             return False
     elif config.is_true("use_horovod"):
         assert BackendEngine.is_tensorflow_selected(), "use_horovod currently assumes TensorFlow"

{returnn-1.20250220.200053.dist-info → returnn-1.20250221.114352.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250220.200053
+Version: 1.20250221.114352
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20250220.200053.dist-info → returnn-1.20250221.114352.dist-info}/RECORD

@@ -1,9 +1,9 @@
-returnn/PKG-INFO,sha256=U58QGiF-75H5Ac8V3JUwKdPkzP3TPwuPkhfzHhpa7Vc,5215
+returnn/PKG-INFO,sha256=VL6JCmrRQYGi4OdodEygiq3hJyM1C-73zvMTSusYpOs,5215
 returnn/__init__.py,sha256=biBtRsM0WZ406vShaeH-9WFoqJ8XwTbn6g0EeFJ7l8E,1012
 returnn/__main__.py,sha256=qBFbuB1yN3adgVM5pXt2-Yq9vorjRNchNPL8kDKx44M,31752
 returnn/__old_mod_loader__.py,sha256=nvsNY-xELdS_IPNkv66Q9Rmvg4dbGW0-EBRDcCmctos,7654
 returnn/__setup__.py,sha256=22kQn2fh11iPM0hLb2Fy5sLmoU1JGvmDxXRYuRgQkwU,4659
-returnn/_setup_info_generated.py,sha256=OojdMrmzo4naqIdlDTwnSiMHtnmVuqlosY9_dqmm20c,77
+returnn/_setup_info_generated.py,sha256=LlW75YDQH_DvPCMrSeF0bY52JGq9l4tJNA5mGTT5MQA,77
 returnn/config.py,sha256=3tmKhB6FnQZaNdtcYsiB61JnEY--iZ2qmJ4yq0b6tE0,29140
 returnn/forward_iface.py,sha256=A_OJiaXsX4MlXQRzST86ylyxSUZbC402PQL1REcqHjM,911
 returnn/learning_rate_control.py,sha256=ZvWryAn_tv9DhV8sh1LV3eE34Yltl3On3mYZAG4hR9s,34684
@@ -206,7 +206,7 @@ returnn/tf/util/ken_lm.py,sha256=R60UAoywriuDIeQ2Hk3Vm_waf2Hxxc88ofzEw6X6Sd4,173
 returnn/tf/util/open_fst.py,sha256=sZRDw4TbxvhGqpGdUJWy1ebvlZm4_RPhygpRw9uLAOQ,11265
 returnn/torch/README.md,sha256=jzJ2FpOHW02vxN69yKaV97C9LI-hmvjBglKfdZXIDdc,85
 returnn/torch/__init__.py,sha256=MHEUyNHB20Vy89uKAqZoj6FxJKF1Gq3HW-i6ra1pNcI,24
-returnn/torch/distributed.py,sha256=i13cUVjI7GxpO0TAresrNyCM0ZBAaf-cXNr09Fmg_2k,6266
+returnn/torch/distributed.py,sha256=skFyutdVztxgTEk3HHJ8S83qRWbNpkNT8Tj16Ic0_hE,6981
 returnn/torch/engine.py,sha256=8BIpdcrpbJL9HrvCX-hISh-14zW9aSrHGvRWT9s0zOk,77103
 returnn/torch/updater.py,sha256=GqtBvZpElPVMm0lq84JPl4NVLFFETZAzAbR0rTomSao,28249
 returnn/torch/data/__init__.py,sha256=6cLNEi8KoGI12PF6akN7mI_mtjlx-0hcQAfMYoExwik,132
@@ -233,7 +233,7 @@ returnn/torch/util/gradient_checkpoint.py,sha256=iLy-FB65DC8O6LxzmMvFjnSdpIVpko8
 returnn/torch/util/module.py,sha256=MXHIrF9Isu575DDJIa81212ULKwdqu1oOLxDVZecVSk,1693
 returnn/torch/util/scaled_gradient.py,sha256=3585VuNypBty-pW6r3BKK047H3MqZQSdMjXeYAb4cmU,3192
 returnn/util/__init__.py,sha256=UIG1qw4idqhW71BV60ha7h9PktxvEVcBIu0lYRossK8,336
-returnn/util/basic.py,sha256=nhCfxWwGL7FchgFW5x9V2OgXD0HtpN885NASdwfeKYg,142339
+returnn/util/basic.py,sha256=__rtDp8crZfm0mEeAKsRxNCdWuBHh9OeOm8UO-X4CJU,142380
 returnn/util/better_exchook.py,sha256=MVMnuu6KoyqgvlMeQLQNTfdspcPR9MwigCXOpeTVqCI,62956
 returnn/util/bpe.py,sha256=LWFhICZsEOnMwNws0lybPNzKRX6rSr8yKCvP65vjl9Y,19656
 returnn/util/debug.py,sha256=wuRzdg9zB84WWCGyTjmRR_zYypu8gXxlc0nZ6si9OC8,28224
@@ -253,8 +253,8 @@ returnn/util/sig_proc.py,sha256=Tjz0VOAVyqu2qDCF5HZ1JjALjcFsHcNkcd96WgZeKfE,7265
 returnn/util/task_system.py,sha256=y4sMVXQ25Qd2z0rx03uOlXlkE-jbCYC1Sjfn-XlraVU,26003
 returnn/util/train_proc_manager.py,sha256=Pjht28k6uz6BNQ47uW6Gf880iyq5q4wx7P_K2tmoAM8,3266
 returnn/util/watch_memory.py,sha256=BR5P2kvBN6UI81cE0_1WAA6Hd1SByLbBaiDxvLhPOew,4213
-returnn-1.20250220.200053.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20250220.200053.dist-info/METADATA,sha256=U58QGiF-75H5Ac8V3JUwKdPkzP3TPwuPkhfzHhpa7Vc,5215
-returnn-1.20250220.200053.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-returnn-1.20250220.200053.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20250220.200053.dist-info/RECORD,,
+returnn-1.20250221.114352.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20250221.114352.dist-info/METADATA,sha256=VL6JCmrRQYGi4OdodEygiq3hJyM1C-73zvMTSusYpOs,5215
+returnn-1.20250221.114352.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+returnn-1.20250221.114352.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20250221.114352.dist-info/RECORD,,