Skip to content

fix: reject unsupported Unsloth RL gradient accumulation - #1137

Open
monsterlag wants to merge 1 commit into
OpenPipe:mainfrom
monsterlag:fix/issue-1121-0079
Open

monsterlag wants to merge 1 commit into
OpenPipe:mainfrom
monsterlag:fix/issue-1121-0079

Conversation

@monsterlag

Copy link
Copy Markdown

Note

AI-generated, not reviewed by a human. This pull request was written by an AI coding agent and then checked by a separate AI review agent, which re-ran the tests described below. No human has reviewed the code. It is submitted by @monsterlag, who runs these agents. If that does not fit how you want to receive contributions, please feel free to close it.

Summary

Closes #1121.

Unsloth's RL producer waits for an optimizer-step result after each input. With gradient_accumulation_steps > 1, the trainer needs another input before producing that result, so both sides wait indefinitely. Reject this unsupported configuration with an actionable error instead of hanging.

Changes

  • Share RL accumulation validation between the service resolver and public training entry point.
  • Reject unsupported values before reading packed tensors, starting background training, or mutating optimizer state.
  • Preserve default accumulation, existing invalid/mismatched configuration checks, and SFT behavior. This does not implement RL accumulation or change loss normalization.
  • Add regression coverage for service/runtime validation and the default queue flow, including warmup.

Test plan

  • Reproduced the stall using the real ART producer and asyncio queues with a CPU trainer stand-in that emits one result per optimizer step. Accumulation=2 stalled after one input with and without warmup; accumulation=1 controls completed. No GPU training was run.
  • New regression tests against the original source: 10 failed, 9 passed. All 19 pass with the fix.
  • Relevant unit tests using the project's manual-test options: 174 passed:
    uv run pytest --nbval --current-env --tb=short tests/unit/test_unsloth_gradient_accumulation.py tests/unit/test_unsloth_asyncio_isolation.py tests/unit/test_unsloth_metrics.py tests/unit/test_unsloth_autocast_dtype.py tests/unit/test_unsloth_logprobs.py tests/unit/test_loss_ratio.py tests/unit/test_loss_tis.py tests/unit/test_local_sft.py tests/unit/test_sft.py
  • uv run ty check src tests/unit/test_unsloth_gradient_accumulation.py: passed.
  • uv run prek run --all-files: Ruff, formatting, and lockfile checks passed; repository-wide ty reports 18 diagnostics, identical to the untouched baseline in the same environment.
  • uv run pytest --nbval --current-env --tb=short tests/unit: blocked by 15 collection errors from missing optional Tinker/Megatron/Triton dependencies; the untouched baseline has the same collection errors.

Fail fast instead of deadlocking the one-input/one-result producer when gradient_accumulation_steps exceeds one. Validate both service resolution and the public runtime entry point before optimizer state changes.

Closes OpenPipe#1121

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

LocalBackend/Unsloth: gradient_accumulation_steps > 1 hangs the trainer (producer waits for one result per sequence)

1 participant