perf(train): release consumed grads at the optimizer-step boundary

[Yuanjia-Lee]

Summary

optimizer_step() zeroed grads only on the skip path; committed updates left
them alive. The colocate loops then carry them through the whole cycle:
offload()/onload() move params, grads, and optimizer state together, and
the next batch's anchor prep runs with the previous update's grads on GPU.
Nothing reads .grad between updates: the only two callers (stack/base.py,
unified_model_stack.py) go straight to metric aggregation, and accumulation
never spans an optimizer step.

A/B over a full FSDP2 train/offload/rollout/onload/anchor cycle, arms
differing only in the release line: parameters, AdamW state, and both
updates' losses are bit-identical. The baseline arm alone still moves the
previous update's grads through offload and onload; with the release,
onload's allocated delta drops from params+grads to params only. On the 4B
full-FT AR recipe the sharded fp32 grads reach 2 GB/rank, crossing PCIe
twice per rollout.

The release is the skip path's own zero_grad(set_to_none=True), placed
after the EMA/count consumers. Accumulation windows are unaffected:
non-stepping parts run zero_grad=False, the stepping update zeroes at
entry. Recorded in unirl/train/readme.md Gotchas.

Related Issue

Related to #429: move count there, payload size here.

Test Plan

• SKIP=no-commit-to-branch pre-commit run --files unirl/train/backend/base_backend.py unirl/train/readme.md --show-diff-on-failure
  → all hooks pass.
• Distributed FSDP2 harness on one 4090 node (two ranks): backward → real
  optimizer_step → real fsdp_offload → real fsdp_onload → no-grad
  anchor forward → next update. Param shards and AdamW state match
  bit-for-bit (max diff 0.0), losses equal.

Compatibility / Risk

• Grads free earlier; the between-updates footprint drops by the local grad
  shard. Clipping and deferred sync are unchanged: grads are fully
  reduce-scattered before optimizer.step().

Reviewer Notes

• Duplicate check: no open PR changes the optimizer-step boundary (feat(bagel): add native vLLM-Omni T2TI UniGRPO rollout #218 adds
  an opt-in parking method for external-engine rollouts; refactor(diffusion): make role residency one choice per role, not per phase #428 reorders
  trainer offload events).

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated tests, docs, and configs where needed, or explained why not.

[CjhHa1]

LGTM

[CjhHa1]

@Yuanjia-Lee I've shorten the docstrings as the method itself is worth adding