diff --git a/.gitignore b/.gitignore index 88569d58cc..b1912e65ce 100644 --- a/.gitignore +++ b/.gitignore @@ -177,3 +177,5 @@ wandb/ outputs/ tests/ local/ +**/rollout_data/ +**/buffer_stats/ \ No newline at end of file diff --git a/docs/en/agent_training.md b/docs/en/agent_training.md new file mode 100644 index 0000000000..0941ee83b0 --- /dev/null +++ b/docs/en/agent_training.md @@ -0,0 +1,362 @@ +# Agent Rollout Usage Documentation + +### Starting Agent RL Training + +First, you need to configure the Slime runtime environment according to the [Readme](../../README.md) documentation and cd to the Slime project directory. + +#### Download Dataset and Model + +```bash + +huggingface-cli download --repo-type dataset zhuzilin/dapo-math-17k \ + --local-dir /root/dapo-math-17k + +# Ensure you are in the Slime root directory + +python ./slime_plugins/rollout_buffer/tools/assign_instance_id.py --input_path /root/dapo-math-17k/dapo-math-17k.jsonl + +# Download model +mkdir /root/hf_models/ +mkdir /root/megatron_model/ + +huggingface-cli download deepseek-ai/DeepSeek-R1-Distill-Qwen-7B --local-dir /root/hf_models/deepseek-ai--DeepSeek-R1-Distill-Qwen-7B --local-dir-use-symlinks False + +# Convert model +PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ + --hf-checkpoint /root/hf_models/deepseek-ai--DeepSeek-R1-Distill-Qwen-7B \ + --save /root/megatron_model/DeepSeek-R1-Distill-Qwen-7B-25.02 +``` + +#### Start slime +```bash +chmod +x ./scripts/run_agent.sh +./scripts/run_agent.sh +``` + +## Overview + +Agent Rollout is a specialized module in the Slime framework for agent task data generation, implementing a **fully asynchronous** data generation workflow based on external APIs. Unlike traditional reinforcement learning Rollout, Agent Rollout generates training data by interacting with external agent services (Rollout Buffer component in Slime Plugins). + +**The main workflow is as follows:** + +1. Slime sends a request to start Rollout to the Rollout Buffer via HTTP request +2. After receiving the request, Rollout Buffer starts executing Rollout tasks by accessing the SGLang Server launched in Slime +3. After Rollout Buffer begins executing Rollout tasks, Slime periodically calls the `/get_rollout_data` interface to obtain corresponding Rollout data and log information +4. When the amount of obtained data reaches the preset threshold, this data is stored in Slime's Data Buffer +5. Slime uses this data for model training + +## Important Notes + +- **Buffer Terminology Distinction**: "Rollout Buffer" in the documentation refers to the Buffer component in Slime Plugins, while "Slime Buffer" or "Data Buffer" refers to the data buffer within the Slime framework. +- **Parameter Passing Mechanism**: Start Rollout passes all parameter configurations from Slime to Rollout Buffer, and the specific Start Rollout implementation logic needs to be completed in Rollout Buffer. +- **Data Retrieval Retry Mechanism**: Get Rollout Data retrieves data from the Rollout Buffer port at fixed time intervals, and if errors occur, the retry count is increased. The retry count limit can be set through the `--fetch-trajectory-retry-times` parameter, with a default value of -1, indicating unlimited retries. +- **Return Data Format**: The default implementation requires that the return value of Get Rollout Data must contain multiple specific fields. Although this can be customized, it is strongly recommended to retain these key fields. +- **Unique Identifier**: Instance ID serves as a unique identifier for different data (similar to a primary key), typically of string type. +- **Data Format Conversion**: In the `generate_agent_rollout` function, when converting Rollout Data to Slime's internal storage universal format (Sample), the `prompt` field is not used during Agent Training. Since subsequent data sorting needs to be performed based on Index, it is essential to ensure that the Instance ID field is assigned to the Index field in Sample. +- **Statistical Information Recording**: When obtaining Rollout Data, in addition to basic Message, Reward and other results, statistical information can also be returned (such as data count before filtering, average reward before filtering, etc.). This additional information can be customized in the Log Raw Info function. +- **Log Prefix Convention**: When recording logs in this file, please use `rollout/` as a prefix. If you need to add new prefixes, you need to modify the initialization part accordingly, otherwise it may cause log information anomalies. +- **Asynchronous Training Characteristics**: Since Agent Training adopts a pure asynchronous mode, it may cause initially high data Reward (because the earliest generated data will be prioritized for training), which is normal. + +## Core Components + +### Main Functions + +- **`generate_agent_rollout`**: Core data generation function, responsible for coordinating the entire Rollout process. It sequentially calls Start Rollout, Get Rollout Data, and Log Raw Info functions, and returns results after adding data to Slime Buffer. +- **`get_rollout_data`**: Retrieves generated data from Rollout Buffer, polling at fixed time intervals until sufficient data is obtained or the retry limit is reached. +- **`start_rollout`**: Starts Rollout tasks for external agent services. Slime is responsible for sending parameter configurations (including sampling parameters, Server URL, etc.) and notifying Rollout Buffer to start Rollout tasks. +- **`log_raw_info`**: Records and statistics meta-information during the Rollout process. Since different tasks and users have different logging requirements, this function can record required information in WandB, with specific information sourced from the return results of Get Rollout Data. + +## Overall Workflow + +### 1. Initialization Phase + +When this round is the first training (if it's resume training, it will read previously trained data IDs, and then Rollout Buffer will skip these data), Agent Rollout will: +- Call `start_rollout()` to send a startup request to Rollout Buffer +- Pass necessary configuration parameters, including model URL, task type, input files, etc. (ensure file paths are accessible to Rollout Buffer) +- Rollout Buffer starts executing data generation tasks + +### 2. Data Collection Phase + +```python +while len(results) < args.rollout_batch_size * args.n_samples_per_prompt: + time.sleep(5) # polling interval + data, meta_info = await get_rollout_data(api_base_url=base_url) + results.extend(data) + if meta_info: + all_meta_info.append(meta_info) +``` + +- Periodically poll Rollout Buffer to retrieve generated data +- Accumulate data until reaching the expected batch size +- Collect meta-information for monitoring and logging + +### 3. Data Processing Phase + +```python +for record in results: + oai_messages = record["messages"] + + mask_generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type=args.loss_mask_type) + token_ids, loss_mask = mask_generator.get_loss_mask(oai_messages) + response_length = mask_generator.get_response_lengths([loss_mask])[0] + + sample = Sample( + index=record["instance_id"], + prompt=record["uid"], + tokens=token_ids, + response_length=response_length, + reward=record["reward"], + truncated=False, + loss_mask=loss_mask[-response_length:], + metadata={**record["extra_info"], "raw_reward": record["raw_reward"]} + ) +``` + +- Convert OpenAI format messages returned by external services to Token sequences +- Generate appropriate Loss Mask, supporting multi-turn conversations +- Create `Sample` objects that conform to Slime specifications + +## Adding Custom Agent Tasks + +If you need to customize Agent Rollout for specific agent tasks, you need to understand and possibly modify four core functions. The following details the purpose and interface specifications of each function: + +### Core Function Details + +#### 1. `generate_agent_rollout` - Main Control Function + +**Function Description**: Serves as the main entry function for Agent Rollout, coordinating the entire data generation process, managing interactions with external services and converting data formats. + +**Function Signature**: +```python +async def generate_agent_rollout( + args, + rollout_id: int, + data_buffer: Buffer, + evaluation: bool = False +) -> List[Sample] +``` + +**Input Parameters**: +- `args`: Object containing all configuration parameters +- `rollout_id`: Unique identifier for the current Rollout, indicating how many times Rollout has been performed, also used as identifier for saving Checkpoint and resume training +- `data_buffer`: Slime's global data buffer for storing and managing sampled data samples +- `evaluation`: Whether in evaluation mode (not supported in current version), recommend evaluating Checkpoint separately + +**Output**: +- `List[Sample]`: List of samples conforming to Slime specifications, each sample contains: + - `tokens`: Token sequence + - `response_length`: Response length (length after excluding Prompt Length), in multi-turn conversations it's the length from the first non-zero position of Loss Mask to the end + - `reward`: Reward value. All Samples obtained in Slime should ensure they are valid Samples, specific filtering operations need to be handled in Rollout Buffer + - `loss_mask`: Loss mask, currently only calculates loss for all Assistant type messages + - `metadata`: Metadata information, containing various custom information + +**Core Logic**: + +1. Determine whether to start a new Rollout task by checking the attributes of Start Rollout (note that you cannot use whether the data length in Slime Buffer is 0 to judge, because according to different Slime Buffer Filter rules, Slime Buffer may also be 0 during training) +2. Poll to get data generated by external services, implemented through Get Rollout Data function + +#### 2. `get_rollout_data` - Data Retrieval Function + +**Function Description**: Asynchronously retrieves generated data from Rollout Buffer, supporting batch retrieval and meta-information collection. + +**Function Signature**: +```python +async def get_rollout_data( + api_base_url: str, + num: Optional[int] = None, + timeout: float = 100.0 +) -> Tuple[List[Dict[str, Any]], Dict[str, Any]] +``` + +**Input Parameters**: +- `api_base_url`: URL of Rollout Buffer +- `num`: Optional, specifies the number of data items to retrieve, default is None, meaning retrieve all data +- `timeout`: Request timeout time (seconds) + +**Output**: +- `Tuple[List[Dict], Dict]`: Tuple containing data list and meta-information + - Each dictionary in the data list must contain: + ```python + { + "uid": "unique id", + "instance_id": "instance id", + "messages": [openai format messages], + "reward": 0.85, + "extra_info": {...} + } + ``` + - Meta-information dictionary supports customization, just needs to be consistent with subsequent logging operations + +**Error Handling**: +- Throws `aiohttp.ClientError` on network errors +- Throws `ValueError` on data format errors +- Throws `asyncio.TimeoutError` on timeout + +#### 3. `start_rollout` - Task Startup Function + +**Function Description**: Sends startup request to Rollout Buffer, configures task parameters and starts data generation process. + +**Function Signature**: +```python +def start_rollout(api_base_url: str, args, metadata) -> Dict[str, Any] +``` + +**Input Parameters**: +- `api_base_url`: URL of Rollout Buffer +- `args`: Parameter object containing all configurations +- `metadata`: Metadata of data buffer, containing list of completed instance IDs for resume training + +**Output**: +- `Dict[str, Any]`: Startup confirmation information returned by external service + +**Sent Payload Format Example**: +```python +{ + "num_process": 1024, # number of parallel processes + "num_epoch": 3, # generation epochs + "remote_engine_url": "http://sglang-router:port", + "task_type": "math", # task type + "input_file": "/path/to/input.jsonl", # input file path + "num_repeat_per_sample": "8", # number of repetitions per sample + "max_tokens": 8192, # maximum number of tokens + "sampling_params": { # sampling parameters + "max_tokens": 8192, + "temperature": 0.8, + "top_p": 0.9 + }, + "tokenizer_path": "/path/to/tokenizer", # tokenizer path + "skip_instance_ids": ["id1", "id2"] # instance IDs to skip, skip previously trained instance IDs during resume training +} +``` + +#### 4. `log_raw_info` - Logging Function + +**Function Description**: Collects, statistics and records key metrics during the Rollout process, supports WandB integration, content is fully customizable. + +**Function Signature**: +```python +def log_raw_info(args, all_meta_info: List[Dict], rollout_id: int) -> None +``` + +**Input Parameters**: +- `args`: Configuration parameter object, including log-related settings +- `all_meta_info`: List of meta-information collected from all data retrieval requests +- `rollout_id`: Current Rollout ID + +**Log Format Example**: +```python +{ + "rollout/no_filter/total_samples": 100, + "rollout/no_filter/avg_reward": 0.75, +} +``` + +## Custom Agent Task Implementation Steps + +### 1. Prepare Input Data Format + +Your input file (`--rollout-input-file`) should contain task-related data, for example: + +```jsonl +{"instance_id": "math_001", "prompt": [{"role":"user","content":"Solve the equation x^2 + 5x + 6 = 0"}]} +{"instance_id": "math_002", "prompt": [{"role":"user","content":"Prove Fermat's Last Theorem"}]} +``` + +Where `instance_id` is a required field, you can use `slime_plugins/rollout_buffer/tools/assign_instance_id.py` to automatically generate Instance ID for each data item. + +### 2. Custom Loss Mask Generator + +If your task has special conversation formats, you need to add a custom Mask generator in `slime/utils/mask_utils.py`: + +```python +class CustomTaskLossMaskGenerator(MultiTurnLossMaskGenerator): + def get_custom_multi_turn_loss_mask(self, messages): + # return (token_ids, loss_mask) + pass + + def get_loss_mask(self, messages: List[Dict]) -> List[int]: + # ... existing code ... + elif self.tokenizer_type == "custom_task": + return xxx + # ... existing code ... +``` + +Then use: +```bash +--loss-mask-type custom_task +``` + +Currently, Loss Mask generation for multi-turn conversations in Qwen and Distill Qwen conversation formats has been implemented. + +### 3. Implement External Service Interface + +In Rollout Buffer, you need to implement the key function `run_rollout()` for your task. For specific implementation, please refer to [Rollout Buffer Usage Documentation](./rollout_buffer_usage.md). + +## Parameter Configuration + +### Agent Rollout Core Parameters + +- **`--rollout-function-path`**: Specify using Agent Rollout function + ```bash + --rollout-function-path slime.rollout.agent_rollout.generate_rollout + ``` + +- **`--agent-rollout-buffer-url`**: API address of Rollout Buffer + ```bash + --agent-rollout-buffer-url http://0.0.0.0:8889 + ``` + For specific deployment methods, please check Rollout Buffer documentation + +- **`--rollout-input-file`**: Input data file path, containing task data to be processed + ```bash + --rollout-input-file ./deepscaler_rl_buffer_instance_id.jsonl + ``` + +- **`--rollout-num-process`**: Number of parallel processes + ```bash + --rollout-num-process 1024 + ``` + +- **`--loss-mask-type`**: Loss Mask type, used for Token masking in multi-turn conversations, supports customization, custom path is in `slime/utils/mask_utils.py` + ```bash + --loss-mask-type distill_qwen + ``` + +### Rollout Generation Parameters + +- **`--num-rollout`**: Total number of Rollouts (maximum Rollout ID count) + ```bash + --num-rollout 3000 + ``` + +- **`--rollout-batch-size`**: Batch size for each Rollout + ```bash + --rollout-batch-size 16 + ``` + +- **`--rollout-max-response-len`**: Maximum response length (single turn) + ```bash + --rollout-max-response-len 8192 + ``` + +- **`--rollout-temperature`**: Sampling temperature parameter + ```bash + --rollout-temperature 0.8 + ``` + +- **`--n-samples-per-prompt`**: Number of samples generated per Prompt + ```bash + --n-samples-per-prompt 8 + ``` + +### Filtering and Reward Configuration + +- **`--buffer-filter-path`**: Data buffer filter path, supports customization, defaults to using the latest batch of data in Buffer for updates + ```bash + --buffer-filter-path slime.rollout.filter_hub.buffer_filters.pop_first + ``` + +- **`--disable-rewards-normalization`**: Disable reward normalization, if rewards are already normalized in Rollout Buffer, please enable this parameter + ```bash + --disable-rewards-normalization + ``` \ No newline at end of file diff --git a/docs/en/rollout_buffer_usage.md b/docs/en/rollout_buffer_usage.md new file mode 100644 index 0000000000..84e2e5fb33 --- /dev/null +++ b/docs/en/rollout_buffer_usage.md @@ -0,0 +1,323 @@ +# Rollout Buffer Usage Documentation + +## Overview + +Rollout Buffer is an independent component in the Slime framework for agent trajectory generation, with the main function of using the LLM OpenAI Server launched by Slime training to generate agent trajectories. + +### Design Philosophy + +The main reasons we made Rollout Buffer independent include: + +1. **Framework Decoupling**: Different Agent tasks depend on different Agent Frameworks and tools, and are likely to reuse third-party Agent Frameworks +2. **Flexible Extension**: If all components are encapsulated within Slime, it would lead to architectural chaos and be unfavorable for extension and maintenance +3. **Responsibility Separation**: Rollout Buffer is only responsible for generating corresponding trajectories by calling the Server launched in Slime, with no restrictions on what framework is used +4. **Complete Decoupling**: Trajectory generation logic and Slime training logic are completely decoupled, supporting the introduction of various complex Agent Frameworks + +### Workflow + +``` +Slime Training Process ←─── HTTP API ───→ Rollout Buffer + ↓ ↓ + LLM Server ←─────── HTTP Requests ─────── Agent Framework + ↓ ↓ + Model Response ──────────────────────→ Trajectory Generation +``` + +For each different Agent task, there should be a corresponding independent Generator class, responsible for generating trajectories for that type of task. Rollout Buffer automatically reads and loads different types of Generators. + +## Quick Start + +### Basic Usage Process + +1. **Copy Template**: Copy `base_generator.py` as a template +2. **Modify Task Type**: Change `TASK_TYPE` to your task name (cannot duplicate with other Generators) +3. **Implement Core Function**: Implement the `run_rollout()` function +4. **Optional Customization**: Rewrite five optional functions as needed +5. **Start Training**: Follow the startup process in [Agent Training Documentation](./agent_training.md) to start Agent training + +### File Structure Standards + +Generator files must end with `_generator.py` and be placed in the `generator/` directory: + +``` +generator/ +├── base_generator.py # Math task implementation (default template) +└── your_task_generator.py # Your custom task +``` + +## Core Components + +### Required Components + +Each Generator file must contain the following components: + +#### 1. `TASK_TYPE` Constant +Define the unique identifier for the task type: +```python +TASK_TYPE = "your_task_name" +``` + +#### 2. `run_rollout()` Function +Entry function for core data generation logic: +```python +def run_rollout(data: dict): + # Implement your trajectory generation logic + pass +``` + +### Optional Components + +In addition to required components, Rollout Buffer also provides five customizable functions to meet special needs of different tasks. If no custom implementation is provided, the system will use default implementations (located in `slime_plugins/rollout_buffer/generator/utils/default_func.py`): + +1. **`normalize_group_data()`**: Reward normalization function +2. **`pad_group_data()`**: Data padding strategy function +3. **`is_valid_group()`**: Group data validity verification function +4. **`get_group_data_meta_info()`**: Meta-information statistics function +5. **`filter_item()`**: Individual data item filtering function + +## Parameter Configuration + +### Generator Core Parameters + +The main parameters received by the `run_rollout(data: dict)` function are as follows (the incoming `data` needs to be consistent with parameters sent from Slime): + +| Parameter Name | Type | Description | +|----------------|------|-------------| +| `remote_engine_url` | string | Inference engine service address, usually the SGLang Router address in Slime | +| `remote_buffer_url` | string | Rollout Buffer service address, usually a port on the Master node (default 8889) | +| `input_file` | string | Input data file path | +| `task_type` | string | Task type identifier, defined in each `_generator.py` file | +| `num_repeat_per_sample` | int | Number of repeated generations per sample (Group Size) | +| `num_epoch` | int | Number of dataset traversal rounds (default 10) | +| `sampling_params` | dict | Model sampling parameters (including max_tokens, temperature, etc.) | +| `num_process` | int | Number of parallel processes | +| `skip_instance_ids` | list | List of instance IDs to skip, used for resume training to skip previously processed instances | + +### Buffer Control Parameters + +Buffer behavior is controlled by the following key parameters, which directly affect data collection, validation, and output strategies: + +#### Core Control Parameters + +| Parameter Name | Default Value | Description | +|----------------|---------------|-------------| +| `group_size` | - | Target number of data items per group, usually equal to `num_repeat_per_sample` | +| `min_valid_group_size_ratio` | 1.0 | Minimum data ratio for a group to be considered "valid" (100%) | +| `min_valid_item_size_ratio` | 0.7 | Minimum ratio of valid data within a group after filtering (70%) | + +**Important Notes**: +- `group_size`: All data will eventually be padded to this size, directly affecting the number of samples per instance during training +- `min_valid_group_size_ratio`: Recommended to set to 1.0, invalid data can also be written and filtered through subsequent steps (such as assigning extreme Rewards) +- `min_valid_item_size_ratio`: Minimum ratio of valid data within a group after filtering, should be greater than 0.5, used to filter groups with poor quality + +#### Timeout Control Parameters + +| Parameter Name | Default Value | Description | +|----------------|---------------|-------------| +| `group_timeout_seconds` | 300 | Group timeout time (5 minutes), prevents some groups from being stuck for a long time | +| `min_timeout_group_size_ratio` | 0.7 | Minimum data ratio threshold for timeout groups (70%) | + +#### System Resource Parameters + +| Parameter Name | Default Value | Description | +|----------------|---------------|-------------| +| `max_buffer_size` | 1,000,000,000 | Maximum Buffer capacity (1 billion), prevents memory overflow | + +## Data Processing Flow + +### Complete Processing Flow + +When retrieving a batch of training data from Rollout Buffer, five optional functions execute in the following fixed order: + +``` +buffer.read(batch_size) call + ↓ +1. 📊 get_group_data_meta_info() + └── Collect statistics (progress, reward distribution, etc.) + ↓ +2. ✅ is_valid_group() + └── Determine if each group is complete and valid + ↓ +3. 🔍 filter_item() + └── Filter each data item in valid groups + ↓ +4. ⚖️ normalize_group_data() + └── Perform reward normalization on filtered group data + ↓ +5. 📦 pad_group_data() + └── Pad normalized data to target group_size + ↓ +📤 Return processed batch data +``` + +### Processing Step Details + +#### Step 1: Meta-information Statistics - `get_group_data_meta_info()` + +**Function**: Collect statistical information for all raw group data in the current Buffer +- **Input**: All raw group data in Buffer (including invalid groups and invalid trajectories) +- **Output**: Dictionary containing statistical information, used for logging and monitoring, such as recording average rewards and other information + +#### Step 2: Group Validity Verification - `is_valid_group()` + +**Function**: Determine which groups can be used for training +- **Input**: Complete data for each group `(instance_id, group_data)` +- **Output**: `(is_valid, is_finished)` tuple +- **Logical Relationship**: `Valid Groups ⊆ Finished Groups ⊆ All Groups`, where instances in finished groups will be skipped during resume training, and qualified groups in valid groups will be used for model training + +#### Step 3: Individual Data Filtering - `filter_item()` + +**Function**: Perform fine-grained filtering on each data item within valid groups +- **Input**: Individual data items within groups +- **Output**: Boolean value determining whether the item should be retained, as data written to Rollout Buffer may contain invalid items that need to be filtered + +#### Step 4: Reward Normalization - `normalize_group_data()` + +**Function**: Perform standardized processing on group reward values +- **Note**: If normalization is performed here, reward normalization needs to be disabled in Slime. The default implementation normalizes only valid data items and performs scaling +- **Other**: Original reward values are saved to the `raw_reward` field for convenient logging + +#### Step 5: Data Padding - `pad_group_data()` + +**Function**: Pad data to standard `group_size` +- **Strategy**: Maintain total reward consistency through reward scaling +- **Output**: Fixed-size group data that can be directly used for training +- **Note**: The returned data count **must** be a multiple of Group Size + +### Important Mechanism Descriptions + +#### Data Storage Strategy +- **Full Storage**: All data should be stored in Buffer regardless of whether trajectory generation is successful +- **Subsequent Filtering**: Filter out useful Groups and Items through filtering mechanisms +- **Failure Handling**: Assign special Reward values to failed trajectories for easy identification + +#### Timeout Cleanup Mechanism +- **Automatic Cleanup**: Check timestamps each time `get_rollout_data` is executed +- **Decision Logic**: Timeout groups decide whether to retrieve or discard based on valid data count +- **Prevent Accumulation**: Effectively prevent excessive data accumulation in Buffer + +## Implementation Examples + +### Basic Implementation Template + +Using Math task as an example, showing complete Generator implementation: + +```python +TASK_TYPE = "math" + +def run_rollout(data: dict): + + print(f"Starting math rollout with data: {data}") + + rollout_func = query_single_turn + reward_func = get_rule_based_math_reward + + print(f"Waiting for 10 seconds for buffer server to start") + time.sleep(10) + global SAMPLING_PARAMS + for k, v in data["sampling_params"].items(): + SAMPLING_PARAMS[k] = v + print(f"Set {k} to {v}", type(v)) + + generator = BaseGenerator( + data["remote_engine_url"], + data["remote_buffer_url"], + num_repeat_per_sample=int(data["num_repeat_per_sample"]), + queue_size=1000000, + max_tokens=int(data["sampling_params"]["max_tokens"]), + num_process=int(data.get("num_process", 100)), + task_type=data["task_type"], + skip_instance_ids=data.get("skip_instance_ids", None), + ) + + generator.entry(data["input_file"], rollout_func, reward_func, int(data.get("num_epoch", 1))) + + +``` + +### Trajectory Generation Function Example + +```python +def query_single_turn(client, messages, sampling_params, tools=None): + base_payload = { + "messages": messages, + **sampling_params, + "model": "custom", + "stream": False, + "seed": random.randint(1, 10000000), + "tools": tools, + } + + text = None + accumulated_tokens = 0 + + for attempt in range(6): + try: + # Create a fresh payload for each attempt + current_payload = copy.deepcopy(base_payload) + + if text is not None: + # Update messages with current progress + current_messages = copy.deepcopy(messages) + current_messages.append({"role": "assistant", "content": text}) + current_payload["messages"] = current_messages + + # Adjust max_tokens based on accumulated tokens + if "max_tokens" in sampling_params: + current_payload["max_tokens"] = max(0, sampling_params["max_tokens"] - accumulated_tokens) + + # Add continue flag for partial rollouts + current_payload["extra_body"] = {"continue_final_message": True} + if current_payload["max_tokens"] == 0: + break + response = client.chat.completions.create(**current_payload) + + if len(response.choices) > 0: + if response.choices[0].finish_reason == "abort": + print( + f"query failed, reason: {response.choices[0].finish_reason}, currently generated: {response.usage.completion_tokens}" + ) + + accumulated_tokens += response.usage.completion_tokens + + if text is None: + text = response.choices[0].message.content + else: + text += response.choices[0].message.content + + sleep(10) + continue + if text is None: + text = response.choices[0].message.content + elif response.choices[0].message.content is not None: + text += response.choices[0].message.content + break + else: + print(f"Error in query, status code: {response.status_code}") + continue + except Exception as e: + print(f"query failed in single turn, error: {e}") + continue + + # Update final messages + if len(messages) > 0 and messages[-1]["role"] == "assistant": + messages = messages[:-1] + messages.append({"role": "assistant", "content": text}) + + return messages + +``` + +## Frequently Asked Questions + +### Q: How to handle failed generation data? +A: Store failed data in Buffer as well, but assign special Reward values (such as -1), and handle through subsequent filtering mechanisms. + +### Q: How to debug data quality issues? +A: Use the `get_group_data_meta_info()` function to collect detailed statistical information and monitor reward distribution and data quality. + +### Q: How does the timeout mechanism work? +A: When a group's last data generation time exceeds `group_timeout_seconds`, the system will decide whether to use that group's data based on `min_timeout_group_size_ratio`. + +### Q: How to implement resume training? +A: Slime will pass a list of processed instance IDs through the `skip_instance_ids` parameter, and the Generator will automatically skip these instances. All completed groups will be automatically skipped. \ No newline at end of file diff --git a/docs/zh/agent_training.md b/docs/zh/agent_training.md new file mode 100644 index 0000000000..c37efd53dd --- /dev/null +++ b/docs/zh/agent_training.md @@ -0,0 +1,364 @@ +## Agent Rollout 使用文档 + + +### 启动 Agent RL Training + +首先需要根据[Readme](../../README_zh.md)文档中配置好 Slime的运行环境并且 cd 到 Slime项目的目录下。 + +#### 下载数据集和模型 + +```bash + +huggingface-cli download --repo-type dataset zhuzilin/dapo-math-17k \ + --local-dir /root/dapo-math-17k + +#确保在 Slime 根目录下 + +python ./slime_plugins/rollout_buffer/tools/assign_instance_id.py --input_path /root/dapo-math-17k/dapo-math-17k.jsonl + +# 下载模型 +mkdir /root/hf_models/ +mkdir /root/megatron_model/ + +huggingface-cli download deepseek-ai/DeepSeek-R1-Distill-Qwen-7B --local-dir /root/hf_models/deepseek-ai--DeepSeek-R1-Distill-Qwen-7B --local-dir-use-symlinks False + +# 转换模型 +PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ + --hf-checkpoint /root/hf_models/deepseek-ai--DeepSeek-R1-Distill-Qwen-7B \ + --save /root/megatron_model/DeepSeek-R1-Distill-Qwen-7B-25.02 +``` + +#### 启动slime +```bash +chmod +x ./scripts/run_agent.sh +./scripts/run_agent.sh +``` +### 流程概述 + +Agent Rollout 是 Slime 框架中用于智能体任务数据生成的专用模块，实现了基于外部 API 的**完全异步**数据生成流程。与传统的强化学习 Rollout 不同，Agent Rollout 通过与外部智能体服务（Slime Plugins 中的 Rollout Buffer 组件）进行交互来生成训练数据。 + +**主要工作流程如下：** + +1. Slime 通过 HTTP 请求向 Rollout Buffer 发送开始 Rollout 的指令 +2. Rollout Buffer 接收到请求后，通过访问 Slime 中启动的 SGLang Server 开始执行 Rollout 任务 +3. 在 Rollout Buffer 开始执行 Rollout 任务后，Slime 会定期调用 `/get_rollout_data` 接口获取对应的 Rollout 数据和日志信息 +4. 当获取的数据量达到预设阈值时，这些数据会被存储到 Slime 的 Data Buffer 中 +5. Slime 使用这些数据进行模型训练 + +### 重要说明 + +- **Buffer 术语区分**：文档中的 "Rollout Buffer" 指 Slime Plugins 中的 Buffer 组件，而 "Slime Buffer" 或 "Data Buffer" 指 Slime 框架内部的数据缓冲区。 +- **参数传递机制**：Start Rollout 会将 Slime 的所有参数配置传递给 Rollout Buffer，具体的 Start Rollout 实现逻辑需要在 Rollout Buffer 中完成。 +- **数据获取重试机制**：Get Rollout Data 按固定时间间隔从 Rollout Buffer 端口获取数据，如果发生错误则会增加重试次数。重试次数上限可通过 `--fetch-trajectory-retry-times` 参数设置，默认值为 -1，表示无限重试。 +- **返回数据格式**：默认实现要求 Get Rollout Data 的返回值必须包含多个特定的字段，虽然可以自定义修改，但强烈建议保留这几个关键字段。 +- **唯一标识符**：Instance ID 作为不同数据的唯一标识符（类似主键），通常为字符串类型。 +- **数据格式转换**：在 `generate_agent_rollout` 函数中，将 Rollout Data 转换为 Slime 内部存储的通用格式（Sample）时，`prompt` 字段在 Agent Training 过程中不会被使用。由于后续需要根据 Index 对数据进行排序，因此必须确保将 Instance ID 字段赋值给 Sample 中的 Index 字段。 +- **统计信息记录**：获取 Rollout Data 时，除了基本的 Message、Reward 等结果外，还可以返回统计信息（如过滤前的数据数量、过滤前的奖励均值等），这些额外信息可以在 Log Raw Info 函数中自定义记录方式。 +- **日志前缀规范**：在该文件中记录日志时，请使用 `rollout/` 作为前缀。如需添加新的前缀，需要相应修改初始化部分，否则可能导致日志信息异常。 +- **异步训练特性**：由于 Agent Training 采用纯异步模式，可能导致最初的数据 Reward 较高（因为最早生成完成的数据会优先用于训练），这是正常现象。 + +### 核心组件 + +#### 主要函数 + +- **`generate_agent_rollout`**：核心数据生成函数，负责协调整个 Rollout 流程。依次调用 Start Rollout、Get Rollout Data 和 Log Raw Info 函数，并将数据加入 Slime Buffer 后返回结果。 +- **`get_rollout_data`**：从 Rollout Buffer 获取生成的数据，按固定时间间隔轮询，直到获取足够数据或达到重试次数上限。 +- **`start_rollout`**：启动外部智能体服务的 Rollout 任务。Slime 负责发送参数配置（包括采样参数、Server URL 等），通知 Rollout Buffer 开始 Rollout 任务。 +- **`log_raw_info`**：记录和统计 Rollout 过程中的元信息。由于不同任务和用户的日志需求不同，可通过该函数在 WandB 中记录所需信息，具体信息来源于 Get Rollout Data 的返回结果。 + +### 整体流程 + +#### 1. 初始化阶段 + +当该轮次为第一次训练时（如果是续训将会读取之前训练过的数据 ID，随后 Rollout Buffer 会跳过这些数据），Agent Rollout 会： +- 调用 `start_rollout()` 向 Rollout Buffer 发送启动请求 +- 传递必要的配置参数，包括模型 URL、任务类型、输入文件等（确保文件路径是 Rollout Buffer 可以访问的） +- Rollout Buffer 开始执行数据生成任务 + +#### 2. 数据收集阶段 + +```python +while len(results) < args.rollout_batch_size * args.n_samples_per_prompt: + time.sleep(5) # 轮询间隔 + data, meta_info = await get_rollout_data(api_base_url=base_url) + results.extend(data) + if meta_info: + all_meta_info.append(meta_info) +``` + +- 定期轮询 Rollout Buffer 获取生成的数据 +- 累积数据直到达到预期的批次大小 +- 收集元信息用于监控和日志记录 + +#### 3. 数据处理阶段 + +```python +for record in results: + oai_messages = record["messages"] + + mask_generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type=args.loss_mask_type) + token_ids, loss_mask = mask_generator.get_loss_mask(oai_messages) + response_length = mask_generator.get_response_lengths([loss_mask])[0] + + sample = Sample( + index=record["instance_id"], + prompt=record["uid"], + tokens=token_ids, + response_length=response_length, + reward=record["reward"], + truncated=False, + loss_mask=loss_mask[-response_length:], + metadata={**record["extra_info"], "raw_reward": record["raw_reward"]} + ) +``` + +- 将外部服务返回的 OpenAI 格式消息转换为 Token 序列 +- 生成适当的 Loss Mask，支持多轮对话 +- 创建符合 Slime 规范的 `Sample` 对象 + +### 添加自定义 Agent Task + +如果您需要为特定的智能体任务定制 Agent Rollout，需要理解并可能修改四个核心函数。以下详细说明每个函数的目的和接口规范： + +#### 核心函数详解 + +##### 1. `generate_agent_rollout` - 主控制函数 + +**功能描述**：作为 Agent Rollout 的主入口函数，协调整个数据生成流程，管理与外部服务的交互并转换数据格式。 + +**函数签名**： +```python +async def generate_agent_rollout( + args, + rollout_id: int, + data_buffer: Buffer, + evaluation: bool = False +) -> List[Sample] +``` + +**输入参数**： +- `args`: 包含所有配置参数的对象 +- `rollout_id`: 当前 Rollout 的唯一标识符，表示当前是第几次 Rollout，同时也会被用作保存 Checkpoint 和续训的标识符 +- `data_buffer`: Slime 的全局数据缓冲区，用于存储和管理采样的数据样本 +- `evaluation`: 是否为评估模式（当前版本不支持），建议单独对 Checkpoint 进行评估 + +**输出**： +- `List[Sample]`: 符合 Slime 规范的样本列表，每个样本包含： + - `tokens`: Token 序列 + - `response_length`: 响应长度（长度是剔除 Prompt Length 后的总长度），在多轮对话中为 Loss Mask 第一个非零位置到最后的长度 + - `reward`: 奖励值。Slime 中所有获取的 Sample 请确保都是有效的 Sample，具体过滤操作需要在 Rollout Buffer 中进行处理 + - `loss_mask`: 损失掩码，目前只对所有 Assistant 类型的消息进行损失计算 + - `metadata`: 元数据信息，包含各种自定义信息 + +**核心逻辑**： + +1. 判断是否需要启动新的 Rollout 任务，通过查看 Start Rollout 的属性来判断（注意此处不可以使用 Slime Buffer 中的数据长度是否为 0 来判断，因为根据不同的 Slime Buffer Filter 规则，可能会导致 Slime Buffer 在训练过程中也可能为 0） +2. 轮询获取外部服务生成的数据，通过 Get Rollout Data 函数实现 + +##### 2. `get_rollout_data` - 数据获取函数 + +**功能描述**：异步从 Rollout Buffer 中获取已生成的数据，支持批量获取和元信息收集。 + +**函数签名**： +```python +async def get_rollout_data( + api_base_url: str, + num: Optional[int] = None, + timeout: float = 100.0 +) -> Tuple[List[Dict[str, Any]], Dict[str, Any]] +``` + +**输入参数**： +- `api_base_url`: Rollout Buffer 的 URL +- `num`: 可选，指定获取的数据条数，默认为 None，意味着获取全部数据 +- `timeout`: 请求超时时间（秒） + +**输出**： +- `Tuple[List[Dict], Dict]`: 包含数据列表和元信息的元组 + - 数据列表中每个字典必须包含： + ```python + { + "uid": "唯一标识符", + "instance_id": "实例ID", + "messages": [OpenAI格式的对话], + "reward": 0.85, # 奖励值 + "extra_info": {...} # 额外信息 + } + ``` + - 元信息字典支持自定义，只需要和后续日志记录等操作保持一致即可 + +**错误处理**： +- 网络错误时抛出 `aiohttp.ClientError` +- 数据格式错误时抛出 `ValueError` +- 超时时抛出 `asyncio.TimeoutError` + +##### 3. `start_rollout` - 任务启动函数 + +**功能描述**：向 Rollout Buffer 发送启动请求，配置任务参数并开始数据生成流程。 + +**函数签名**： +```python +def start_rollout(api_base_url: str, args, metadata) -> Dict[str, Any] +``` + +**输入参数**： +- `api_base_url`: Rollout Buffer 的 URL +- `args`: 包含所有配置的参数对象 +- `metadata`: 数据缓冲区的元数据，包含已完成的实例 ID 列表，用于续训 + +**输出**： +- `Dict[str, Any]`: 外部服务返回的启动确认信息 + +**发送的载荷格式示例**： +```python +{ + "num_process": 1024, # 并行进程数 + "num_epoch": 3, # 生成轮次 + "remote_engine_url": "http://sglang-router:port", + "task_type": "math", # 任务类型 + "input_file": "/path/to/input.jsonl", # 输入文件路径 + "num_repeat_per_sample": "8", # 每个样本重复次数 + "max_tokens": 8192, # 最大 Token 数 + "sampling_params": { # 采样参数 + "max_tokens": 8192, + "temperature": 0.8, + "top_p": 0.9 + }, + "tokenizer_path": "/path/to/tokenizer", # Tokenizer 路径 + "skip_instance_ids": ["id1", "id2"] # 跳过的实例 ID，续训时跳过之前已经训练过的 Instance ID +} +``` + +##### 4. `log_raw_info` - 日志记录函数 + +**功能描述**：收集、统计和记录 Rollout 过程中的关键指标，支持 WandB 集成，内容均可自定义。 + +**函数签名**： +```python +def log_raw_info(args, all_meta_info: List[Dict], rollout_id: int) -> None +``` + +**输入参数**： +- `args`: 配置参数对象，包含日志相关设置 +- `all_meta_info`: 从所有数据获取请求中收集的元信息列表 +- `rollout_id`: 当前 Rollout 的 ID + +**日志格式示例**： +```python +{ + "rollout/no_filter/total_samples": 100, + "rollout/no_filter/avg_reward": 0.75, +} +``` + +### 自定义 Agent Task 实现步骤 + +#### 1. 准备输入数据格式 + +您的输入文件（`--rollout-input-file`）应包含任务相关的数据，例如： + +```jsonl +{"instance_id": "math_001", "prompt": [{"role":"user","content":"求解方程 x^2 + 5x + 6 = 0"}]} +{"instance_id": "math_002", "prompt": [{"role":"user","content":"证明费马大定理"}]} +``` + +其中 `instance_id` 为必需字段，可以使用 `slime_plugins/rollout_buffer/tools/assign_instance_id.py` 自动生成每条数据的 Instance ID。 + +#### 2. 自定义 Loss Mask 生成器 + +如果您的任务有特殊的对话格式，需要在 `slime/utils/mask_utils.py` 中添加自定义的 Mask 生成器： + +```python +class CustomTaskLossMaskGenerator(MultiTurnLossMaskGenerator): + def get_custom_multi_turn_loss_mask(self, messages): + # 自定义实现 + # 返回 (token_ids, loss_mask) + pass + + def get_loss_mask(self, messages: List[Dict]) -> List[int]: + # ... existing code ... + elif self.tokenizer_type == "custom_task": + return xxx + # ... existing code ... +``` + +然后使用： +```bash +--loss-mask-type custom_task +``` + +目前已经实现了 Qwen 和 Distill Qwen 两种对话格式的多轮对话 Loss Mask 生成。 + +#### 3. 实现外部服务接口 + +在 Rollout Buffer 中需要针对您的任务实现关键函数 `run_rollout()`，具体实现请参见 [Rollout Buffer 使用文档](./rollout_buffer_usage.md)。 + +### 参数配置 + +#### Agent Rollout 核心参数 + +- **`--rollout-function-path`**：指定使用 Agent Rollout 函数 + ```bash + --rollout-function-path slime.rollout.agent_rollout.generate_rollout + ``` + +- **`--agent-rollout-buffer-url`**：Rollout Buffer 的 API 地址 + ```bash + --agent-rollout-buffer-url http://0.0.0.0:8889 + ``` + 具体部署方式请查看 Rollout Buffer 文档 + +- **`--rollout-input-file`**：输入数据文件路径，包含待处理的任务数据 + ```bash + --rollout-input-file ./deepscaler_rl_buffer_instance_id.jsonl + ``` + +- **`--rollout-num-process`**：并行处理的进程数 + ```bash + --rollout-num-process 1024 + ``` + +- **`--loss-mask-type`**：Loss Mask 类型，用于多轮对话的 Token 掩码，支持自定义，自定义路径在 `slime/utils/mask_utils.py` + ```bash + --loss-mask-type distill_qwen + ``` + +#### Rollout 生成参数 + +- **`--num-rollout`**：总的 Rollout 次数（最大 Rollout ID 数） + ```bash + --num-rollout 3000 + ``` + +- **`--rollout-batch-size`**：每次 Rollout 的批次大小 + ```bash + --rollout-batch-size 16 + ``` + +- **`--rollout-max-response-len`**：最大响应长度（单轮） + ```bash + --rollout-max-response-len 8192 + ``` + +- **`--rollout-temperature`**：采样温度参数 + ```bash + --rollout-temperature 0.8 + ``` + +- **`--n-samples-per-prompt`**：每个 Prompt 生成的样本数 + ```bash + --n-samples-per-prompt 8 + ``` + +#### 过滤与奖励配置 + +- **`--buffer-filter-path`**：数据缓冲区过滤器路径，支持自定义，默认采用 Buffer 中最新的一批数据进行更新 + ```bash + --buffer-filter-path slime.rollout.filter_hub.buffer_filters.pop_first + ``` + +- **`--disable-rewards-normalization`**：禁用奖励归一化，如果 Rollout Buffer 中已经归一化奖励，请启用此参数 + ```bash + --disable-rewards-normalization + ``` + diff --git a/docs/zh/rollout_buffer_usage.md b/docs/zh/rollout_buffer_usage.md new file mode 100644 index 0000000000..7081c82ac0 --- /dev/null +++ b/docs/zh/rollout_buffer_usage.md @@ -0,0 +1,324 @@ +# Rollout Buffer 使用文档 + +## 概述 + +Rollout Buffer 是 Slime 框架中用于智能体轨迹生成的独立组件，其主要功能是使用 Slime 训练启动的 LLM OpenAI Server 进行智能体轨迹的生成。 + +### 设计理念 + +我们将 Rollout Buffer 独立出来的主要原因包括： + +1. **框架解耦**：不同 Agent 任务所依赖的 Agent Framework 和工具都不相同，很可能会复用第三方的 Agent Framework +2. **灵活扩展**：如果将所有组件都封装到 Slime 内部会导致架构混乱，不利于扩展和维护 +3. **职责分离**：Rollout Buffer 只负责通过调用 Slime 中启动的 Server 生成对应的轨迹，具体使用什么框架没有任何限制 +4. **完全解耦**：轨迹生成逻辑和 Slime 训练逻辑完全解耦，支持引入各种复杂的 Agent Framework + +### 工作流程 + +``` +Slime Training Process ←─── HTTP API ───→ Rollout Buffer + ↓ ↓ + LLM Server ←─────── HTTP Requests ─────── Agent Framework + ↓ ↓ + Model Response ──────────────────────→ Trajectory Generation +``` + +对于每一个不同的 Agent 任务，都应该对应一个独立的 Generator 类，负责生成该类任务的轨迹。Rollout Buffer 会自动读取并加载不同类型的 Generator。 + +## 快速开始 + +### 基本使用流程 + +1. **复制模板**：将 `base_generator.py` 作为模板进行复制 +2. **修改任务类型**：将 `TASK_TYPE` 修改为您的任务名称（不能与其他 Generator 重复） +3. **实现核心函数**：实现 `run_rollout()` 函数 +4. **可选定制**：根据需要重写五个可选函数 +5. **启动训练**：按照 [Agent Training 文档](./agent_training.md) 中的启动流程启动 Agent 训练 + +### 文件结构规范 + +Generator 文件必须以 `_generator.py` 结尾，并放置在 `generator/` 目录下： + +``` +generator/ +├── base_generator.py # Math 任务实现（默认模板） +└── your_task_generator.py # 您的自定义任务 +``` + +## 核心组件 + +### 必需组件 + +每个 Generator 文件必须包含以下组件： + +#### 1. `TASK_TYPE` 常量 +定义任务类型的唯一标识符： +```python +TASK_TYPE = "your_task_name" +``` + +#### 2. `run_rollout()` 函数 +核心数据生成逻辑的入口函数： +```python +def run_rollout(data: dict): + # 实现您的轨迹生成逻辑 + pass +``` + +### 可选组件 + +除了必需组件外，Rollout Buffer 还提供了五个可自定义的函数来满足不同任务的特殊需求。如果不提供自定义实现，系统将使用默认实现（位于 `slime_plugins/rollout_buffer/generator/utils/default_func.py`）： + +1. **`normalize_group_data()`**：奖励归一化函数 +2. **`pad_group_data()`**：数据填充策略函数 +3. **`is_valid_group()`**：组数据有效性验证函数 +4. **`get_group_data_meta_info()`**：元信息统计函数 +5. **`filter_item()`**：单个数据项过滤函数 + +## 参数配置 + +### Generator 核心参数 + +`run_rollout(data: dict)` 函数接收的主要参数如下（传入的 `data` 需要与 Slime 中发送的参数保持一致）： + +| 参数名 | 类型 | 描述 | +|--------|------|------| +| `remote_engine_url` | string | 推理引擎服务地址，通常为 Slime 中的 SGLang Router 地址 | +| `remote_buffer_url` | string | Rollout Buffer 服务地址，通常为 Master 节点的某个端口（默认 8889） | +| `input_file` | string | 输入数据文件路径 | +| `task_type` | string | 任务类型标识符，定义在每个 `_generator.py` 文件中 | +| `num_repeat_per_sample` | int | 每个样本重复生成次数（Group Size） | +| `num_epoch` | int | 数据集遍历轮次（默认为 10） | +| `sampling_params` | dict | 模型采样参数（包含 max_tokens、temperature 等） | +| `num_process` | int | 并行进程数 | +| `skip_instance_ids` | list | 要跳过的实例 ID 列表，用于续训时跳过之前已处理的实例 | + +### Buffer 控制参数 + +Buffer 的行为由以下关键参数控制，这些参数直接影响数据的收集、验证和输出策略： + +#### 核心控制参数 + +| 参数名 | 默认值 | 描述 | +|--------|--------|------| +| `group_size` | - | 每组的目标数据数量，通常等于 `num_repeat_per_sample` | +| `min_valid_group_size_ratio` | 1.0 | 组被认为"有效"的最小数据比例（100%） | +| `min_valid_item_size_ratio` | 0.7 | 过滤后组内有效数据的最小比例（70%） | + +**重要说明**： +- `group_size`：所有数据最终会被填充到这个大小，直接影响训练时每个实例的采样数量 +- `min_valid_group_size_ratio`：建议设为 1.0，无效数据也可以写入，通过后续步骤过滤（如赋予极端 Reward） +- `min_valid_item_size_ratio`：过滤后组内有效数据的最小比例，应大于 0.5，用于过滤质量过差的组 + +#### 超时控制参数 + +| 参数名 | 默认值 | 描述 | +|--------|--------|------| +| `group_timeout_seconds` | 300 | 组超时时间（5分钟），防止部分组长时间卡住 | +| `min_timeout_group_size_ratio` | 0.7 | 超时组的最小数据比例阈值（70%） | + +#### 系统资源参数 + +| 参数名 | 默认值 | 描述 | +|--------|--------|------| +| `max_buffer_size` | 1,000,000,000 | Buffer 最大容量（10亿），防止内存溢出 | + +## 数据处理流程 + +### 完整处理流程 + +当从 Rollout Buffer 中获取一批训练数据时，五个可选函数按照以下固定顺序执行： + +``` +buffer.read(batch_size) 调用 + ↓ +1. 📊 get_group_data_meta_info() + └── 收集统计信息（进度、奖励分布等） + ↓ +2. ✅ is_valid_group() + └── 判断每个组是否完成且有效 + ↓ +3. 🔍 filter_item() + └── 对有效组中的每个数据项进行过滤 + ↓ +4. ⚖️ normalize_group_data() + └── 对过滤后的组数据进行奖励归一化 + ↓ +5. 📦 pad_group_data() + └── 将归一化后的数据填充至目标 group_size + ↓ +📤 返回处理完成的批次数据 +``` + +### 处理步骤详解 + +#### 第1步：元信息统计 - `get_group_data_meta_info()` + +**功能**：收集当前 Buffer 中所有原始组数据的统计信息 +- **输入**：Buffer 中所有原始组数据（包含无效组和无效轨迹） +- **输出**：包含统计信息的字典，用于日志记录和监控，比如可以记录平均奖励等信息 + +#### 第2步：组有效性验证 - `is_valid_group()` + +**功能**：确定哪些组可以用于训练 +- **输入**：每个组的完整数据 `(instance_id, group_data)` +- **输出**：`(is_valid, is_finished)` 元组 +- **逻辑关系**：`有效组 ⊆ 已完成组 ⊆ 所有组`，其中已完成组中的实例将会在续训时被跳过，有效组中的符合要求的组将会被用于训练模型 + +#### 第3步：单项数据过滤 - `filter_item()` + +**功能**：对有效组内的每个数据项进行精细化过滤 +- **输入**：组内的单个数据项 +- **输出**：布尔值，决定该项是否保留，因为写入 Rollout Buffer 的数据可能存在无效项，需要将其过滤 + +#### 第4步：奖励归一化 - `normalize_group_data()` + +**功能**：对组内奖励值进行标准化处理 +- **注意**：如果在此处进行归一化，需要在 Slime 中禁用奖励归一化，这里默认实现归一化的方式是只对于有效的数据 item 进行归一化并进行缩放 +- **其他**：原始奖励值会保存到 `raw_reward` 字段，方便进行日志记录 + +#### 第5步：数据填充 - `pad_group_data()` + +**功能**：将数据填充至标准的 `group_size` +- **策略**：通过奖励缩放保持总奖励一致性 +- **输出**：固定大小的组数据，可直接用于训练 +- **注意**：返回的数据数量**必须**要是 Group Size 的整数倍 + +### 重要机制说明 + +#### 数据存储策略 +- **全量存储**：无论轨迹生成是否成功，都应将所有数据存入 Buffer +- **后续过滤**：通过过滤机制筛选出有用的 Group 和 Item +- **失败处理**：为失败的轨迹分配特殊的 Reward 值便于识别 + +#### 超时清理机制 +- **自动清理**：每次执行 `get_rollout_data` 时检查时间戳 +- **判断逻辑**：超时组根据有效数据数量决定取出或丢弃 +- **防止积累**：有效防止数据在 Buffer 中过度积累 + +## 实现示例 + +### 基础实现模板 + +以 Math 任务为例，展示完整的 Generator 实现： + +```python +TASK_TYPE = "math" + +def run_rollout(data: dict): + + print(f"Starting math rollout with data: {data}") + + rollout_func = query_single_turn + reward_func = get_rule_based_math_reward + + print(f"Waiting for 10 seconds for buffer server to start") + time.sleep(10) + global SAMPLING_PARAMS + for k, v in data["sampling_params"].items(): + SAMPLING_PARAMS[k] = v + print(f"Set {k} to {v}", type(v)) + + generator = BaseGenerator( + data["remote_engine_url"], + data["remote_buffer_url"], + num_repeat_per_sample=int(data["num_repeat_per_sample"]), + queue_size=1000000, + max_tokens=int(data["sampling_params"]["max_tokens"]), + num_process=int(data.get("num_process", 100)), + task_type=data["task_type"], + skip_instance_ids=data.get("skip_instance_ids", None), + ) + + generator.entry(data["input_file"], rollout_func, reward_func, int(data.get("num_epoch", 1))) + + +``` + +### 轨迹生成函数示例 + +```python + +def query_single_turn(client, messages, sampling_params, tools=None): + base_payload = { + "messages": messages, + **sampling_params, + "model": "custom", + "stream": False, + "seed": random.randint(1, 10000000), + "tools": tools, + } + + text = None + accumulated_tokens = 0 + + for attempt in range(6): + try: + # Create a fresh payload for each attempt + current_payload = copy.deepcopy(base_payload) + + if text is not None: + # Update messages with current progress + current_messages = copy.deepcopy(messages) + current_messages.append({"role": "assistant", "content": text}) + current_payload["messages"] = current_messages + + # Adjust max_tokens based on accumulated tokens + if "max_tokens" in sampling_params: + current_payload["max_tokens"] = max(0, sampling_params["max_tokens"] - accumulated_tokens) + + # Add continue flag for partial rollouts + current_payload["extra_body"] = {"continue_final_message": True} + if current_payload["max_tokens"] == 0: + break + response = client.chat.completions.create(**current_payload) + + if len(response.choices) > 0: + if response.choices[0].finish_reason == "abort": + print( + f"query failed, reason: {response.choices[0].finish_reason}, currently generated: {response.usage.completion_tokens}" + ) + + accumulated_tokens += response.usage.completion_tokens + + if text is None: + text = response.choices[0].message.content + else: + text += response.choices[0].message.content + + sleep(10) + continue + if text is None: + text = response.choices[0].message.content + elif response.choices[0].message.content is not None: + text += response.choices[0].message.content + break + else: + print(f"Error in query, status code: {response.status_code}") + continue + except Exception as e: + print(f"query failed in single turn, error: {e}") + continue + + # Update final messages + if len(messages) > 0 and messages[-1]["role"] == "assistant": + messages = messages[:-1] + messages.append({"role": "assistant", "content": text}) + + return messages + +``` + +## 常见问题 + +### Q: 如何处理生成失败的数据？ +A: 将失败数据也存入 Buffer，但分配特殊的 Reward 值（如 -1），通过后续过滤机制处理。 + +### Q: 如何调试数据质量问题？ +A: 利用 `get_group_data_meta_info()` 函数收集详细统计信息，监控奖励分布和数据质量。 + +### Q: 超时机制如何工作？ +A: 当组的最后一次数据生成时间超过 `group_timeout_seconds` 时，系统会根据 `min_timeout_group_size_ratio` 决定是否使用该组数据。 + +### Q: 如何实现续训？ +A: Slime 将通过 `skip_instance_ids` 参数传递已处理的实例 ID 列表，Generator 会自动跳过这些实例。所有已完成的组都会自动的被跳过。 \ No newline at end of file diff --git a/scripts/agent-example.sh b/scripts/agent-example.sh new file mode 100644 index 0000000000..23d0a65873 --- /dev/null +++ b/scripts/agent-example.sh @@ -0,0 +1,152 @@ +#!/bin/bash + +# for rerun the task +pkill -9 sglang +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +export PYTHONBUFFERED=16 + +export TP_SIZE=2 +export PP_SIZE=1 +export CP_SIZE=1 + +export HF_MODEL_PATH=/root/hf_models/deepseek-ai--DeepSeek-R1-Distill-Qwen-7B +export MCORE_MODEL_PATH=/root/megatron_model/DeepSeek-R1-Distill-Qwen-7B-25.02 +export PROMPT_DATA=/root/dapo-math-17k/dapo-math-17k_processed.jsonl +export MCORE_MODEL_PATH_SAVE=/root/megatron_model/DeepSeek-R1-Distill-Qwen-7B-25.02_save + +# DeepSeek-R1-Distill-Qwen-7B +MODEL_ARGS=( + --swiglu + --num-layers 28 + --hidden-size 3584 + --ffn-hidden-size 18944 + --num-attention-heads 28 + --group-query-attention + --num-query-groups 4 + --max-position-embeddings 131072 + --seq-length 4096 + --use-rotary-position-embeddings + --disable-bias-linear + --add-qkv-bias + --normalization "RMSNorm" + --norm-epsilon 1e-06 + --rotary-base 10000 + --vocab-size 152064 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --moe-token-dispatcher-type alltoall + --untie-embeddings-and-output-weights + --attention-dropout 0.0 + --hidden-dropout 0.0 +) + +CKPT_ARGS=( + --hf-checkpoint ${HF_MODEL_PATH} + --ref-load ${MCORE_MODEL_PATH} + --save-interval 100 + --save ${MCORE_MODEL_PATH_SAVE} +) + +ROLLOUT_ARGS=( + --rollout-function-path slime.rollout.agent_rollout.generate_rollout + --rm-type deepscaler + --prompt-data ${PROMPT_DATA} + --label-key label + --num-rollout 3000 + --rollout-batch-size 128 + --rollout-max-response-len 8192 + --rollout-temperature 0.8 + --rollout-shuffle + --n-samples-per-prompt 8 + --global-batch-size 1024 + --micro-batch-size 8 + --ref-micro-batch-size 8 + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 + --balance-data +) + +DISTRIBUTED_ARGS=( + --tensor-model-parallel-size ${TP_SIZE} + --pipeline-model-parallel-size ${PP_SIZE} + --context-parallel-size ${CP_SIZE} + --sequence-parallel +) + +PERF_ARGS=( + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.001 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.00 +) + +OPTIMIZER_ARGS=( + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +WANDB_ARGS=( + # --use-wandb \ +) + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +export MASTER_PORT=${MASTER_PORT:-"12345"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json='{ + "env_vars": { + "PYTHONPATH": "/root/Megatron-LM/", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_CUMEM_ENABLE": "0" + } + }' \ + -- python3 train_agent_async.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 4 \ + --rollout-num-gpus 4 \ + --rollout-num-gpus-per-engine 1 \ + --sglang-mem-fraction-static 0.8 \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${DISTRIBUTED_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + --agent-rollout-buffer-url http://${MASTER_ADDR}:8889 \ + --keep-old-actor \ + --update-rollout-weights-interval 1 \ + --buffer-filter-path slime.rollout.filter_hub.buffer_filters.pop_first \ + --disable-rewards-normalization \ + --offload-rollout \ + --offload-ref \ + --rollout-input-file ${PROMPT_DATA} \ + --rollout-num-process 1024 \ + --loss-mask-type distill_qwen \ + --sglang-log-level error \ + --input-key prompt \ + --log-passrate diff --git a/scripts/run_agent.sh b/scripts/run_agent.sh new file mode 100755 index 0000000000..8b07040e5f --- /dev/null +++ b/scripts/run_agent.sh @@ -0,0 +1,22 @@ + +#!/bin/bash + +set -e + +SESSION_NAME="slime_run" +WINDOW_1="slime" +WINDOW_2="buffer" + +if tmux has-session -t $SESSION_NAME 2>/dev/null; then + echo "Killing existing tmux session: $SESSION_NAME" + tmux kill-session -t $SESSION_NAME +fi + +tmux new-session -d -s $SESSION_NAME -n $WINDOW_1 +tmux send-keys -t ${SESSION_NAME}:${WINDOW_1} "cd $(pwd)" C-m +tmux send-keys -t ${SESSION_NAME}:${WINDOW_1} "bash ./scripts/agent-example.sh" C-m + +tmux new-window -t $SESSION_NAME -n $WINDOW_2 +tmux send-keys -t ${SESSION_NAME}:${WINDOW_2} "sleep 30 && cd slime_plugins/rollout_buffer && python buffer.py" C-m + +tmux attach-session -t $SESSION_NAME \ No newline at end of file diff --git a/slime/backends/megatron_utils/__init__.py b/slime/backends/megatron_utils/__init__.py index 9ba67ba26a..9843547171 100644 --- a/slime/backends/megatron_utils/__init__.py +++ b/slime/backends/megatron_utils/__init__.py @@ -4,6 +4,8 @@ get_batch, get_data_iterator, log_eval_data, + log_multi_turn_data, + log_passrate, log_perf_data, log_rollout_data, process_rollout_data, @@ -25,6 +27,8 @@ "set_metadata", "get_log_probs_and_entropy", "log_rollout_data", + "log_passrate", + "log_multi_turn_data", "log_eval_data", "log_perf_data", "compute_advantages_and_returns", diff --git a/slime/backends/megatron_utils/cp_utils.py b/slime/backends/megatron_utils/cp_utils.py index 36e606c19e..2742d238c1 100644 --- a/slime/backends/megatron_utils/cp_utils.py +++ b/slime/backends/megatron_utils/cp_utils.py @@ -57,7 +57,7 @@ def sum_of_sample_mean(x: torch.Tensor): if loss_masks is not None: res = sum( [ - (x_i * loss_mask_i).sum() / torch.max(loss_mask_i.sum(), 1) + (x_i * loss_mask_i).sum() / torch.clamp_min(loss_mask_i.sum(), 1) for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks) ] ) @@ -97,7 +97,7 @@ def sum_of_sample_mean(x): if loss_masks is not None: x = sum( [ - (x_i * chunked_loss_mask).sum() / torch.max(loss_mask.sum(), 1) + (x_i * chunked_loss_mask).sum() / torch.clamp_min(loss_mask.sum(), 1) for x_i, chunked_loss_mask, loss_mask in zip( x.split(cp_chunk_lengths, dim=0), chunked_loss_masks, loss_masks ) diff --git a/slime/backends/megatron_utils/data.py b/slime/backends/megatron_utils/data.py index 6537b7c71c..2b769bf25a 100644 --- a/slime/backends/megatron_utils/data.py +++ b/slime/backends/megatron_utils/data.py @@ -1,14 +1,16 @@ +import math from typing import Any, Optional +import numpy as np import ray import torch import torch.distributed as dist import torch.nn.functional as F -import wandb from megatron.core import mpu from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.utils import get_model_config +import wandb from slime.utils.flops_utils import calculate_fwd_flops from slime.utils.memory_utils import clear_memory from slime.utils.seqlen_balancing import get_seqlen_balanced_partitions @@ -300,6 +302,7 @@ def get_partition(val): "rewards", "truncated", "loss_masks", + "round_number", ]: if key not in data: continue @@ -371,6 +374,133 @@ def log_rollout_data(rollout_id, args): group=mpu.get_data_parallel_group(with_context_parallel=True), ) + if args.log_multi_turn: + log_multi_turn_data(rollout_id, args) + if args.log_passrate: + log_passrate(rollout_id, args) + + +def log_multi_turn_data(rollout_id, args): + if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage(): + cp_size = mpu.get_context_parallel_world_size() + log_dict = {} + response_lengths = get_local_storage("response_lengths") + for key, val in get_local_storage().items(): + if key == "loss_masks": + if val: # Check if val is not empty + device = val[0].device # Get device from first tensor + + # Vectorized length calculation using torch + raw_response_lengths = torch.tensor([v.shape[0] for v in val], dtype=torch.float32, device=device) + log_dict["raw_response_length/response_length_mean"] = raw_response_lengths.mean().item() + log_dict["raw_response_length/response_length_max"] = raw_response_lengths.max().item() + log_dict["raw_response_length/response_length_min"] = raw_response_lengths.min().item() + log_dict["raw_response_length/response_length_clip_ratio"] = ( + (raw_response_lengths > args.rollout_max_response_len).float().mean().item() + ) + + # Vectorized sum calculation using torch - stay on GPU + wo_obs_response_lengths = torch.tensor( + [v.sum().item() for v in val], dtype=torch.float32, device=device + ) + log_dict["wo_obs_response_length/response_length_mean"] = wo_obs_response_lengths.mean().item() + log_dict["wo_obs_response_length/response_length_max"] = wo_obs_response_lengths.max().item() + log_dict["wo_obs_response_length/response_length_min"] = wo_obs_response_lengths.min().item() + if key == "round_number": + # Use numpy for vectorized round number statistics + round_number_array = np.array(val) + log_dict["multi_turn_metric/round_number_mean"] = np.mean(round_number_array) + log_dict["multi_turn_metric/round_number_max"] = np.max(round_number_array) + log_dict["multi_turn_metric/round_number_min"] = np.min(round_number_array) + if mpu.get_data_parallel_rank(with_context_parallel=True) == 0: + gathered_log_dict = [None] * mpu.get_data_parallel_world_size(with_context_parallel=True) + # Not sure if this will be a performance bottleneck. + dist.gather_object( + log_dict, + gathered_log_dict, + dst=mpu.get_data_parallel_src_rank(with_context_parallel=True), + group=mpu.get_data_parallel_group(with_context_parallel=True), + ) + dp_size = mpu.get_data_parallel_world_size(with_context_parallel=True) + reduced_log_dict = { + f"multi_turn/{key}": sum([d[key] for d in gathered_log_dict]) / dp_size for key in log_dict + } + print(f"multi_turn {rollout_id}: {reduced_log_dict}") + if args.use_wandb: + wandb.log(reduced_log_dict) + else: + dist.gather_object( + log_dict, + None, + dst=mpu.get_data_parallel_src_rank(with_context_parallel=True), + group=mpu.get_data_parallel_group(with_context_parallel=True), + ) + + +def log_passrate(rollout_id, args): + if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage(): + cp_size = mpu.get_context_parallel_world_size() + log_dict = {} + response_lengths = get_local_storage("response_lengths") + for key, val in get_local_storage().items(): + if key != "raw_reward": + continue + + group_size = args.n_samples_per_prompt + group_number = args.rollout_batch_size + assert len(val) == group_number * group_size + pass_rate_name_list = [2**i for i in range(int(math.log2(group_size)) + 1)] + + val = np.array(val).reshape(group_number, group_size) + + def estimate_pass_at_k(num_samples, num_correct, k): + """ + Estimates pass@k of each problem and returns them in an array. + """ + + def estimator(n, c, k): + """ + Calculates 1 - comb(n - c, k) / comb(n, k). + """ + if n - c < k: + return 1.0 + return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1)) + + return np.array([estimator(int(n), int(c), k) for n, c in zip(num_samples, num_correct)]) + + for k in pass_rate_name_list: + num_correct = np.sum(val == 1, axis=1) + num_samples = np.full(group_number, group_size) + + pass_k_estimates = estimate_pass_at_k(num_samples, num_correct, k) + + pass_k = np.mean(pass_k_estimates) + log_dict[f"pass@{k}"] = pass_k + + if mpu.get_data_parallel_rank(with_context_parallel=True) == 0: + gathered_log_dict = [None] * mpu.get_data_parallel_world_size(with_context_parallel=True) + # Not sure if this will be a performance bottleneck. + dist.gather_object( + log_dict, + gathered_log_dict, + dst=mpu.get_data_parallel_src_rank(with_context_parallel=True), + group=mpu.get_data_parallel_group(with_context_parallel=True), + ) + dp_size = mpu.get_data_parallel_world_size(with_context_parallel=True) + reduced_log_dict = { + f"passrate/{key}": sum([d[key] for d in gathered_log_dict]) / dp_size for key in log_dict + } + print(f"passrate {rollout_id}: {reduced_log_dict}") + if args.use_wandb: + wandb.log(reduced_log_dict) + else: + dist.gather_object( + log_dict, + None, + dst=mpu.get_data_parallel_src_rank(with_context_parallel=True), + group=mpu.get_data_parallel_group(with_context_parallel=True), + ) + def log_eval_data(rollout_id, args, data_buffer): if ( diff --git a/slime/backends/megatron_utils/initialize.py b/slime/backends/megatron_utils/initialize.py index 257b99540b..d84e1d321c 100644 --- a/slime/backends/megatron_utils/initialize.py +++ b/slime/backends/megatron_utils/initialize.py @@ -4,11 +4,12 @@ import numpy as np import torch import torch.distributed as dist -import wandb from megatron.core import mpu, tensor_parallel from megatron.core.num_microbatches_calculator import init_num_microbatches_calculator from megatron.training.global_vars import _build_tokenizer, set_args +import wandb + def _set_random_seed( seed_: int, @@ -128,6 +129,8 @@ def init(args): wandb.define_metric("train/*", step_metric="train/step") wandb.define_metric("rollout/step") wandb.define_metric("rollout/*", step_metric="rollout/step") + wandb.define_metric("multi_turn/*", step_metric="rollout/step") + wandb.define_metric("passrate/*", step_metric="rollout/step") wandb.define_metric("eval/step") wandb.define_metric("eval/*", step_metric="eval/step") wandb.define_metric("perf/step") diff --git a/slime/ray/buffer.py b/slime/ray/buffer.py index b5b211baf2..a1e8af1788 100644 --- a/slime/ray/buffer.py +++ b/slime/ray/buffer.py @@ -5,9 +5,9 @@ import ray import torch -import wandb from transformers import AutoTokenizer +import wandb from slime.utils.data import JsonlDataset from slime.utils.misc import load_function from slime.utils.types import Sample @@ -37,6 +37,8 @@ def convert_samples_to_train_data(samples: list[Sample]): train_data["loss_masks"].append(sample.loss_mask) if samples[0].metadata and "raw_reward" in samples[0].metadata: train_data["raw_reward"] = [sample.metadata["raw_reward"] for sample in samples] + if samples[0].metadata and "round_number" in samples[0].metadata: + train_data["round_number"] = [sample.metadata["round_number"] for sample in samples] return train_data diff --git a/slime/ray/ppo_actor.py b/slime/ray/ppo_actor.py index 6a4d4bbbc5..64c702f75c 100644 --- a/slime/ray/ppo_actor.py +++ b/slime/ray/ppo_actor.py @@ -108,6 +108,18 @@ def init(self, args, role, with_ref=False): self.rollout_engines = None self.rollout_engine_lock = None self.data_buffer = None + if self.args.keep_old_actor: + old_args = args.load, args.no_load_optim, args.no_load_rng + args.load = args.ref_load + args.no_load_optim = True + args.no_load_rng = True + self.old_actor, _, _, _ = megatron_utils.initialize_model_and_optimizer(args, with_optimizer=False) + args.load, args.no_load_optim, args.no_load_rng = old_args + if self.args.offload_rollout: + for model_module in self.old_actor: + model_module.to(device="cpu", non_blocking=True) + else: + self.old_actor = None self.rollout_data_postprocess = None if self.args.rollout_data_postprocess_path is not None: @@ -263,7 +275,7 @@ def train(self, rollout_id, with_data_fetching=True): # For debug rollout, we just log the data and return. if with_data_fetching: self.get_rollout_data(rollout_id) - megatron_utils.log_rollout_data(rollout_id, self.args, self.data_buffer) + megatron_utils.log_rollout_data(rollout_id, self.args) megatron_utils.log_perf_data(rollout_id, self.args) Timer().start("train_wait") return @@ -284,27 +296,28 @@ def train(self, rollout_id, with_data_fetching=True): print_memory("begin train") - if self.ref is not None: - if self.args.offload_ref: - for model_module in self.ref: - model_module.to(device=torch.cuda.current_device(), non_blocking=True) - print_memory("after load ref model") - - with timer("ref_log_probs"): - megatron_utils.forward_only( - self.args, - self.ref, - log_probs_data_iterator, - log_probs_num_microbatches, - store_prefix="ref_", - ) + with timer("ref_log_probs"): + if self.ref is not None: + if self.args.offload_ref: + for model_module in self.ref: + model_module.to(device=torch.cuda.current_device(), non_blocking=True) + print_memory("after load ref model") + + with timer("ref_log_probs_forward"): + megatron_utils.forward_only( + self.args, + self.ref, + log_probs_data_iterator, + log_probs_num_microbatches, + store_prefix="ref_", + ) - # TODO: we should offload ref model here. But some how CuMemAllocator will raise error. - if self.args.offload_ref: - for model_module in self.ref: - model_module.to(device="cpu", non_blocking=True) - clear_memory() - print_memory("after offload ref model") + # TODO: we should offload ref model here. But some how CuMemAllocator will raise error. + if self.args.offload_ref: + for model_module in self.ref: + model_module.to(device="cpu", non_blocking=True) + clear_memory() + print_memory("after offload ref model") # reset data iterator for data_iterator in log_probs_data_iterator: @@ -312,24 +325,49 @@ def train(self, rollout_id, with_data_fetching=True): # Calculate logprob, the results will be stored on last pp stage. with timer("log_probs"): - if self.args.offload: - self.wake_up("model") + if self.old_actor is not None: + if self.args.offload_rollout: + for model_module in self.old_actor: + model_module.to(device=torch.cuda.current_device(), non_blocking=True) + print_memory("after load rollout model") + with timer("rollout_log_probs_forward"): + megatron_utils.forward_only( + self.args, + self.old_actor, + log_probs_data_iterator, + log_probs_num_microbatches, + ) - megatron_utils.forward_only( - self.args, - self.model, - log_probs_data_iterator, - log_probs_num_microbatches, - ) + if self.args.offload_rollout: + for model_module in self.old_actor: + model_module.to(device="cpu", non_blocking=True) + torch.cuda.empty_cache() + print_memory("after offload rollout model") + + if self.args.offload: + self.wake_up("model") + else: + if self.args.offload: + self.wake_up("model") + + megatron_utils.forward_only( + self.args, + self.model, + log_probs_data_iterator, + log_probs_num_microbatches, + ) # Calculate adv and returns. Need to performed before training (instead of on the fly), # because we may need normalize the whole rollout. - megatron_utils.compute_advantages_and_returns(self.args) + with timer("compute_advantages_and_returns"): + megatron_utils.compute_advantages_and_returns(self.args) if self.rollout_data_postprocess is not None: - self.rollout_data_postprocess(self.args) + with timer("rollout_data_postprocess"): + self.rollout_data_postprocess(self.args) - megatron_utils.log_rollout_data(rollout_id, self.args) + with timer("log_rollout_data"): + megatron_utils.log_rollout_data(rollout_id, self.args) # Train with timer("actor_train"): @@ -342,7 +380,8 @@ def train(self, rollout_id, with_data_fetching=True): train_num_microbatches, ) - megatron_utils.log_perf_data(rollout_id, self.args) + with timer("log_perf_data"): + megatron_utils.log_perf_data(rollout_id, self.args) Timer().start("train_wait") megatron_utils.data.clear_local_storage() @@ -540,6 +579,17 @@ def update_weights(self): clear_memory() print_memory("after update_weights") + if getattr(self.args, "keep_old_actor", False): + print("update rollout model on cpu using actor model") + for cpu_module, src_module in zip(self.old_actor, self.model): + cpu_module.load_state_dict(src_module.state_dict(), strict=True) + + def get_old_actor_state_dict(self): + if self.old_actor is not None: + return self.old_actor.state_dict() + else: + return None + class RayTrainGroup: """ diff --git a/slime/rollout/agent_rollout.py b/slime/rollout/agent_rollout.py new file mode 100644 index 0000000000..b9d3b1d207 --- /dev/null +++ b/slime/rollout/agent_rollout.py @@ -0,0 +1,317 @@ +import asyncio +import time +from typing import Any, Dict, List, Optional + +import aiohttp +import requests +from transformers import AutoTokenizer + +import wandb +from slime.ray.buffer import Buffer +from slime.utils.async_utils import run +from slime.utils.mask_utils import MultiTurnLossMaskGenerator +from slime.utils.types import Sample + + +__all__ = ["generate_agent_rollout"] + + +# Global variables for evaluation +TOKENIZER = None +START_ROLLOUT = True + + +def select_rollout_data(args, results, need_length): + """ + Select the most recent groups when there are too many samples. + Groups all samples by instance_id, sorts groups by timestamp. + + Args: + args: Arguments containing configuration + results: List of rollout data items with timestamps + + Returns: + Selected samples from the newest groups based on timestamp cutoff + """ + if not results: + return results + + # Group samples by instance_id + groups = {} + for item in results: + assert "instance_id" in item, "instance_id must be in item" + instance_id = item["instance_id"] + if instance_id not in groups: + groups[instance_id] = [] + groups[instance_id].append(item) + + print(f"📊 Total groups: {len(groups)}, total samples: {len(results)}") + + # If we don't have too many samples, return all + assert need_length < len(results), "need_length must be smaller than results length" + + # Get timestamp for each group (use the latest timestamp in the group) + def get_group_timestamp(group_items): + timestamps = [] + for item in group_items: + if "timestamp" in item: + timestamps.append(float(item["timestamp"])) + elif "extra_info" in item and "timestamp" in item["extra_info"]: + timestamps.append(float(item["extra_info"]["timestamp"])) + return max(timestamps) if timestamps else 0 + + # Create list of (group_id, timestamp, samples) and sort by timestamp + group_data = [] + for group_id, group_items in groups.items(): + group_timestamp = get_group_timestamp(group_items) + group_data.append((group_id, group_timestamp, group_items)) + + # Sort groups by timestamp (newest first) + group_data.sort(key=lambda x: x[1], reverse=True) + + selected_groups = group_data[:need_length] + + # Flatten selected groups back to sample list + selected_results = [] + for group_id, timestamp, group_items in selected_groups: + selected_results.extend(group_items) + + # Statistics for monitoring + if selected_groups: + newest_ts = selected_groups[0][1] + oldest_ts = selected_groups[-1][1] + print(f"📈 Selected {len(selected_groups)} groups with {len(selected_results)} samples") + print(f"📈 Group timestamp range: {oldest_ts:.2f} to {newest_ts:.2f}") + print(f"📈 Time span: {newest_ts - oldest_ts:.2f} seconds") + + return selected_results + + +def log_raw_info(args, all_meta_info, rollout_id): + final_meta_info = {} + if all_meta_info: + final_meta_info = { + "total_samples": sum(meta["total_samples"] for meta in all_meta_info if "total_samples" in meta) + } + + total_samples = final_meta_info["total_samples"] + if total_samples > 0: + weighted_reward_sum = sum( + meta["avg_reward"] * meta["total_samples"] + for meta in all_meta_info + if "avg_reward" in meta and "total_samples" in meta + ) + + final_meta_info.update( + { + "avg_reward": weighted_reward_sum / total_samples, + } + ) + if hasattr(args, "use_wandb") and args.use_wandb: + log_dict = { + f"rollout/no_filter/total_samples": final_meta_info["total_samples"], + f"rollout/no_filter/avg_reward": final_meta_info["avg_reward"], + } + try: + if args.use_wandb: + log_dict["rollout/step"] = ( + rollout_id + if not args.wandb_always_use_train_step + else rollout_id + * args.rollout_batch_size + * args.n_samples_per_prompt + // args.global_batch_size + ) + wandb.log(log_dict) + print(f"no filter rollout log {rollout_id}: {log_dict}") + except Exception as e: + print(f"Failed to log to wandb: {e}") + print(f"no filter rollout log {rollout_id}: {final_meta_info}") + else: + print(f"no filter rollout log {rollout_id}: {final_meta_info}") + + +async def get_rollout_data( + api_base_url: str, num: Optional[int] = None, timeout: float = 100.0 +) -> tuple[List[Dict[str, Any]], Dict[str, Any]]: + + url = f"{api_base_url}/get_rollout_data" + payload = {} + + if num is not None: + payload["batch_size"] = num + print(url) + try: + start_time = time.time() + async with aiohttp.ClientSession() as session: + while True: + async with session.post(url, json=payload, timeout=aiohttp.ClientTimeout(total=timeout)) as response: + response.raise_for_status() + resp_json = await response.json() + if resp_json["success"]: + break + await asyncio.sleep(3) + if time.time() - start_time > 30: + print("rollout data is not ready, have been waiting for 30 seconds") + # Reset start_time to continue waiting or handle timeout differently + start_time = time.time() # Or raise an exception, or return empty list + + data = resp_json["data"] + meta_info = {} + if type(data) is list: + if "data" in data: + data = [item["data"] for item in data] + elif type(data) is dict: + if "data" in data: + meta_info = data["meta_info"] + data = data["data"] + print(f"Meta info: {meta_info}") + required_keys = {"uid", "instance_id", "messages", "reward", "extra_info"} + for item in data: + if not required_keys.issubset(item.keys()): + raise ValueError(f"Missing required keys in response item: {item}") + + return data, meta_info + + except aiohttp.ClientError as e: + print(f"[ERROR] Request failed: {e}") + raise + except ValueError as ve: + # print(f"[ERROR] Invalid data format: {ve}") + raise + except asyncio.TimeoutError: + print(f"[ERROR] Request timed out after {timeout} seconds") + raise + + +def start_rollout(api_base_url: str, args, metadata): + url = f"{api_base_url}/start_rollout" + if args.rollout_input_file is None: + raise ValueError("rollout_input_file is required") + print(f"metadata: {metadata}") + finished_groups_instance_id_list = [item for sublist in metadata.values() for item in sublist] + payload = { + "num_process": str(getattr(args, "rollout_num_process", 100)), + "num_epoch": str(getattr(args, "rollout_num_epoch", 3)), + "remote_engine_url": f"http://{args.sglang_router_ip}:{args.sglang_router_port}", + "remote_buffer_url": args.agent_rollout_buffer_url, + "task_type": args.rollout_task_type, + "input_file": args.rollout_input_file, + "num_repeat_per_sample": str(args.n_samples_per_prompt), + "max_tokens": str(args.rollout_max_response_len), + "sampling_params": { + "max_tokens": args.rollout_max_response_len, + "temperature": args.rollout_temperature, + "top_p": args.rollout_top_p, + }, + "tokenizer_path": args.hf_checkpoint, + "skip_instance_ids": finished_groups_instance_id_list, + } + print("start rollout with payload: ", payload) + + while True: + try: + resp = requests.post(url, json=payload, timeout=10) + resp.raise_for_status() + data = resp.json() + print(f"[start_rollout] Success: {data}") + return data + except Exception as e: + print(f"[start_rollout] Failed to send rollout config: {e}") + + +async def generate_agent_rollout( + args, rollout_id: int, data_buffer: Buffer, evaluation: bool = False +) -> Dict[str, Any]: + + global START_ROLLOUT + if evaluation: + raise NotImplementedError("Evaluation rollout is not implemented") + + if START_ROLLOUT: + metadata = data_buffer.get_metadata() + start_inform = start_rollout(args.agent_rollout_buffer_url, args, metadata) + print(f"start rollout with payload: {start_inform}") + print(f"start rollout id: {rollout_id}") + START_ROLLOUT = False + + data_number_to_fetch = args.rollout_batch_size * args.n_samples_per_prompt - data_buffer.get_buffer_length() + if data_number_to_fetch <= 0: + print( + f"❕buffer length: {data_buffer.get_buffer_length()}, buffer has enough data, return {args.rollout_batch_size * args.n_samples_per_prompt} samples" + ) + return await data_buffer.get_samples(args.rollout_batch_size * args.n_samples_per_prompt) + assert ( + data_number_to_fetch % args.n_samples_per_prompt == 0 + ), "data_number_to_fetch must be a multiple of n_samples_per_prompt" + print(f"INFO: buffer length: {data_buffer.get_buffer_length()}, data_number_to_fetch: {data_number_to_fetch}") + base_url = args.agent_rollout_buffer_url + tokenizer = AutoTokenizer.from_pretrained(args.hf_checkpoint, trust_remote_code=True) + retry_times = 0 + results = [] + all_meta_info = [] + + if args.fetch_trajectory_retry_times == -1: + print( + f"⚠️ [get_rollout_data] Fetch trajectory retry times set to -1, will retry indefinitely until sufficient data is collected" + ) + while args.fetch_trajectory_retry_times == -1 or retry_times < args.fetch_trajectory_retry_times: + try: + while len(results) < data_number_to_fetch: + time.sleep(5) + data, meta_info = await get_rollout_data(api_base_url=base_url) + results.extend(data) + if meta_info: + all_meta_info.append(meta_info) + print(f"get rollout data with length: {len(results)}") + break + except Exception as err: + print(f"[get_rollout_data] Failed to get rollout data: {err}, retry times: {retry_times}") + retry_times += 1 + + log_raw_info(args, all_meta_info, rollout_id) + + # Apply group-based data selection if there are too many samples + results = select_rollout_data(args, results, data_number_to_fetch // args.n_samples_per_prompt) + + if len(all_meta_info) > 0 and "finished_groups" in all_meta_info[0]: + finished_groups_instance_id_list = [] + for item in all_meta_info: + finished_groups_instance_id_list.extend(item["finished_groups"]) + + data_buffer.update_metadata({str(rollout_id): finished_groups_instance_id_list}) + + print("finally get rollout data with length: ", len(results)) + sample_results = [] + for i, record in enumerate(results): + oai_messages = record["messages"] + + mask_generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type=args.loss_mask_type) + token_ids, loss_mask = mask_generator.get_loss_mask(oai_messages) + response_length = mask_generator.get_response_lengths([loss_mask])[0] + + loss_mask = loss_mask[-response_length:] + + sample_results.append( + Sample( + index=record["instance_id"], + prompt=record["uid"], + tokens=token_ids, + response_length=response_length, + reward=record["reward"], + truncated=False, + loss_mask=loss_mask, + metadata={**record["extra_info"], "raw_reward": record["raw_reward"]}, + ) + ) + final_return_results = [] + + await data_buffer.add_samples(sample_results) + final_return_results = await data_buffer.get_samples(args.rollout_batch_size * args.n_samples_per_prompt) + + return final_return_results + + +def generate_rollout(args, rollout_id, data_buffer, evaluation=False): + """Generate rollout for both training and evaluation.""" + return run(generate_agent_rollout(args, rollout_id, data_buffer, evaluation)) diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 041d13afac..2d05492e19 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -645,6 +645,55 @@ def add_reward_model_arguments(parser): ) return parser + def add_agent_rollout_arguments(parser): + parser.add_argument( + "--agent-rollout-buffer-url", + type=str, + default=None, + help="URL for the agent rollout buffer", + ) + parser.add_argument( + "--update-rollout-weights-interval", + type=int, + default=1, + help="Interval for updating the weights of the agent", + ) + parser.add_argument( + "--fetch-trajectory-retry-times", + type=int, + default=-1, + help="Number of times to retry fetching trajectory, -1 means unlimited retry", + ) + parser.add_argument( + "--keep-old-actor", + action="store_true", + help="Whether to keep the rollout model on training process", + ) + parser.add_argument( + "--offload-rollout", + action="store_true", + help="Whether to update the rollout model on cpu", + ) + parser.add_argument( + "--min-batch-collection-ratio", + type=float, + default=1, + help="Minimum batch collection ratio", + ) + parser.add_argument( + "--rollout-task-type", + type=str, + default="math", + ) + parser.add_argument( + "--loss-mask-type", + type=str, + default="qwen", + choices=["qwen", "distill_qwen"], + help="Loss mask type", + ) + return parser + def add_custom_megatron_plugins_arguments(parser): """ Add custom Megatron plugins arguments. @@ -682,6 +731,7 @@ def add_custom_megatron_plugins_arguments(parser): parser = add_sglang_arguments(parser) parser = add_network_arguments(parser) parser = add_reward_model_arguments(parser) + parser = add_agent_rollout_arguments(parser) parser = add_custom_megatron_plugins_arguments(parser) # For megatron parser.add_argument("--padded-vocab-size", type=int, default=None) diff --git a/slime/utils/mask_utils.py b/slime/utils/mask_utils.py new file mode 100644 index 0000000000..bddb432618 --- /dev/null +++ b/slime/utils/mask_utils.py @@ -0,0 +1,101 @@ +from typing import Dict, List, Tuple + +from transformers import AutoTokenizer + + +class MultiTurnLossMaskGenerator: + def __init__(self, tokenizer: AutoTokenizer, tokenizer_type: str = "qwen"): + self.tokenizer = tokenizer + self.system_message_length, self.gen_token_length = self.get_system_message_length() + self.tokenizer_type = tokenizer_type + + def get_response_lengths(self, loss_masks: List[List[int]]) -> List[int]: + return [len(mask[mask.index(1) :]) if 1 in mask else 0 for mask in loss_masks] + + def find_all_sublist_indices(self, main_list, sublist): + sublist_len = len(sublist) + indices = [] + for i in range(len(main_list) - sublist_len + 1): + if main_list[i : i + sublist_len] == sublist: + indices.append(i) + return indices + + def get_system_message_length(self) -> Tuple[int, int]: + test_string = "FOR TESTING ONLY" + test_messages = [ + {"role": "user", "content": test_string}, + {"role": "user", "content": test_string}, + ] + raw_token_ids = self.tokenizer(test_string, add_special_tokens=False)["input_ids"] + chat_template_token = self.tokenizer.apply_chat_template( + test_messages, add_special_tokens=False, tokenize=False + ) + chat_template_token_ids = self.tokenizer(chat_template_token, add_special_tokens=False)["input_ids"] + idx_1, idx_2 = self.find_all_sublist_indices(chat_template_token_ids, raw_token_ids) + end_interval = len(chat_template_token_ids) - len(raw_token_ids) - idx_2 + gen_token_length = len( + self.tokenizer.apply_chat_template( + test_messages, add_special_tokens=False, tokenize=True, add_generation_prompt=True + ) + ) - len(chat_template_token_ids) + + system_message_length = idx_1 - ((idx_2 - idx_1) - end_interval - len(raw_token_ids)) + return system_message_length, gen_token_length + + def gen_multi_turn_loss_mask_qwen(self, messages: List[Dict]) -> Tuple[List[int], List[int]]: + all_loss_masks = [] + all_token_ids = [] + + for i, message in enumerate(messages): + message_ids = self.tokenizer.apply_chat_template([message], tokenize=True) + + if message["role"] != "system" and i > 0: + message_ids = message_ids[self.system_message_length :] + + if message["role"] == "assistant": + loss_mask = [0] * self.gen_token_length + [1] * (len(message_ids) - self.gen_token_length) + else: + loss_mask = [0] * len(message_ids) + + all_loss_masks.extend(loss_mask) + all_token_ids.extend(message_ids) + + return all_token_ids, all_loss_masks + + def gen_multi_turn_loss_mask_distill_qwen(self, messages: List[Dict]) -> Tuple[List[int], List[int]]: + prompt = self.tokenizer.apply_chat_template(messages[:1], tokenize=False, add_generation_prompt=True) + response = messages[-1]["content"] + prompt_tokens = self.tokenizer(prompt, add_special_tokens=False)["input_ids"] + response_tokens = self.tokenizer(response, add_special_tokens=False)["input_ids"] + + response_length = len(response_tokens) + token_ids = prompt_tokens + response_tokens + loss_mask = [0] * len(prompt_tokens) + [1] * response_length + return token_ids, loss_mask + + def get_loss_mask(self, messages: List[Dict]) -> List[int]: + if self.tokenizer_type == "qwen": + if "<｜Assistant｜>" in self.tokenizer.get_added_vocab(): + return self.gen_multi_turn_loss_mask_distill_qwen(messages) + + return self.gen_multi_turn_loss_mask_qwen(messages) + elif self.tokenizer_type == "distill_qwen": + return self.gen_multi_turn_loss_mask_distill_qwen(messages) + else: + raise ValueError(f"Unsupported tokenizer type: {self.tokenizer_type}") + + def get_text_from_loss_mask(self, token_ids: List[int], loss_masks: List[int]) -> List[str]: + selected_texts = [] + current_tokens = [] + + for idx, mask in enumerate(loss_masks): + if mask == 1: + current_tokens.append(token_ids[idx]) + elif current_tokens: + selected_texts.append(self.tokenizer.decode(current_tokens)) + current_tokens = [] + + if current_tokens: + selected_texts.append(self.tokenizer.decode(current_tokens)) + + return selected_texts diff --git a/slime_plugins/rollout_buffer/buffer.py b/slime_plugins/rollout_buffer/buffer.py new file mode 100644 index 0000000000..70ada2d870 --- /dev/null +++ b/slime_plugins/rollout_buffer/buffer.py @@ -0,0 +1,611 @@ +# [IMPORTANT] Normalize Process: [raw reward -> normalized reward (only perform on valid reward) -> padded normalized reward] +# Please note we multiply the normalized reward by **group_size / valid_size** to align with GRPO reward +# You can see `normalize_group_data` function at utils.py + +import copy +import glob +import importlib.util +import json +import os +import pathlib +import threading +import time +from datetime import datetime +from typing import Any, Dict, Optional + +import uvicorn +from fastapi import BackgroundTasks, FastAPI, HTTPException, Request +from generator.utils.default_func import ( + default_filter_item, + default_get_group_data_meta_info, + default_is_valid_group, + default_normalize_group_data, + default_pad_group_data, +) +from pydantic import BaseModel + +from tools.visualizer import BufferStatsVisualizer + +app = FastAPI(title="Rollout Buffer Server", debug=True) + +MAX_SIZE = 1000_000_000 + + +def discover_generators(): + """ + Automatically discover generator modules in the generator directory. + Returns a dictionary mapping task_type to module with run_rollout function. + """ + generator_map = {} + generator_dir = pathlib.Path(__file__).parent / "generator" + + # Find all files ending with _generator.py + generator_files = glob.glob(str(generator_dir / "*_generator.py")) + + for file_path in generator_files: + try: + # Load the module + spec = importlib.util.spec_from_file_location("generator_module", file_path) + if spec is None or spec.loader is None: + print(f"Warning: Could not load spec for {file_path}") + continue + + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + # Check if module has TASK_TYPE constant + if not hasattr(module, "TASK_TYPE"): + print(f"Warning: {file_path} does not define TASK_TYPE constant") + continue + + # Check if module has run_rollout function + if not hasattr(module, "run_rollout"): + print(f"Warning: {file_path} does not define run_rollout function") + continue + + task_type = getattr(module, "TASK_TYPE") + generator_info = { + "module": module, + "file_path": file_path, + "run_rollout": getattr(module, "run_rollout"), + } + + # Check for optional functions and use defaults if not present + optional_functions = [ + "normalize_group_data", + "pad_group_data", + "is_valid_group", + "get_group_data_meta_info", + "filter_item", + ] + + for func_name in optional_functions: + if hasattr(module, func_name): + generator_info[func_name] = getattr(module, func_name) + print(f"Found custom {func_name} for {task_type}") + else: + # Use default functions + default_func_name = f"default_{func_name}" + if default_func_name in globals(): + generator_info[func_name] = globals()[default_func_name] + else: + print(f"Warning: No default function found for {func_name}") + + generator_map[task_type] = generator_info + print(f"Discovered generator: {task_type} -> {file_path}") + + except Exception as e: + print(f"Error loading generator from {file_path}: {str(e)}") + continue + + return generator_map + + +@app.middleware("http") +async def set_body_size(request: Request, call_next): + request._body_size_limit = 1_073_741_824 # 1GB + response = await call_next(request) + return response + + +class BufferResponse(BaseModel): + success: bool + message: str = "" + data: Optional[Dict[str, Any]] = None + + +class BufferQueue: + def __init__( + self, + group_size, + min_valid_group_size_ratio=1, + min_valid_item_size_ratio=1, + max_buffer_size=None, + task_type="math", + normalize_group_data_func=None, + pad_group_data_func=None, + is_valid_group_func=None, + get_group_data_meta_info_func=None, + group_timeout_seconds=300, # 5 minutes default timeout + min_timeout_group_size_ratio=0.7, # minimum ratio for timeout groups (default 70%) + filter_item_func=None, + ): + self.data = {} + self.temp_data = {} + self.group_timestamps = {} + self.group_size = group_size + self.min_valid_group_size = int(group_size * min_valid_group_size_ratio) + self.min_valid_item_size = int(group_size * min_valid_item_size_ratio) + self.max_buffer_size = max_buffer_size + self.task_type = task_type + self.group_timeout_seconds = group_timeout_seconds + self.min_timeout_group_size_ratio = min_timeout_group_size_ratio + + # Set up function handlers with defaults + self.normalize_group_data = normalize_group_data_func or default_normalize_group_data + self.pad_group_data_func = pad_group_data_func or default_pad_group_data + self.is_valid_group_func = is_valid_group_func or default_is_valid_group + self.get_group_data_meta_info_func = get_group_data_meta_info_func or default_get_group_data_meta_info + self.filter_item_func = filter_item_func or default_filter_item + + def filter_group_items(self, group_data): + """ + Filter individual items in a group before normalization. + + Args: + group_data (tuple): Tuple of (instance_id, items) + + Returns: + tuple: Filtered (instance_id, items) + """ + instance_id, items = group_data + filtered_items = [item for item in items if self.filter_item_func(item, self.task_type)] + return (instance_id, filtered_items) + + def popleft(self): + if len(self.data) == 0: + return None + return self.data.pop(0) + + def append(self, item): + instance_id = item["instance_id"] + current_time = time.time() + + # Update timestamp for this group + self.group_timestamps[instance_id] = current_time + + if instance_id not in self.temp_data: + self.temp_data[instance_id] = [copy.deepcopy(item)] + else: + self.temp_data[instance_id].append(copy.deepcopy(item)) + + if instance_id not in self.data: + self.data[instance_id] = [item] + else: + self.data[instance_id].append(item) + + def _is_group_timed_out(self, instance_id): + """Check if a group has timed out""" + if instance_id not in self.group_timestamps: + return False + + current_time = time.time() + last_update = self.group_timestamps[instance_id] + return (current_time - last_update) > self.group_timeout_seconds + + def _get_valid_groups_with_timeout(self, del_data=False): + """Get valid groups including timeout-based groups""" + valid_groups = {} + timed_out_groups = {} + finished_groups = [] + + for instance_id, group_data in self.data.items(): + group_size = len(group_data) + is_normally_valid, is_finished = self.is_valid_group_func( + (instance_id, group_data), self.min_valid_group_size, self.task_type + ) + is_timed_out = self._is_group_timed_out(instance_id) + + # Ensure that valid groups are always finished (valid groups ⊆ finished groups) + assert not ( + is_normally_valid and not is_finished + ), f"Group {instance_id} is valid but not finished - this should not happen" + + if is_finished: + finished_groups.append(instance_id) + # If finished and valid, include in processing + if is_normally_valid: + valid_groups[instance_id] = group_data + + continue + + if is_timed_out: + # All timed out groups are considered finished + finished_groups.append(instance_id) + + actual_ratio = group_size / self.group_size + if actual_ratio >= self.min_timeout_group_size_ratio: + # Timed out but has enough items based on ratio, include for processing + timed_out_groups[instance_id] = group_data + + # Remove finished groups and timed out groups with insufficient data + if del_data: + for instance_id in finished_groups: + self.data.pop(instance_id, None) + self.group_timestamps.pop(instance_id, None) + print(f"Removed finished group {instance_id}") + + # Combine normal valid groups and timeout groups + all_valid_groups = {**valid_groups, **timed_out_groups} + + return all_valid_groups, finished_groups + + def get_batch(self, batch_size=1): + output = {"data": [], "meta_info": {}} + + # Get meta information about temp data before processing + meta_info = self.get_group_data_meta_info_func(self.temp_data) + output["meta_info"] = meta_info + + valid_groups, finished_groups = self._get_valid_groups_with_timeout(del_data=True) + batch_count = sum([len(v) for v in valid_groups.values()]) + + output["meta_info"]["finished_groups"] = finished_groups + + print(f"Batch meta info: {json.dumps(meta_info, indent=2)}") + print(f"Found {len(valid_groups)} valid groups and {len(finished_groups)} finished groups") + + valid_groups = list(valid_groups.items()) + + if batch_count < batch_size: + print(f"Not enough valid data: {batch_count} < {batch_size}") + return output + + for instance_id, group in valid_groups: + # First filter individual items + filtered_group = self.filter_group_items((instance_id, group)) + # Only proceed with normalization if we have enough valid items + if len(filtered_group[1]) >= self.min_valid_item_size: + norm_group = self.normalize_group_data(filtered_group) + pad_group = self.pad_group_data_func(norm_group, self.group_size) + output["data"].extend(pad_group[1]) + else: + print( + f"instance_id: {instance_id} has {len(filtered_group[1])} items, which is less than {self.min_valid_item_size}" + ) + + if instance_id in self.data: + self.data.pop(instance_id) + if len(output["data"]) >= batch_size: + break + + return output + + def __len__(self): + valid_groups, _ = self._get_valid_groups_with_timeout() + num = sum([len(v) for v in valid_groups.values()]) + num_of_all_groups = sum([len(v) for v in self.data.values()]) + print(f"valid_groups: {len(valid_groups)}, num: {num}, num_of_all_groups: {num_of_all_groups}") + return num + + +class RolloutBuffer: + def __init__( + self, + group_size=16, + min_valid_group_size_ratio=1, + min_valid_item_size_ratio=1, + max_size=None, + task_type="math", + normalize_group_data_func=None, + pad_group_data_func=None, + is_valid_group_func=None, + get_group_data_meta_info_func=None, + group_timeout_seconds=300, # 5 minutes default + min_timeout_group_size_ratio=0.7, # minimum ratio for timeout groups (default 10%) + filter_item_func=None, + ): + self.buffer = BufferQueue( + group_size=group_size, + min_valid_group_size_ratio=min_valid_group_size_ratio, + min_valid_item_size_ratio=min_valid_item_size_ratio, + max_buffer_size=max_size, + task_type=task_type, + normalize_group_data_func=normalize_group_data_func, + pad_group_data_func=pad_group_data_func, + is_valid_group_func=is_valid_group_func, + get_group_data_meta_info_func=get_group_data_meta_info_func, + group_timeout_seconds=group_timeout_seconds, + min_timeout_group_size_ratio=min_timeout_group_size_ratio, + filter_item_func=filter_item_func, + ) + self.lock = threading.RLock() + self.not_empty = threading.Condition(self.lock) + self.not_full = threading.Condition(self.lock) + self.max_size = max_size + self.total_written = 0 + self.total_read = 0 + self.task_type = task_type + + # Initialize the visualizer + self.visualizer = BufferStatsVisualizer(time_window=60) # 60 second window + # Set args for filename generation + self.visualizer.set_args( + { + "task_type": task_type, + "group_size": group_size, + "num_repeat_per_sample": group_size, + } + ) + + print( + f"set group_size = {group_size}, timeout = {group_timeout_seconds}s, min_timeout_ratio = {min_timeout_group_size_ratio}" + ) + + def write(self, data): + with self.not_full: + while self.max_size and len(self.buffer) >= self.max_size: + print(f"Buffer is full, waiting for space") + self.not_full.wait() + + self.buffer.append(data) + self.total_written += 1 + + # Update visualization stats - just increment the counter for current window + self.visualizer.add_data_point(1) + + self.not_empty.notify_all() + return data + + def read(self, batch_size=-1, wait=True, timeout=10): + with self.not_empty: + if len(self.buffer) < batch_size and wait: + self.not_empty.wait(timeout=timeout) + + if len(self.buffer) == 0: + return {"data": [], "meta_info": {}} + + actual_size = min(batch_size, len(self.buffer)) if batch_size != -1 else len(self.buffer) + # Don't clear temp_data for regular read operations + result = self.buffer.get_batch(batch_size=actual_size) + self.total_read += len(result["data"]) + + self.not_full.notify_all() + return result + + def peek(self, batch_size=1): + with self.lock: + if len(self.buffer) == 0: + return {"data": [], "meta_info": {}} + actual_size = min(batch_size, len(self.buffer)) + # Note: peek doesn't actually get the batch, so we don't get real meta_info + return {"data": list(self.buffer)[:actual_size], "meta_info": {}} + + def get_stats(self): + with self.lock: + return { + "current_size": len(self.buffer), + "max_size": self.max_size, + "total_written": self.total_written, + "total_read": self.total_read, + } + + def count(self): + with self.lock: + return len(self.buffer) + + def close(self): + """Close the buffer and clean up""" + if hasattr(self, "visualizer"): + self.visualizer.close() + + +buffer = RolloutBuffer() + + +@app.post("/buffer/write", response_model=BufferResponse) +async def write_to_buffer(request: Request): + try: + data = await request.json() + item = buffer.write(data) + return BufferResponse( + success=True, + message="Data has been successfully written to buffer", + data={"data": [item], "meta_info": "write to buffer"}, + ) + except Exception as e: + print(f"Write failed: {str(e)}") + import traceback + + traceback.print_exc() + raise HTTPException(status_code=500, detail=f"Write failed: {str(e)}") + + +@app.post("/buffer/read", response_model=BufferResponse) +async def read_from_buffer(request: Request): + data = await request.json() + try: + items = buffer.read( + batch_size=data["batch_size"], + timeout=data["timeout"], + ) + + if not items["data"] and data.get("wait", True): + return BufferResponse( + success=False, + message="Timeout waiting, no data available to read", + data={"data": [], "meta_info": items["meta_info"]}, + ) + + return BufferResponse( + success=True, + message=f"Successfully read {len(items['data'])} items", + data=items, # Return the complete items dictionary + ) + + except Exception as e: + raise HTTPException(status_code=500, detail=f"Read failed: {str(e)}") + + +@app.post("/get_rollout_data", response_model=BufferResponse) +async def get_rollout_data(request: Request): + data = await request.json() + current_size = buffer.count() + + if not "batch_size" in data.keys(): + data["batch_size"] = -1 + + if data["batch_size"] > 0 and current_size < data["batch_size"]: + return BufferResponse( + success=False, + message=f"Not enough data. Requested {data['batch_size']} items but only {current_size} available.", + data={"data": [], "meta_info": {}}, + ) + + try: + # Clear temp_data only for get_rollout_data operations + items = buffer.read(batch_size=data["batch_size"], timeout=600) + except TimeoutError as e: + print(f"TimeoutError: {e}") + + if not items["data"]: + return BufferResponse( + success=False, + message="No data available to read", + data={"data": [], "meta_info": items["meta_info"]}, + ) + + if items["data"]: + print(f"return {len(items['data'])} items and save them to local") + save_data_to_local(items["data"]) + buffer.buffer.temp_data = {} + + return BufferResponse( + success=True, + message=f"Successfully read {len(items['data'])} items", + data=items, + ) + + +def save_data_to_local(data): + try: + save_dir = "rollout_data" + os.makedirs(save_dir, exist_ok=True) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = f"{save_dir}/rollout_data_{timestamp}.json" + with open(filename, "w", encoding="utf-8") as f: + json.dump(data, f, ensure_ascii=False, indent=2) + except Exception as e: + import traceback + + traceback.print_exc() + print(f"Error saving data to local: {str(e)}") + + +def run_rollout(data: dict): + global buffer + # Auto-discover generators + generator_map = discover_generators() + + task_type = data["task_type"] + if task_type not in generator_map: + print(f"Error: No generator found for task_type '{task_type}'") + print(f"Available generators: {list(generator_map.keys())}") + return + + generator_info = generator_map[task_type] + print(f"Using generator: {generator_info['file_path']} for task_type: {task_type}") + + # Extract processing functions from generator + normalize_func = generator_info.get("normalize_group_data") + pad_func = generator_info.get("pad_group_data") + is_valid_func = generator_info.get("is_valid_group") + get_meta_info_func = generator_info.get("get_group_data_meta_info") + filter_item_func = generator_info.get("filter_item") + + if "min_valid_group_size_ratio" not in data.keys(): + data["min_valid_group_size_ratio"] = 1 + + # Add timeout configuration + group_timeout_seconds = data.get("group_timeout_seconds", 300) # 5 minutes default + min_timeout_group_size_ratio = data.get("min_timeout_group_size_ratio", 0.7) + if "min_valid_item_size_ratio" not in data.keys(): + data["min_valid_item_size_ratio"] = 0.7 + + buffer = RolloutBuffer( + max_size=MAX_SIZE, + group_size=int(data["num_repeat_per_sample"]), + min_valid_group_size_ratio=data["min_valid_group_size_ratio"], + min_valid_item_size_ratio=data["min_valid_item_size_ratio"], + min_timeout_group_size_ratio=min_timeout_group_size_ratio, + task_type=task_type, + normalize_group_data_func=normalize_func, + pad_group_data_func=pad_func, + is_valid_group_func=is_valid_func, + get_group_data_meta_info_func=get_meta_info_func, + filter_item_func=filter_item_func, + group_timeout_seconds=group_timeout_seconds, + ) + + try: + # Call the run_rollout function from the appropriate generator module + generator_info["run_rollout"](data) + print(f"Rollout completed successfully for task_type: {task_type}") + except Exception as e: + print(f"Error running rollout for task_type '{task_type}': {str(e)}") + import traceback + + traceback.print_exc() + finally: + # Save the visualization when rollout is complete + buffer.close() + + +@app.post("/start_rollout") +async def start_rollout(request: Request, background: BackgroundTasks): + payload = await request.json() + background.add_task(run_rollout, payload) + return {"message": "Rollout started"} + + +@app.get("/buffer/peek", response_model=BufferResponse) +async def peek_buffer(request: Request): + data = await request.json() + try: + items = buffer.peek(batch_size=data["batch_size"]) + return BufferResponse( + success=True, + message=f"Successfully previewed {len(items['data'])} items", + data=items, # Return the complete items dictionary + ) + except Exception as e: + raise HTTPException(status_code=500, detail=f"Peek method failed: {str(e)}") + + +@app.get("/buffer/stats") +async def get_buffer_stats(): + stats = buffer.get_stats() + return BufferResponse( + success=True, + message="Buffer stats retrieved successfully", + data={ + "data": [stats], + "meta_info": {}, + }, + ) + + +@app.get("/") +async def root(): + return {"message": "Rollout Buffer Server is running"} + + +if __name__ == "__main__": + uvicorn.run( + app, + host="0.0.0.0", + port=8889, + limit_concurrency=1000, # Connection concurrency limit + # limit_max_requests=1000000, # Maximum request limit + timeout_keep_alive=5, # Keep-alive timeout, + ) diff --git a/slime_plugins/rollout_buffer/generator/__init__.py b/slime_plugins/rollout_buffer/generator/__init__.py new file mode 100644 index 0000000000..9bbf2a0d55 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/__init__.py @@ -0,0 +1,10 @@ +from .base_generator import BaseGenerator, query_single_turn +from .reward_utils import get_rule_based_math_reward +from .utils.arguments import add_arguments + +__all__ = [ + "BaseGenerator", + "query_single_turn", + "get_rule_based_math_reward", + "add_arguments", +] diff --git a/slime_plugins/rollout_buffer/generator/base_generator.py b/slime_plugins/rollout_buffer/generator/base_generator.py new file mode 100644 index 0000000000..a55229b56d --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/base_generator.py @@ -0,0 +1,342 @@ +import copy +import json +import random +import time +import uuid +from functools import partial +from multiprocessing import Process, Queue +from time import sleep +from typing import List, Optional + +import requests +from generator.reward_utils.math_utils import get_rule_based_math_reward +from openai import OpenAI +from tqdm import tqdm + +TASK_TYPE = "math" + +SAMPLING_PARAMS = { + "top_p": 1, +} + + +def query_single_turn(client, messages, sampling_params, tools=None): + base_payload = { + "messages": messages, + **sampling_params, + "model": "custom", + "stream": False, + "seed": random.randint(1, 10000000), + "tools": tools, + } + + text = None + accumulated_tokens = 0 + + for attempt in range(6): + try: + # Create a fresh payload for each attempt + current_payload = copy.deepcopy(base_payload) + + if text is not None: + # Update messages with current progress + current_messages = copy.deepcopy(messages) + current_messages.append({"role": "assistant", "content": text}) + current_payload["messages"] = current_messages + + # Adjust max_tokens based on accumulated tokens + if "max_tokens" in sampling_params: + current_payload["max_tokens"] = max(0, sampling_params["max_tokens"] - accumulated_tokens) + + # Add continue flag for partial rollouts + current_payload["extra_body"] = {"continue_final_message": True} + if current_payload["max_tokens"] == 0: + break + response = client.chat.completions.create(**current_payload) + + if len(response.choices) > 0: + if response.choices[0].finish_reason == "abort": + print( + f"query failed, reason: {response.choices[0].finish_reason}, currently generated: {response.usage.completion_tokens}" + ) + + accumulated_tokens += response.usage.completion_tokens + + if text is None: + text = response.choices[0].message.content + else: + text += response.choices[0].message.content + + sleep(10) + continue + if text is None: + text = response.choices[0].message.content + elif response.choices[0].message.content is not None: + text += response.choices[0].message.content + break + else: + print(f"Error in query, status code: {response.status_code}") + continue + except Exception as e: + print(f"query failed in single turn, error: {e}") + continue + + # Update final messages + if len(messages) > 0 and messages[-1]["role"] == "assistant": + messages = messages[:-1] + messages.append({"role": "assistant", "content": text}) + + return messages + + +def worker_process(task_queue, done_queue, rollout_func, reward_func, client, sampling_params): + + for line in iter(task_queue.get, "STOP"): + if isinstance(line, str): + item = json.loads(line) + else: + item = line + + # try: + messages = rollout_func(client, item["prompt"], sampling_params) + + item["uid"] = str(uuid.uuid4()) + item["messages"] = messages + reward = reward_func(item) + item["rollout_index"] = 1 + item["reward"] = reward + item["extra_info"] = {} + item.update(sampling_params) + item["timestamp"] = str(time.time()) + item["round_number"] = len([_ for _ in item["messages"] if _["role"] == "assistant"]) + + output_item = { + "uid": item.pop("uid"), + "messages": messages, + "reward": reward, + "instance_id": item.pop("instance_id"), + "extra_info": item, + } + + done_queue.put(output_item) + + done_queue.put("COMPLETE") + + +class BaseGenerator: + def __init__( + self, + remote_engine_url, + remote_buffer_url, + num_repeat_per_sample=1, + queue_size=1000000, + num_process=10, + task_type="math", + max_tokens=4096, + num_repeats=10, + skip_instance_ids: Optional[List[str]] = None, + ): + self.queue_size = queue_size + self.num_process = num_process + self.remote_engine_url = remote_engine_url + self.remote_buffer_url = remote_buffer_url + self.num_repeat_per_sample = num_repeat_per_sample + self.task_type = task_type + self.max_tokens = max_tokens + self.num_repeats = num_repeats + # Ensure skip_instance_ids is a mutable list (copy to avoid modifying original) + self.skip_instance_ids = list(skip_instance_ids) if skip_instance_ids is not None else None + + if self.skip_instance_ids is not None: + print(f"BaseGenerator initialized with {len(self.skip_instance_ids)} instance_ids to skip") + self.skip_instance_ids = self.skip_instance_ids * self.num_repeat_per_sample + + if "/v1" in remote_engine_url: + self.client = OpenAI(api_key="test", base_url=remote_engine_url) + else: + remote_engine_url = remote_engine_url.strip("/") + "/v1" + self.client = OpenAI(api_key="test", base_url=remote_engine_url) + + def send_data_to_buffer(self, data): + if "/buffer/write" not in self.remote_buffer_url: + remote_buffer_url = self.remote_buffer_url.rstrip("/") + "/buffer/write" + else: + remote_buffer_url = self.remote_buffer_url + + for _ in range(2): + try: + response = requests.post(remote_buffer_url, json=data) + if response.status_code == 200: + break + else: + print(f"send data to buffer failed, status code: {response.status_code}") + continue + except Exception as e: + print(f"send data to buffer failed, error: {e}") + continue + + def run(self, input_file, rollout_func, reward_func): + task_queue, done_queue = Queue(maxsize=self.queue_size), Queue(maxsize=self.queue_size) + + def read_data_into_queue(): + cnt = 0 + items = [] + skipped_count = 0 + with open(input_file, "r") as r: + print("read files") + for line in r: + item = json.loads(line) + items.append(item) + print("read files done: ", len(items)) + random.shuffle(items) + + for _ in range(self.num_repeats): + + for item in items: + for internal_idx in range(self.num_repeat_per_sample): + item_repeat = copy.deepcopy(item) + if "instance_id" not in item_repeat: + raise ValueError(f"instance_id not in item: {item}, the input data must have instance_id") + + if "uid" not in item_repeat: + item_repeat["uid"] = str(uuid.uuid4()) + + # Check if instance_id should be skipped + if self.skip_instance_ids is not None and item_repeat["instance_id"] in self.skip_instance_ids: + print(f"Skipping instance_id: {item_repeat['instance_id']}") + # Remove from skip list to handle potential duplicates in multiple epochs + self.skip_instance_ids.remove(item_repeat["instance_id"]) + skipped_count += 1 + continue + + task_queue.put(item_repeat) + cnt += 1 + time.sleep(300) + + if skipped_count > 0: + remaining_skip_count = len(self.skip_instance_ids) if self.skip_instance_ids is not None else 0 + print( + f"Rollout summary: skipped {skipped_count} instance_ids, {remaining_skip_count} still in skip list" + ) + + for _ in range(self.num_process): + task_queue.put("STOP") + + processes = [] + SAMPLING_PARAMS["max_tokens"] = self.max_tokens + + for _ in range(self.num_process): + process = Process( + target=partial(worker_process, client=self.client, sampling_params=SAMPLING_PARAMS), + args=(task_queue, done_queue, rollout_func, reward_func), + ) + process.start() + processes.append(process) + + process = Process(target=read_data_into_queue) + process.start() + + progress_bar = tqdm() + print("----- GOGOGOGOGOGOGO !!!!!") + + num_finished = 0 + while num_finished < self.num_process: + item = done_queue.get() + if item == "COMPLETE": + num_finished += 1 + else: + # print(f'save {num_save} examples to {output_file}', end='\r') + assert "reward" in item, f"reward not in item: {item}" + assert "instance_id" in item, f"instance_id not in item: {item}" + self.send_data_to_buffer(item) + progress_bar.update(1) + + progress_bar.close() + + return "finished" + + def entry(self, input_file, rollout_func, reward_func, num_epoch=1): + for _ in range(num_epoch): + status = self.run(input_file, rollout_func, reward_func) + + +def run_rollout(data: dict): + + print(f"Starting math rollout with data: {data}") + + rollout_func = query_single_turn + reward_func = get_rule_based_math_reward + + print(f"Waiting for 10 seconds for buffer server to start") + time.sleep(10) + global SAMPLING_PARAMS + for k, v in data["sampling_params"].items(): + SAMPLING_PARAMS[k] = v + print(f"Set {k} to {v}", type(v)) + + generator = BaseGenerator( + data["remote_engine_url"], + data["remote_buffer_url"], + num_repeat_per_sample=int(data["num_repeat_per_sample"]), + queue_size=1000000, + max_tokens=int(data["sampling_params"]["max_tokens"]), + num_process=int(data.get("num_process", 100)), + task_type=data["task_type"], + skip_instance_ids=data.get("skip_instance_ids", None), + ) + + generator.entry(data["input_file"], rollout_func, reward_func, int(data.get("num_epoch", 1))) + + +def normalize_group_data(group, epsilon=1e-8, algo="grpo"): + print(f"Using math-specific normalization for group {group[0]}") + + assert algo == "grpo", "Only 'grpo' is supported for now." + + instance_id = group[0] + data = group[1] + rewards = [item["reward"] for item in data] + + valid_rewards = [r for r in rewards if 1 >= r >= 0] + + if set(valid_rewards) == {0}: + normalized_rewards = rewards + else: + mean_reward = sum(valid_rewards) / len(valid_rewards) + std_reward = (sum((r - mean_reward) ** 2 for r in valid_rewards) / len(valid_rewards)) ** 0.5 + + if std_reward < epsilon: + print(f"[Math Info] Zero variance in group {instance_id}, setting all to 0.") + normalized_rewards = [0.0 if 1 >= r >= 0 else r for r in rewards] + else: + normalized_rewards = [(r - mean_reward) / (std_reward + epsilon) if 1 >= r >= 0 else r for r in rewards] + + for i, item in enumerate(data): + item["reward"] = normalized_rewards[i] + item["raw_reward"] = rewards[i] + + return (instance_id, data) + + +def is_valid_group(group, min_valid_group_size, task_type="math"): + # Handle both tuple and list inputs + if isinstance(group, tuple): + instance_id, items = group + else: + items = group + + # Count valid items (non-empty responses) + valid_indices = [] + for i, item in enumerate(items): + if item["messages"][-1]["content"].strip(): + valid_indices.append(i) + + group_size = len(items) + valid_count = len(valid_indices) + + # A group is finished if it has reached the target size + is_finished = group_size >= min_valid_group_size + + is_valid = is_finished and valid_count >= min_valid_group_size + + return is_valid, is_finished diff --git a/slime_plugins/rollout_buffer/generator/reward_utils/__init__.py b/slime_plugins/rollout_buffer/generator/reward_utils/__init__.py new file mode 100644 index 0000000000..a60effe5a7 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/reward_utils/__init__.py @@ -0,0 +1,3 @@ +from .math_utils import get_rule_based_math_reward + +__all__ = ["get_rule_based_math_reward"] diff --git a/slime_plugins/rollout_buffer/generator/reward_utils/math_utils.py b/slime_plugins/rollout_buffer/generator/reward_utils/math_utils.py new file mode 100644 index 0000000000..f535e5fc69 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/reward_utils/math_utils.py @@ -0,0 +1,557 @@ +# from https://github.com/agentica-project/deepscaler/blob/e6080ccd974eb64bd3430f0b36108244a6fee330/deepscaler/rewards/math_utils/utils.py +""" +Answer checker API that uses sympy to simplify expressions and check for equality. + +Call grade_answer(given_answer: str, ground_truth: str). +""" +import re +from typing import Optional + +import sympy +from pylatexenc import latex2text +from sympy.parsing import sympy_parser + + +# Dan Hendrycks' code +def mathd_normalize_answer(answer: Optional[str]) -> Optional[str]: + if answer is None: + return None + answer = answer.strip() + try: + # Remove enclosing `\text{}`. + # m = re.search("^\\\\text\{(?P.+?)\}$", answer) + m = re.search(r"^\\text\{(?P.+?)\}$", answer) + if m is not None: + answer = m.group("text").strip() + return _strip_string(answer) + except: + return answer + + +def _strip_string(string): + def _fix_fracs(string): + substrs = string.split("\\frac") + new_str = substrs[0] + if len(substrs) > 1: + substrs = substrs[1:] + for substr in substrs: + new_str += "\\frac" + if substr[0] == "{": + new_str += substr + else: + try: + assert len(substr) >= 2 + except: + return string + a = substr[0] + b = substr[1] + if b != "{": + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}{" + b + "}" + post_substr + else: + new_str += "{" + a + "}{" + b + "}" + else: + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}" + b + post_substr + else: + new_str += "{" + a + "}" + b + string = new_str + return string + + def _fix_a_slash_b(string): + if len(string.split("/")) != 2: + return string + a = string.split("/")[0] + b = string.split("/")[1] + try: + a = int(a) + b = int(b) + assert string == "{}/{}".format(a, b) + new_string = "\\frac{" + str(a) + "}{" + str(b) + "}" + return new_string + except: + return string + + def _remove_right_units(string): + # "\\text{ " only ever occurs (at least in the val set) when describing units + if "\\text{ " in string: + splits = string.split("\\text{ ") + assert len(splits) == 2 + return splits[0] + else: + return string + + def _fix_sqrt(string): + if "\\sqrt" not in string: + return string + splits = string.split("\\sqrt") + new_string = splits[0] + for split in splits[1:]: + if split[0] != "{": + a = split[0] + new_substr = "\\sqrt{" + a + "}" + split[1:] + else: + new_substr = "\\sqrt" + split + new_string += new_substr + return new_string + + # linebreaks + string = string.replace("\n", "") + # print(string) + + # remove inverse spaces + string = string.replace("\\!", "") + # print(string) + + # replace \\ with \ + string = string.replace("\\\\", "\\") + # print(string) + + # replace tfrac and dfrac with frac + string = string.replace("tfrac", "frac") + string = string.replace("dfrac", "frac") + # print(string) + + # remove \left and \right + string = string.replace("\\left", "") + string = string.replace("\\right", "") + # print(string) + + # Remove circ (degrees) + string = string.replace("^{\\circ}", "") + string = string.replace("^\\circ", "") + + # remove dollar signs + string = string.replace("\\$", "") + + # remove units (on the right) + string = _remove_right_units(string) + + # remove percentage + string = string.replace("\\%", "") + # string = string.replace("\%", "") + string = string.replace("%", "") # no back-slash needed + + # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string + string = string.replace(" .", " 0.") + string = string.replace("{.", "{0.") + # if empty, return empty string + if len(string) == 0: + return string + if string[0] == ".": + string = "0" + string + + # to consider: get rid of e.g. "k = " or "q = " at beginning + if len(string.split("=")) == 2: + if len(string.split("=")[0]) <= 2: + string = string.split("=")[1] + + # fix sqrt3 --> sqrt{3} + string = _fix_sqrt(string) + + # remove spaces + string = string.replace(" ", "") + + # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b} + string = _fix_fracs(string) + + # manually change 0.5 --> \frac{1}{2} + if string == "0.5": + string = "\\frac{1}{2}" + + # NOTE: X/Y changed to \frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y + string = _fix_a_slash_b(string) + + return string + + +# sympy might hang -- we don't care about trying to be lenient in these cases +BAD_SUBSTRINGS = ["^{", "^("] +# BAD_REGEXES = ["\^[0-9]+\^", "\^[0-9][0-9]+"] +BAD_REGEXES = ["\\^[0-9]+\\^", "\\^[0-9][0-9]+"] +TUPLE_CHARS = "()[]" + + +def _sympy_parse(expr: str): + """Parses an expression with sympy.""" + py_expr = expr.replace("^", "**") + return sympy_parser.parse_expr( + py_expr, + transformations=(sympy_parser.standard_transformations + (sympy_parser.implicit_multiplication_application,)), + ) + + +def _parse_latex(expr: str) -> str: + """Attempts to parse latex to an expression sympy can read.""" + expr = expr.replace("\\tfrac", "\\frac") + expr = expr.replace("\\dfrac", "\\frac") + expr = expr.replace("\\frac", " \\frac") # Play nice with mixed numbers. + expr = latex2text.LatexNodes2Text().latex_to_text(expr) + + # Replace the specific characters that this parser uses. + expr = expr.replace("√", "sqrt") + expr = expr.replace("π", "pi") + expr = expr.replace("∞", "inf") + expr = expr.replace("∪", "U") + expr = expr.replace("·", "*") + expr = expr.replace("×", "*") + + return expr.strip() + + +def _is_float(num: str) -> bool: + try: + float(num) + return True + except ValueError: + return False + + +def _is_int(x: float) -> bool: + try: + return abs(x - int(round(x))) <= 1e-7 + except: + return False + + +def _is_frac(expr: str) -> bool: + return bool(re.search(r"^-?[0-9]+.?/0*[1-9][0-9]*.?$", expr)) + + +def _str_is_int(x: str) -> bool: + try: + x = _strip_properly_formatted_commas(x) + x = float(x) + return abs(x - int(round(x))) <= 1e-7 + except: + return False + + +def _str_to_int(x: str) -> bool: + x = x.replace(",", "") + x = float(x) + return int(x) + + +def _inject_implicit_mixed_number(step: str): + """ + Automatically make a mixed number evalable + e.g. 7 3/4 => 7+3/4 + """ + p1 = re.compile("([0-9]) +([0-9])") + step = p1.sub("\\1+\\2", step) ## implicit mults + return step + + +def _strip_properly_formatted_commas(expr: str): + # We want to be careful because we don't want to strip tuple commas + # p1 = re.compile("(\d)(,)(\d\d\d)($|\D)") + p1 = re.compile(r"(\d)(,)(\d\d\d)($|\D)") + while True: + next_expr = p1.sub("\\1\\3\\4", expr) + if next_expr == expr: + break + expr = next_expr + return next_expr + + +def _normalize(expr: str) -> str: + """Normalize answer expressions.""" + if expr is None: + return None + + # Remove enclosing `\text{}`. + # m = re.search("^\\\\text\{(?P.+?)\}$", expr) + m = re.search(r"^\\text\{(?P.+?)\}$", expr) + if m is not None: + expr = m.group("text") + + expr = expr.replace("\\%", "%") + expr = expr.replace("\\$", "$") + expr = expr.replace("$", "") + expr = expr.replace("%", "") + expr = expr.replace(" or ", " , ") + expr = expr.replace(" and ", " , ") + + expr = expr.replace("million", "*10^6") + expr = expr.replace("billion", "*10^9") + expr = expr.replace("trillion", "*10^12") + + for unit in [ + "degree", + "cm", + "centimeter", + "meter", + "mile", + "second", + "minute", + "hour", + "day", + "week", + "month", + "year", + "foot", + "feet", + "inch", + "yard", + ]: + # expr = re.sub(f"{unit}(es)?(s)? *(\^[0-9]+)?", "", expr) + expr = re.sub(f"{unit}(es)?(s)? *(\\^[0-9]+)?", "", expr) + + # expr = re.sub(f"\^ *\\\\circ", "", expr) + expr = re.sub(f"\\^ *\\\\circ", "", expr) + + if len(expr) > 0 and expr[0] == "{" and expr[-1] == "}": + expr = expr[1:-1] + + expr = re.sub(",\\\\! *", "", expr) + if _is_float(expr) and _is_int(float(expr)): + expr = str(int(round(float(expr)))) + if "\\" in expr: + try: + expr = _parse_latex(expr) + except: + pass + + # edge case with mixed numbers and negative signs + expr = re.sub("- *", "-", expr) + + expr = _inject_implicit_mixed_number(expr) + expr = expr.replace(" ", "") + + # if we somehow still have latex braces here, just drop them + expr = expr.replace("{", "") + expr = expr.replace("}", "") + + # don't be case sensitive for text answers + expr = expr.lower() + + if _str_is_int(expr): + expr = str(_str_to_int(expr)) + + return expr + + +def count_unknown_letters_in_expr(expr: str): + expr = expr.replace("sqrt", "") + expr = expr.replace("frac", "") + letters_in_expr = set([x for x in expr if x.isalpha()]) + return len(letters_in_expr) + + +def should_allow_eval(expr: str): + # we don't want to try parsing unknown text or functions of more than two variables + if count_unknown_letters_in_expr(expr) > 2: + return False + + for bad_string in BAD_SUBSTRINGS: + if bad_string in expr: + return False + + for bad_regex in BAD_REGEXES: + if re.search(bad_regex, expr) is not None: + return False + + return True + + +def are_equal_under_sympy(ground_truth_normalized: str, given_normalized: str): + are_equal = False + try: + expr = f"({ground_truth_normalized})-({given_normalized})" + if should_allow_eval(expr): + sympy_diff = _sympy_parse(expr) + simplified = sympy.simplify(sympy_diff) + if simplified == 0: + are_equal = True + except: + pass + return are_equal + + +def split_tuple(expr: str): + """ + Split the elements in a tuple/interval, while handling well-formatted commas in large numbers + """ + expr = _strip_properly_formatted_commas(expr) + if len(expr) == 0: + return [] + if ( + len(expr) > 2 + and expr[0] in TUPLE_CHARS + and expr[-1] in TUPLE_CHARS + and all([ch not in expr[1:-1] for ch in TUPLE_CHARS]) + ): + elems = [elem.strip() for elem in expr[1:-1].split(",")] + else: + elems = [expr] + return elems + + +def last_boxed_only_string(string): + idx = string.rfind("\\boxed") + if idx < 0: + idx = string.rfind("\\fbox") + if idx < 0: + return None + + i = idx + right_brace_idx = None + num_left_braces_open = 0 + while i < len(string): + if string[i] == "{": + num_left_braces_open += 1 + if string[i] == "}": + num_left_braces_open -= 1 + if num_left_braces_open == 0: + right_brace_idx = i + break + i += 1 + + if right_brace_idx == None: + retval = None + else: + retval = string[idx : right_brace_idx + 1] + + return retval + + +def remove_boxed(s): + left = "\\boxed{" + try: + assert s[: len(left)] == left + assert s[-1] == "}" + return s[len(left) : -1] + except: + return None + + +def extract_boxed_answer(solution: str) -> str: + """Extract the answer from inside a LaTeX \\boxed{} command""" + solution = last_boxed_only_string(solution) + solution = remove_boxed(solution) + return solution + + +def grade_answer_sympy(given_answer: str, ground_truth: str) -> bool: + ground_truth_normalized = _normalize(ground_truth) + given_normalized = _normalize(given_answer) + + if ground_truth_normalized is None: + return False + + if ground_truth_normalized == given_normalized: + return True + + if len(given_normalized) == 0: + return False + + ground_truth_elems = split_tuple(ground_truth_normalized) + given_elems = split_tuple(given_normalized) + + if len(ground_truth_elems) > 1 and ( + ground_truth_normalized[0] != given_normalized[0] or ground_truth_normalized[-1] != given_normalized[-1] + ): + is_correct = False + elif len(ground_truth_elems) != len(given_elems): + is_correct = False + else: + for ground_truth_elem, given_elem in zip(ground_truth_elems, given_elems): + if _is_frac(ground_truth_elem) and _is_frac(given_elem): + # if fractions aren't reduced, then shouldn't be marked as correct + # so, we don't want to allow sympy.simplify in this case + is_correct = ground_truth_elem == given_elem + elif _str_is_int(ground_truth_elem) != _str_is_int(given_elem): + # if the ground truth answer is an integer, we require the given answer to be a strict match (no sympy.simplify) + is_correct = False + else: + is_correct = are_equal_under_sympy(ground_truth_elem, given_elem) + if not is_correct: + break + + return is_correct + + +def grade_answer_mathd(given_answer: str, ground_truth: str) -> bool: + ground_truth_normalized_mathd = mathd_normalize_answer(ground_truth) + given_answer_normalized_mathd = mathd_normalize_answer(given_answer) + + # be at least as lenient as mathd + if ground_truth_normalized_mathd == given_answer_normalized_mathd: + return True + return False + + +def extract_answer(passage: str) -> str: + if "\\boxed" in passage: + return extract_boxed_answer(passage) + return None + + +def grade_answer_verl(solution_str, ground_truth): + if not ground_truth: + return False + if "\\boxed" in ground_truth: + ground_truth = extract_answer(ground_truth) + given_answer = extract_answer(solution_str) + if given_answer is None: + return False + return grade_answer_mathd(given_answer, ground_truth) or grade_answer_sympy(given_answer, ground_truth) + + +def get_deepscaler_rule_based_reward(response, label): + if "" in response: + model_solution = response.split("")[1] + elif "###Response" in response: + model_solution = response.split("###Response")[1] + else: + model_solution = response + + model_answer = extract_answer(model_solution) + if model_answer is None: + return 0 + if label == "": + return 0 + + # Convert single answer to list for uniform processing + assert isinstance(label, (str, float, int)) + ground_truths = [label] + + # Process each ground truth + processed_ground_truths = [] + for truth in ground_truths: + truth = str(truth) + if "\\boxed" in truth: + processed_truth = extract_answer(truth) + if processed_truth is not None: + processed_ground_truths.append(processed_truth) + else: + processed_ground_truths.append(truth) + + if not processed_ground_truths: + return 0 + + # Check against all possible correct answers + for ground_truth in processed_ground_truths: + is_correct = grade_answer_mathd(model_answer, ground_truth) or grade_answer_sympy(model_answer, ground_truth) + if is_correct: + return 1 + + return 0 + + +def get_rule_based_math_reward(item): + messages = item["messages"] + label = item["label"] + assert messages[-1]["role"] == "assistant", "last message must be assistant, but got {}".format( + messages[-1]["role"] + ) + + response = messages[-1]["content"] + if response is None or len(response) == 0: + return 0 + + reward = get_deepscaler_rule_based_reward(response, label) + return reward diff --git a/slime_plugins/rollout_buffer/generator/utils/arguments.py b/slime_plugins/rollout_buffer/generator/utils/arguments.py new file mode 100644 index 0000000000..1ea2d5d4a8 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/utils/arguments.py @@ -0,0 +1,20 @@ +import argparse + + +def add_arguments(add_task_arguments=None): + parser = argparse.ArgumentParser() + parser.add_argument("--task_type", type=str, default="math") + parser.add_argument("--input_file", type=str, default="None") + parser.add_argument("--num_epoch", type=int, default=1) + parser.add_argument("--num_repeat_per_sample", type=int, default=4) + parser.add_argument("--num_process", type=int, default=4) + parser.add_argument("--remote_engine_url", type=str, default="http://0.0.0.0:8000/v1") + parser.add_argument("--remote_buffer_url", type=str, default="http://localhost:8888") + parser.add_argument("--max_tokens", type=int, default=4096) + parser.add_argument("--num_repeats", type=int, default=20) + + if add_task_arguments is not None: + parser = add_task_arguments(parser) + + args = parser.parse_args() + return args diff --git a/slime_plugins/rollout_buffer/generator/utils/default_func.py b/slime_plugins/rollout_buffer/generator/utils/default_func.py new file mode 100644 index 0000000000..37395e92b1 --- /dev/null +++ b/slime_plugins/rollout_buffer/generator/utils/default_func.py @@ -0,0 +1,208 @@ +""" +Default functions for the Rollout Buffer System + +This module contains the default implementations of data processing functions +that are used when generators don't provide their own custom implementations. +""" + +import copy +from typing import Any, Dict, List, Tuple + + +def is_valid_reward(reward: float) -> bool: + """ + Check if a reward value is valid. + + Args: + reward: The reward value to check + + Returns: + bool: True if reward is valid (between 0 and 1), False otherwise + """ + return 1 >= reward >= 0 + + +def default_normalize_group_data(group: Tuple[str, List[Dict[str, Any]]], epsilon=1e-8, algo="grpo"): + """ + Default normalize rewards in a group using z-score normalization. + If all rewards are 0 -> skip normalization. + If std is very small -> set normalized reward to 0. + + Args: + group: (instance_id, [sample_dicts]) + epsilon: Numerical stability parameter + algo: Algorithm type, only "grpo" is supported for now + + Returns: + (instance_id, normalized_data) tuple + """ + assert algo == "grpo", "Only 'grpo' is supported for now." + + instance_id = group[0] + data = group[1] + rewards = [item["reward"] for item in data] + valid_rewards = [r for r in rewards if is_valid_reward(r)] + + if set(valid_rewards) == {0}: + print(f"[Info] All rewards zero in group {instance_id}, skipping normalization.") + normalized_rewards = rewards + else: + mean_reward = sum(valid_rewards) / len(valid_rewards) + std_reward = (sum((r - mean_reward) ** 2 for r in valid_rewards) / len(valid_rewards)) ** 0.5 + + if std_reward < epsilon: + print(f"[Info] Zero variance in group {instance_id} (non-zero constant rewards), setting all to 0.") + normalized_rewards = [0.0 if is_valid_reward(r) else r for r in rewards] + else: + normalized_rewards = [ + (r - mean_reward) / (std_reward + epsilon) if is_valid_reward(r) else r for r in rewards + ] + + for i, item in enumerate(data): + item["reward"] = normalized_rewards[i] + item["raw_reward"] = rewards[i] + + return (instance_id, data) + + +def default_pad_group_data(batch, group_size): + """ + Default padding strategy for group data. + Input batch: (instance_id, [data_1, ..., data_n]) + We multiply the normalized reward by group_size / valid_size to keep the reward range + + Args: + batch: (instance_id, data) tuple + group_size: Target group size + + Returns: + (instance_id, padded_data) tuple + """ + instance_id = batch[0] + data = batch[1] + + # to ensure the padding is equal to dummy padding + for item in data: + item["reward"] = item["reward"] * group_size / len(data) + + pad_count = group_size - len(data) + + assert pad_count <= len(data), "pad_count should be less than or equal to the length of data" + + if pad_count > 0: + print(f"padding {pad_count} items") + data = data + copy.deepcopy(data[:pad_count]) + for i in range(pad_count): + data[i]["reward"] /= 2 + data[-(i + 1)]["reward"] /= 2 + + return (instance_id, data) + + +def default_is_valid_group(group_data, min_valid_group_size, task_type): + """ + Default implementation for checking if a group is valid and finished. + + Logic: + - finished groups are a superset of valid groups + - all valid groups are finished + - some finished groups may not be valid (discarded due to quality issues) + + Args: + group_data: Tuple of (instance_id, items) + min_valid_group_size: Minimum required group size + task_type: Task type for task-specific validation + + Returns: + tuple: (is_valid, is_finished) + """ + instance_id, items = group_data + + group_size = len(items) + reward_list = [item["reward"] for item in items] + + # A group is finished if it has reached the minimum size + is_finished = group_size >= min_valid_group_size + + has_reward_diversity = len(set(reward_list)) > 1 + + is_valid = is_finished and has_reward_diversity + + return is_valid, is_finished + + +def default_get_group_data_meta_info( + temp_data: Dict[str, List[Dict[str, Any]]], +) -> Dict[str, Any]: + """ + Default implementation for getting meta information about the temporary data + collected between get_batch calls. + """ + if not temp_data: + return { + "total_samples": 0, + "num_groups": 0, + "avg_group_size": 0, + "avg_reward": 0, + "reward_std": 0, + "reward_min": 0, + "reward_max": 0, + } + + meta_info = {"total_samples": 0, "num_groups": len(temp_data)} + + all_rewards = [] + all_raw_rewards = [] + # Calculate per-group statistics + for instance_id, samples in temp_data.items(): + group_size = len(samples) + group_rewards = [s["reward"] for s in samples] # Calculate group reward standard deviation + meta_info["total_samples"] += group_size + all_rewards.extend(group_rewards) + # Calculate global statistics + meta_info["avg_group_size"] = meta_info["total_samples"] / meta_info["num_groups"] + + if all_rewards: + meta_info["avg_reward"] = sum(all_rewards) / len(all_rewards) + meta_info["reward_min"] = min(all_rewards) + meta_info["reward_max"] = max(all_rewards) + # Calculate global reward standard deviation + squared_diff_sum = sum((r - meta_info["avg_reward"]) ** 2 for r in all_rewards) + meta_info["reward_std"] = (squared_diff_sum / (len(all_rewards) - 1)) ** 0.5 if len(all_rewards) > 1 else 0 + else: + meta_info["avg_reward"] = 0 + meta_info["reward_std"] = 0 + meta_info["reward_min"] = 0 + meta_info["reward_max"] = 0 + return meta_info + + +def default_filter_item(item: dict, task_type: str) -> bool: + """ + Default function to filter individual items before normalization. + Returns True if the item is valid, False otherwise. + + Args: + item (dict): Single data item to validate + task_type (str): Type of task for task-specific validation + + Returns: + bool: True if item is valid, False otherwise + """ + # Basic validation that all items should have + required_fields = {"instance_id", "reward", "messages"} + if not all(field in item for field in required_fields): + return False + + # Validate reward is a number + if not isinstance(item["reward"], (int, float)): + return False + if item["reward"] < 0 or item["reward"] > 1: + return False + + # Validate messages is a list + if not isinstance(item["messages"], list): + return False + + # Validate messages is a list of dicts + return True diff --git a/slime_plugins/rollout_buffer/tools/assign_instance_id.py b/slime_plugins/rollout_buffer/tools/assign_instance_id.py new file mode 100644 index 0000000000..86ea3a5ea2 --- /dev/null +++ b/slime_plugins/rollout_buffer/tools/assign_instance_id.py @@ -0,0 +1,54 @@ +import argparse +import json +from pathlib import Path + + +def main(input_path, task_type="math", output_path=None): + input_path = Path(input_path) + if output_path is None: + output_path = str(input_path).replace(".jsonl", "_processed.jsonl") + used_ids = set() + processed = [] + + # First pass: load all lines and collect existing instance_ids + with open(input_path, "r", encoding="utf-8") as f: + for line in f: + item = json.loads(line) + if "instance_id" in item: + used_ids.add(item["instance_id"]) + processed.append(item) + + # Second pass: assign missing instance_ids + counter = 0 + for item in processed: + if "instance_id" not in item: + # Find unused id + while True: + candidate_id = f"{task_type}_{counter}" + counter += 1 + if candidate_id not in used_ids: + item["instance_id"] = candidate_id + used_ids.add(candidate_id) + break + + # Save to new jsonl file + with open(output_path, "w", encoding="utf-8") as f: + for item in processed: + f.write(json.dumps(item, ensure_ascii=False) + "\n") + + print(f"✅ Processed {len(processed)} items. Saved to {output_path}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--input_path", type=str, help="Path to input JSONL file") + parser.add_argument( + "--task_type", + type=str, + default="math", + help="Task type prefix for new instance_id", + ) + parser.add_argument("--output_path", type=str, default=None, help="Optional path to output file") + args = parser.parse_args() + + main(args.input_path, args.task_type, args.output_path) diff --git a/slime_plugins/rollout_buffer/tools/visualizer.py b/slime_plugins/rollout_buffer/tools/visualizer.py new file mode 100644 index 0000000000..a238f79ded --- /dev/null +++ b/slime_plugins/rollout_buffer/tools/visualizer.py @@ -0,0 +1,123 @@ +import os +import threading +import time + +import matplotlib.pyplot as plt + + +class BufferStatsVisualizer: + def __init__(self, time_window=60): + """ + Initialize the buffer statistics visualizer + + Args: + time_window (int): Time window in seconds for each data point (default: 60s) + """ + self.time_window = time_window + self.data_points = [] # List to store data points + self.timestamps = [] # List to store timestamps + self.start_time = time.time() + self.last_window_start = self.start_time + self.window_count = 0 # Counter for current window + self.args_dict = None # Store args for filename + + # Initialize the plot + plt.ion() # Enable interactive mode + self.fig, self.ax = plt.subplots(figsize=(12, 6)) + (self.line,) = self.ax.plot([], [], "b-", label="Data Points per Window") + + # Set up the plot + self.ax.set_xlabel("Time (minutes)") + self.ax.set_ylabel("Data Points per 60s Window") + self.ax.set_title("Buffer Statistics - Data Points per 60s Window") + self.ax.grid(True) + self.ax.legend() + + # Start the update thread + self.running = True + self.update_thread = threading.Thread(target=self._update_plot) + self.update_thread.daemon = True + self.update_thread.start() + + def set_args(self, args_dict): + """Set the args dictionary for filename generation""" + self.args_dict = args_dict + + def add_data_point(self, _): + """Add a new data point to the statistics""" + current_time = time.time() + self.window_count += 1 # Increment counter for current window + + # Check if we've reached the end of a time window + if current_time - self.last_window_start >= self.time_window: + # Calculate the time in minutes since start + time_in_minutes = (current_time - self.start_time) / 60 + + # Add the data point and timestamp + self.data_points.append(self.window_count) + self.timestamps.append(time_in_minutes) + + # Save the plot after adding new data point + self.save_plot() + + # Reset for next window + self.last_window_start = current_time + self.window_count = 0 + + def _update_plot(self): + """Update the plot periodically""" + while self.running: + if self.data_points: # Only update if we have data + self.line.set_data(self.timestamps, self.data_points) + self.ax.relim() + self.ax.autoscale_view() + self.fig.canvas.draw() + self.fig.canvas.flush_events() + time.sleep(1) # Update every second + + def save_plot(self): + """Save the current plot to a file""" + timestamp = self.start_time + + # Create filename based on args and timestamp + if self.args_dict: + # Extract key parameters from args + key_params = [] + for key in ["task_type", "num_repeat_per_sample", "group_size"]: + if key in self.args_dict: + key_params.append(f"{key}_{self.args_dict[key]}") + + filename = f"buffer_stats_{'_'.join(key_params)}_{timestamp}.png" + else: + filename = f"buffer_stats_{timestamp}.png" + + # Create directory if it doesn't exist + os.makedirs("buffer_stats", exist_ok=True) + filepath = os.path.join("buffer_stats", filename) + + # Save the plot + plt.savefig(filepath, dpi=300, bbox_inches="tight") + print(f"Plot saved to {filepath}") + + def close(self): + """Close the visualizer and clean up""" + self.running = False + if self.update_thread.is_alive(): + self.update_thread.join() + plt.close(self.fig) + + +# Example usage: +if __name__ == "__main__": + visualizer = BufferStatsVisualizer(time_window=60) + visualizer.set_args({"task_type": "test", "num_repeat_per_sample": 16}) + + # Simulate some data + try: + for i in range(1000): + visualizer.add_data_point(1) # Just increment the counter + time.sleep(0.1) + except KeyboardInterrupt: + pass + finally: + visualizer.close() diff --git a/train_agent_async.py b/train_agent_async.py new file mode 100644 index 0000000000..0cb855dd8c --- /dev/null +++ b/train_agent_async.py @@ -0,0 +1,87 @@ +import ray + +from slime.ray.placement_group import create_actor_group, create_placement_groups, create_rollout_group +from slime.utils.arguments import parse_args + + +def add_my_custom_args(parser): + parser.add_argument( + "--rollout-num-process", + type=int, + default=32, + help="Number of processes to rollout", + ) + parser.add_argument( + "--rollout-num-epoch", + type=int, + default=3, + help="Number of epochs to rollout", + ) + parser.add_argument( + "--rollout-input-file", + type=str, + default=None, + help="Input file for rollout", + ) + return parser + + +def train(args): + # allocate the GPUs + pgs = create_placement_groups(args) + + actor_model = create_actor_group(args, pgs["actor"]) + + # create the rollout generator, with sglang engines inside. + rollout_generator = create_rollout_group(args, pgs["rollout"]) + + # sync the initialization (model initalization, load checkpoint, etc.) + start_rollout_ids = ray.get( + actor_model.async_init(args, role="actor", with_ref=args.kl_coef != 0 or args.use_kl_loss) + ) + assert len(set(start_rollout_ids)) == 1 + if args.start_rollout_id is None: + args.start_rollout_id = start_rollout_ids[0] + + if args.rollout_global_dataset: + ray.get(rollout_generator.data_buffer.load.remote(args.start_rollout_id - 1)) + + # initialize the connection for weight update during training + ray.get(actor_model.async_init_weight_update_connections(rollout_generator)) + + # always update weight first so that sglang has the loaded weights from training. + ray.get(actor_model.async_update_weights()) + + generation_handles = rollout_generator.async_generate(args.start_rollout_id) + # async train loop. + for rollout_id in range(args.start_rollout_id, args.num_rollout): + + ray.get(generation_handles) + actor_model.get_rollout_data(rollout_id) + + actor_handles = actor_model.async_train(rollout_id, with_data_fetching=False) + + ray.get(actor_handles) + if ( + args.update_rollout_weights_interval is not None + and (rollout_id + 1) % args.update_rollout_weights_interval == 0 + ): + ray.get(actor_model.async_update_weights()) + + generation_handles = rollout_generator.async_generate(rollout_id + 1) + + if args.eval_interval is not None and (rollout_id + 1) % args.eval_interval == 0: + ray.get(rollout_generator.async_generate(rollout_id, evaluation=True)) + ray.get(actor_model.async_eval(rollout_id)) + + if args.save_interval is not None and (rollout_id + 1) % args.save_interval == 0: + ray.get(actor_model.async_save_model(rollout_id)) + if args.rollout_global_dataset: + ray.get(rollout_generator.data_buffer.save.remote(rollout_id)) + + +if __name__ == "__main__": + + args = parse_args(add_my_custom_args) + + train(args)