Add batched inference support in finetune_pubmed.py.

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 <workdir> 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)
 
 

dpsynth/text/model.py

@@ -18,13 +18,16 @@
 
 from collections.abc import Callable, Sequence
 import dataclasses
+import itertools
+import time
 from typing import Any, Literal
 
 from absl import logging
 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
 
@@ -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,18 +212,15 @@ 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)
 
   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