add code of klotski and fate.

research/SYSU/Fate/.gitignore

@@ -0,0 +1 @@
+model_weights/

research/SYSU/Fate/dataset.py

@@ -0,0 +1,132 @@
+import datasets
+import random
+import os
+
+def get_line_from_dataset(dataset, line):
+    # there are some topics and None in dataset
+    if len(dataset['train'][line]['text']) < 128:
+        line += 1
+        return get_line_from_dataset(dataset, line)
+    return dataset['train'][line]['text']
+
+def get_inputs():
+    if not os.path.exists('datasets/wikitext-103-v1'):
+        os.makedirs('datasets/wikitext-103-v1')
+        # load from HF
+        dataset = datasets.load_dataset('wikitext', 'wikitext-103-v1')
+        dataset.save_to_disk('datasets/wikitext-103-v1')
+    else:
+        # load from local
+        dataset = datasets.load_from_disk('datasets/wikitext-103-v1')
+
+    # get_line = []
+    # for _ in range(num_prompts):
+    #     get_line.append(random.randint(0, 1801350-1))
+
+    # inputs = ()
+    # for line in get_line:
+    #     input_ids = tokenizer(get_line_from_dataset(dataset, line), return_tensors="pt").input_ids
+    #     inputs = inputs + (input_ids,) 
+    # return inputs
+
+# sum
+def get_xsum_inputs():
+    if not os.path.exists('datasets/xsum'):
+        os.makedirs('datasets/xsum')
+        # load from HF
+        dataset = datasets.load_dataset('xsum')
+        dataset.save_to_disk('datasets/xsum')
+    else:
+        # load from local
+        dataset = datasets.load_from_disk('datasets/xsum')
+
+    return dataset['validation']
+
+# sum
+def get_samsum_inputs():
+    if not os.path.exists('datasets/samsum'):
+        os.makedirs('datasets/samsum')
+        # load from HF
+        dataset = datasets.load_dataset('Samsung/samsum')
+        dataset.save_to_disk('datasets/samsum')
+    else:
+        # load from local
+        dataset = datasets.load_from_disk('datasets/samsum')
+
+    return dataset['test']
+
+# translate
+def get_wmt_inputs():
+    if not os.path.exists('datasets/wmt16'):
+        os.makedirs('datasets/wmt16')
+        # load from HF
+        dataset = datasets.load_dataset("wmt16", "ro-en")
+        dataset.save_to_disk('datasets/wmt16')
+    else:
+        # load from local
+        dataset = datasets.load_from_disk('datasets/wmt16')
+
+    return dataset['test']
+
+# mmlu
+def get_mmlu_inputs():
+    if not os.path.exists('datasets/mmlu'):
+        os.makedirs('datasets/mmlu')
+        # load from HF
+        dataset = datasets.load_dataset("lighteval/mmlu", "all")
+        dataset.save_to_disk('datasets/mmlu')
+    else:
+        # load from local
+        dataset = datasets.load_from_disk('datasets/mmlu')
+
+    return dataset['test']
+
+# gsm8k
+def get_gsm8k_inputs():
+    if not os.path.exists('datasets/gsm8k'):
+        os.makedirs('datasets/gsm8k', exist_ok=True)
+        # load from HF
+        config = datasets.DownloadConfig(resume_download=True, max_retries=100)
+        dataset = datasets.load_dataset("gsm8k", "main", download_config=config)
+        dataset.save_to_disk('datasets/gsm8k')
+    else:
+        # load from local
+        dataset = datasets.load_from_disk('datasets/gsm8k')
+
+    return dataset['test']
+
+# ChatGPT-prompts
+def get_ChatGPT_prompts_inputs():
+    if not os.path.exists('datasets/ChatGPT-prompts'):
+        os.makedirs('datasets/ChatGPT-prompts')
+        # load from HF
+        dataset = datasets.load_dataset("MohamedRashad/ChatGPT-prompts")
+        dataset.save_to_disk('datasets/ChatGPT-prompts')
+    else:
+        # load from local
+        dataset = datasets.load_from_disk('datasets/ChatGPT-prompts')
+
+    return dataset['train']
+
+# openai_humaneval
+def get_openai_humaneval_inputs():
+    if not os.path.exists('datasets/openai_humaneval'):
+        os.makedirs('datasets/openai_humaneval')
+        # load from HF
+        dataset = datasets.load_dataset("openai/openai_humaneval")
+        dataset.save_to_disk('datasets/openai_humaneval')
+    else:
+        # load from local
+        dataset = datasets.load_from_disk('datasets/openai_humaneval')
+
+    return dataset['test']
+
+if __name__ == "__main__":
+   dataset = get_openai_humaneval_inputs()
+#    dataset = dataset['validation']
+   for i in range(len(dataset)):
+       if i < 5:
+        # print(dataset[i])
+        print(dataset[i])
+
+

research/SYSU/Fate/expert_ARC_cache.py

@@ -0,0 +1,99 @@
+class ARC_Cache:
+    def __init__(self, max_size):
+        self.max_size = max_size
+        self.p = 0  # 自适应参数，用于调整T1和T2的大小
+
+        self.T1 = []  # 最近使用列表，保存专家ID
+        self.T2 = []  # 频繁使用列表，保存专家ID
+        self.B1 = []  # T1的Ghost列表，保存专家ID
+        self.B2 = []  # T2的Ghost列表，保存专家ID
+
+        for expert_id in range(max_size):
+            self.update(expert_id)
+
+    def is_evicted(self, expert_id):
+        if (expert_id in self.T1) or (expert_id in self.T2):
+            return False
+        else:
+            return True
+
+    def update_list(self, expert_list):
+        evicted_list = []
+        for expert_id in expert_list:
+            evicted_id = self.update(expert_id)
+            if (evicted_id is not None) and (evicted_id not in expert_list) and (evicted_id not in evicted_list):
+                evicted_list.append(evicted_id)
+        return evicted_list
+
+    def update(self, expert_id):
+        if expert_id in self.T1:
+            self.T1.remove(expert_id)
+            self.T2.append(expert_id)
+        elif expert_id in self.T2:
+            self.T2.remove(expert_id)
+            self.T2.append(expert_id)
+        elif expert_id in self.B1:
+            self._adjust_p(min(len(self.T1), self.max_size))
+            evicted_expert_id = self._replace(expert_id)
+            self.B1.remove(expert_id)
+            self.T2.append(expert_id)
+            return evicted_expert_id
+        elif expert_id in self.B2:
+            self._adjust_p(-min(len(self.T2), self.max_size))
+            evicted_expert_id = self._replace(expert_id)
+            self.B2.remove(expert_id)
+            self.T2.append(expert_id)
+            return evicted_expert_id
+        else:
+            evicted_expert_id = None
+            if len(self.T1) + len(self.B1) == self.max_size:
+                if len(self.T1) < self.max_size:
+                    self.B1.pop(0)
+                    evicted_expert_id = self._replace(expert_id)
+                else:
+                    evicted_expert_id = self.T1.pop(0)
+            elif len(self.T1) + len(self.T2) + len(self.B1) + len(self.B2) >= self.max_size:
+                if len(self.T1) + len(self.T2) + len(self.B1) + len(self.B2) >= 2 * self.max_size:
+                    if len(self.B1) > 0:
+                        self.B1.pop(0)
+                    else:
+                        self.B2.pop(0)
+                evicted_expert_id = self._replace(expert_id)
+            self.T1.append(expert_id)
+            return evicted_expert_id
+
+    def _adjust_p(self, delta):
+        self.p = min(self.max_size, max(0, self.p + delta))
+
+    def _replace(self, expert_id):
+        if (len(self.T1) > 0) and ((expert_id in self.B2 and len(self.T1) > self.p) or len(self.T1) > self.p):
+            evicted_expert_id = self.T1.pop(0)
+            self.B1.append(evicted_expert_id)
+        else:
+            if len(self.T2) > 0:
+                evicted_expert_id = self.T2.pop(0)
+                self.B2.append(evicted_expert_id)
+            else:
+                evicted_expert_id = self.T1.pop(0)
+                self.B1.append(evicted_expert_id)
+        return evicted_expert_id
+
+
+if __name__ == "__main__":
+    arc_cache = ARC_Cache(max_size=3)
+    arc_cache.initialize_cache([1,2,3])
+
+    evicted_id = arc_cache.update(4)
+    print(f"Evicted Expert: {evicted_id}")
+
+    evicted_id = arc_cache.update(2)
+    print(f"Evicted Expert: {evicted_id}")
+
+    evicted_id = arc_cache.update(5)
+    print(f"Evicted Expert: {evicted_id}")
+
+    evicted_id = arc_cache.update(3)
+    print(f"Evicted Expert: {evicted_id}")
+
+    evicted_id = arc_cache.update(6)
+    print(f"Evicted Expert: {evicted_id}")

research/SYSU/Fate/main.py

@@ -0,0 +1,44 @@
+import argparse
+from utils import str2bool
+from mindnlp.transformers import AutoTokenizer
+from models.Qwen.modeling_qwen_moe import Qwen2MoeForCausalLM
+import mindspore as ms
+from mindnlp.core import set_default_dtype
+
+def add_parser_arguments(parser):
+    parser.add_argument("--model", type=str, default="Qwen/Qwen1.5-MoE-A2.7B", help="The model name.")
+    parser.add_argument("--path", type=str, default="model_weights", help="The path to the model weights.")
+    parser.add_argument("--early_stopping", type=str2bool, nargs='?', const=True, default=True)
+    parser.add_argument("--min_length", type=int, default=1)
+    parser.add_argument("--max_length", type=int, default=256)
+    parser.add_argument("--pin-weight", type=str2bool, nargs="?", const=True, default=True)
+    parser.add_argument("--memory_budget", type=int, default=0, help="GB")
+    parser.add_argument("--device", type=str, default='0')
+    parser.add_argument("--overlap", type=str2bool, nargs='?', const=True, default=True)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    add_parser_arguments(parser)
+    args = parser.parse_args()
+    set_default_dtype(ms.float16)
+    import mindspore.context as context
+    context.set_context(pynative_synchronize=True)
+
+    model_name = args.model
+    if model_name == "Qwen/Qwen1.5-MoE-A2.7B":
+        tokenizer = AutoTokenizer.from_pretrained("model_weights/qwen1.5-moe-a2.7b/tokenizer")
+        model = Qwen2MoeForCausalLM(args)
+
+    model.eval()
+
+    input_prompt = "Hey, are you conscious? Can you talk to me?"
+    input_tokenizer = tokenizer(input_prompt, return_tensors="ms")
+    input_ids = input_tokenizer.input_ids
+    attention_mask = input_tokenizer.attention_mask
+
+    (output_ids, prefill_time) = model.generate(input_ids, attention_mask=attention_mask, expriment_mode="decoding")
+
+    outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+    print(outputs)
+    # outputs = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    # print(outputs)

research/SYSU/Fate/models/Qwen/configuration_qwen.py

@@ -0,0 +1,92 @@
+from mindnlp.transformers.configuration_utils import PretrainedConfig
+
+class Qwen2MoeConfig(PretrainedConfig):
+
+    model_type = "qwen2_moe"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=2048,
+        intermediate_size=5632,
+        num_hidden_layers=24,
+        num_attention_heads=16,
+        num_key_value_heads=16,
+        hidden_act="silu",
+        max_position_embeddings=8192,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        # tie_word_embeddings=False,
+        rope_theta=1000000.0,
+        use_sliding_window=False,
+        sliding_window=32768,
+        max_window_layers=21,
+        attention_dropout=0.0,
+        decoder_sparse_step=1,
+        moe_intermediate_size=1408,
+        shared_expert_intermediate_size=5632,
+        num_experts_per_tok=4,
+        num_experts=60,
+        norm_topk_prob=False,
+        output_router_logits=False,
+        router_aux_loss_coef=0.001,
+        mlp_only_layers=None,
+        eos_token_id=151643,
+        bos_token_id=151643,
+        pad_token_id=151643,
+        # **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window if use_sliding_window else None
+        self.max_window_layers = max_window_layers
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.attention_dropout = attention_dropout
+        self.eos_token_id = eos_token_id
+        self.bos_token_id = bos_token_id
+        self.pad_token_id = pad_token_id
+
+        # MoE arguments
+        self.decoder_sparse_step = decoder_sparse_step
+        self.moe_intermediate_size = moe_intermediate_size
+        self.shared_expert_intermediate_size = shared_expert_intermediate_size
+        self.num_experts_per_tok = num_experts_per_tok
+        self.num_experts = num_experts
+        self.norm_topk_prob = norm_topk_prob
+        self.output_router_logits = output_router_logits
+        self.router_aux_loss_coef = router_aux_loss_coef
+        self.mlp_only_layers = [] if mlp_only_layers is None else mlp_only_layers
+
+        # super().__init__(
+        #     tie_word_embeddings=tie_word_embeddings,
+        #     **kwargs,
+        # )
+
+def get_Qwen_config(name):
+    if "/" in name:
+        name = name.split("/")[1]
+    name = name.lower()
+
+    if name == "qwen1.5-moe-a2.7b":
+        config = Qwen2MoeConfig()
+    else:
+        raise ValueError(f"Invalid model name: {name}")
+    return config
+
+if __name__ == "__main__":
+   config = get_Qwen_config("qwen1.5-moe-a2.7b")
+   config.offload = True
+   print(config.d_ff)