Allow position_ids_start=2 on DataCollatorWithFlattening for RoBERTa etc.

[tomaarsen]

Fixes #47496

What does this PR do?

Pull Request overview

• Add a position_ids_start argument to DataCollatorWithFlattening
• Infer FlashAttention sequence boundaries from position_ids == position_ids.min() instead of position_ids == 0

Details

As described in #47496, DataCollatorWithFlattening numbers positions from 0, but RoBERTa-like architectures (roberta, xlm_roberta, camembert, mpnet, esm with absolute position embeddings, and more) compute their position ids starting at padding_idx + 1 (usually 2) when none are passed. Feeding these models the collator's 0-based position_ids silently shifts the position embeddings and degrades the outputs. This affects some very common checkpoints, e.g. xlm-roberta.

This PR adds a position_ids_start argument (default 0, existing behavior is unchanged) so users can match the numbering of their model, e.g. position_ids_start=2 for XLM-R. The collator cannot reliably derive this from a config: mpnet for example hardcodes padding_idx = 1 regardless of config.pad_token_id. The docstring now documents all arguments and warns about the pitfall.

The new argument alone turned out to be insufficient on the path where FlashAttention infers the varlen kwargs from position_ids at runtime (i.e. without return_flash_attn_kwargs=True, like the reproduction in the issue). prepare_fa_kwargs_from_position_ids detected sequence starts via position_ids == 0, which never matches offset positions, after which it crashes on an empty cu_seq_lens.diff().max(). I changed the detection to position_ids == position_ids.min(), mirroring _is_packed_sequence which already handles a non-zero start via position_ids.min(). Whenever a 0 position is present, this is behavior-identical, so existing models are unaffected.

I verified that with position_ids_start=2 the collator output exactly matches RobertaEmbeddings.create_position_ids_from_input_ids on unpadded sequences, and that the RoBERTa embedding layer produces identical outputs for the flattened batch and for the individual sequences. The flattening collator tests now cover the new argument as well as the boundary inference on collator output.

Because of this prepare_fa_kwargs_from_position_ids patch, the PR got a bit bigger than I originally intended in #47496 (comment), but it does seem important.

Code Agent Policy

• (First-time contributors only): I confirm that this PR description and code is not written by an LLM or code agent

The tests and docstrings were written by Fable and reviewed by myself.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline and the
  Pull Request checks?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes according to the guidelines?
• Did you write any new necessary tests?

Who can review?

@vasqu

• Tom Aarsen

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 30105737205
Result: success | Grafana metrics are not available yet.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[vasqu]

Can we add an integration test for roberta that tests this case implicitly? It's ok to require FA or kernels