Fix shared loss gradient accumulation (#8245)
## Problem
Fixes #8224.
With gradient accumulation enabled, a loss can combine a forward pass
through the DeepSpeed engine with a forward pass through the model
inside the engine. The engine-output hook divides only the gradient from
the engine forward by the accumulation count. The gradient from the
forward pass on the model inside the engine remains unscaled, producing
an incorrect gradient for that parameter.
## Approach
Apply gradient-accumulation scaling to the complete loss passed through
managed `engine.backward`. During that call, mark the loss graph as
already scaled so engine-output hooks do not apply the scaling again,
and restore all managed-backward state on every exit.
If a managed backward is interrupted, preserve the existing ZeRO retry
behavior by not running a direct-backward epilogue over incomplete
reduction state. Direct tensor backward, pipeline output hooks, and
`scale_wrt_gas=False` otherwise retain their existing behavior.
## Testing
- `pytest -q tests/unit/runtime/zero/test_zero_shared_loss_gradient.py`
- Verified the existing ZeRO-3 exception-retry lifecycle behavior.
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
