From 29503ffe019b1ce441667a90ff4ef99c83fc6747 Mon Sep 17 00:00:00 2001 From: Qiyuan Chen Date: Thu, 29 Oct 2026 12:01:00 +0800 Subject: [PATCH 1/2] Keep an explicit zero optimizer step out of the driver budget. A reported 0 was treated as a missing count, so a batch with no optimizer update consumed one max_steps slot and disaggregated runs stopped early. Co-authored-by: Cursor --- .../async_rl/streaming_driver.py | 15 ++- .../test_streaming_driver_budget.py | 116 ++++++++++++++++++ 2 files changed, 125 insertions(+), 6 deletions(-) create mode 100644 tests/twinkle_agentic/test_streaming_driver_budget.py diff --git a/src/twinkle_agentic/async_rl/streaming_driver.py b/src/twinkle_agentic/async_rl/streaming_driver.py index 569e8f71e..608428d62 100644 --- a/src/twinkle_agentic/async_rl/streaming_driver.py +++ b/src/twinkle_agentic/async_rl/streaming_driver.py @@ -797,7 +797,8 @@ def _train_pull(self, records: List[BufferRecord], ready: List[BufferRecord], if group_counts is not None: group_counts.pop(record.group_key, None) steps = self._consume_batch(records) - steps = int(steps) if steps else 1 + # None means the consumer did not report a count. An explicit 0 is a real report. + steps = 1 if steps is None else int(steps) self._last_pull_steps = steps if steps > 0: current = self._current_version() @@ -1153,16 +1154,18 @@ def __next__(self) -> Any: raise StopIteration return batch - def commit(self, optimizer_steps: int = 1) -> None: + def commit(self, optimizer_steps: Optional[int] = 1) -> None: """Release the driver thread blocked in the bridge, reporting the real optimizer steps just run. The recipe calls this after training the batch ``__next__`` handed it. The count becomes the bridge's consume return, which drives the publish cadence (twinkle publishes after consume returns, - so the trained weights -- not stale ones -- are what the next admission window samples from). A falsy - count is treated as one step, matching :meth:`_train_pull`. The same count also feeds the internal - ``max_steps`` budget, so a driver constructed with a numeric budget exits on its own. + so the trained weights -- not stale ones -- are what the next admission window samples from). Omitting + the count, or passing ``None``, means one step. An explicit ``0`` means this batch did not update the + optimizer and does not consume the ``max_steps`` budget. The same count also feeds that budget, so a + driver constructed with a numeric budget exits on its own. """ - self._committed_steps = int(optimizer_steps) if optimizer_steps else 1 + # None means the caller omitted the count. An explicit 0 does not consume budget. + self._committed_steps = 1 if optimizer_steps is None else int(optimizer_steps) self._budget_steps += self._committed_steps self._ack_event.set() diff --git a/tests/twinkle_agentic/test_streaming_driver_budget.py b/tests/twinkle_agentic/test_streaming_driver_budget.py new file mode 100644 index 000000000..0bc12b487 --- /dev/null +++ b/tests/twinkle_agentic/test_streaming_driver_budget.py @@ -0,0 +1,116 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""A reported optimizer step of 0 must not be rewritten into a budget step.""" + +import types + +from twinkle_agentic.async_rl.streaming_driver import BufferRecord, StreamingDriver + + +def _driver(**overrides): + kwargs = dict( + prompt_stream=[], + max_staleness=0, + num_generations=1, + groups_per_partition=1, + buffer_depth=1, + submit=lambda prompt, index, policy: None, + poll=lambda handles: [], + collect=lambda handle: None, + assembly_ready=lambda ready: None, + consume=lambda records: 0, + max_steps=3, + ) + kwargs.update(overrides) + return StreamingDriver(**kwargs) + + +def test_commit_zero_does_not_exhaust_a_three_step_budget(): + driver = _driver() + for reported in (0, 1, 1): + driver.commit(reported) + assert driver._committed_steps == 1 + assert driver._budget_steps == 2 + assert driver._default_reached_max() is False + driver.commit(1) + assert driver._budget_steps == 3 + assert driver._default_reached_max() is True + + +def test_omitted_commit_still_counts_as_one_step(): + driver = _driver() + driver.commit() + driver.commit(None) + assert driver._budget_steps == 2 + assert driver._committed_steps == 1 + + +def test_train_pull_keeps_explicit_zero_out_of_the_stats(): + driver = _driver() + record = BufferRecord(sample=object(), version=0, group_key='g') + ready = [record] + driver._consume = lambda records: 0 + driver._train_pull([record], ready, {'g': 1}) + assert ready == [] + assert driver._last_pull_steps == 0 + assert driver.stats.optimizer_steps == 0 + assert driver.stats.train_pulls == 0 + assert driver.stats.consumed == 0 + + +def test_train_pull_missing_report_defaults_to_one_step(): + driver = _driver() + record = BufferRecord(sample=object(), version=0, group_key='g') + ready = [record] + driver._consume = lambda records: None + driver._train_pull([record], ready, {'g': 1}) + assert ready == [] + assert driver._last_pull_steps == 1 + assert driver.stats.optimizer_steps == 1 + assert driver.stats.train_pulls == 1 + assert driver.stats.consumed == 1 + + +def test_iterator_runs_until_three_real_updates(): + """The disaggregated loop reports [0, 1, 1, 1]; only the real updates fill max_steps=3.""" + active = {} + next_handle = 0 + + def submit(prompt, index, policy): + nonlocal next_handle + next_handle += 1 + active[next_handle] = policy.version + return next_handle + + def collect(handle): + active.pop(handle) + return types.SimpleNamespace() + + driver = StreamingDriver( + prompt_stream=list(range(32)), + max_staleness=0, + num_generations=1, + groups_per_partition=1, + buffer_depth=4, + submit=submit, + poll=lambda handles: list(handles), + collect=collect, + group_pull=True, + need_rows=4, + version_control=True, + max_steps=3, + consume=None, + ) + reported = iter((0, 1, 1, 1)) + seen = [] + try: + for batch in driver: + assert len(batch) == 4 + steps = next(reported) + seen.append(steps) + driver.commit(steps) + finally: + driver.abort() + assert seen == [0, 1, 1, 1] + assert driver._budget_steps == 3 + assert driver.stats.optimizer_steps == 3 + assert driver.stats.consumed == 12 From dd14ad1ce541bf134df5701bce5b714a08534340 Mon Sep 17 00:00:00 2001 From: Qiyuan Chen Date: Fri, 30 Oct 2026 13:52:26 +0800 Subject: [PATCH 2/2] Refresh the Native FSDP restore device at the start of each offload cycle. Lazy wrapping can move an initially CPU model onto CUDA between rounds, so reusing the first recorded device left non-persistent buffers on CPU after reload. Co-authored-by: Cursor --- .../transformers/strategy/native_fsdp.py | 17 ++- .../transformers/test_native_fsdp_offload.py | 138 ++++++++++++++++++ 2 files changed, 148 insertions(+), 7 deletions(-) create mode 100644 tests/transformers/test_native_fsdp_offload.py diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index 630c8546d..da49bce79 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/src/twinkle/model/transformers/strategy/native_fsdp.py @@ -344,23 +344,26 @@ def offload_to_cpu(self, model, optimizer=None, *, offload_model: bool = True) - """ import torch from twinkle.model.offload import ( - move_grads, move_optimizer_state, parameter_device, record_offload, release_caching_allocator) + move_grads, move_optimizer_state, offload_record, parameter_device, record_offload, + release_caching_allocator) module = self.unwrap_model(model) if not offload_model and optimizer is None: - record_offload(self, model=False, optimizer=False) release_caching_allocator() return - if getattr(self, '_offload_device', None) is None: + moved_model, moved_optimizer = offload_record(self) + # Lazy FSDP wrapping may place an initially CPU model on CUDA between rounds. + # Refresh the target for each new round, but retain it for repeated offloads. + if not moved_model and not moved_optimizer: self._offload_device = parameter_device(module) - if offload_model: + if offload_model and not moved_model: # Blocking copy: FSDP's ``_apply`` may read the CPU tensors before a queued # non_blocking transfer has finished. module.to('cpu', non_blocking=False) move_grads(module, torch.device('cpu')) - if optimizer is not None: + if optimizer is not None and not moved_optimizer: move_optimizer_state(optimizer, torch.device('cpu')) - record_offload(self, model=offload_model, optimizer=optimizer is not None) + record_offload(self, model=moved_model or offload_model, optimizer=moved_optimizer or optimizer is not None) if torch.cuda.is_available(): torch.cuda.empty_cache() @@ -382,9 +385,9 @@ def reload_to_gpu(self, model, optimizer=None) -> None: move_grads(module, device) if moved_optimizer and optimizer is not None: move_optimizer_state(optimizer, device) - clear_offload_record(self) if torch.cuda.is_available(): torch.cuda.synchronize() + clear_offload_record(self) def _ep_gather_context(self): """Return (ep_group, ep_world_size) for the EP submesh, or (None, 1) when EP is off.""" diff --git a/tests/transformers/test_native_fsdp_offload.py b/tests/transformers/test_native_fsdp_offload.py new file mode 100644 index 000000000..797eee98e --- /dev/null +++ b/tests/transformers/test_native_fsdp_offload.py @@ -0,0 +1,138 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""CPU placement probes for the lazy-wrap / colocated-offload boundary.""" +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from twinkle.model import offload +from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy + + +class PlacementModule: + """Track placement without allocating CUDA tensors or claiming to run FSDP.""" + + def __init__(self, device='cpu'): + self.param = SimpleNamespace(device=torch.device(device), grad=None) + self.inv_freq = SimpleNamespace(device=torch.device(device)) + self.moves = [] + + def parameters(self): + yield self.param + + def to(self, device, **kwargs): + self.param.device = self.inv_freq.device = torch.device(device) + self.moves.append(self.param.device) + return self + + +@pytest.fixture +def strategy(monkeypatch): + monkeypatch.setattr(torch.cuda, 'is_available', lambda: False) + monkeypatch.setattr(offload, 'release_caching_allocator', lambda: None) + return NativeFSDPStrategy(enable_ep=False) + + +@pytest.mark.parametrize('offload_model,with_optimizer', [(True, False), (False, True), (True, True)]) +def test_reload_uses_each_cycles_device_after_lazy_wrap(strategy, monkeypatch, offload_model, with_optimizer): + module = PlacementModule() + optimizer = object() if with_optimizer else None + move_state = Mock() + monkeypatch.setattr(offload, 'move_optimizer_state', move_state) + + # The initial CPU placement is legitimate; wrapping has not happened yet. + strategy.offload_to_cpu(module, optimizer, offload_model=offload_model) + strategy.reload_to_gpu(module, optimizer) + assert module.param.device == module.inv_freq.device == torch.device('cpu') + + # Simulate lazy wrapping, then two further cycles (including a changed rank device). + for device in (torch.device('cuda:0'), torch.device('cuda:1')): + module.to(device) + strategy.offload_to_cpu(module, optimizer, offload_model=offload_model) + expected = torch.device('cpu') if offload_model else device + assert module.param.device == module.inv_freq.device == expected + strategy.reload_to_gpu(module, optimizer) + assert module.param.device == module.inv_freq.device == device + if with_optimizer: + assert move_state.call_args.args == (optimizer, device) + assert offload.offload_record(strategy) == (False, False) + + +@pytest.mark.parametrize('first_model,first_optimizer', [(True, False), (False, True), (True, True)]) +def test_repeated_offload_keeps_restore_target_and_pending_pieces(strategy, monkeypatch, first_model, first_optimizer): + module = PlacementModule('cuda:1') + optimizer = object() + move_state = Mock() + monkeypatch.setattr(offload, 'move_optimizer_state', move_state) + + strategy.offload_to_cpu(module, optimizer if first_optimizer else None, offload_model=first_model) + strategy.offload_to_cpu(module, optimizer) + # A narrower repeated request must not forget the optimizer or capture CPU as the target. + strategy.offload_to_cpu(module) + strategy.offload_to_cpu(module, offload_model=False) + strategy.reload_to_gpu(module, optimizer) + assert module.param.device == module.inv_freq.device == torch.device('cuda:1') + assert move_state.call_args.args == (optimizer, torch.device('cuda:1')) + assert offload.offload_record(strategy) == (False, False) + + moves = list(module.moves) + calls = move_state.call_count + strategy.reload_to_gpu(module, optimizer) + assert module.moves == moves + assert move_state.call_count == calls + + +def test_failed_reload_can_be_retried(strategy, monkeypatch): + module = PlacementModule('cuda:0') + strategy.offload_to_cpu(module) + original_to = module.to + monkeypatch.setattr(module, 'to', Mock(side_effect=RuntimeError('copy failed'))) + with pytest.raises(RuntimeError, match='copy failed'): + strategy.reload_to_gpu(module) + assert offload.offload_record(strategy) == (True, False) + monkeypatch.setattr(module, 'to', original_to) + strategy.reload_to_gpu(module) + assert module.param.device == module.inv_freq.device == torch.device('cuda:0') + + +def test_failed_synchronize_keeps_pending_reload(strategy, monkeypatch): + module = PlacementModule('cuda:0') + monkeypatch.setattr(torch.cuda, 'is_available', lambda: True) + monkeypatch.setattr(torch.cuda, 'empty_cache', lambda: None) + synchronize = Mock(side_effect=[RuntimeError('synchronize failed'), None]) + monkeypatch.setattr(torch.cuda, 'synchronize', synchronize) + strategy.offload_to_cpu(module) + with pytest.raises(RuntimeError, match='synchronize failed'): + strategy.reload_to_gpu(module) + assert offload.offload_record(strategy) == (True, False) + strategy.reload_to_gpu(module) + assert offload.offload_record(strategy) == (False, False) + + +def test_cpu_roundtrip_preserves_real_nonpersistent_buffer_gradients_and_adam(strategy): + module = torch.nn.Linear(2, 2) + module.register_buffer('inv_freq', torch.ones(2), persistent=False) + optimizer = torch.optim.AdamW(module.parameters(), foreach=False) + module(torch.ones(1, 2)).sum().backward() + optimizer.step() + before_params = [param.detach().clone() for param in module.parameters()] + before_grads = [param.grad.clone() for param in module.parameters()] + before_buffer = module.inv_freq.clone() + before_state = {param: {name: value.clone() for name, value in state.items()} + for param, state in optimizer.state.items()} + + strategy.offload_to_cpu(module, optimizer) + strategy.reload_to_gpu(module, optimizer) + + assert 'inv_freq' not in module.state_dict() + torch.testing.assert_close(module.inv_freq, before_buffer) + for param, before, grad in zip(module.parameters(), before_params, before_grads): + torch.testing.assert_close(param, before) + torch.testing.assert_close(param.grad, grad) + for name, value in optimizer.state[param].items(): + torch.testing.assert_close(value, before_state[param][name]) + optimizer.zero_grad() + module(torch.ones(1, 2)).sum().backward() + optimizer.step() + assert any(not torch.equal(param, before) for param, before in zip(module.parameters(), before_params))