github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-02 17:46:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

[autoparallel] added solver option dataclass (#1588)

Browse Source
This commit is contained in:
Frank Lee
2022-09-13 14:47:09 +08:00
committed by GitHub
parent 82d4376c23
commit 219f66c571
6 changed files with 43 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_auto_parallel/test_cost_graph.py
View File

@@ -1,3 +1,4 @@
from pickletools import optimize
import torch
from torch.fx import GraphModule
import torch.nn as nn
@@ -10,6 +11,7 @@ from colossalai.tensor.shape_consistency import ShapeConsistencyManager
from colossalai.device.device_mesh import DeviceMesh
from colossalai.auto_parallel.solver.strategies_constructor import StrategiesConstructor
from colossalai.auto_parallel.solver.cost_graph import CostGraph
from colossalai.auto_parallel.solver.options import SolverOptions
from copy import deepcopy
@@ -52,7 +54,7 @@ def test_cost_graph():
gm = GraphModule(model, graph, model.__class__.__name__)
gm.recompile()
solver_options = {'fast_mode': True}
solver_options = SolverOptions(fast=True)
strategies_constructor = StrategiesConstructor(graph, device_mesh, shape_consistency_manager, solver_options)
strategies_constructor.build_strategies_and_cost()

5
tests/test_auto_parallel/test_solver.py
View File

@@ -11,6 +11,7 @@ from colossalai.auto_parallel.solver.cost_graph import CostGraph
from colossalai.auto_parallel.solver.graph_analysis import GraphAnalyser
from copy import deepcopy
from colossalai.auto_parallel.solver import Solver
from colossalai.auto_parallel.solver.options import SolverOptions
class ConvModel(nn.Module):
@@ -39,7 +40,6 @@ def test_solver():
# [[0, 1]
# [2, 3]]
device_mesh = DeviceMesh(physical_mesh_id, mesh_shape)
entire_shape = torch.Size((4, 16, 64, 64))
shape_consistency_manager = ShapeConsistencyManager()
tracer = ColoTracer()
@@ -57,9 +57,8 @@ def test_solver():
# return relu
graph = tracer.trace(root=model, meta_args=input_sample)
gm = GraphModule(model, graph, model.__class__.__name__)
gm.recompile()
solver_options = {'fast_mode': True}
solver_options = SolverOptions(fast=True)
strategies_constructor = StrategiesConstructor(graph, device_mesh, shape_consistency_manager, solver_options)
strategies_constructor.build_strategies_and_cost()

3
tests/test_auto_parallel/test_strategies_constructor.py
View File

@@ -11,6 +11,7 @@ from colossalai.auto_parallel.solver.sharding_strategy import ShardingStrategy,
from colossalai.tensor.shape_consistency import ShapeConsistencyManager
from colossalai.device.device_mesh import DeviceMesh
from colossalai.auto_parallel.solver.strategies_constructor import StrategiesConstructor
from colossalai.auto_parallel.solver.options import SolverOptions
from copy import deepcopy
@@ -47,7 +48,7 @@ def test_strategies_constructor():
gm = GraphModule(model, graph, model.__class__.__name__)
gm.recompile()
solver_options = {'fast_mode': True}
solver_options = SolverOptions(fast=True)
strategies_constructor = StrategiesConstructor(graph, device_mesh, shape_consistency_manager, solver_options)
assert strategies_constructor.leaf_strategies == []