update data flow for auto parallel (#105)

* update data flow for auto parallel

* not scale_samples by default for acceleration

Author: ghostplant

tests/test_tutel.py

@@ -37,7 +37,9 @@ def run(
         if helloworld_file == 'helloworld':
             command = 'python3 -m torch.distributed.launch --nproc_per_node=' + str(nproc_per_node) + ' tutel/examples/helloworld.py --top ' + str(top) + ' --dtype ' + dtype + ' --num_local_experts ' + str(num_local_experts) + ' --hidden_size ' + str(hidden_size) + ' --batch_size ' + str(batch_size) + ' --a2a_ffn_overlap_degree ' + str(a2a_ffn_overlap_degree) + ' --num_steps ' + str(num_steps)
             if use_model_parallel:
-                command += ' --use_model_parallel'
+                command += ' --parallel_type model'
+            else:
+                command += ' --parallel_type data'
 
         p = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=new_env)
         losses = []

tutel/examples/helloworld.py

@@ -32,7 +32,7 @@
 parser.add_argument('--l_aux_wt', type=float, default=0.0)
 parser.add_argument('--a2a_ffn_overlap_degree', type=int, default=1)
 parser.add_argument('--num_steps', type=int, default=100)
-parser.add_argument('--use_model_parallel', default=False, action='store_true')
+parser.add_argument('--parallel_type', type=str, default='auto')
 parser.add_argument('--save_load_checkpoint', default=False, action='store_true')
 args = parser.parse_args()
 
@@ -72,8 +72,8 @@ def __init__(self):
             scan_expert_func = lambda name, param: setattr(param, 'skip_allreduce', True),
             seeds = (1, dist_rank + 1, 1),
             a2a_ffn_overlap_degree = a2a_ffn_overlap_degree,
-            use_model_parallel = args.use_model_parallel,
-        ).to(device)
+            parallel_type = args.parallel_type,
+        )
 
         # Summary of different parameter types: gate, local_experts
         local_count = sum([torch.numel(param) for name, param in self._moe_layer.get_parameter_iterator(param_type='local_experts')])
@@ -85,7 +85,7 @@ def forward(self, input):
         result = F.log_softmax(torch.sum(result, dim=2), dim=1)
         return result
 
-model = ExampleModel()
+model = ExampleModel().to(device)
 dist_print(model)
 
 if args.save_load_checkpoint:
@@ -101,8 +101,8 @@ def forward(self, input):
 x = torch.tensor(torch.randn([batch_size, num_tokens, model_dim], dtype=torch.float32, device='cpu').detach().numpy(), dtype=torch.get_default_dtype(), requires_grad=True, device=device)
 y = torch.LongTensor(batch_size).random_(1).to(device)
 
-tuples = (dist_world_size, args.dtype, model_dim, hidden_size, batch_size * num_tokens, num_local_experts, top_value, a2a_ffn_overlap_degree, args.use_model_parallel)
-dist_print('[Benchmark] world_size = %s, dtype = %s, model_dim = %s, hidden_size = %s, samples = %s, num_local_experts = %s, topK = %s, a2a_ffn_overlap_degree = %s, use_model_parallel = %s' % tuples)
+tuples = (dist_world_size, args.dtype, model_dim, hidden_size, batch_size * num_tokens, num_local_experts, top_value, a2a_ffn_overlap_degree, args.parallel_type)
+dist_print('[Benchmark] world_size = %s, dtype = %s, model_dim = %s, hidden_size = %s, samples = %s, num_local_experts = %s, topK = %s, a2a_ffn_overlap_degree = %s, parallel_type = `%s`' % tuples)
 
 average_time, num_steps = 0, args.num_steps
 

tutel/examples/helloworld_amp.py

@@ -71,7 +71,7 @@ def __init__(self):
             scan_expert_func = lambda name, param: setattr(param, 'skip_allreduce', True),
             seeds = (1, dist_rank + 1, 1),
             a2a_ffn_overlap_degree = a2a_ffn_overlap_degree,
-        ).to(device)
+        )
 
         # Summary of different parameter types: gate, local_experts
         local_count = sum([torch.numel(param) for name, param in self._moe_layer.get_parameter_iterator(param_type='local_experts')])
@@ -83,7 +83,7 @@ def forward(self, input):
         result = F.log_softmax(torch.sum(result, dim=2), dim=1)
         return result
 
-model = ExampleModel()
+model = ExampleModel().to(device)
 dist_print(model)
 
 optimizer = torch.optim.SGD(model.parameters(), lr=1e-5)

tutel/examples/helloworld_ddp.py

@@ -72,7 +72,7 @@ def __init__(self):
             scan_expert_func = lambda name, param: setattr(param, 'skip_allreduce', True),
             seeds = (1, dist_rank + 1, 1),
             a2a_ffn_overlap_degree = a2a_ffn_overlap_degree,
-        ).to(device)
+        )
 
         # Summary of different parameter types: gate, local_experts
         local_count = sum([torch.numel(param) for name, param in self._moe_layer.get_parameter_iterator(param_type='local_experts')])
@@ -88,7 +88,7 @@ def add_param_to_skip_allreduce(self, param_name):
         self._ddp_params_and_buffers_to_ignore.append(param_name)
 
 
-model = ExampleModel()
+model = ExampleModel().to(device)
 
 for name, param in model.named_parameters():
     if hasattr(param, 'skip_allreduce'):

tutel/examples/helloworld_deepspeed.py

@@ -97,7 +97,7 @@ def __init__(self):
                 num_experts = num_local_experts * dist_world_size,
                 k = top_value,
                 use_tutel = args.use_tutel
-        ).to(device)
+        )
 
         for name, param in self._moe_layer.named_parameters():
             if '.experts.' in name:
@@ -113,7 +113,7 @@ def forward(self, input):
         result = F.log_softmax(torch.sum(result, dim=2), dim=1)
         return result
 
-model = ExampleModel()
+model = ExampleModel().to(device)
 dist_print(model)
 
 optimizer = torch.optim.SGD(model.parameters(), lr=1e-5)

tutel/examples/helloworld_megatron.py

@@ -64,7 +64,7 @@ def __init__(self):
             model_dim = model_dim,
             scan_expert_func = lambda name, param: setattr(param, 'skip_allreduce', True),
             seeds = (1, dist_rank + 1, 1),
-        ).to(device)
+        )
 
         # Summary of different parameter types: gate, local_experts
         local_count = sum([torch.numel(param) for name, param in self._moe_layer.get_parameter_iterator(param_type='local_experts')])
@@ -76,7 +76,7 @@ def forward(self, input):
         result = F.log_softmax(torch.sum(result, dim=2), dim=1)
         return result
 
-model = ExampleModel()
+model = ExampleModel().to(device)
 dist_print(model)
 
 optimizer = torch.optim.SGD(model.parameters(), lr=1e-5)

tutel/examples/helloworld_sharded_experts.py

@@ -66,7 +66,7 @@ def __init__(self):
             model_dim = model_dim,
             scan_expert_func = lambda name, param: setattr(param, 'skip_allreduce', True),
             seeds = (1, dist_rank + 1, 1),
-        ).to(device)
+        )
 
         # Summary of different parameter types: gate, local_experts
         local_count = sum([torch.numel(param) for name, param in self._moe_layer.get_parameter_iterator(param_type='local_experts')])
@@ -78,7 +78,7 @@ def forward(self, input):
         result = F.log_softmax(torch.sum(result, dim=2), dim=1)
         return result
 
-model = ExampleModel()
+model = ExampleModel().to(device)
 dist_print(model)
 
 optimizer = torch.optim.SGD(model.parameters(), lr=1e-5)