[unknownhackerworld]

Description

When several calibration datasets are combined (for example --dataset "NeelNanda/pile-10k,BAAI/CCI3-HQ" or with :num=), _get_dataset_impl in auto_round/calib_dataset.py can produce fewer than nsamples samples and split them unevenly. This PR fixes the two bugs behind that in the mixing block.

Root cause

The mixing code sorts the datasets by length (smallest first) and then hands out the budget so that whatever a small dataset cannot provide flows to the larger ones after it. Two things break that:

1. lens is not reordered. After res = sorted(zip(indices, lens), ...), datasets and dataset_names are reordered, but lens is left in the original order. target_cnt = min(target_cnt, lens[i]) then caps each dataset by another dataset's length. When a small dataset is listed after a large one, it is asked for more rows than it has: Dataset.select raises IndexError, select_dataset falls back to the slow iterate-and-pandas path, which returns only the rows that exist, and the large dataset is capped at the small one's length. Example: with big (100 rows) listed before small (5 rows) and nsamples=20, tracing the code gives 10 samples (5 + 5) instead of 20 (15 + 5).

2. Wrong divisor when num= is used. With any num= entry, the remaining budget is always divided by len(datasets) - len(data_lens), the total number of unspecified datasets, even though cnt already includes what earlier datasets took. Later datasets therefore get too little. Example: a:num=4,b,c (100 rows each) with nsamples=20 gives 4 + 8 + 4 = 16 instead of 4 + 8 + 8 = 20. Without num= the code divides by len(datasets) - i, which is the intended "datasets still left" count.

Fix

• Reorder lens together with datasets and dataset_names.
• Count the unspecified datasets up front and divide the remaining budget by how many of them are still left, decrementing as each one is filled. Without num= this gives the same result as the old len(datasets) - i, so that path is unchanged.

Alternatives considered

• Not sorting at all and splitting evenly. Simpler, but then leftover budget from a small dataset would be lost, which the sort was clearly added to avoid.
• Pulling the allocation into a separate helper function. Cleaner for testing, but a bigger diff. I kept the change local to the existing block.

Type of Change

Bug fix

Related Issues

None found.

Checklist Before Submitting

• My code has been tested locally.
• Documentation has been updated as needed. (No doc change needed.)
• New or updated tests are included where applicable.
• The CUDA CI has passed. (Not triggered by me; this change only touches CPU-side dataset handling. Happy for a maintainer to run it.)

Testing

• Added TestMixedDatasetSampleAllocation to test/unit/common/utils/test_calib_dataset_helpers.py. It registers small in-memory datasets in CALIB_DATASETS and calls _get_dataset_impl directly:
  • test_small_dataset_listed_last_does_not_shrink_total: big (100) + small (5), nsamples=20, expects 15 + 5.
  • test_remaining_budget_is_shared_with_num_option: a:num=4,b,c, nsamples=20, expects 4 + 8 + 8.
• pytest test/unit/common/utils/test_calib_dataset_helpers.py: 40 passed (CPU, torch CPU wheel). With the fix stashed and only the tests applied, both new tests fail.
• black (line length 120) leaves both files unchanged.

Disclosure: this change and description were prepared with the help of an AI coding assistant.

[Copilot]

🟢 Approval recommended

All reviewed changes are covered by regression tests with no unresolved review issues.

No open issues 🎉

What changed in this PR

Fixes calibration sample allocation when mixing datasets, including :num= configurations.

Changes:

• Keeps dataset lengths aligned after sorting.
• Correctly distributes remaining samples.
• Adds regression tests for both allocation bugs.

File	Description
test/​unit/​common/​utils/​test_calib_dataset_helpers.py	Adds allocation regression tests.
auto_round/​calib_dataset.py	Fixes mixed-dataset sample allocation.

_{Review effort: Lite}

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).