perf: batch trajectory generation in a single model.generate call (issue #6 item 4)

[hai-pilgrim]

Summary

Addresses item 4 from issue #6 (Performance: training pipeline optimizations).

Previously training_step called _generate_trajectory in a Python loop, issuing one model.generate call per (message × trajectory) pair — B × T GPU round trips per step. This PR replaces that loop with a single batched call.

Changes

• Added _generate_trajectories_batched(user_messages, num_trajectories) to PromptBakingTrainer
  • Repeats each user message num_trajectories times
  • Tokenizes all prompts with left-padding (required for generation so real tokens align at the right)
  • Issues one model.generate call for the entire batch
  • Decodes and pairs outputs back to source messages, filtering empty responses
  • Returns early on empty input without touching the tokenizer
• Updated training_step to call _generate_trajectories_batched instead of the sequential loop

Impact

Scenario	Before	After
batch_size=1, num_traj=4	4 generate calls	1 generate call
batch_size=4, num_traj=4	16 generate calls	1 generate call

The speedup is most significant when batch_size > 1 or num_trajectories > 1, which are the common configurations.

Test plan

• 5 new tests: pair types, message identity, trajectory repeat count, empty input guard, empty-response filtering
• All 13 trainer tests pass

🤖 Generated with Claude Code

[marksverdhei]

Review

The core batched generation logic is sound — left-padding for model.generate is correct, and the single-call approach is a clear win over B×T sequential calls.

Issues to address:

1. Scope creep. This PR bundles ~160 lines of unrelated test additions (deps tests, data tests, KL tests) that overlap with dedicated test PRs (#18, #24, #29 — now merged). Please rebase onto main and drop the test additions that are already covered. The PR should only contain the batched generation changes + its own 5 trainer tests.

2. Misleading variable name. prompt_lengths (line ~188) is actually the padded sequence length (same for all samples since left-padding makes them uniform). Rename to padded_prompt_length or prompt_end_idx to avoid confusion.

3. Conflicts with #26. PR #26 tests training_step by mocking _generate_trajectory in a loop, but this PR replaces that loop with _generate_trajectories_batched. Once this merges, #26's tests break. This is noted for awareness — #26 will need updating.

Please rebase to resolve the overlapping test additions, and rename the variable. The core trainer changes look good.

[marksverdhei]

The batched generation logic is solid and the variable rename looks good. However, this now conflicts with #20 (just merged) which refactored compute_loss/prediction_step into shared helpers (_compute_batched_kl, _prepare_pairs, etc.).

The conflict is in trainer.py — #20 restructured the methods that #17's training_step delegates to. Please rebase onto main. The _generate_trajectories_batched method and training_step changes themselves should apply cleanly; only the context around them shifted.

Will approve and merge once rebased.

[marksverdhei]

Closing — we're pivoting away from on-the-fly trajectory generation toward SFT datasets. This optimization targets code that's slated for removal. Thanks for the contribution.