guarin/xiaomi-lens-train-repro

Reproduction of stage 2 of the Xiaomi-Lens-0 training protocol (paper: arXiv:2610.23279, code: SeerRay-Lab/Xiaomi-Lens-0, data recipe: examples/train/), run twice on the toy dataset:

1. run1-original — the original repo's train.py + XiaomiLensTrainer with the SeerRay-Lab/Xiaomi-Lens-0 checkpoint (transformers 5.4.0, torch 2.10.0).
2. run2-pr — transformers PR #50028 (commit 25eeec1ba7421667e00260002ccb326a769ea4f9) with guarin/Xiaomi-Lens-0, loss computed by the PR model's internal loss_function, same data pipeline and batch sampler as run 1.

Result

The PR implementation reproduces the original's loss to within bf16/kernel tolerance. On 6 fixed batches (24 toy rows, byte-identical inputs verified by input_ids checksums, dataloader order verified identical), forward-only loss parity:

batch	loss (original)	loss (PR)	abs diff
0	0.74830	0.75255	0.00425
1	0.80895	0.80760	0.00135
2	0.71534	0.71545	0.00012
3	0.74080	0.74175	0.00095
4	0.90019	0.89674	0.00346
5	0.79423	0.78993	0.00430

Max abs diff 0.0043 (~0.5% relative). Full artifact: parity/parity.json. The full 4,622,144,192-parameter checkpoint also loads under the PR with 0 missing / 0 unexpected keys (automatic legacy→HF weight-key conversion), and the PR's own tests/models/xiaomi_lens0 modeling tests pass on the subset run.

Training curves (15 optimizer steps, 1 epoch, both runs identical batches in identical order):

• run1-original: loss 0.7552 → 0.2741 (lm_loss 0.6839 → 0.2330, mask_loss 0.0713 → 0.0411)
• run2-pr: loss 0.8417 → 0.2774 (mask_loss 0.0872 → 0.0438)

Both descend with the same shape and land within 0.003 at the final step. Per-step offsets (up to ~0.09 early) are larger than the 0.0043 fixed-input parity tolerance; they come from (a) different accumulation contracts (below) and (b) gradient-path divergence: once any update differs, later steps diverge stochastically while remaining statistically equivalent.

Accumulation contracts

HF Trainer has two modes for gradient accumulation, selected by the model's accepts_loss_kwargs attribute (trainer.py:517):

• Contract A (num_items path) — Trainer computes num_items_in_batch = the number of supervised label positions (shifted labels != -100) over the whole accumulation window (all ranks, average_tokens_across_devices=True), passes it to the forward, and skips its own /gradient_accumulation_steps division (trainer.py:2088-2093). The model must produce losses that sum to the correct window quantity.
• Contract B (default path) — Trainer passes no num_items and divides every microbatch loss by gradient_accumulation_steps itself (trainer.py:2091).

Original uses Contract A (modeling_xiaomi_lens.py: accepts_loss_kwargs = True): the LM term is token_sum_CE × world_size / num_items_in_batch (lm_loss.py:: _token_sum_scale) — the exact per-token mean over the window — and the model divides the mask term by gradient_accumulation_steps itself, so the window mask contribution is the equal-weight average of the per-microbatch mask means.

PR uses Contract B (modeling_xiaomi_lens0.py:1948: accepts_loss_kwargs = False): every microbatch loss is a plain per-microbatch mean and Trainer applies the /G division. The mask term therefore ends up with the same accumulation semantics as the original (per-microbatch mask mean averaged over the window; identical DETR-style Σ(dice + 2·bce) / num_valid_masks reduction and cross-rank N/W count normalization). The observable difference is the LM term: an equal-weight mean of microbatch means instead of the exact window token mean — only visible when microbatches carry unequal supervised-token counts (toy batches are near-uniform, hence the match).

Latent trap in the PR: its loss_function already accepts num_items_in_batch via **kwargs, so activating Contract A (flipping accepts_loss_kwargs = True, or computing num_items in a custom compute_loss — the standard HF recipe for exact accumulation) window-normalizes the CE while the mask term loses its /G division: the mask gradient becomes G× overweighted. Matching the original exactly would require three changes in the PR: (1) accepts_loss_kwargs = True; (2) accept gradient_accumulation_steps in forward and divide the mask term by max(gradient_accumulation_steps or 1, 1) when num_items_in_batch is not None — mirroring the original's guard; (3) plumbing, since Trainer passes num_items_in_batch but never gradient_accumulation_steps — either a documented two-line compute_loss override or an upstream Trainer enhancement. A cleaner long-term alternative is a Trainer mechanism that also reduces the window's valid-mask count (num_masks_in_batch) so the mask loss can be a true window mean instead of a mean of microbatch means; that changes behavior relative to the original, so it is a follow-up proposal, not an equivalence fix.

Protocol

Toy data: examples/train/make_toy_data.py --num-images 48 --seed 0 → 240 stage-2 rows, manifest examples/train/stage2.yaml. Stage-2 flags from scripts/train_stage2.sh: 1 epoch, per-device batch 4 × grad-accum 4 (effective 16), LR 5e-6 cosine, 3% warmup, weight decay 0, full-parameter training, bf16, model_max_length 8192, max_pixels 4194304, min_pixels 65536, mask loss weights BCE 2.0 / Dice 1.0 / chunk 8, seed 42. Both runs consume identical batches in identical order: the repo's BatchSamplerShard (seeded randperm over the 240 rows) with group_by_modality_length disabled.

Deviations from scripts/train_stage2.sh (both runs, documented):

• no DeepSpeed (single GPU; ZeRO-2 on one rank is plain DDP)
• attn_implementation sdpa instead of flash_attention_2 (no kernel build)
• run 1: lm_loss_impl hf instead of the liger fused_linear_ce default, so both sides use the same HF reference LM-loss path
• run 2: HF TrainingArguments (warmup_steps 0.03 as a float ratio, per the repo's argparse semantics)

Files

• run1_stage2.py — run-1 driver
• run2_stage2.py — run-2 driver (includes the BatchSamplerShard dataloader override and the verified collator→PR mask remap)
• parity_stage2.py — fixed-batch parity job; parity/parity.json — verdict
• remap_check.py — proof that the collator→PR mask remap is pixel-exact against the repo's own rasterizer (repo LazyCollator emits packed-bit canvases at 16 × patch grid; the PR forward takes per-image mask_labels lists — no resize involved)
• run1/, run2/ — curves.json, trainer_state.json, tensorboard events

Runs executed on H200 (each under 4 min wall including setup); metrics dashboard: guarin/xiaomi-lens-stage2-repro-trackio.