[bugfix] slice rows per rank and add data_config.min_batch_size by tiankongdeguiji · Pull Request #677 · alibaba/TorchEasyRec

tiankongdeguiji · 2026-09-19T04:14:11Z

Problem

An hourly training job failed with ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 1024]) inside a BatchNorm1d in train mode. calc_slice_position split each table into num_workers * world_size slices, and the drop_redundant_bs_eq_one guard only handled the residue where the base slice is a whole number of batches. When floor(rows / slices) % batch_size == 1, every worker ended with a 1-row tail batch. With batch_size=1024 and 8 workers that is roughly one pass in a thousand, so the same config runs fine most hours and fails at random.

Change

Rank-level slicing. plan_rank_worker_intervals in tzrec/datasets/utils.py replaces calc_slice_position and the pre_total_remain chaining. All tables or sessions of an input are planned as one stream. Ranks split the stream as if its rows were dealt round-robin, laid out contiguously inside every source: each source stays spread over all ranks (partition order is kept under is_orderby_partition) while the rank totals differ by at most one row over the whole pass, so 33 + 47 rows over 8 ranks are exactly 10 rows each with nothing lost. Inside a rank the stream is cut into whole batches plus one tail and dealt source by source: the rows completing the open batch go to its holder, whole batches go as contiguous blocks to the least loaded workers, and the tail opens the next batch. Only the holder ever buffers a partial batch, so a pass produces at most one partial batch per rank, including across sources (150 + 150 rows at batch 100 are exactly three full batches). Equal step count across ranks is the only invariant the dataloader owes the trainer, and it holds by construction: the ranks holding an extra row drop it only when (rows // world_size) % batch_size == 0, where it would buy them a step.

Reader API. to_batches(worker_id, num_workers) keeps its meaning of "the worker_id-th of num_workers even shares", so tools such as hitrate.py and faiss_util.build_faiss_index are unchanged. BaseDataset, the one caller that knows the DataLoader workers of each rank, passes the new optional world_size so the planner can group slices rank-major; the reader never infers it from the process group.

data_config.min_batch_size (uint32, default 0, no effect). Training tail batches with fewer rows are dropped on every rank identically. Models with BatchNorm, in-batch negatives, or listwise losses set it to 2 or higher. drop_remainder=true is now implemented as min_batch_size=batch_size, and dropped rows are never read.

Eval and predict never drop. Both knobs apply to Mode.TRAIN only; rank equalization stays on for train and eval, which run collectives per step, and is off in predict. Previously drop_remainder=true also dropped the final partial buffer per worker in eval and predict through _arrow_reader_iter, so eval metrics and predict output silently lost rows; that is fixed here as a side effect and is a behavior change worth noting.

ODPS session record counts are fetched once per rank when sessions are created or restored, outside any collective, and cached on the reader, so dataloader workers issue no metadata RPCs. The 1-row BatchNorm failure and the knob are documented as FAQ Q21 rather than a runtime warning.

Untouched: Kafka (its consume loop never ends, so both knobs are inert for it), odps_dataset_v1, the CSV file split, batch_cost_size, and main.py.

Alternatives considered

Detecting BatchNorm in the model and deriving the drop automatically. Rejected: it couples the dataloader to the model, the flag would have to be known before create_dataloader forks its workers, and it misses non-BatchNorm minimum-row requirements. An explicit threshold expresses the real constraint.
Dropping the rows % world_size extras of every session unconditionally. Simpler, but drops rows that step equality does not require. The cumulative planner keeps them whenever no rank would gain a step.

Test Plan

tzrec/datasets/utils_test.py: brute-force property test over batch size, world size, worker count, 1 to 5 sources, threshold, and both modes, asserting equal steps per rank when equalizing, at most one partial batch per rank, no batch above batch_size or below the threshold, every row read at most once, zero drops in predict, and at most world_size - 1 dropped rows per pass (plus a short tail per rank) otherwise; a uniform-source case checks worker skew stays within a batch regardless of source count. Fixed cases cover 8200 and 8201 rows at 8 workers, 33 to 36 rows at 4 workers and 2 ranks with the threshold on and off, an empty pass, two sources whose tails sum to batch_size + 1, 150 + 150 rows forming three full batches, 6 + 7 rows forming 4/4/4/1, 33 + 47 rows over 8 ranks losing nothing, drop_remainder, calc_slice_intervals at 2 ranks x 2 workers with and without a checkpoint state, and even-share slices asserted to tile the remaining intervals disjointly, including a two-source resume. The old test_calc_slice_position, which encoded the previous carry semantics, is replaced. All 54 tests pass.
tzrec/datasets/parquet_dataset_test.py: a two-rank gloo harness, with one and two workers per rank, asserting the expected per-rank batch-size lists for residue-hitting row counts and that a bare to_batches() reads the whole table on every rank; create_dataloader with 8 workers on 8200 rows yields nine batches with an 8-row tail for threshold 0 and 2, and eight batches with drop_remainder; eval and predict modes return all 8201 rows with drop_remainder=true and min_batch_size=2; drop_remainder=true alone drops the 8-row tail; min_batch_size > batch_size is rejected. All 19 tests pass, as do csv_dataset_test and dataset_test.
tzrec/datasets/odps_dataset_test.py: the orderby-partition resume test now sets min_batch_size=2; it needs ODPS credentials and runs in that CI lane.
End to end on CPU: a create_dataloader with 8 workers on 8193 rows at batch_size=1024, feeding each batch's row count through a train-mode BatchNorm1d, reproduces the ValueError with min_batch_size=0 and completes 8 full steps with min_batch_size=2.
tzrec/tests/match_integration_test.py::test_dssm_with_fg_train_eval_export, which failed in the first CI run because build_faiss_index calls to_batches() under a 2-rank torchrun, passes locally after the world_size change; the two-rank parquet test now also asserts a bare to_batches() reads the whole table on every rank.
pre-commit run and pyrefly check are clean.

🤖 Generated with Claude Code

https://claude.ai/code/session_01163rPXLGHiM74LxdbWjMox

Training crashed with "Expected more than 1 value per channel" from a BatchNorm1d in train mode when a table's row count left every dataloader worker with a single leftover row. calc_slice_position split a table into num_workers * world_size slices and its drop_redundant_bs_eq_one guard only covered the residue where the base slice is a whole number of batches; when floor(rows / slices) % batch_size == 1 the guard did nothing and each worker emitted a 1-row tail batch. The slicing now works per rank: every rank takes the same number of rows of each table or session, whole batches are spread over the rank's workers and the partial tail goes to one worker, so a pass yields one partial batch per rank instead of one per worker. Ranks stay in lockstep by construction and extra rows are only skipped when they would buy a rank an additional step, which removes the residue case analysis and the pre_total_remain chaining. Models that need at least N rows per batch set data_config.min_batch_size, which drops a shorter final batch on every rank identically; it defaults to 0 and replaces drop_remainder's mechanism, which now equals min_batch_size=batch_size. Both apply to train and eval only, so predict reads every row. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01163rPXLGHiM74LxdbWjMox

github-actions · 2026-09-19T14:42:26Z

-        drop_redundant_bs_eq_one=drop_redundant_bs_eq_one,
-        pre_total_remain=pre_total_remain,
+    world_size = _get_world_size()
+    assert num_workers % world_size == 0, (


The ambient world-size sniffing + this divisibility assert changes the contract for standalone to_batches() callers that don't go through BaseDataset.get_worker_info:

tzrec/tools/hitrate.py → faiss_util.build_faiss_index (faiss_util.py:70) calls reader.to_batches() with the defaults (0, 1) after init_process_group(). Under the documented launch (torchrun --nproc-per-node=2 -m tzrec.tools.hitrate, docs/source/quick_start/local_tutorial_u2i_vec.md:113, docs/source/models/mind.md:212), _get_world_size() returns 2 and this assert fires on every rank. Pre-PR, (0, 1) planned the full range, so each rank built the complete faiss index. Same pattern in tzrec/tools/create_online_infer_data.py:73 (to_batches() after predict() initialized dist).

The _get_world_size() env fallback also disagrees with get_worker_info (dataset.py:313-320), which hardcodes world=1 when dist is uninitialized. A process with a stale WORLD_SIZE builds ids with the world-1 view and re-divides them here with the env value — crash when they don't divide (e.g. the num_workers=3 cases in utils_test.py now fail on any machine that happens to export WORLD_SIZE=2).

Suggestion: give standalone callers an escape — an explicit world_size parameter (defaulting to the sniffed value), or have the tools pass rank-major ids the way hitrate.py:361 already does.

Fixed in 6a17bd9: the reader no longer reads the world size from the process group. to_batches(worker_id, num_workers) keeps its even-share meaning for standalone callers (hitrate, faiss_util, create_online_infer_data), and BaseDataset passes the new optional world_size explicitly.

github-actions · 2026-09-19T14:42:40Z

+        layout[-1] = (order[-1], layout[-1][1] + tail)
+        layouts.append(layout)
+        last_workers.append(order[-1])
+        cursor = (cursor + full + (1 if tail else 0)) % num_workers


The cursor advances by full + (1 if tail), which can be ≡ 0 (mod num_workers) — and when every source has the same per-rank row count (e.g. is_orderby_partition over uniformly sized daily partitions), it is ≡ 0 for every source at once, so the rotation never happens: the same workers get the q+1-batch share and the same worker gets the tail of every source, and the imbalance accumulates linearly with the session count.

Small example: bs=4, num_workers=2, sources=[13, 13] → advance = 3+1 = 4 ≡ 0 (mod 2) → worker totals 16/10, whereas a rotating cursor would give 13/13. Over 90 partitions one worker per rank does nearly all the IO while its peers idle; the old calc_slice_position split each session ±1 row across all workers, so this shape had near-zero skew before.

Consider decoupling the advance from the source sizes (e.g. also add the source index t) so uniform multi-session inputs can't resonate.

Reworked in the latest commit. The per-source layout and cursor are gone: a rank's sources are planned as one stream, ranks split it by a round-robin row count so nothing is lost across sources, and inside a rank whole batches are dealt source by source to the least loaded workers with the open batch topped up first. Uniform sources now keep worker skew within one batch regardless of the source count, and a pass ends with at most one partial batch per rank.

github-actions · 2026-09-19T14:42:45Z

+                sources.append(
+                    (
+                        f"{input_path}#{sess_req.session_id}",
+                        _get_session_record_count(client, sess_req),


All session record counts are now fetched serially before the first batch of every epoch, per worker (the old code fetched session k's count when reading reached session k, hiding the RPC behind streaming). Total RPC volume is unchanged, but the timing is front-loaded: with is_orderby_partition over N partitions, every worker of every rank issues N get_read_session calls at epoch start — N=90 with 16 ranks × 8 workers is ~11.5k identical metadata RPCs concentrated in the first seconds, and _get_session_record_count has no retry/backoff (unlike _read_rows_arrow_with_retry), so one flow-control error during the burst kills the job.

A session's record count is an immutable snapshot — consider fetching counts once in _init_session/_restore_sessions (rank 0 already creates the sessions and broadcasts the ids; counts could ride along) and caching them on the reader, which would also remove the per-epoch/per-worker refetch entirely.

Done: record counts are fetched once on rank 0 when sessions are created (thread pool, broadcast with the session ids) and taken from the get_read_session response on restore; workers read them from the cache.

github-actions · 2026-09-19T14:43:00Z

+        ]
+        for p in procs:
+            p.start()
+        results = dict(queue.get() for _ in procs)


The unbounded queue.get() runs before the p.join()/exitcode checks: if a child dies before its put (e.g. rank 0 raises in ParquetReader.__init__ before the all_gather_object collective, rank 1 blocks in gloo until the ~30 min default timeout and then exits without putting), the parent blocks forever on the second get() — a real regression becomes an opaque CI hang instead of a failure with a message. The neighboring test_parquet_reader avoids this by asserting inside the children and only joining. Consider queue.get(timeout=...) while polling p.is_alive()/p.exitcode, or move the assertions into _reader_worker.

Done: the children assert their own expectations and the parent only joins and checks exit codes.

github-actions · 2026-09-19T14:43:05Z

+                min_batch_size=min_batch_size,
+                equalize_rank_steps=True,
+            )
+            queue.put((rank, [len(b["id_a"]) for b in reader.to_batches(rank, 2)]))


to_batches(rank, 2) exercises the gloo harness at world=2 × 1 local worker, where the rank-major glue (get_worker_info's rank * nw + id ↔ calc_slice_intervals's divmod(worker_id, num_workers // world_size)) is trivially (rank, 0). The production shape of the bug being fixed — world>1 and multiple dataloader workers per rank — isn't exercised end-to-end anywhere: the utils_test simulation calls plan_rank_worker_intervals directly, and the calc_slice_intervals resume tests use batch_size=1, where the planner's tail logic is a no-op. One 2-rank × 2-local-worker case here (or a calc_slice_intervals test with world_size 2 and a realistic batch size) would pin the id decomposition and interval mapping down.

Added: the gloo harness now runs a 2 ranks x 2 workers case, and utils_test pins calc_slice_intervals at world_size=2 with two local workers, with and without a checkpoint state.

github-actions · 2026-09-19T14:43:18Z

+       first worker across sources) and the partial tail goes to the worker laid
+       out last, so a worker never manufactures a partial batch of its own.
+    3. A cumulative tail smaller than ``min_batch_size`` is dropped.
+    4. With ``equalize_rank_steps`` every rank yields the same batch-size list:


Small accuracy nit: equalize_rank_steps guarantees the same number of batches per rank, not "the same batch-size list" — with kept extras the final partial sizes differ across ranks (the PR's own expectation at parquet_dataset_test.py:517 has rank 0 tail [2] vs rank 1 tail [1]). The same wording appears in the Args below, in calc_slice_intervals, dataset.py:554, and both reader docstrings. Suggest "every rank yields the same number of batches" so nobody later relies on cross-rank shape equality.

github-actions · 2026-09-19T14:43:23Z


+    // drop train/eval tail batches with fewer rows than this, 0 disables.
+    // drop_remainder=true is equivalent to min_batch_size=batch_size.
+    optional uint32 min_batch_size = 29 [default = 0];


Question on the default: with min_batch_size=0 the exact crash from the PR description is still reachable by default — any table where (rows // world_size) % batch_size == 1 yields a 1-row tail batch on every rank at the same step (the planner concentrates the rank's residue onto one worker but doesn't drop it; [[34], 2, 4, 4, 0, 5, [1], 34] in utils_test enshrines this). The protection for BatchNorm/listwise models is opt-in, and I see the PR body rejects auto-deriving it from the model. Middle ground worth considering: a one-time WARN from the reader when a TRAIN-mode plan produces a tail of 1 row (or when min_batch_size < 2 in train), so users get a pointer to this knob instead of a ValueError in a random hourly job.

Decision: no runtime warning (it would also fire in eval, where the knob does not apply). The failure and the knob are documented instead as FAQ Q21 in docs/source/faq.md.

github-actions · 2026-09-19T14:43:45Z

+
+- 训练和评估时，丢弃掉每个数据读取进程最后一个行数小于`min_batch_size`的batch，默认为0（不丢弃）
+- 使用BatchNorm等要求batch内至少2行样本的模型时，建议设置为2，避免样本表行数恰好使最后一个batch只剩1行导致训练失败
+- 注：OdpsDataset和ParquetDataset按行切分数据时，会保证每个`proc`（rank）读取相同的步数，以避免同步训练时卡住；为此每张表（或每个分区）最多有`nproc - 1`行样本不会被读取，预测时不受影响


两个小的准确性问题：

这条「最多 nproc - 1 行不会被读取」的注放在 min_batch_size 小节里，但丢行来自 rank 步数对齐，它在默认配置的所有训练/评估（含 eval 指标统计）中都生效，与是否设置 min_batch_size 无关。不改这个开关的用户不会读到这条说明——建议在 drop_remainder 小节或 OdpsDataset/ParquetDataset 文档处也提一句，或注明默认配置即生效。

单位不精确：ParquetDataset 的规划单位是整个 input_path（所有文件合并为一个 source，见 parquet_dataset.py:252），最多丢 nproc - 1 行是全部输入合计；「每张表（或每个分区）」只对 OdpsDataset 成立（每个 input_path 一个 session；仅 is_orderby_partition 时按分区）。

github-actions · 2026-09-19T14:44:07Z

            shuffle_buffer_size,
            sample_cost_field=sample_cost_field,
            batch_cost_size=batch_cost_size,
+            **kwargs,


Nit: CsvReader is the only reader whose docstring wasn't updated — it still says drop_remainder (bool): drop last batch. (line 90) and doesn't mention min_batch_size, which CsvDataset now passes (line 74) through this newly added **kwargs. BaseReader, OdpsReader, and ParquetReader all got the new wording in this PR.

Already updated in 91e147d: the CsvReader docstring now describes drop_remainder as min_batch_size=batch_size and lists min_batch_size.

github-actions · 2026-09-19T14:44:54Z

Review summary

Ran five parallel sub-reviews (code quality, performance, test coverage, documentation accuracy, security/robustness); all completed and their findings were verified by hand.

The core planner looks sound. I traced plan_rank_worker_intervals through several configurations, including multi-source and checkpoint-resume: per-rank row accounting is exact, chunks are contiguous and disjoint (extras included), shrink and extras never collide under equalize_rank_steps, and cross-rank step equality holds by construction. calc_slice_position is cleanly removed (no remaining callers), and predict-keeps-every-row is verified at planner, reader, and dataloader level. The brute-force invariant test is what makes this fix trustworthy — nice work.

Blocking (1):

The new num_workers % world_size == 0 assert + ambient world-size sniffing in calc_slice_intervals breaks standalone to_batches() callers under torchrun: tzrec.tools.hitrate → faiss_util.build_faiss_index calls to_batches() with the defaults (0, 1) after init_process_group(), which now asserts under the documented --nproc-per-node=2 launch (pre-PR each rank read the full range and built the complete faiss index). Same pattern in create_online_infer_data.py. Details in the inline comment on utils.py:992.

Worth fixing:

Worker-load resonance: when every source has identical per-rank rows (uniform partitions under is_orderby_partition), the cursor advance is ≡ 0 mod num_workers, the rotation never happens, and skew accumulates linearly with the session count (utils.py:889).
ODPS session record counts are now all fetched serially up front, per worker per epoch — an RPC burst at epoch start with many partitions × workers, and _get_session_record_count has no retry (odps_dataset.py:632).
test_parquet_reader_equal_steps can hang CI forever when a child dies before its queue.put (parquet_dataset_test.py:560), and the world>1 × >1-worker-per-rank glue — the production shape of the original bug — is not exercised end-to-end anywhere (parquet_dataset_test.py:540).

Minor / questions:

With default min_batch_size=0 the original 1-row BN crash is still reachable whenever (rows // world_size) % batch_size == 1; a one-time WARN pointing at the new knob would close the gap without changing defaults (data.proto:161).
equalize_rank_steps docstrings overstate the guarantee (equal step counts, not equal batch sizes); the data.md row-drop note sits under min_batch_size though it applies to default train/eval, and its per-table unit doesn't match ParquetDataset's single-source planning; CsvReader docstring missed the update — inline comments on each.
FYI: duplicate ODPS paths in input_path produce duplicate prefix keys — the returned dict silently collapses them, so _combined_reader would read one planned slice twice and never read the other. Edge-case config, but a uniqueness check would fail louder.

calc_slice_intervals read the world size from the process group and required the slice count to be a multiple of it, so a tool calling to_batches() under torchrun to read the whole table on every rank, as build_faiss_index does in the hitrate job, failed the assertion. The slice geometry is now an explicit optional world_size argument of to_batches: without it every slice is an independent even share as before, and only BaseDataset, which knows the dataloader workers of each rank, passes it. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01163rPXLGHiM74LxdbWjMox

Eval metrics and predict output must cover every row, and BatchNorm runs on running statistics outside training, so a short final batch is harmless there. Both tail-drop knobs now take effect only in Mode.TRAIN; rank step equalization still applies to train and eval. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01163rPXLGHiM74LxdbWjMox

…ow counts The planner laid every table or session out on its own: each source was split by world_size, losing up to world_size - 1 rows per source, and each source could leave a partial batch even though a worker's buffer carries its tail into the next source, so a multi-source pass ended with one partial batch per source rather than per rank. The sources of a rank are now planned as one stream: ranks split it as if rows were dealt round-robin, so every source stays spread over all ranks and the rank totals differ by at most one row over the whole pass; inside a rank the stream is cut into whole batches plus one tail, dealt source by source to the least loaded workers, with the rows completing the open batch going to its holder. Extra rows and the min_batch_size cut are decided once per pass on the stream length. ODPS session record counts are fetched once when the sessions are created or restored and cached on the reader instead of by every dataloader worker at the start of every epoch. calc_slice_intervals returns a list aligned with its sources so duplicate input paths no longer collapse to one plan, and the two-rank parquet test asserts inside the children instead of blocking on a queue. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01163rPXLGHiM74LxdbWjMox

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01163rPXLGHiM74LxdbWjMox

…licing # Conflicts: # tzrec/version.py

github-actions · 2026-09-20T05:07:11Z

+    if equalize_rank_steps and min_batch_size < 2 and num_rows % batch_size == 1:
+        logger.warning(
+            "The final training batch of this pass has a single row; set "
+            "data_config.min_batch_size >= 2 if the model needs more than one "
+            "row per batch, e.g. with BatchNorm."
+        )


This warning also fires in EVAL: readers are built with equalize_rank_steps = mode != Mode.PREDICT and BaseDataset forces min_batch_size = 0 outside TRAIN. In eval none of the advice applies — it is not a training batch, a 1-row batch does not trip BatchNorm in eval mode (running stats are used), and setting data_config.min_batch_size cannot change eval behavior. Suggest rewording mode-neutrally, e.g. "The final batch of this pass has a single row; for training, set data_config.min_batch_size >= 2 if the model needs more than one row per batch (e.g. BatchNorm)."

Agreed; the warning is removed entirely rather than reworded, and the guidance moved to FAQ Q21.

github-actions · 2026-09-20T05:07:15Z

+    remaining = [
+        calc_remaining_intervals(checkpoint_state, prefix, total_rows)
+        for prefix, total_rows in sources
+    ]


Resume path, non-blocking: calc_remaining_intervals scans every key of the (globally merged, all-ranks) checkpoint state, and this now runs once per source, eagerly in to_batches before the first batch is yielded. For ODPS with is_orderby_partition (hundreds of sessions × global workers × per-interval keys), that is O(sources × keys) string parsing per dataloader worker, all paid as first-batch latency — the old code did the same total work but lazily per session, overlapped with network reads. Since calc_slice_intervals now receives all sources at once, it could bucket the checkpoint keys by source prefix in a single O(keys) pass and hand each source only its own entries.

Done: calc_slice_intervals buckets the checkpoint keys by source prefix in one pass and hands each source only its own entries.

github-actions · 2026-09-20T05:07:19Z

+        # the tail opens the next batch on the least loaded worker, whose block
+        # is laid out last so that block and tail form one interval
+        owner = min(range(num_workers), key=lambda w: (loads[w], w))
+        for w in order:
+            if counts[w] and (w != owner or not tail):
+                chunks.append((w, pos, pos + counts[w]))
+                pos += counts[w]
+        if tail:
+            chunks.append((owner, pos, pos + counts[owner] + tail))
+            loads[owner] += tail
+            holder, carry = owner, tail


Optional (perf, many-session ODPS): the carry-completion chunk (holder, 0, pos) is laid out first, but the holder's own full-batch block is laid out at its order position — only the owner's block is forced last. Unless the holder's block happens to land immediately after pos, the adjacency merge below cannot fuse them, so the holder reads two intervals for that source: a tiny <batch_size stream plus its block. That is one extra read_rows_arrow stream open (full connection/retry setup) per source boundary per rank. Since counts are fixed before layout, laying the holder's block immediately after the carry chunk (owner's block still last for the tail merge) keeps loads and determinism unchanged and recovers the merge; the exact-value tests would need updating.

Done: the topped-up worker's block is now laid out right after its carry chunk (the tail owner's block stays last), so both read a single interval per source. loads/owner selection unchanged; the brute-force sweep still reports 0 violations.

Simplified further: since a worker's buffer only depends on how many rows it takes from a source, every worker now reads exactly one contiguous chunk per source (whole batches, plus the top-up for the holder, plus the tail for the least loaded worker). The layout ordering and adjacency merge are gone; a 20k-plan sweep confirms at most one interval per worker per source.

github-actions · 2026-09-20T05:07:22Z

+        equalize_rank_steps (bool): make every rank yield the same number of
+            batches,            honored by readers that slice rows by count.


Docstring artifact: batches, honored by readers that slice rows by count. — the run of spaces mid-sentence reads like a deleted word, and ruff-format does not normalize docstring interiors, so it persists in the rendered API docs. Suggest: "batches; honored by readers that slice rows by count."

Fixed, and the sentence now also notes that the guarantee does not hold when batch_cost_size cuts batches by cost.

github-actions · 2026-09-20T05:07:40Z

+- 仅在训练时生效，评估和预测时不会丢弃任何数据；等价于`min_batch_size`设置为`batch_size`
+
+### min_batch_size
+
+- 训练时，丢弃掉每个数据读取进程最后一个行数小于`min_batch_size`的batch，默认为0（不丢弃）；评估和预测时不生效
+- 使用BatchNorm等要求batch内至少2行样本的模型时，建议设置为2，避免样本表行数恰好使最后一个batch只剩1行导致训练失败


Three accuracy nits in this block:

L337 "评估和预测时不会丢弃任何数据" contradicts the notes this PR adds above (L36/L80: 训练和评估时最多 nproc - 1 行不会被读取) — rank equalization drops in EVAL too (equalize_rank_steps = mode != PREDICT). Scope the claim to the knob, e.g. "该参数仅在训练时生效，评估/预测不会因batch不足而丢弃数据（训练/评估时为对齐各rank步数仍可能有少量样本不被读取，见上文各Dataset注意事项）".

L341 "每个数据读取进程" overstates the drop granularity: for Odps/Parquet the sub-threshold tail is dropped once per rank at plan time (at most one partial batch per rank, rows never read); only CsvDataset drops per reading process. As written, users overestimate dropped rows by ~num_workers.

Consider documenting the validated constraint min_batch_size <= batch_size (ValueError in TRAIN), and that with min_batch_size = m > 0 each rank may additionally drop up to m - 1 tail rows — the L36/L80 "最多 nproc - 1 行" bound assumes m = 0.

Done: the drop_remainder claim is scoped to the knob, min_batch_size documents per-proc granularity for Odps/Parquet vs per-process for Csv and the min_batch_size <= batch_size constraint. The nproc-1 dataset notes were removed on the maintainer's request as an implementation detail.

github-actions · 2026-09-20T05:07:44Z

        Args:
            worker_id: Worker ID
            num_workers: Total number of workers
+            world_size: Unused, partitions are split by global worker id


This adds a world_size Args entry to _reader's docstring, but _reader(self, worker_id, num_workers) has no such parameter — only to_batches gained one (L638), whose docstring already carries the identical note. Delete this line.

(Minor, optional: KafkaDataset.__init__ still passes raw data_config.drop_remainder and KafkaReader.__init__ does not forward its **kwargs to super(). Both are benign today — the endless consume loop never reaches _arrow_reader_iter's tail-drop branch, so the knobs are inert as the PR description says — but a one-line "drop_remainder/min_batch_size are inert for Kafka" in the reader docstring would document the exemption next to the code.)

Done: the stray world_size line is removed from _reader, and the KafkaReader docstring now states that drop_remainder and min_batch_size are inert because the consume loop never ends.

github-actions · 2026-09-20T05:07:50Z


    def to_batches(
-        self, worker_id: int = 0, num_workers: int = 1
+        self, worker_id: int = 0, num_workers: int = 1, world_size: Optional[int] = None


V1 accepting and ignoring world_size is fine, but the new min_batch_size knob is a silent no-op for OdpsDatasetV1: the dataset never passes it, and even if it did, OdpsReaderV1.to_batches yields from _iter_one_table/_pa_read and never reaches _arrow_reader_iter where the drop lives. A V1 user following the new docs (min_batch_size: 2 for BatchNorm models) still crashes on a 1-row tail. Consider logging a warning at init when min_batch_size > 0 and the reader cannot honor it.

Related (pre-existing, out of scope, but now contradicting the contract this PR documents): L171 allow_smaller_final_batch=self._drop_remainder is inverted — per _pa_read, drop_remainder=true keeps the smaller final batch and false drops it, the opposite of "drop last training batch". Worth a follow-up.

OdpsDatasetV1 is deprecated, so no change here.

github-actions · 2026-09-20T05:07:55Z

+            sess_infos = [(x, None) for x in session_ids]
+            if int(os.environ.get("RANK", 0)) == 0:
+                sess_reqs = [SessionRequest(session_id=x) for x in session_ids]
+                with ThreadPoolExecutor(max_workers=8) as executor:
+                    record_counts = executor.map(
+                        _get_session_record_count, [client] * len(sess_reqs), sess_reqs
+                    )
+                sess_infos = list(zip(session_ids, record_counts))
            if self._pg is not None:
-                dist.broadcast_object_list(session_ids, group=self._pg)
+                dist.broadcast_object_list(sess_infos, group=self._pg)


Non-blocking robustness note: this moves the session-readiness wait inside rank 0's pre-broadcast critical section. _get_session_record_count polls (while True … sleep(1)) until each session leaves INIT, while all other ranks are already blocked in broadcast_object_list; before this change the readiness wait happened per worker in to_batches, outside any collective. Two consequences: (a) if a count fetch raises (the exception surfaces when list(zip(...)) consumes the lazy executor.map), rank 0 dies with the real error while the other ranks sit in the broadcast until the pg timeout and then fail with an unrelated one; (b) a session that legitimately takes longer than the pg timeout to become ready now aborts an otherwise healthy job. Cheap hardening: broadcast session_ids first as before, then fetch counts and broadcast sess_infos as a second collective — or broadcast an error sentinel so all ranks fail fast with the root cause.

Done: session ids are broadcast first as before, then every rank fetches its own record counts with a bounded thread pool outside any collective, so a readiness wait or fetch error never sits inside a broadcast.

github-actions · 2026-09-20T05:08:05Z

+                label_fields=["label"],
+                num_workers=8,
+                min_batch_size=min_batch_size,
+                drop_remainder=min_batch_size == 1024,


The only drop_remainder=True case also sets min_batch_size=1024, so the reader-level mapping self._min_batch_size = batch_size if drop_remainder else min_batch_size (BaseReader.init) is never exercised on its own — deleting the if drop_remainder branch would keep the whole suite green, and a user setting drop_remainder: true alone would silently stop dropping. Add a case with drop_remainder=True, min_batch_size=0 expecting [1024] * 8.

Added a drop_remainder=True, min_batch_size=0 dataloader case expecting [1024] * 8.

github-actions · 2026-09-20T05:08:08Z

+    @parameterized.expand(
+        [[Mode.EVAL], [Mode.PREDICT]], name_func=parameterized_name_func
+    )
+    def test_create_dataloader_eval_predict_keep_every_row(self, mode):


"keep every row" is verified single-process only, where equalization provably cannot drop (world_size=1 → cut = total - base = 0). Multi-rank EVAL still equalizes (equalize_rank_steps = mode != PREDICT) and drops up to W - 1 rows when (rows // W) % batch_size == 0 — intended and documented (data.md L36/L80), but no test pins it and the name overclaims. The two-rank gloo harness in ParquetReaderTest below makes an EVAL case cheap; alternatively scope the test name/comment to single-rank and note the multi-rank eval row loss.

Renamed to test_create_dataloader_eval_predict_no_tail_drop with a comment; the multi-rank eval lockstep drop is pinned by test_parquet_reader_equal_steps (33 rows, equalize on, threshold 0 is exactly the eval reader configuration).

github-actions · 2026-09-20T05:08:14Z

        # Worker 0 gets first half of total rows
-        result_w0, _ = calc_slice_intervals(
-            total_rows=1000,
+        result_w0 = calc_slice_intervals(


Both world_size=None tests assert sums only (total_w0 + total_w1 == 800 at L340, topology total == 400) — slices that overlap and skip rows would still pass, and this is exactly the path faiss_util.build_faiss_index / hitrate rely on. Assert pairwise disjointness and union == the remaining intervals, as test_calc_slice_intervals_two_ranks_two_workers does with exact values.

Also: (a) the wrapper's multi-source path (per-source zip(remaining, plan) + _map_logical_range — what the rewritten OdpsReader.to_batches now depends on) has no credential-free test; a two-source case with a checkpoint on the first source would cover it cheaply. (b) Minor: the bound asserts in test_plan_rank_worker_intervals_invariants (all(b <= batch_size) / all(b >= min_batch_size), L109-110) can never fail, since _rank_batches constructs the batch list to satisfy them by design — they don't add the coverage they appear to.

Done: the world_size=None tests now assert the slices are pairwise disjoint and tile the remaining intervals, a two-source case with a checkpoint on the first source covers the wrapper's multi-source path, and the harness asserts the raw worker residue instead of the constructed batch list.

github-actions · 2026-09-20T05:08:36Z

Review summary

Reviewed via five parallel areas (code quality, performance, test coverage, documentation, multi-process safety); all five completed. Findings are posted inline — none are blocking correctness bugs in the new planner.

Verified clean (independently re-checked, not just agent-reported):

plan_rank_worker_intervals is deterministic across ranks: pure integer arithmetic, no dict/set iteration-order dependence, total-order tie-breaks (loads[w], w). The equal-step invariant holds in all three cut branches (including batch_size == 1, empty passes, and rows < world_size), and the "at most one partial batch per rank" invariant holds across sources because _arrow_reader_iter's buffer spans the whole combined reader.
Dropped rows are never read (plan-time cut shrinks shares before any interval is emitted); ODPS session counts are fetched once on rank 0 and piggyback the existing broadcast, so workers do pure dict lookups.
The old calc_slice_position / drop_redundant_bs_eq_one / pre_total_remain plumbing is fully removed with no dangling references; world_size=None preserves the legacy even-share semantics tools rely on.
The Kafka exemption claim checks out: the endless consume loop never reaches the tail-drop branch, so both knobs are genuinely inert there.

Inline findings (details on each line): the single-row warning also fires in EVAL where its advice is a no-op; data.md's "评估和预测时不会丢弃任何数据" contradicts the new nproc-1 notes and overstates drop granularity; min_batch_size is a silent no-op for OdpsDatasetV1 (it never reaches _arrow_reader_iter); rank-0's unbounded session-readiness wait now sits inside the pre-broadcast critical section; the resume path rescans the whole checkpoint state once per source eagerly; two carry/layout optimizations; a garbled BaseReader docstring line; a docstring line for a parameter _reader doesn't have; and three test gaps (drop_remainder equivalence untested on its own, multi-rank EVAL row loss unpinned, sums-only assertions on the world_size=None path).

One minor note not worth an inline: the equalize_rank_steps contract ("every rank yields the same number of batches") does not hold when batch_cost_size is set — batches are then cut by cumulative cost, so step counts depend on which rows each rank received. The old row-based guard had the same tension, so this is not a regression, but a one-time warning (or a docstring/doc exception) when both are set would keep the new flag from overclaiming.

Overall this is a well-engineered fix: the planner's invariants are provable from the code, the property-test sweep is unusually thorough, and the PR description's claims all survived verification.

…and tests Rank 0 fetched every session's record count inside the pre-broadcast critical section, so a slow or failing readiness poll left the other ranks blocked in the collective; now the session ids are broadcast first and every rank fetches its own counts outside any collective. The planner lays the topped-up worker's block right after its carry chunk so it reads one interval per source, and calc_slice_intervals hands each source only its own checkpoint entries instead of rescanning the whole state per source. The single-row warning is replaced by a FAQ entry because it also fired in eval, where the knob does not apply; docs and docstrings are corrected and the tests assert disjoint tiling, the raw worker residue, and drop_remainder on its own. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01163rPXLGHiM74LxdbWjMox

A worker's buffer only depends on how many rows it takes from a source, not on where they sit, so the planner now gives every worker a single chunk per source: its whole batches, plus the rows completing the open batch for the holder, plus the tail for the least loaded worker. This drops the layout ordering and adjacency merge that previously left the holder with two intervals when it also owned the tail. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01163rPXLGHiM74LxdbWjMox

calc_remaining_intervals inferred each keyed chunk's end from the next key and attributed the rows after the last key to it, but rows before the first key were treated as consumed. A worker with fewer batches in one source reaches the next source before its peers, so a checkpoint can hold a key for a later chunk of that source while an earlier chunk has no key yet; on resume those earlier rows were skipped. Rows before the first key can only lack a key because no batch ever read them, so they are now returned as remaining. The planner tests also assert that every worker gets at most one interval per source. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01163rPXLGHiM74LxdbWjMox

tiankongdeguiji added the claude-review Let Claude Review label Sep 19, 2026

github-actions Bot removed the claude-review Let Claude Review label Sep 19, 2026

github-actions Bot reviewed Sep 19, 2026

View reviewed changes

tiankongdeguiji and others added 5 commits September 20, 2026 09:43

[chore] bump version to 1.4.12

b1bad3c

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01163rPXLGHiM74LxdbWjMox

Merge remote-tracking branch 'origin/master' into bugfix/rank-level-s…

0207b22

…licing # Conflicts: # tzrec/version.py

tiankongdeguiji added the claude-review Let Claude Review label Sep 20, 2026

github-actions Bot removed the claude-review Let Claude Review label Sep 20, 2026

github-actions Bot reviewed Sep 20, 2026

View reviewed changes

tiankongdeguiji and others added 3 commits September 20, 2026 15:03

eric-gecheng approved these changes Sep 20, 2026

View reviewed changes

tiankongdeguiji merged commit 02a4650 into alibaba:master Sep 20, 2026
7 checks passed

tiankongdeguiji deleted the bugfix/rank-level-slicing branch September 20, 2026 10:42

		equalize_rank_steps (bool): make every rank yield the same number of
		batches, honored by readers that slice rows by count.

Conversation

tiankongdeguiji commented Sep 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Change

Alternatives considered

Test Plan

Uh oh!

github-actions Bot Sep 19, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 19, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 19, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 19, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 19, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 19, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 19, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 19, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 19, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Sep 19, 2026

Review summary

Uh oh!

github-actions Bot Sep 20, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 20, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 20, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 20, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 20, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 20, 2026

tiankongdeguiji commented Sep 19, 2026 •

edited

Loading