From a404526b411d0d0ef4f60d942a1584740e31e405 Mon Sep 17 00:00:00 2001 From: monsterlag <64719697+monsterlag@users.noreply.github.com> Date: Wed, 21 Oct 2026 23:19:37 +0000 Subject: [PATCH] fix: reject unsupported Unsloth RL gradient accumulation 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 #1121 --- src/art/unsloth/service.py | 4 +- src/art/unsloth/train.py | 12 ++ .../test_unsloth_gradient_accumulation.py | 164 ++++++++++++++++++ 3 files changed, 178 insertions(+), 2 deletions(-) create mode 100644 tests/unit/test_unsloth_gradient_accumulation.py diff --git a/src/art/unsloth/service.py b/src/art/unsloth/service.py index 2cfd45f7b..b5790bebe 100644 --- a/src/art/unsloth/service.py +++ b/src/art/unsloth/service.py @@ -35,6 +35,7 @@ ) from .train import ( UnslothTrainContext, + _validate_rl_gradient_accumulation, create_unsloth_train_context, gc_and_empty_cuda_cache, run_unsloth_rl_training, @@ -441,8 +442,7 @@ async def resolve_global_grad_accumulation_sequences( configured = int( self.config.get("trainer_args", {}).get("gradient_accumulation_steps", 1) ) - if configured < 1: - raise ValueError("Unsloth gradient accumulation must be >= 1") + _validate_rl_gradient_accumulation(configured) requested = config.grad_accumulation_sequences if requested is not None and requested != configured: raise ValueError( diff --git a/src/art/unsloth/train.py b/src/art/unsloth/train.py index b196fa511..0f51df2e2 100644 --- a/src/art/unsloth/train.py +++ b/src/art/unsloth/train.py @@ -794,6 +794,17 @@ def _precalculate_new_logprobs( ).to("cpu") +def _validate_rl_gradient_accumulation(steps: int) -> None: + if steps < 1: + raise ValueError("Unsloth gradient accumulation must be >= 1") + # The RL producer waits for an optimizer-step result after each input. + if steps != 1: + raise ValueError( + "Unsloth RL training does not support gradient accumulation; " + f"set gradient_accumulation_steps=1 (got {steps})" + ) + + async def run_unsloth_rl_training( ctx: UnslothTrainContext, disk_packed_tensors: DiskPackedTensors, @@ -801,6 +812,7 @@ async def run_unsloth_rl_training( _config: dev.TrainConfig, verbose: bool = False, ) -> AsyncIterator[dict[str, float]]: + _validate_rl_gradient_accumulation(ctx.trainer.args.gradient_accumulation_steps) _reset_optimizer_if_mode_changed(ctx, "rl") optimizer = _get_trainer_optimizer(ctx) for param_group in optimizer.param_groups: diff --git a/tests/unit/test_unsloth_gradient_accumulation.py b/tests/unit/test_unsloth_gradient_accumulation.py new file mode 100644 index 000000000..397826baa --- /dev/null +++ b/tests/unit/test_unsloth_gradient_accumulation.py @@ -0,0 +1,164 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from art import dev +from art.preprocessing.pack import DiskPackedTensors +from art.types import TrainConfig +from art.unsloth.service import UnslothService +import art.unsloth.train as train + +_DISK_PACKED_TENSORS = DiskPackedTensors( + dir="unused", num_sequences=2, sequence_length=1 +) + + +@pytest.mark.parametrize("configured", [2, 4]) +@pytest.mark.parametrize("requested", [None, 2]) +async def test_service_rejects_unsupported_rl_accumulation( + tmp_path, configured, requested +): + service = UnslothService( + "test", + "test-model", + {"trainer_args": {"gradient_accumulation_steps": configured}}, + str(tmp_path), + ) + + with pytest.raises(ValueError, match="set gradient_accumulation_steps=1"): + await service.resolve_global_grad_accumulation_sequences( + TrainConfig(grad_accumulation_sequences=requested) + ) + assert "_state" not in service.__dict__ + + +@pytest.mark.parametrize("configured", [None, 1]) +@pytest.mark.parametrize("requested", [None, 1]) +async def test_service_accepts_default_rl_accumulation(tmp_path, configured, requested): + config: dev.InternalModelConfig = {} + if configured is not None: + config["trainer_args"] = {"gradient_accumulation_steps": configured} + service = UnslothService("test", "test-model", config, str(tmp_path)) + + assert ( + await service.resolve_global_grad_accumulation_sequences( + TrainConfig(grad_accumulation_sequences=requested) + ) + == 1 + ) + + +@pytest.mark.parametrize("configured", [-1, 0]) +async def test_service_rejects_invalid_accumulation(tmp_path, configured): + service = UnslothService( + "test", + "test-model", + {"trainer_args": {"gradient_accumulation_steps": configured}}, + str(tmp_path), + ) + + with pytest.raises(ValueError, match="must be >= 1"): + await service.resolve_global_grad_accumulation_sequences(TrainConfig()) + + +async def test_service_rejects_mismatched_accumulation(tmp_path): + service = UnslothService("test", "test-model", {}, str(tmp_path)) + + with pytest.raises(ValueError, match="grad_accumulation_sequences=1, got 2"): + await service.resolve_global_grad_accumulation_sequences( + TrainConfig(grad_accumulation_sequences=2) + ) + + +def _context(steps, warmup): + trainer = Mock() + trainer.args.gradient_accumulation_steps = steps + trainer.optimizer = SimpleNamespace( + state={"sentinel": {}}, param_groups=[{"weight_decay": 0.0}] + ) + return train.UnslothTrainContext( + model=Mock(), + tokenizer=Mock(), + peft_model=Mock(), + trainer=trainer, + inputs_queue=asyncio.Queue(), + results_queue=asyncio.Queue(), + warmup_pending=warmup, + last_training_mode="sft", + ) + + +@pytest.mark.parametrize("steps", [0, 2, 4]) +@pytest.mark.parametrize("warmup", [False, True]) +async def test_runtime_rejects_accumulation_before_mutating_context( + monkeypatch, steps, warmup +): + ctx = _context(steps, warmup) + load_tensors = Mock(side_effect=AssertionError("packed tensors must not be read")) + start_training = Mock(side_effect=AssertionError("trainer must not start")) + monkeypatch.setattr(train, "packed_tensors_from_dir", load_tensors) + monkeypatch.setattr(train, "train", start_training) + message = "must be >= 1" if steps < 1 else "set gradient_accumulation_steps=1" + + with pytest.raises(ValueError, match=message): + await anext( + train.run_unsloth_rl_training(ctx, _DISK_PACKED_TENSORS, TrainConfig(), {}) + ) + + load_tensors.assert_not_called() + start_training.assert_not_called() + assert ctx.train_task is None + assert ctx.inputs_queue.empty() + assert ctx.results_queue.empty() + assert ctx.warmup_pending is warmup + assert ctx.last_training_mode == "sft" + assert ctx.trainer.optimizer.state == {"sentinel": {}} + assert ctx.trainer.optimizer.param_groups == [{"weight_decay": 0.0}] + + +@pytest.mark.parametrize("warmup", [False, True]) +async def test_runtime_default_accumulation_still_produces_results(monkeypatch, warmup): + ctx = _context(1, warmup) + consumed = [] + + async def trainer_loop(*, trainer, results_queue): + while True: + for _ in range(trainer.args.gradient_accumulation_steps): + consumed.append(await ctx.inputs_queue.get()) + results_queue.put_nowait({"loss": 1.0}) + + monkeypatch.setattr(train, "train", trainer_loop) + monkeypatch.setattr( + train, + "packed_tensors_from_dir", + lambda **kwargs: {"tokens": torch.zeros((2, 1))}, + ) + monkeypatch.setattr( + train, + "create_train_inputs", + lambda packed, offset, config, dev_config, warmup: (offset, warmup), + ) + monkeypatch.setattr(train, "gc_and_empty_cuda_cache", lambda: None) + + async def collect(): + return [ + result + async for result in train.run_unsloth_rl_training( + ctx, _DISK_PACKED_TENSORS, TrainConfig(), {} + ) + ] + + try: + results = await asyncio.wait_for(collect(), timeout=5) + assert results == [{"loss": 1.0}, {"loss": 1.0}] + assert consumed == ([(0, True)] if warmup else []) + [(0, False), (1, False)] + assert ctx.warmup_pending is False + assert ctx.inputs_queue.empty() + await asyncio.wait_for(ctx.results_queue.join(), timeout=5) + finally: + if ctx.train_task is not None: + ctx.train_task.cancel() + await asyncio.gather(ctx.train_task, return_exceptions=True)