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,23 @@ 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:
+            assert token_ids.size(0) == 1, 'Packed BERT input should use dummy batch size 1'
+            assert (
+                packed_seq_params.cu_seqlens_q is not None
+            ), '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)
+            seq_lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.long)
+            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 +313,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 +332,15 @@ def forward(
 
         inference_context = deprecate_inference_params(inference_context, inference_params)
 
-        extended_attention_mask = self.bert_extended_attention_mask(attention_mask)
+        extended_attention_mask = (
+            None
+            if packed_seq_params is not None
+            else self.bert_extended_attention_mask(attention_mask)
+        )
 
         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 +359,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
@@ -92,6 +93,67 @@ def test_post_process_forward(self):
         assert logits[0].shape[1] == sequence_length
         assert logits[0].shape[2] == self.bert_model.vocab_size
 
+    @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
+
 
 class TestBertModelAttentionDimensions: