Test-Time Scaling Evaluation

This repository contains code for the paper "Revisiting the Test-Time Scaling of o1-like Models: Do they Truly Possess Test-Time Scaling Capabilities?"

This code provides scripts to reproduce the experiments and results presented in the paper, including evaluating the test-time scaling capabilities of o1-like models on mathematical reasoning datasets.

🛠️ Environment Setup

To set up the Python environment, please follow these steps:

pip install -r requirements.txt

Alternatively, you can install the latest version of sglang by following the sglang documentation.

📚 Datasets

The evaluation datasets are located in the ./data directory.

• MATH Dataset: Originally from the 'lighteval/MATH' dataset on Hugging Face, which is no longer available. Alternatively, we can download MATH dataset from:
  • DigitalLearningGmbH/MATH-lighteval.
• AIME Dataset: Downloaded from AI-MO/aimo-validation-aime.
• GPQA Dataset: Downloaded from Idavidrein/gpqa.
• Omni-Math Dataset: Downloaded from KbsdJames/Omni-MATH.

🚀 Running Rollout

This section describes how to run the rollout process using the sglang server and client.

1. Launch sglang Server

bash run_server.sh -d ${data-parallel-size} -t ${tensor-parallel-size} -m ${model-name-or-path}

For running R1-full, which requires specific configurations due to its memory requirements, use the following command:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False
python -m sglang.launch_server --dist-init-addr ${MASTER_ADDR}:5000  --model-path ${model_path} --tp 16 --mem-fraction-static 0.8 --chunked-prefill-size 8192 --node-rank ${RANK} --trust-remote-code --nnodes 2

Server Launch Parameters:

• ${data-parallel-size}: Specify the data parallelism size. Adjust based on your GPU resources.
• ${tensor-parallel-size}: Specify the tensor parallelism size. For R1-full, it is recommended to set it to 16.
• ${model-name-or-path}: Path to the model or model identifier from Hugging Face Model Hub.
• ${MASTER_ADDR}: Master address for distributed training.
• ${model_path}: Path to the R1 model.
• ${RANK}: Node rank in distributed training.

2. Launch sglang Client

python rollout.py --dataset ${dataset} --subset test --output_file ${output_path} --n_sample 5 --overwrite --model_class ${model_class} --host ${sglang-server-ip}

Client Parameters:

• ${dataset}: Choose the dataset for evaluation: math, aime, gpqa, or omini.
• ${model_class}: Use deepseep-v3 for R1 models and qwen for Qwen models. For other models, refer to sglang's chat template to find the correct ${model_class}.
• ${sglang-server-ip}: This is the IP address of your sglang server. If running locally, use 127.0.0.1. If running on a remote server, use the server's IP address.

After launching the sglang server and client, the rollout process will begin, generating model predictions for the specified dataset. The output will be saved in the ${output_path} file.

📊 Running Evaluation

To evaluate the rollout results, use the eval.py script:

python eval.py --input_file ${output_file_from_rollout}

${output_file_from_rollout}: Specify the path to the output file generated by the rollout.py script.

The eval.py script outputs the following evaluation metrics:

• Overall Accuracy: The overall accuracy of the model on the dataset.
• Overall Length: The average length of the generated solutions.
• Grouped Accuracy and Length: Accuracy and average length categorized by solution length groups. Solutions are sorted by length, and then grouped into different ranks (e.g., shortest, medium, longest) to analyze performance across different solution lengths.
• Frequency of "wait" and "alternatively": Counts of "wait" and "alternatively" keywords in different solution groups, indicating potential test-time scaling behavior.
• Average Length of Correct and Incorrect Solutions: Compares the average length of correctly answered solutions versus incorrectly answered solutions.

Running eval.py will calculate and print these evaluation metrics to the console, providing insights into the model's performance and test-time scaling behavior.

📈 Running Sequential Scaling Evaluation

To evaluate the sequential-scaling performance of Qwen and R1 models via self-revision, run the seq_search.py script:

python seq_search.py --input_file ${output_file_from_rollout} --search_file ${output_file_of_seq_scaling} --epoch ${epoch} --resume --model_class ${model_class} --host {sglang-server-ip}

Sequential Scaling Parameters:

• ${output_file_from_rollout}: Same as defined in the Rollout section.
• ${output_file_of_seq_scaling}: Specify the output file path for sequential scaling results.
• ${epoch}: Specify the number of sequential scaling iterations (epochs).

⚠️ Memory Considerations: Running sequential scaling is memory-intensive. If you encounter out-of-memory errors, consider increasing the --tp value in your sglang server launch command.

The seq_search.py script will perform sequential scaling evaluation and save the results to ${output_file_of_seq_scaling}.

⚖️ Shortest Majority Vote Implementation

The implementation of the Shortest Majority Vote algorithm is available in the visualize.ipynb notebook.

The notebook compares the performance of majority vote, shortest solution, and shortest majority vote.

def calculate_weighted_majority_vote_accuracy(predictions, lengths, reference_answers):
    correct_count = 0

    for pred, length, ref in zip(predictions, lengths, reference_answers):
        filtered = [(p, l) for p, l in zip(pred, length) if p is not None]
        if not filtered:
            continue

        filtered_pred, filtered_length = zip(*filtered)

        option_data = defaultdict(lambda: {'count': 0, 'lengths': []})
        for p, l in zip(filtered_pred, filtered_length):
            option_data[p]['count'] += 1
            option_data[p]['lengths'].append(l)

        option_weights = []
        for option, data in option_data.items():
            count = data['count']
            avg_length = sum(data['lengths']) / len(data['lengths'])  # 计算length均值
            if avg_length > 0:
                weight = count / math.log(avg_length)
                option_weights.append((option, weight))

        if option_weights:
            option_weights.sort(key=lambda x: x[1], reverse=True)
            best_option = option_weights[0][0]

            if eval_func(best_option, ref):
                correct_count += 1

    accuracy = correct_count / len(reference_answers) if reference_answers else 0
    return accuracy

If you have any questions or issues, please open an issue in this repository.