beam_search: parallel auto_agent workers in a beam_search orchestrator loop#59
beam_search: parallel auto_agent workers in a beam_search orchestrator loop#59ligh-svg wants to merge 2 commits into
Conversation
| with open(self.baseline_code_path, "r") as f: | ||
| self.baseline_code = f.read() | ||
|
|
||
| def _prepare_grouped_harness(self, candidate_paths: List[str]) -> str: |
There was a problem hiding this comment.
This _prepare_grouped_harness code is actually broken. It simply concatenates all candidate codes, which results in repeated main functions in the output python file. In the simulation evaluator it's fine because we just ask LLM to predict the performance by showing all codes. It won't work if we run this on real TPU backend.
I'll added a comment to clarify this.
|
|
||
| # Step 1: Plan | ||
| logging.info(f"[{self.name}] Running PlanKernelAgent...") | ||
| async for event in self.plan_agent.run_async(ctx): |
There was a problem hiding this comment.
For the planning agent, there are a few different strategies in AutoComp. Select from the optimization menu and the current version of the kernel code, or combine from best candidates from previous iteration. Are we capturing those combinations here?
There was a problem hiding this comment.
It implements the random selection from the optimization menu, but doesn't support combination of best candidates yet.
|
|
||
| #### **Pros & Cons:** | ||
| * **Pros**: | ||
| * Clean separation: The execution loop logic for "tuning search" and "correctness validation" remain separate. |
There was a problem hiding this comment.
This is great! +1 to this initial design
| We modify the base `AutonomousPipelineAgent` to accept `None` for `autotune_agent` and `profile_agent`, and adjust the control flow logic internally. | ||
|
|
||
| #### **Changes to `AutonomousPipelineAgent`:** | ||
| 1. Make subagents optional in constructor: |
There was a problem hiding this comment.
In AutoComp, the initial search strategies are less rigorous than latter ones (even allowing sub-optimal candidates to survive). Should we mimic this in the new design where perhaps in the initial iterations we can set a shallower iteration in the autoagent compared to latter.
I might also oppose the above due to increased stochasticity of auto-agent with smaller N
There was a problem hiding this comment.
The option A & B here is a detail about how to implement the worker agent that modifies the current auto-agent to do coding-only-but-no-evaluation.
It sounds to me the idea of start-shallower-then-go-deeper assumes we run the full auto-agent cycle in every beam search branch. So far I set the iteration to 1 for simplicity (the auto agent will stop as soon as it gets a successful profiling results). We need more tweaks in the orchestrator to adjust the max_iterations per auto-agent as the orchestrator enters later rounds.
2 participants
Part 3/3 of the Beam Search prototype. It introduces BeamWorkerPipeline which exits on compilation success, and the AgenticSearchOrchestrator that runs parallel workers (support both full-cycle AutonomousPipelineAgent and the new BeamWorkerPipeline), groups candidate and runs batch evaluation (if using BeamWorkerPipeline), sorts and prunes the beams and iterates on optimization.
The prototype was tested in 2 modes, summarized in:
The test runs were merely for verification of the orchestration flow. All the performance numbers are either simulated or fake. Don't rely on them to assess the pipeline's efficiency.