From 38d8554c68b4dfe162be526abf5699577e5ea0e3 Mon Sep 17 00:00:00 2001 From: Yuanjia-Lee Date: Mon, 14 Sep 2026 12:06:12 +0800 Subject: [PATCH 1/2] perf(train): release consumed grads at the optimizer-step boundary optimizer_step() zeroed grads only on the skip path; committed updates left them alive across offload()/onload(), which move "params + grads + optimizer" as one train state, and across the next batch's anchor prep. Release them after the commit so dead tensors stop riding the rollout boundary. Accumulation windows are unaffected: non-stepping parts run zero_grad=False and the stepping update still zeroes at entry. --- unirl/train/backend/base_backend.py | 4 ++++ unirl/train/readme.md | 6 ++++++ 2 files changed, 10 insertions(+) diff --git a/unirl/train/backend/base_backend.py b/unirl/train/backend/base_backend.py index 936b504c1..878d1452c 100644 --- a/unirl/train/backend/base_backend.py +++ b/unirl/train/backend/base_backend.py @@ -308,6 +308,10 @@ def optimizer_step(self, *, max_grad_norm: float) -> float: if self.ema is not None: self.ema.step(self._optimizer_step_count) self._optimizer_step_count += 1 + # Grads are consumed. offload()/onload() move "params + grads + optimizer" + # as one train state, so releasing here keeps dead tensors off the bus + # across the rollout boundary and anchor prep. + self.optimizer.zero_grad(set_to_none=True) return grad_norm def on_rollout_end(self) -> None: diff --git a/unirl/train/readme.md b/unirl/train/readme.md index b311fbd3c..a08d3ca48 100644 --- a/unirl/train/readme.md +++ b/unirl/train/readme.md @@ -68,6 +68,12 @@ in `backend/base.py`; a multi-update-capable algorithm sets - **`optimizer_step` silently *skips* (does not crash) on a non-finite grad norm** and zeroes grads — a flat loss curve with a logged warning means grads went non-finite. +- **`optimizer_step` releases grads on both exits** — the skip path always did; + the committed-update path now does too. Nothing reads `.grad` between updates, + but `offload()`/`onload()` move "params + grads + optimizer" as one train + state, so consumed grads otherwise ride the rollout boundary and the next + batch's anchor prep. Accumulation windows are unaffected: non-stepping parts + run `zero_grad=False`, the stepping update still zeroes at entry. - **`master_dtype` defaults to `None`, so the optimizer master follows `param_dtype`** — a bf16-loaded base then keeps a bf16 LoRA master and the ~1e-6 AdamW steps round away (the policy drifts into a degenerate reward-hack). An fp32-loaded model gets an From 696764c38093944b5072c703c390de3186adb1d6 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Thu, 17 Sep 2026 09:19:29 +0800 Subject: [PATCH 2/2] chore(train): trim redundant grad-release commentary --- unirl/train/backend/base_backend.py | 4 +--- unirl/train/readme.md | 6 ------ 2 files changed, 1 insertion(+), 9 deletions(-) diff --git a/unirl/train/backend/base_backend.py b/unirl/train/backend/base_backend.py index 878d1452c..549b2e4a2 100644 --- a/unirl/train/backend/base_backend.py +++ b/unirl/train/backend/base_backend.py @@ -308,9 +308,7 @@ def optimizer_step(self, *, max_grad_norm: float) -> float: if self.ema is not None: self.ema.step(self._optimizer_step_count) self._optimizer_step_count += 1 - # Grads are consumed. offload()/onload() move "params + grads + optimizer" - # as one train state, so releasing here keeps dead tensors off the bus - # across the rollout boundary and anchor prep. + # Release consumed grads before train-state offload. self.optimizer.zero_grad(set_to_none=True) return grad_norm diff --git a/unirl/train/readme.md b/unirl/train/readme.md index a08d3ca48..b311fbd3c 100644 --- a/unirl/train/readme.md +++ b/unirl/train/readme.md @@ -68,12 +68,6 @@ in `backend/base.py`; a multi-update-capable algorithm sets - **`optimizer_step` silently *skips* (does not crash) on a non-finite grad norm** and zeroes grads — a flat loss curve with a logged warning means grads went non-finite. -- **`optimizer_step` releases grads on both exits** — the skip path always did; - the committed-update path now does too. Nothing reads `.grad` between updates, - but `offload()`/`onload()` move "params + grads + optimizer" as one train - state, so consumed grads otherwise ride the rollout boundary and the next - batch's anchor prep. Accumulation windows are unaffected: non-stepping parts - run `zero_grad=False`, the stepping update still zeroes at entry. - **`master_dtype` defaults to `None`, so the optimizer master follows `param_dtype`** — a bf16-loaded base then keeps a bf16 LoRA master and the ~1e-6 AdamW steps round away (the policy drifts into a degenerate reward-hack). An fp32-loaded model gets an