github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Merge remote-tracking branch 'upstream/feature/zerobubble' into dev/zero_bubble

pull/6065/head
duanjunwen 2024-09-19 03:52:13 +00:00
parent af2c2f8092 11ae6848c6
commit 1f5c7258aa
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tests/test_pipeline/test_schedule/test_zerobubble_pp.py
View File

@ -18,8 +18,7 @@ from colossalai.pipeline.schedule.zero_bubble_pp import ZeroBubbleVPipeScheduler
from colossalai.pipeline.stage_manager import PipelineStageManager
from colossalai.tensor.d_tensor.api import clear_layout_converter
from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
# from tests.kit.model_zoo import model_zoo
from tests.kit.model_zoo import model_zoo
class MlpModel(nn.Module):