212 lines
7.1 KiB
Python
212 lines
7.1 KiB
Python
#!/usr/bin/env python3
|
|
|
|
|
|
"""
|
|
A set of primitive functions for performing collective ops.
|
|
|
|
Each should also handle single rank scenario.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any, Callable, cast, Generic, List, Optional, Tuple, TypeVar, Union
|
|
|
|
import torch.distributed as dist
|
|
|
|
T = TypeVar("T")
|
|
|
|
@dataclass
|
|
class SyncPayload(Generic[T]):
|
|
stage_name: Optional[str]
|
|
success: bool
|
|
payload: T
|
|
exception: Optional[Exception] = None
|
|
|
|
def broadcast(
|
|
data_or_fn: Union[T, Callable[[], T]],
|
|
*,
|
|
success: bool = True,
|
|
stage_name: Optional[str] = None,
|
|
rank: int = 0,
|
|
pg: Optional[dist.ProcessGroup] = None,
|
|
) -> T:
|
|
"""
|
|
Broadcasts the data payload from rank 0 to all other ranks.
|
|
Or if a function is passed, execute it in rank 0 and broadcast result to all other ranks.
|
|
|
|
Can be used to broadcast a failure signal to stop all ranks.
|
|
|
|
If the function raises an exception, all ranks will raise.
|
|
|
|
Args:
|
|
data_or_fn: the data to broadcast or function to execute and broadcast result.
|
|
success: False to stop all ranks.
|
|
stage_name: the name of the logical stage for synchronization and debugging
|
|
rank: rank to broadcast data or execute function and broadcast resutls.
|
|
pg: the process group for sync
|
|
Throws:
|
|
RuntimeError from original exception trace
|
|
Returns:
|
|
the value after synchronization
|
|
|
|
Example usage:
|
|
>> id = broadcast(data_or_fn=allocate_id, rank=0, pg=ext_pg.my_pg)
|
|
"""
|
|
|
|
if not success and data_or_fn is not None:
|
|
raise AssertionError("Data or Function is expected to be None if not successful")
|
|
|
|
payload: Optional[T] = None
|
|
exception : Optional[Exception] = None
|
|
# if no pg is passed then execute if rank is 0
|
|
if (pg is None and rank == 0) or (pg is not None and pg.rank() == rank):
|
|
# determine if it is an executable function or data payload only
|
|
if callable(data_or_fn):
|
|
try:
|
|
payload = data_or_fn()
|
|
except Exception as e:
|
|
success = False
|
|
exception = e
|
|
else:
|
|
payload = data_or_fn
|
|
|
|
# broadcast the exception type if any to all ranks for failure categorization
|
|
sync_obj = SyncPayload(
|
|
stage_name=stage_name,
|
|
success=success,
|
|
payload=payload,
|
|
exception=exception,
|
|
)
|
|
|
|
if pg is not None:
|
|
broadcast_list = [sync_obj]
|
|
dist.broadcast_object_list(broadcast_list, src=rank, group=pg)
|
|
assert len(broadcast_list) == 1
|
|
sync_obj = broadcast_list[0]
|
|
|
|
# failure in any rank will trigger a throw in every rank.
|
|
if not sync_obj.success:
|
|
error_msg = f"Rank {rank} failed"
|
|
if stage_name is not None:
|
|
error_msg += f": stage {sync_obj.stage_name}"
|
|
if sync_obj.exception is not None:
|
|
error_msg += f": exception {sync_obj.exception}"
|
|
raise RuntimeError(error_msg) from sync_obj.exception
|
|
|
|
return cast(T, sync_obj.payload)
|
|
|
|
|
|
def all_gather(
|
|
data_or_fn: Union[T, Callable[[], T]],
|
|
stage_name: Optional[str] = None,
|
|
pg: Optional[dist.ProcessGroup] = None,
|
|
) -> List[T]:
|
|
"""
|
|
A simple all_gather primitive with basic synchronization guard logic,
|
|
by checking payload from all ranks has the same stage name.
|
|
|
|
Args:
|
|
data_or_fn: the data to be all gathered across ranks or function to be executed
|
|
stage_name: the sync stage name for out-of-sync protection
|
|
pg: the process group for sync
|
|
Throws:
|
|
RuntimeError from original exception trace
|
|
Returns:
|
|
a list of synced data from all ranks
|
|
|
|
Example usage:
|
|
>> all_ids = all_gather(data_or_fn=allocate_id, pg=ext_pg.my_pg)
|
|
"""
|
|
payload: Optional[T] = None
|
|
exception : Optional[Exception] = None
|
|
success = True
|
|
# determine if it is an executable function or data payload only
|
|
if callable(data_or_fn):
|
|
try:
|
|
payload = data_or_fn()
|
|
except Exception as e:
|
|
success = False
|
|
exception = e
|
|
else:
|
|
payload = data_or_fn
|
|
|
|
sync_obj = SyncPayload(
|
|
stage_name=stage_name,
|
|
success=success,
|
|
payload=payload,
|
|
exception=exception,
|
|
)
|
|
|
|
if pg is not None:
|
|
# List of success/failure across all ranks.
|
|
total_list = [None] * dist.get_world_size(pg)
|
|
all_gather_object_enforce_type(pg, total_list, sync_obj)
|
|
# Each rank will throw RuntimeError in case of failure on any rank.
|
|
stage_name = cast(SyncPayload[T], total_list[0]).stage_name
|
|
exception_list: List[Tuple[int, Exception]] = []
|
|
ret_list: List[T] = []
|
|
error_msg: str = ""
|
|
|
|
for i, sp in enumerate(cast(List[SyncPayload[T]], total_list)):
|
|
if sp.stage_name != stage_name:
|
|
error_msg += (
|
|
f"Unexpected stage name received from rank {i}: {sp.stage_name} "
|
|
)
|
|
continue
|
|
if not sp.success and sp.exception is not None:
|
|
exception_list.append((i, sp.exception))
|
|
continue
|
|
ret_list.append(sp.payload)
|
|
|
|
if len(exception_list) > 0:
|
|
raise RuntimeError( # type: ignore[misc]
|
|
error_msg, exception_list) from exception_list[0]
|
|
return ret_list
|
|
else:
|
|
if not sync_obj.success:
|
|
raise RuntimeError(
|
|
f"all_gather failed with exception {sync_obj.exception}",
|
|
) from sync_obj.exception
|
|
return [sync_obj.payload] # type: ignore[list-item]
|
|
|
|
|
|
# Note: use Any for typing for now so users can pass in
|
|
# either a list of None or target type placeholders
|
|
# otherwise pyre would complain
|
|
def all_gather_object_enforce_type(
|
|
pg: dist.ProcessGroup,
|
|
# pyre-fixme[2]: Parameter must have a type that does not contain `Any`
|
|
object_list: List[Any],
|
|
# pyre-fixme[2]: Parameter must have a type other than `Any`
|
|
obj: Any,
|
|
# pyre-fixme[2]: Parameter must have a type that does not contain `Any`
|
|
type_checker: Callable[[Any, Any], bool] = lambda x, y: type(x) == type(y),
|
|
) -> None:
|
|
"""
|
|
Similar to plain all_gather_object but with additional type checking
|
|
AFTER gather is done to ensure basic consistency.
|
|
If check does not pass, all ranks will fail with exception.
|
|
|
|
This is generally to prevent conditional logic leading to
|
|
unexpected messages being received. This is considered fatal code error,
|
|
but due to logic stacks this might happen implicitly in practice.
|
|
|
|
The default check does not check sub type (considered different)
|
|
or covariance (considered same) but users can pass in custom checker
|
|
if more complicated check is needed.
|
|
"""
|
|
dist.all_gather_object(object_list, obj, group=pg)
|
|
|
|
# conservative check
|
|
list_len = len(object_list)
|
|
if list_len == 0:
|
|
return
|
|
first_obj = object_list[0]
|
|
for i in range(1, list_len):
|
|
if not type_checker(first_obj, object_list[i]):
|
|
raise TypeError(
|
|
f"Object type at index {i} is {type(object_list[i])}, "
|
|
f"while first object type is {type(first_obj)}"
|
|
)
|