diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 5540f058..340b353e 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -562,7 +562,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec raise ValueError('sampling replay does not support sequence or context parallelism') loss_require_values = getattr(loss_instance, 'require_values', False) assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' - inputs: Dict[str, Any] = processor( + inputs: Dict[str, Any] = processor.process_local( inputs, sp_strategy=self.sp_strategy, model=self.model, @@ -687,7 +687,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T if getattr(self, '_enable_sp', False) or cp_world_size > 1: raise ValueError('sampling replay does not support sequence or context parallelism') loss_require_values = getattr(loss_instance, 'require_values', False) - inputs: Dict[str, Any] = processor( + inputs: Dict[str, Any] = processor.process_local( inputs, sp_strategy=self.sp_strategy, model=self.model, diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 131ae81c..816c47d5 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -86,6 +86,17 @@ def __init__(self, @remote_function() def __call__(self, inputs: Union[InputFeature, List[InputFeature]], **kwargs) -> Union[InputFeature, List[InputFeature]]: + return self.process_local(inputs, **kwargs) + + def process_local(self, inputs: Union[InputFeature, List[InputFeature]], + **kwargs) -> Union[InputFeature, List[InputFeature]]: + """Process inputs on this worker without remote dispatch. + + Use this for worker-owned inputs such as the model passed by Transformers. + This method is not remotely exported; use ``__call__`` for remote calls. + """ + if hasattr(self, '_actors'): + raise ValueError('Local processing requires a worker-owned object, not a remote handle.') for pipe in self.process_pipeline: inputs = pipe(inputs, **kwargs) return inputs diff --git a/tests/model/test_safe_mode.py b/tests/model/test_safe_mode.py new file mode 100644 index 00000000..7fcb3ea0 --- /dev/null +++ b/tests/model/test_safe_mode.py @@ -0,0 +1,137 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""CPU training through production methods, processors and Ray-mode decorators.""" +from types import SimpleNamespace + +import pytest +import torch + +import twinkle.infra as infra +from twinkle.loss import CrossEntropyLoss +from twinkle.model.transformers.transformers import OptimizerGroup, TransformersModel +from twinkle.processor import InputProcessor +from twinkle.template import Template +from twinkle.utils import Platform + + +class _TinyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.embedding = torch.nn.Embedding(8, 8) + + def forward(self, input_ids, **kwargs): + return {'logits': self.embedding(input_ids)} + + +@pytest.fixture +def cpu_model(monkeypatch): + monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', '0') + monkeypatch.setattr(infra, '_mode', 'ray') + monkeypatch.setattr(Platform, 'get_local_device', lambda *args, **kwargs: torch.device('cpu')) + + def build(model_cls=TransformersModel): + # Construct worker-owned objects without Ray driver handles or tokenizer + # downloads. Keep the full processor pipeline, loss and backward methods. + with monkeypatch.context() as local: + local.setattr(infra, '_mode', 'local') + processor = InputProcessor(padding_free=False) + group = OptimizerGroup( + processor=processor, + template=object.__new__(Template), + loss_instance=CrossEntropyLoss(), + ) + model = object.__new__(model_cls) + torch.nn.Module.__init__(model) + model.model = _TinyModel() + model.optimizer_group = {'': group} + model.device_mesh = None + model.mixed_precision = 'no' + model.sp_strategy = None + model.hf_config = SimpleNamespace() + model._model_wrapped = True + model._router_replay_setup = lambda *args: lambda: None + model.strategy = SimpleNamespace(unwrap_model=lambda obj: obj) + return model, group + + return build + + +def _features(): + return [{'input_ids': [1, 2, 3, 4], 'labels': [2, 3, 4, 5]}] + + +def _assert_gradients(model): + grad = model.model.embedding.weight.grad + assert grad is not None and torch.isfinite(grad).all() + assert grad.norm().item() > 0 + + +@pytest.mark.parametrize('mode', ['local', 'ray']) +@pytest.mark.parametrize('trust', ['0', '1']) +@pytest.mark.parametrize('entry', ['forward_backward', 'micro_batch', 'forward_only']) +def test_training_entries(cpu_model, monkeypatch, mode, trust, entry): + monkeypatch.setattr(infra, '_mode', mode) + monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', trust) + model, _ = cpu_model() + if entry == 'forward_only': + model.forward_only(inputs=_features()) + loss = model.calculate_loss() + assert model.model.embedding.weight.grad is None + else: + kwargs = {'micro_batch_size': 1} if entry == 'micro_batch' else {} + inputs = _features() * 2 if kwargs else _features() + loss = model.forward_backward(inputs=inputs, **kwargs)['loss'] + _assert_gradients(model) + assert torch.isfinite(torch.tensor(loss)) + + +@pytest.mark.parametrize('trust', ['0', '1']) +def test_tinker_forward_backward(cpu_model, monkeypatch, trust): + import numpy as np + types = pytest.importorskip('tinker').types + from twinkle.server.model.backends.transformers_model import TwinkleCompatFullTransformersModel + + monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', trust) + model, _ = cpu_model(model_cls=TwinkleCompatFullTransformersModel) + datum = types.Datum( + model_input=types.ModelInput.from_ints([1, 2, 3, 4]), + loss_fn_inputs={ + 'target_tokens': types.TensorData.from_numpy(np.array([2, 3, 4, 5], dtype=np.int64)), + 'weights': types.TensorData.from_numpy(np.ones(4, dtype=np.float32)), + }, + ) + results, loss = model.tinker_forward_backward(inputs=[datum], adapter_name='', loss_fn='cross_entropy') + assert results and torch.isfinite(torch.tensor(loss)) + _assert_gradients(model) + + +@pytest.mark.parametrize('value', [lambda: None, type, {'nested': [lambda: None]}]) +@pytest.mark.parametrize('remote_handle', [False, True]) +def test_public_processor_rejects_callable(cpu_model, value, remote_handle): + _, group = cpu_model() + if remote_handle: + group.processor._actors = [] + with pytest.raises(ValueError, match='Callable or Type'): + group.processor(_features(), external=value) + + +def test_internal_processing_keeps_public_safety_checks(cpu_model): + _, group = cpu_model() + _, other = cpu_model() + + def pipeline(inputs, **kwargs): + other.processor(inputs, external=lambda: None) + + group.processor.process_pipeline = [pipeline] + with pytest.raises(ValueError, match='Callable or Type'): + group.processor.process_local(_features(), model=lambda: None) + + +def test_local_call_rejects_remote_handle(cpu_model): + _, group = cpu_model() + group.processor._actors = [] + with pytest.raises(ValueError, match='remote handle'): + group.processor.process_local(_features()) + + +def test_internal_processor_entry_is_not_exported(): + assert not hasattr(InputProcessor.process_local, '_execute')