Add scripts to search context length capacity on given settings.

[chenyushuo]

Description

As the title says.

Checklist

Please check the following items before code is ready to be reviewed.

• Code has passed all tests
• Docstrings have been added/updated in Google Style
• Documentation has been updated
• Code is ready for review

[gemini-code-assist]

Summary of Changes

Hello @chenyushuo, 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 a robust framework for automatically assessing the context length capacity of large language models under various hardware and configuration constraints. By systematically increasing context length and monitoring for OOM conditions, it provides critical insights into model scalability and memory efficiency, which is essential for optimizing model deployment and performance.

Highlights

• Automated Context Length Testing Script: Introduced a new Python script, search_context_length_capacity.py, designed to automatically determine the maximum context length a large language model can handle before encountering Out-of-Memory (OOM) errors.
• Configurable Testing Environment: The testing framework is highly configurable, allowing for exploration of different GPU counts, sequence parallelism settings, model paths, and other training parameters via environment variables and a dedicated YAML configuration file.
• Iterative OOM Detection: The script iteratively increases the context length, runs training jobs, and monitors the output for OOM errors, providing a systematic way to find the capacity limits of models.
• Dummy Experience Workflow: A new dummy_exp_workflow.py has been added to generate synthetic experiences, enabling the context length tests to run without requiring real dataset processing.

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]

Summary of Changes

Hello @chenyushuo, 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 establishes a robust and automated framework for systematically evaluating the context length scalability of large language models. It provides a Python script that orchestrates distributed training jobs, incrementally increasing context lengths until Out-of-Memory errors are triggered. This system is designed to help researchers and engineers understand and optimize model performance under varying hardware and parallelism settings, utilizing a dedicated dummy workflow to isolate memory usage related to context length, and logging all experimental outcomes for analysis.

Highlights

• Automated Context Length Testing: Introduces a Python script (search_context_length_capacity.py) to automatically determine the maximum context length a large language model can handle before encountering Out-of-Memory (OOM) errors under various distributed training configurations.
• Configurable Test Parameters: The testing script allows for extensive configuration, including specifying the model path, GPU counts, sequence parallel sizes, starting context length, and different training strategies.
• Dynamic ROPE Scaling: The system intelligently applies Rope Scaling configurations when the tested context length exceeds the model's original maximum position embeddings, enabling testing beyond inherent limits.
• Dummy Experiment Workflow: A new dummy_exp_workflow.py is added, providing a synthetic training data generator. This allows for memory footprint testing focused purely on context length, without the overhead of real data processing.
• Trinity Configuration Integration: A new YAML configuration file (context_length.yaml) is included, serving as a template for the trinity run command, which orchestrates the experiments launched by the testing script.

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 introduces new scripts for testing the context length capacity of models. The main script, search_context_length_capacity.py, automates the process of finding the maximum context length by running experiments with increasing lengths. My review focuses on improving the security, robustness, and maintainability of this new script. Key suggestions include refactoring the use of subprocess.Popen to avoid security risks associated with shell=True, replacing rm -rf with a safer Python alternative, and improving the structure of the YAML configuration to reduce duplication. I've also pointed out areas for documentation improvement and code style enhancements.

[gemini-code-assist]

Code Review

This pull request introduces scripts for automated context length capacity testing. The changes include a new configuration file, a Python script to run the tests, and a dummy workflow. My review focuses on improving the security and maintainability of the new script.

I've identified several high-severity security vulnerabilities in search_context_length_capacity.py:

• The use of rm -rf is unsafe and can be replaced with a safer Python-native alternative.
• The use of shell=True in subprocess.Popen poses a command injection risk.
• There's a path traversal vulnerability when constructing log file paths.

Additionally, I've suggested a refactoring in context_length.yaml to reduce configuration duplication, which will improve maintainability.

Please address these points to make the script more robust and secure.