| # Adapted from https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/distributed/communication_op.py | |
| import torch | |
| import torch.distributed as dist | |
| def broadcast_item(item, src: int = 0): | |
| if not dist.is_initialized(): | |
| return item | |
| item_list = [item] | |
| dist.broadcast_object_list(item_list, src=src) | |
| return item_list[0] | |
| def all_ranks_finite(value: torch.Tensor) -> bool: | |
| """Make the same finite-value decision on every rank before backward/step.""" | |
| from torch.distributed.tensor import DTensor | |
| if isinstance(value, DTensor): | |
| value = value.to_local() | |
| finite = torch.isfinite(value.detach()).all().to(dtype=torch.int32) | |
| if dist.is_initialized(): | |
| dist.all_reduce(finite, op=dist.ReduceOp.MIN) | |
| return bool(finite.item()) | |