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

[tutorial] fixed pipeline bug for sequence parallel (#1943)

Browse Source
This commit is contained in:
Frank Lee
2022-11-14 18:06:57 +08:00
committed by GitHub
parent e52f9d9109
commit c6ea65011f
1 changed files with 12 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
examples/tutorial/sequence_parallel/train.py
View File

@@ -35,6 +35,17 @@ def parse_args():
return parser.parse_args() return parser.parse_args()
def pipeline_data_process_func(stage_output, micro_batch_data):
tokens, types, sentence_order, loss_mask, lm_labels, padding_mask = micro_batch_data
if gpc.is_first_rank(ParallelMode.PIPELINE):
data = (tokens, padding_mask, types, lm_labels)
label = (loss_mask, sentence_order)
else:
data = (stage_output, padding_mask, types, lm_labels)
label = (loss_mask, sentence_order)
return data, label
def main(): def main():
# initialize # initialize
args = parse_args() args = parse_args()
@@ -155,6 +166,7 @@ def main():
if use_pipeline: if use_pipeline:
train_data_iter = SequenceParallelDataIterator(trainloader) train_data_iter = SequenceParallelDataIterator(trainloader)
valid_data_iter = SequenceParallelDataIterator(validloader) valid_data_iter = SequenceParallelDataIterator(validloader)
engine.schedule.data_process_func = pipeline_data_process_func
logger.info("start training") logger.info("start training")