NVIDIA / NVIDIA/Megatron-LM

[Docs-only] Clarify seq_length behavior when variable_seq_lengths=True under pipeline parallelism (PP>1)

Open
#2,064 0 comments 0 reactions 0 assignees View on GitHub
enhancement module: documentation
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

*Is your feature request related to a problem? Please describe.*

In megatron/core/pipeline_parallel/schedules.py#get_forward_backward_func, current comments/docstrings say seq_length is ignored when variable_seq_lengths=True. That’s correct for PP=1, but for PP>1 the pipelined schedules still use seq_length to size P2P activation tensors—e.g., forward_backward_pipelining_without_interleaving(...) calls get_tensor_shapes(seq_length, ...), and forward_backward_pipelining_with_interleaving(...) builds tensor_shape = [seq_length, micro_batch_size, hidden_size]. This mismatch confuses users and can cause shape errors (if seq_length is too small) or wasted memory/bandwidth (if too large).

*Describe the solution you'd like*

Documentation-only: update comments/docstrings near get_forward_backward_func (and the two pipelined schedule functions) to state that with variable_seq_lengths=True PP=1 ignores seq_length, while PP>1 requires it as the per-step maximum sequence length used to size P2P tensors (actual microbatches may be ≤ seq_length).

Contributor guide

Open the contributing guide

Research direction

Read megatron/core/pipeline_parallel/schedules.py, starting at get_forward_backward_func and the forward_backward_pipelining_without_interleaving and forward_backward_pipelining_with_interleaving functions. Update the nearby comments or docstrings to distinguish PP=1 from PP>1 behavior when variable_seq_lengths=True, including that PP>1 uses the per-step maximum sequence length for P2P tensor sizing.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, documentation, machine-learning
Issue type
Documentation
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.