[serge] Fix 8 integration tests for model grounding_dino failing with output_mismatch (tensor values differ (6), oth (#49886)

* [serge] Fix 8 integration tests for model `grounding_dino` failing with `output_mismatch` (tensor values differ (6), oth

* what pablo asked

* fix repo

* and the pos id thing seems valid

* trf

* Address review: required special_tokens, regression tests, new expectations

- `generate_masks_with_special_tokens_and_transfer_map` now requires `special_tokens`
- drop the buffer comment
- add fast tests for phrase masks / position ids and for float text position embeddings
- update detection-head expectations for the position-id fix (grounding_dino + mm_grounding_dino)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* no more flaky asserts and update loss

* add expectations from runner here

---------

Co-authored-by: serge[bot] <serge[bot]@users.noreply.github.com>
Co-authored-by: raushan <raushan@huggingface.co>
Co-authored-by: Tarek Ziade <tarek@ziade.org>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Author: 5 people

src/transformers/models/grounding_dino/modeling_grounding_dino.py

@@ -1049,20 +1049,16 @@ def get_text_position_embeddings(
         text_position_embedding: torch.Tensor | None,
         text_position_ids: torch.Tensor | None,
     ) -> Tensor:
-        batch_size, seq_length, _ = text_features.shape
-        if text_position_embedding is None and text_position_ids is None:
-            text_position_embedding = torch.arange(seq_length, device=text_features.device)
-            text_position_embedding = text_position_embedding.float()
-            text_position_embedding = text_position_embedding.unsqueeze(0).unsqueeze(-1)
-            text_position_embedding = text_position_embedding.repeat(batch_size, 1, 1)
-            text_position_embedding = encode_sinusoidal_position_embedding(
-                text_position_embedding, num_pos_feats=self.d_model
-            )
-        if text_position_ids is not None:
-            text_position_embedding = encode_sinusoidal_position_embedding(
-                text_position_ids[..., None], num_pos_feats=self.d_model
-            )
-
+        if text_position_embedding is not None:
+            return text_position_embedding
+
+        if text_position_ids is None:
+            batch_size, seq_length, _ = text_features.shape
+            text_position_ids = torch.arange(seq_length, device=text_features.device)
+            text_position_ids = text_position_ids.expand(batch_size, -1)
+        text_position_embedding = encode_sinusoidal_position_embedding(
+            text_position_ids[..., None].to(text_features.dtype), num_pos_feats=self.d_model
+        )
         return text_position_embedding
 
     def forward(
@@ -1337,6 +1333,8 @@ def _init_weights(self, module):
         if isinstance(module, GroundingDinoLearnedPositionEmbedding):
             init.uniform_(module.row_embeddings.weight)
             init.uniform_(module.column_embeddings.weight)
+        elif isinstance(module, GroundingDinoModel):
+            init.copy_(module.special_tokens, torch.tensor(SPECIAL_TOKENS))
         elif isinstance(module, GroundingDinoMultiscaleDeformableAttention):
             init.constant_(module.sampling_offsets.weight, 0.0)
             default_dtype = torch.get_default_dtype()
@@ -1809,11 +1807,15 @@ def custom_forward(*inputs):
 SPECIAL_TOKENS = [101, 102, 1012, 1029]
 
 
-def generate_masks_with_special_tokens_and_transfer_map(input_ids: torch.LongTensor) -> tuple[Tensor, Tensor]:
+def generate_masks_with_special_tokens_and_transfer_map(
+    input_ids: torch.LongTensor, special_tokens: torch.LongTensor
+) -> tuple[Tensor, Tensor]:
     """Generate attention mask between each pair of special tokens and positional ids.
     Args:
         input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
             Indices of input sequence tokens in the vocabulary.
+        special_tokens (`torch.LongTensor` of shape `(num_special_tokens,)`):
+            Ids of the delimiter tokens that split the text into phrases.
     Returns:
         `tuple(torch.Tensor)` comprising attention mask between each special tokens and position_ids:
         - **attention_mask** (`torch.BoolTensor` of shape `(batch_size, sequence_length, sequence_length)`)
@@ -1823,14 +1825,15 @@ def generate_masks_with_special_tokens_and_transfer_map(input_ids: torch.LongTen
     device = input_ids.device
 
     # Identify special token positions
-    special_mask = torch.isin(input_ids, torch.tensor(SPECIAL_TOKENS, device=device))
+    special_mask = torch.isin(input_ids, special_tokens)
 
     # For each position, find the previous and next special token indices
     indices = torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1)
 
     # Previous special token: cummax of special token indices
     prev_special = torch.where(special_mask, indices, torch.full((), -1, device=device))
     prev_special = torch.cummax(prev_special, dim=1)[0]
+    prev_special = torch.cat([prev_special.new_full((batch_size, 1), -1), prev_special[:, :-1]], dim=1)
 
     # Next special token: flip, cummin, flip back
     next_special = torch.where(special_mask, indices, torch.tensor(seq_len, device=device))
@@ -1933,6 +1936,10 @@ def __init__(self, config: GroundingDinoConfig):
         else:
             self.reference_points = nn.Embedding(config.num_queries, 4)
 
+        self.special_tokens = nn.Buffer(
+            torch.tensor(SPECIAL_TOKENS, dtype=torch.long),
+            persistent=False,
+        )
         self.post_init()
 
     def freeze_backbone(self):
@@ -2060,7 +2067,9 @@ def forward(
         )
         return_dict = return_dict if return_dict is not None else self.config.return_dict
 
-        text_self_attention_masks, position_ids = generate_masks_with_special_tokens_and_transfer_map(input_ids)
+        text_self_attention_masks, position_ids = generate_masks_with_special_tokens_and_transfer_map(
+            input_ids, special_tokens=self.special_tokens
+        )
 
         if attention_mask is None:
             attention_mask = torch.ones_like(input_ids)

src/transformers/models/mm_grounding_dino/modeling_mm_grounding_dino.py

@@ -24,6 +24,7 @@
 from ..auto import AutoConfig
 from ..auto.modeling_auto import AutoModel
 from ..grounding_dino.modeling_grounding_dino import (
+    SPECIAL_TOKENS,
     GroundingDinoContrastiveEmbedding,
     GroundingDinoConvEncoder,
     GroundingDinoConvModel,
@@ -247,6 +248,10 @@ def __init__(self, config: MMGroundingDinoConfig):
         )
         self.encoder_output_class_embed = MMGroundingDinoContrastiveEmbedding(config)
 
+        self.special_tokens = nn.Buffer(
+            torch.tensor(SPECIAL_TOKENS, dtype=torch.long),
+            persistent=False,
+        )
         self.post_init()