Request for Communication Performance Evaluation and Subset Parallelization Support

[huangzx02]

First of all, thank you for your amazing work on the nnScaler project. It has been incredibly inspiring, and I’ve been learning and using the contents from this repository in my own work.
I have a few suggestions and questions:

1. Performance Evaluation:

It would be great if you could provide some sample code or a reference implementation to evaluate the computation and communication performance with the profiling tool.

2. Subset Parallelization Support:

Following the README, I have successfully implemented end-to-end parallelization of an entire model. However, I’m wondering if it is possible to perform parallelization on a subset of the model and generate parallelized code only for that portion.
If this functionality already exists, could you provide documentation or examples on how to use it?

[zyeric]

Thanks for your suggestions.

1. You add the following code around training step to print communication cost in nnScaler. Note that if the allreduce operation is aync, the communication time may be inaccurate.

from nnscaler.profiler.timer import CudaTimer

for i, batches in data_iter:
    CudaTimer().clear()
    CudaTimer(enable=True, predefined=True).start('e2e')

    loss = model(batches)

    CudaTimer().stop('e2e')
    CudaTimer().print_all(times=1, rank_only=0)

2. By design, parallelize can accept any pytorch module as input. We will add examples in the future. By the way, could you provide a model architecture that requires a non-end-to-end usage of parallelize?

[huangzx02]

Thanks for the reply. The second problem aims to use nnScaler to generate parallelized pytorch code based on given partition rules automatically. For example, I have a neural network model that contains two linear layers. I would like to apply a parallelization strategy to specific tensors, such as splitting the weights of the first linear layer along the 1st dimension and the second linear layer along the 2nd dimension. How should I input my sub-model and the corresponding splitting strategy to generate the tensor parallel code for these submodules?

[zyeric]

nnscaler's plan searcher autodist has an input parameter partition_constraints_path that can control the partition strategy of a specified torch module by its name. We will add its doc soon.

In autodist, we provide a set of partition constraints to control the distributed plan of the model.

Here is an example that will generate a data parallel plan for retnet

- allowed_partition_dims:
  - 0,0
  name: torch.nn.functional.linear
  parent_module: MultiScaleRetention
  replica_allowed: false
- allowed_partition_dims:
  - 0,0
  name: cube.runtime.function.embedding
  parent_module: LMDecoder
  replica_allowed: false
- allowed_partition_dims:
  - 0,0
  name: apex.normalization.fused_layer_norm.FusedLayerNormAffineFunction.apply
  parent_module: LMDecoder
  replica_allowed: false
- allowed_partition_dims:
  - 0,0
  - 0,1
  name: torch.matmul
  parent_module: MultiScaleRetention
  replica_allowed: false
- allowed_partition_dims:
  - 0,0
  name: torch.nn.functional.linear
  parent_module: FeedForwardNetwork
  replica_allowed: false

Each partition constraint has 4 fields: name, parent_module, allowed_partition_dims, and replica_allowed.

• name is the name of the corresponding operator in the model. It equals to the signature field in the IRFwOperation in cube. Note: signature is the full name of the operator, for example, you should provide torch.nn.functional.linear instead of linear.
• parent_module is the closest father module name of the operator. You can provide two partition constraints with a same name but different module to control the partition of the same operator in different modules. Similar to recompute_modules, Module name can be any suffix of the full module name, e.g., module1 will match x.module1, y.module1, x.module1 will match x.module1 but not y.module1.
• allowed_partition_dims is a list of allowed partition dimensions of input tensors. Each element in the list is a list of two integers, which are the index of the partitioned tensor among inputs and the partitioned dimension of that tensor. An empty list means no partition is allowed, note that in yaml, you should give an empty list explicitly, i.e., allowed_partition_dims: [].
• replica_allowed is a boolean value. If it is true, it is allowed to replicate the operator across devices.

[huangzx02]

Thank you very much, I'm looking forward to the following update.

[zyeric]

You can check examples by our mini-trainer under https://github.com/microsoft/nnscaler/tree/main/examples/, thanks.