Support messages-type data for ERNIE-VL (#1371)

Author: BossPi

ernie/dataset/text_sft_reader/finetuning.py

@@ -77,6 +77,7 @@ def __init__(
         simplify=False,
         use_train_part_sharding=False,
         rope_3d=False,
+        chat_template="ernie_vl",
         **kwargs,
     ):
         self.task_group = copy.deepcopy(task_group)
@@ -104,6 +105,7 @@ def __init__(
         self.simplify = simplify
         self.use_train_part_sharding = use_train_part_sharding
         self.rope_3d = rope_3d
+        self.chat_template = chat_template
         self.place = paddle.set_device(device)
 
         # setup special tokens
@@ -331,7 +333,10 @@ def _read_jsonl(self, input_file):
                         if "<think>" in last_tgt and "</think>" in last_tgt:
                             data["prefix"] = ""
                         else:
-                            data["prefix"] = "<think>\n\n</think>\n\n"
+                            if self.chat_template == "ernie_vl_thinking":
+                                data["prefix"] = "\n<think>\n\n</think>\n\n"
+                            else:
+                                data["prefix"] = "<think>\n\n</think>\n\n"
                             data["label"] = [0] * len(data["tgt"])
                             data["label"][-1] = 1
                     else:
@@ -744,6 +749,25 @@ def _convert_example_to_record(self, example, max_seq_length, tokenizer, index):
                 prefix_token = tokenizer.tokenize(example.prefix)
                 cur_tokens = tokens_src + prefix_token + tokens_target
                 extra_loss_mask = [0] * len(prefix_token)
+            elif (
+                "</think>" in tgt.strip() and self.chat_template == "ernie_vl_thinking"
+            ):
+                reasoning_content = (
+                    tgt.strip()
+                    .split("</think>")[0]
+                    .rstrip("\n")
+                    .split("<think>")[-1]
+                    .lstrip("\n")
+                )
+                content = tgt.strip().split("</think>")[-1].lstrip("\n")
+                tokens_target = (
+                    tokenizer.tokenize("\n<think>\n")
+                    + tokenizer.tokenize(reasoning_content.strip("\n"))
+                    + tokenizer.tokenize("\n</think>\n\n")
+                    + tokenizer.tokenize(content)
+                )
+                cur_tokens = tokens_src + tokens_target
+                extra_loss_mask = []
             else:
                 cur_tokens = tokens_src + tokens_target
                 extra_loss_mask = []
@@ -1025,7 +1049,7 @@ def _convert_example_to_record(self, example, max_seq_length, tokenizer, index):
                 # User
                 if "user" in example.messages[index - 1]["role"]:
                     src = example.messages[index - 1]["content"]
-                    tokens_src = self.begin_of_query + tokenizer.tokenize(src)
+                    tokens_src = self.begin_of_query + tokenizer.tokenize(src.strip())
 
                 # Tool
                 if "tool" in example.messages[index - 1]["role"]:
@@ -1041,30 +1065,39 @@ def _convert_example_to_record(self, example, max_seq_length, tokenizer, index):
                     tokens_src = tokens_src + tokenizer.tokenize("\n</tool_output>\n")
 
                 # Assistant
-                if "</think>" in turn["content"]:
+                if "</think>" in turn["content"].strip():
                     reasoning_content = (
                         turn["content"]
+                        .strip()
                         .split("</think>")[0]
                         .rstrip("\n")
                         .split("<think>")[-1]
                         .lstrip("\n")
                     )
-                    content = turn["content"].split("</think>")[-1].lstrip("\n")
+                    content = turn["content"].strip().split("</think>")[-1].lstrip("\n")
                 else:
                     reasoning_content = ""
-                    content = turn["content"]
+                    content = turn["content"].strip()
 
                 tokens_target = []
                 if reasoning_content:
                     tokens_src = tokens_src + self.begin_of_response
-                    tokens_src = tokens_src + tokenizer.tokenize("\n<think>\n")
+                    if self.chat_template == "ernie_vl_thinking":
+                        tokens_target = tokens_target + tokenizer.tokenize(
+                            "\n<think>\n"
+                        )
+                    else:
+                        tokens_target = tokens_target + tokenizer.tokenize("<think>\n")
                     tokens_target = tokens_target + tokenizer.tokenize(
                         reasoning_content.strip("\n")
                     )
                     tokens_target = tokens_target + tokenizer.tokenize("\n</think>\n\n")
                 else:
                     tokens_src = tokens_src + self.begin_of_response
-                    tokens_src = tokens_src + tokenizer.tokenize("\n<think>\n")
+                    if self.chat_template == "ernie_vl_thinking":
+                        tokens_src = tokens_src + tokenizer.tokenize("\n<think>\n")
+                    else:
+                        tokens_src = tokens_src + tokenizer.tokenize("<think>\n")
                     tokens_src = tokens_src + tokenizer.tokenize("\n</think>\n\n")
 
                 if len(content) > 0:
@@ -1122,12 +1155,12 @@ def _convert_example_to_record(self, example, max_seq_length, tokenizer, index):
 
                 previous_cur_len += len(cur_tokens) + len(break_token_multi_turn)
 
-                if len(tokens) <= 4:
-                    return []
+        if len(tokens) <= 4:
+            return []
 
-                if tokens[0] != self.begin_token:
-                    tokens = [self.begin_token] + tokens
-                    loss_mask = [0] + loss_mask
+        if tokens[0] != self.begin_token:
+            tokens = [self.begin_token] + tokens
+            loss_mask = [0] + loss_mask
 
         assert len(tokens) <= self.max_seq_len, f"{len(tokens)}-{self.max_seq_len}"
         assert (
@@ -1157,7 +1190,7 @@ def _convert_example_to_record(self, example, max_seq_length, tokenizer, index):
         assert len(pos_ids) == len(pos_ids_extra)
 
         if sum(loss_mask) == 0:
-            print("[BAD CASE] loss_mask all 0", example.src, example.tgt)
+            print("[BAD CASE] loss_mask all 0", example.messages)
             return []
 
         records = []
@@ -1211,6 +1244,9 @@ def _read_jsonl(self, input_file):
                     ]
                     Example = namedtuple("Example", names)
 
+                if "tools" not in data:
+                    data["tools"] = ""
+
                 # 自动生成label
                 if "label" not in data:
                     data["label"] = []

erniekit/train/ocr_vl_sft/workflow.py

@@ -152,16 +152,6 @@ def run_ocr_vl_sft(
 
         PipelineParallel.timer_printer = lambda _: None
 
-    # checkpoint O1 quantization is open by default.
-    if (
-        not finetuning_args.disable_ckpt_quant
-        and finetuning_args.ckpt_quant_stage == "O0"
-        and not model_args.lora
-    ):
-        finetuning_args.ckpt_quant_stage = "O1"
-    elif finetuning_args.disable_ckpt_quant:
-        finetuning_args.ckpt_quant_stage = "O0"
-
     finetuning_args.resume_from_checkpoint = get_resume_checkpoint_path(finetuning_args)
     if (
         finetuning_args.resume_from_checkpoint is not None

erniekit/train/sft/workflow.py

@@ -129,16 +129,6 @@ def run_sft(
             if finetuning_args.release_grads is True:
                 finetuning_args.release_grads = False
 
-    # checkpoint O1 quantization is open by default.
-    if (
-        not finetuning_args.disable_ckpt_quant
-        and finetuning_args.ckpt_quant_stage == "O0"
-        and not model_args.lora
-    ):
-        finetuning_args.ckpt_quant_stage = "O1"
-    elif finetuning_args.disable_ckpt_quant:
-        finetuning_args.ckpt_quant_stage = "O0"
-
     finetuning_args.print_config(model_args, "Model")
     finetuning_args.print_config(data_args, "Data")
 

erniekit/train/vl_sft/workflow.py

@@ -182,16 +182,6 @@ def run_vl_sft(
 
         PipelineParallel.timer_printer = lambda _: None
 
-    # checkpoint O1 quantization is open by default.
-    if (
-        not finetuning_args.disable_ckpt_quant
-        and finetuning_args.ckpt_quant_stage == "O0"
-        and not model_args.lora
-    ):
-        finetuning_args.ckpt_quant_stage = "O1"
-    elif finetuning_args.disable_ckpt_quant:
-        finetuning_args.ckpt_quant_stage = "O0"
-
     finetuning_args.resume_from_checkpoint = get_resume_checkpoint_path(finetuning_args)
     if (
         finetuning_args.resume_from_checkpoint is not None
@@ -706,6 +696,7 @@ def compute_metrics(p):
                 "max_shot": finetuning_args.max_shot,
                 "use_train_part_sharding": finetuning_args.text_use_train_part_sharding,
                 "rope_3d": model_args.rope_3d,
+                "chat_template": preprocess_args.chat_template,
             }
 
             text_sft_train_reader = create_pyreader(config_dataset_text)

examples/configs/ERNIE-4.5-VL-28B-A3B-Thinking/sft/run_sft_128k.yaml

@@ -17,6 +17,7 @@ moe_aux_loss_lambda: 0.0
 moe_use_aux_free: true
 moe_use_hard_gate: true
 moe_multimodal_dispatch_use_allgather: v2-alltoall-unpad-text
+pp_seg_method: layer:Ernie4_5_DecoderLayer|ErnieDecoderLayer|EmptyLayer
 
 # data
 train_dataset_path: "examples/data/sft_vl-train_demo1.jsonl"

examples/configs/ERNIE-4.5-VL-28B-A3B-Thinking/sft/run_sft_32k.yaml

@@ -17,6 +17,7 @@ moe_aux_loss_lambda: 0.0
 moe_use_aux_free: true
 moe_use_hard_gate: true
 moe_multimodal_dispatch_use_allgather: v2-alltoall-unpad-text
+pp_seg_method: layer:Ernie4_5_DecoderLayer|ErnieDecoderLayer|EmptyLayer
 
 # data
 train_dataset_path: "examples/data/sft_vl-train_demo1.jsonl"

examples/configs/ERNIE-4.5-VL-28B-A3B-Thinking/sft/run_sft_8k.yaml

@@ -17,6 +17,7 @@ moe_aux_loss_lambda: 0.0
 moe_use_aux_free: true
 moe_use_hard_gate: true
 moe_multimodal_dispatch_use_allgather: v2-alltoall-unpad-text
+pp_seg_method: layer:Ernie4_5_DecoderLayer|ErnieDecoderLayer|EmptyLayer
 
 # data
 train_dataset_path: "examples/data/sft_vl-train_demo1.jsonl"

examples/configs/ERNIE-4.5-VL-28B-A3B-Thinking/sft/run_sft_lora_32k.yaml

@@ -18,6 +18,7 @@ moe_aux_loss_lambda: 0.0
 moe_use_aux_free: true
 moe_use_hard_gate: true
 moe_multimodal_dispatch_use_allgather: v2-alltoall-unpad-text
+pp_seg_method: layer:Ernie4_5_DecoderLayer|ErnieDecoderLayer|EmptyLayer
 
 # data
 train_dataset_path: "examples/data/sft_vl-train_demo1.jsonl"

examples/configs/ERNIE-4.5-VL-28B-A3B-Thinking/sft/run_sft_lora_8k.yaml

@@ -18,6 +18,7 @@ moe_aux_loss_lambda: 0.0
 moe_use_aux_free: true
 moe_use_hard_gate: true
 moe_multimodal_dispatch_use_allgather: v2-alltoall-unpad-text
+pp_seg_method: layer:Ernie4_5_DecoderLayer|ErnieDecoderLayer|EmptyLayer
 
 # data
 train_dataset_path: "examples/data/sft_vl-train_demo1.jsonl"

examples/configs/ERNIE-4.5-VL-28B-A3B-Thinking/sft_function_call/run_sft_128k.yaml

@@ -17,6 +17,7 @@ moe_aux_loss_lambda: 0.0
 moe_use_aux_free: true
 moe_use_hard_gate: true
 moe_multimodal_dispatch_use_allgather: v2-alltoall-unpad-text
+pp_seg_method: layer:Ernie4_5_DecoderLayer|ErnieDecoderLayer|EmptyLayer
 
 # data
 dataset_name: "FunctionCallSFTReader"