From 14a70617709cc8b7db5337525d24fe4ea88e4cfe Mon Sep 17 00:00:00 2001 From: Ryan McKenna Date: Thu, 2 Jul 2026 11:11:04 -0700 Subject: [PATCH] Add DP supervised fine-tuning for Gemma (dpsynth.text): DPFineTuner returns a FineTuneResult with the trained model and a ChatSampler factory. Includes an end-to-end bioRxiv abstract example. Install the LLM extras with `pip install dpsynth[text]`. PiperOrigin-RevId: 941766080 --- dpsynth/examples/finetune_biorxiv.py | 223 +++++++++++++++++++++++++++ dpsynth/text/dp_sft.py | 190 +++++++++++++++++++++++ dpsynth/text/dp_trainer.py | 160 +++++++++++++++++++ dpsynth/text/model.py | 215 ++++++++++++++++++++++++++ tests/text/dp_sft_test.py | 198 ++++++++++++++++++++++++ tests/text/dp_trainer_test.py | 114 ++++++++++++++ 6 files changed, 1100 insertions(+) create mode 100644 dpsynth/examples/finetune_biorxiv.py create mode 100644 dpsynth/text/dp_sft.py create mode 100644 dpsynth/text/dp_trainer.py create mode 100644 dpsynth/text/model.py create mode 100644 tests/text/dp_sft_test.py create mode 100644 tests/text/dp_trainer_test.py diff --git a/dpsynth/examples/finetune_biorxiv.py b/dpsynth/examples/finetune_biorxiv.py new file mode 100644 index 0000000..b0ff56f --- /dev/null +++ b/dpsynth/examples/finetune_biorxiv.py @@ -0,0 +1,223 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Differentially private fine-tuning of Gemma on bioRxiv abstracts. + +This example demonstrates the end-to-end ``DPFineTuner`` workflow for generating +differentially private synthetic text: + + 1. Download real bioRxiv paper abstracts from the public bioRxiv API. + 2. Frame each abstract as a supervised fine-tuning example: a fixed + instruction prompt -> the abstract text. + 3. Differentially privately fine-tune a small Gemma model with DP-SGD. + 4. Generate *synthetic* abstracts by sampling from the fine-tuned model. + +The end result is a generator of biology-paper abstracts that reflects the +aggregate style and content of the training corpus while satisfying a formal +(epsilon, delta)-DP guarantee with respect to the individual papers. + +Because every training example shares the same fixed instruction prompt, the +model learns the *unconditional* distribution of abstracts: at generation time +we feed that same prompt and sample multiple times (with a non-zero +temperature) to draw diverse synthetic abstracts. + +This is a long-running job (model loading + DP-SGD + sampling on an +accelerator), and hence is provided as a binary rather than a colab notebook. + +Example (single GPU/TPU host): + + blaze run -c opt \ + //third_party/py/dpsynth/colabs:finetune_biorxiv -- \ + --alsologtostderr \ + --num_abstracts=256 \ + --iterations=200 \ + --epsilon=8.0 +""" + +from collections.abc import Sequence +import json +import time +import urllib.request + +from absl import app +from absl import flags +from absl import logging +from dpsynth.text import dp_sft +from dpsynth.text import model +from gemma import gm +import jax +from jax_privacy import execution_plan +import optax + +_MODEL = flags.DEFINE_enum( + 'model', + 'gemma3_270m_it', + [ + 'gemma3_270m_it', + 'gemma3_1b_it', + 'gemma3_4b_it', + 'gemma4_e2b_it', + 'gemma4_e4b_it', + ], + 'Gemma model variant to fine-tune.', +) +_NUM_ABSTRACTS = flags.DEFINE_integer( + 'num_abstracts', 256, 'Number of bioRxiv abstracts to train on.' +) +_START_DATE = flags.DEFINE_string( + 'start_date', '2024-01-01', 'Start date for the bioRxiv query (YYYY-MM-DD).' +) +_END_DATE = flags.DEFINE_string( + 'end_date', '2024-03-31', 'End date for the bioRxiv query (YYYY-MM-DD).' +) +_ITERATIONS = flags.DEFINE_integer( + 'iterations', 200, 'Number of DP-SGD training iterations.' +) +_EPSILON = flags.DEFINE_float( + 'epsilon', 8.0, 'Target (epsilon, delta)-DP epsilon.' +) +_DELTA = flags.DEFINE_float('delta', 1e-5, 'Target (epsilon, delta)-DP delta.') +_LORA_RANK = flags.DEFINE_integer('lora_rank', 16, 'LoRA rank.') +_MAX_SEQ_LENGTH = flags.DEFINE_integer( + 'max_seq_length', 512, 'Max token sequence length per example.' +) +_LEARNING_RATE = flags.DEFINE_float( + 'learning_rate', 1e-4, 'AdamW learning rate.' +) +_NUM_SAMPLES = flags.DEFINE_integer( + 'num_samples', 8, 'Number of synthetic abstracts to generate at the end.' +) +_MAX_OUT_LENGTH = flags.DEFINE_integer( + 'max_out_length', 512, 'Max tokens to generate per synthetic abstract.' +) +_TEMPERATURE = flags.DEFINE_float( + 'temperature', + 1.0, + 'Sampling temperature for synthetic abstract generation.', +) + +# A fixed instruction prompt prepended to every training example. Because it is +# identical for all examples, the model learns to produce an abstract whenever +# it sees this prompt -- i.e. an unconditional abstract generator. +_INSTRUCTION = 'Write the abstract of a biology research paper.' + +_API_BASE = 'https://api.biorxiv.org/details/biorxiv' + + +def fetch_abstracts( + start_date: str, end_date: str, max_records: int +) -> list[str]: + """Downloads bioRxiv abstracts (text only) via the public paginated API. + + Args: + start_date: Start of the query date range (YYYY-MM-DD). + end_date: End of the query date range (YYYY-MM-DD). + max_records: Maximum number of abstracts to return. + + Returns: + A list of abstract strings (at most ``max_records``). + """ + abstracts: list[str] = [] + cursor = 0 + + while len(abstracts) < max_records: + url = f'{_API_BASE}/{start_date}/{end_date}/{cursor}/json' + logging.info('Fetching bioRxiv page: %s', url) + req = urllib.request.Request(url) + with urllib.request.urlopen(req, timeout=30) as response: + data = json.loads(response.read().decode('utf-8')) + + collection = data.get('collection', []) + if not collection: + logging.info('No more records at cursor %d; stopping.', cursor) + break + + for record in collection: + abstract = record.get('abstract', '').strip() + if abstract: + abstracts.append(abstract) + if len(abstracts) >= max_records: + break + + cursor += len(collection) + logging.info('Collected %d abstracts so far.', len(abstracts)) + time.sleep(1.0) # Be polite to the public API. + + logging.info('Downloaded %d bioRxiv abstracts.', len(abstracts)) + return abstracts + + +def main(argv: Sequence[str]) -> None: + if len(argv) > 1: + raise app.UsageError('Too many command-line arguments.') + + logging.info('Devices (%d): %s', len(jax.devices()), jax.devices()) + + # 1. Download real bioRxiv abstracts. + abstracts = fetch_abstracts( + _START_DATE.value, _END_DATE.value, _NUM_ABSTRACTS.value + ) + if not abstracts: + raise app.UsageError('No abstracts were downloaded; check the date range.') + + # 2. Frame each abstract as a (fixed instruction prompt -> abstract) SFT pair. + train_data = [(_INSTRUCTION, abstract) for abstract in abstracts] + + # 3. Differentially privately fine-tune Gemma with DP-SGD. + model_name: model.ModelName = _MODEL.value # type: ignore[assignment] + model_variant = model.GemmaModel.default(model_name) + config = execution_plan.BandMFConfig.default( + num_bands=1, + iterations=_ITERATIONS.value, + expected_participations=1.0, + ) + fine_tuner = dp_sft.DPFineTuner( + model_variant=model_variant, + mechanism_config=config, + lora_rank=_LORA_RANK.value, + max_seq_length=_MAX_SEQ_LENGTH.value, + optimizer=optax.adamw(_LEARNING_RATE.value), + performance_flags=execution_plan.PerformanceFlags(microbatch_size=1), + ).calibrate(epsilon=_EPSILON.value, delta=_DELTA.value) + + logging.info( + 'DP fine-tuning on %d abstracts for %d iterations at (eps=%.1f, ' + 'delta=%.0e)...', + len(train_data), + _ITERATIONS.value, + _EPSILON.value, + _DELTA.value, + ) + result = fine_tuner(rng=0, data=train_data) + logging.info('Fine-tuning complete.') + + # 4. Generate synthetic abstracts by sampling the fine-tuned model. We use the + # same instruction prompt with a non-zero temperature and a different rng per + # draw to obtain diverse outputs. + sampler = result.sampler( + max_out_length=_MAX_OUT_LENGTH.value, + sampling=gm.text.RandomSampling(temperature=_TEMPERATURE.value), + ) + + logging.info('Generating %d synthetic abstracts...', _NUM_SAMPLES.value) + for i in range(_NUM_SAMPLES.value): + synthetic_abstract = sampler.chat(_INSTRUCTION, rng=i) + logging.info('=' * 70) + logging.info('SYNTHETIC ABSTRACT %d:\n%s', i + 1, synthetic_abstract) + logging.info('=' * 70) + logging.info('Done. Generated %d synthetic abstracts.', _NUM_SAMPLES.value) + + +if __name__ == '__main__': + app.run(main) diff --git a/dpsynth/text/dp_sft.py b/dpsynth/text/dp_sft.py new file mode 100644 index 0000000..9865e92 --- /dev/null +++ b/dpsynth/text/dp_sft.py @@ -0,0 +1,190 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Differentially private fine-tuning for Gemma language models. + +``DPFineTuner`` is a ``DPMechanism`` that takes raw text strings and handles +everything end-to-end: tokenization, model loading, LoRA application, and +DP-SGD training via ``DPTrainer``. + +Example usage:: + + config = execution_plan.BandMFConfig.default( + num_bands=1, iterations=100, expected_participations=1.0, + ) + fine_tuner = DPFineTuner( + model_variant=model.GemmaModel.default('gemma3_270m_it'), + mechanism_config=config, + ).configure(zcdp_rho=0.5) + + data = [("What is 2+2?", "4"), ("Capital of France?", "Paris")] + final_state = fine_tuner(rng=42, data=data) + +For the underlying DP-SGD implementation, see ``dp_trainer``. +For model loading and loss construction, see ``model``. +""" + +from __future__ import annotations + +from collections.abc import Sequence +import dataclasses +import math + +import dp_accounting +from dpsynth.local_mode import primitives +from dpsynth.text import dp_trainer +from dpsynth.text import model +from gemma import gm +from gemma import peft +from jax_privacy import execution_plan +from jax_privacy import training +import optax + + +@dataclasses.dataclass +class FineTuneResult: + """Result of running ``DPFineTuner``. + + Private training state (noise state, optimizer state) is intentionally + excluded. + + Attributes: + model: The model architecture (LoRA-wrapped, with adapters folded in). + params: Pretrained + trained LoRA params, merged and ready for sampling. + """ + + model: gm.nn.TransformerLike + params: training.Params + + def sampler(self, **kwargs) -> gm.text.ChatSampler: + """Builds a ChatSampler from the fine-tuned model. + + Args: + **kwargs: Arguments forwarded to ``gm.text.ChatSampler`` (e.g. + ``max_out_length``, ``cache_length``). + + Returns: + A ``ChatSampler`` ready for text generation. + """ + return gm.text.ChatSampler( + model=self.model, + params=self.params, + **kwargs, + ) + + +@dataclasses.dataclass +class DPFineTuner(primitives.DPMechanism): + """Differentially private fine-tuning of Gemma models via DP-SGD. + + 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. + """ + + model_variant: model.GemmaModel + mechanism_config: execution_plan.BandMFConfig + lora_rank: int = 16 + max_seq_length: int = 512 + optimizer: optax.GradientTransformation = dataclasses.field( + default_factory=lambda: optax.adamw(1e-4) + ) + performance_flags: execution_plan.PerformanceFlags | None = None + + def configure(self, *, zcdp_rho: float, delta: float = 0.0) -> DPFineTuner: + """Returns a copy with noise_multiplier calibrated to the zCDP budget. + + Sets the noise_multiplier to satisfy ``zcdp_rho`` under a **loose upper + bound** that ignores Poisson subsampling amplification: the unamplified + composition of ``T`` Gaussian mechanisms with noise_multiplier ``sigma`` + has zCDP cost ``T / (2 * sigma**2)``, so we set + ``sigma = sqrt(T / (2 * rho))``. + + This is deliberately conservative. The ``dp_event`` property returns the + **full** event including subsampling amplification, so downstream callers + using ``dp_accounting`` (e.g. ``calibrate_dp_mechanism`` over a + heterogeneous composition) will get tight PLD-based accounting from the + raw events -- not from this loose zCDP bound. + + Args: + zcdp_rho: The zCDP privacy budget (rho). + delta: Unused. Accepted for interface compatibility. + + Returns: + A new ``DPFineTuner`` with calibrated + ``mechanism_config.noise_multiplier``. + """ + num_bands = len(self.mechanism_config.strategy) + rounds = math.ceil(self.mechanism_config.iterations / num_bands) + noise_multiplier = math.sqrt(rounds / (2.0 * zcdp_rho)) + calibrated_config = dataclasses.replace( + self.mechanism_config, + noise_multiplier=noise_multiplier, + ) + return dataclasses.replace(self, mechanism_config=calibrated_config) + + @property + def dp_event(self) -> dp_accounting.DpEvent: + """The DpEvent characterizing the privacy cost of DP-SGD training.""" + if self.mechanism_config.noise_multiplier is None: + raise ValueError('noise_multiplier is not set. Call calibrate() first.') + return self.mechanism_config.make( + performance_flags=self.performance_flags + ).dp_event + + def __call__( + self, + rng: int, + data: Sequence[tuple[str, str]], + ) -> 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. + + Returns: + A ``FineTuneResult`` with trained LoRA parameters and a sampler factory. + 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, + ) + + lora_config = model.LoraConfig(rank=self.lora_rank) + module, frozen_params, trainable_params = model.load_gemma( + self.model_variant, + lora_config, + seq_length=self.max_seq_length, + ) + + def loss_fn(trainable_params, batch, prng): + del prng # Deterministic forward pass (no dropout during DP training). + full_params = peft.merge_params(frozen_params, trainable_params) + return model.sft_loss_fn(module, full_params, batch) + + trainer = dp_trainer.DPTrainer( + mechanism_config=self.mechanism_config, + init_params=trainable_params, + loss_fn=loss_fn, + optimizer=self.optimizer, + performance_flags=self.performance_flags, + ) + state = trainer(rng=rng, data=dataset) + + merged = peft.merge_params(frozen_params, state.params) # pytype: disable=wrong-arg-types + return FineTuneResult(model=module, params=merged) diff --git a/dpsynth/text/dp_trainer.py b/dpsynth/text/dp_trainer.py new file mode 100644 index 0000000..d2f5887 --- /dev/null +++ b/dpsynth/text/dp_trainer.py @@ -0,0 +1,160 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""General-purpose differentially private training for JAX models. + +``DPTrainer`` is a ``DPMechanism`` that runs DP-SGD on an arbitrary JAX loss +function and parameter pytree. It wraps ``jax_privacy`` — which provides +per-example gradient clipping, noise addition, and privacy accounting — and +exposes the same three-phase configure → calibrate → run API as other dpsynth +mechanisms (``DPHistogram``, ``DPMarginals``, etc.). + +**Design rationale.** The key design decision is framework-agnosticism: +``DPTrainer`` operates on pure JAX pytrees and callables, with no dependency on +Flax, NNX, or any model-specific concepts like LoRA or tokenization. This +preserves the full generality of ``jax_privacy`` (which itself is +framework-agnostic) while fitting into dpsynth's ``DPMechanism`` protocol for +uniform calibration and composition. Supervised fine-tuning of Flax NNX models +is supported as a special case: the caller splits the ``nnx.Module`` into +trainable and frozen state, closes over the frozen state in the loss function, +and passes the trainable pytree to ``DPTrainer``. The higher-level +``DPTextTrainer`` handles this NNX plumbing, along with Gemma model loading, +LoRA application, tokenization, and checkpoint management — keeping this module +focused on the DP-SGD training loop itself. + +Usage:: + + trainer = DPTrainer.default( + init_params=my_params, + loss_fn=my_loss, + iterations=100, + batch_size=8, + num_examples=1000, + ).configure(zcdp_rho=0.5) + + final_state = trainer(rng=42, data=dataset) +""" + +from __future__ import annotations + +import dataclasses +import json +import math + +from absl import logging +import dp_accounting +from dpsynth.local_mode import primitives +from jax_privacy import execution_plan +from jax_privacy.experimental import training +import optax + + +@dataclasses.dataclass(frozen=True, slots=True, kw_only=True) +class DPTrainer(primitives.DPMechanism): + """General-purpose differentially private training via DP-SGD. + + Wraps ``jax_privacy.DPTrainer`` in the ``DPMechanism`` protocol, providing + calibration, privacy accounting, and a clean ``(rng, data)`` call signature. + + The caller is responsible for preparing ``init_params`` and ``loss_fn``. For + Flax NNX models, this means splitting the model into trainable and frozen + state and closing over the frozen state in the loss function. For pure JAX + code, the params and loss function can be used directly. + + Mechanism can be configured via `jax_privacy.execution_plan.BandMFConfig`. + To configure a basic 4-epoch DP-SGD, use: + + >>> config = execution_plan.BandMFConfig.default( + ... num_bands=1, iterations=500, expected_participations=4, l2_clip_norm=1 + ... ) + + Attributes: + mechanism_config: ``BandMFConfig`` specifying the DP-SGD mechanism. + init_params: Initial trainable parameter pytree. + loss_fn: ``(params, batch, prng) -> scalar`` loss function. Must be + compatible with ``jax.grad`` w.r.t. the first argument. + optimizer: Optax gradient transformation. + performance_flags: Optional ``PerformanceFlags`` for microbatching, SPMD, + compute dtype, and noise generation. + callback: Optional ``(step, state, aux) -> None`` called after each step. + """ + + init_params: training.Params + loss_fn: training.LossFn + mechanism_config: execution_plan.BandMFConfig + optimizer: optax.GradientTransformation + performance_flags: execution_plan.PerformanceFlags | None = None + callback: training.CallbackFn | None = None + + def configure(self, *, zcdp_rho: float, delta: float = 0.0) -> DPTrainer: + """Returns a copy with noise calibrated to the zCDP budget. + + Uses a loose upper bound ignoring subsampling amplification: + ``sigma = sqrt(T / (2 * rho))``. The ``dp_event`` property returns the + full event with amplification for tight downstream accounting. + + Args: + zcdp_rho: The zCDP privacy budget (rho). + delta: Unused. Accepted for interface compatibility. + + Returns: + A new ``DPTrainer`` with calibrated ``config.noise_multiplier``. + """ + num_bands = len(self.mechanism_config.strategy) # pyrefly: ignore[bad-argument-type] + rounds = math.ceil(self.mechanism_config.iterations / num_bands) + noise_multiplier = math.sqrt(rounds / (2.0 * zcdp_rho)) + calibrated_config = dataclasses.replace( + self.mechanism_config, + noise_multiplier=noise_multiplier, + ) + return dataclasses.replace(self, mechanism_config=calibrated_config) + + @property + def dp_event(self) -> dp_accounting.DpEvent: + """The DpEvent characterizing the privacy cost of DP-SGD training.""" + if self.mechanism_config.noise_multiplier is None: + raise ValueError('noise_multiplier is not set. Call calibrate() first.') + return self._make_plan().dp_event + + def _make_plan(self) -> execution_plan.DPExecutionPlan: + return self.mechanism_config.make(performance_flags=self.performance_flags) + + def __call__(self, rng: int, data: training.Batch) -> training.TrainingState: + """Runs DP-SGD training on the given dataset. + + Args: + rng: Random seed for batch selection and noise generation. + data: Training data as a dict of JAX arrays. The first axis of each array + is the example axis. + + Returns: + Final ``TrainingState`` containing the trained parameters. + """ + if self.mechanism_config.noise_multiplier is None: + raise ValueError('noise_multiplier is not set. Call calibrate() first.') + + d = dataclasses.asdict(self.mechanism_config) + d['strategy'] = self.mechanism_config.strategy.tolist() # JSON/numpy hack. + logging.info('DPTrainer config:\n%s', json.dumps(d, indent=2)) + + plan = self._make_plan() + dp_trainer = training.DPTrainer( + plan=plan, loss_fn=self.loss_fn, optimizer=self.optimizer + ) + return dp_trainer.fit( + data, + self.init_params, + callback=self.callback, + rng_or_seed=rng, + ) diff --git a/dpsynth/text/model.py b/dpsynth/text/model.py new file mode 100644 index 0000000..4763c1a --- /dev/null +++ b/dpsynth/text/model.py @@ -0,0 +1,215 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Gemma model loading and SFT loss function.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +import dataclasses +from typing import Any, Literal + +from absl import logging +from gemma import gm +from gemma import peft +import jax +import jax.numpy as jnp +import numpy as np +import optax + +# Type alias for Gemma model params (nested dict of jax.Array). +Params = Any + +# Supported model names for GemmaModel.default(). +ModelName = Literal[ + 'gemma3_270m_it', + 'gemma3_1b_it', + 'gemma3_4b_it', + 'gemma4_e2b_it', + 'gemma4_e4b_it', +] + +_DEFAULTS: dict[str, tuple[Callable[..., Any], str, Callable[..., Any]]] = { + 'gemma3_270m_it': ( + gm.nn.Gemma3_270M, + gm.ckpts.CheckpointPath.GEMMA3_270M_IT, + gm.text.Gemma3Tokenizer, + ), + 'gemma3_1b_it': ( + gm.nn.Gemma3_1B, + gm.ckpts.CheckpointPath.GEMMA3_1B_IT, + gm.text.Gemma3Tokenizer, + ), + 'gemma3_4b_it': ( + gm.nn.Gemma3_4B, + gm.ckpts.CheckpointPath.GEMMA3_4B_IT, + gm.text.Gemma3Tokenizer, + ), + 'gemma4_e2b_it': ( + gm.nn.Gemma4_E2B, + gm.ckpts.CheckpointPath.GEMMA4_E2B_IT, + gm.text.Gemma4Tokenizer, + ), + 'gemma4_e4b_it': ( + gm.nn.Gemma4_E4B, + gm.ckpts.CheckpointPath.GEMMA4_E4B_IT, + gm.text.Gemma4Tokenizer, + ), +} + + +@dataclasses.dataclass(frozen=True) +class GemmaModel: + """Specification for a Gemma model variant.""" + + model_class: Callable[..., Any] + checkpoint_path: str + tokenizer_class: Callable[..., Any] + + @classmethod + def default(cls, name: ModelName) -> GemmaModel: + """Constructs a GemmaModel from a preset name.""" + if name not in _DEFAULTS: + raise ValueError(f'Unknown model {name!r}. Options: {list(_DEFAULTS)}') + model_class, checkpoint_path, tokenizer_class = _DEFAULTS[name] + return cls(model_class, checkpoint_path, tokenizer_class) + + +@dataclasses.dataclass(frozen=True) +class LoraConfig: + """Configuration for LoRA adaptation.""" + + rank: int = 16 + dtype: Any = jnp.bfloat16 + + +def load_gemma( + model_variant: GemmaModel, + lora_config: LoraConfig, + *, + seq_length: int = 64, +) -> tuple[Any, Params, Params]: + """Loads a pretrained Gemma model with LoRA adapters. + + Args: + model_variant: Which Gemma variant to load. + lora_config: LoRA adapter configuration. + seq_length: Sequence length for model initialization. + + Returns: + ``(module, frozen_params, trainable_params)`` tuple. + """ + base_model = model_variant.model_class() + model = gm.nn.LoRA( + rank=lora_config.rank, + model=base_model, + dtype=lora_config.dtype, + ) + + dummy_tokens = jnp.ones((1, seq_length), dtype=jnp.int32) + variables = model.init(jax.random.key(0), tokens=dummy_tokens) + + params, lora_params = peft.split_params(variables['params']) + pt_params = gm.ckpts.load_params(model_variant.checkpoint_path, params=params) + + num_trainable = optax.tree.size(lora_params) + num_frozen = optax.tree.size(pt_params) + logging.info( + 'Loaded Gemma model w/ LoRA (rank=%d): %d trainable (%.4f%%), %d frozen', + lora_config.rank, + num_trainable, + 100.0 * num_trainable / (num_trainable + num_frozen), + num_frozen, + ) + + return model, pt_params, lora_params + + +def sft_loss_fn( + module: Any, + full_params: Params, + data: dict[str, jax.Array], +) -> tuple[jax.Array, dict[str, jax.Array]]: + """Cross-entropy next-token-prediction loss for supervised fine-tuning. + + Args: + module: LoRA-wrapped Gemma model. + full_params: Full parameter dict (frozen + trainable, merged). + data: Dict with ``'input_tokens'`` and ``'loss_mask'`` (int32 ``[B, L]``). + + Returns: + ``(loss, aux)`` where ``aux`` contains ``'loss'``. + """ + input_tokens = data['input_tokens'] + loss_mask = data['loss_mask'] + + out = module.apply({'params': full_params}, tokens=input_tokens) + logits = out.logits[:, :-1, :] + targets = input_tokens[:, 1:] + mask = loss_mask[:, 1:] + + pt_losses = optax.softmax_cross_entropy_with_integer_labels(logits, targets) + loss = jnp.sum(pt_losses * mask) / jnp.maximum(jnp.sum(mask), 1.0) + return loss, {'loss': loss} + + +def tokenize_texts( + examples: Sequence[tuple[str, str]], + model_variant: GemmaModel, + max_seq_length: int, +) -> dict[str, np.ndarray]: + """Tokenizes (prompt, response) pairs for supervised fine-tuning. + + Prompt tokens are masked out (``loss_mask=0``) so only the response + contributes to the training loss. Turn formatting follows the Gemma + dialog template (forked from ``gemma/gm/data/_tasks.py``). + + Args: + examples: Sequence of ``(prompt, response)`` string pairs. + model_variant: Determines which tokenizer and turn format to use. + max_seq_length: Maximum sequence length (including special tokens). + + Returns: + Dict with ``'input_tokens'`` and ``'loss_mask'`` (int32 ``[N, L]``). + """ + tokenizer = model_variant.tokenizer_class() + sp = tokenizer.special_tokens + sot = tokenizer.tokens[sp.START_OF_TURN] + eot = tokenizer.tokens[sp.END_OF_TURN] + + 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 = f'{sot}user\n{prompt}{eot}\n{sot}model\n' + response_str = f'{response}{eot}' + 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 + + logging.info( + 'Tokenized %d examples (max_seq_length=%d)', + len(examples), + max_seq_length, + ) + + return {'input_tokens': tokens, 'loss_mask': mask} diff --git a/tests/text/dp_sft_test.py b/tests/text/dp_sft_test.py new file mode 100644 index 0000000..6c7eceb --- /dev/null +++ b/tests/text/dp_sft_test.py @@ -0,0 +1,198 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for dpsynth.text.model and dpsynth.text.dp_sft. + +These tests verify the loss function, configuration, and calibration logic +without loading any real model checkpoints. +""" + +import dataclasses +import math + +from absl.testing import absltest +from dpsynth.text import dp_sft +from dpsynth.text import model +import jax +import jax.numpy as jnp +from jax_privacy import execution_plan + + +def _default_config(): + return execution_plan.BandMFConfig.default( + num_bands=1, + iterations=100, + expected_participations=1.0, + ) + + +def _make_tiny_model_and_params(vocab_size=8, embed_dim=4, seq_len=4): + """Creates a minimal Flax Linen model and random params for testing.""" + import flax.linen as nn # pylint: disable=g-import-not-at-top + + @dataclasses.dataclass + class _Output: + logits: jnp.ndarray + + class TinyModel(nn.Module): + vocab_size: int + embed_dim: int + + @nn.compact + def __call__(self, tokens, **kwargs): + x = nn.Embed(self.vocab_size, self.embed_dim)(tokens) + logits = nn.Dense(self.vocab_size)(x) + return _Output(logits=logits) + + module = TinyModel(vocab_size=vocab_size, embed_dim=embed_dim) + dummy = jnp.ones((1, seq_len), dtype=jnp.int32) + params = module.init(jax.random.key(0), dummy)['params'] + return module, params + + +class SftLossFnTest(absltest.TestCase): + + def test_loss_is_finite_and_positive(self): + module, params = _make_tiny_model_and_params() + data = { + 'input_tokens': jnp.array([[1, 2, 3, 4, 0]], dtype=jnp.int32), + 'loss_mask': jnp.array([[1, 1, 1, 1, 0]], dtype=jnp.int32), + } + loss, aux = model.sft_loss_fn(module, params, data) + self.assertTrue(jnp.isfinite(loss)) + self.assertGreater(float(loss), 0.0) + self.assertIn('loss', aux) + + def test_loss_decreases_with_gradient_descent(self): + """Gradient descent on a tiny model should reduce loss.""" + module, params = _make_tiny_model_and_params() + tokens = jnp.array([[1, 2, 3, 1]], dtype=jnp.int32) + mask = jnp.ones_like(tokens) + data = {'input_tokens': tokens, 'loss_mask': mask} + + loss_before = float(model.sft_loss_fn(module, params, data)[0]) + + # Overfit on this single example. + for _ in range(200): + + def step_fn(p): + return model.sft_loss_fn(module, p, data) + + (_, _), grads = jax.value_and_grad(step_fn, has_aux=True)(params) + params = jax.tree.map(lambda p, g: p - 0.1 * g, params, grads) + + loss_after = float(model.sft_loss_fn(module, params, data)[0]) + self.assertLess(loss_after, loss_before) + + def test_masked_tokens_do_not_contribute(self): + module, params = _make_tiny_model_and_params() + tokens = jnp.array([[1, 2, 3, 4]], dtype=jnp.int32) + mask_short = jnp.array([[1, 1, 0, 0]], dtype=jnp.int32) + mask_full = jnp.array([[1, 1, 1, 1]], dtype=jnp.int32) + + loss_short = float( + model.sft_loss_fn( + module, params, {'input_tokens': tokens, 'loss_mask': mask_short} + )[0] + ) + loss_full = float( + model.sft_loss_fn( + module, params, {'input_tokens': tokens, 'loss_mask': mask_full} + )[0] + ) + # Different masks should give different losses. + self.assertNotAlmostEqual(loss_short, loss_full, places=3) + + +class GemmaModelTest(absltest.TestCase): + + def test_default_preset(self): + m = model.GemmaModel.default('gemma3_270m_it') + self.assertIsNotNone(m.checkpoint_path) + self.assertTrue(callable(m.model_class)) + self.assertTrue(callable(m.tokenizer_class)) + + def test_custom_checkpoint_path(self): + m = dataclasses.replace( + model.GemmaModel.default('gemma3_270m_it'), + checkpoint_path='/custom/path', + ) + self.assertEqual(m.checkpoint_path, '/custom/path') + + def test_unknown_name_raises(self): + with self.assertRaises(ValueError): + model.GemmaModel.default('nonexistent') + + +class LoraConfigTest(absltest.TestCase): + + def test_defaults(self): + cfg = model.LoraConfig() + self.assertEqual(cfg.rank, 16) + self.assertEqual(cfg.dtype, jnp.bfloat16) + + def test_custom_values(self): + cfg = model.LoraConfig(rank=8, dtype=jnp.float32) + self.assertEqual(cfg.rank, 8) + self.assertEqual(cfg.dtype, jnp.float32) + + +class DPFineTunerTest(absltest.TestCase): + + def test_creates_valid_mechanism_config(self): + mechanism = dp_sft.DPFineTuner( + model_variant=model.GemmaModel.default('gemma3_270m_it'), + mechanism_config=_default_config(), + ) + self.assertEqual(mechanism.mechanism_config.iterations, 100) + self.assertIsNone(mechanism.mechanism_config.noise_multiplier) + + def test_calibrate_sets_noise_multiplier(self): + mechanism = dp_sft.DPFineTuner( + model_variant=model.GemmaModel.default('gemma3_270m_it'), + mechanism_config=_default_config(), + ).configure(zcdp_rho=0.5) + self.assertIsNotNone(mechanism.mechanism_config.noise_multiplier) + self.assertGreater(mechanism.mechanism_config.noise_multiplier, 0.0) + + def test_calibrate_noise_formula(self): + mechanism = dp_sft.DPFineTuner( + model_variant=model.GemmaModel.default('gemma3_270m_it'), + mechanism_config=_default_config(), + ).configure(zcdp_rho=0.5) + # Single band: rounds = iterations, sigma = sqrt(T / (2*rho)). + expected = math.sqrt(100 / (2.0 * 0.5)) + self.assertAlmostEqual( + mechanism.mechanism_config.noise_multiplier, expected + ) + + def test_dp_event_before_calibration_raises(self): + mechanism = dp_sft.DPFineTuner( + model_variant=model.GemmaModel.default('gemma3_270m_it'), + mechanism_config=_default_config(), + ) + with self.assertRaises(ValueError): + _ = mechanism.dp_event + + def test_dp_event_after_calibration(self): + mechanism = dp_sft.DPFineTuner( + model_variant=model.GemmaModel.default('gemma3_270m_it'), + mechanism_config=_default_config(), + ).configure(zcdp_rho=0.5) + event = mechanism.dp_event + self.assertIsNotNone(event) + + +if __name__ == '__main__': + absltest.main() diff --git a/tests/text/dp_trainer_test.py b/tests/text/dp_trainer_test.py new file mode 100644 index 0000000..43ca586 --- /dev/null +++ b/tests/text/dp_trainer_test.py @@ -0,0 +1,114 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from absl.testing import absltest +from dpsynth.text import dp_trainer +from flax import nnx +import jax.numpy as jnp +import jax_privacy +import optax + + +def _dummy_params_and_loss(): + """Creates a trivial pytree and loss function for testing.""" + params = {'w': jnp.ones((4, 4))} + + def loss_fn(params, batch, prng): + del prng + return jnp.sum(params['w'] * batch['x']), () + + return params, loss_fn + + +class DPTrainerTest(absltest.TestCase): + + @property + def dpsgd_config(self): + return jax_privacy.execution_plan.BandMFConfig.default( + iterations=100, + expected_participations=4, + num_bands=1, + l2_clip_norm=1, + ) + + def test_default_creates_valid_config(self): + params, loss_fn = _dummy_params_and_loss() + trainer = dp_trainer.DPTrainer( + init_params=params, + loss_fn=loss_fn, + mechanism_config=self.dpsgd_config, + optimizer=optax.adamw(1e-4), + ) + self.assertEqual(trainer.mechanism_config.iterations, 100) + self.assertIsNone(trainer.mechanism_config.noise_multiplier) + + def test_calibrate_sets_noise_multiplier(self): + params, loss_fn = _dummy_params_and_loss() + trainer = dp_trainer.DPTrainer( + init_params=params, + loss_fn=loss_fn, + mechanism_config=self.dpsgd_config, + optimizer=optax.adamw(1e-4), + ).configure(zcdp_rho=0.5) + # Single band: sigma = sqrt(T / (2 * rho)) = 10. + self.assertAlmostEqual(trainer.mechanism_config.noise_multiplier, 10.0) + self.assertIsNotNone(trainer.dp_event) + + def test_raises_before_calibration(self): + params, loss_fn = _dummy_params_and_loss() + trainer = dp_trainer.DPTrainer( + init_params=params, + loss_fn=loss_fn, + mechanism_config=self.dpsgd_config, + optimizer=optax.adamw(1e-4), + ) + with self.assertRaises(ValueError): + _ = trainer.dp_event + + with self.assertRaises(ValueError): + trainer(rng=42, data={'x': jnp.ones((10, 4, 4))}) + + def test_nnx_split_merge_round_trip(self): + """Demonstrates the NNX split/merge pattern for LoRA fine-tuning.""" + base = nnx.Linear(4, 4, rngs=nnx.Rngs(0)) + lora_model = nnx.LoRA(4, 2, 4, base_module=base, rngs=nnx.Rngs(0)) + + graphdef, trainable, frozen = nnx.split(lora_model, nnx.LoRAParam, ...) + + def loss_fn(params, batch, prng): + del prng + model = nnx.merge(graphdef, params, frozen) + x = batch['x'] + return jnp.mean(model(x) ** 2), () + + full_batch_config = jax_privacy.execution_plan.BandMFConfig.default( + num_bands=1, + iterations=5, + expected_participations=5, + l2_clip_norm=1.0, + ) + trainer = dp_trainer.DPTrainer( + init_params=trainable, + loss_fn=loss_fn, + mechanism_config=full_batch_config, + optimizer=optax.adamw(1e-4), + ).configure(zcdp_rho=1.0) + + train_state = trainer(rng=42, data={'x': jnp.ones((10, 4, 4))}) + self.assertIsNotNone(train_state) + + +if __name__ == '__main__': + absltest.main()