[Bug] NaFlex linear loss scaling uses different coefficients within one accumulation window

Summary

With NaFlex linear loss scaling and gradient accumulation, each replay in the accumulation loop scales a loss computed over the whole cached window by its own microbatch size. Uneven microbatches therefore weight parts of the same update differently.

Changing only the microbatch split changes the gradient, even when the examples, cached features, and effective batch are identical.

System Info

Environment: OpenCLIP commit 602d4af7, FP32 on CPU, through _train_step_eager.

The results below are from the pinned version. Source inspection on 2026-09-10 found the same calculation in 2d534609; the full experiment has not been rerun on that commit.

The separate regression on current source baseline 2d534609 and PR #1212 ran on Ubuntu 26.04 LTS (x86_64), Python 3.11.15, PyTorch 2.13.0+cpu, and timm 1.0.28. The linked test uses the native NaFlex model and training step.

Reproduction

The test sets naflex_loss_scale='linear', accum_freq=2, and reference batch_size=3, keeps the same three ordered examples and model initialization, and compares the splits [1, 2] and [2, 1].

The first divergence is in the loss coefficients, before any gradient is computed. A common window scale removes the coefficient mismatch and brings the gradient difference down to the control level.

The regression code is in PR #1212. From a checkout of that PR with its test dependencies installed:

python -m pytest -q tests/test_grad_accum.py

The CPU run passed 26 tests. The two uneven linear-scaling cases fail on the unmodified baseline; equal-size and scale-disabled controls pass.

Expected behavior

Every replay in the same all-NaFlex accumulation window should use one common linear loss multiplier. Repartitioning the same cached batch as [1, 2] or [2, 1] should then give matching gradients within numerical tolerance.

Root cause

In the accumulation replay loop, _train_step_eager recomputes one microbatch while using the other cached features as negatives, then calls get_naflex_loss_scale(batch_j, args, task) for that replay.

The loss sees the whole window, but its multiplier depends on batch_j alone. Different microbatch sizes therefore assign different coefficients to the replayed gradient contributions.

Proposed fix

For an all-NaFlex accumulation window in linear mode, compute one scale before entering the replay loop:

sum(window_batch_sizes) / (len(window_batches) * reference_batch_size)

and use it for every replay. In the reproduction this gives 3 / (2 * 3) = 0.5 regardless of the split. The tested patch makes this change inside _train_step_eager and checks that the reference batch size is positive.