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.
| 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
| 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""" |
|
|
||
| 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 sampling to refill batches after filtering, refactor data preparation into aggregated chunks with safer iterators and enriched metrics, and adjust save/eval flow.
args.active_samplingandargs.active_sampling_max_attempts; removeargs.fill_completions.data_preparation_threadto iteratively accumulate, score, filter (zero-std), optionally refill with new prompts, aggregate chunks, and pack sequences.prompt_lengths/response_lengths.accumulate_inference_batchesnow returns(GenerationResult, Batch, prompt_lengths, response_lengths)and preservesindicesfield.combine_reward_metrics,FilteredChunk, andAggregationStateto merge multi-iteration results and stats.ShufflingIteratornow guards empty data, centralizes effective-size computation, and updates it on state restore.val/total_filtered_groupsto submit extra prompt batches when filtering shrinks data; avoid step collisions via fractionaltraining_stepoffsets.eval_on_step_0is true.Written by Cursor Bugbot for commit f8181e5. This will update automatically on new commits. Configure here.