[fix] Return origin_samples instead of False in filter_long_prompt

[guapisolo]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @guapisolo, I'm Gemini Code Assist¹! 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 introduces several key improvements across distributed training, evaluation, and utility functions. It addresses a critical type-related bug in prompt filtering, enhances the flexibility of parameter quantization ignore rules, and significantly improves the robustness of distributed weight updates through better synchronization. Furthermore, it adds a powerful feature for specifying custom generation functions on a per-sample basis, which is crucial for advanced evaluation setups. Minor adjustments to logging and configuration validation also contribute to overall system stability and usability.

Highlights

• Bug Fix: filter_long_prompt Return Type: The filter_long_prompt function now correctly returns the original list of samples (origin_samples) instead of False when no filtering is applied (e.g., max_length is None or prompt is not a string), aligning with its type annotation.
• Enhanced Quantization Ignore Rules: The logic for ignoring parameters during quantization has been extended to include prefix-based matching (name.startswith(r)), providing more flexible control over which parameters are quantized.
• Improved Distributed Weight Update Synchronization: Synchronization mechanisms for distributed weight updates have been refined. This includes reordering the pausing and continuing of generation engines and adding dist.barrier calls to ensure proper state management and data consistency, especially for quantized models.
• Per-Sample Custom Generation Functions: New functionality has been introduced to allow specifying custom generation functions on a per-sample basis. This is supported by adding generate_function_path to the Sample class and EvalDatasetConfig, enabling more dynamic and flexible evaluation scenarios like tool calling.
• Pipeline Model Parallelism Validation: A new assertion has been added to validate that decoder_first_pipeline_num_layers and decoder_last_pipeline_num_layers are None when pipeline_model_parallel_size is 1, preventing misconfigurations when pipeline parallelism is not in use.
• Log Readability Improvement: Sample labels in rollout logging messages are now truncated to the first 100 characters, improving log readability by preventing excessively long label outputs.

Changelog

• miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py
  • Expanded parameter ignore rules in quantize_params_compressed_tensors to include prefix matching (name.startswith(r)).
• miles/backends/megatron_utils/update_weight/update_weight_from_distributed.py
  • Reordered ray.get([engine.continue_generation.remote()]) to occur after post_process_weights.
  • Added dist.barrier(group=get_gloo_group()) before the rank 0 specific post-processing block.
• miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py
  • Added ray.get([engine.pause_generation.remote()]) before flush_cache.remote() for rank 0.
  • Introduced dist.barrier(group=get_gloo_group()) before the int4/fp4 post_process block.
  • Moved ray.get([engine.continue_generation.remote()]) to after post_process_weights for rank 0.
• miles/rollout/sglang_rollout.py
  • Added logic to update sample.loss_mask during generation when partial rollout and off-policy masking are enabled.
  • Modified generate_and_rm to prioritize sample.generate_function_path over args.custom_generate_function_path.
  • Truncated sample.label to 100 characters in logger info messages for generate_rollout_async.
  • Assigned dataset_cfg.custom_generate_function_path to sample.generate_function_path in eval_rollout_single_dataset.
• miles/utils/arguments.py
  • Added an assertion to ensure decoder_first_pipeline_num_layers and decoder_last_pipeline_num_layers are None when pipeline_model_parallel_size is 1.
• miles/utils/data.py
  • Fixed filter_long_prompt to return origin_samples instead of False when max_length is None or prompt is not a string.
• miles/utils/eval_config.py
  • Added custom_generate_function_path: str | None field to EvalDatasetConfig.
• miles/utils/types.py
  • Added generate_function_path: str | None field to the Sample class.

Activity

• No specific activity (comments, reviews, or progress updates) has been recorded for this pull request yet.

Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request correctly fixes a bug in filter_long_prompt where it was returning False instead of a list of samples, which violated the function's type hint. It also introduces a valuable feature for specifying a custom generation function on a per-sample basis, enhancing flexibility for evaluations. Furthermore, it includes several important correctness and synchronization fixes for the weight update process, especially concerning quantization and distributed operations. The changes are well-implemented and improve both the robustness and capabilities of the codebase. My review includes suggestions to handle empty list inputs gracefully, optimize performance by caching loaded functions, and a micro-optimization for rule checking.

[gemini-code-assist]

This function does not handle the case where origin_samples is an empty list. Accessing origin_samples[0] later on line 85 would raise an IndexError. It's safer to handle this case at the beginning of the function.

Suggested change

	if max_length is None:
	return False
	return origin_samples
	if max_length is None or not origin_samples:
	return origin_samples

[gemini-code-assist]

For a slight performance improvement, it's good practice to place cheaper, more common checks before more expensive ones like regex matching. Due to short-circuit evaluation of or, this can avoid some unnecessary startswith('re:') and re.match calls.

Suggested change

	(r.startswith("re:") and re.match(r[3:], name)) or r == name or name.startswith(r) for r in ignore_rules
	r == name or name.startswith(r) or (r.startswith("re:") and re.match(r[3:], name)) for r in ignore_rules