fix: prevent hang in in download_model_weights (#991)

Signed-off-by: Hemil Desai <[email protected]>

Author: hemildesai

nemo_automodel/_transformers/auto_model.py

@@ -21,7 +21,6 @@
 from typing import List, Optional, Union
 
 import torch
-import torch.distributed as dist
 from torch.nn.attention import SDPBackend, sdpa_kernel
 from transformers import (
     AutoConfig,
@@ -38,7 +37,6 @@
 from nemo_automodel import __version__
 from nemo_automodel._transformers.registry import ModelRegistry
 from nemo_automodel.components.distributed.init_utils import (
-    get_local_rank_preinit,
     get_local_world_size_preinit,
     get_world_size_safe,
 )
@@ -219,13 +217,11 @@ def _verify_sdpa_support(model, is_hf_model, cp_size):
 
 
 def _download_model_weights(hf_config, pretrained_model_name_or_path):
-    if (not dist.is_initialized() or get_local_rank_preinit() == 0) and not os.path.isdir(
-        pretrained_model_name_or_path
-    ):
+    if not os.path.isdir(pretrained_model_name_or_path):
         num_nodes = (get_world_size_safe() % get_local_world_size_preinit()) + 1  # 1-indexed
         if num_nodes > 1:
             logging.info(
-                f"""Downloading model weights on {num_nodes} nodes. This incurs high storage usage. 
+                f"""Downloading model weights on {num_nodes} nodes. This incurs high storage usage.
                 It is recommended to download once with `hf download` and pass in the downloaded path to the `pretrained_model_name_or_path` argument."""
             )
         # Import via module reference (vs bound name) so unit tests can patch

tests/unit_tests/_transformers/test_auto_model.py

@@ -149,9 +149,6 @@ def test_from_pretrained_registry_downloads_checkpoint_files_rank0(self):
             patch.object(transformers.AutoModelForCausalLM, "from_pretrained") as mock_hf_loader,
             patch("nemo_automodel._transformers.auto_model._get_resolved_checkpoint_files") as mock_get_files,
             patch("nemo_automodel._transformers.auto_model.os.path.isdir", return_value=False),
-            patch("nemo_automodel._transformers.auto_model.dist.is_initialized", return_value=True),
-            patch("nemo_automodel._transformers.auto_model.dist.get_world_size", return_value=1),
-            patch("nemo_automodel._transformers.auto_model.dist.get_rank", return_value=0),
             patch("nemo_automodel.components.distributed.utils.FirstRankPerNode") as mock_barrier,
         ):
             # Prepare a fake config with architectures and commit hash
@@ -187,9 +184,6 @@ def test_from_pretrained_registry_downloads_when_dist_uninitialized(self):
             patch.object(transformers.AutoModelForCausalLM, "from_pretrained") as mock_hf_loader,
             patch("nemo_automodel._transformers.auto_model._get_resolved_checkpoint_files") as mock_get_files,
             patch("nemo_automodel._transformers.auto_model.os.path.isdir", return_value=False),
-            patch("nemo_automodel._transformers.auto_model.dist.is_initialized", return_value=False),
-            patch("nemo_automodel._transformers.auto_model.dist.get_world_size", return_value=1),
-            patch("nemo_automodel._transformers.auto_model.dist.barrier") as mock_barrier,
         ):
             # Prepare a fake config with architectures and commit hash
             cfg = Mock()
@@ -213,8 +207,6 @@ def test_from_pretrained_registry_downloads_when_dist_uninitialized(self):
             _, kwargs = mock_get_files.call_args
             assert kwargs["pretrained_model_name_or_path"] == "dummy/repo-id"
             assert kwargs["commit_hash"] == "commit456"
-            # No barrier when dist not initialized
-            mock_barrier.assert_not_called()
 
     def test_from_config_happy_path(self):
         """Test the basic from_config functionality works."""