The offloaded log-softmax Function recomputes its output in backward, then called
torch.autograd.backward(output, grad_output) and returned lm_head.grad. That writes the
leaf .grad and returns it as the op's gradient contribution, so the outer engine's
AccumulateGrad adds it again: lm_head gradients came out 2x too large for a single use,
and more when lm_head is shared across chunks (the normal GRPO path), and any grad already
accumulated from a previous micro-batch was corrupted. Switch to torch.autograd.grad, which
returns each input's local gradient without touching leaf .grad. hidden_states.grad was
already correct (its reloaded tensor is a private leaf) and is unchanged. Only reachable on
the offload path (long-completion / large-batch GRPO) with a trainable lm_head or tied
embeddings; frozen lm_head (standard LoRA) returned None and was unaffected.

Optimize the CPU-offloading log-softmax Function used by the large-batch or
long-completion GRPO logprob path, on top of the local-input-grads fix:

- Offload hidden_states through a pinned host buffer on a side CUDA stream so
  the D2H copy runs asynchronously and overlaps the no_grad forward; backward
  waits on the recorded copy event before the H2D reload, which is required
  for correctness with pinned non_blocking copies.
- Skip the CPU round-trip entirely when the GPU has generous headroom (free
  memory at least 4x the tensor size); the detached alias adds no allocation
  and the big logits are still checkpointed via recompute.
- Drop the dead post-reload dtype cast: the saved tensor already keeps the
  forward dtype.
- Fall back to the original pageable copy on CPU tensors, missing CUDA, or
  pinned allocation failure, keeping Windows/Mac/CPU behavior unchanged.

Outputs and gradients are bitwise identical to the previous implementation
for both frozen and trainable lm_head, and fwd+bwd time on
B=16,T=2048,H=2048,V=32000 bf16 with 256 chunks is equal or slightly better
in all measured cases. Add CPU tests pinning the fallback numerics and the
event-sync and is_cuda-guard source invariants.

The padded GRPO path calls Unsloth_Offloaded_Log_Softmax once per
zipped_inputs chunk before any backward runs. The per-call check
tensor_bytes * 4 <= free_bytes only sees current free memory, so with N
chunks each sized at the threshold the retained total approaches
free * (1 - (3/4)^N), i.e. nearly all initially free memory, starving
backward's logit recompute.

Track bytes kept this step in a function-local counter and require
4 * (tensor_bytes + retained) <= free_bytes, capping retention at 1/4 of
measured free memory. The counter is recreated on every
grpo_accumulated_loss call, so no state can leak across steps, eval
passes, or exceptions. Single-chunk behavior is unchanged.