2021-10-28 16:21:23 +00:00
|
|
|
#!/usr/bin/env python
|
|
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
|
2021-12-30 07:56:46 +00:00
|
|
|
from typing import List, Tuple, Union, Callable
|
|
|
|
import inspect
|
2021-10-28 16:21:23 +00:00
|
|
|
import torch.cuda
|
|
|
|
from torch import Tensor
|
|
|
|
|
2022-01-07 05:22:22 +00:00
|
|
|
import colossalai.communication as comm
|
2021-10-28 16:21:23 +00:00
|
|
|
from colossalai.context.parallel_mode import ParallelMode
|
|
|
|
from colossalai.core import global_context as gpc
|
Develop/experiments (#59)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
* Split conv2d, class token, positional embedding in 2d, Fix random number in ddp
Fix convergence in cifar10, Imagenet1000
* Integrate 1d tensor parallel in Colossal-AI (#39)
* fixed 1D and 2D convergence (#38)
* optimized 2D operations
* fixed 1D ViT convergence problem
* Feature/ddp (#49)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* support torch ddp
* fix loss accumulation
* add log for ddp
* change seed
* modify timing hook
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* Feature/pipeline (#40)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* optimize communication of pipeline parallel
* fix grad clip for pipeline
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* optimized 3d layer to fix slow computation ; tested imagenet performance with 3d; reworked lr_scheduler config definition; fixed launch args; fixed some printing issues; simplified apis of 3d layers (#51)
* Update 2.5d layer code to get a similar accuracy on imagenet-1k dataset
* update api for better usability (#58)
update api for better usability
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
2021-12-09 07:08:29 +00:00
|
|
|
from colossalai.amp.naive_amp import NaiveAMPModel
|
2021-12-30 07:56:46 +00:00
|
|
|
from colossalai.utils.cuda import get_current_device
|
Develop/experiments (#59)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
* Split conv2d, class token, positional embedding in 2d, Fix random number in ddp
Fix convergence in cifar10, Imagenet1000
* Integrate 1d tensor parallel in Colossal-AI (#39)
* fixed 1D and 2D convergence (#38)
* optimized 2D operations
* fixed 1D ViT convergence problem
* Feature/ddp (#49)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* support torch ddp
* fix loss accumulation
* add log for ddp
* change seed
* modify timing hook
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* Feature/pipeline (#40)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* optimize communication of pipeline parallel
* fix grad clip for pipeline
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* optimized 3d layer to fix slow computation ; tested imagenet performance with 3d; reworked lr_scheduler config definition; fixed launch args; fixed some printing issues; simplified apis of 3d layers (#51)
* Update 2.5d layer code to get a similar accuracy on imagenet-1k dataset
* update api for better usability (#58)
update api for better usability
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
2021-12-09 07:08:29 +00:00
|
|
|
from colossalai.zero import (ZeroRedundancyOptimizer_Level_2,
|
|
|
|
ZeroRedundancyOptimizer_Level_3)
|
2021-12-30 07:56:46 +00:00
|
|
|
from colossalai.utils import switch_virtual_pipeline_parallel_rank
|
2021-10-28 16:21:23 +00:00
|
|
|
from ._base_schedule import BaseSchedule
|
|
|
|
|
|
|
|
|
|
|
|
def squeeze(x: Union[Tensor, tuple, list]):
|
|
|
|
if isinstance(x, (tuple, list)):
|
|
|
|
return x[0]
|
|
|
|
else:
|
|
|
|
return x
|
|
|
|
|
|
|
|
|
|
|
|
class PipelineSchedule(BaseSchedule):
|
|
|
|
"""A helper schedule class for pipeline parallelism running environment.
|
|
|
|
It uses non-interleaved 1F1B strategy. Other properties are similar as
|
Develop/experiments (#59)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
* Split conv2d, class token, positional embedding in 2d, Fix random number in ddp
Fix convergence in cifar10, Imagenet1000
* Integrate 1d tensor parallel in Colossal-AI (#39)
* fixed 1D and 2D convergence (#38)
* optimized 2D operations
* fixed 1D ViT convergence problem
* Feature/ddp (#49)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* support torch ddp
* fix loss accumulation
* add log for ddp
* change seed
* modify timing hook
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* Feature/pipeline (#40)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* optimize communication of pipeline parallel
* fix grad clip for pipeline
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* optimized 3d layer to fix slow computation ; tested imagenet performance with 3d; reworked lr_scheduler config definition; fixed launch args; fixed some printing issues; simplified apis of 3d layers (#51)
* Update 2.5d layer code to get a similar accuracy on imagenet-1k dataset
* update api for better usability (#58)
update api for better usability
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
2021-12-09 07:08:29 +00:00
|
|
|
:class:`NonPipelineSchedule`.
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
:param num_microbatches: The number of microbatches
|
|
|
|
:type num_microbatches: int
|
2021-12-30 07:56:46 +00:00
|
|
|
:param batch_data_process_func: The preprocessing function which receives a batch of data, and it will be executed in `load_batch`
|
|
|
|
:type batch_data_process_func: Callable
|
2022-01-07 05:22:22 +00:00
|
|
|
:param scatter_gather_tensors: If set to `True`, communication will be reduced over pipeline when using 1D tensor parallelization
|
|
|
|
:type scatter_gather_tensors: bool
|
2021-10-28 16:21:23 +00:00
|
|
|
"""
|
|
|
|
|
|
|
|
def __init__(self,
|
|
|
|
num_microbatches,
|
2021-12-30 07:56:46 +00:00
|
|
|
batch_data_process_func: Callable = None,
|
2022-01-07 05:22:22 +00:00
|
|
|
tensor_shape: Union[torch.Size, List[int], Tuple[int]] = None,
|
|
|
|
scatter_gather_tensors: bool = False):
|
2021-12-30 07:56:46 +00:00
|
|
|
super().__init__(batch_data_process_func=batch_data_process_func)
|
2021-10-28 16:21:23 +00:00
|
|
|
self.num_microbatches = num_microbatches
|
2021-12-20 15:26:19 +00:00
|
|
|
self.dtype = torch.float
|
2021-12-30 07:56:46 +00:00
|
|
|
self.tensor_shape = tensor_shape
|
2022-01-07 05:22:22 +00:00
|
|
|
self.scatter_gather_tensors = False
|
|
|
|
if gpc.is_initialized(ParallelMode.PARALLEL_1D) and gpc.get_world_size(ParallelMode.PARALLEL_1D) > 1:
|
|
|
|
self.scatter_gather_tensors = scatter_gather_tensors
|
2021-10-28 16:21:23 +00:00
|
|
|
|
2021-11-18 11:45:06 +00:00
|
|
|
def load_batch(self, data_iter):
|
2021-12-30 07:56:46 +00:00
|
|
|
# Pipeline schedule just puts data in memory
|
|
|
|
self.batch_data, self.batch_label = super().load_batch(data_iter, to_gpu=False)
|
|
|
|
self.microbatch_offset = 0
|
2022-01-04 12:52:31 +00:00
|
|
|
if isinstance(self.batch_data, torch.Tensor):
|
|
|
|
batch_size = self.batch_data.size(0)
|
|
|
|
else:
|
|
|
|
batch_size = next(iter(self.batch_data.values())).size(0)
|
|
|
|
assert batch_size % self.num_microbatches == 0, \
|
2021-10-28 16:21:23 +00:00
|
|
|
"Batch size should divided by the number of microbatches"
|
2022-01-04 12:52:31 +00:00
|
|
|
self.microbatch_size = batch_size // self.num_microbatches
|
2021-10-28 16:21:23 +00:00
|
|
|
|
2021-12-30 07:56:46 +00:00
|
|
|
def _get_data_slice(self, data, offset):
|
|
|
|
if isinstance(data, torch.Tensor):
|
|
|
|
return data[offset: offset + self.microbatch_size]
|
|
|
|
else:
|
|
|
|
return {k: v[offset:offset + self.microbatch_size] for k, v in data.items()}
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
def load_micro_batch(self):
|
2021-12-30 07:56:46 +00:00
|
|
|
data = self._get_data_slice(self.batch_data, self.microbatch_offset)
|
|
|
|
label = self._get_data_slice(self.batch_label, self.microbatch_offset)
|
|
|
|
self.microbatch_offset += self.microbatch_size
|
|
|
|
return self._move_to_device(data), self._move_to_device(label)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
Develop/experiments (#59)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
* Split conv2d, class token, positional embedding in 2d, Fix random number in ddp
Fix convergence in cifar10, Imagenet1000
* Integrate 1d tensor parallel in Colossal-AI (#39)
* fixed 1D and 2D convergence (#38)
* optimized 2D operations
* fixed 1D ViT convergence problem
* Feature/ddp (#49)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* support torch ddp
* fix loss accumulation
* add log for ddp
* change seed
* modify timing hook
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* Feature/pipeline (#40)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* optimize communication of pipeline parallel
* fix grad clip for pipeline
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* optimized 3d layer to fix slow computation ; tested imagenet performance with 3d; reworked lr_scheduler config definition; fixed launch args; fixed some printing issues; simplified apis of 3d layers (#51)
* Update 2.5d layer code to get a similar accuracy on imagenet-1k dataset
* update api for better usability (#58)
update api for better usability
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
2021-12-09 07:08:29 +00:00
|
|
|
def pre_processing(self, engine):
|
|
|
|
if isinstance(engine.optimizer, (ZeroRedundancyOptimizer_Level_2, ZeroRedundancyOptimizer_Level_3)):
|
2021-10-28 16:21:23 +00:00
|
|
|
raise TypeError(
|
|
|
|
"Pipeline schedule is currently not compatible with ZeRO Level 2 and Level 3"
|
|
|
|
)
|
2021-12-30 07:56:46 +00:00
|
|
|
model = engine.model
|
|
|
|
if isinstance(model, NaiveAMPModel):
|
2021-12-20 15:26:19 +00:00
|
|
|
self.dtype = torch.half
|
2021-12-30 07:56:46 +00:00
|
|
|
model = model.model
|
|
|
|
sig = inspect.signature(model.forward)
|
|
|
|
for p in sig.parameters.values():
|
|
|
|
assert p.kind != inspect.Parameter.VAR_POSITIONAL, '*args is not supported'
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
def _call_engine(model, input_tensor, batch_data):
|
|
|
|
if isinstance(model, NaiveAMPModel):
|
|
|
|
sig = inspect.signature(model.model.forward)
|
|
|
|
else:
|
|
|
|
sig = inspect.signature(model.forward)
|
|
|
|
if isinstance(batch_data, torch.Tensor):
|
|
|
|
if input_tensor is None:
|
|
|
|
return model(batch_data)
|
|
|
|
elif len(sig.parameters) > 1:
|
|
|
|
return model(input_tensor, batch_data)
|
|
|
|
else:
|
|
|
|
return model(input_tensor)
|
|
|
|
else:
|
|
|
|
filter_batch = True
|
|
|
|
for p in sig.parameters.values():
|
|
|
|
if p.kind == inspect.Parameter.VAR_KEYWORD:
|
|
|
|
filter_batch = False
|
|
|
|
if filter_batch:
|
|
|
|
batch_data = {k: v for k, v in batch_data.items() if k in sig.parameters}
|
|
|
|
if input_tensor is None:
|
|
|
|
return model(**batch_data)
|
|
|
|
else:
|
|
|
|
return model(input_tensor, **batch_data)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
2021-12-30 07:56:46 +00:00
|
|
|
def forward_step(self, engine, input_tensor, return_tensors, return_output_label=True, accum_loss=None):
|
2021-10-28 16:21:23 +00:00
|
|
|
"""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_tensor is used.
|
|
|
|
Returns output tensor. This is a helper function and can be ignored by users.
|
2021-12-13 14:07:01 +00:00
|
|
|
|
|
|
|
:param engine: your engine object
|
|
|
|
:type engine: colossalai.engine.Engine
|
|
|
|
:param input_tensor: input tensor for this pipeline stage
|
|
|
|
:type input_tensor: :class:`torch.Tensor`
|
|
|
|
:param return_tensors: a list of tensors to return
|
|
|
|
:type return_tensors: List[:class:`torch.Tensor`]
|
2021-12-20 15:26:19 +00:00
|
|
|
|
2021-12-13 14:07:01 +00:00
|
|
|
:return: output or the loss value of the current pipeline stage
|
|
|
|
:rtype: :class:`torch.Tensor`
|
2021-10-28 16:21:23 +00:00
|
|
|
"""
|
2021-12-30 07:56:46 +00:00
|
|
|
data, label = self.load_micro_batch()
|
|
|
|
output_tensor = self._call_engine(engine.model, input_tensor, data)
|
2021-10-28 16:21:23 +00:00
|
|
|
output_tensor = squeeze(output_tensor)
|
|
|
|
|
|
|
|
if gpc.is_last_rank(ParallelMode.PIPELINE):
|
2021-12-30 07:56:46 +00:00
|
|
|
if return_output_label:
|
|
|
|
return_tensors.append(tuple((output_tensor, label)))
|
|
|
|
if accum_loss is not None:
|
|
|
|
loss_reduced = self._call_engine_criterion(engine, output_tensor, label) / self.num_microbatches
|
|
|
|
accum_loss.add_(loss_reduced.detach())
|
2021-10-28 16:21:23 +00:00
|
|
|
return loss_reduced
|
|
|
|
else:
|
|
|
|
return output_tensor
|
|
|
|
else:
|
|
|
|
return output_tensor
|
|
|
|
|
Develop/experiments (#59)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
* Split conv2d, class token, positional embedding in 2d, Fix random number in ddp
Fix convergence in cifar10, Imagenet1000
* Integrate 1d tensor parallel in Colossal-AI (#39)
* fixed 1D and 2D convergence (#38)
* optimized 2D operations
* fixed 1D ViT convergence problem
* Feature/ddp (#49)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* support torch ddp
* fix loss accumulation
* add log for ddp
* change seed
* modify timing hook
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* Feature/pipeline (#40)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* optimize communication of pipeline parallel
* fix grad clip for pipeline
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* optimized 3d layer to fix slow computation ; tested imagenet performance with 3d; reworked lr_scheduler config definition; fixed launch args; fixed some printing issues; simplified apis of 3d layers (#51)
* Update 2.5d layer code to get a similar accuracy on imagenet-1k dataset
* update api for better usability (#58)
update api for better usability
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
2021-12-09 07:08:29 +00:00
|
|
|
def backward_step(self, engine, input_tensor, output_tensor, output_tensor_grad):
|
2021-10-28 16:21:23 +00:00
|
|
|
"""Backward step through the passed-in output tensor. If it is the last stage, the
|
|
|
|
output_tensor_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.
|
2021-12-13 14:07:01 +00:00
|
|
|
|
|
|
|
:param engine: your engine object
|
|
|
|
:type engine: colossalai.engine.Engine
|
|
|
|
:param input_tensor: input tensor for this pipeline stage
|
|
|
|
:type input_tensor: :class:`torch.Tensor`
|
|
|
|
:param output_tensor: output tensor for this pipeline stage
|
|
|
|
:type output_tensor: :class:`torch.Tensor`
|
|
|
|
:param output_tensor_grad: gradient of output tensor for this pipeline stage
|
|
|
|
:type output_tensor_grad: :class:`torch.Tensor`
|
|
|
|
|
|
|
|
:return: gradient of input tensor
|
|
|
|
:rtype: :class:`torch.Tensor`
|
2021-10-28 16:21:23 +00:00
|
|
|
"""
|
|
|
|
|
|
|
|
# Retain the grad on the input_tensor.
|
|
|
|
if input_tensor is not None:
|
|
|
|
input_tensor.retain_grad()
|
|
|
|
|
|
|
|
# Backward pass.
|
Develop/experiments (#59)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
* Split conv2d, class token, positional embedding in 2d, Fix random number in ddp
Fix convergence in cifar10, Imagenet1000
* Integrate 1d tensor parallel in Colossal-AI (#39)
* fixed 1D and 2D convergence (#38)
* optimized 2D operations
* fixed 1D ViT convergence problem
* Feature/ddp (#49)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* support torch ddp
* fix loss accumulation
* add log for ddp
* change seed
* modify timing hook
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* Feature/pipeline (#40)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* optimize communication of pipeline parallel
* fix grad clip for pipeline
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* optimized 3d layer to fix slow computation ; tested imagenet performance with 3d; reworked lr_scheduler config definition; fixed launch args; fixed some printing issues; simplified apis of 3d layers (#51)
* Update 2.5d layer code to get a similar accuracy on imagenet-1k dataset
* update api for better usability (#58)
update api for better usability
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
2021-12-09 07:08:29 +00:00
|
|
|
if output_tensor_grad is None:
|
|
|
|
engine.backward(output_tensor)
|
|
|
|
else:
|
|
|
|
engine.backward_by_grad(output_tensor, output_tensor_grad)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
# Collect the grad of the input_tensor.
|
|
|
|
input_tensor_grad = None
|
|
|
|
if input_tensor is not None:
|
|
|
|
input_tensor_grad = input_tensor.grad
|
|
|
|
|
|
|
|
return input_tensor_grad
|
|
|
|
|
2021-11-18 11:45:06 +00:00
|
|
|
def forward_backward_step(self,
|
Develop/experiments (#59)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
* Split conv2d, class token, positional embedding in 2d, Fix random number in ddp
Fix convergence in cifar10, Imagenet1000
* Integrate 1d tensor parallel in Colossal-AI (#39)
* fixed 1D and 2D convergence (#38)
* optimized 2D operations
* fixed 1D ViT convergence problem
* Feature/ddp (#49)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* support torch ddp
* fix loss accumulation
* add log for ddp
* change seed
* modify timing hook
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* Feature/pipeline (#40)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* optimize communication of pipeline parallel
* fix grad clip for pipeline
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* optimized 3d layer to fix slow computation ; tested imagenet performance with 3d; reworked lr_scheduler config definition; fixed launch args; fixed some printing issues; simplified apis of 3d layers (#51)
* Update 2.5d layer code to get a similar accuracy on imagenet-1k dataset
* update api for better usability (#58)
update api for better usability
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
2021-12-09 07:08:29 +00:00
|
|
|
engine,
|
2021-11-18 11:45:06 +00:00
|
|
|
data_iter,
|
|
|
|
forward_only=False,
|
2021-12-30 07:56:46 +00:00
|
|
|
return_loss=True,
|
|
|
|
return_output_label=True):
|
2021-10-28 16:21:23 +00:00
|
|
|
"""Runs non-interleaved 1F1B schedule, with communication between pipeline stages.
|
|
|
|
Returns a tuple with losses if the last stage, an empty tuple otherwise.
|
2021-11-18 11:45:06 +00:00
|
|
|
|
2021-12-13 14:07:01 +00:00
|
|
|
:param engine: your engine object
|
|
|
|
:type engine: colossalai.engine.Engine
|
|
|
|
:param data_iter: dataloader as the form of an iterator, obtained by calling iter(dataloader)
|
|
|
|
:type data_iter: Iterable
|
|
|
|
:param forward_only: whether run forward step only. Default is false. If true, no backward will be run.
|
|
|
|
:type forward_only: bool
|
|
|
|
:param return_loss: whether returns the loss value. Default is true.
|
|
|
|
:type return_loss: bool
|
2021-12-30 07:56:46 +00:00
|
|
|
:param return_output_label: If False, the output and label won't be returned
|
|
|
|
:type return_output_label: bool
|
2021-12-13 14:07:01 +00:00
|
|
|
|
2021-10-28 16:21:23 +00:00
|
|
|
:return: (output, label, loss)
|
2021-12-13 14:07:01 +00:00
|
|
|
:rtype: Tuple[:class:`torch.Tensor`]
|
2021-10-28 16:21:23 +00:00
|
|
|
"""
|
|
|
|
|
|
|
|
assert forward_only or return_loss, \
|
|
|
|
'The argument \'return_loss\' has to be True when \'forward_only\' is False, but got False.'
|
|
|
|
|
2021-11-18 11:45:06 +00:00
|
|
|
self.load_batch(data_iter)
|
2021-10-28 16:21:23 +00:00
|
|
|
num_warmup_microbatches = \
|
|
|
|
(gpc.get_world_size(ParallelMode.PIPELINE) -
|
|
|
|
gpc.get_local_rank(ParallelMode.PIPELINE) - 1)
|
|
|
|
num_warmup_microbatches = min(num_warmup_microbatches,
|
|
|
|
self.num_microbatches)
|
|
|
|
num_microbatches_remaining = self.num_microbatches - num_warmup_microbatches
|
|
|
|
|
|
|
|
# Input, output tensors only need to be saved when doing backward passes
|
|
|
|
input_tensors = None
|
|
|
|
output_tensors = None
|
|
|
|
if not forward_only:
|
|
|
|
input_tensors = []
|
|
|
|
output_tensors = []
|
|
|
|
return_tensors = []
|
2021-12-30 07:56:46 +00:00
|
|
|
if return_loss and gpc.is_pipeline_last_stage(ignore_virtual=True):
|
|
|
|
accum_loss = torch.zeros(1, device=get_current_device())
|
|
|
|
else:
|
|
|
|
accum_loss = None
|
2021-10-28 16:21:23 +00:00
|
|
|
# Used for tensor meta information communication
|
2021-12-30 07:56:46 +00:00
|
|
|
ft_shape = self.tensor_shape
|
2021-10-28 16:21:23 +00:00
|
|
|
bt_shape = None
|
2021-12-30 07:56:46 +00:00
|
|
|
fs_checker = self.tensor_shape is None
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
# Run warmup forward passes.
|
|
|
|
for i in range(num_warmup_microbatches):
|
|
|
|
if not gpc.is_first_rank(ParallelMode.PIPELINE):
|
2022-01-07 05:22:22 +00:00
|
|
|
ft_shape = comm.recv_tensor_meta(ft_shape)
|
|
|
|
input_tensor = comm.recv_forward(ft_shape, dtype=self.dtype,
|
|
|
|
scatter_gather_tensors=self.scatter_gather_tensors)
|
2021-11-18 11:45:06 +00:00
|
|
|
output_tensor = self.forward_step(
|
Develop/experiments (#59)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
* Split conv2d, class token, positional embedding in 2d, Fix random number in ddp
Fix convergence in cifar10, Imagenet1000
* Integrate 1d tensor parallel in Colossal-AI (#39)
* fixed 1D and 2D convergence (#38)
* optimized 2D operations
* fixed 1D ViT convergence problem
* Feature/ddp (#49)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* support torch ddp
* fix loss accumulation
* add log for ddp
* change seed
* modify timing hook
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* Feature/pipeline (#40)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* optimize communication of pipeline parallel
* fix grad clip for pipeline
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* optimized 3d layer to fix slow computation ; tested imagenet performance with 3d; reworked lr_scheduler config definition; fixed launch args; fixed some printing issues; simplified apis of 3d layers (#51)
* Update 2.5d layer code to get a similar accuracy on imagenet-1k dataset
* update api for better usability (#58)
update api for better usability
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
2021-12-09 07:08:29 +00:00
|
|
|
engine, input_tensor, return_tensors,
|
2021-12-30 07:56:46 +00:00
|
|
|
return_output_label=return_output_label,
|
|
|
|
accum_loss=accum_loss
|
2021-11-18 11:45:06 +00:00
|
|
|
)
|
2021-10-28 16:21:23 +00:00
|
|
|
if not gpc.is_last_rank(ParallelMode.PIPELINE):
|
|
|
|
bt_shape = output_tensor.shape
|
2022-01-07 05:22:22 +00:00
|
|
|
fs_checker = comm.send_tensor_meta(output_tensor, fs_checker)
|
|
|
|
comm.send_forward(output_tensor, scatter_gather_tensors=self.scatter_gather_tensors)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
if not forward_only:
|
|
|
|
input_tensors.append(input_tensor)
|
|
|
|
output_tensors.append(output_tensor)
|
|
|
|
|
|
|
|
# Before running 1F1B, need to receive first forward tensor.
|
|
|
|
# If all microbatches are run in warmup / cooldown phase, then no need to
|
|
|
|
# receive this tensor here.
|
|
|
|
if num_microbatches_remaining > 0:
|
|
|
|
if not gpc.is_first_rank(ParallelMode.PIPELINE):
|
2022-01-07 05:22:22 +00:00
|
|
|
ft_shape = comm.recv_tensor_meta(ft_shape)
|
|
|
|
input_tensor = comm.recv_forward(ft_shape, dtype=self.dtype,
|
|
|
|
scatter_gather_tensors=self.scatter_gather_tensors)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
# Run 1F1B in steady state.
|
|
|
|
for i in range(num_microbatches_remaining):
|
|
|
|
last_iteration = (i == (num_microbatches_remaining - 1))
|
|
|
|
|
2021-11-18 11:45:06 +00:00
|
|
|
output_tensor = self.forward_step(
|
Develop/experiments (#59)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
* Split conv2d, class token, positional embedding in 2d, Fix random number in ddp
Fix convergence in cifar10, Imagenet1000
* Integrate 1d tensor parallel in Colossal-AI (#39)
* fixed 1D and 2D convergence (#38)
* optimized 2D operations
* fixed 1D ViT convergence problem
* Feature/ddp (#49)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* support torch ddp
* fix loss accumulation
* add log for ddp
* change seed
* modify timing hook
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* Feature/pipeline (#40)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* optimize communication of pipeline parallel
* fix grad clip for pipeline
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* optimized 3d layer to fix slow computation ; tested imagenet performance with 3d; reworked lr_scheduler config definition; fixed launch args; fixed some printing issues; simplified apis of 3d layers (#51)
* Update 2.5d layer code to get a similar accuracy on imagenet-1k dataset
* update api for better usability (#58)
update api for better usability
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
2021-12-09 07:08:29 +00:00
|
|
|
engine, input_tensor, return_tensors,
|
2021-12-30 07:56:46 +00:00
|
|
|
return_output_label=return_output_label,
|
|
|
|
accum_loss=accum_loss
|
2021-11-18 11:45:06 +00:00
|
|
|
)
|
2021-10-28 16:21:23 +00:00
|
|
|
if forward_only:
|
2022-01-07 05:22:22 +00:00
|
|
|
comm.send_forward(output_tensor, scatter_gather_tensors=self.scatter_gather_tensors)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
if not last_iteration:
|
2022-01-07 05:22:22 +00:00
|
|
|
input_tensor = comm.recv_forward(ft_shape, dtype=self.dtype,
|
|
|
|
scatter_gather_tensors=self.scatter_gather_tensors)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
else:
|
2022-01-07 05:22:22 +00:00
|
|
|
output_tensor_grad = comm.send_forward_recv_backward(
|
|
|
|
output_tensor, bt_shape, dtype=self.dtype, scatter_gather_tensors=self.scatter_gather_tensors)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
# Add input_tensor and output_tensor to end of list.
|
|
|
|
input_tensors.append(input_tensor)
|
|
|
|
output_tensors.append(output_tensor)
|
|
|
|
|
|
|
|
# Pop input_tensor and output_tensor from the start of the list for
|
|
|
|
# the backward pass.
|
|
|
|
input_tensor = input_tensors.pop(0)
|
|
|
|
output_tensor = output_tensors.pop(0)
|
|
|
|
|
2021-11-18 11:45:06 +00:00
|
|
|
input_tensor_grad = self.backward_step(
|
Develop/experiments (#59)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
* Split conv2d, class token, positional embedding in 2d, Fix random number in ddp
Fix convergence in cifar10, Imagenet1000
* Integrate 1d tensor parallel in Colossal-AI (#39)
* fixed 1D and 2D convergence (#38)
* optimized 2D operations
* fixed 1D ViT convergence problem
* Feature/ddp (#49)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* support torch ddp
* fix loss accumulation
* add log for ddp
* change seed
* modify timing hook
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* Feature/pipeline (#40)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* optimize communication of pipeline parallel
* fix grad clip for pipeline
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* optimized 3d layer to fix slow computation ; tested imagenet performance with 3d; reworked lr_scheduler config definition; fixed launch args; fixed some printing issues; simplified apis of 3d layers (#51)
* Update 2.5d layer code to get a similar accuracy on imagenet-1k dataset
* update api for better usability (#58)
update api for better usability
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
2021-12-09 07:08:29 +00:00
|
|
|
engine,
|
2021-11-18 11:45:06 +00:00
|
|
|
input_tensor, output_tensor,
|
|
|
|
output_tensor_grad
|
|
|
|
)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
if last_iteration:
|
|
|
|
input_tensor = None
|
2022-01-07 05:22:22 +00:00
|
|
|
comm.send_backward(input_tensor_grad, scatter_gather_tensors=self.scatter_gather_tensors)
|
2021-10-28 16:21:23 +00:00
|
|
|
else:
|
2022-01-07 05:22:22 +00:00
|
|
|
input_tensor = comm.send_backward_recv_forward(
|
|
|
|
input_tensor_grad, ft_shape, dtype=self.dtype, scatter_gather_tensors=self.scatter_gather_tensors)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
# Run cooldown backward passes.
|
|
|
|
if not forward_only:
|
|
|
|
for i in range(num_warmup_microbatches):
|
|
|
|
input_tensor = input_tensors.pop(0)
|
|
|
|
output_tensor = output_tensors.pop(0)
|
|
|
|
|
2022-01-07 05:22:22 +00:00
|
|
|
output_tensor_grad = comm.recv_backward(bt_shape, dtype=self.dtype,
|
|
|
|
scatter_gather_tensors=self.scatter_gather_tensors)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
2021-11-18 11:45:06 +00:00
|
|
|
input_tensor_grad = self.backward_step(
|
Develop/experiments (#59)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
* Split conv2d, class token, positional embedding in 2d, Fix random number in ddp
Fix convergence in cifar10, Imagenet1000
* Integrate 1d tensor parallel in Colossal-AI (#39)
* fixed 1D and 2D convergence (#38)
* optimized 2D operations
* fixed 1D ViT convergence problem
* Feature/ddp (#49)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* support torch ddp
* fix loss accumulation
* add log for ddp
* change seed
* modify timing hook
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* Feature/pipeline (#40)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* optimize communication of pipeline parallel
* fix grad clip for pipeline
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* optimized 3d layer to fix slow computation ; tested imagenet performance with 3d; reworked lr_scheduler config definition; fixed launch args; fixed some printing issues; simplified apis of 3d layers (#51)
* Update 2.5d layer code to get a similar accuracy on imagenet-1k dataset
* update api for better usability (#58)
update api for better usability
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
2021-12-09 07:08:29 +00:00
|
|
|
engine,
|
2021-11-18 11:45:06 +00:00
|
|
|
input_tensor, output_tensor,
|
|
|
|
output_tensor_grad
|
|
|
|
)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
2022-01-07 05:22:22 +00:00
|
|
|
comm.send_backward(input_tensor_grad, scatter_gather_tensors=self.scatter_gather_tensors)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
if len(return_tensors) > 0:
|
2021-12-30 07:56:46 +00:00
|
|
|
output, label = tuple(map(list, zip(*return_tensors)))
|
|
|
|
return (torch.cat(output, dim=0),
|
|
|
|
torch.cat(label, dim=0),
|
|
|
|
accum_loss)
|
2021-10-28 16:21:23 +00:00
|
|
|
else:
|
2021-12-30 07:56:46 +00:00
|
|
|
return tuple((None, None, accum_loss))
|
2021-12-20 15:26:19 +00:00
|
|
|
|
|
|
|
|
|
|
|
class InterleavedPipelineSchedule(PipelineSchedule):
|
2021-12-30 07:56:46 +00:00
|
|
|
def __init__(self,
|
|
|
|
num_microbatches,
|
|
|
|
num_model_chunks,
|
|
|
|
batch_data_process_func: Callable = None,
|
2022-01-07 05:22:22 +00:00
|
|
|
tensor_shape: Union[torch.Size, List[int], Tuple[int]] = None,
|
|
|
|
scatter_gather_tensors: bool = False):
|
2021-12-30 07:56:46 +00:00
|
|
|
"""A helper schedule class for pipeline parallelism running environment.
|
|
|
|
It uses interleaved 1F1B strategy. Other properties are similar as
|
|
|
|
:class:`NonPipelineSchedule`.
|
|
|
|
|
|
|
|
:param num_microbatches: The number of microbatches
|
|
|
|
:type num_microbatches: int
|
|
|
|
:param num_model_chunks: The number of model chunks
|
|
|
|
:type num_model_chunks: int
|
|
|
|
:param batch_data_process_func: The preprocessing function which receives a batch of data, and it will be executed in `load_batch`
|
|
|
|
:type batch_data_process_func: Callable
|
2022-01-07 05:22:22 +00:00
|
|
|
:param scatter_gather_tensors: If set to `True`, communication will be reduced over pipeline when using 1D tensor parallelization
|
|
|
|
:type scatter_gather_tensors: bool
|
2021-12-30 07:56:46 +00:00
|
|
|
"""
|
2021-12-20 15:26:19 +00:00
|
|
|
assert num_microbatches % gpc.get_world_size(ParallelMode.PIPELINE) == 0, \
|
|
|
|
'num_microbatches must be an integer multiple of pipeline parallel world size'
|
2022-01-07 05:22:22 +00:00
|
|
|
super().__init__(num_microbatches, batch_data_process_func=batch_data_process_func,
|
|
|
|
tensor_shape=tensor_shape, scatter_gather_tensors=scatter_gather_tensors)
|
2021-12-20 15:26:19 +00:00
|
|
|
gpc.set_virtual_pipeline_parallel_size(num_model_chunks)
|
|
|
|
gpc.set_virtual_pipeline_parallel_rank(0)
|
2021-12-30 07:56:46 +00:00
|
|
|
self.num_model_chunks = num_model_chunks
|
2021-12-20 15:26:19 +00:00
|
|
|
|
|
|
|
def pre_processing(self, engine):
|
|
|
|
if isinstance(engine.optimizer, (ZeroRedundancyOptimizer_Level_2, ZeroRedundancyOptimizer_Level_3)):
|
|
|
|
raise TypeError(
|
|
|
|
"Pipeline schedule is currently not compatible with ZeRO Level 2 and Level 3"
|
|
|
|
)
|
|
|
|
|
|
|
|
if isinstance(engine.model[0], NaiveAMPModel):
|
|
|
|
self.dtype = torch.half
|
|
|
|
|
2021-12-30 07:56:46 +00:00
|
|
|
for model in engine.model:
|
|
|
|
if isinstance(model, NaiveAMPModel):
|
|
|
|
model = model.model
|
|
|
|
sig = inspect.signature(model.forward)
|
|
|
|
for p in sig.parameters.values():
|
|
|
|
assert p.kind != inspect.Parameter.VAR_POSITIONAL, '*args is not supported'
|
|
|
|
|
|
|
|
def load_batch(self, data_iter):
|
|
|
|
super().load_batch(data_iter)
|
|
|
|
# overwrite microbatch_offset, since model chunks load the same microbatch, and should tract the offset
|
|
|
|
self.microbatch_offset = [0 for _ in range(self.num_model_chunks)]
|
|
|
|
|
|
|
|
def load_micro_batch(self, model_chunk_id):
|
|
|
|
data = self._get_data_slice(self.batch_data, self.microbatch_offset[model_chunk_id])
|
|
|
|
label = self._get_data_slice(self.batch_label, self.microbatch_offset[model_chunk_id])
|
|
|
|
self.microbatch_offset[model_chunk_id] += self.microbatch_size
|
|
|
|
return self._move_to_device(data), self._move_to_device(label)
|
|
|
|
|
|
|
|
def forward_step(self, engine, model_chunk_id, input_tensor, return_tensors, return_output_label=True, accum_loss=None):
|
2021-12-20 15:26:19 +00:00
|
|
|
"""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_tensor is used.
|
|
|
|
Returns output tensor. This is a helper function and can be ignored by users.
|
|
|
|
"""
|
2021-12-30 07:56:46 +00:00
|
|
|
data, label = self.load_micro_batch(model_chunk_id)
|
|
|
|
output_tensor = self._call_engine(engine.model[model_chunk_id], input_tensor, data)
|
2021-12-20 15:26:19 +00:00
|
|
|
output_tensor = squeeze(output_tensor)
|
|
|
|
|
|
|
|
if gpc.is_pipeline_last_stage():
|
2021-12-30 07:56:46 +00:00
|
|
|
if return_output_label:
|
|
|
|
return_tensors.append(tuple(output_tensor, label))
|
|
|
|
if accum_loss is not None:
|
|
|
|
loss_reduced = self._call_engine_criterion(engine, output_tensor, label) / self.num_microbatches
|
|
|
|
accum_loss.add_(loss_reduced.detach())
|
2021-12-20 15:26:19 +00:00
|
|
|
return loss_reduced
|
|
|
|
else:
|
|
|
|
return output_tensor
|
|
|
|
else:
|
|
|
|
return output_tensor
|
|
|
|
|
2021-12-30 07:56:46 +00:00
|
|
|
def forward_backward_step(self, engine, data_iter, forward_only=False, return_loss=True, return_output_label=True):
|
2021-12-20 15:26:19 +00:00
|
|
|
"""Run interleaved 1F1B schedule (model split into model chunks), with
|
|
|
|
communication between pipeline stages as needed.
|
|
|
|
|
|
|
|
Returns dictionary with losses if the last stage, empty dict otherwise."""
|
|
|
|
assert forward_only or return_loss, \
|
|
|
|
'The argument \'return_loss\' has to be True when \'forward_only\' is False, but got False.'
|
|
|
|
self.load_batch(data_iter)
|
|
|
|
model = engine.model
|
|
|
|
input_tensors = [[] for _ in range(len(model))]
|
|
|
|
output_tensors = [[] for _ in range(len(model))]
|
|
|
|
return_tensors = []
|
|
|
|
if not forward_only:
|
|
|
|
output_tensor_grads = [[] for _ in range(len(model))]
|
2021-12-30 07:56:46 +00:00
|
|
|
if return_loss and gpc.is_pipeline_last_stage(ignore_virtual=True):
|
|
|
|
accum_loss = torch.zeros(1, device=get_current_device())
|
|
|
|
else:
|
|
|
|
accum_loss = None
|
2021-12-20 15:26:19 +00:00
|
|
|
|
|
|
|
# Used for tensor meta information communication
|
2021-12-30 07:56:46 +00:00
|
|
|
input_tensor_shapes = [self.tensor_shape for _ in range(len(model))]
|
2021-12-20 15:26:19 +00:00
|
|
|
output_tensor_shapes = [None for _ in range(len(model))]
|
2021-12-30 07:56:46 +00:00
|
|
|
send_tensor_shape_flags = [self.tensor_shape is None for _ in range(len(model))]
|
2021-12-20 15:26:19 +00:00
|
|
|
|
|
|
|
pipeline_parallel_size = gpc.get_world_size(ParallelMode.PIPELINE)
|
|
|
|
pipeline_parallel_rank = gpc.get_local_rank(ParallelMode.PIPELINE)
|
|
|
|
|
|
|
|
# Compute number of warmup and remaining microbatches.
|
|
|
|
num_model_chunks = len(model)
|
|
|
|
num_microbatches = self.num_microbatches * num_model_chunks
|
|
|
|
all_warmup_microbatches = False
|
|
|
|
if forward_only:
|
|
|
|
num_warmup_microbatches = num_microbatches
|
|
|
|
else:
|
|
|
|
# Run all forward passes and then all backward passes if number of
|
|
|
|
# microbatches is just the number of pipeline stages.
|
|
|
|
# Otherwise, perform (num_model_chunks-1)*pipeline_parallel_size on
|
|
|
|
# all workers, followed by more microbatches after depending on
|
|
|
|
# stage ID (more forward passes for earlier stages, later stages can
|
|
|
|
# immediately start with 1F1B).
|
|
|
|
if self.num_microbatches == pipeline_parallel_size:
|
|
|
|
num_warmup_microbatches = num_microbatches
|
|
|
|
all_warmup_microbatches = True
|
|
|
|
else:
|
|
|
|
num_warmup_microbatches = \
|
|
|
|
(pipeline_parallel_size - pipeline_parallel_rank - 1) * 2
|
|
|
|
num_warmup_microbatches += (
|
|
|
|
num_model_chunks - 1) * pipeline_parallel_size
|
|
|
|
num_warmup_microbatches = min(num_warmup_microbatches,
|
|
|
|
num_microbatches)
|
|
|
|
num_microbatches_remaining = \
|
|
|
|
num_microbatches - num_warmup_microbatches
|
|
|
|
|
|
|
|
def get_model_chunk_id(microbatch_id, forward):
|
|
|
|
"""Helper method to get the model chunk ID given the iteration number."""
|
|
|
|
microbatch_id_in_group = microbatch_id % (pipeline_parallel_size * num_model_chunks)
|
|
|
|
model_chunk_id = microbatch_id_in_group // pipeline_parallel_size
|
|
|
|
if not forward:
|
|
|
|
model_chunk_id = (num_model_chunks - model_chunk_id - 1)
|
|
|
|
return model_chunk_id
|
|
|
|
|
|
|
|
def forward_step_helper(microbatch_id):
|
|
|
|
"""Helper method to run forward step with model split into chunks
|
|
|
|
(run set_virtual_pipeline_model_parallel_rank() before calling
|
|
|
|
forward_step())."""
|
|
|
|
model_chunk_id = get_model_chunk_id(microbatch_id, forward=True)
|
|
|
|
gpc.set_virtual_pipeline_parallel_rank(model_chunk_id)
|
|
|
|
|
|
|
|
# forward step
|
|
|
|
if gpc.is_pipeline_first_stage():
|
|
|
|
if len(input_tensors[model_chunk_id]) == \
|
|
|
|
len(output_tensors[model_chunk_id]):
|
|
|
|
input_tensors[model_chunk_id].append(None)
|
|
|
|
input_tensor = input_tensors[model_chunk_id][-1]
|
2021-12-30 07:56:46 +00:00
|
|
|
output_tensor = self.forward_step(engine, model_chunk_id, input_tensor,
|
|
|
|
return_tensors, return_output_label=return_output_label, accum_loss=accum_loss)
|
2021-12-20 15:26:19 +00:00
|
|
|
output_tensors[model_chunk_id].append(output_tensor)
|
|
|
|
|
|
|
|
# if forward-only, no need to save tensors for a backward pass
|
|
|
|
if forward_only:
|
|
|
|
input_tensors[model_chunk_id].pop()
|
|
|
|
output_tensors[model_chunk_id].pop()
|
|
|
|
|
|
|
|
return output_tensor
|
|
|
|
|
|
|
|
def backward_step_helper(microbatch_id):
|
|
|
|
"""Helper method to run backward step with model split into chunks
|
|
|
|
(run set_virtual_pipeline_model_parallel_rank() before calling
|
|
|
|
backward_step())."""
|
|
|
|
model_chunk_id = get_model_chunk_id(microbatch_id, forward=False)
|
|
|
|
gpc.set_virtual_pipeline_parallel_rank(model_chunk_id)
|
|
|
|
|
|
|
|
if gpc.is_pipeline_last_stage():
|
|
|
|
if len(output_tensor_grads[model_chunk_id]) == 0:
|
|
|
|
output_tensor_grads[model_chunk_id].append(None)
|
|
|
|
input_tensor = input_tensors[model_chunk_id].pop(0)
|
|
|
|
output_tensor = output_tensors[model_chunk_id].pop(0)
|
|
|
|
output_tensor_grad = output_tensor_grads[model_chunk_id].pop(0)
|
|
|
|
input_tensor_grad = self.backward_step(engine, input_tensor, output_tensor, output_tensor_grad)
|
|
|
|
|
|
|
|
return input_tensor_grad
|
|
|
|
|
|
|
|
# Run warmup forward passes.
|
|
|
|
gpc.set_virtual_pipeline_parallel_rank(0)
|
|
|
|
if not gpc.is_pipeline_first_stage():
|
2022-01-07 05:22:22 +00:00
|
|
|
input_tensor_shapes[0] = comm.recv_tensor_meta(input_tensor_shapes[0])
|
|
|
|
input_tensors[0].append(comm.recv_forward(input_tensor_shapes[0], dtype=self.dtype,
|
|
|
|
scatter_gather_tensors=self.scatter_gather_tensors))
|
2021-12-20 15:26:19 +00:00
|
|
|
|
|
|
|
for k in range(num_warmup_microbatches):
|
|
|
|
model_chunk_id = get_model_chunk_id(k, forward=True)
|
|
|
|
output_tensor = forward_step_helper(k)
|
|
|
|
if not gpc.is_pipeline_last_stage():
|
|
|
|
output_tensor_shapes[model_chunk_id] = output_tensor.shape
|
2022-01-07 05:22:22 +00:00
|
|
|
send_tensor_shape_flags[model_chunk_id] = comm.send_tensor_meta(
|
2021-12-20 15:26:19 +00:00
|
|
|
output_tensor, send_tensor_shape_flags[model_chunk_id])
|
|
|
|
# Determine if tensor should be received from previous stage.
|
|
|
|
next_forward_model_chunk_id = get_model_chunk_id(k+1, forward=True)
|
|
|
|
recv_prev = True
|
|
|
|
if gpc.is_pipeline_first_stage(ignore_virtual=True):
|
|
|
|
if next_forward_model_chunk_id == 0:
|
|
|
|
recv_prev = False
|
|
|
|
if k == (num_microbatches - 1):
|
|
|
|
recv_prev = False
|
|
|
|
|
|
|
|
# Don't send tensor downstream if on last stage.
|
|
|
|
if gpc.is_pipeline_last_stage():
|
|
|
|
output_tensor = None
|
|
|
|
|
|
|
|
with switch_virtual_pipeline_parallel_rank(next_forward_model_chunk_id):
|
|
|
|
if not gpc.is_pipeline_first_stage():
|
2022-01-07 05:22:22 +00:00
|
|
|
input_tensor_shapes[next_forward_model_chunk_id] = comm.recv_tensor_meta(
|
2021-12-20 15:26:19 +00:00
|
|
|
input_tensor_shapes[next_forward_model_chunk_id])
|
|
|
|
# Send and receive tensors as appropriate (send tensors computed
|
|
|
|
# in this iteration; receive tensors for next iteration).
|
|
|
|
input_shape = input_tensor_shapes[next_forward_model_chunk_id] if recv_prev else None
|
|
|
|
if k == (num_warmup_microbatches - 1) and not forward_only and \
|
|
|
|
not all_warmup_microbatches:
|
|
|
|
input_tensor_grad = None
|
|
|
|
recv_next = True
|
|
|
|
if gpc.is_pipeline_last_stage(ignore_virtual=True):
|
|
|
|
recv_next = False
|
|
|
|
output_shape = output_tensor_shapes[num_model_chunks-1] if recv_next else None
|
|
|
|
input_tensor, output_tensor_grad = \
|
2022-01-07 05:22:22 +00:00
|
|
|
comm.send_forward_backward_recv_forward_backward(
|
2021-12-20 15:26:19 +00:00
|
|
|
output_tensor, input_tensor_grad,
|
|
|
|
input_shape,
|
|
|
|
output_shape,
|
|
|
|
recv_prev=recv_prev, recv_next=recv_next,
|
2022-01-07 05:22:22 +00:00
|
|
|
dtype=self.dtype,
|
|
|
|
scatter_gather_tensors=self.scatter_gather_tensors)
|
2021-12-20 15:26:19 +00:00
|
|
|
output_tensor_grads[num_model_chunks-1].append(output_tensor_grad)
|
|
|
|
else:
|
|
|
|
input_tensor = \
|
2022-01-07 05:22:22 +00:00
|
|
|
comm.send_forward_recv_forward(
|
2021-12-20 15:26:19 +00:00
|
|
|
output_tensor,
|
|
|
|
input_shape,
|
|
|
|
recv_prev=recv_prev,
|
2022-01-07 05:22:22 +00:00
|
|
|
dtype=self.dtype,
|
|
|
|
scatter_gather_tensors=self.scatter_gather_tensors)
|
2021-12-20 15:26:19 +00:00
|
|
|
input_tensors[next_forward_model_chunk_id].append(input_tensor)
|
|
|
|
|
|
|
|
# Run 1F1B in steady state.
|
|
|
|
for k in range(num_microbatches_remaining):
|
|
|
|
# Forward pass.
|
|
|
|
forward_k = k + num_warmup_microbatches
|
|
|
|
output_tensor = forward_step_helper(forward_k)
|
|
|
|
|
|
|
|
# Backward pass.
|
|
|
|
backward_k = k
|
|
|
|
input_tensor_grad = backward_step_helper(backward_k)
|
|
|
|
|
|
|
|
# Send output_tensor and input_tensor_grad, receive input_tensor
|
|
|
|
# and output_tensor_grad.
|
|
|
|
|
|
|
|
# Determine if current stage has anything to send in either direction,
|
|
|
|
# otherwise set tensor to None.
|
|
|
|
forward_model_chunk_id = get_model_chunk_id(forward_k, forward=True)
|
|
|
|
gpc.set_virtual_pipeline_parallel_rank(forward_model_chunk_id)
|
|
|
|
if gpc.is_pipeline_last_stage():
|
|
|
|
output_tensor = None
|
|
|
|
|
|
|
|
backward_model_chunk_id = get_model_chunk_id(backward_k, forward=False)
|
|
|
|
gpc.set_virtual_pipeline_parallel_rank(backward_model_chunk_id)
|
|
|
|
if gpc.is_pipeline_first_stage():
|
|
|
|
input_tensor_grad = None
|
|
|
|
|
|
|
|
# Determine if peers are sending, and where in data structure to put
|
|
|
|
# received tensors.
|
|
|
|
recv_prev = True
|
|
|
|
if gpc.is_pipeline_first_stage(ignore_virtual=True):
|
|
|
|
# First stage is ahead of last stage by (pipeline_parallel_size - 1).
|
|
|
|
next_forward_model_chunk_id = get_model_chunk_id(
|
|
|
|
forward_k - (pipeline_parallel_size - 1), forward=True)
|
|
|
|
if next_forward_model_chunk_id == (num_model_chunks - 1):
|
|
|
|
recv_prev = False
|
|
|
|
next_forward_model_chunk_id += 1
|
|
|
|
else:
|
|
|
|
next_forward_model_chunk_id = get_model_chunk_id(forward_k + 1,
|
|
|
|
forward=True)
|
|
|
|
|
|
|
|
recv_next = True
|
|
|
|
if gpc.is_pipeline_last_stage(ignore_virtual=True):
|
|
|
|
# Last stage is ahead of first stage by (pipeline_parallel_size - 1).
|
|
|
|
next_backward_model_chunk_id = get_model_chunk_id(
|
|
|
|
backward_k - (pipeline_parallel_size - 1), forward=False)
|
|
|
|
if next_backward_model_chunk_id == 0:
|
|
|
|
recv_next = False
|
|
|
|
next_backward_model_chunk_id -= 1
|
|
|
|
else:
|
|
|
|
next_backward_model_chunk_id = get_model_chunk_id(backward_k + 1,
|
|
|
|
forward=False)
|
|
|
|
|
|
|
|
# If last iteration, don't receive; we already received one extra
|
|
|
|
# before the start of the for loop.
|
|
|
|
if k == (num_microbatches_remaining - 1):
|
|
|
|
recv_prev = False
|
|
|
|
|
|
|
|
input_shape = input_tensor_shapes[next_forward_model_chunk_id] if recv_prev else None
|
|
|
|
output_shape = output_tensor_shapes[next_backward_model_chunk_id] if recv_next else None
|
|
|
|
# Communicate tensors.
|
|
|
|
input_tensor, output_tensor_grad = \
|
2022-01-07 05:22:22 +00:00
|
|
|
comm.send_forward_backward_recv_forward_backward(
|
2021-12-20 15:26:19 +00:00
|
|
|
output_tensor, input_tensor_grad,
|
|
|
|
input_shape,
|
|
|
|
output_shape,
|
|
|
|
recv_prev=recv_prev, recv_next=recv_next,
|
2022-01-07 05:22:22 +00:00
|
|
|
dtype=self.dtype,
|
|
|
|
scatter_gather_tensors=self.scatter_gather_tensors)
|
2021-12-20 15:26:19 +00:00
|
|
|
|
|
|
|
# Put input_tensor and output_tensor_grad in data structures in the
|
|
|
|
# right location.
|
|
|
|
if recv_prev:
|
|
|
|
input_tensors[next_forward_model_chunk_id].append(input_tensor)
|
|
|
|
if recv_next:
|
|
|
|
output_tensor_grads[next_backward_model_chunk_id].append(
|
|
|
|
output_tensor_grad)
|
|
|
|
|
|
|
|
# Run cooldown backward passes (flush out pipeline).
|
|
|
|
if not forward_only:
|
|
|
|
if all_warmup_microbatches:
|
|
|
|
output_tensor_grads[num_model_chunks-1].append(
|
2022-01-07 05:22:22 +00:00
|
|
|
comm.recv_backward(output_tensor_shapes[num_model_chunks-1], scatter_gather_tensors=self.scatter_gather_tensors))
|
2021-12-20 15:26:19 +00:00
|
|
|
for k in range(num_microbatches_remaining, num_microbatches):
|
|
|
|
input_tensor_grad = backward_step_helper(k)
|
|
|
|
next_backward_model_chunk_id = get_model_chunk_id(k+1, forward=False)
|
|
|
|
recv_next = True
|
|
|
|
if gpc.is_pipeline_last_stage(ignore_virtual=True):
|
|
|
|
if next_backward_model_chunk_id == (num_model_chunks - 1):
|
|
|
|
recv_next = False
|
|
|
|
if k == (num_microbatches - 1):
|
|
|
|
recv_next = False
|
|
|
|
output_shape = output_tensor_shapes[next_backward_model_chunk_id] if recv_next else None
|
|
|
|
output_tensor_grads[next_backward_model_chunk_id].append(
|
2022-01-07 05:22:22 +00:00
|
|
|
comm.send_backward_recv_backward(
|
2021-12-20 15:26:19 +00:00
|
|
|
input_tensor_grad,
|
|
|
|
output_shape,
|
|
|
|
recv_next=recv_next,
|
2022-01-07 05:22:22 +00:00
|
|
|
dtype=self.dtype,
|
|
|
|
scatter_gather_tensors=self.scatter_gather_tensors))
|
2021-12-20 15:26:19 +00:00
|
|
|
|
|
|
|
if len(return_tensors) > 0:
|
2021-12-30 07:56:46 +00:00
|
|
|
output, label = tuple(map(list, zip(*return_tensors)))
|
|
|
|
return (torch.cat(output, dim=0),
|
|
|
|
torch.cat(label, dim=0),
|
|
|
|
accum_loss)
|
2021-12-20 15:26:19 +00:00
|
|
|
else:
|
2021-12-30 07:56:46 +00:00
|
|
|
return tuple((None, None, accum_loss))
|