Add GRPO trainer

.gitignore

@@ -0,0 +1,10 @@
+*pyc
+*pth
+*checkpoint*
+*workspace
+*txt
+core*
+*safetensors
+*swp
+*jsonl
+west-slm/

.vscode/launch.json

@@ -0,0 +1,93 @@
+{
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Train Qwen2-Whisper",
+            "type": "debugpy",
+            "module":"torch.distributed.launch",
+            "request": "launch",
+            "env": {
+                "HF_ENDPOINT": "https://hf-mirror.com",
+                "PYTHONPATH":"${workspaceRoot}:$PYTHONPATH",
+                "CUDA_VISIBLE_DEVICES":"1"
+            },
+            "args": [
+                "--use_env",
+                "--nnodes=1",
+                "--nproc_per_node=1",
+                "train.py",
+                "--grpo",
+                "--llm_model_name_or_path", "Qwen/Qwen2-1.5B-Instruct",
+                "--whisper_model_name_or_path", "tiny",
+                "--temperature","0.5",
+                "--data_path", "/ceph2/user-data/chenzhongliang/west/aishell/train.jsonl",
+                "--bf16", "True",
+                "--projector_model_path", "Qwen-1.5B-Instruct-whisper-tiny/checkpoint-1170/model.safetensors",
+                "--output_dir", "Qwen/Qwen2-1.5B-Instruct-whisper-tiny",
+                "--num_train_epochs", "5",
+                "--per_device_train_batch_size", "4",
+                "--per_device_eval_batch_size", "1",
+                "--gradient_accumulation_steps", "2",
+                "--evaluation_strategy", "no",
+                "--save_strategy", "steps",
+                "--save_steps", "10000",
+                "--save_total_limit", "10",
+                "--learning_rate", "3e-4",
+                "--weight_decay", "0.01",
+                "--adam_beta2", "0.95",
+                "--warmup_ratio", "0.01",
+                "--lr_scheduler_type", "cosine",
+                "--logging_steps", "1",
+                "--report_to", "none",
+                "--model_max_length", "512",
+                "--gradient_checkpointing",
+                "--dataloader_num_workers", "4",
+                "--dataloader_prefetch_factor", "10",
+                "--deepspeed", "ds_config_zero3.json"
+            ],
+            "console": "integratedTerminal",
+            "justMyCode": false
+        },
+        {
+            "name": "export",
+            "type": "debugpy",
+            "module":"torch.distributed.launch",
+            "request": "launch",
+            "env": {
+                "HF_ENDPOINT": "https://hf-mirror.com",
+                "PYTHONPATH":"${workspaceRoot}:$PYTHONPATH",
+                "CUDA_VISIBLE_DEVICES":"1"
+            },
+            "args": [
+                "export.py",
+                "--llm_model_name_or_path", "Qwen/Qwen2-1.5B-Instruct",
+                "--whisper_model_name_or_path", "tiny",
+                "--data_path", "/ceph2/user-data/chenzhongliang/west/aishell/train.jsonl",
+                "--bf16", "True",
+                "--projector_model_path", "Qwen-1.5B-Instruct-whisper-tiny/checkpoint-1170/model.safetensors",
+                "--output_dir", "Qwen/Qwen2-1.5B-Instruct-whisper-tiny",
+                "--num_train_epochs", "5",
+                "--per_device_train_batch_size", "4",
+                "--per_device_eval_batch_size", "1",
+                "--gradient_accumulation_steps", "2",
+                "--evaluation_strategy", "no",
+                "--save_strategy", "steps",
+                "--save_steps", "10000",
+                "--save_total_limit", "10",
+                "--learning_rate", "3e-4",
+                "--weight_decay", "0.01",
+                "--adam_beta2", "0.95",
+                "--warmup_ratio", "0.01",
+                "--lr_scheduler_type", "cosine",
+                "--logging_steps", "1",
+                "--report_to", "none",
+                "--model_max_length", "512",
+                "--gradient_checkpointing",
+                "--dataloader_num_workers", "4",
+                "--dataloader_prefetch_factor", "10",
+            ],
+            "console": "integratedTerminal",
+            "justMyCode": false
+        }
+    ]
+}

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "git.ignoreLimitWarning": true
+}

dataset.py

@@ -12,6 +12,7 @@
 import torchaudio
 import transformers
 import whisper
+from tqdm import tqdm
 
 
 @dataclass
@@ -33,16 +34,34 @@ def __init__(
         tokenizer: transformers.PreTrainedTokenizer,
         config,  # model config
         inference: bool = False,
+        grpo = False,
     ):
         super(SpeechDataset, self).__init__()
         print("Formatting inputs...")
         self.tokenizer = tokenizer
         self.config = config
         self.inference = inference
         self.raw_data = []
+        i = 0
         with open(data_path, "r") as f:
-            for line in f:
-                self.raw_data.append(json.loads(line))
+            for line in tqdm(f):
+                i += 1
+                if i > 100000:
+                    break
+                if not line.startswith('{'):
+                    key, wav, txt, txt2, start, end, dur, u1, _, _, _ = line.split('\t')
+                    obj = {}
+                    obj['wav'] = wav
+                    obj['key'] = key
+                    obj['txt'] = txt.replace('▁',' ')
+                    obj['start'] = round(float(start))
+                    obj['end'] = round(float(end))
+                    self.raw_data.append(obj)
+                else:
+                    self.raw_data.append(json.loads(line))
+        self.grpo = grpo
+        if self.grpo:
+            self.inference = True
 
     def __len__(self):
         return len(self.raw_data)
@@ -123,6 +142,13 @@ def __getitem__(self, i) -> Dict[str, torch.Tensor]:
             'mel': mel,
             'mel_len': mel_len,
         }
+        if self.grpo:
+            ret['prompt'] = instruction
+            if 'key' in msg:
+                ret['key'] = msg['key']
+            ret['txt'] = msg['txt']
+            ret['wav'] = msg['wav']
+
         if not self.inference:
             ret['labels'] = target_ids
             ret['ctc_ids'] = ctc_ids

ds_config_zero3.json

@@ -25,7 +25,7 @@
         "params": {
             "warmup_min_lr": "auto",
             "warmup_max_lr": "auto",
-            "warmup_num_steps": "auto"
+            "warmup_num_steps": 2500
         }
     },
 
@@ -39,12 +39,13 @@
             "device": "none",
             "pin_memory": true
         },
-        "overlap_comm": true,
+        "overlap_comm": false,
         "contiguous_gradients": true,
         "sub_group_size": 1e9,
+        "reduce_scatter": false,
         "reduce_bucket_size": "auto",
         "stage3_prefetch_bucket_size": "auto",
-        "stage3_param_persistence_threshold": "auto",
+        "stage3_param_persistence_threshold": 1e10,
         "stage3_max_live_parameters": 1e9,
         "stage3_max_reuse_distance": 1e9,
         "stage3_gather_16bit_weights_on_model_save": true

export.py

@@ -0,0 +1,24 @@
+from accelerate import Accelerator
+from speech_llm import init_model, ModelArguments
+import transformers
+from dataset import DataArguments
+from train import TrainingArguments
+from trl.models import unwrap_model_for_generation
+
+def export_model(model, output_dir):
+    accelerator = Accelerator()
+    with unwrap_model_for_generation(model, accelerator) as unwrapped_model:
+        unwrapped_model.save_pretrained(output_dir)
+
+if __name__ == '__main__':
+    parser = transformers.HfArgumentParser(
+        (ModelArguments, DataArguments,TrainingArguments ))
+    (
+        model_args,
+        data_args,
+        _
+    ) = parser.parse_args_into_dataclasses()
+
+    model = init_model(model_args)
+    model.freeze_llm()
+    export_model(model, './west-slm')

recognize.py

@@ -11,7 +11,8 @@
 
 from dataset import SpeechDataset, DataArguments
 from speech_llm import init_model, ModelArguments
-
+from transformers import GenerationConfig
+
 
 @dataclass
 class DecodeArguments:
@@ -32,6 +33,7 @@ def main():
     if decode_args.llm_type == 'qwen2':
         eos_token_id = tokenizer.convert_tokens_to_ids(
             ['<|endoftext|>', '<|im_end|>'])
+        decode_args.pad_token_id = tokenizer.pad_token_id
     else:
         tokenizer.pad_token = '<|finetune_right_pad_id|>'
         eos_token_id = tokenizer.convert_tokens_to_ids(
@@ -52,11 +54,20 @@ def main():
         decode_func = model.generate
     else:
         decode_func = model.decode_ctc
+    generation_config = GenerationConfig(
+            do_sample=False,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            max_new_tokens=100,
+            num_beams=1,
+        )
     with torch.no_grad():
         for item in tqdm(data_loader):
             generated_ids = decode_func(**item,
-                                        eos_token_id=eos_token_id,
-                                        decode_config=decode_args)
+                                        decode_config=generation_config,
+                                        repetition_penalty=1.2,
+                                        no_repeat_ngram_size=3,
+                                        )
             text = tokenizer.batch_decode(generated_ids,
                                           skip_special_tokens=True)
             print(text)

requirements.txt

@@ -1,8 +1,15 @@
-deepspeed==0.14.4
-openai-whisper==20231117
-peft==0.12.0
+deepspeed==0.16.4
+openai-whisper
+#==20231117
 tensorboardX==2.6.2.2
-torch>=2.2.2
-torchaudio>=2.2.2
-transformers==4.43.3
-git+https://github.com/wenet-e2e/wenet.git
+#torch
+#>=2.2.2
+#peft>=0.12.0
+#torchaudio
+#>=2.2.2
+transformers==4.49.0
+accelerate>=1.4.0
+#trl>=0.16
+#git+https://github.com/wenet-e2e/wenet.git
+jieba
+editdistance

reward_funcs.py

@@ -0,0 +1,31 @@
+
+import jieba
+import editdistance
+import numpy as np
+# for simple demo
+# def reward_len(completions, **kwargs):
+#     return [-abs(20 - len(completion)) for completion in completions]
+
+def editdistance_score(completions, **kwargs):
+    # return [0] * len(completions)
+    diff = []
+    for hyp, lab in zip(completions,kwargs['txt']):
+        # diff.append(-1.0 * editdistance.eval(hyp, lab) / len(lab))
+        # diff.append(np.log(1e-9 + editdistance.eval(hyp, lab)))
+        diff.append(-editdistance.eval(hyp, lab))
+    return diff
+
+def word_count(completions, **kwargs):
+    c = []
+    for i, completion in enumerate(completions):
+        words = jieba.cut(completion.strip())
+        words = [w.strip() for w in words]
+        words = [w for w in words if w != '']
+        if i==0:
+            print(kwargs['wav'][i],words)
+        c.append(-abs(0 - len(completion)))
+    # print('-'*120)
+    return c
+    # return [len(list(jieba.cut(completion))) for completion in completions]
+
+active_reward_func = [editdistance_score]