fix(fsdp): import os for safe_get_rank fallback

Author: fallintoplace

megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py

@@ -16,6 +16,7 @@
 import inspect
 import logging
 import operator
+import os
 from contextlib import nullcontext
 from functools import reduce
 from importlib.metadata import version

tests/unit_tests/distributed/megatron_fsdp/test_mcore_tensor_parallelism_detect.py

@@ -6,6 +6,7 @@
 from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel
 from megatron.core.distributed.fsdp.src.megatron_fsdp.utils import (
     get_mcore_tensor_parallel_partition_dim,
+    safe_get_rank,
     is_mcore_tensor_parallel_duplicated,
     using_tensor_parallel,
 )
@@ -79,6 +80,13 @@ def test_using_tensor_parallel_false_when_mesh_size_one():
     assert using_tensor_parallel(dist_index) is False
 
 
+def test_safe_get_rank_should_fall_back_to_rank_env_if_distributed_is_not_initialized(monkeypatch):
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
+    monkeypatch.setenv("RANK", "7")
+
+    assert safe_get_rank() == 7
+
+
 class DummyConfig:
     # Just enough attributes for __init__ to run if needed in future tests.
     gradient_accumulation_fusion = False