Add pipeline shared module wrapper

This feature is especially useful when train GPT/BERT whose word embedding layer is used at the first and the end.

Usage:

pipeline_size = gpc.get_world_size(ParallelMode.PIPELINE)
pipeline_rank = gpc.get_local_rank(ParallelMode.PIPELINE)
rank = gpc.get_global_rank()
wrapper = PipelineSharedModuleWrapper([0, pipeline_size - 1])
parts = _partition_uniform(num_layers, pipeline_size, num_chunks)[pipeline_rank]
models = []
for start, end in parts:
    kwargs['num_layers'] = end - start
    kwargs['first'] = start == 0
    kwargs['last'] = end == num_layers
    print(f'==> Rank{rank} build layer {start}-{end}, total {num_layers}')
    chunk = PipelineGPT1D(**kwargs).to(device)
    if start == 0:
        wrapper.register_module(chunk.embedding.word_embeddings)
    elif end == num_layers:
        wrapper.register_module(chunk.head)
    models.append(chunk)
if len(models) == 1:
    model = models[0]
else:
    model = nn.ModuleList(models)

PipelineSharedModuleWrapper must be initialized in all ranks, and PipelineSharedModuleWrapper.register_module() should be called in the ranks which share the module. Modules have to be moved to corresponding device based on your distributed backend before calling register_module().

Update load_batch() of schedule

We update the rule of loading batch in schedule to support GPT/BERT training with pipeline parallelism.

Now please make sure the item your dataset returned is a tuple of (data, label), and the type of data and label must be torch.Tensor or dict. When you set sync_data to True in schedule, you must make sure the values of the dict are torch.Tensor. Note that when your dataset returns dict, the keys must be the same as arguments in your model.forward() or loss_function.forward().

When using pipeline parallelism, the input of first layer is from dataloader. For other layers, the first argument of forward is the output of the previous pipeline stage and other arguments are from dataloader. Note that each layer can only return one tensor in forward().

Optimize GPU memory usage of pipeline schedule

Add a argument (return_output_label) in schedule.forward_backward_step, trainer.fit and trainer.evaluate. The output of model and labels won't be returned, which can further reduce GPU memory usage especially when using pipeline parallelism.

Optimize loss accumulation in pipeline schedule. Use loss.detach() when accumulating it to avoid unexpected large memory usage.

Example:

trainer.fit(
    train_dataloader=train_dataloader,
    epochs=num_epochs,
    test_interval=1,
    hooks=hook_list,
    display_progress=True,
    return_output_label=False
)

Reduce communication of pipeline schedule

Add a argument tensor_shape for PipelineSchedule and InterleavedPipelineSchedule. You can set this argument to a Union[torch.Size, List[int], Tuple[int]] if the tensor shapes transmitted along pipeline are the same and fixed during training. By setting this, the communication will be further reduced.

Example:

schedule = InterleavedPipelineSchedule(num_micro_batches,
                                               num_model_chunks, tensor_shape=tensor_shape)