packed seq lengths token count correction #348
Conversation
Summary of ChangesHello @djsaunde, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request refines the token counting mechanism in the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a correction for token counting when using packed sequences, which is an important fix for correct loss calculation. The overall logic is sound, but I've identified a potential edge case in the implementation where the token count could be incorrectly modified. I've provided a suggestion to make the calculation more robust against this edge case.
djsaunde
force-pushed
the
packed-token-count
branch
from
0c1e04f to
dd7fe22
Compare
djsaunde
force-pushed
the
packed-token-count
branch
from
dd7fe22 to
1e5b0b7
Compare
| and not isinstance(trainer.data_collator, DataCollatorForSeq2Seq) | ||
| and not packing_enabled | ||
| ): | ||
| trainer.data_collator = DataCollatorForSeq2Seq(tokenizer = tokenizer) |
There was a problem hiding this comment.
oh @danielhanchen I wanted to ask why we do this, and if there's anything that might be breaking about my change? FWIW I haven't encountered any issues with this change when running all the notebooks in the unsloth main repo README.
There was a problem hiding this comment.
It's because we must use DataCollatorForSeq2Seq for training on completions only since DataCollatorForSeq2Seq allows us to have input_ids and labels separately. The other DataCollatorForLanguageModel doesn't allow labels.
There was a problem hiding this comment.
hmm I think we should actually use DataCollatorForLanguageModel and pass completion_only_loss = True there. I don't think we should use a custom (non-default) collator here. that breaks packing
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| count = token_count.sum() | ||
| seq_lengths = x.get("packed_seq_lengths") | ||
| if seq_lengths is not None: | ||
| count -= torch.count_nonzero(seq_lengths > 0).item() - 1 |
There was a problem hiding this comment.
When packed_seq_lengths is present (packing enabled), the token count is reduced by torch.count_nonzero(seq_lengths > 0).item() - 1, but labels[..., 1:] has already removed one token per sample. For batch sizes larger than 1 this double‑subtracts batch_size-1 tokens (total sequences minus only one), so num_items_in_batch becomes too small and the loss gets normalized over an artificially low denominator whenever packing is used with multiple examples per batch.
Useful? React with 👍 / 👎.
Should be merged alongside #3525.