github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] polish colossalai/auto_parallel/tensor_shard/deprecated/op_handler/strategy_generator.py code style (#2695) 

Co-authored-by: shenggan <csg19971016@gmail.com>
pull/2708/head
Liu Ziming 2023-02-14 21:30:25 +08:00 committed by GitHub
parent 534f68c83c
commit 6427c406cf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
colossalai/auto_parallel/tensor_shard/deprecated/op_handler/strategy_generator.py
View File

@ -1,6 +1,7 @@
from dataclasses import dataclass
from abc import ABC, abstractmethod
from typing import List, Dict
from dataclasses import dataclass
from typing import Dict, List
from colossalai.device.device_mesh import DeviceMesh
__all__ = ['IntermediateStrategy', 'StrategyGenerator']
@ -9,7 +10,7 @@ __all__ = ['IntermediateStrategy', 'StrategyGenerator']
@dataclass
class IntermediateStrategy:
"""
IntermediateStrategy contains the subset of meta information for ShardingStrategy. It is
IntermediateStrategy contains the subset of meta information for ShardingStrategy. It is
to store the essential information regarding the tensor sharding and leave other meta information to OperatorHandler.
Args:
@ -24,7 +25,7 @@ class IntermediateStrategy:
class StrategyGenerator(ABC):
"""
StrategyGenerator is used to generate the same group of sharding strategies.
StrategyGenerator is used to generate the same group of sharding strategies.
"""
def __init__(self, device_mesh: DeviceMesh):
@ -39,7 +40,7 @@ class StrategyGenerator(ABC):
@abstractmethod
def validate(self, *args, **kwargs) -> bool:
"""
Validate if the operands are of desired shape.
Validate if the operands are of desired shape.
If True, means this generator can be used for the current operation.
"""
pass