Reject long untruncated calibration samples when max_seq_length is unset - #3012
Reject long untruncated calibration samples when max_seq_length is unset#3012rishabhsinha17 wants to merge 1 commit into
Conversation
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Merge Protections🔴 2 of 2 protections blocking · waiting on 👀 reviews
🔴 Require one maintainer reviewWaiting for any of
This rule is failing.All PRs must have at least one approving review from a maintainer before merging.
🔴 Require two reviewsWaiting for
This rule is failing.PRs labelled "two-reviews" must have at least two approving reviews before merging.
|
|
👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review. Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed. |
![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 warning mechanism to alert users when max_seq_length is unset and the calibration dataset contains long untruncated sequences, which could lead to GPU out-of-memory errors. It also adds corresponding unit tests. The feedback suggests optimizing the sequence length calculation for large datasets using PyArrow to prevent potential CPU memory spikes or OOM errors.
| else: | ||
| return | ||
|
|
||
| lengths = [len(sample) for sample in dataset[feature_name]] |
There was a problem hiding this comment.
For large datasets, loading the entire tokenized column into Python memory using dataset[feature_name] can be extremely slow and cause a massive CPU memory spike (or even OOM), as it instantiates Python list and integer objects for every single token in the dataset.
Since Hugging Face Dataset is backed by PyArrow, we can compute the sequence lengths extremely efficiently in C++ using pyarrow.compute.list_value_length without loading the nested lists into Python memory. This reduces memory usage from gigabytes to megabytes and speeds up the check from seconds/minutes to milliseconds.
We can fall back to the list comprehension if PyArrow is unavailable or if an exception occurs.
try:
import pyarrow.compute as pc
column = dataset.data[feature_name]
if hasattr(dataset, "_indices") and dataset._indices is not None:
column = pc.take(column, dataset._indices.arrow_array)
lengths = pc.list_value_length(column).to_pylist()
except Exception:
lengths = [len(sample) for sample in dataset[feature_name]]There was a problem hiding this comment.
Done in the amended commit: lengths now come from pyarrow list_value_length on the arrow column, with the indices mapping applied for shuffled/selected datasets and the list comprehension kept as fallback. Verified the arrow path returns identical lengths to the fallback on an indexed dataset.
e901a74 to
837e303
Compare
|
@brian-dellabetta #3020 looks like the right fix for the failure path, and I agree the error site beats a warning for anyone who never hits this. Flagging where this PR sits relative to it so you can take one, both, or neither. They fire at different points and in different places:
The part that is not duplicated: the warning reports the measured longest sample and how many exceed the threshold, so the user sees which lever applies to their data and by how much, rather than a list of levers to try. It also covers pre-tokenized datasets, which skip the tokenization path entirely. If you would rather carry a single mechanism, say so and I will close this. If both are wanted I am happy to rebase on #3020 once it lands and trim anything the error message makes redundant. |
|
Hi @rishabhsinha17 , yes the error is favorable to the warning. logs are inevitably noisy with warnings from transformers and ourselves, and the error will display it at the proper location in logs. if you'd like to update your PR to look more like mine -- raising an error, fixing the merge conflict, and displaying max sequence length in error message -- we can close mine in favor of yours |
837e303 to
c446a9d
Compare
c446a9d to
eda22e2
Compare
rishabhsinha17
changed the title
|
@brian-dellabetta Done, reworked as you described. |
When `max_seq_length` is unset and the tokenized calibration dataset contains samples longer than SEQ_LEN_ERROR_THRESHOLD (2048) tokens, format_calibration_data() now raises a ValueError instead of warning. Calibrating with long untruncated samples runs out of GPU memory with the OOM raised from attention or attention mask expansion, which is easily mistaken for the model not fitting on the device, and a warning for it is buried in noisy logs while an error surfaces at the failure point (per review). The message reports how many samples exceed the threshold and the longest length, and gives both remedies: set `max_seq_length` to truncate, or set it to at least the longest length to calibrate on full-length samples intentionally. Sample lengths are measured on the arrow column to avoid materializing the tokenized dataset in Python memory. Also fold in vllm-project#3020's guidance so it can close in favor of this PR: the sequential pipeline's OOM message now also suggests reducing `num_calibration_samples` or `max_seq_length`. Fixes vllm-project#3011 Co-authored-by: Brian Dellabetta <bdellabe@redhat.com> Signed-off-by: Rishabh Sinha <rsinha17@terpmail.umd.edu>
eda22e2 to
626173b
Compare
2 participants
SUMMARY:
Resolves #3011
When
oneshot()runs with a text dataset and nomax_seq_length, tokenization usesmax_length=tokenizer.model_max_length, so long samples reach calibration untruncated. On a 32 GB card this OOMs inside transformers attention mask expansion with an error that points at GPU capacity, not sequence length, sending users towardpipeline="basic"or CPU offload, which do not help.Reworked per the review discussion with @brian-dellabetta:
format_calibration_datanow raises a ValueError, instead of warning, whenmax_seq_lengthis unset and the tokenized calibration dataset contains samples longer than 2048 tokens. The error reports how many samples exceed the threshold and the longest sample length, and gives both remedies: setmax_seq_lengthto truncate, or set it to at least the longest length to calibrate on the full-length samples intentionally (truncation to a bound no sample exceeds is a no-op, so that is the explicit opt-in). Checking the tokenized dataset also covers pre-tokenized datasets, which skip the tokenization path entirely. Sample lengths are measured on the arrow column to avoid materializing the tokenized dataset in Python memory.This also carries over #3020's improvement so it can close in favor of this PR: the sequential pipeline's OOM message now also suggests reducing
num_calibration_samplesormax_seq_length(co-author credit in the commit).TEST PLAN:
Unit tests in
tests/llmcompressor/datasets/test_max_seq_length.pyupdated to assert the ValueError: trigger paths (input_ids,decoder_input_ids) check the measured numbers and bothmax_seq_lengthsuggestions in the message; non-trigger paths (samples at or below threshold,max_seq_lengthset, dataset without input ids) complete without raising.pytest tests/llmcompressor/datasets/test_max_seq_length.pypasses (10 tests).ruff checkandruff format --checkare clean on the changed files with the repo-pinned ruff (0.4.10). Branch rebased onto latest main.