|
|
|
@ -1,18 +1,21 @@
|
|
|
|
|
import builtins
|
|
|
|
|
import math
|
|
|
|
|
import operator
|
|
|
|
|
from copy import deepcopy
|
|
|
|
|
from typing import Dict, List
|
|
|
|
|
|
|
|
|
|
import torch
|
|
|
|
|
from torch.fx import Graph, Node
|
|
|
|
|
from colossalai.tensor.sharding_spec import ShardingSpec
|
|
|
|
|
|
|
|
|
|
from colossalai.device.device_mesh import DeviceMesh
|
|
|
|
|
from colossalai.tensor.shape_consistency import ShapeConsistencyManager
|
|
|
|
|
from colossalai.tensor.sharding_spec import ShardingSpec
|
|
|
|
|
|
|
|
|
|
from ._utils import generate_resharding_costs, generate_sharding_spec
|
|
|
|
|
from .constants import *
|
|
|
|
|
from .op_handler import *
|
|
|
|
|
from .options import SolverOptions
|
|
|
|
|
from .sharding_strategy import ShardingStrategy, StrategiesVector
|
|
|
|
|
from .op_handler import *
|
|
|
|
|
from .constants import *
|
|
|
|
|
from copy import deepcopy
|
|
|
|
|
import math
|
|
|
|
|
import torch
|
|
|
|
|
import operator
|
|
|
|
|
from typing import Dict, List
|
|
|
|
|
from ._utils import generate_sharding_spec, generate_resharding_costs
|
|
|
|
|
import builtins
|
|
|
|
|
|
|
|
|
|
__all__ = ['StrategiesConstructor']
|
|
|
|
|
|
|
|
|
|