Mask loss at document boundaries - #4075
Conversation
meta-cla
Bot
added
the
CLA Signed
|
|
| label = x[1:].clone() | ||
| # Mask EOS tokens in the label to avoid predicting from | ||
| # next-document context. | ||
| label[input == self._tokenizer.eos_id] = IGNORE_INDEX |
There was a problem hiding this comment.
Nice catch!
Instead of doing this, where we always "waste" the training when input is eos_id as we do input, label split on the packed buffer, how about let's follow the ChatDataloader and do input, label split per each sample
https://github.com/pytorch/torchtitan/blob/main/torchtitan/hf_datasets/text_datasets.py#L482-L483
|
Hi @tianyu-l
Seems reasonable to me. With this change, the shift now happens at the document level rather than the row level. I’ve also added a few tests with some basic checks. A few things to be aware of:
|
There was a problem hiding this comment.
lgtm -- the numerics drift in CI is expected, could you update the loss files?
@felipemello1 to take another look
|
Warning: Unknown label
Please add the new label to .github/pytorch-probot.yml |
| if "labels_buffer" in state_dict: | ||
| self._inputs_buffer = state_dict["inputs_buffer"] | ||
| self._labels_buffer = state_dict["labels_buffer"] | ||
| self._positions_buffer = state_dict["positions_buffer"] | ||
| else: | ||
| # Checkpoints written before labels were buffered separately hold a | ||
| # raw, unshifted token stream in 'inputs_buffer'. It cannot be split | ||
| # into input/label pairs here without reintroducing the cross-document | ||
| # targets this shift exists to avoid, so drop the partial sample and | ||
| # resume at the next document boundary. | ||
| logger.warning( | ||
| "Checkpoint missing 'positions_buffer'. Falling back to empty buffer. " | ||
| "RoPE positions may be incorrect with block_causal attention." | ||
| "Checkpoint missing 'labels_buffer'. Dropping the buffered partial " | ||
| "sample; iteration resumes at the next document boundary." | ||
| ) |
There was a problem hiding this comment.
i dont think that we need backwards compatibility. All datasets from here on should have labels buffer and positions buffer, right? should we just delete it @francesco-bertolotti ?
@tianyu-l can you confirm.
There was a problem hiding this comment.
yeah BC is secondary concern, I don't care that much
There was a problem hiding this comment.
Sounds reasonable
|
I’ve updated Regarding the numerics drift in CI, I have updated |
|
seems there are still CI errors |
|
Sorry, Let's hope this time works.
|
4 participants
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.