From 2253de47552803cfd3749746e1d13d74a95fd050 Mon Sep 17 00:00:00 2001 From: Ryan McKenna Date: Wed, 9 Sep 2026 14:34:57 -0700 Subject: [PATCH] Support pygrain.MapDataset in dpsynth.text DPFineTuner. PiperOrigin-RevId: 978757926 --- dpsynth/text/dp_sft.py | 31 +++++-- dpsynth/text/model.py | 52 ++++++++---- tests/text/dp_sft_test.py | 171 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 232 insertions(+), 22 deletions(-) diff --git a/dpsynth/text/dp_sft.py b/dpsynth/text/dp_sft.py index 348fda2a..d3796ad8 100644 --- a/dpsynth/text/dp_sft.py +++ b/dpsynth/text/dp_sft.py @@ -47,8 +47,10 @@ from dpsynth.text import model from gemma import gm from gemma import peft +import grain.python as pygrain from jax_privacy import execution_plan from jax_privacy import training +import numpy as np import optax @@ -75,7 +77,7 @@ class DPFineTuner(api.DPMechanism): A ``DPMechanism`` that wraps ``DPTrainer`` with tokenization, model loading, and LoRA handling. All configuration is specified at construction time; ``__call__`` takes ``(rng, data)`` where ``data`` is a sequence of text - strings. + string pairs or a pygrain MapDataset. """ model_variant: model.GemmaModel @@ -134,23 +136,36 @@ def dp_event(self) -> dp_accounting.DpEvent: def __call__( self, rng: int, - data: Sequence[tuple[str, str]], + data: Sequence[tuple[str, str]] | pygrain.MapDataset, ) -> FineTuneResult: """Tokenizes text, loads the model, and runs DP-SGD fine-tuning. Args: rng: Random seed for batch selection and noise generation. - data: Sequence of ``(prompt, response)`` string pairs. + data: Sequence of ``(prompt, response)`` string pairs or a + ``pygrain.MapDataset``. Returns: A ``FineTuneResult`` with the trained model and merged LoRA parameters. Private training state (noise, optimizer) is not exposed. """ - dataset = model.tokenize_texts( - data, - model_variant=self.model_variant, - max_seq_length=self.max_seq_length, - ) + if isinstance(data, pygrain.MapDataset): + tokenizer = self.model_variant.tokenizer_class() + tokenized = data.map( + lambda ex: model.tokenize_example(ex, tokenizer, self.max_seq_length) + ) + # before DPTrainer. + elements = list(tokenized) + dataset = { + 'input_tokens': np.stack([x['input_tokens'] for x in elements]), + 'loss_mask': np.stack([x['loss_mask'] for x in elements]), + } + else: + dataset = model.tokenize_texts( + data, + model_variant=self.model_variant, + max_seq_length=self.max_seq_length, + ) lora_config = model.LoraConfig(rank=self.lora_rank) module, frozen_params, trainable_params = model.load_gemma( diff --git a/dpsynth/text/model.py b/dpsynth/text/model.py index 0f053654..1c3cc609 100644 --- a/dpsynth/text/model.py +++ b/dpsynth/text/model.py @@ -197,6 +197,40 @@ def format_response(response: str, tokenizer: Any) -> str: return f'{response}{eot}' +def tokenize_example( + example: tuple[str, str], + tokenizer: Any, + max_seq_length: int, +) -> dict[str, np.ndarray]: + """Tokenizes a single (prompt, response) pair for supervised fine-tuning. + + Prompt tokens are masked out (``loss_mask=0``) so only the response + contributes to the training loss. + + Args: + example: ``(prompt, response)`` string pair. + tokenizer: Gemma tokenizer instance. + max_seq_length: Maximum sequence length (including special tokens). + + Returns: + Dict with ``'input_tokens'`` and ``'loss_mask'`` (int32 ``[L]``). + """ + prompt, response = example + prompt_str = format_prompt(prompt, tokenizer) + response_str = format_response(response, tokenizer) + prompt_ids = tokenizer.encode(prompt_str, add_bos=True) + response_ids = tokenizer.encode(response_str, add_eos=True) + + ids = prompt_ids + response_ids + length = min(len(ids), max_seq_length) + tokens = np.zeros(max_seq_length, dtype=np.int32) + mask = np.zeros(max_seq_length, dtype=np.int32) + tokens[:length] = ids[:length] + mask[min(len(prompt_ids), length) : length] = 1 + + return {'input_tokens': tokens, 'loss_mask': mask} + + def tokenize_texts( examples: Sequence[tuple[str, str]], model_variant: GemmaModel, @@ -221,20 +255,10 @@ def tokenize_texts( tokens = np.zeros((len(examples), max_seq_length), dtype=np.int32) mask = np.zeros((len(examples), max_seq_length), dtype=np.int32) - for i, (prompt, response) in enumerate(examples): - # Embed turn tags as strings so SentencePiece handles tokenization - # boundaries correctly (encoding pieces separately can shift BPE merges). - prompt_str = format_prompt(prompt, tokenizer) - response_str = format_response(response, tokenizer) - prompt_ids = tokenizer.encode(prompt_str, add_bos=True) - response_ids = tokenizer.encode(response_str, add_eos=True) - - ids = prompt_ids + response_ids - length = min(len(ids), max_seq_length) - tokens[i, :length] = ids[:length] - # Mask: 0 for prompt, 1 for response. - resp_start = min(len(prompt_ids), length) - mask[i, resp_start:length] = 1 + for i, example in enumerate(examples): + tokenized = tokenize_example(example, tokenizer, max_seq_length) + tokens[i] = tokenized['input_tokens'] + mask[i] = tokenized['loss_mask'] logging.info( 'Tokenized %d examples (max_seq_length=%d)', diff --git a/tests/text/dp_sft_test.py b/tests/text/dp_sft_test.py index 6c7eceb6..9ed4d3a4 100644 --- a/tests/text/dp_sft_test.py +++ b/tests/text/dp_sft_test.py @@ -20,13 +20,18 @@ import dataclasses import math +from unittest import mock from absl.testing import absltest from dpsynth.text import dp_sft +from dpsynth.text import dp_trainer from dpsynth.text import model +from gemma import peft +import grain.python as pygrain import jax import jax.numpy as jnp from jax_privacy import execution_plan +import numpy as np def _default_config(): @@ -193,6 +198,172 @@ def test_dp_event_after_calibration(self): event = mechanism.dp_event self.assertIsNotNone(event) + @mock.patch.object(peft, 'merge_params', autospec=True) + @mock.patch.object(dp_trainer, 'DPTrainer', autospec=True) + @mock.patch.object(model, 'load_gemma', autospec=True) + def test_call_with_map_dataset_materializes_to_numpy( + self, mock_load, mock_trainer_cls, mock_merge + ): + mock_load.return_value = (mock.MagicMock(), {}, {}) + mock_trainer = mock_trainer_cls.return_value + mock_trainer.return_value = mock.MagicMock(params={}) + mock_merge.return_value = {} + + variant = model.GemmaModel( + model_class=mock.MagicMock(), + checkpoint_path='/mock/path', + tokenizer_class=_MockTokenizer, + ) + pairs = [('prompt1', 'resp1'), ('prompt2', 'resp2')] + ds = pygrain.MapDataset.source(pairs) + + fine_tuner = dp_sft.DPFineTuner( + model_variant=variant, + mechanism_config=_default_config(), + max_seq_length=16, + ).configure(zcdp_rho=0.5) + + res = fine_tuner(rng=0, data=ds) + self.assertIsInstance(res, dp_sft.FineTuneResult) + mock_trainer.assert_called_once() + passed_dataset = mock_trainer.call_args.kwargs['data'] + self.assertIsInstance(passed_dataset['input_tokens'], np.ndarray) + self.assertIsInstance(passed_dataset['loss_mask'], np.ndarray) + self.assertEqual(passed_dataset['input_tokens'].shape, (2, 16)) + self.assertEqual(passed_dataset['loss_mask'].shape, (2, 16)) + + expected_dataset = model.tokenize_texts(pairs, variant, max_seq_length=16) + np.testing.assert_array_equal( + passed_dataset['input_tokens'], expected_dataset['input_tokens'] + ) + np.testing.assert_array_equal( + passed_dataset['loss_mask'], expected_dataset['loss_mask'] + ) + + @mock.patch.object(peft, 'merge_params', autospec=True) + @mock.patch.object(dp_trainer, 'DPTrainer', autospec=True) + @mock.patch.object(model, 'load_gemma', autospec=True) + def test_call_with_sequence(self, mock_load, mock_trainer_cls, mock_merge): + mock_load.return_value = (mock.MagicMock(), {}, {}) + mock_trainer = mock_trainer_cls.return_value + mock_trainer.return_value = mock.MagicMock(params={}) + mock_merge.return_value = {} + + variant = model.GemmaModel( + model_class=mock.MagicMock(), + checkpoint_path='/mock/path', + tokenizer_class=_MockTokenizer, + ) + pairs = [('prompt1', 'resp1'), ('prompt2', 'resp2')] + fine_tuner = dp_sft.DPFineTuner( + model_variant=variant, + mechanism_config=_default_config(), + max_seq_length=16, + ).configure(zcdp_rho=0.5) + + res = fine_tuner(rng=0, data=pairs) + self.assertIsInstance(res, dp_sft.FineTuneResult) + mock_trainer.assert_called_once() + passed_dataset = mock_trainer.call_args.kwargs['data'] + expected_dataset = model.tokenize_texts(pairs, variant, max_seq_length=16) + np.testing.assert_array_equal( + passed_dataset['input_tokens'], expected_dataset['input_tokens'] + ) + np.testing.assert_array_equal( + passed_dataset['loss_mask'], expected_dataset['loss_mask'] + ) + + +class _MockSpecialTokens: + START_OF_TURN = 0 + END_OF_TURN = 1 + + +class _MockTokenizer: + + def __init__(self): + self.special_tokens = _MockSpecialTokens() + self.tokens = ['', ''] + + def encode(self, text, add_bos=False, add_eos=False): + del add_bos, add_eos + return [len(text) % 10 + 1, len(text) % 5 + 1] + + +class TokenizeExampleTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self.tokenizer = _MockTokenizer() + + def test_tokenize_example_shapes_and_types(self): + example = ('Hello', 'World') + result = model.tokenize_example(example, self.tokenizer, max_seq_length=16) + self.assertIn('input_tokens', result) + self.assertIn('loss_mask', result) + self.assertEqual(result['input_tokens'].shape, (16,)) + self.assertEqual(result['loss_mask'].shape, (16,)) + self.assertEqual(result['input_tokens'].dtype, np.int32) + self.assertEqual(result['loss_mask'].dtype, np.int32) + + def test_tokenize_example_masking(self): + example = ('Hi', 'There') + result = model.tokenize_example(example, self.tokenizer, max_seq_length=8) + np.testing.assert_array_equal(result['loss_mask'][:2], [0, 0]) + np.testing.assert_array_equal(result['loss_mask'][2:4], [1, 1]) + np.testing.assert_array_equal(result['loss_mask'][4:], [0, 0, 0, 0]) + + def test_tokenize_example_truncation(self): + example = ('Hi', 'There') + result = model.tokenize_example(example, self.tokenizer, max_seq_length=3) + self.assertEqual(result['input_tokens'].shape, (3,)) + self.assertEqual(result['loss_mask'].shape, (3,)) + np.testing.assert_array_equal(result['loss_mask'], [0, 0, 1]) + + +class TokenizeTextsTest(absltest.TestCase): + + def test_parity_with_tokenize_example(self): + variant = model.GemmaModel( + model_class=mock.MagicMock(), + checkpoint_path='/mock/path', + tokenizer_class=_MockTokenizer, + ) + tokenizer = variant.tokenizer_class() + pairs = [('Hello', 'World'), ('How are you?', 'I am fine.')] + seq_result = model.tokenize_texts( + pairs, model_variant=variant, max_seq_length=16 + ) + for i, pair in enumerate(pairs): + ex_result = model.tokenize_example(pair, tokenizer, max_seq_length=16) + np.testing.assert_array_equal( + seq_result['input_tokens'][i], ex_result['input_tokens'] + ) + np.testing.assert_array_equal( + seq_result['loss_mask'][i], ex_result['loss_mask'] + ) + + def test_parity_between_sequence_and_numpy_array(self): + variant = model.GemmaModel( + model_class=mock.MagicMock(), + checkpoint_path='/mock/path', + tokenizer_class=_MockTokenizer, + ) + pairs = [('Hello', 'World'), ('How are you?', 'I am fine.')] + seq_result = model.tokenize_texts( + pairs, model_variant=variant, max_seq_length=16 + ) + np_result = model.tokenize_texts( + np.asarray(pairs), model_variant=variant, max_seq_length=16 + ) + + np.testing.assert_array_equal( + seq_result['input_tokens'], np_result['input_tokens'] + ) + np.testing.assert_array_equal( + seq_result['loss_mask'], np_result['loss_mask'] + ) + if __name__ == '__main__': absltest.main()