returnn 1.20240124.211635__tar.gz → 1.20240202.232607__tar.gz

{returnn-1.20240124.211635/returnn.egg-info → returnn-1.20240202.232607}/PKG-INFO

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

returnn-1.20240202.232607/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240202.232607'
+long_version = '1.20240202.232607+git.d2d5015'

{returnn-1.20240124.211635 → returnn-1.20240202.232607}/returnn/__main__.py

@@ -304,12 +304,21 @@ def init_backend_engine(*, config_opts: Optional[Dict[str, Any]] = None):
             # We actually also do this in init_config, but it might have been updated here.
             BehaviorVersion.set(config.int("behavior_version", None))
 
-    if config.value("PYTORCH_CUDA_ALLOC_CONF", None):
+    if config.value("PYTORCH_CUDA_ALLOC_CONF", None) is not None:
         # Set this very early, before *any* `import torch`, thus also before select_engine.
         # (It would not hurt if this is set for any non-PT engine.)
-        value = config.value("PYTORCH_CUDA_ALLOC_CONF", "")
-        print(f"Set PYTORCH_CUDA_ALLOC_CONF={value!r}.")
-        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = value
+        if config.typed_value("torch_distributed") is not None:
+            # torchrun (torch.distributed.run) will likely already have imported torch in the worker procs
+            # before our code runs here, so it is not safe to assign the env var now.
+            # Otherwise, you will likely get such an error:
+            #   RuntimeError: config[i] == get()->name() INTERNAL ASSERT FAILED
+            #   at "../c10/cuda/CUDACachingAllocator.cpp":1058, please report a bug to PyTorch.
+            #   Allocator backend parsed at runtime != allocator backend parsed at load time
+            print("Warning: PYTORCH_CUDA_ALLOC_CONF does not work with torch_distributed, not used", file=log.v2)
+        else:
+            value = config.value("PYTORCH_CUDA_ALLOC_CONF", "")
+            print(f"Set PYTORCH_CUDA_ALLOC_CONF={value!r}.", file=log.v3)
+            os.environ["PYTORCH_CUDA_ALLOC_CONF"] = value
 
     BackendEngine.select_engine(config=config)
     if BackendEngine.is_tensorflow_selected():

{returnn-1.20240124.211635 → returnn-1.20240202.232607/returnn.egg-info}/PKG-INFO

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

{returnn-1.20240124.211635 → returnn-1.20240202.232607}/tools/torch_avg_checkpoints.py

@@ -67,7 +67,11 @@ def merge_checkpoints(in_ckpts: Sequence[str], out_ckpt: str, extra_state: Optio
     out_state: Dict[str, Any] = {"model": out_model_state, "merged_epochs": [], "merged_steps": []}
     for in_ckpt in in_ckpts:
         print("read ckpt:", in_ckpt)
-        in_state = torch.load(in_ckpt, map_location=torch.device("cpu"), mmap=True)
+        torch_version = tuple(int(s) for s in str(torch.__version__).split(".")[:2])
+        load_kwargs = dict(map_location=torch.device("cpu"), mmap=True)
+        if torch_version < (2, 1):
+            load_kwargs.pop("mmap")  # mmap flag only introduced from 2.1.0 onwards
+        in_state = torch.load(in_ckpt, **load_kwargs)
         assert isinstance(in_state, dict)
 
         assert "model" in in_state and isinstance(in_state["model"], dict)

returnn-1.20240124.211635/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240124.211635'
-long_version = '1.20240124.211635+git.780e425'