Merge pull request #187 from InfiniTensor/issue/186

issue/186 support longrope

csrc/models/llama/llama_config.hpp

@@ -7,6 +7,8 @@
 
 #include "../infinilm_model.hpp"
 
+#include <infinicore/nn/rope.hpp>
+
 namespace infinilm::models::llama {
 
 /**
@@ -34,6 +36,8 @@ struct LlamaConfig : public InfinilmModel::Config {
     size_t max_position_embeddings = 2048; // Maximum sequence length
     double rope_theta = 10000.0;           // RoPE base frequency
 
+    std::shared_ptr<infinicore::nn::RoPE::ScalingConfig> rope_scaling = nullptr; // RoPE scaling type
+
     // Normalization
     double rms_norm_eps = 1e-6; // RMSNorm epsilon
 

csrc/models/llama/llama_model.cpp

@@ -34,7 +34,7 @@ LlamaModel::LlamaModel(const LlamaConfig &config,
     // Use GPT-J-style inverse frequencies (default) and GPT_NEOX rotation pairing
     INFINICORE_NN_MODULE_INIT(rotary_emb, config.head_dim, config.max_position_embeddings,
                               config.rope_theta, infinicore::nn::RoPE::Algo::GPT_NEOX,
-                              dtype, device);
+                              dtype, device, config.rope_scaling);
 
     for (auto &layer : layers_) {
         if (layer) {

csrc/pybind11/models/llama.hpp

@@ -6,6 +6,7 @@
 #include "../../models/llama/llama_attention.hpp"
 #include "infinicore/device.hpp"
 #include "infinicore/nn/module.hpp"
+#include "infinicore/nn/rope.hpp"
 #include "infinicore/tensor.hpp"
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
@@ -69,7 +70,8 @@ inline void bind_llama(py::module &m) {
         .def_readwrite("pretraining_tp", &LlamaConfig::pretraining_tp)
         .def_readwrite("name_or_path", &LlamaConfig::name_or_path)
         .def_readwrite("pad_token_id", &LlamaConfig::pad_token_id)
-        .def_property("bos_token_id", [](const LlamaConfig &self) {
+        .def_property(
+            "bos_token_id", [](const LlamaConfig &self) {
                 // Always return as list to match Python config format
                 return py::cast(self.bos_token_id); }, [](LlamaConfig &self, py::object value) {
                 // Accept both single int and list
@@ -80,7 +82,8 @@ inline void bind_llama(py::module &m) {
                 } else {
                     throw py::type_error("bos_token_id must be int or list of ints");
                 } })
-        .def_property("eos_token_id", [](const LlamaConfig &self) {
+        .def_property(
+            "eos_token_id", [](const LlamaConfig &self) {
                 // Always return as list to match Python config format
                 return py::cast(self.eos_token_id); }, [](LlamaConfig &self, py::object value) {
                 // Accept both single int and list
@@ -91,6 +94,86 @@ inline void bind_llama(py::module &m) {
                 } else {
                     throw py::type_error("eos_token_id must be int or list of ints");
                 } })
+        .def_property(
+            "rope_scaling",
+
+            // ---------- getter ----------
+            [](const LlamaConfig &self) -> py::object {
+                if (!self.rope_scaling) {
+                    return py::none();
+                }
+
+                using ScalingConfig = infinicore::nn::RoPE::ScalingConfig;
+                using LongRopeConfig = infinicore::nn::RoPE::LongRopeConfig;
+
+                py::dict d;
+
+                if (auto *lr = dynamic_cast<const LongRopeConfig *>(self.rope_scaling.get())) {
+                    d["type"] = "longrope";
+                    d["rope_type"] = "longrope";
+                    d["factor"] = lr->factor();
+                    d["original_max_position_embeddings"] = lr->original_max_position_embeddings();
+                    d["short_factor"] = lr->short_factor();
+                    d["long_factor"] = lr->long_factor();
+                } else {
+                    throw std::runtime_error("Unknown RoPE scaling type");
+                }
+
+                return std::move(d);
+            },
+
+            // ---------- setter ----------
+            [](LlamaConfig &self, py::object value) {
+                if (value.is_none()) {
+                    self.rope_scaling.reset();
+                    return;
+                }
+
+                if (!py::isinstance<py::dict>(value)) {
+                    throw py::type_error("rope_scaling must be a dict or None");
+                }
+
+                py::dict d = value.cast<py::dict>();
+
+                auto get_str = [&](const char *k) {
+                    if (!d.contains(k)) {
+                        throw py::key_error(k);
+                    }
+                    return py::cast<std::string>(d[k]);
+                };
+
+                std::string type = d.contains("rope_type")
+                                     ? py::cast<std::string>(d["rope_type"])
+                                     : get_str("type");
+
+                if (type == "longrope") {
+                    using LongRopeConfig = infinicore::nn::RoPE::LongRopeConfig;
+
+                    if (!d.contains("short_factor") || !d.contains("long_factor") || !d.contains("original_max_position_embeddings")) {
+                        throw py::value_error(
+                            "longrope requires short_factor, long_factor, "
+                            "original_max_position_embeddings");
+                    }
+
+                    std::vector<float> short_factor = py::cast<std::vector<float>>(d["short_factor"]);
+                    std::vector<float> long_factor = py::cast<std::vector<float>>(d["long_factor"]);
+
+                    size_t original_max_position_embeddings = py::cast<size_t>(d["original_max_position_embeddings"]);
+
+                    float factor = 1.0f;
+                    if (d.contains("factor")) {
+                        factor = py::cast<float>(d["factor"]);
+                    }
+
+                    self.rope_scaling = std::make_shared<LongRopeConfig>(
+                        std::move(short_factor),
+                        std::move(long_factor),
+                        original_max_position_embeddings,
+                        factor);
+                } else {
+                    throw py::value_error("Unsupported rope_scaling type: " + type);
+                }
+            })
         .def("validate", &LlamaConfig::validate)
         .def("kv_dim", &LlamaConfig::kv_dim)
         // Add __dir__ to make attributes discoverable via dir() in Python
@@ -108,6 +191,7 @@ inline void bind_llama(py::module &m) {
             dir_list.append("hidden_act");
             dir_list.append("model_type");
             dir_list.append("rope_theta");
+            dir_list.append("rope_scaling");
             dir_list.append("attention_bias");
             dir_list.append("attention_output_bias");
             dir_list.append("mlp_bias");

test/bench/test_benchmark.py

@@ -368,7 +368,7 @@ def render_ceval(_tokenizer, conversation):
 def render_mmlu(_tokenizer, question, choices):
     """Render MMLU question and choices to input content"""
     choices_text = "\n".join(
-        [f"{chr(65+i)}. {choice}" for i, choice in enumerate(choices)]
+        [f"{chr(65 + i)}. {choice}" for i, choice in enumerate(choices)]
     )
     instruction = (
         "You are a multiple-choice question solver. "
@@ -924,7 +924,9 @@ def test():
                 splits_to_load = (
                     ["test"]
                     if split == "test"
-                    else ["validation"] if split == "val" else ["validation", "test"]
+                    else ["validation"]
+                    if split == "val"
+                    else ["validation", "test"]
                 )
                 # Load each subject individually from hardcoded list, excluding "all"
                 for subject_name in mmlu_subjects:
@@ -946,7 +948,9 @@ def test():
                 splits_to_load = (
                     ["test"]
                     if split == "test"
-                    else ["validation"] if split == "val" else ["validation", "test"]
+                    else ["validation"]
+                    if split == "val"
+                    else ["validation", "test"]
                 )
                 records = []
                 for sp in splits_to_load:
@@ -980,14 +984,13 @@ def test():
     all_results = []
 
     for subj in subject_list:
-        print(f"\n{'='*60}")
+        print(f"\n{'=' * 60}")
         print(f"Evaluating subject: {subj}")
-        print(f"{'='*60}\n")
+        print(f"{'=' * 60}\n")
 
         try:
             samples, actual_subj_name = load_subject_samples(subj)
             print(f"Loaded {len(samples)} samples for subject: {actual_subj_name}")
-
             # Limit number of samples if specified
             if num_samples is not None and num_samples > 0:
                 original_count = len(samples)
@@ -996,37 +999,9 @@ def test():
                     f"Limited to {len(samples)} samples for validation (from {original_count} total)"
                 )
 
-            # Test with first sample if available
-            if len(samples) > 0:
-                sample = samples[0]
-                if benchmark == "ceval":
-                    input_content = f"'question':{sample['question']},'A': {sample['A']}, 'B':{sample['B']}, 'C': {sample['C']},'D': {sample['D']}。"
-                    test_conversation = [
-                        {
-                            "role": "system",
-                            "content": "请从question的A，B，C，D四个选项中选择正确的选项。例如，标准答案：A。",
-                        },
-                        {"role": "user", "content": input_content},
-                    ]
-                    test_output = model.generate(
-                        test_conversation,
-                        max_steps=max_new_tokens,
-                        topp_=1.0,
-                        topk_=1,
-                        temperature_=1.0,
-                    )
-                elif benchmark == "mmlu":
-                    question = sample["question"]
-                    choices = sample["choices"]
-                    test_output = model.generate(
-                        question,
-                        choices,
-                        max_steps=max_new_tokens,
-                        topp_=1.0,
-                        topk_=1,
-                        temperature_=1.0,
-                    )
-                print(f"\nTest output: {test_output}\n")
+            if len(samples) == 0:
+                print(f"No samples found for subject: {actual_subj_name}")
+                continue
 
             # Evaluate samples for this subject
             result = evaluate_samples(
@@ -1044,13 +1019,22 @@ def test():
     model.destroy_model_instance()
 
     # Calculate overall results
+    print(f"\n{'=' * 60}")
+    print("OVERALL RESULTS")
+    print(f"{'=' * 60}")
+    if len(all_results) == 0:
+        print("No tests were run.")
+        return
+    elif len(all_results) > 1:
+        for r in all_results:
+            print(
+                f"Subject '{r['subject']}': {r['correct']}/{r['total']} = {r['accuracy']:.2%}"
+            )
     overall_correct = sum(r["correct"] for r in all_results)
     overall_total = sum(r["total"] for r in all_results)
     overall_accuracy = overall_correct / overall_total if overall_total > 0 else 0.0
 
-    print(f"\n{'='*60}")
-    print("OVERALL RESULTS")
-    print(f"{'='*60}")
+    print(f"{'=' * 60}")
     if benchmark == "ceval":
         print(
             f"Overall 成绩: {overall_correct}/{overall_total} = {overall_accuracy:.2%}"
@@ -1062,7 +1046,7 @@ def test():
 
     print(f"Total Latency: {TOTAL_TIME} seconds")
     print(f"Total Tokens Processed: {TOTAL_TOKENS} tokens")
-    print(f"Overall Throughput: {TOTAL_TOKENS/TOTAL_TIME:.2f} tokens/s")
+    print(f"Overall Throughput: {TOTAL_TOKENS / TOTAL_TIME:.2f} tokens/s")
 
     # Write CSV if output path is specified
     if output_csv: