diff --git a/dpsynth/examples/finetune_pubmed.py b/dpsynth/examples/finetune_pubmed.py index 2aa9f5b..d393780 100644 --- a/dpsynth/examples/finetune_pubmed.py +++ b/dpsynth/examples/finetune_pubmed.py @@ -38,7 +38,9 @@ from etils import epath from gemma import gm from gemma import peft +import jax from jax_privacy import execution_plan +from kauldron import kd import optax _MODEL = flags.DEFINE_enum( @@ -75,8 +77,12 @@ 'learning_rate', 1e-4, 'AdamW learning rate.' ) _NUM_SAMPLES = flags.DEFINE_integer('num_samples', 64, 'Abstracts to generate.') -_MAX_OUT_LENGTH = flags.DEFINE_integer( - 'max_out_length', 512, 'Max output tokens.' +_SAMPLE_BATCH_SIZE = flags.DEFINE_integer( + 'sample_batch_size', + 32, + 'Batch size for sampling across devices; must be divisible by the number of' + ' available devices.', + lower_bound=1, ) _TEMPERATURE = flags.DEFINE_float('temperature', 1.0, 'Sampling temperature.') _SEED = flags.DEFINE_integer( @@ -149,28 +155,36 @@ def load_finetuned() -> dp_sft.FineTuneResult: _model_variant(), model.LoraConfig(rank=_LORA_RANK.value), seq_length=_MAX_SEQ_LENGTH.value, + sharding=kd.sharding.FSDPSharding(), ) template = peft.merge_params(frozen, trainable) - params = gm.ckpts.load_params(_ckpt_dir(), params=template) + params = gm.ckpts.load_params( + _ckpt_dir(), + params=template, + ) return dp_sft.FineTuneResult(model=module, params=params) def sample(result: dp_sft.FineTuneResult) -> None: """Generates synthetic abstracts and writes them to as JSONL.""" - sampler = gm.text.ChatSampler( + sampler = model.GemmaSampler( model=result.model, - params=typing.cast(typing.Mapping[str, typing.Any], result.params), - max_out_length=_MAX_OUT_LENGTH.value, - sampling=gm.text.RandomSampling(temperature=_TEMPERATURE.value), + params=result.params, + max_seq_length=_MAX_SEQ_LENGTH.value, + temperature=_TEMPERATURE.value, ) + prompts = [_INSTRUCTION] * _NUM_SAMPLES.value + responses = sampler( + prompts, + rng=_SEED.value, + batch_size=_SAMPLE_BATCH_SIZE.value, + ) + out_path = epath.Path(_WORKDIR.value) / 'synthetic_abstracts.jsonl' with out_path.open('w') as f: - # One abstract per call is simple but slow. For higher throughput, pass a - # *list* to the batched sampler (`sampler.sampler.sample` for Gemma 3, - # `sampler.gemma4_sampler.sample` for Gemma 4); chunk it to fit HBM. - for i in range(_NUM_SAMPLES.value): - abstract = sampler.chat(_INSTRUCTION, rng=_SEED.value + i) - logging.info('Synthetic abstract %d:\n%s', i + 1, abstract) + for i, abstract in enumerate(responses): + if i % _SAMPLE_BATCH_SIZE.value == 0: + logging.info('Synthetic abstract %d:\n%s', i + 1, abstract) f.write(json.dumps({'abstract': abstract}) + '\n') logging.info('Wrote %d abstracts to %s.', _NUM_SAMPLES.value, out_path) @@ -198,6 +212,11 @@ def main(_) -> None: result = load_finetuned() if _SAMPLE.value: + if _SAMPLE_BATCH_SIZE.value % len(jax.devices()) != 0: + raise app.UsageError( + f'--sample_batch_size ({_SAMPLE_BATCH_SIZE.value}) must be divisible' + f' by the number of available devices ({len(jax.devices())}).' + ) sample(result) diff --git a/dpsynth/text/model.py b/dpsynth/text/model.py index 4763c1a..dfe3d59 100644 --- a/dpsynth/text/model.py +++ b/dpsynth/text/model.py @@ -18,6 +18,8 @@ from collections.abc import Callable, Sequence import dataclasses +import itertools +import time from typing import Any, Literal from absl import logging @@ -25,6 +27,7 @@ from gemma import peft import jax import jax.numpy as jnp +from kauldron import kd import numpy as np import optax @@ -99,13 +102,15 @@ def load_gemma( lora_config: LoraConfig, *, seq_length: int = 64, -) -> tuple[Any, Params, Params]: + sharding: Any = None, +) -> tuple[Any, Any, Any]: """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. + sharding: Optional sharding tree for restoring pretrained parameters. Returns: ``(module, frozen_params, trainable_params)`` tuple. @@ -121,7 +126,11 @@ def load_gemma( 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) + pt_params = gm.ckpts.load_params( + model_variant.checkpoint_path, + params=params if sharding is None else None, + sharding=sharding, + ) num_trainable = optax.tree.size(lora_params) num_frozen = optax.tree.size(pt_params) @@ -164,6 +173,25 @@ def sft_loss_fn( return loss, {'loss': loss} +def format_prompt( + prompt: str, + tokenizer: Any, +) -> str: + """Formats a prompt string with the Gemma user and model turn tags.""" + sp = tokenizer.special_tokens + sot, eot = ( + tokenizer.tokens[sp.START_OF_TURN], + tokenizer.tokens[sp.END_OF_TURN], + ) + return f'{sot}user\n{prompt}{eot}\n{sot}model\n' + + +def format_response(response: str, tokenizer: Any) -> str: + """Formats a model response string with the end-of-turn tag.""" + eot = tokenizer.tokens[tokenizer.special_tokens.END_OF_TURN] + return f'{response}{eot}' + + def tokenize_texts( examples: Sequence[tuple[str, str]], model_variant: GemmaModel, @@ -184,9 +212,6 @@ def tokenize_texts( 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) @@ -194,8 +219,8 @@ def tokenize_texts( 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_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) @@ -213,3 +238,68 @@ def tokenize_texts( ) return {'input_tokens': tokens, 'loss_mask': mask} + + +class GemmaSampler: + """Batched inference sampler for generating synthetic text.""" + + def __init__( + self, + *, + model: Any, + params: Params, + max_seq_length: int, + temperature: float, + ): + sampling_method = ( + gm.text.RandomSampling(temperature=temperature) + if temperature > 0 + else gm.text.Greedy() + ) + self._sampler = gm.text.Sampler( + model=model, + params=params, + cache_length=max_seq_length, + max_out_length=max_seq_length, + sampling=sampling_method, + ) + + def __call__( + self, + prompts: Sequence[str], + *, + rng: int = 0, + batch_size: int = 32, + ) -> list[str]: + """Formats prompts, batches them across devices, and samples responses. + + Args: + prompts: Sequence of prompt instruction strings. + rng: Base random seed for sampling (default 0). The random seed is set per + batch (``rng + batch_idx``), so fixing the seed and changing the + ``batch_size`` will change the sampled outputs. + batch_size: Inference batch size (default 32). + + Returns: + List of generated response strings corresponding to each prompt. + """ + formatted = [format_prompt(p, self._sampler.tokenizer) for p in prompts] + + results: list[str] = [] + for i, batch_items in enumerate(itertools.batched(formatted, batch_size)): + cur_size = len(batch_items) + batch = list(batch_items) + [batch_items[-1]] * (batch_size - cur_size) + t0 = time.perf_counter() + responses = self._sampler.sample( + batch, sharding=kd.sharding.FIRST_DIM, rng=rng + i + ) + elapsed = time.perf_counter() - t0 + logging.info( + 'Batch %d: %d samples in %.2fs (%.2f samples/s)', + i + 1, + cur_size, + elapsed, + cur_size / elapsed if elapsed > 0 else 0.0, + ) + results.extend([str(r) for r in responses[:cur_size]]) + return results