[Fix] Reduce ChunkLoss memory usage with detached LM heads

[matrix72c]

Summary

ChunkLoss currently allocates and retains a full-sized LM-head gradient tensor even when the supplied head_weight is detached and does not require gradients. This happens in the MTP path when detach_mtp_lm_head_weight=True and can waste a substantial amount of accelerator memory because the tensor scales with vocab_size * hidden_size.

This change makes the custom autograd function respect head_weight.requires_grad throughout forward and backward:

• allocate the accumulated head-weight gradient only when the weight is trainable;
• avoid allocating a zero-filled head-gradient tensor for every chunk when it is detached;
• save only the tensors needed by backward in the detached case; and
• return None for the head-weight gradient when no gradient was requested.

The trainable-head path keeps its existing loss and gradient behavior.

Problem

Before this change, ChunkLoss.forward() unconditionally created grad_weight = torch.zeros_like(head_weight). Each chunk also produced or allocated a tensor with the same shape, and the accumulated tensor was saved for backward.

For a detached LM head this work cannot contribute to optimization: PyTorch does not request a gradient for that input, and the custom backward should return None. Nevertheless, the implementation still paid the memory cost of one or more vocabulary-sized weight-gradient buffers. Depending on model dimensions and dtype, this can consume hundreds of MiB or more.

The existing MTP loss intentionally exercises this case by detaching the LM-head weight when detach_mtp_lm_head_weight is enabled.

Implementation

• Cache whether head_weight requires gradients at the start of the custom forward.
• Initialize and accumulate grad_weight only for the trainable-head path.
• Save grad_inputs alone for the detached-head path and record the condition on the autograd context.
• In backward, scale and return the saved head gradient only when it exists; otherwise return None for head_weight.

This is deliberately local to ChunkLoss: it does not alter the chunking algorithm, numerical loss definition, MTP configuration, or public API.

Tests

Added focused regression tests covering both modes:

• the trainable-head path matches a non-chunked reference for loss, hidden-state gradients, and head-weight gradients;
• the detached-head path matches the reference for loss and hidden-state gradients, leaves head_weight.grad unset, and verifies that no full-sized zero head-gradient buffer is allocated.

Validation performed:

PYTHONPATH=. python -m pytest -q tests/loss/test_chunk_loss.py
2 passed

git diff --check upstream/main...fix/chunk-loss-detached-head-memory