[Feature] A systematic tensor parallelism verifier
- Dominant language
- Python
- Stars
- 153
- Forks
- 17
- PR merge metrics
- No merged PRs in 30d
Description
We need a systematic approach to verify whether a scheduled model with tensor parallelism still produces the same results (i.e., outputs and gradients). It should check the following:
- Whether `.shard` and `.sync` are correctly specified to maintain the **shape correctness**. For example, shard the weight of a linear layer by its output feature dimension results in partitioned outputs. In this case, we either need an `all-gather` right after the linear, or the next linear must shard its weight by input feature dimension.
- Whether `.shard` and `.sync` are correctly specified to maintain the **functional correctness**. For example, shard the weight of a linear layer by its input feature dimension results in the same shape but partial sum outputs. In this case, `all-reduce` is required.
- Whether the random seeds of each dropout is property configured. In the case that the input tensor of dropout is not partitioned (i.e., replica and partial sum), the random seed on each device should be the same, because all devices are supposed to do redundant computations. On the other hand, if the input tensor is partitioned (i.e., the shape on each device is divided by TP group size), the random seed on each device should be different to avoid repeated dropout patterns that hurt the convergence.
We use this issue to discuss possible solutions and track the progress. At the first glance, we may consider a compiler approach that perform type inference on static graphs. We could use TorchDynamo or LazyTensor to capture static graph and apply this analysis. However, this approach won't work if the module cannot be captured as a static graph due to coding style limitations.
Contributor guide
Research direction
No files, tests, or entry points are named. Start by locating the scheduling behavior for .shard and .sync, then review how tensor shapes, partial sums, all-gather, all-reduce, and dropout seeds are represented; done means a systematic verifier can detect shape and functional mismatches for outputs and gradients.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 28/100