modify qwen2 (#2023)

Author: cui0523

mindnlp/core/ops/array.py

@@ -533,7 +533,7 @@ def range_(start, length):
                 updates = moveaxis(
                     updates, range_(batch_start, batch_size), range(batch_size)
                 )
-            tensor = ops.tensor_scatter_update(tensor, stacked_indices, updates)
+            tensor = ops.tensor_scatter_update(tensor, stacked_indices, updates.to(tensor.dtype))
             if range(len(dims)) != dims:
                 tensor = moveaxis(tensor, range(len(dims)), dims)
             return strided_slice_update(

mindnlp/transformers/generation/beam_search.py

@@ -329,7 +329,7 @@ def finalize(
             decoded[i, : sent_lengths[i]] = hypo
 
             if indices is not None:
-                indices[i, : len(best_idx)] = best_idx
+                indices[i, : len(best_idx)] = mindspore.Tensor(best_idx, dtype=indices.dtype)
 
             if sent_lengths[i] < sent_max_len:
                 # inserting only the first eos_token_id
@@ -832,7 +832,7 @@ def finalize(
             decoded[i, : sent_lengths[i]] = hypo
 
             if indices is not None:
-                indices[i, : len(best_idx)] = best_idx
+                indices[i, : len(best_idx)] = mindspore.Tensor(best_idx, dtype=indices.dtype)
 
             if sent_lengths[i] < sent_max_len:
                 # inserting only the first eos_token_id

mindnlp/transformers/models/qwen2/modeling_qwen2.py

@@ -758,11 +758,12 @@ def forward(
             shift_logits = logits[..., :-1, :]
             shift_labels = labels[..., 1:]
             # Flatten the tokens
-            loss_fct = CrossEntropyLoss()
+            loss_fct = mindspore.ops.SoftmaxCrossEntropyWithLogits()
             shift_logits = shift_logits.view(-1, self.config.vocab_size)
-            shift_labels = shift_labels.view(-1)
+            shift_labels = nn.functional.one_hot(shift_labels.view(-1), self.config.vocab_size)
             # Enable model parallelism
-            loss = loss_fct(shift_logits, shift_labels)
+            loss, _ = loss_fct(shift_logits, shift_labels.to(shift_logits.dtype))
+            loss = loss.mean()
 
         if not return_dict:
             output = (logits,) + outputs[1:]
@@ -934,8 +935,10 @@ def forward(
                 else:
                     loss = loss_fct(pooled_logits, labels)
             elif self.config.problem_type == "single_label_classification":
-                loss_fct = CrossEntropyLoss()
-                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+                loss_fct = mindspore.ops.SoftmaxCrossEntropyWithLogits()
+                labels = nn.functional.one_hot(labels.view(-1), self.num_labels)
+                loss, _ = loss_fct(pooled_logits.view(-1, self.num_labels), labels.to(pooled_logits.dtype))
+                loss = loss.mean()
             elif self.config.problem_type == "multi_label_classification":
                 loss_fct = BCEWithLogitsLoss()
                 loss = loss_fct(pooled_logits, labels)
@@ -1014,8 +1017,10 @@ def forward(
 
         loss = None
         if labels is not None:
-            loss_fct = CrossEntropyLoss()
-            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            loss_fct = mindspore.ops.SoftmaxCrossEntropyWithLogits()
+            labels = nn.functional.one_hot(labels.view(-1), self.num_labels)
+            loss, _= loss_fct(logits.view(-1, self.num_labels), labels.to(logits.dtype))
+            loss = loss.mean()
 
         if not return_dict:
             output = (logits,) + outputs[2:]

tests/transformers/generation/test_utils.py

@@ -1496,7 +1496,7 @@ def test_left_padding_compatibility(self):
         def _prepare_model_kwargs(input_ids, attention_mask, signature):
             model_kwargs = {"input_ids": input_ids, "attention_mask": attention_mask}
             if "position_ids" in signature:
-                position_ids = ops.cumsum(attention_mask, dim=-1) - 1
+                position_ids = ops.cumsum(attention_mask.int(), dim=-1) - 1
                 position_ids = position_ids.masked_fill(attention_mask == 0, 1)
                 model_kwargs["position_ids"] = position_ids
             if "cache_position" in signature:
@@ -3286,4 +3286,4 @@ def test_generate_from_inputs_embeds_with_bos_token_id_is_none(self):
 
         # bos_token_id is required when no input ids nor inputs_embeds is passed
         with self.assertRaises(ValueError):
-            model.generate(max_length=20, bos_token_id=None)
+            model.generate(max_length=20, bos_token_id=None)

tests/transformers/models/qwen2/test_modeling_qwen2.py

@@ -37,7 +37,6 @@
 if is_mindspore_available():
     import mindspore
     from mindnlp.core import ops, nn, no_grad
-
     from mindnlp.transformers import (
         Qwen2ForCausalLM,
         Qwen2ForSequenceClassification,
@@ -482,4 +481,4 @@ def test_speculative_generation(self):
         self.assertEqual(EXPECTED_TEXT_COMPLETION, text)
 
         del model
-        gc.collect()
+        gc.collect()