[pipeline/rpc] update outstanding mechanism | optimize dispatching strategy (#1497)

* 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

* [pipeline/rpc] implement distributed optimizer | test with assert_close

* [pipeline/rpc] implement distributed optimizer | test with assert_close

* [pipeline/rpc] update outstanding mechanism | optimize dispatching strategy

* [pipeline/rpc] update outstanding mechanism | optimize dispatching strategy

* [pipeline/rpc] update outstanding mechanism | optimize dispatching strategy

tests/test_pipeline/rpc_test_utils.py

@@ -5,13 +5,9 @@ import torch
 from torch import nn
 import torch.multiprocessing as mp
 import torch.distributed.rpc as rpc
-from torch import autograd
 from torch.optim import SGD, Adam, RMSprop, Optimizer
 from colorama import Back, Style
 
-from colossalai.pipeline.rpc.PipelineBase import FillDrainPipelineEngine, OneFOneBPipelineEngine
-from colossalai.testing import assert_close
-
 
 def color_debug(text, prefix=' ', color='blue'):
     color = color.upper()
@@ -43,13 +39,13 @@ class RpcTestModel(nn.Module):
 
 def parse_args():
     parser = argparse.ArgumentParser()
+    parser.add_argument('--epoch', type=int, default=1)
     parser.add_argument('--world_size', type=int, default=2)
     parser.add_argument('--num_microbatches', type=int, default=2)
     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('--optimizer', type=str, choices=['SGD', 'Adam', 'RMSprop'], default='SGD')
-    parser.add_argument('--device', type=str, default='cuda')
+    parser.add_argument('--device', type=str, choices=['cpu', 'cuda'], default='cuda')
     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)

tests/test_pipeline/test_cuda_rpc_optimizer.py

-import os
-import argparse
-
 import torch
 from torch import nn
-import torch.multiprocessing as mp
-import torch.distributed.rpc as rpc
 from torch import autograd
 from torch.optim import SGD, Adam, RMSprop, Optimizer
-from colorama import Back, Style
 
 from colossalai.pipeline.rpc.PipelineBase import FillDrainPipelineEngine, OneFOneBPipelineEngine
 from colossalai.testing import assert_close
@@ -21,7 +15,6 @@ def run_master(args):
     stage_num = args.world_size
     chunk = args.chunk
     actual_stage_num = stage_num * chunk
-    use_interleave = args.use_interleave
     use_checkpoint = args.use_checkpoint
     num_microbatches = args.num_microbatches
     optimizer_class = globals()[args.optimizer]
@@ -45,7 +38,6 @@ def run_master(args):
                                     num_microbatches=num_microbatches,
                                     device=device,
                                     chunk=chunk,
-                                    use_interleave=use_interleave,
                                     checkpoint=use_checkpoint)
 
     engine.initialize_optimizer(optimizer_class, lr=lr)

tests/test_pipeline/test_cuda_rpc_pipeline.py

-import os
-import argparse
-
 import torch
 from torch import nn
-import torch.multiprocessing as mp
-import torch.distributed.rpc as rpc
 
 from colossalai.pipeline.rpc.PipelineBase import FillDrainPipelineEngine, OneFOneBPipelineEngine
 from rpc_test_utils import rpc_run, parse_args, RpcTestModel
@@ -13,12 +8,12 @@ from rpc_test_utils import rpc_run, parse_args, RpcTestModel
 def run_master(args):
     torch.manual_seed(100)
 
+    epoch = args.epoch
     device = args.device
     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
@@ -38,10 +33,10 @@ def run_master(args):
                                     num_microbatches=num_microbatches,
                                     device=device,
                                     chunk=chunk,
-                                    use_interleave=use_interleave,
                                     checkpoint=use_checkpoint)
 
-    _ = engine.forward_backward(input_sample)
+    for _ in range(epoch):
+        _ = engine.forward_backward(input_sample, forward_only=False)
 
 
 if __name__ == "__main__":

tests/test_pipeline/test_cuda_rpc_value_correctness.py

-import os
-import argparse
-
 import torch
 from torch import nn
-import torch.multiprocessing as mp
-import torch.distributed.rpc as rpc
 from torch import autograd
-from colorama import Back, Style
 
 from colossalai.pipeline.rpc.PipelineBase import FillDrainPipelineEngine, OneFOneBPipelineEngine
 from colossalai.testing import assert_close
@@ -20,7 +14,6 @@ def run_master(args):
     stage_num = args.world_size
     chunk = args.chunk
     actual_stage_num = stage_num * chunk
-    use_interleave = args.use_interleave
     use_checkpoint = args.use_checkpoint
     num_microbatches = args.num_microbatches
 
@@ -41,7 +34,6 @@ def run_master(args):
                                     num_microbatches=num_microbatches,
                                     device=device,
                                     chunk=chunk,
-                                    use_interleave=use_interleave,
                                     checkpoint=use_checkpoint)
 
     forward_result = engine.forward_backward(input_sample)