Add support for packed thd in BERT language module.

megatron/core/models/bert/bert_model.py

@@ -14,6 +14,7 @@
 from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
 from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
 from megatron.core.models.common.language_module.language_module import LanguageModule
+from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.attention import SelfAttentionSubmodules
 from megatron.core.transformer.dot_product_attention import (
@@ -269,8 +270,35 @@ def bert_extended_attention_mask(self, attention_mask: Tensor) -> Tensor:
 
         return extended_attention_mask
 
-    def bert_position_ids(self, token_ids):
+    def bert_position_ids(
+        self, token_ids: Tensor, packed_seq_params: PackedSeqParams | None = None
+    ) -> Tensor:
         """Position ids for bert model"""
+        if packed_seq_params is not None:
+            if token_ids.size(0) != 1:
+                raise ValueError('Packed BERT input should use dummy batch size 1')
+            if packed_seq_params.cu_seqlens_q is None:
+                raise ValueError(
+                    'packed_seq_params.cu_seqlens_q must be provided for packed BERT input'
+                )
+            cu_seqlens = packed_seq_params.cu_seqlens_q.to(device=token_ids.device)
+            if cu_seqlens.dim() != 1 or cu_seqlens.numel() < 2:
+                raise ValueError('packed_seq_params.cu_seqlens_q must be a 1D tensor')
+            if cu_seqlens[0].item() != 0:
+                raise ValueError('packed_seq_params.cu_seqlens_q must start at 0')
+            if cu_seqlens[-1].item() != token_ids.size(1):
+                raise ValueError(
+                    'packed_seq_params.cu_seqlens_q must end at the packed sequence length'
+                )
+            seq_lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.long)
+            if torch.any(seq_lengths < 0):
+                raise ValueError('packed_seq_params.cu_seqlens_q must be monotonically increasing')
+            seq_starts = torch.repeat_interleave(cu_seqlens[:-1].to(torch.long), seq_lengths)
+            token_positions = torch.arange(
+                token_ids.size(1), dtype=torch.long, device=token_ids.device
+            )
+            return (token_positions - seq_starts).unsqueeze(0)
+
         # Create position ids
         seq_length = token_ids.size(1)
         position_ids = torch.arange(seq_length, dtype=torch.long, device=token_ids.device)
@@ -297,10 +325,11 @@ def set_input_tensor(self, input_tensor: Tensor) -> None:
     def forward(
         self,
         input_ids: Tensor,
-        attention_mask: Tensor,
+        attention_mask: Tensor | None,
         tokentype_ids: Tensor = None,
         lm_labels: Tensor = None,
         inference_context=None,
+        packed_seq_params: PackedSeqParams | None = None,
         *,
         inference_params: Optional[BaseInferenceContext] = None,
     ):
@@ -315,11 +344,20 @@ def forward(
 
         inference_context = deprecate_inference_params(inference_context, inference_params)
 
-        extended_attention_mask = self.bert_extended_attention_mask(attention_mask)
+        if packed_seq_params is None:
+            if attention_mask is None:
+                raise ValueError('attention_mask must be provided when packed_seq_params is None')
+            extended_attention_mask = self.bert_extended_attention_mask(attention_mask)
+        else:
+            if packed_seq_params.qkv_format != 'thd':
+                raise ValueError("BERT packed sequence support requires qkv_format='thd'")
+            if attention_mask is not None:
+                raise ValueError('attention_mask must be None when using packed BERT input')
+            extended_attention_mask = None
 
         if parallel_state.is_pipeline_first_stage():
             input_ids = input_ids
-            position_ids = self.bert_position_ids(input_ids)
+            position_ids = self.bert_position_ids(input_ids, packed_seq_params)
         else:
             position_ids = None
             input_ids = None
@@ -338,16 +376,21 @@ 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.encoder, encoder_input, self.config
+                inference_context, self.encoder, encoder_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',
+                cp_group=packed_seq_params.cp_group if packed_seq_params is not None else None,
             )
-            rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len)
 
         # Run encoder.
         hidden_states = self.encoder(
             hidden_states=encoder_input,
             attention_mask=extended_attention_mask,
             inference_context=inference_context,
             rotary_pos_emb=rotary_pos_emb,
+            packed_seq_params=packed_seq_params,
         )
         if not self.post_process:
             return hidden_states

tests/unit_tests/models/test_bert_model.py

@@ -12,6 +12,7 @@
     get_bert_layer_with_transformer_engine_submodules,
 )
 from megatron.core.models.bert.bert_model import BertModel
+from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer.enums import AttnBackend, AttnMaskType
 from megatron.core.transformer.spec_utils import ModuleSpec
@@ -128,6 +129,132 @@ def test_qk_layernorm_from_config_fallback(self):
         assert isinstance(attention.q_layernorm, te_pytorch.LayerNorm)
         assert isinstance(attention.k_layernorm, te_pytorch.LayerNorm)
 
+    @pytest.mark.internal
+    def test_packed_forward_uses_cu_seqlens_positions_and_no_attention_mask(self, mocker):
+        config = TransformerConfig(
+            num_layers=2,
+            hidden_size=12,
+            num_attention_heads=4,
+            use_cpu_initialization=True,
+            perform_initialization=True,
+            attention_backend=AttnBackend.unfused,
+        )
+        bert_model = BertModel(
+            config=config,
+            num_tokentypes=0,
+            transformer_layer_spec=get_bert_layer_with_transformer_engine_spec(),
+            vocab_size=100,
+            max_sequence_length=4,
+            position_embedding_type='rope',
+            post_process=False,
+        )
+        sequence_length = 6
+        input_ids = torch.arange(sequence_length, dtype=torch.int64).unsqueeze(0)
+        cu_seqlens = torch.tensor([0, 2, 6], dtype=torch.int32)
+        packed_seq_params = PackedSeqParams(
+            qkv_format='thd',
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_kv=cu_seqlens,
+            max_seqlen_q=2,
+            max_seqlen_kv=4,
+        )
+        encoder_input = torch.ones(sequence_length, 1, config.hidden_size)
+        hidden_states = torch.zeros_like(encoder_input)
+        rotary_pos_emb = torch.ones(4, 1, 1, config.hidden_size // config.num_attention_heads)
+
+        extended_attention_mask = mocker.patch.object(bert_model, 'bert_extended_attention_mask')
+        embedding_forward = mocker.patch.object(
+            bert_model.embedding, 'forward', return_value=encoder_input
+        )
+        rotary_forward = mocker.patch.object(
+            bert_model.rotary_pos_emb, 'forward', return_value=rotary_pos_emb
+        )
+        encoder_forward = mocker.patch.object(
+            bert_model.encoder, 'forward', return_value=hidden_states
+        )
+
+        output = bert_model.forward(
+            input_ids=input_ids, attention_mask=None, packed_seq_params=packed_seq_params
+        )
+
+        extended_attention_mask.assert_not_called()
+        assert torch.equal(
+            embedding_forward.call_args.kwargs['position_ids'],
+            torch.tensor([[0, 1, 0, 1, 2, 3]], dtype=torch.int64),
+        )
+        rotary_forward.assert_called_once_with(
+            4, packed_seq=True, cp_group=packed_seq_params.cp_group
+        )
+        assert encoder_forward.call_args.kwargs['attention_mask'] is None
+        assert encoder_forward.call_args.kwargs['packed_seq_params'] is packed_seq_params
+        assert encoder_forward.call_args.kwargs['rotary_pos_emb'] is rotary_pos_emb
+        assert output is hidden_states
+
+    @pytest.mark.internal
+    def test_forward_validates_dense_attention_mask_and_packed_format(self):
+        sequence_length = self.bert_model.max_sequence_length
+        input_ids = torch.arange(sequence_length, dtype=torch.int64).unsqueeze(0)
+        attention_mask = torch.ones((1, sequence_length), dtype=bool)
+        cu_seqlens = torch.tensor([0, sequence_length], dtype=torch.int32)
+
+        with pytest.raises(ValueError, match='attention_mask must be provided'):
+            self.bert_model.forward(input_ids=input_ids, attention_mask=None)
+
+        with pytest.raises(ValueError, match='qkv_format'):
+            self.bert_model.forward(
+                input_ids=input_ids,
+                attention_mask=None,
+                packed_seq_params=PackedSeqParams(
+                    qkv_format='sbhd',
+                    cu_seqlens_q=cu_seqlens,
+                    cu_seqlens_kv=cu_seqlens,
+                    max_seqlen_q=sequence_length,
+                    max_seqlen_kv=sequence_length,
+                ),
+            )
+
+        with pytest.raises(ValueError, match='attention_mask must be None'):
+            self.bert_model.forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                packed_seq_params=PackedSeqParams(
+                    qkv_format='thd',
+                    cu_seqlens_q=cu_seqlens,
+                    cu_seqlens_kv=cu_seqlens,
+                    max_seqlen_q=sequence_length,
+                    max_seqlen_kv=sequence_length,
+                ),
+            )
+
+    @pytest.mark.internal
+    @pytest.mark.parametrize(
+        ('token_ids', 'cu_seqlens', 'error_match'),
+        [
+            (
+                torch.arange(4, dtype=torch.int64).repeat((2, 1)),
+                torch.tensor([0, 4]),
+                'dummy batch',
+            ),
+            (
+                torch.arange(4, dtype=torch.int64).unsqueeze(0),
+                None,
+                'cu_seqlens_q must be provided',
+            ),
+            (torch.arange(4, dtype=torch.int64).unsqueeze(0), torch.tensor([1, 4]), 'start at 0'),
+            (torch.arange(4, dtype=torch.int64).unsqueeze(0), torch.tensor([0, 3]), 'end at'),
+            (
+                torch.arange(4, dtype=torch.int64).unsqueeze(0),
+                torch.tensor([0, 3, 2, 4]),
+                'monotonically',
+            ),
+        ],
+    )
+    def test_packed_position_ids_validate_cu_seqlens(self, token_ids, cu_seqlens, error_match):
+        with pytest.raises(ValueError, match=error_match):
+            self.bert_model.bert_position_ids(
+                token_ids, PackedSeqParams(qkv_format='thd', cu_seqlens_q=cu_seqlens)
+            )
+
 
 class TestBertModelAttentionDimensions: