Beam_search when n = 256 is really very slow

[weizier]

I spent almost 10h to get only 30 questions beam_search completions with n = 256 on LLaMA3.1-8B, is that normal? At the same time on the same environment, best_of_256 is about 30 minutes with batch_size = 1.

[Wloner0809]

I spent almost 10h to get only 30 questions beam_search completions with n = 256 on LLaMA3.1-8B, is that normal? At the same time on the same environment, best_of_256 is about 30 minutes with batch_size = 1.

I think this could actually happen. One optimization method is to use vllm&fastchat server instead of vllm offline inference

[lewtun]

Yes, beam search is the slowest of the three algorithms due to its iterative nature and the fact we cannot subsample outputs from a single run (e.g. Best of N and DVTS allow for a single run at fixed N and then subsample M < N).

To speed things up you can either:

• Chunk the dataset and run each chunk in parallel (this is what we did)
• Reduce the number of iterations

In practice we found beam search was best at relatively small N, so you're likely better off using DVTS at large N.

[iseesaw]

Yes, beam search is the slowest of the three algorithms due to its iterative nature and the fact we cannot subsample outputs from a single run (e.g. Best of N and DVTS allow for a single run at fixed N and then subsample M < N).

To speed things up you can either:

• Chunk the dataset and run each chunk in parallel (this is what we did)
• Reduce the number of iterations

In practice we found beam search was best at relatively small N, so you're likely better off using DVTS at large N.

Hi there! I have a few questions about the beam search implementation, as it seems to cause confusion and potentially slows down the process:

1. Beam search for batch prompts

In theory, if we have num_prompts prompts and config.n beams per prompt, we should have num_prompts * config.n beams in total:

search-and-learn/src/sal/search/beam_search.py

Lines 42 to 61 in 179c75b

	beams: list[Beam] = []
	for prompt in batch_of_prompts:
	for i in range(config.n):
	beams.append(
	Beam(
	prompt=prompt,
	index=i,
	current_text="",
	next_texts=None,
	lookahead_texts=None,
	pruned=False,
	completed=False, # New flag to track completion
	stop_reasons=None,
	history=[],
	best_scores=[],
	all_scores=[],
	previous_text=None,
	completion_tokens=0,
	)
	)

However, the code duplicates active beams to ensure there are always config.n beams in each iteration:

search-and-learn/src/sal/search/beam_search.py

Lines 65 to 84 in 179c75b

	for i in tqdm(range(config.num_iterations), desc="Beam search iterations"):
	if i == 0:
	active_beams = [b for b in beams if not b.pruned]
	else:
	active_beams = [b for b in active_beams if not b.pruned]
	# Duplicate active beams to ensure that we have config.n beams per iteration
	if len(active_beams) != config.n:
	repeats = (config.n // len(active_beams)) + 1
	logger.debug(
	f"Extending active_beams with {repeats} repetitions to reach size {config.n}"
	)
	extended_active_beams = [
	copy.deepcopy(b) for b in (active_beams * repeats)[: config.n]
	]
	active_beams = extended_active_beams
	if len(active_beams) != config.n:
	raise ValueError(
	f"Expected {config.n} active beams, but got {len(active_beams)}"
	)

This part is confusing because it suggests that all prompts share the same pool of beams, rather than keeping beams per prompt.

2. Hyperparameters for beam search (config.n and beam_width)

Typically, we would create config.n beams for each prompt, then sample beam_width candidates per beam, resulting in config.n * beam_width total candidates. Next, we would keep only the top config.n among these and repeat. In many standard beam search implementations, config.n equals beam_width. However, in the code, it seems only the first candidate (next_texts[0]) is selected for each beam, which looks more like a greedy approach than a full beam search.

search-and-learn/src/sal/search/beam_search.py

Lines 112 to 133 in 179c75b

	gen_results = generate_k_steps(
	templated_convs, lookahead, llm, sampling_params, 1
	)
	prompts, completions = [], []
	for beam, gen_result in zip(active_beams, gen_results, strict=True):
	beam.next_texts = gen_result.next_texts
	beam.stop_reasons = gen_result.stop_reasons
	beam.lookahead_texts = gen_result.lookahead_texts
	beam.completion_tokens += gen_result.completion_tokens
	beam.current_text += beam.next_texts[0]
	beam.history.append(beam.next_texts[0])
	if (
	beam.stop_reasons[0] == "EOS"
	or beam.stop_reasons[0] == "length"
	or beam.next_texts[0] == ""
	):
	beam.completed = True
	completed_beams.append(beam)
	prompts.append(beam.prompt)
	completions.append([beam.current_text])

Overall, the current implementation is difficult to follow and appears not to leverage parallelization. If I’m misunderstanding something or missing a key detail, I would appreciate any clarification.

Thank you!