[vnayakde]

Fixes #111

This PR fixes a loss-scaling issue when micro_acc_steps > 1, where the accumulated gradients and logged metrics (fvu, auxk, multi_topk) were scaled down by micro_acc_steps.

Background

In PR #138 (commit f19b274), micro-batching was restored by chunking activations and normalizing each chunk using the total_variance of the full batch.

Since all chunks use the full-batch variance, their losses already sum up to the loss of the full batch. However, in sparsify/trainer.py, acc_steps was still calculated as:

acc_steps = self.cfg.grad_acc_steps * self.cfg.micro_acc_steps

Each chunk was then divided by acc_steps before backward:

loss.div(acc_steps).backward()

This caused the loss to be divided by micro_acc_steps an extra time.

On the unpatched main, I observed the following gradient norms compared to an unchunked batch:

• micro_acc_steps=2: 0.50x
• micro_acc_steps=3: 0.33x
• micro_acc_steps=8: 0.125x

The logged FVU was also under-reported by a factor of micro_acc_steps.

The existing test_micro_acc_steps_matches_unchunked_update did not catch this because it compared parameter weights after an Adam step. Adam's first update is largely insensitive to gradient scaling.

Changes Made

1. sparsify/trainer.py: Changed acc_steps to divide only by grad_acc_steps.
2. tests/test_micro_acc_steps.py:
   • Added tests comparing raw parameter gradients before the optimizer step across even and uneven chunks, multi-step accumulation, multi_topk, and auxk.
   • Added a test to verify that logged FVU matches between chunked and unchunked runs.
   • Updated comments about small chunks diverging, since they now produce closely matching results with shared variance.
   • Added a device fallback ("cuda" if torch.cuda.is_available() else "cpu") so tests can run on non-CUDA environments without crashing.
3. Documentation: Updated config.py and README.md to clarify that micro_acc_steps only applies to loss_fn="fvu", since ce and kl require full-sequence reconstruction.

Testing

• On unpatched main, the new tests reproduced the issue: 7 failed, 4 passed.
• With this fix, all 11 tests pass.
• pre-commit and ruff check both pass.

Note on CPU / Triton:
In environments without an NVIDIA GPU, tests can be run on CPU using SPARSIFY_DISABLE_TRITON=1 ./venv/bin/pytest tests/test_micro_acc_steps.py -v. This leverages the repository's built-in eager decoder fallback (sparsify/utils.py) so CUDA/Triton drivers aren't required to verify correctness.
Closes #111

Related: #131, #138