Check dist.is_initialized

thanh-lam · thanh-lam · commit 57fd859bb3ee · 2026-03-20T11:08:29.000-04:00
diff --git a/aiu_fms_testing_utils/scripts/drive_paged_programs.py b/aiu_fms_testing_utils/scripts/drive_paged_programs.py
@@ -1397,8 +1397,10 @@ def main() -> None:
     distributed_kwargs: Dict[str, Any] = _get_distributed_kwargs(
         is_distributed=args.distributed, dist_timeout=args.dist_timeout
     )
+    # Only rank0 does the "save" ==> Set bool=1. All other ranks get bool=0.
+    # In case of single AIU ==> Only rank0 exists ==> Set bool=1 for rank0.
     args.save_validation_info_outputs = args.save_validation_info_outputs and (
-        dist.get_rank() == 0
+        dist.get_rank() == 0 if dist.is_initialized() else 1
     )
     model_config: DPPRunnerConfig = DPPRunnerConfig()
     world_size = (

Original file line number	Diff line number	Diff line change
`@@ -1397,8 +1397,10 @@ def main() -> None:`
`1397`	`1397`	`distributed_kwargs: Dict[str, Any] = _get_distributed_kwargs(`
`1398`	`1398`	`is_distributed=args.distributed, dist_timeout=args.dist_timeout`
`1399`	`1399`	`)`
	`1400`	`+ # Only rank0 does the "save" ==> Set bool=1. All other ranks get bool=0.`
	`1401`	`+ # In case of single AIU ==> Only rank0 exists ==> Set bool=1 for rank0.`
`1400`	`1402`	`args.save_validation_info_outputs = args.save_validation_info_outputs and (`
`1401`		`- dist.get_rank() == 0`
	`1403`	`+ dist.get_rank() == 0 if dist.is_initialized() else 1`
`1402`	`1404`	`)`
`1403`	`1405`	`model_config: DPPRunnerConfig = DPPRunnerConfig()`
`1404`	`1406`	`world_size = (`