github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-03 18:46:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

[feat] fix func name & ci; add comments;

Browse Source
This commit is contained in:
duanjunwen 2024-08-28 03:40:50 +00:00
parent b5f7b4d228
commit 582ba0d6ff
1 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -36,8 +36,8 @@ def get_model_numel(model: torch.nn.Module) -> Tuple[int, int]:
return num_params, num_params_trainable
# Test iter input & multiple microbatch
def test_run_fwd_bwd_iter_input(
# Test manual v_schedule with multiple microbatch
def run_fwd_bwd_iter_input(
rank: int,
world_size: int,
port: int,
@ -474,8 +474,8 @@ def test_run_fwd_bwd_iter_input(
assert_close(local_chunk[1].weight.grad, model_base.layers[4].weight.grad)
# T
def test_run_fwd_bwd_with_vschedule(
# Test v_schedule generated by graph with multiple microbatch
def run_fwd_bwd_with_vschedule(
rank: int,
world_size: int,
port: int,
@ -623,7 +623,7 @@ def test_run_fwd_bwd_with_vschedule(
@rerun_if_address_is_in_use()
def test_pp(num_microbatch: int, batch_size: int, num_model_chunk: int):
spawn(
test_run_fwd_bwd_with_vschedule,
run_fwd_bwd_with_vschedule,
nprocs=4,
num_microbatch=num_microbatch,
batch_size=batch_size,