perf(distillation): skip dW/dB calculation when parameters are frozen

[djl11]

Summary

LigerFusedLinearDistillationBase, which backs LigerFusedLinearJSDLoss and LigerFusedLinearCosineSimilarityLoss, always allocates a full grad_weight buffer and differentiates the student weight (and bias) on every chunk, even when they do not require grad. That is the usual case when distilling into LoRA/PEFT adapters with a frozen lm_head: the weight gradient is computed on every chunk and then discarded by autograd.

This applies to the distillation base the change #1441 made to the preference bases. The argnums passed to torch.func.grad_and_value are built from requires_grad, so a frozen weight or bias gets no gradient buffer and no per-chunk GEMM, and a call where nothing requires grad runs forward only. Trainable parameters take the same path as before.

Details

Forward + backward in bf16, LigerFusedLinearJSDLoss with its defaults (compiled, chunk_size=1024), B·T = 4096, H = 4096. Each figure is the median of 10 timed iterations after 3 warm-ups, then the median over separate processes, with main and this branch run alternately.

V	student head	main	this PR
128,256	frozen	42.3 ms, 6.51 GiB	32.0 ms, 5.53 GiB (−24%)
151,936	frozen	70.8 ms, 7.98 GiB	59.1 ms, 6.82 GiB (−17%)
128,256	trainable	42.1 ms	42.0 ms
151,936	trainable	71.4 ms	71.3 ms

Loss and input gradients are unchanged; the new tests compare frozen and trainable runs for both losses.

#1489 touches the same function (it skips the hard loss when it is neither weighted nor returned). The two are independent, but whichever lands second needs a small rebase.

Testing Done

• Hardware Type: NVIDIA H100 80GB HBM3 (torch 2.14.0+cu130, triton 3.8.0)
• run make test to ensure correctness: I ran the test files for the changed base rather than the full suite. test/chunked_loss/test_jsd_loss.py and test/chunked_loss/test_cosine_loss.py: 120 passed on this branch (104 on main, plus 16 new). They also pass on CPU.
• run make checkstyle to ensure code style (ruff 0.15.22 check and format --check on the changed files)
• run make test-convergence to ensure convergence: not run; the convergence tests do not exercise the chunked losses.

New tests, for both losses: a frozen vs a trainable weight and bias, compiled and eager; a trainable weight with a frozen bias; nothing requiring grad; and a check that a frozen weight is left out of argnums, so its gradient is skipped rather than computed and discarded.

[kolehma8]

create enums for the numbers so they are easier to follow

[kolehma8]

constants would also do, but enum is bit cleaner in my opinion

[kolehma8]

same here, lets use enums for the 0 etc.

[kolehma8]

assert chunk_grad_weight is not None?

[kolehma8]

Thank you @djl11 for your contribution. I had few stylistic comments around the argument order constants but otherwise the PR looks good. Once you address those I am happy to merge the PR.