feat(train): add a shared epoch-aware multi-task sampling policy #6097

Context and compatibility

Follow-up to #6096 and the task-sampling discussion.

For #6096, the compatibility decision is to preserve PT's existing default task-weighting behavior and align PT_expt with it. The broader change proposed below is deferred to this issue, not a requirement for that PR.

Problem

When training uses numb_steps without explicit model_prob, the legacy task-size heuristic does not represent the number of training batches:

• PT counts retained frames for LMDB, whereas PT_expt has counted one unit per LMDB file.
• NPY data uses system counts, not frame counts. Mixed-format inputs therefore combine different units.
• LMDB filtering affects retained-frame counts, but these counts still do not capture resolved batch sizes, internal resampling, or grouped sampling probabilities.
• A zero-probability group can contribute to the legacy task-size estimate despite never being sampled.

Aligning PT_expt with PT removes the backend discrepancy but intentionally retains these limitations. For example, equal-sized LMDB tasks with batch sizes 1 and 10 receive equal task probabilities, although they need different numbers of steps per nominal epoch.

Proposed follow-up

Introduce one shared policy for PT and PT_expt that can derive task probabilities from each task's nominal epoch length:

p_t = E_t / sum(E)

Here, E_t is the number of task-local optimizer steps per nominal epoch, resolved after filtering, batching, internal sampling, group probabilities, and distributed batching. Reuse the epoch-length contract used by the training-horizon resolver rather than introducing another independent size heuristic.

With a fixed total step budget, this policy gives tasks equal expected nominal epoch counts. It does not guarantee exact visitation counts or equal contributions from every frame, nor invariance to storage changes that alter batching/padding.

Requirements:

• Use a shared implementation and identical resolved task probabilities on every distributed rank.
• Keep task selection, group selection, and within-group system sampling distinct.
• Preserve explicit model_prob behavior and the existing num_epoch_dict scheduling contract.
• Exclude zero-probability groups from nominal epoch accounting and handle empty/invalid workloads explicitly.
• Define a compatibility strategy before changing defaults. An explicit opt-in policy and a versioned default transition are alternatives to evaluate, not decisions already made. Existing NPY defaults must not silently change to frame-based or epoch-based weighting.

Acceptance criteria

• Cross-backend tests cover NPY, LMDB, and mixed-format tasks.
• Tests cover unequal batch sizes, filtering, internal resampling, multiple recursively discovered sources, unequal group probabilities, and zero-probability groups.
• Expected task probabilities follow the documented nominal epoch lengths, including uneven batches and distributed tails.
• Distributed ranks resolve identical probabilities.
• Explicit model_prob and num_epoch_dict retain their documented semantics.
• Compatibility tests preserve legacy defaults until an explicitly selected migration policy applies.
• Documentation explains the three sampling levels, expected nominal epoch semantics, and migration path.

Related: #5814, #5869, #5879.