ColossalAI/__init__.py at dev/gpt2_metainfo_patch - ColossalAI - Gitea: Git with a cup of tea at home

github/ColossalAI

mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-11 04:41:54 +00:00

Kirigaya Kazuto 9708638ded

[pipeline/pytree] add pytree to process args and kwargs | provide data_process_func to process args and kwargs after forward (#1642 )

* [pipeline/tuning] improve dispatch performance both time and space cost

* [pipeline/converge] add interface for testing convergence

* [NFC] polish colossalai/utils/multi_tensor_apply/multi_tensor_apply.py code style

* Update PipelineBase.py

* [pipeline/chimera] reconstruct PipelineBase and Worker to support more feasible custom schedule | finish Chimera

* [pipeline/chimera] test chimera | fix bug of initializing

* [pipeline/pytree] add pytree to process args and kwargs | provide  to process args and kwargs after forward

2022-09-29 10:58:58 +08:00

4 lines

236 B

Python

Raw Permalink Blame History

 from ._pipeline_schedule import FillDrainPipelineEngine, OneFOneBPipelineEngine, ChimeraPipelineEngine
 from .utils import pytree_map
 __all__ = ['FillDrainPipelineEngine', 'OneFOneBPipelineEngine', 'ChimeraPipelineEngine', 'pytree_map']