support jetmoe & fix python id() caused bugs (#998)

Author: lvyufeng

llm/inference/jetmoe/run_jetmoe.py

@@ -0,0 +1,21 @@
+import mindspore
+from mindnlp.transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+# Initialize the model and tokenizer
+model_name = "jetmoe/jetmoe-8b-chat"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, ms_dtype=mindspore.float16)
+# Encode input context
+messages = [
+    {
+        "role": "system",
+        "content": "You are a friendly chatbot",
+    },
+    {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
+ ]
+tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="ms")
+print(tokenized_chat)
+# Generate text
+output = model.generate(tokenized_chat, max_length=500, num_return_sequences=1, no_repeat_ngram_size=2)
+# Decode the generated text
+generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+print(generated_text)

mindnlp/diffusers/__init__.py

@@ -19,7 +19,9 @@
 from typing import OrderedDict
 from functools import reduce, partial
 import math
+from uuid import uuid4
 from packaging import version
+
 import numpy as np
 import mindspore
 import mindspore.common.dtype as mstype
@@ -410,6 +412,12 @@ def _initialize(self, init_method):
 
 Parameter.initialize = _initialize
 
+old_param_init = Parameter.__init__
+def _param_new_init(self, default_input, name=None, requires_grad=True, layerwise_parallel=False, parallel_optimizer=True):
+    old_param_init(self, default_input, name, requires_grad, layerwise_parallel, parallel_optimizer)
+    self.uuid = uuid4().hex
+
+Parameter.__init__ = _param_new_init
 
 old_repeat = Tensor.repeat
 def new_repeat_interleave(input, repeats, axis=None):

mindnlp/diffusers/loaders/__init__.py

@@ -1048,17 +1048,16 @@ def empty_initializer(init, shape=None, dtype=mindspore.float32):
 
         # These are all the pointers of shared tensors.
         tied_params = [names for _, names in ptrs.items() if len(names) > 1]
-
         def load_ckpt(resolved_archive_file):
             if 'ckpt' not in resolved_archive_file:
                 if use_safetensors or 'safetensors' in resolved_archive_file:
                     from safetensors.numpy import load_file
                     origin_state_dict = load_file(resolved_archive_file)
                     if use_fp16:
                         logger.warning_once("MindSpore do not support bfloat16 dtype, we will automaticlly convert to float16")
-                    state_dict = {k: Parameter(v.astype(usage_dtype)) for k, v in origin_state_dict.items()}
+                    new_state_dict = {k: Parameter(Tensor.from_numpy(v.astype(usage_dtype))) for k, v in origin_state_dict.items()}
                 else:
-                    state_dict = load(resolved_archive_file)
+                    new_state_dict = load(resolved_archive_file)
             else:
                 try:
                     state_dict = load_checkpoint(str(resolved_archive_file))
@@ -1067,12 +1066,12 @@ def load_ckpt(resolved_archive_file):
                         f"Unable to load weights from mindspore checkpoint file '{resolved_archive_file}'. "
                     ) from exc
 
-            new_state_dict = {}
-            for key, value in state_dict.items():
-                key = key.replace('gamma', 'weight').replace('beta', 'bias').replace('embedding_table', 'weight')
-                value.name = value.name.replace('gamma', 'weight').replace('beta', 'bias')\
-                    .replace('embedding_table', 'weight')
-                new_state_dict[key] = value
+                new_state_dict = {}
+                for key, value in state_dict.items():
+                    key = key.replace('gamma', 'weight').replace('beta', 'bias').replace('embedding_table', 'weight')
+                    value.name = value.name.replace('gamma', 'weight').replace('beta', 'bias')\
+                        .replace('embedding_table', 'weight')
+                    new_state_dict[key] = value
             return new_state_dict
 
         keys_missing = list(model.parameters_dict().keys())
@@ -1114,7 +1113,7 @@ def load_param_into_net(model: nn.Cell, param_dict: dict, prefix: str, dtype_gro
                 else:
                     param_name = pname_in_net
 
-                if id(param) in param_id_set:
+                if param.uuid in param_id_set:
                     # for tied params
                     if param_name in keys_unexpected:
                         keys_unexpected.remove(param_name)
@@ -1161,7 +1160,7 @@ def load_param_into_net(model: nn.Cell, param_dict: dict, prefix: str, dtype_gro
                         param.set_data(new_param)
                     keys_unexpected.remove(param_name)
                     keys_missing.remove(pname_in_net)
-                    param_id_set.add(id(param))
+                    param_id_set.add(param.uuid)
                 else:
                     # fix missing value parameter dtype cast.
                     if ms_dtype and ms_dtype != param.dtype:
@@ -1358,7 +1357,7 @@ def num_parameters(self, only_trainable=False):
         total = 0
         param_set = set()
         for param in self.get_parameters():
-            param_id = id(param)
+            param_id = param.uuid
             if param_id not in param_set and (only_trainable or param.requires_grad):
                 total += param.size
             param_set.add(param_id)

mindnlp/diffusers/models/__init__.py

@@ -71,6 +71,7 @@
     gpt_pangu,
     graphormer,
     hubert,
+    jetmoe,
     layoutlm,
     layoutlmv2,
     llama,
@@ -163,6 +164,7 @@
 from .gpt2 import *
 from .graphormer import *
 from .hubert import *
+from .jetmoe import *
 from .layoutlm import *
 from .layoutlmv2 import *
 from .llama import *
@@ -255,6 +257,7 @@
 __all__.extend(gpt2.__all__)
 __all__.extend(graphormer.__all__)
 __all__.extend(hubert.__all__)
+__all__.extend(jetmoe.__all__)
 __all__.extend(layoutlm.__all__)
 __all__.extend(layoutlmv2.__all__)
 __all__.extend(llama.__all__)