diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index 630c8546..da49bce7 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/src/twinkle_agentic/async_rl/streaming_driver.py b/src/twinkle_agentic/async_rl/streaming_driver.py index 05abf9ed..f541cc05 100644 --- a/src/twinkle_agentic/async_rl/streaming_driver.py +++ b/src/twinkle_agentic/async_rl/streaming_driver.py @@ -797,6 +797,7 @@ 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) + # 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 current = self._current_version() @@ -1153,18 +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 zero - count reports no optimizer update. 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. """ - if optimizer_steps < 0: - raise ValueError('optimizer_steps must be non-negative') - self._committed_steps = int(optimizer_steps) + # 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/transformers/test_native_fsdp_offload.py b/tests/transformers/test_native_fsdp_offload.py new file mode 100644 index 00000000..797eee98 --- /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)) 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 00000000..0bc12b48 --- /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