fix: report true rank in logs via dist-aware helpers
- get_rank/get_world_size fall back to RANK/WORLD_SIZE env instead of hardcoded 0/1, matching torchrun's env-before-init contract - log filter reuses the helpers so initialized groups show real ranks; local-spawn children previously logged rank=0/8
This commit is contained in:
@@ -30,15 +30,13 @@ def get_current_device():
|
||||
def get_world_size() -> int:
|
||||
if dist.is_available() and dist.is_initialized():
|
||||
return dist.get_world_size()
|
||||
else:
|
||||
return 1
|
||||
return int(os.environ.get("WORLD_SIZE", "1"))
|
||||
|
||||
|
||||
def get_rank() -> int:
|
||||
if dist.is_available() and dist.is_initialized():
|
||||
return dist.get_rank()
|
||||
else:
|
||||
return 0
|
||||
return int(os.environ.get("RANK", "0"))
|
||||
|
||||
|
||||
@contextmanager
|
||||
|
||||
Reference in New Issue
Block a user