active sampling or whatever we call it #1111
Conversation
Summary of ChangesHello @mnoukhov, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly refactors the data preparation and sampling pipeline within the GRPO training framework. The core change is the introduction of an "active sampling" strategy, which intelligently replaces filtered-out prompts with new ones to maintain batch size and improve training efficiency. This involves updates to the data iterator, batch processing, and the addition of new configuration options and metrics. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a significant and valuable feature: active sampling. This refactors the data preparation pipeline to dynamically request new prompts when the initial batch is heavily filtered, which should improve training efficiency and stability. The changes are extensive, particularly in the data_preparation_thread.
My review focuses on a few key areas:
- A fragile implementation detail in the main training loop for generating unique step IDs.
- Potential state management issues in the
ShufflingIterator. - An unused parameter that suggests an incomplete or abandoned feature implementation.
- A minor typo in a log message.
Overall, the active sampling logic itself appears well-designed, with good safeguards against infinite loops. The changes to handle truncated completions by masking instead of filtering are also a solid improvement.
| i = 0.1 | ||
| while filtered_prompts_count > args.num_unique_prompts_rollout: | ||
| another_batch = next_batch(next(iter_dataloader), train_dataset) | ||
| # little hack to make sure that we don't have the same training step, otherwise vllm requests can break | ||
| # putting both batches in the same training step might also break if we accidentally sample the same dataset index twice | ||
| split_and_insert_batch( | ||
| another_batch, | ||
| iter_dataloader.epoch_number, | ||
| training_step + i, | ||
| pending_queries_map, | ||
| param_prompt_Q, | ||
| generation_configs["train"], | ||
| is_eval=False, | ||
| ) | ||
| filtered_prompts_count -= args.num_unique_prompts_rollout | ||
| i += 0.1 |
There was a problem hiding this comment.
The use of training_step + i with a floating-point increment i to generate unique request identifiers is fragile and not best practice. This can lead to precision issues and subtle bugs, especially if the step logic changes in the future. A more robust approach should be used for generating unique IDs.
For example, you could introduce a separate request counter that is initialized before the main training loop and incremented for each batch submitted to the queue. This would provide a simple, unique, and reliable integer identifier for each request.
open_instruct/grpo_fast.py
Outdated
| if prompts_kept_this_iter == 0 and fill_iteration > args.active_sampling_max_attempts: | ||
| logger.warning( | ||
| "[Active fill completions] Unable to collect non-zero advantage prompts in iteration %d; " | ||
| "set as max by args.active_fill_max_attempts, proceeding with existing batch of size %d.", | ||
| fill_iteration, | ||
| len(aggregated_responses), | ||
| ) | ||
| break |
There was a problem hiding this comment.
There appears to be a typo in the warning message. It references args.active_fill_max_attempts, but the relevant argument defined in the Args class is args.active_sampling_max_attempts. This could cause confusion when debugging.
| if prompts_kept_this_iter == 0 and fill_iteration > args.active_sampling_max_attempts: | |
| logger.warning( | |
| "[Active fill completions] Unable to collect non-zero advantage prompts in iteration %d; " | |
| "set as max by args.active_fill_max_attempts, proceeding with existing batch of size %d.", | |
| fill_iteration, | |
| len(aggregated_responses), | |
| ) | |
| break | |
| if prompts_kept_this_iter == 0 and fill_iteration > args.active_sampling_max_attempts: | |
| logger.warning( | |
| "[Active fill completions] Unable to collect non-zero advantage prompts in iteration %d; " | |
| "set as max by args.active_sampling_max_attempts, proceeding with existing batch of size %d.", | |
| fill_iteration, | |
| len(aggregated_responses), | |
| ) | |
| break |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
|
|
||
| # Ensure the effective dataset size is divisible by batch_size | ||
| self.effective_size = len(self.data) - (len(self.data) % batch_size) | ||
| self._update_effective_size() |
There was a problem hiding this comment.
Why don't we just remove this and always have batches of size batch_size?
Either by throwing away the extra data or refilling across epoch numbers.
open_instruct/grpo_fast.py
Outdated
|
|
||
| return batch | ||
|
|
||
| def exclude_indices(self, exclude_list: List) -> None: |
There was a problem hiding this comment.
Why would we use this?
There was a problem hiding this comment.
my bad, this is holdovers from my next PR which will remove resampling of batches with 90% solve rate
open_instruct/grpo_fast.py
Outdated
| stop_rate = sum(int(finish_reason == "stop") for finish_reason in result.finish_reasons) / len( | ||
| result.finish_reasons | ||
| ) | ||
| def combine_reward_metrics(metric_records: list[tuple[Dict[str, Any], int]]) -> Dict[str, Any]: |
There was a problem hiding this comment.
Can you add a docstring?
open_instruct/grpo_fast.py
Outdated
| result.finish_reasons | ||
| ) | ||
| def combine_reward_metrics(metric_records: list[tuple[Dict[str, Any], int]]) -> Dict[str, Any]: | ||
| if not metric_records: |
There was a problem hiding this comment.
I don't think we need this, everything else should be empty and not iterate if metric_records is empty
open_instruct/grpo_fast.py
Outdated
|
|
||
| buckets: Dict[str, list[tuple[Any, int]]] = {} | ||
| for metrics, weight in metric_records: | ||
| if not metrics: |
There was a problem hiding this comment.
Similarly, this is not needed as the for loop will not execute if metrics is empty
open_instruct/grpo_fast.py
Outdated
| if isinstance(sample_value, np.ndarray): | ||
| combined[key] = np.concatenate([np.asarray(value) for value, _ in records]) | ||
| elif isinstance(sample_value, (list, tuple)): | ||
| concatenated: list[Any] = [] |
There was a problem hiding this comment.
what about:
combined[key] = [x for value, _ in records for x in value]
open_instruct/grpo_fast.py
Outdated
| if total_weight == 0: | ||
| combined[key] = float(sample_value) | ||
| else: | ||
| combined[key] = sum(float(value) * weight for value, weight in records) / total_weight |
There was a problem hiding this comment.
what about:
weighted_sum = 0
total_weight = 0
for value, weight in records:
weighted_sum += float(value) * weight
total_weight += weight
combined[key] = weighted_sum / total_weight if total_weight < 1e-6 else sample_value
this way we 1) don't compare a sum of floats directly to 0 and 2) we avoid iterating over records twice.
open_instruct/grpo_fast.py
Outdated
| oe_eval_beaker_image: Optional[str] = None | ||
| """the docker image for evaluation for oe-eval""" | ||
| oe_eval_gpu_multiplier: Optional[int] = 1 | ||
| """gpu mulitplier for eval jobs""" |
open_instruct/grpo_fast.py
Outdated
| """the max generation length for evaluation for oe-eval""" | ||
| oe_eval_beaker_image: Optional[str] = None | ||
| """the docker image for evaluation for oe-eval""" | ||
| oe_eval_gpu_multiplier: Optional[int] = 1 |
There was a problem hiding this comment.
What does this do? Why would I want to multiply GPUs for eval jobs?
There was a problem hiding this comment.
also holdover from another PR
open_instruct/grpo_fast.py
Outdated
|
|
||
| active_sampling: bool = False | ||
| """Whether to refill the batch with *new prompts/completions* after filtering.""" | ||
| active_sampling_max_attempts: int = 3 |
There was a problem hiding this comment.
Do we actually want a max number of attempts? Let's have it wait indefinitely.
Note
Introduce active refill of filtered prompts with chunked aggregation, add truncated-completion masking and richer metrics, tweak iterator robustness, and adjust training/save logic.
stopoutputs.prompt_lengthsandresponse_lengthsfromaccumulate_inference_batchesand propagate through packing and utilization metrics.combine_reward_metrics,FilteredChunk, andAggregationStateto merge per-chunk data and metrics.fill_completions; addactive_fill_completionsandactive_fill_max_attemptsto control active refill.ShufflingIteratorto recompute effective size via_update_effective_size, update on state restore, and handle empty data withStopIteration.filtered_prompts_countand queue extra prompt batches (with fractional training_step) to backfill after filtering.eval_on_step_0is set.Written by Cursor Bugbot for commit 79282ce. This will update automatically on new commits. Configure here.