Skip to content

Add sequence packing support for hybrid model #2913

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
duncanriach merged 24 commits into from
Jan 16, 2026
Merged

Add sequence packing support for hybrid model #2913

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
941c056
Add sequence packing support for hybrid model
duncanriach Sep 25, 2025
e84041a
Fix for packed_seq + CP>1 + PP>1
duncanriach Oct 7, 2025
2c24331
Fix for packed_seq + PP>2
duncanriach Oct 13, 2025
5064681
Fix packed sequence info broadcast for PP>1
duncanriach Nov 8, 2025
38c9bc9
Resolve conflict with MR 3963 -> MR 4013
duncanriach Dec 4, 2025
5aab3b2
Prevent training with packed sequences when use_mem_eff_path==False
duncanriach Dec 4, 2025
aa9c757
Remove statistics code
duncanriach Dec 5, 2025
cdffa88
Improve function signatures wrt packed_seq_params
duncanriach Dec 5, 2025
b18c657
Add error check to prevent SFTDataset being used with GPTModel
duncanriach Dec 5, 2025
308ff65
Improve comment
duncanriach Dec 5, 2025
64941be
Confirm that mamba-ssm and causal-conv1d support sequence packing
duncanriach Dec 8, 2025
ac7bafd
Fix packed sequence changes
duncanriach Dec 11, 2025
191aab1
Add mamba and conv1d version checking for training
duncanriach Dec 11, 2025
2a0d5ad
Encapsulate utility function better. Add commit references
duncanriach Dec 12, 2025
864001b
Fix small bug in get_batch
duncanriach Dec 25, 2025
5955579
Make MambaMixer methods private
duncanriach Dec 25, 2025
9fd2216
Fix typo in refactoring
duncanriach Dec 30, 2025
6230b63
Add unit test for hybrid model packed sequence
duncanriach Jan 12, 2026
5949222
Remove comment
duncanriach Jan 12, 2026
75004f0
Fix formatting errors
duncanriach Jan 15, 2026
94aaa51
Merge branch 'main' into duncan/hybrid-packed-sequence-for-main
ericharper Jan 15, 2026
ebf2d31
Merge branch 'main' into duncan/hybrid-packed-sequence-for-main
duncanriach Jan 15, 2026
fc28a1e
Fix GPT error check for packed sequence
duncanriach Jan 16, 2026
a2f0e20
Merge branch 'main' into duncan/hybrid-packed-sequence-for-main
Phlip79 Jan 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions megatron/core/datasets/gpt_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,9 @@ class GPTDatasetConfig(BlendedMegatronDatasetConfig):
Check --per-dataset-sequences-path
"""

context_parallel_size: Optional[int] = None
"""The size of the context parallel group. Needed for padding in packed sequences."""

Comment thread
duncanriach marked this conversation as resolved.
def __post_init__(self) -> None:
"""Do asserts and set fields post init"""
super().__post_init__()
Expand Down
11 changes: 6 additions & 5 deletions megatron/core/models/mamba/mamba_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,9 +191,6 @@ def forward(
processing layer (optional).

It either returns the Loss values if labels are given or the final hidden units

packed_seq_params is unused but included to maintain compatibility with
GPTModel's forward signature.
"""
# If decoder_input is provided (not None), then input_ids and position_ids are ignored.
# Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input.
Expand Down Expand Up @@ -227,9 +224,12 @@ def forward(
rotary_pos_emb = None
if self.position_embedding_type == 'rope':
rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
inference_context, self.decoder, decoder_input, self.config
inference_context, self.decoder, decoder_input, self.config, packed_seq_params
)
rotary_pos_emb = self.rotary_pos_emb(
rotary_seq_len,
packed_seq=packed_seq_params is not None and packed_seq_params.qkv_format == 'thd',
)
rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len)

# Wrap decoder_input to allow the decoder (MambaBlock) to delete the
# reference held by this caller function, enabling early garbage collection
Expand All @@ -253,6 +253,7 @@ def forward(
attention_mask=attention_mask,
inference_context=inference_context,
rotary_pos_emb=rotary_pos_emb,
packed_seq_params=packed_seq_params,
)

if not self.post_process:
Expand Down
31 changes: 10 additions & 21 deletions megatron/core/models/multimodal/llava_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -924,27 +924,16 @@ def forward(
)
)

if isinstance(self.language_model, MambaModel):
output = self.language_model(
input_ids=None,
position_ids=None,
attention_mask=attention_mask,
decoder_input=combined_embeddings,
labels=new_labels,
inference_context=inference_context,
runtime_gather_output=runtime_gather_output,
)
else:
output = self.language_model(
input_ids=None,
position_ids=None,
attention_mask=attention_mask,
decoder_input=combined_embeddings,
labels=new_labels,
inference_context=inference_context,
runtime_gather_output=runtime_gather_output,
packed_seq_params=packed_seq_params,
)
output = self.language_model(
input_ids=None,
position_ids=None,
attention_mask=attention_mask,
decoder_input=combined_embeddings,
labels=new_labels,
inference_context=inference_context,
runtime_gather_output=runtime_gather_output,
packed_seq_params=packed_seq_params,
)

return output, new_loss_mask

Expand Down
4 changes: 4 additions & 0 deletions megatron/core/ssm/mamba_block.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from megatron.core.extensions.transformer_engine import TENorm
from megatron.core.fp8_utils import get_fp8_context
from megatron.core.inference.contexts import BaseInferenceContext
from megatron.core.packed_seq_params import PackedSeqParams
from megatron.core.process_groups_config import ProcessGroupCollection
from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols as LayerSymbols
from megatron.core.ssm.mamba_hybrid_layer_allocation import allocate_layers
Expand Down Expand Up @@ -206,6 +207,7 @@ def forward(
rotary_pos_emb: Optional[Tensor] = None,
*,
inference_params: Optional[BaseInferenceContext] = None,
packed_seq_params: Optional[PackedSeqParams] = None,
):
"""
Forward function of the MambaStack class.
Expand Down Expand Up @@ -287,12 +289,14 @@ def forward(
inference_context=inference_context,
rotary_pos_emb=rotary_pos_emb,
sequence_len_offset=sequence_len_offset,
packed_seq_params=packed_seq_params,
)
else: # MambaLayer
hidden_states = layer(
hidden_states=hidden_states,
attention_mask=attention_mask,
inference_context=inference_context,
packed_seq_params=packed_seq_params,
)

# The attention layer (currently a simplified transformer layer)
Expand Down
120 changes: 92 additions & 28 deletions megatron/core/ssm/mamba_context_parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,14 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

from megatron.core.packed_seq_params import PackedSeqParams
from megatron.core.tensor_parallel import all_to_all
from megatron.core.utils import is_te_min_version

try:
from einops import repeat
Expand All @@ -13,6 +17,16 @@
except ImportError:
HAVE_EINOPS = False

try:
# Register the TE CUDA kernels
import transformer_engine # pylint: disable=unused-import

# Alias the PyTorch wrapper so we can call tex.* APIs
import transformer_engine_torch as tex
except ImportError:
# TE isn’t installed or the torch wrapper is missing
tex = None


class MambaContextParallel:
"""
Expand Down Expand Up @@ -116,7 +130,9 @@ def __init__(
# and also `nheads_local_tpcp = nheads_local_tp // cp_size` whilst ngroups_local_tpcp is
# either 1 or `ngroups_local_tp // cp_size`

def pre_conv_ssm(self, input_: torch.Tensor) -> torch.Tensor:
def pre_conv_ssm(
self, input_: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None
) -> torch.Tensor:
"""Method to be applied before the convolution and SSM"""
if self.cp_size == 1:
return input_
Expand Down Expand Up @@ -171,17 +187,20 @@ def pre_conv_ssm(self, input_: torch.Tensor) -> torch.Tensor:

output = torch.cat([z, x, B, C, dt], dim=-1)
# TODO(duncan): for hybrid models, consider isolating load-balancing to attention layers
output = _undo_attention_load_balancing(output, self.cp_size)
output = _undo_attention_load_balancing(output, self.cp_size, packed_seq_params)

return output

def post_conv_ssm(self, input_: torch.Tensor) -> torch.Tensor:
def post_conv_ssm(
self, input_: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None
) -> torch.Tensor:
"""Method to be applied after the convolution and SSM"""
if self.cp_size == 1:
return input_
else:
return _all_to_all_hp2cp(
_redo_attention_load_balancing(input_, self.cp_size), self.cp_group
_redo_attention_load_balancing(input_, self.cp_size, packed_seq_params),
self.cp_group,
)

def conv1d(self, input_: torch.Tensor) -> torch.Tensor:
Expand Down Expand Up @@ -357,33 +376,78 @@ def _all_to_all_hp2cp(
return output


def _undo_attention_load_balancing(input_: torch.Tensor, cp_size: int) -> torch.Tensor:
def _undo_attention_load_balancing(
input_: torch.Tensor, cp_size: int, packed_seq_params: Optional[PackedSeqParams] = None
) -> torch.Tensor:
"""
Undoes the context parallel attention load balancing
For example, for cp_size=3, converts 162534 to 123456 for sequential
processing by the convolution and SSM.
Undoes the context parallel attention load balancing.
For example (non-packed), for cp_size=3, converts 162534 to 123456 for
sequential processing by the convolution and SSM.
"""
num_chunks_div_2 = cp_size
num_chunks = num_chunks_div_2 * 2
chunks = torch.chunk(input_, chunks=num_chunks, dim=0)
order = [2 * i for i in range(num_chunks_div_2)] + [
num_chunks - 2 * i - 1 for i in range(num_chunks_div_2)
]
reordered_chunks = [chunks[i] for i in order]
return torch.cat(reordered_chunks, dim=0)
if packed_seq_params is None:
num_chunks_div_2 = cp_size
num_chunks = num_chunks_div_2 * 2
chunks = torch.chunk(input_, chunks=num_chunks, dim=0)
order = [2 * i for i in range(num_chunks_div_2)] + [
num_chunks - 2 * i - 1 for i in range(num_chunks_div_2)
]
reordered_chunks = [chunks[i] for i in order]
return torch.cat(reordered_chunks, dim=0)
else:
assert tex is not None and is_te_min_version("1.10.0"), (
"Please update Transformer Engine to >= 1.10 to use "
"Context Parallel with THD format data"
)
if packed_seq_params.cu_seqlens_q_padded is not None:
cu_seqlens = packed_seq_params.cu_seqlens_q_padded
else:
cu_seqlens = packed_seq_params.cu_seqlens_q
total_tokens = input_.size(0)
assert total_tokens % cp_size == 0
seqlen_per_rank = total_tokens // cp_size
output = torch.empty_like(input_)
for cp_rank in range(cp_size):
start = cp_rank * seqlen_per_rank
end = start + seqlen_per_rank
index = tex.thd_get_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank)
output[index] = input_[start:end]
return output


def _redo_attention_load_balancing(input_: torch.Tensor, cp_size: int) -> torch.Tensor:
def _redo_attention_load_balancing(
input_: torch.Tensor, cp_size: int, packed_seq_params: Optional[PackedSeqParams] = None
) -> torch.Tensor:
"""
Redo the context parallel attention load balancing
For example, for cp_size=3, converts 123456 to 162534 for efficient
processing by attention.
Redo the context parallel attention load balancing.
For example (non-packed), for cp_size=3, converts 123456 to 162534 for
efficient processing by attention.
"""
num_chunks_div_2 = cp_size
num_chunks = num_chunks_div_2 * 2
chunks = torch.chunk(input_, chunks=num_chunks, dim=0)
order = [None] * num_chunks
order[::2] = range(num_chunks_div_2) # order[even]
order[1::2] = reversed(range(num_chunks_div_2, num_chunks)) # order[odd]
reordered_chunks = [chunks[i] for i in order]
return torch.cat(reordered_chunks, dim=0)
if packed_seq_params is None:
num_chunks_div_2 = cp_size
num_chunks = num_chunks_div_2 * 2
chunks = torch.chunk(input_, chunks=num_chunks, dim=0)
order = [None] * num_chunks
order[::2] = range(num_chunks_div_2) # order[even]
order[1::2] = reversed(range(num_chunks_div_2, num_chunks)) # order[odd]
reordered_chunks = [chunks[i] for i in order]
return torch.cat(reordered_chunks, dim=0)
else:
assert tex is not None and is_te_min_version("1.10.0"), (
"Please update Transformer Engine to >= 1.10 to use "
"Context Parallel with THD format data"
)
if packed_seq_params.cu_seqlens_q_padded is not None:
cu_seqlens = packed_seq_params.cu_seqlens_q_padded
else:
cu_seqlens = packed_seq_params.cu_seqlens_q
total_tokens = input_.size(0)
assert total_tokens % cp_size == 0
seqlen_per_rank = total_tokens // cp_size
index = torch.empty(total_tokens, device=input_.device, dtype=torch.int32)
for cp_rank in range(cp_size):
start = cp_rank * seqlen_per_rank
end = start + seqlen_per_rank
index[start:end] = tex.thd_get_partitioned_indices(
cu_seqlens, total_tokens, cp_size, cp_rank
)
return input_.index_select(0, index)
6 changes: 5 additions & 1 deletion megatron/core/ssm/mamba_layer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
from megatron.core.dist_checkpointing.mapping import ShardedStateDict
from megatron.core.dist_checkpointing.utils import apply_prefix_mapping
from megatron.core.inference.contexts import BaseInferenceContext
from megatron.core.packed_seq_params import PackedSeqParams
from megatron.core.process_groups_config import ProcessGroupCollection
from megatron.core.transformer.identity_op import IdentityOp
from megatron.core.transformer.module import GraphableMegatronModule
Expand Down Expand Up @@ -96,6 +97,7 @@ def forward(
rotary_pos_emb: Optional[Tensor] = None, # Not used in MambaLayer
*,
inference_params: Optional[BaseInferenceContext] = None,
packed_seq_params: Optional[PackedSeqParams] = None,
):
"""
Perform a forward pass through the Mamba layer.
Expand Down Expand Up @@ -124,7 +126,9 @@ def forward(
hidden_states = hidden_states.to(dtype=self.config.params_dtype)
hidden_states = self.norm(hidden_states)

mixer_out_with_bias = self.mixer(hidden_states, inference_context=inference_context)
mixer_out_with_bias = self.mixer(
hidden_states, inference_context=inference_context, packed_seq_params=packed_seq_params
)

with self.bias_dropout_add_exec_handler():
hidden_states = self.mamba_bda(
Expand Down
Loading
Loading