github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-12-05 15:46:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

[pipeline/chimera] reconstruct PipelineBase and Worker to support more feasible custom schedule | finish Chimera (#1595)

Browse Source 
* [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
This commit is contained in:
Kirigaya Kazuto
2022-09-19 11:44:18 +08:00
committed by GitHub
parent c9e8ce67b8
commit edc9e419ad
8 changed files with 614 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_pipeline/test_cuda_rpc_optimizer.py
View File

@@ -3,7 +3,7 @@ from torch import nn
from torch import autograd
from torch.optim import SGD, Adam, RMSprop, Optimizer
from colossalai.pipeline.rpc.PipelineBase import FillDrainPipelineEngine, OneFOneBPipelineEngine
from colossalai.pipeline.rpc._pipeline_schedule import FillDrainPipelineEngine, OneFOneBPipelineEngine
from colossalai.testing import assert_close
from rpc_test_utils import rpc_run, parse_args, RpcTestModel