feat(data): support deterministic weighted dataset mixing

[lkxdsb]

Summary

• add deterministic weighted mixing for any number of map-style SFT datasets (K >= 2)
• support repeatable --dataset-source NAME=PATH:WEIGHT inputs and a versioned JSON manifest
• define stop, cycle, and renormalize exhaustion policies with optional per-epoch sample budgets
• preserve source identity and report deterministic plan hashes plus per-source scheduled, filtered, trained, and target-token counts
• add CPU/adversarial coverage and a two-T4 Kaggle validation report with captured evidence

Motivation

AReno previously required project-specific loader code to combine datasets, making weighted runs difficult to validate and reproduce. This change keeps mixing inside the existing dataset-loader and SFT contracts and validates every source before model or worker initialization.

Usage

Dataset mixing is opt-in for SFT. Use either repeated inline sources or a JSON manifest; both forms are mutually exclusive with --dataset-path.

Inline sources

areno train \
  --algo sft \
  --ckpt /path/to/model \
  --dataset-source math=examples/sft/mixed/math.jsonl:0.7 \
  --dataset-source code=examples/sft/mixed/code.jsonl:0.3 \
  --dataset-mix-seed 42 \
  --dataset-mix-exhaustion cycle \
  --dataset-mix-samples-per-epoch 1000 \
  --dataset-loader-fn examples/sft/mixed/dataset_loader.py \
  --world-size 1 --tp-size 1 \
  --batch-size 2 --mini-bs 1 \
  --epochs 1 \
  --metrics-log-dir outputs/mixed-sft/metrics

Weights must be finite and positive and are normalized automatically. Inline sources default to seed 42, cycle, per-source shuffling, and an epoch budget equal to the total loaded row count.

JSON manifest

{
  "version": 1,
  "seed": 42,
  "exhaustion": "cycle",
  "shuffle_within_sources": true,
  "samples_per_epoch": 1000,
  "sources": [
    {"name": "math", "path": "math.jsonl", "weight": 0.7},
    {"name": "code", "path": "code.jsonl", "weight": 0.3}
  ]
}

areno train \
  --algo sft \
  --ckpt /path/to/model \
  --dataset-mix-config examples/sft/mixed/mix.json \
  --dataset-loader-fn examples/sft/mixed/dataset_loader.py \
  --world-size 1 --tp-size 1

Policy semantics:

• stop: end at the first selected-source exhaustion without repeating records
• cycle: restart exhausted sources until the configured epoch budget is reached
• renormalize: remove exhausted sources and continue with normalized remaining weights, emitting each record once

--dataset-mix-samples-per-epoch bounds scheduled rows, not downloads or accepted rows. Length filtering can consume extra scheduled rows, and sample weights need not equal final target-token contribution.

Observability

• stage=dataset_mix_plan: normalized contract and deterministic hashes
• stage=dataset_mix_progress: per-source scheduled, filtered, trained, and target-token counts
• dataset_mix_plan.<pid>.epoch-<n>.json: sample-free per-epoch plan under --metrics-log-dir

Rows retain source identity in reserved __areno_meta__ metadata. Invalid weights, duplicate names, empty/incompatible sources, unknown fields, and reserved-field collisions fail before expensive initialization.

Validation

Automated

• dataset-mixing and SFT trainer suite — 66 passed
• dataset-mixing CLI coverage — 8 passed
• full Apple Silicon CPU run after merging current main — 703 passed, 48 skipped; the remaining 9 failures reproduce unchanged on a clean origin/main worktree and are unrelated platform/network test assumptions
• pre-commit run --all-files
• python -m compileall -q areno tests
• Sphinx dummy build with warnings as errors
• conflict-resolution compatibility with the current backend/LoRA configuration contract and multimodal SFT records (prompt + response, tokens + prompt_mask, and image-backed responses)
• deterministic matrix across 2, 3, 8, and 32 sources, all policies, multiple budgets/seeds, and different PYTHONHASHSEED values

Kaggle GPU

Validated with Qwen3-0.6B FP16 on two Tesla T4 GPUs (world_size=2, tp_size=2, dp_size=1), PyTorch 2.10.0+cu128, native attention, and the compiled AReno CUDA extension:

• three datasets scheduled at 0.6 / 0.3 / 0.1 produced the same hash across the 128-token, 256-token, and 500-step runs
• increasing the token limit from 128 to 256 reduced the 50-step filter rate from 39.39% to 14.53%
• the final run completed 500 steps: 1,171 scheduled, 171 filtered, 1,000 trained, and 77,533 target tokens
• post-filter trained-row shares were 53.3% / 36.1% / 10.6%; no NaN, Inf, CUDA, or worker errors occurred
• first-50 vs. last-50 mean loss: 2.088 vs. 1.530 (26.71% decrease)

See the Kaggle validation guide and its 12 notebook screenshots.

Limitations

• map-style datasets only; the epoch index schedule is precomputed
• weights apply before tokenization/length filtering, so effective row and token shares can differ
• the schedule hash covers source names and selected indices, not immutable remote dataset contents
• deterministic replay starts at an epoch boundary; exact mid-epoch optimizer/data-cursor resume is out of scope

Closes #198

[lkxdsb]

Hi maintainers, a gentle follow-up on this PR.

I have merged the latest main, resolved the remaining conflicts, and confirmed that the PR is mergeable again. The focused dataset-mixing and CLI tests, pre-commit checks, compile checks, and documentation build have been rerun, and the validation details in the PR description are up to date.

The fork workflows are currently awaiting maintainer approval. When convenient, could you please take a look and approve the CI runs/review the implementation? I would be happy to address any feedback.

Thank you for your time!