perf(recipes): skip scale_grads when gradient_accumulation_steps == 1

[n-dlms]

Summary

When gradient_accumulation_steps == 1, the loss is already normalized per-token by the loss function (returning mean CE per chunk). The recipe currently multiplies it back by the token count before backward(), then scales all parameter gradients via scale_grads(1 / num_tokens) to undo that multiplication. When accumulation == 1 these two steps cancel — the multiply and the full-parameter gradient scaling pass are redundant, costing ~32 GB of memory traffic (~16ms on H100 at 2TB/s) per optimizer step.

This PR skips both for single-device recipes only: use mean loss directly for backward, no scale_grads pass needed.

Why single-device only?

Distributed recipes cannot safely skip this pass. The current_num_tokens (per-rank valid-token count) need not equal each other across ranks (ragged batches are normal with variable-length instruction fine-tuning). The existing distributed path normalizes the loss by the all-reduced global token count, which cancels out the per-rank multiplication only when every rank has the same number of valid tokens. Skipping this path would up-weight ranks with fewer tokens — a silent convergence shift, not a crash. See the NOTE comments in full_finetune_distributed.py and lora_finetune_distributed.py.

The 3 other distributed recipes (knowledge_distillation_distributed, qat_distributed, qat_lora_finetune_distributed) were not in the original PR scope and remain untouched.

Changed files

• recipes/full_finetune_single_device.py — skip when grad_accum == 1 and not optimizer_in_bwd; scale_grads gated out; clip_grad_norm moved outside the scaling gate; loss_value denominator adjusted
• recipes/lora_finetune_single_device.py — skip when grad_accum == 1 (no in-bwd path in LoRA); logging denominator adjusted
• recipes/qat_single_device.py — skip applied; optimizer.step/zero_grad kept outside the skip gate; grad_norm = None initialized for robustness
• recipes/full_finetune_distributed.py and recipes/lora_finetune_distributed.py — no logic change. Only a NOTE comment explaining why the skip is invalid in distributed mode (rank-local vs global token counts)
• tests/recipes/test_full_finetune_single_device.py — added parametrize case (8, 1, False) to exercise the new skip path

Verification

Math equivalence verified by an 8-test CPU suite:

1. Single-device skip == non-skip — gradient values identical (fp32 exact) for n_valid in {1,3,5,8}
2. All-masked batch — both paths raise the identical RuntimeError (no regression)
3. Distributed imbalanced [5,3] — skip gradients differ by 25% from correct (proves skip is invalid)
4. Distributed extreme imbalance [1,64] — skip gradients differ by 242% (confirms the reviewer's concern)
5. Distributed balanced [4,4] — skip == correct (the only safe case, proving the condition is equality of per-rank token counts)
6. Gradient accumulation > 1 — two-pass accumulation == combined pass (regression test for original path)
7. Logging equivalence — loss_to_log identical between skip and non-skip paths
8. Grad clipping — clip_grad_norm_ produces identical norms whether skip or non-skip (after fix)

[ErenAta16]

The single-device half of this looks fine to me. world_size is 1, there is no all-reduce, and num_tokens == current_num_tokens when accumulation is 1, so the multiply-then-divide really is a round trip and skipping it is free.

I am less sure about full_finetune_distributed.py, and I would like the author to confirm the reasoning before it lands, because the two token counts there are not the same quantity.

current_num_tokens is this rank's local count. num_tokens is all-reduced immediately before the scaler runs:

if (batch_count + 1) % self._gradient_accumulation_steps == 0:
    if not self._optimizer_in_bwd:
        # Get total number of tokens across all ranks to normalize gradients
        torch.distributed.all_reduce(num_tokens)
        ...
        self._grad_scaler(
            list(self._model.parameters()),
            self.world_size / num_tokens,
            ...
        )

So with accumulation at 1 the existing path applies current_num_tokens_local on the loss and world_size / num_tokens_global on the gradients. Those cancel to 1 only when every rank contributed the same number of tokens, since that is the case where num_tokens_global == world_size * current_num_tokens_local. With ragged batches, which is the normal case for instruction fine-tuning with variable-length samples, they do not cancel: the current code normalizes by the global token count, and skipping both steps normalizes each rank by its own local count instead.

The direction of the error is that ranks with shorter batches get up-weighted relative to ranks with longer ones, which is exactly the imbalance the global normalization exists to remove. It would not crash or NaN, it would show up as a small shift in convergence that is very hard to attribute later.

Two things that would settle it:

If sequences are padded to a fixed length in every supported distributed config, then token counts per rank are equal by construction and the cancellation is exact. I could not convince myself that holds, given current_num_tokens is computed from a mask sum rather than from the padded shape, but you would know.

If it does not hold, the optimization can still be kept by scaling by world_size alone rather than skipping the call, since that is the factor that does not cancel. That keeps the saved read-write pass over num_tokens out of the equation while preserving the global normalization.

Worth noting that a single-GPU smoke test cannot distinguish these, because world_size is 1 there and every version agrees. If you have a two-rank run available, comparing gradient norms for a step with deliberately uneven sequence lengths across ranks against main would answer it directly.

The perf argument itself is sound and the saving is real, so this is about scoping rather than about whether to do it.

[n-dlms]

Thanks for the careful review — you were absolutely right about the distributed case. I've addressed both points:

Distributed (your concern)

The distributed skip has been fully reverted. Both full_finetune_distributed.py and lora_finetune_distributed.py now keep the original loss multiply + scale_grads path, with only a NOTE comment explaining why the skip is invalid (rank-local current_num_tokens vs global all-reduced num_tokens). The optimization is now restricted to single-device recipes where world_size == 1 and the two token counts are identical.

I verified this with focused math tests — on a simulated 2-rank setup with imbalanced batches [5, 3], skipping produces gradients that differ by 25%, and with extreme imbalance [1, 64] the difference is 242%.

Single-device bugs found and fixed

While adding the verification tests I found three regressions in the original skip logic, all now fixed:

1. qat_single_device: self._optimizer.step() and zero_grad() were inside the and not skip_loss_scale condition — with accum=1 the optimizer would never step. Restored step/zero_grad outside the skip gate.
2. full_finetune_single_device: clip_grad_norm_ was inside the skip gate — gradient clipping would be silently disabled with accum=1. Moved clip to its own conditional block.
3. All 3 single-device recipes: loss_to_log = running_loss / num_tokens — but with the skip, running_loss accumulates mean loss (not sum), so dividing by num_tokens again gave mean/n. Fixed by using denominator 1.0 when skipping.

The existing test_loss parametrizes don't excercise the accum=1 non-inbwd skip path in all recipes, so this went unnoticed. I've added (8, 1, False) to test_full_finetune_single_device.py:test_loss to ensure it is covered.

[ErenAta16]

Checked the revert rather than taking it on trust, and it is complete. Both distributed recipes now carry only the explanatory comment, with no code change at all:

recipes/full_finetune_distributed.py   + comment only, 0 code lines
recipes/lora_finetune_distributed.py   + comment only, 0 code lines

so the loss multiply and scale_grads run exactly as they did before, and skip_loss_scale appears only in the single-device recipes where world_size == 1 collapses the two token counts into the same number.

The NOTE is the right artifact to leave behind. Stating that current_num_tokens is rank-local while num_tokens is all-reduced, and that skipping would normalise each rank by its own count, is the whole reason the optimisation cannot generalise. Without it someone reads the single-device version, sees an obvious win, and copies it across.

The numbers are the part I could not produce myself, having no multi-rank hardware, and they are more convincing than the argument was. A 25% gradient difference on a mildly imbalanced [5, 3] split is well inside what real instruction data produces, and 242% on [1, 64] shows it is not a rounding-scale effect. Those two figures belong in the PR description, since they answer the obvious "does this actually matter in practice" question that a reviewer would otherwise have to take on faith.

Finding three further single-device regressions while writing the verification tests is the useful kind of side effect, and it is an argument for the tests existing independently of this optimisation.

Scoping the change to single-device is the right resolution. The saving is real there and the invariant that makes it safe is checkable in one line, world_size == 1, rather than resting on an assumption about batch composition.

Nothing further from me.