Mask loss at document boundaries

[francesco-bertolotti]

This is probably not a significant issue, but it seems like the more correct way to handle document packing.

Currently, when training with document packing, the last token of one document is used to predict the first token of the next document.

Because the tokenizer always appends an EOS token and prepends a BOS token, the model is effectively trained to predict a BOS token after every EOS token. Given this behavior, the current implementation does not introduce an issue. However, without these special tokens, the last token of one document would be trained to predict the first token of an unrelated document, which is clearly undesirable.

Although this is not a problem in the current setup, I believe the correct approach is to ignore the loss for the last token of each document. This avoids learning cross-document transitions and aligns the training objective with document boundaries. Additionally, implementing this requires only a one-line change.

---

This PR probably breaks some tests. Before fixing those too, I would like to know what do you think about this.

[pytorch-bot]

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

• ciflow/8gpu

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.