|
|
from typing import List, Callable, Optional, Union, TypeVar, Dict, Any, cast |
|
|
import torch.distributed as dist |
|
|
from .api import ( |
|
|
CheckpointException, |
|
|
_wrap_exception, |
|
|
_is_wrapped_exception, |
|
|
WRAPPED_EXCEPTION |
|
|
) |
|
|
|
|
|
import torch |
|
|
|
|
|
from torch.distributed._shard.sharded_tensor import ( |
|
|
ShardedTensor, |
|
|
) |
|
|
|
|
|
from torch.distributed._shard.sharded_tensor.shard import Shard |
|
|
|
|
|
from .metadata import ( |
|
|
STATE_DICT_TYPE, |
|
|
MetadataIndex, |
|
|
) |
|
|
|
|
|
|
|
|
T = TypeVar('T') |
|
|
R = TypeVar('R') |
|
|
|
|
|
def _get_failure_dict(results: List[Union[T, WRAPPED_EXCEPTION]]) -> Dict[int, WRAPPED_EXCEPTION]: |
|
|
return cast(Dict[int, WRAPPED_EXCEPTION], {i: err for i, err in enumerate(results) if _is_wrapped_exception(err)}) |
|
|
|
|
|
class _DistWrapper: |
|
|
""" |
|
|
This is a wrapper around PG that provides a series of features around object collectives. |
|
|
|
|
|
It works without distributed initialized, where most collectives turns into nops. |
|
|
|
|
|
All variants that take functions are exception robust, meaning that if one or more |
|
|
ranks raise errors, all ranks will observe those. |
|
|
""" |
|
|
def __init__(self, group: Optional[dist.ProcessGroup], use_dist: bool, coordinator_rank: int): |
|
|
self.group = group |
|
|
self.use_dist = use_dist |
|
|
self.coordinator_rank = coordinator_rank |
|
|
if self.use_dist: |
|
|
self.rank = dist.get_rank(group) |
|
|
self.is_coordinator = self.rank == coordinator_rank |
|
|
else: |
|
|
self.rank = 0 |
|
|
self.is_coordinator = True |
|
|
|
|
|
def get_rank(self) -> int: |
|
|
return self.rank |
|
|
|
|
|
def get_world_size(self) -> int: |
|
|
if self.use_dist: |
|
|
return dist.get_world_size(self.group) |
|
|
return 1 |
|
|
|
|
|
def broadcast_object(self, object: Optional[T]) -> T: |
|
|
""" |
|
|
Same as c10d::broadcast_object_list but works without distributed enabled. |
|
|
""" |
|
|
object_list = [object] |
|
|
if self.use_dist: |
|
|
dist.broadcast_object_list( |
|
|
object_list=object_list, |
|
|
group=self.group, |
|
|
src=self.coordinator_rank) |
|
|
return cast(T, object_list[0]) |
|
|
|
|
|
def gather_object(self, object: T) -> Optional[List[T]]: |
|
|
""" |
|
|
Same as c10d::gather_object but works without distributed enabled. |
|
|
""" |
|
|
if self.use_dist: |
|
|
gather_objs = cast(List[T], [None] * dist.get_world_size(self.group)) if self.is_coordinator else None |
|
|
|
|
|
dist.gather_object( |
|
|
obj=object, |
|
|
object_gather_list=gather_objs if self.is_coordinator else None, |
|
|
dst=self.coordinator_rank, |
|
|
group=self.group |
|
|
) |
|
|
result = gather_objs |
|
|
else: |
|
|
result = [object] |
|
|
return result |
|
|
|
|
|
def all_gather_object(self, object: T) -> List[T]: |
|
|
""" |
|
|
Same as c10d::all_gather_object but works without distributed enabled. |
|
|
""" |
|
|
if self.use_dist: |
|
|
gather_objs = cast(List[T], [None] * dist.get_world_size(self.group)) |
|
|
|
|
|
dist.all_gather_object( |
|
|
object_list=gather_objs, |
|
|
obj=object, |
|
|
group=self.group |
|
|
) |
|
|
else: |
|
|
gather_objs = [object] |
|
|
return gather_objs |
|
|
|
|
|
def scatter_object(self, object_list: Optional[List[T]]) -> T: |
|
|
""" |
|
|
Same as c10d::scatter_object but works without distributed enabled. |
|
|
""" |
|
|
if self.use_dist: |
|
|
gather_result = cast(List[T], [None]) |
|
|
dist.scatter_object_list( |
|
|
scatter_object_output_list=gather_result, |
|
|
scatter_object_input_list=object_list if self.is_coordinator else None, |
|
|
src=self.coordinator_rank, |
|
|
group=self.group |
|
|
) |
|
|
|
|
|
local_reply = gather_result[0] |
|
|
else: |
|
|
assert object_list is not None |
|
|
local_reply = object_list[0] |
|
|
return local_reply |
|
|
|
|
|
def reduce_scatter( |
|
|
self, |
|
|
step: str, |
|
|
map_fun: Callable[[], T], |
|
|
reduce_fun: Callable[[List[T]], List[R]] |
|
|
) -> R: |
|
|
""" |
|
|
Compute a value on each rank, then do centralized reduce on a single rank, followed by a scatter. |
|
|
|
|
|
This method operates in the following way: |
|
|
Run ``map_fun`` on all ranks |
|
|
Gather results on rank 0 |
|
|
Call ``reduce_fun`` on all those values |
|
|
Scatter to each rank part of the result. |
|
|
""" |
|
|
local_data: Union[WRAPPED_EXCEPTION, T] |
|
|
try: |
|
|
local_data = map_fun() |
|
|
except BaseException as e: |
|
|
local_data = _wrap_exception(e) |
|
|
|
|
|
all_data = self.gather_object(local_data) |
|
|
all_results: Optional[List[Union[R, CheckpointException]]] = None |
|
|
if self.is_coordinator: |
|
|
assert all_data is not None |
|
|
node_failures = _get_failure_dict(all_data) |
|
|
|
|
|
if len(node_failures) == 0: |
|
|
try: |
|
|
|
|
|
all_results = cast(List[Union[R, CheckpointException]], reduce_fun(cast(List[T], all_data))) |
|
|
except BaseException as e: |
|
|
node_failures[self.rank] = _wrap_exception(e) |
|
|
|
|
|
if len(node_failures) > 0: |
|
|
all_results = [CheckpointException(step, node_failures)] * self.get_world_size() |
|
|
|
|
|
result = self.scatter_object(all_results) |
|
|
if isinstance(result, CheckpointException): |
|
|
raise result |
|
|
return result |
|
|
|
|
|
def all_reduce( |
|
|
self, |
|
|
step: str, |
|
|
map_fun: Callable[[], T], |
|
|
reduce_fun: Callable[[List[T]], R] |
|
|
) -> R: |
|
|
""" |
|
|
Compute a value on each rank, then do centralized reduce on a single rank, followed by a broadcast. |
|
|
|
|
|
This method operates in the following way: |
|
|
Run ``map_fun`` on all ranks |
|
|
Gather results on rank 0 |
|
|
Call ``reduce_fun`` on all those values |
|
|
Broadcast the reduced value to all ranks. |
|
|
""" |
|
|
local_data: Union[T, WRAPPED_EXCEPTION] |
|
|
try: |
|
|
local_data = map_fun() |
|
|
except BaseException as e: |
|
|
local_data = _wrap_exception(e) |
|
|
|
|
|
all_data = self.gather_object(local_data) |
|
|
result: Optional[Union[R, CheckpointException]] = None |
|
|
if self.is_coordinator: |
|
|
assert all_data is not None |
|
|
node_failures = _get_failure_dict(all_data) |
|
|
if len(node_failures) == 0: |
|
|
try: |
|
|
result = reduce_fun(cast(List[T], all_data)) |
|
|
except BaseException as e: |
|
|
node_failures[self.rank] = _wrap_exception(e) |
|
|
|
|
|
if len(node_failures) > 0: |
|
|
result = CheckpointException(step, node_failures) |
|
|
|
|
|
final_result = self.broadcast_object(result) |
|
|
if isinstance(final_result, CheckpointException): |
|
|
raise final_result |
|
|
return cast(R, final_result) |
|
|
|
|
|
def all_gather( |
|
|
self, |
|
|
step: str, |
|
|
map_fun: Callable[[], T], |
|
|
) -> List[T]: |
|
|
""" |
|
|
Compute a value on each rank, then all_gather them. |
|
|
|
|
|
This method operates in the following way: |
|
|
Run ``map_cp`` on all ranks |
|
|
all_gather the values to all ranks |
|
|
""" |
|
|
result: Union[T, WRAPPED_EXCEPTION] |
|
|
try: |
|
|
result = map_fun() |
|
|
except BaseException as e: |
|
|
result = _wrap_exception(e) |
|
|
|
|
|
all_results = self.all_gather_object(result) |
|
|
|
|
|
node_failures = _get_failure_dict(all_results) |
|
|
if len(node_failures) > 0: |
|
|
raise CheckpointException(step, node_failures) |
|
|
return cast(List[T], all_results) |
|
|
|
|
|
def broadcast( |
|
|
self, |
|
|
step: str, |
|
|
map_fun: Callable[[], T], |
|
|
) -> T: |
|
|
""" |
|
|
Compute a value on rank 0 and broadcast it. |
|
|
|
|
|
This method operates in the following way: |
|
|
Run ``map_cp`` on rank 0 |
|
|
broadcast the value |
|
|
""" |
|
|
result: Optional[Union[T, CheckpointException]] = None |
|
|
if self.is_coordinator: |
|
|
try: |
|
|
result = map_fun() |
|
|
except BaseException as e: |
|
|
result = CheckpointException(step, {self.rank: _wrap_exception(e)}) |
|
|
final_result = self.broadcast_object(result) |
|
|
if isinstance(final_result, CheckpointException): |
|
|
raise final_result |
|
|
return cast(T, final_result) |
|
|
|
|
|
def _find_shard(tensor: ShardedTensor, index: MetadataIndex) -> Shard: |
|
|
if index.offset is None: |
|
|
raise ValueError(f"Cannot lookup {index.fqn} since its a ShardedTensor and no offset was provided") |
|
|
|
|
|
shards = tensor.local_shards() |
|
|
|
|
|
if index.index is not None: |
|
|
if len(shards) > index.index and torch.Size(shards[index.index].metadata.shard_offsets) == index.offset: |
|
|
return shards[index.index] |
|
|
|
|
|
for shard in shards: |
|
|
if torch.Size(shard.metadata.shard_offsets) == index.offset: |
|
|
return shard |
|
|
raise ValueError(f"Could not find shard at '{index.offset}' for FQN: '{index.fqn}'") |
|
|
|
|
|
def find_tensor_shard(tensor: torch.Tensor, index: MetadataIndex) -> torch.Tensor: |
|
|
if isinstance(tensor, ShardedTensor): |
|
|
return _find_shard(tensor, index).tensor |
|
|
if index.offset is not None: |
|
|
|
|
|
if index.offset == torch.Size([0] * len(tensor.size())): |
|
|
return tensor |
|
|
raise ValueError(f"FQN: '{index.fqn}' is not a ShardedTensor, can't find by offset: '{index.offset}'") |
|
|
return tensor |
|
|
|
|
|
def find_state_dict_object(state_dict: STATE_DICT_TYPE, index: MetadataIndex) -> Any: |
|
|
if index.fqn not in state_dict: |
|
|
raise ValueError(f"Could not find FQN: '{index.fqn}'") |
|
|
obj = state_dict[index.fqn] |
|
|
if isinstance(obj, torch.Tensor): |
|
|
return find_tensor_shard(obj, index) |
|
|
elif index.offset is not None: |
|
|
raise ValueError(f"FQN: '{index.fqn}' is not a ShardedTensor, can't find by offset: '{index.offset}'") |
|
|
return obj |
|
|
|
