diff --git a/colossalai/builder/pipeline.py b/colossalai/builder/pipeline.py index d6394be6ceb8..a25e8990d10a 100644 --- a/colossalai/builder/pipeline.py +++ b/colossalai/builder/pipeline.py @@ -1,7 +1,6 @@ import copy import heapq - from colossalai.builder import build_model, build_layer from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc @@ -40,6 +39,7 @@ def _binary_partition(weights, st, ed): def _heap_addition(weights, intervals, add_cnt): """ """ + def _heap_push(heap, st, ed): value = weights[ed - 1] if st > 0: @@ -162,7 +162,10 @@ def count_layer_params(layers): return param_counts -def build_pipeline_model_from_cfg(config, num_chunks: int = 1, partition_method: str = 'parameter', verbose: bool = False): +def build_pipeline_model_from_cfg(config, + num_chunks: int = 1, + partition_method: str = 'parameter', + verbose: bool = False): """An initializer to split the model into different stages for pipeline parallelism. An example for the model config is shown below. The class VisionTransformerFromConfig should @@ -218,7 +221,7 @@ def build_pipeline_model_from_cfg(config, num_chunks: int = 1, partition_method: log_str += f'\n===== stage={stage}, layers={num_layers} =====\n' for st, ed in parts[stage]: - for idx, layer in enumerate(layers[st: ed]): + for idx, layer in enumerate(layers[st:ed]): log_str += f'\t{idx + st:2d}: {layer}\n' logger.info(log_str, ranks=[0])