mirror of https://github.com/hpcaitech/ColossalAI
aibig-modeldata-parallelismdeep-learningdistributed-computingfoundation-modelsheterogeneous-traininghpcinferencelarge-scalemodel-parallelismpipeline-parallelism
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
22 lines
559 B
22 lines
559 B
from abc import abstractmethod |
|
from typing import Any, Callable, Iterator, Optional |
|
|
|
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: Optional[OptimizerWrapper] = None, |
|
return_loss: bool = True, |
|
return_outputs: bool = False, |
|
) -> dict: |
|
pass
|
|
|