github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-12-02 19:56:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

[pipeline/rpc] support interleaving | fix checkpoint bug | change logic when dispatch data in work_list to ensure steady 1F1B (#1483)

Browse Source 
* support p2p communication with any type of object | pass test

* reconstruct pipeline schedule with p2p_v2.py(support communication with List[Any]) | pass test

* [engin/schedule] use p2p_v2 to recontruct pipeline_schedule

* [pipeline/rpc] implement a demo for PP with cuda rpc framework

* [pipeline/rpc] support interleaving | fix checkpoint bug | change logic when dispatch data in work_list to ensure steady 1F1B
This commit is contained in:
Kirigaya Kazuto
2022-08-24 11:19:46 +08:00
committed by GitHub
parent d6e3dca436
commit a6c8749198
3 changed files with 366 additions and 139 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
tests/test_pipeline/test_cuda_rpc_pipeline.py
View File

@@ -11,14 +11,14 @@ from colossalai.pipeline.rpc.PipelineBase import FillDrainPipelineEngine, OneFOn
class TestModel(nn.Module):
def __init__(self, rank, world_size, feat_num, h) -> None:
def __init__(self, stage_id, actual_stage_num, feat_num, h) -> None:
super().__init__()
self.rank = rank
self.is_last_rank = rank == world_size - 1
self.linear_name = f'linear_{rank}'
if rank == 0:
self.rank = stage_id
self.is_last_rank = stage_id == actual_stage_num - 1
self.linear_name = f'linear_{stage_id}'
if stage_id == 0:
setattr(self, self.linear_name, nn.Linear(feat_num, h))
elif rank == world_size - 1:
elif stage_id == actual_stage_num - 1:
setattr(self, self.linear_name, nn.Linear(h, 1))
else:
setattr(self, self.linear_name, nn.Linear(h, h))
@@ -35,32 +35,35 @@ class TestModel(nn.Module):
def run_main(args):
torch.manual_seed(100)
sample_num = 128
feat_num = 10000
h = 10000
device = args.device
world_size = args.world_size
batch_size = 128
stage_num = args.world_size
chunk = args.chunk
num_microbatches = args.num_microbatches
actual_stage_num = stage_num * chunk
use_interleave = args.use_interleave
use_checkpoint = args.use_checkpoint
sample_num = 1024
feat_num = 10
h = 10
batch_size = 1024
assert sample_num % batch_size == 0
batch_num = sample_num // batch_size
num_microbatches = world_size
input_sample = torch.randn((sample_num, feat_num), device=device)
module_partitions = [TestModel(rank, world_size, feat_num, h) for rank in range(world_size)]
module_partitions = [TestModel(pp_rank, actual_stage_num, feat_num, h) for pp_rank in range(actual_stage_num)]
engine = OneFOneBPipelineEngine(module_partitions=module_partitions,
chunk=1,
world_size=world_size,
stage_num=stage_num,
num_microbatches=num_microbatches,
device=args.device,
max_outstanding=world_size,
use_interleave=False,
checkpoint=False)
device=device,
chunk=chunk,
use_interleave=use_interleave,
checkpoint=use_checkpoint)
for i in range(batch_num):
batch = input_sample[i * batch_size:(i + 1) * batch_size]
engine.forward_backward(batch)
_ = engine.forward_backward(input_sample)
def run_worker(rank, args):
@@ -88,7 +91,11 @@ def run_worker(rank, args):
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--world_size', type=int, default=2)
parser.add_argument('--num_microbatches', type=int, default=2)
parser.add_argument('--device', type=str, default='cuda')
parser.add_argument('--chunk', type=int, default=1)
parser.add_argument('--use_checkpoint', action='store_true')
parser.add_argument('--use_interleave', action='store_true')
parser.add_argument('--master_addr', type=str, default='localhost')
parser.add_argument('--master_port', type=str, default='29020')
parser.add_argument('--num_worker_threads', type=str, default=128)