Use indices as position_ids in modernebert (#41789)

remi-or · web-flow · commit 3c16c1ae43ab · 2025-11-03T12:10:24.000+01:00
* Use indices as position_ids in modernebert

* Move position_ids init to the branch
diff --git a/src/transformers/models/modernbert/modeling_modernbert.py b/src/transformers/models/modernbert/modeling_modernbert.py
@@ -905,6 +905,8 @@ def forward(
                     inputs_embeds, indices, cu_seqlens, max_seqlen, *_ = _unpad_modernbert_input(
                         inputs=inputs_embeds, attention_mask=attention_mask
                     )
+            if position_ids is None:
+                position_ids = indices.unsqueeze(0)
         else:
             if position_ids is None:
                 position_ids = torch.arange(seq_len, device=device).unsqueeze(0)
diff --git a/src/transformers/models/modernbert/modular_modernbert.py b/src/transformers/models/modernbert/modular_modernbert.py
@@ -1014,6 +1014,8 @@ def forward(
                     inputs_embeds, indices, cu_seqlens, max_seqlen, *_ = _unpad_modernbert_input(
                         inputs=inputs_embeds, attention_mask=attention_mask
                     )
+            if position_ids is None:
+                position_ids = indices.unsqueeze(0)
         else:
             if position_ids is None:
                 position_ids = torch.arange(seq_len, device=device).unsqueeze(0)

Original file line number	Diff line number	Diff line change
`@@ -905,6 +905,8 @@ def forward(`
`905`	`905`	`inputs_embeds, indices, cu_seqlens, max_seqlen, *_ = _unpad_modernbert_input(`
`906`	`906`	`inputs=inputs_embeds, attention_mask=attention_mask`
`907`	`907`	`)`
	`908`	`+ if position_ids is None:`
	`909`	`+ position_ids = indices.unsqueeze(0)`
`908`	`910`	`else:`
`909`	`911`	`if position_ids is None:`
`910`	`912`	`position_ids = torch.arange(seq_len, device=device).unsqueeze(0)`
Original file line number	Diff line number	Diff line change
`@@ -1014,6 +1014,8 @@ def forward(`
`1014`	`1014`	`inputs_embeds, indices, cu_seqlens, max_seqlen, *_ = _unpad_modernbert_input(`
`1015`	`1015`	`inputs=inputs_embeds, attention_mask=attention_mask`
`1016`	`1016`	`)`
	`1017`	`+ if position_ids is None:`
	`1018`	`+ position_ids = indices.unsqueeze(0)`
`1017`	`1019`	`else:`
`1018`	`1020`	`if position_ids is None:`
`1019`	`1021`	`position_ids = torch.arange(seq_len, device=device).unsqueeze(0)`