feat(examples,scripts): add micro-batched single invocation multi-model evaluation and interleaved orchestration

Author: droot

examples/sft/gsm8k/gsm8k_sft.py

@@ -40,8 +40,8 @@ def make_datum(row: dict) -> tinker.Datum:
 
     eval_dataset = load_dataset("openai/gsm8k", "main", split="test[:16]")
     return (
-        SupervisedDatasetFromHFDataset(dataset, self.batch_size, map_fn=make_datum),
-        SupervisedDatasetFromHFDataset(eval_dataset, self.batch_size, map_fn=make_datum),
+      SupervisedDatasetFromHFDataset(dataset, self.batch_size, map_fn=make_datum),
+      SupervisedDatasetFromHFDataset(eval_dataset, self.batch_size, map_fn=make_datum),
     )
 
 

examples/sft/gsm8k/vllm_eval.py

@@ -22,43 +22,59 @@ def main() -> None:
   from tinker_cookbook.tokenizer_utils import get_tokenizer
 
   parser = argparse.ArgumentParser()
-  parser.add_argument("--path", required=True)
+  parser.add_argument("--path", required=True, action="append", help="One or more URI paths to evaluate concurrently")
   parser.add_argument("--base-model", default="Qwen/Qwen2.5-0.5B")
   parser.add_argument("--base-url", default=os.getenv("TINKER_BASE_URL", os.getenv("BASE_URL", "http://127.0.0.1:8000")))
   parser.add_argument("--data", default="gsm8k_test.json")
   parser.add_argument("--gpu-memory-utilization", type=float, default=0.85)
+  parser.add_argument("--microbatch-size", type=int, default=10, help="Number of evaluation problems to dispatch per micro-batch")
   parser.add_argument("--min-accuracy", type=float, default=0.0, help="exit nonzero if accuracy falls below this fraction")
   args = parser.parse_args()
 
   with open(args.data) as f:
     data = json.load(f)
 
+  paths = args.path if isinstance(args.path, list) else [args.path]
   client = ServiceClient(api_key=os.getenv("TINKER_API_KEY", "tml-dummy-key"), base_url=args.base_url)
-  sampler = client.create_sampling_client(args.path)
+  samplers = [client.create_sampling_client(p) for p in paths]
   tokenizer = get_tokenizer(args.base_model)
 
   sampling_params = types.SamplingParams(temperature=0.0, max_tokens=256)
   start = time.time()
-  
-  outputs = []
-  for datum in data:
-    prompt_tokens = tokenizer.encode(datum["prompt"], add_special_tokens=False)
-    seqs = sampler.sample(
-        prompt=types.ModelInput.from_ints(tokens=prompt_tokens),
-        num_samples=1,
-        sampling_params=sampling_params,
-    ).result().sequences
-    outputs.append(tokenizer.decode(seqs[0].tokens) if seqs else "")
 
-  elapsed = time.time() - start
-  correct = sum(int(extract(text) == datum["gold"]) for datum, text in zip(data, outputs, strict=True))
-  accuracy = correct / len(data)
+  import asyncio
+
+  async def run_evals():
+    outputs_by_sampler = [[] for _ in paths]
+    batch_size = args.microbatch_size
+    for i in range(0, len(data), batch_size):
+      chunk = data[i : i + batch_size]
+      for s_idx, sampler in enumerate(samplers):
+        tasks = [
+          sampler.sample_async(
+            prompt=types.ModelInput.from_ints(tokens=tokenizer.encode(datum["prompt"], add_special_tokens=False)),
+            num_samples=1,
+            sampling_params=sampling_params,
+          )
+          for datum in chunk
+        ]
+        res_list = await asyncio.gather(*tasks)
+        for res in res_list:
+          seqs = res.sequences
+          outputs_by_sampler[s_idx].append(tokenizer.decode(seqs[0].tokens) if seqs else "")
+    return outputs_by_sampler
 
-  print("***************************************************************")
-  print(f"[SAMPLER] {args.path} 0-shot GSM8K acc = {accuracy:.1%} on {len(data)} problems in {elapsed:.1f}s")
-  print("***************************************************************")
-  if accuracy < args.min_accuracy:
-    raise SystemExit(f"GSM8K accuracy {accuracy:.1%} is below the required {args.min_accuracy:.1%}")
+  outputs_by_sampler = asyncio.run(run_evals())
+
+  elapsed = time.time() - start
+  for path, outputs in zip(paths, outputs_by_sampler, strict=True):
+    correct = sum(int(extract(text) == datum["gold"]) for datum, text in zip(data, outputs, strict=True))
+    accuracy = correct / len(data)
+    print("***************************************************************")
+    print(f"[SAMPLER] {path} 0-shot GSM8K acc = {accuracy:.1%} on {len(data)} problems in {elapsed:.1f}s")
+    print("***************************************************************")
+    if accuracy < args.min_accuracy:
+      raise SystemExit(f"GSM8K accuracy {accuracy:.1%} for {path} is below the required {args.min_accuracy:.1%}")
 
 
 if __name__ == "__main__":

scripts/run_training_e2e.py

@@ -396,12 +396,15 @@ def run_gsm8k_train(config: RunConfig, base_url: str, watch: list[ManagedProcess
   return run_example(config, ["examples/sft/gsm8k/gsm8k_sft.py"], defaults, watch=watch, prefix=prefix)
 
 
-def run_gsm8k_eval(config: RunConfig, model_path: str) -> None:
+def run_gsm8k_eval(config: RunConfig, model_path: str | list[str]) -> None:
+  paths = model_path if isinstance(model_path, list) else [model_path]
+  path_args = []
+  for p in paths:
+    path_args.extend(["--path", p])
   run_command(
     ["uv", "--project", "examples", "run", "python", "examples/sft/gsm8k/vllm_eval.py"]
+    + path_args
     + [
-      "--path",
-      model_path,
       "--base-url",
       config.base_url or "http://127.0.0.1:8000",
       "--data",
@@ -472,10 +475,12 @@ def train(job: str) -> None:
       raise RuntimeError(f"gsm8k {job} failed") from result
 
   check_snapshot_interleaving(config)
-  for job, result in sorted(results.items()):
+  eval_paths = []
+  for _, result in sorted(results.items()):
     assert isinstance(result, str)
-    print(f"[training-e2e] evaluating {job}")
-    run_gsm8k_eval(config, resolve_eval_model_path(result))
+    eval_paths.append(resolve_eval_model_path(result))
+  print(f"[training-e2e] evaluating jobs in single micro-batched invocation: {eval_paths}")
+  run_gsm8k_eval(config, eval_paths)
 
 
 def run_tiny_fft_rl_x2(config: RunConfig, base_url: str, watch: list[ManagedProcess]) -> None: