From 4049e002386b3d6237bd152f9c79fff173ba5893 Mon Sep 17 00:00:00 2001 From: Qianran Date: Thu, 16 Feb 2023 12:50:56 +0800 Subject: [PATCH 1/2] [NFC] polish colossalai/engine/schedule/_pipeline_schedule.py code style --- colossalai/engine/schedule/_pipeline_schedule.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/colossalai/engine/schedule/_pipeline_schedule.py b/colossalai/engine/schedule/_pipeline_schedule.py index 97571fa024ba..712ae8242409 100644 --- a/colossalai/engine/schedule/_pipeline_schedule.py +++ b/colossalai/engine/schedule/_pipeline_schedule.py @@ -4,8 +4,9 @@ import inspect from typing import Callable, List, Tuple, Union -import colossalai.communication as comm import torch.cuda + +import colossalai.communication as comm from colossalai.amp.naive_amp import NaiveAMPModel from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc @@ -72,9 +73,9 @@ class PipelineSchedule(BaseSchedule): tensor_shape (torch.Size, optional): Specified shape in pipeline communication. scatter_gather_tensors (bool, optional): If set to `True`, communication will be reduced over pipeline when using 1D tensor parallelization. - + Example: - + # this shows an example of customized data_process_func def data_process_func(stage_output, dataloader_output): output1, output2 = stage_output @@ -157,6 +158,7 @@ def load_micro_batch(self): def pre_processing(self, engine): from colossalai.zero.sharded_model.sharded_model_v2 import ShardedModelV2 + # TODO: remove this after testing new zero with pipeline parallelism model = engine.model if isinstance(model, NaiveAMPModel): @@ -229,7 +231,7 @@ def _get_data_label_for_current_step(self, stage_output, micro_batch_data, crite return data, label def _forward_step(self, engine, input_obj, return_tensors, return_output_label=True, accum_loss=None): - """Forward step for passed-in model. If it is the first stage, the input tensor + """Forward step for passed-in model. If it is the first stage, the input tensor is obtained from data_iterator, otherwise the passed-in input_obj is used. Returns output tensor. This is a helper function and can be ignored by users. @@ -266,7 +268,7 @@ def _forward_step(self, engine, input_obj, return_tensors, return_output_label=T return output_obj def _backward_step(self, engine, input_obj, output_obj, output_obj_grad): - """Backward step through the passed-in output tensor. If it is the last stage, the + """Backward step through the passed-in output tensor. If it is the last stage, the output_obj_grad is None, otherwise it is the gradients with respect to stage's output tensor. Returns the gradients with respect to the input tensor (None if first stage). This is a helper function and can be ignored by users. @@ -511,7 +513,7 @@ def _forward_step(self, return_tensors, return_output_label=True, accum_loss=None): - """Forward step for passed-in model. If it is the first stage, the input tensor + """Forward step for passed-in model. If it is the first stage, the input tensor is obtained from data_iterator, otherwise the passed-in input_obj is used. Returns output tensor. This is a helper function and can be ignored by users. From ac59bc35719efd97e5ac4cfc88bf408f2e7190ff Mon Sep 17 00:00:00 2001 From: Qianran Ma Date: Wed, 29 Mar 2023 17:43:02 +0800 Subject: [PATCH 2/2] [NFC] polish colossalai/fx/tracer/_tracer_utils.py code style --- colossalai/fx/tracer/_tracer_utils.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/colossalai/fx/tracer/_tracer_utils.py b/colossalai/fx/tracer/_tracer_utils.py index 0ec49a90a133..e160497a7444 100644 --- a/colossalai/fx/tracer/_tracer_utils.py +++ b/colossalai/fx/tracer/_tracer_utils.py @@ -1,6 +1,8 @@ -from typing import List, Union, Any -from ..proxy import ColoProxy, ColoAttribute +from typing import Any, List, Union + import torch + +from ..proxy import ColoAttribute, ColoProxy from .meta_patch import meta_patched_function, meta_patched_module __all__ = ['is_element_in_list', 'extract_meta']