[Afloat16]

Problem

The reranker collators flatten each query's positive and negative passages, but the model reshapes their scores using the configured training batch size. A final DataLoader batch with fewer queries can therefore mix candidates from different queries, raise a reshape error, or create one-candidate groups with zero cross-entropy. The same grouping also affects teacher targets and layerwise distillation.

Changes

• Preserve the number of passages per query as group_size in both reranker collators.
• Infer the actual local query count from the encoded score rows and that group size in the shared and layerwise reranker training forwards.
• Keep the configured batch-size fallback for direct callers that omit the new metadata.
• Add native, download-free regression tests using tiny BERT, GPT-2, and the repository's LayerWiseMiniCPM models with real collators and DataLoader batches.

The group-size metadata stays valid when tensors are split at complete query-group boundaries. This change does not address existing DataParallel behavior for Python-list teacher scores or tensor splits through the middle of a query group.

Validation

• On the unchanged upstream base, the 27 new cases produce 15 failures and 12 passes; with this patch all 27 pass.
• The tests compare cross-entropy, external distillation, and layerwise self-distillation losses and gradients against explicit per-query calculations, including one-query/two-query final batches, full batches, and legacy direct calls.
• Six additional cases within that set simulate per-device tensor splitting with replicated scalar metadata for all three model families.
• New tests plus existing Trainer/import and offline tokenizer-compatibility tests: 39 passed, 1 skipped on PyTorch 2.15.0+cpu and Transformers 4.57.1. The skipped import check requires Transformers 5.x. Download-dependent model/tokenizer tests and CUDA execution were not run.
• Ruff on the new test file and git diff --check pass.