[NPU] InternVL2 supports npu training (PaddlePaddle#714)

Co-authored-by: LokeZhou <aishenghuoaiqq@163.com>
Co-authored-by: nifeng <nemonameless@qq.com>

Author: 3 people

paddlemix/datasets/internvl_dataset.py

@@ -532,6 +532,7 @@ def preprocess_internlm(
     ).input_ids
     targets = input_ids.clone()
 
+    new_targets = []
     # print('tokenizer.pad_token_id:\n', tokenizer.pad_token_id) # 151643
     # print('targets', targets, targets.shape, targets.sum().item())
     # [[151644, 8948  , 198   , ..., 103978, 1773  , 151645]]   [1, 1918]   281157253
@@ -569,10 +570,14 @@ def preprocess_internlm(
                 target[:] = IGNORE_TOKEN_ID
                 print(f'WARNING: tokenization mismatch: {cur_len} vs. {total_len}. This dataset is {ds_name}.')
                 sys.stdout.flush()
+        
+        new_targets.append(target)
+    
+    new_targets = paddle.stack(new_targets, axis=0)
 
     return dict(
         input_ids=input_ids,
-        labels=targets,
+        labels=new_targets,
         attention_mask=input_ids.not_equal(paddle.to_tensor(tokenizer.pad_token_id)),
     )
 

paddlemix/examples/internvl2/README.md

@@ -53,6 +53,8 @@ sh paddlemix/examples/internvl2/shell/internvl2.0/2nd_finetune/internvl2_2b_inte
 sh paddlemix/examples/internvl2/shell/internvl2.0/2nd_finetune/internvl2_8b_internlm2_7b_dynamic_res_2nd_finetune_full.sh
 ```
 
+## 5 NPU硬件训练
+请参照[tools](../../tools/README.md)进行NPU硬件Paddle安装和环境变量设置，配置完成后可直接执行微调命令进行训练或预测。
 
 ### 参考文献
 ```BibTeX

paddlemix/examples/llava/README.md

@@ -104,6 +104,9 @@ python paddlemix/tools/supervised_finetune.py paddlemix/config/llava/v1_5/lora_s
 python paddlemix/tools/supervised_finetune.py paddlemix/config/llava/v1_5/sft_argument.json
 ```
 
+## 5 NPU硬件训练
+请参照[tools](../../tools/README.md)进行NPU硬件Paddle安装和环境变量设置，配置完成后可直接执行微调命令进行训练或预测。
+
 ### 参考文献
 ```BibTeX
 @misc{liu2024llavanext,

paddlemix/models/internvl2/internvl_chat/modeling_intern_vit.py

@@ -19,6 +19,7 @@
 # --------------------------------------------------------
 
 from typing import Optional, Tuple, Union
+import numpy as np
 
 import paddle
 import paddle.nn as nn
@@ -63,7 +64,7 @@ def drop_path(x, drop_prob: float = 0.0, training: bool = False, scale_by_keep:
         return x
     keep_prob = 1 - drop_prob
     shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
-    random_tensor = paddle.bernoulli(paddle.full(shape, keep_prob, dtype=x.dtype))
+    random_tensor = paddle.bernoulli(paddle.full(shape, keep_prob, dtype='float32')).cast(x.dtype)
     if keep_prob > 0.0 and scale_by_keep:
         random_tensor = paddle.divide(random_tensor, paddle.to_tensor(keep_prob))
     return x * random_tensor
@@ -386,9 +387,9 @@ def forward(
             hidden_states = hidden_states.cast(original_dtype)
             return hidden_states
         else:
-            hidden_states = hidden_states + self.drop_path1(self.attn(self.norm1(hidden_states.cast("bfloat16"))) * self.ls1)
-            hidden_states = hidden_states + self.drop_path2(self.mlp(self.norm2(hidden_states.cast("bfloat16"))) * self.ls2)
-            return hidden_states.cast("bfloat16")
+            hidden_states = hidden_states + self.drop_path1(self.attn(self.norm1(hidden_states)) * self.ls1)
+            hidden_states = hidden_states + self.drop_path2(self.mlp(self.norm2(hidden_states)) * self.ls2)
+            return hidden_states
 
 
 class InternVisionEncoder(nn.Layer):
@@ -405,7 +406,7 @@ def __init__(self, config: InternVisionConfig):
         super().__init__()
         self.config = config
         # stochastic depth decay rule
-        dpr = [x.item() for x in paddle.linspace(0, config.drop_path_rate, config.num_hidden_layers)]
+        dpr = [x.item() for x in np.linspace(0, config.drop_path_rate, config.num_hidden_layers)]
         self.layers = nn.LayerList([
             InternVisionEncoderLayer(config, dpr[idx]) for idx in range(config.num_hidden_layers)])
         self.gradient_checkpointing = True

paddlemix/models/internvl2/internvl_chat/modeling_internvl_chat.py

@@ -31,7 +31,7 @@
 from paddlenlp.generation import GenerationConfig
 from paddlenlp.transformers import LlamaForCausalLM, Qwen2ForCausalLM
 from paddlenlp.transformers.model_outputs import CausalLMOutputWithPast
-from paddlemix.models.model_utils import MixPretrainedModel
+from paddlemix.models.model_utils import MixPretrainedModel, NPUCrossEntropyLoss
 from paddlenlp.transformers.model_utils import PretrainedModel
 from .configuration_internvl_chat import InternVLChatConfig
 from .modeling_intern_vit import InternVisionModel
@@ -201,7 +201,7 @@ def forward(
             shift_logits = logits[..., :-1, :]
             shift_labels = labels[..., 1:]
             # Flatten the tokens
-            loss_fct = CrossEntropyLoss()
+            loss_fct = NPUCrossEntropyLoss() if "npu" in paddle.get_device() else CrossEntropyLoss()
             shift_logits = shift_logits.reshape([-1, self.language_model.config.vocab_size])
             shift_labels = shift_labels.reshape([-1])
             # Enable model parallelism

paddlemix/models/model_utils.py

@@ -27,7 +27,7 @@
 from paddlemix.utils.env import MODEL_HOME
 from paddlemix.utils.log import logger
 
-__all__ = ["MixPretrainedModel"]
+__all__ = ["MixPretrainedModel", "NPUCrossEntropyLoss"]
 
 
 def resolve_cache_dir(pretrained_model_name_or_path: str, cache_dir: Optional[str] = None) -> str:
@@ -489,3 +489,26 @@ def _find_mismatched_keys(
             )
 
         return missing_keys, unexpected_keys, mismatched_keys
+
+
+class NPUCrossEntropyLoss(paddle.nn.Layer):
+    """
+    Make cross_entropy_loss compatible with npu device
+    """
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.reduction = kwargs.get('reduction', 'mean')
+        kwargs['reduction'] = 'none'
+        self.nll_loss = paddle.nn.NLLLoss(**kwargs)
+        self.log_softmax = paddle.nn.functional.log_softmax
+    
+    def forward(self, logits, labels):
+        loss = self.nll_loss(self.log_softmax(logits, axis=-1), labels)
+        if self.reduction == 'mean':
+            return loss.mean()
+        elif self.reduction == 'sum':
+            return loss.sum()
+        elif self.reduction == 'none':
+            return loss
+        else:
+            raise ValueError(f"Unexcepted reduction method: {self.reduction}")

paddlemix/tools/README.md

@@ -193,7 +193,9 @@ python paddlemix/tools/merge_lora_params.py \
 
 **NPU硬件训练**
 
-PaddleMIX支持在NPU硬件上训练，只需要在config配置文件中增加`device`字段制定设备即可：
+PaddleMIX支持在NPU硬件上进行训练：
+1. 请先参照[PaddleCustomDevice](https://github.com/PaddlePaddle/PaddleCustomDevice/blob/develop/backends/npu/README_cn.md)安装NPU硬件Paddle
+2. 在config配置文件中增加`device`字段指定设备：
 ```json
 {
     ...
@@ -203,4 +205,13 @@ PaddleMIX支持在NPU硬件上训练，只需要在config配置文件中增加`d
     ...
 }
 ```
+3. 启动训练前请设置如下环境变量用于性能加速和精度对齐
+```shell
+export FLAGS_use_stride_kernel=0
+export FLAGS_npu_storage_format=0 # 关闭私有格式
+export FLAGS_npu_jit_compile=0 # 关闭即时编译
+export FLAGS_npu_scale_aclnn=True # aclnn加速
+export FLAGS_npu_split_aclnn=True # aclnn加速
+export CUSTOM_DEVICE_BLACK_LIST=set_value,set_value_with_tensor # set_value加入黑名单
+```
 目前支持NPU训练的模型可以参考此[文档](../examples/README.md)