[WIP]Support ray

[tastelikefeet]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

1. Support ray
2. Refactor sampling to support multi-node
3. Support yaml config file

Experiment results

Paste your experiment result here(if needed).

[gemini-code-assist]

Summary of Changes

Hello @tastelikefeet, 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 significantly enhances the system's distributed computing capabilities by integrating Ray. The core purpose is to enable scalable execution of tasks like model sampling and reward model evaluation across multiple devices and nodes. This is achieved through new configuration options, a dedicated argument class for Ray, and the strategic use of Ray worker and function decorators within the sampling pipeline, allowing for efficient resource allocation and parallel processing.

Highlights

• Ray Integration: Introduced comprehensive support for distributed execution using Ray, enabling the system to leverage Ray's capabilities for managing and orchestrating distributed tasks.
• Configuration Management: Added a new YAML configuration mechanism that allows defining Ray-specific parameters, including device groups and worker assignments, which are then parsed and passed as command-line arguments.
• Distributed Sampling and Reward Models: Modified the sampling pipeline to distribute generation, ORM (Output Reward Model), and PRM (Prompt Reward Model) tasks across Ray workers, improving scalability and resource utilization.
• New Ray Module: Created a dedicated swift/ray module with RayHelper and ResourceManager classes to abstract Ray's distributed functionalities, including worker and function decorators, and placement group management.
• Dependency Update: Added omegaconf as a new dependency to facilitate flexible configuration parsing from YAML files.

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 Ray support to the Swift framework, enabling distributed execution for sampling tasks. The changes include adding new arguments for Ray configuration, modifying the CLI to handle these arguments, and integrating Ray into the sampling pipeline. The review focuses on correctness and potential issues arising from the integration of Ray.

[gemini-code-assist]

This loop iterates through argv[1:], but accesses elements using argv[i] and argv[i + 1]. This will cause an IndexError when i is the last index in the loop, as argv[i + 1] will be out of bounds. The loop should iterate through the indices of argv directly.

To fix this, iterate through range(1, len(argv), 2) and adjust the indexing accordingly.

Suggested change

	for i in range(0, len(argv[1:]), 2):
	arg_name = argv[i]
	arg_value = argv[i + 1]
	if arg_name == '--config':
	for i in range(1, len(argv), 2):
	arg_name = argv[i]
	if i + 1 < len(argv):
	arg_value = argv[i + 1]
	else:
	break # Handle the case where there is no value for the last argument

[gemini-code-assist]

The device_groups attribute is loaded as a JSON string, but there's no error handling if the string is malformed. If the string is not a valid JSON, json.loads will throw an exception, crashing the program. Add a try-except block to handle potential json.JSONDecodeError exceptions.

Suggested change

	if isinstance(self.device_groups, str):
	self.device_groups = json.loads(self.device_groups)
	try:
	self.device_groups = json.loads(self.device_groups)
	except json.JSONDecodeError:
	print("Error decoding device_groups JSON string.")
	self.device_groups = None # or some default value

[gemini-code-assist]

The prepare_sampler and read_cache methods are called immediately after the __init__ method, but they are also decorated with @RayHelper.function. This means that when running in a Ray worker, these methods will not be executed, as the RayHelper.function decorator prevents execution within the worker. This could lead to unexpected behavior, as the sampler might not be properly initialized in the Ray workers.

Consider removing the @RayHelper.function decorator from these methods, or ensure that the initialization logic is correctly handled within the Ray worker context.

Suggested change

	self.prepare_sampler()
	self.caches = self.read_cache()
	def __init__(self, *args, **kwargs):
	super().__init__(*args, **kwargs)
	self.prepare_sampler()
	self.caches = self.read_cache()