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:
2026-08-27 13:39:38 +08:00
parent 1c04a0b9fa
commit bbb2d95256
2 changed files with 6 additions and 6 deletions
+2 -4
View File
@@ -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