github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-08 04:24:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

[elixir] moved simulator build to op_builder (#3939)

Browse Source
This commit is contained in:
Frank Lee
2023-06-09 16:38:53 +08:00
committed by GitHub
parent 3b58ff5c73
commit 86ff5c152b
5 changed files with 51 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/elixir/search/simulator.py
View File

@@ -1,10 +1,12 @@
import math
from colossalai.kernel.op_builder import ElixirSimulatorBuilder
from .utils import to_divide
def calc_move_times(param_per_step: list, param_to_chunk: dict, n_blocks: int):
from colossalai.elixir.simulator import move_count
simulator = ElixirSimulatorBuilder().load()
chunk_per_step = list()
for param_set in param_per_step:
@@ -17,7 +19,7 @@ def calc_move_times(param_per_step: list, param_to_chunk: dict, n_blocks: int):
if len(id_set) > 0:
chunk_per_step.append(list(id_set))
return move_count(chunk_per_step, n_blocks)
return simulator.move_count(chunk_per_step, n_blocks)
def find_optimal_chunk_size(