allenai / allenai/open-instruct

`reduce_loss=sum` removed in #1024 without deprecation warning: LR miscalibration and mean-of-microbatch-means under gradient accumulation

Offen
#1,728 3 Kommentare 1 Reaktion 0 zugewiesene Personen Auf GitHub ansehen
Vorherrschende Sprache
Python
Sterne
3.9k
Forks
585
Ø Merge
5 T. 17 Std.
Gemergte PRs (30 T.)
16

Beschreibung

## Context

I have been attempting to reproduce the Tulu 3 SFT results (Llama-3.1-8B, full fine-tuning, `tulu-3-sft-mixture`) using `open-instruct` as of April 2026, following the hyperparameters reported in section 4.3.2 of the Tulu 3 paper. I am observing large, consistent performance gaps across most benchmarks, which I believe are at least partly explained by the fact that the loss computation was changed from sum to mean without a warning against using a learning rate optimal for a summed loss on a mean loss training run.

## Performance gap

The table below compares my reproduction of AllenAI's results on the **released Tulu 3 checkpoints** via `olmes` (confirming my evaluation pipeline is correct) against my own retrain using the current codebase and the paper's published hyperparameters:

| Task | Ai2 Llama-3.1-Tulu-8B-SFT (my `olmes` repro) | My retrain | Δ |
|---|---|---|---|
| GSM8K | 75.6 | 59.9 | -15.7 |
| DROP | 61.1 | 37.1 | -24.0 |
| HumanEval | 86.0 | 66.3 | -19.6 |
| HumanEval+ | 79.5 | 58.6 | -20.9 |
| IFEval | 71.7 | 54.9 | -16.8 |
| PopQA | 29.4 | 26.6 | -2.8 |
| MMLU | 62.1 | 57.8 | -4.3 |
| MATH (4-shot) | 31.1 | 17.3 | -13.8 |
| BBH (3-shot) | 69.2 | 43.5 | -25.7 |
| TriviaQA | 46.9 | 49.8 | +2.9 |

The gaps are large and systematic, yet on reasoning tasks the model does engage in chain-of-thought, just outputs a wrong answer at the end. To rule out an evaluation mismatch, I ran `olmes` on both the checkpoints released by AllenAI on huggingface, and my attempt at retraining.

## My training configuration

- 4 nodes × 4 H100 GPUs (Jean Zay), DeepSpeed ZeRO-3
- `per_device_train_batch_size=2`, `gradient_accumulation_steps=4`
- Effective batch size: **128 sequences / optimizer step** (32 microbatches)
- `learning_rate=5e-6` (Tulu 3 paper recommendation)
- `max_seq_length=4096`, 2 epochs, Llama-3.1-8B, full fine-tuning
- `open-instruct` as of April 2026 (post-#1024)
- chat_template: tulu

## Root cause hypothesis: the training loss?

PR #1024 removed the `reduce_loss` argument and hardcoded `mean` loss. The Tulu 3 models were trained with `reduce_loss=sum` (as confirmed in Issue #995) and all published hyperparameters — including the learning rate of `5e-6` — were calibrated to that setting. Naturally, I used the published LR, unaware that the loss reduction had changed and that this LR is no longer appropriate. Unless I missed it, there is no warning, no error, and no mention of this in the documentation.

## Loss computation with gradient_accumulation > 1: mean-of-microbatch-means

Beyond the LR mismatch, I believe the current `mean` implementation is itself incorrect under gradient accumulation > 1.

The current training step in open_instruct/finetune.py (l. 817-824) is:

```python
with accelerator.accumulate(model):
if args.load_balancing_loss:
outputs = model(**batch, use_cache=False, output_router_logits=True)
total_aux_loss += outputs.aux_loss.detach().float()
else:
outputs = model(**batch, use_cache=False)

loss = outputs.loss
```

`outputs.loss` is the mean cross-entropy over non-`-100` tokens in the current microbatch (HuggingFace default, see [`transformers/loss/loss_utils.py`](https://github.com/huggingface/transformers/blob/main/src/transformers/loss/loss_utils.py)). Under `accelerator.accumulate`, gradients are averaged across the accumulation window, which amounts in this case to mean of microbatch means. A token-weighted mean would compute the sum of the loss across microbatches and divide that by the total of supervised tokens in the batch.

The deleted `reduce_loss=sum` implementation avoided this by computing `CrossEntropyLoss(reduction="sum")` per microbatch, making the per-microbatch gradient proportional to token count. The deleted code even included an explicit comment acknowledging this:

> "this ensures that we weight all tokens in the dataset equally, rather than weighting each overall example equally when using high amounts of gradient accumulation. this can result in > 5 point improvements in AlpacaEval"

That reasoning was removed along with the code, without a replacement (deleted lines 790-819 in [this commit](https://github.com/allenai/open-instruct/commit/bb98dbcac5203ddb4b17731a0fa6157883539792#diff-408d69efc79bb0c4dd513149521ed724e4f8fb12e8db91362600f258a3e22348L793))

## The sequence-parallel path already does this correctly

The SP branch (lines 827–841 of current `finetune.py`) implements the correct token-weighted mean:

```python
total_loss_sp = sum(
losses_per_rank[r] * good_tokens_per_rank[r]
for r in range(args.sequence_parallel_size)
if good_tokens_per_rank[r] > 0
)
loss = total_loss_sp / torch.clamp(total_good_tokens, min=1)
```

The standard non-SP path has no equivalent correction. This asymmetry suggests the omission is unintentional.

### Is there something wrong in my reasoning or did I miss something crucial?
I plan to rerun with the old sum loss and will report back. In the meantime, any feedback would be very welcome.

## References

- Tulu 3 paper, section 4.3.2
- PR #1024 (removal of `reduce_loss`)
- Issue #995 (confirmation that `sum` is no longer supported but that Tulu 3 models were trained using it)
- HuggingFace transformers issue [#24725](https://github.com/huggingface/transformers/issues/24725)

Beitragsleitfaden

Beitragsleitfaden öffnen

Bewertung

Dieses Issue wurde noch nicht bewertet.

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.