From cf8075f319b96c2e009e4ea6bb55378faccbbc80 Mon Sep 17 00:00:00 2001 From: DPSynth Team Date: Thu, 20 Aug 2026 19:11:23 -0700 Subject: [PATCH] Add batched inference support in finetune_pubmed.py. PiperOrigin-RevId: 968193592 --- dpsynth/text/dp_sft.py | 7 +- dpsynth/text/model.py | 145 +++++++++++++++++++++++++++++------- examples/finetune_pubmed.py | 51 ++++++++----- 3 files changed, 157 insertions(+), 46 deletions(-) diff --git a/dpsynth/text/dp_sft.py b/dpsynth/text/dp_sft.py index b5090b1..2e8c145 100644 --- a/dpsynth/text/dp_sft.py +++ b/dpsynth/text/dp_sft.py @@ -60,11 +60,11 @@ class FineTuneResult: excluded. Attributes: - model: The model architecture (LoRA-wrapped, with adapters folded in). - params: Pretrained + trained LoRA params, merged and ready for sampling. + model: The LoRA-wrapped model architecture. + params: Pretrained and trained LoRA parameters merged into a single tree. """ - model: gm.nn.TransformerLike + model: gm.nn.LoRA params: training.Params @@ -153,7 +153,6 @@ def __call__( 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): diff --git a/dpsynth/text/model.py b/dpsynth/text/model.py index 4763c1a..0f05365 100644 --- a/dpsynth/text/model.py +++ b/dpsynth/text/model.py @@ -18,13 +18,17 @@ from collections.abc import Callable, Sequence import dataclasses +import itertools +import time from typing import Any, Literal from absl import logging +from etils import epath from gemma import gm from gemma import peft import jax import jax.numpy as jnp +from kauldron import kd import numpy as np import optax @@ -98,17 +102,21 @@ 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. + checkpoint_path: epath.PathLike | None = None, + sharding: Any = None, +) -> tuple[gm.nn.LoRA, Params, Params]: + """Loads a 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. + checkpoint_path: Checkpoint to restore from. If None, loads pretrained base + weights from ``model_variant.checkpoint_path`` and initializes fresh LoRA + adapters. If a checkpoint is provided, restores base and LoRA params. + sharding: Optional sharding tree to constrain parameters across devices. Returns: - ``(module, frozen_params, trainable_params)`` tuple. + ``(module, base_params, lora_params)`` tuple. """ base_model = model_variant.model_class() model = gm.nn.LoRA( @@ -116,24 +124,32 @@ def load_gemma( 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) + if checkpoint_path is not None: + all_params = gm.ckpts.load_params(checkpoint_path, sharding=sharding) + # LoRA params are present when loading from a lora checkpoint + base_params, lora_params = peft.split_params(all_params) # pyrefly: ignore[bad-argument-type] + else: + # When starting from scratch, initialize LoRA params + init_params = model.init( + jax.random.key(0), + tokens=jnp.ones((1, 64), dtype=jnp.int32), + )['params'] + _, lora_params = peft.split_params(init_params) + lora_params = kd.sharding.with_sharding_constraint(lora_params, sharding) + base_params = gm.ckpts.load_params( + model_variant.checkpoint_path, sharding=sharding + ) + + num_lora = optax.tree.size(lora_params) + num_base = optax.tree.size(base_params) logging.info( - 'Loaded Gemma model w/ LoRA (rank=%d): %d trainable (%.4f%%), %d frozen', + 'Loaded Gemma model w/ LoRA (rank=%d): %d lora (%.4f%%), %d base', lora_config.rank, - num_trainable, - 100.0 * num_trainable / (num_trainable + num_frozen), - num_frozen, + num_lora, + 100.0 * num_lora / (num_lora + num_base), + num_base, ) - - return model, pt_params, lora_params + return model, base_params, lora_params def sft_loss_fn( @@ -164,6 +180,23 @@ 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], + ) + # Universal chat prompt format expected by Gemma instruction-tuned models. + 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 +217,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 +224,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 +243,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 diff --git a/examples/finetune_pubmed.py b/examples/finetune_pubmed.py index 2aa9f5b..beffe76 100644 --- a/examples/finetune_pubmed.py +++ b/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( @@ -145,32 +151,38 @@ def load_finetuned() -> dp_sft.FineTuneResult: Returns: A FineTuneResult holding the reconstructed model and loaded params. """ - module, frozen, trainable = model.load_gemma( + module, base_params, lora_params = model.load_gemma( _model_variant(), model.LoraConfig(rank=_LORA_RANK.value), - seq_length=_MAX_SEQ_LENGTH.value, + checkpoint_path=_ckpt_dir(), + sharding=kd.sharding.FSDPSharding(), + ) + return dp_sft.FineTuneResult( + model=module, + params=peft.merge_params(base_params, lora_params), ) - template = peft.merge_params(frozen, trainable) - 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 +210,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)