Fix and enhance dumper (#8725)

This commit is contained in:
fzyzcjy
2025-09-05 19:51:09 +08:00
committed by GitHub
parent 8e85ee887e
commit 9f00ec44eb
3 changed files with 192 additions and 50 deletions
+11 -3
View File
@@ -53,7 +53,7 @@ class _Dumper:
if self._partial_name is None:
self._partial_name = _get_partial_name()
rank = dist.get_rank()
rank = _get_rank()
full_kwargs = dict(
forward_pass_id=self._forward_pass_id,
rank=rank,
@@ -80,12 +80,20 @@ class _Dumper:
def _get_partial_name():
rank = dist.get_rank()
rank = _get_rank()
object_list = [str(time.time()) if rank == 0 else None]
dist.broadcast_object_list(object_list, device="cuda")
if dist.is_initialized():
dist.broadcast_object_list(object_list, device="cuda")
return object_list[0]
def _get_rank():
if dist.is_initialized():
return dist.get_rank()
else:
return 0
def get_truncated_value(value):
if value is None:
return None