fix: fix broadcast_pyobj breaking VerlEngine (#5997)

This commit is contained in:
Junrong Lin
2025-05-05 04:15:53 +08:00
committed by GitHub
parent 95c231e50d
commit 357fb2dba5
2 changed files with 7 additions and 3 deletions

View File

@@ -897,7 +897,10 @@ def broadcast_pyobj(
src: int = 0,
force_cpu_device: bool = True,
):
"""Broadcast inputs from rank=0 to all other ranks with torch.dist backend."""
"""Broadcast inputs from src rank to all other ranks with torch.dist backend.
The `rank` here refer to the source rank on global process group (regardless
of dist_group argument).
"""
device = torch.device(
"cuda" if torch.cuda.is_available() and not force_cpu_device else "cpu"
)