Fix/priority search #37

chinganc · 2025-09-12T23:49:56Z

This PR finishes the PrioritySearch

It's been tested on the convex optimization and some prompt optimization problems.
It supports

Running multiple optimizers in parallel
Branching out optimizers along their update chains. This is useful for using optimizer with memory and multiple optimizers.

Changes to core in PR:

ModelWrapper is now moved to the top module as Model so that classed decorated by @trace.model can be pickled.
The old implementation of Algorithm.save and Algorithm.load at the abstract class level are commented out. We will leave that to the subclasses.
Some minor bug fixes.

NOTE
-Resuming experiments. will be done in future PRs

Remove mailing list

…into fix/priority_search

…date

…nTrace into fix/priority_search

…emory Add compress_candidate_memory method to PrioritySearch

…y_search

…e Node

…o recently added compress_candidate_memory.

…nTrace into fix/priority_search_pickling

…commented.

add a saving method to pre-save source code

…s for export

- Make the copied modules' parameter nodes have the same as the original one, so that optimizer's memory works. - Add a flag to allow using the same optimizer instance across search - Remove commented code

…nd multi algo BENCH

added GEPA in examples/priority_search_on_convex_fn.py

allenanie

A few files that need to be removed (including one that I accidentally committed/pushed). No immediate implementation problem spotted.

allenanie · 2025-09-27T17:11:57Z

examples/priority_search_on_convex_fn_BENCH.py

Can we rename this file to be gepa_on_convex_fn.py? @doxav

allenanie · 2025-09-27T17:25:41Z

opto/optimizers/textgrad.py

Why is save/load removed here -- is it because this solution doesn't work and needs more testing or more customized implementation?

the base Optimizer class now implements getstate and setstate, which skips parameters, so the optimizers can be pickled directly. There is save and load implemented there. What's done manually in the original code is basically doing that.

allenanie · 2025-09-27T18:58:17Z

opto/trace/modules.py

Is the separation between Model and Module mostly for pickling?

Module is quite high level. It's also the parent of FunModule created by bundle.
Model is what is the user should use when building a traceable agent.

allenanie · 2025-09-27T19:58:18Z

opto/trainer/__init__.py

Should we leave resume for a future feature/implementation?

allenanie · 2025-09-27T19:59:44Z

opto/trainer/algorithms/gepa_algorithms.py

GEPA should not be part of the Fix/priority search branch -- can it be a different one? @doxav any chance you can remove this from this PR? You can commit the convex_fn_BENCH to that PR too if you want.

I guess I don't know if @chinganc wants them in this PR lol -- I'll leave it to him to decide :)

Let's do it in a different PR. This PR is already quite large. Let's follow the principle to keep each PR targeted.
@doxav sorry for the trouble.

Maybe another simpler resolution is to move GEPA under features in the PR?
In this way, the intention that the code is not fully tested is clear?

allenanie · 2025-09-27T20:01:52Z

opto/trainer/loader.py

-            # Reset the index for the next epoch
-            self._i = 0
-            self.n_epochs += 1
+        """Get the next batch of data, always of batch_size. If the dataset is smaller or at the end, the batch will include data from the next epoch after shuffling."""


Wait, is this the logic we want?
This will cause the step vs epoch distinction. I think a lot of trainer does this...so I'm ok if this is the route we want to go with...but just curious if there's any special rationale for this

This is to guarantee that the sampled batch is always of the same size.
In the many applications of generative optimization, what I found is that the training dataset size is much smaller than what we typically see in DL. So we see the boundary effects more often. A special degenerate case is single-node optimization, for which, the original sampler won't honor the batch size at all (even if you specify batch_size>0, the optimizer will only see one feedback as opposed to multiple feedback of batch_size).

I see -- that makes sense!

test.py

allenanie · 2025-09-27T20:05:12Z

tests/llm_optimizers_tests/test_gepa_benchmark.py

I have the same concern as before -- GEPA should be its own PR (even for future references/code checks)

allenanie and others added 28 commits August 18, 2025 12:06

Update README.md

05de9c1

Remove mailing list

update webpage links

712e3b4

merge

0ccfe56

Merge branch 'main' into v0.2.0

9b40d17

Merge branch 'experimental' into v0.2.0

1b7f990

add a testing / example code to see if search methods work

c1bffd1

add an import that was missing

0b196ce

reward is already the negative of the loss -- no need to negate it

e24f121

Merge branch 'experimental' of https://github.com/AgentOpt/OpenTrace …

374cc09

…into fix/priority_search

Remove dependency of SearchTemplate on Minibatch. Fix various bugs.

a9828c2

Allow customizing exploitation criterion.

2934b52

Minor update the convex_fn example

d98dc59

Update priority_search to track optimizer

7e49a17

Add support of multiple optimizers

0cfcaab

Fix some bugs of using multi optimizers due hasing def of ModuleCandi…

4377bee

…date

Add an example of using multi-optimizers

1fcb0bc

Remove subbatch arg in PrioritySearch. Use num_batches instead.

4e539c3

Merge branch 'experimental' into fix/priority_search

3956fd8

Add time as priority

c77dcaa

Fix a bug of testing first params twice.

a8b2479

Add save and resume. Fix bugs in tests.

b79a940

Merge branch 'experimental' into fix/priority_search

1dc3b7e

Update save and load of Optimzier, ModelWrapper, Guide, by implementing

0a8f465

Make ModelWrapper pickleable

16ef98a

Add Model class

98c96a1

Add trace.Model

043b837

add a toy test

3b2c571

Resume seems to work now

15b398d

chinganc requested review from allenanie and adith387 September 12, 2025 23:50

chinganc and others added 28 commits September 17, 2025 20:12

Merge branch 'fix/priority_search' of https://github.com/AgentOpt/Ope…

a5d6d9f

…nTrace into fix/priority_search

Add compress_candidate_memory method to PrioritySearch

22ff61a

Change to saving only score in the long term memory.

fa75803

Merge pull request #41 from AgentOpt/feature/add_compress_candidate_m…

f3f0fe2

…emory Add compress_candidate_memory method to PrioritySearch

Update dataloader to circulate through dataset in creating a batch.

961db94

Merge branch 'feature/add_compress_candidate_memory' into fix/priorit…

13c8612

…y_search

Fix a bug copy.copy Model

7a8f45f

Fix the bug triggered when multiple inputs to MessageNode are the sam…

6d938bc

…e Node

Fix the failure of hash check of candidates in long_term_memory due t…

08e1b1e

…o recently added compress_candidate_memory.

Add self._enforce_using_data_collecting_candidates

8905f8c

Fix the bug of the incorrect fix in 6d938bc

831b034

Merge branch 'fix/bundle_deepcopy' into fix/priority_search

267cfbf

Fix merge error

03e3861

Remove copy overwrites in ModuleCandidate

4ecad9a

Merge branch 'fix/GRAPH_memory_overflow' into fix/priority_search

ce0a6eb

Merge branch 'fix/GRAPH_memory_overflow' into fix/priority_search

e227a37

Merge branch 'experimental' into fix/priority_search

adc451d

Merge branch 'fix/priority_search' of https://github.com/AgentOpt/Ope…

5fcc1b3

…nTrace into fix/priority_search_pickling

push the fix to model.export() -- all tests now pass, unit tests un…

69e5512

…commented.

Merge pull request #39 from AgentOpt/fix/priority_search_pickling

4fae0fc

add a saving method to pre-save source code

fixed the error (due to name shadowing), commented out more test case…

cd4e533

…s for export

Add num_steps to SearchTemplate

4874004

Move assertion into match_candidates_and_samples

55684cc

Add StreamingPrioritySearch

ee76a49

Fix the bug of PS using optimizer with memory

6fe6f25

- Make the copied modules' parameter nodes have the same as the original one, so that optimizer's memory works. - Add a flag to allow using the same optimizer instance across search - Remove commented code

Fix memory due to tracking optimizers in PS

135b33e

Separate priority_search_on_convex_fn into dedicated PrioritySearch a…

de4a77a

…nd multi algo BENCH

Merge pull request #47 from doxav/fix/priority_search

82294b0

added GEPA in examples/priority_search_on_convex_fn.py

allenanie reviewed Sep 27, 2025

View reviewed changes

remove test.py

4890020

Fix/priority search #37

Are you sure you want to change the base?

Fix/priority search #37

Uh oh!

Conversation

chinganc commented Sep 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

allenanie left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

chinganc commented Sep 12, 2025 •

edited

Loading