[YueshenZ]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Related issue: none. Found through code analysis and reproduced with constructed CPU evaluation.

Distributed evaluation repeats dataset rows to fill its final batch. EmbeddingTrainer gathers flattened anchor/positive/negative tensors without removing those padded examples, so the built-in embedding metrics include extra queries and candidate documents. Preserve each example as a group during metric gathering, let Accelerate remove sampler padding, and immediately restore the flattened tensors. This also preserves real repeated rows and supports different negative counts without padding the entire evaluation cache to the largest group.

The change covers the built-in infonce and paired metrics with their default preprocessing, Transformers 4.36 or later, and no sequence parallelism. Custom preprocessing and older Transformers keep the existing path. Loss gathering is unchanged. Add a CPU distributed regression and clarify the scope in both embedding guides.

Experiment results

Tested on CPU with Python 3.13, PyTorch 2.14.0 and Transformers 5.17.0 on main bdb8741699a9cc0647d9759575e0f30cc08516f0 before pushing. The failure reproduction used the original 69872d1 baseline; the embedding trainer was unchanged between those upstream commits.

.venv/bin/python -m pytest tests/utils/test_embedding_evaluation.py tests/utils/test_embedding_metrics.py tests/utils/test_embedding_metrics_dp.py tests/train/test_embedding_loss.py -q
.venv/bin/python tests/run.py --pattern test_embedding_evaluation.py --result_dir /tmp/embedding-eval-results
.venv/bin/python -m pre_commit run --all-files
.venv/bin/python -m pip check
git diff --check

• The new test fails on unmodified main: a three-example evaluation reports InfoNCE margin -0.11395209 instead of the direct-model reference -0.11750958.
• The related suite passes 15 tests and 3 subtests; 4 CUDA-only cases are skipped on this CPU machine. The new test is in tests/utils, which is discoverable from the CI runner’s tests root. The actual tests/run.py command collects and passes one test, and its result is recorded as Success in the runner’s result table. All configured pre-commit hooks, dependency and diff checks pass.
• The regression uses a local randomly initialized tiny Qwen3 and real EmbeddingTrainer.evaluate / predict on two Gloo processes. It covers 3/7/8 examples, varying negative counts across ranks and batches, a genuine repeated row, both built-in metrics, repeated evaluation, flat prediction outputs, loss-only evaluation, custom preprocessing and the old-Transformers fallback. Its reference comes from independent per-example model forwards.
• A separate three-example evaluation with negative counts [1, 2, 1] changes from 4 queries before the fix to the expected 3 afterward. Its mean_pos changes from 0.95677412 to 0.94236547, matching the single-process reference; the margin agrees within 2e-8. This is an evaluation correctness experiment, not a model-quality or training benchmark. GPU/NCCL and sequence-parallel execution were not tested.

The CPU worker disables the project’s NPU-specific model patches before its first Swift model import. It also uses the existing single-device-mode setting to bypass accelerator-count-based model parallelism checks, and bypasses physical set_device only while constructing CPU training arguments. Model forwards, distributed collectives and metrics remain real; the fixture explicitly checks that the trainer uses CPU and two processes. Local fault injection with one reported accelerator and forbidden physical device selection fails before this fixture adjustment and passes afterward. GPU/NPU host compatibility still awaits the updated remote CI run.

AI assistance was used (OpenAI Codex) for investigation, implementation and validation.

Checklist

• I have linked the related issue above, if there is one.
• I have reviewed and understand every changed line, and I can explain how the change works and why it is correct.
• I have run the relevant tests or reproductions locally, and included the commands and results above (if applicable).
• I am responsible for the content of this PR, regardless of how it was produced.