mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2026-04-11 14:43:10 +00:00
261eab02fb
* [amp] add mixed precision optimizer * [plugin] add 3d parallel plugin * [booster] support pipeline * [plugin] 3d parallel plugin support clip grad norm * [shardformer] fix sharder and add plugin test * [plugin] rename 3d parallel plugin * [ci] support testmon core pkg change detection (#4305) * [hotfix] debug testmon * [hotfix] fix llama * [hotfix] fix p2p bugs * [hotfix] fix requirements
22 lines
620 B
Python
22 lines
620 B
Python
from abc import abstractmethod
|
|
from typing import Any, Callable, Iterator
|
|
|
|
import torch
|
|
|
|
from colossalai.interface import ModelWrapper, OptimizerWrapper
|
|
|
|
from .plugin_base import Plugin
|
|
|
|
|
|
class PipelinePluginBase(Plugin):
|
|
|
|
@abstractmethod
|
|
def execute_pipeline(self,
|
|
data_iter: Iterator,
|
|
model: ModelWrapper,
|
|
criterion: Callable[[Any, Any], torch.Tensor],
|
|
optimizer: OptimizerWrapper,
|
|
return_loss: bool = True,
|
|
return_outputs: bool = False) -> dict:
|
|
pass
|