Empty pipeline stages deadlock with custom pipeline-parallel-layout
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
**Describe the bug**
A clear and concise description of what the bug is.
Assume you have a model with high VPP (like an MOE) and a tokenizer / embedding layer you want to isolate to one node. In this case, you may want no transformer layers there except for the embedding/tokenizer layer. similar for MTP or any other loss layers. Assume for the following example you have 16 physical pipeline stages, and two virtual stages per node. I wawnt the last physical pipeline stage to only have the loss layer.
The way to represent this currently to even divide the transformers layers in a DeepSeek style model across all the pipeline stages roughly evenly is ` --pipeline-model-parallel-layout "Et|(tt|)*30L' for a 61 stage DeepSeek model, but what if I want to not have any other transformers on the last pipeline stage. "Et|(tt|)*29L|" errors predictably, but "Et|(tt|)*29|L" also errors. The best we could do in this case is t "Et|(tt|)*28t|t|L" but that still has an MOE layer on the last physical stage. Having an odd number of MOE transformer layers also seems to cause A2A_Overlap to deadlock as reported in an other issue #1810
**Steps/Code to reproduce bug**
Please list *minimal* steps or code snippet for us to be able to reproduce the bug.
A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.
Try to create a DeepSeekV3 model with the previous pipeline stage layout and optionally A2A_Overlap, watch it hang on the first step until it times out.
**Expected behavior**
* Either an assertion error at least or better yet properly support empty pipeline stages to increase the expressivity and allow different number of VPP stages per physical pipelineing layer.
A clear and concise description of what you expected to happen.
**Additional context**
NGC Nemo 25.07
Add any other context about the problem here.
Contributor guide
Assessment
This issue has not been assessed yet.