NVIDIA / NVIDIA/Megatron-LM

Empty pipeline stages deadlock with custom pipeline-parallel-layout

Open
#1,852 0 comments 0 reactions 1 assignee Claimed by @gautham-kollu View on GitHub
bug community-request module: moe
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.