mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-06-21 13:11:27 +00:00
079bf3cb26
* [misc] update pre-commit * [misc] run pre-commit * [misc] remove useless configuration files * [misc] ignore cuda for clang-format
21 lines
635 B
Python
21 lines
635 B
Python
from abc import ABC, abstractmethod
|
|
from typing import List, Optional
|
|
|
|
import torch.distributed as dist
|
|
|
|
from colossalai.legacy.zero.sharded_param.sharded_tensor import ShardedTensor
|
|
|
|
|
|
class BaseShardStrategy(ABC):
|
|
def __init__(self) -> None:
|
|
"""Abstract Shard Strategy. Use to shard a tensors on multiple GPUs."""
|
|
super().__init__()
|
|
|
|
@abstractmethod
|
|
def shard(self, tensor_list: List[ShardedTensor], process_group: Optional[dist.ProcessGroup] = None):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def gather(self, tensor_list: List[ShardedTensor], process_group: Optional[dist.ProcessGroup] = None):
|
|
pass
|