Preserve the final batch in balanced resumable sampling

[Afloat (Afloat16)]

BalancedResumableSampler only visits complete global batches. After the parent sampler has already padded or truncated for distributed divisibility, a final partial batch is silently discarded; small datasets can yield no indices even though len(sampler) is positive.

Visit the final distributed batch as well. Its size remains divisible by world size, so the existing equal-size load balancing preserves equal per-rank sample counts. drop_last continues to control the distributed remainder through the parent sampler.

Regression tests check sample conservation and reported cardinality across ranks, small datasets, distributed drop_last, and resume offsets through the tail. Run python -m unittest discover -s tests -p test_trellis2_balanced_tail.py -v.

Validation: native CPU regression tests pass on this branch; the same tests fail against the pinned original source. Validation run.

Exact-head validation: all 3 native CPU regression tests pass at 167651ac9c01c43b3ded45f0e689a074a2434df0. Restoring the original production files at 75fbf0183001ed9876c8dbb35de6b68552ee08bd makes the same complete suite fail on the original behavior, with no skipped tests or import failures. Candidate and original-source logs.

[Afloat (Afloat16)]

@microsoft-github-policy-service agree