Speech RL recipe update

rl-tutorial/cosyvoice_llm/Dockerfile

@@ -0,0 +1,6 @@
+FROM verlai/verl:app-verl0.4-vllm0.8.5-mcore0.12.2-te2.2
+COPY requirements-cosyvoice.txt /myworkspace/requirements.txt
+RUN pip install -r /myworkspace/requirements.txt
+RUN pip install -U nvidia-pytriton
+RUN git clone https://github.com/yuekaizhang/verl.git /myworkspace/verl -b thread && cd /myworkspace/verl && pip install --no-deps -e .
+RUN git clone https://github.com/yuekaizhang/PytritonSenseVoice.git /myworkspace/PytritonSenseVoice && cd /myworkspace/PytritonSenseVoice && pip install -e .

rl-tutorial/cosyvoice_llm/README.md

@@ -132,9 +132,10 @@ The script synthesizes a tongue-twister using the merged checkpoint and prints t
 | Model                                                 | Seed-TTS `test_zh` CER | Cosyvoice3 `zero_shot_zh` |Comment                                                                        |
 |-|------------------------------------------------------|------------------------|--------------------------------------------------------------------------------|
 | Official CosyVoice2 LLM                               | 1.45 %             |4.08%| See the [paper](https://arxiv.org/abs/2412.10117)                              |
-| SFT (initialized from Qwen2-0.5B-Instruct)            | 1.81 %                 |4.83%| See [PR #1887](https://github.com/k2-fsa/icefall/pull/1887)                    |
-| GRPO (this work, trained on AIShell-3)                | **1.06 %**             |4.03%|                                                                                |
-
+| SFT (initialized from Qwen2-0.5B-Instruct)            | 1.70 %                 |4.26%| See [PR #1887](https://github.com/k2-fsa/icefall/pull/1887)                    |
+| GRPO (this work, trained on AIShell-3)                | 1.06 %             |3.01%| [Commit](https://github.com/nvidia-china-sae/mair-hub/commit/5659ee4d128d5902f2f1a2abb333bdd2e387268d) |
+| GRPO (this work, trained on emilia_zh subset, using top_p=1, temperature=1.0)  |    0.87%          |2.63% (1800 steps)| |
+| DAPO (this work, trained on emilia_zh subset, using top_p=1, temperature=1.0)  |    0.83%          |2.71% (700 steps)| See `run_dapo.sh`|
 ## Acknowledgement
 
 This work is inspired by the implementation in  

rl-tutorial/cosyvoice_llm/filter_data.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+from tqdm import tqdm
+import re
+
+def load_jsonl(file_path: str):
+    """Load data from jsonl file."""
+    data = []
+    count = 0
+    with open(file_path, 'r', encoding='utf-8') as f:
+        for line in tqdm(f):
+            item = json.loads(line.strip())
+            if item["language"] == "zh":
+                # check if there is any english in the text
+                if item["duration"] < 30:
+                    count += 1
+                    item["text"] = item["text"].lower()
+                    if re.search(r'[a-z]', item["text"]):
+                        print(item["text"])
+                        continue
+                    else:
+                        data.append(item)
+            if count > 80000:
+                break
+    print(f"Total data: {len(data)}")
+    return data
+
+if __name__ == "__main__":
+    jsonl_file = "data/emilia_zh.jsonl"
+    data = load_jsonl(jsonl_file)
+    with open(f"./data/{jsonl_file.split('/')[-1].split('.')[0]}-zh-filtered.jsonl", "w", encoding="utf-8") as f:
+        for item in data:
+            f.write(json.dumps(item, ensure_ascii=False) + "\n")

rl-tutorial/cosyvoice_llm/infer_dataset.py

@@ -52,10 +52,6 @@
 except RuntimeError:
     pass
 
-
-TEMPLATE = "{% for message in messages %}{%- if message['role'] == 'user' %}{{- '<|im_start|>' + message['role'] + '\n' + 'Convert the text to speech: ' + message['content'] + '<|im_end|>\n'}}{%- elif message['role'] == 'assistant' %}{{- '<|im_start|>' + message['role'] + '\n' + '<|SPEECH_GENERATION_START|>' + message['content']}}{%- endif %}{%- endfor %}"
-
-
 def audio_decode_cosyvoice2(
     audio_tokens, prompt_text, prompt_speech_16k, codec_decoder
 ):
@@ -197,6 +193,12 @@ def data_collator(batch, tokenizer, s3_tokenizer):
         prompt_text_list.append(prompt_text)
         # Combine prompt and target text
         full_text = prompt_text + target_text
+        # remove the unnecessary punctuation for cosyvoice3 zero_shot_zh dataset
+        puncts = ['"', '(', ')', '“', '”', '‘', '（', '）', '\'']
+        for p in puncts:
+            if p in full_text:
+                full_text = full_text.replace(p, '')
+                print(f"removed {p} from {full_text}")
 
         # get prompt audio for CosyVoice2 (convert to 16kHz)
         ref_audio_org, ref_sr = (
@@ -234,8 +236,9 @@ def data_collator(batch, tokenizer, s3_tokenizer):
             {"role": "user", "content": full_text},
             {"role": "assistant", "content": prompt_audio_cosy2_id_str}
         ]
-        if 'system' in tokenizer.chat_template:
-            tokenizer.chat_template = TEMPLATE
+
+        assert 'system' not in tokenizer.chat_template, "system is not allowed in the chat template"
+
         input_ids = tokenizer.apply_chat_template(
             chat,
             tokenize=True,
@@ -301,7 +304,7 @@ def main():
         prompt_speech_16k = load_wav(args.prompt_speech_path, 16000)
     else:
         prompt_speech_16k = None
-    s3_tokenizer = s3tokenizer.load_model("speech_tokenizer_v2_25hz").to(device) if 'zero' in args.split_name else None
+    s3_tokenizer = s3tokenizer.load_model(f"{args.token2wav_path}/speech_tokenizer_v2.onnx").to(device) if 'zero' in args.split_name else None
     dataset_name = "yuekai/CV3-Eval" if 'zero' in args.split_name else "yuekai/seed_tts_cosy2"
     dataset = load_dataset(
         dataset_name,

rl-tutorial/cosyvoice_llm/prepare_data.py

@@ -23,15 +23,20 @@
 
 from verl.utils.hdfs_io import copy, makedirs
 
+from typing import List
+import random
 
+def code_to_solution_str(code_list: List[int]) -> str:
+    """Convert code list to solution string format."""
+    return ''.join([f"<|s_{code}|>" for code in code_list])
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--train_file", required=True, help="Path to training JSON/JSONL file")
     parser.add_argument("--test_file", required=True, help="Path to test JSON/JSONL file")
     parser.add_argument("--local_dir", default=None, required=True)
     parser.add_argument("--hdfs_dir", default=None)
-    parser.add_argument("--use_custom_template", action="store_true", help="Use custom template for training")
+    parser.add_argument("--use_speech_prefix", action="store_true", help="Use speech prefix")
 
     args = parser.parse_args()
 
@@ -40,17 +45,21 @@
     test_dataset = datasets.load_dataset("json", data_files=args.test_file)['train']
 
     # add a row to each data item that represents a unique id
-    def make_map_fn(split, use_custom_template=False):
+    def make_map_fn(split):
         def process_fn(example, idx):
             text = example.pop("text")
-            if use_custom_template:
-                question = f"Convert the text to speech: {text}"
-                answer = "<|SPEECH_GENERATION_START|>"
-                print(f"use custom template for {split} {idx}")
-            else:
-                # use cosyvoice2 official huggingface compatible checkpoint template
-                question = text
-                answer = ""
+
+            # use cosyvoice2 official huggingface compatible checkpoint template
+            question = text
+            answer = ""
+            # generate a random float between 0 and 1, then convert it to 0 to 5
+            random_number = random.random() * 4 + 1
+            speech_token_len = int(random_number * 25)
+            codes = example.pop("code")
+            prefix_speech_token = codes[:speech_token_len]
+            prefix_speech_str = code_to_solution_str(prefix_speech_token)
+
+            answer = prefix_speech_str if args.use_speech_prefix else ""
 
             data = {
                 "data_source": f"{args.train_file}_{args.test_file}",  # Use file names as data source
@@ -72,12 +81,15 @@ def process_fn(example, idx):
                     "text": text,
                 },
             }
+            if args.use_speech_prefix:
+                data["extra_info"]["prefix_speech_str"] = prefix_speech_str
+
             return data
 
         return process_fn
 
-    train_dataset = train_dataset.map(function=make_map_fn("train", use_custom_template=args.use_custom_template), with_indices=True)
-    test_dataset = test_dataset.map(function=make_map_fn("test", use_custom_template=args.use_custom_template), with_indices=True)
+    train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True)
+    test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True)
 
     local_dir = args.local_dir
     hdfs_dir = args.hdfs_dir

rl-tutorial/cosyvoice_llm/requirements-cosyvoice.txt

@@ -21,3 +21,11 @@ soundfile==0.12.1
 tensorboard==2.14.0
 wget==3.2
 WeTextProcessing==1.0.3
+s3tokenizer
+tensorrt
+sherpa_onnx
+jiwer
+zhon
+numpy==1.25.2
+pypinyin
+openai-whisper

rl-tutorial/cosyvoice_llm/reward_tts.py

@@ -94,8 +94,11 @@ def compute_score(
     """
 
     # Decode token IDs
-    ids = _parse_ids(solution_str)
-
+    if "prefix_speech_str" in extra_info:
+        prefix_speech_str = extra_info["prefix_speech_str"]
+        ids = _parse_ids(prefix_speech_str + solution_str)
+    else:
+        ids = _parse_ids(solution_str)
     # Query remote server for reward
     try:
         reward = _remote_reward(ids, ground_truth)
@@ -107,7 +110,7 @@ def compute_score(
         print(
             f"\033[92m[{data_source}] Remote reward: {reward:.4f}\033[0m"
         )
-
+    reward = {"score": reward}
     return reward
 
 # CLI quick test

rl-tutorial/cosyvoice_llm/run.sh

@@ -19,7 +19,9 @@ if [ $stage -le -1 ] && [ $stop_stage -ge -1 ]; then
   # install verl
   git clone https://github.com/volcengine/verl.git
   cd verl
-  USE_MEGATRON=0 bash scripts/install_vllm_sglang_mcore.sh
+  USE_MEGATRON=0 USE_SGLANG=0 bash scripts/install_vllm_sglang_mcore.sh
+  # manually install flash attn above 2.7.4post1
+  pip install -r requirements.txt
   pip install --no-deps -e .
 
   # install CosyVoice
@@ -50,21 +52,25 @@ if [ $stage -le -1 ] && [ $stop_stage -ge -1 ]; then
   # If you would like to use the official CosyVoice2-0.5B LLM and do RL training, please see run_official.sh
 fi
 
-
+data_dir=data/parquet_emilia_zh_en_removed
 if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
   log "stage 0: prepare data into verl format"
   # yuekai/llasa_cosyvoice2_token_qwen_0.5b is the emilia zh trained model, please set use_custom_template=True to use the custom template
   # yuekai/cosyvoice2_llm is the official cosyvoice2 llm model, please set use_custom_template=False to use the official template
+  mkdir -p $data_dir
+  wget https://huggingface.co/datasets/SparkAudio/voxbox/resolve/main/metadata/emilia_zh.jsonl -O data/emilia_zh.jsonl
+  tail -n 100 data/emilia_zh.jsonl > data/emilia_test.jsonl
+  python3 filter_data.py
   python prepare_data.py \
-    --train_file data/emilia_zh-cosy-tiny-train.jsonl \
-    --test_file data/emilia_zh-cosy-tiny-test.jsonl \
-    --local_dir data/parquet_tiny \
-    --use_custom_template
+    --train_file data/emilia_zh-cosy-zh-filtered.jsonl \
+    --test_file data/emilia_test.jsonl \
+    --local_dir $data_dir
 fi
 
+n_gpus=8
 if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
   log "stage 1: start token2wav asr server for reward function"
-  python3 token2wav_asr_server.py --number-of-devices 8
+  python3 token2wav_asr_server.py --number-of-devices $n_gpus
 
   # log "Test the reward server"
   # python3 reward_tts.py \
@@ -76,18 +82,19 @@ if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
 fi 
 
 sft_model_path=/workspace/rl/llasa_cosyvoice2_token_qwen_0.5b/checkpoint-885000
-
+exp_name=emilia_zh
 if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
   log "stage 2: grpo train"
+  # wandb login
   export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
   export MKL_SERVICE_FORCE_INTEL=TRUE
-  n_gpus_per_node=8
+  n_gpus_per_node=$n_gpus
   micro_batch_size=4
   train_batch_size=32
   python3 -m verl.trainer.main_ppo \
       algorithm.adv_estimator=grpo \
-      data.train_files=data/parquet_aishell3_custom/train.parquet \
-      data.val_files=data/parquet_aishell3_custom/test.parquet \
+      data.train_files=$data_dir/train.parquet \
+      data.val_files=$data_dir/test.parquet \
       data.train_batch_size=$train_batch_size \
       data.max_prompt_length=1024 \
       data.max_response_length=1024 \
@@ -106,54 +113,61 @@ if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
       actor_rollout_ref.rollout.name=vllm \
       actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
       actor_rollout_ref.rollout.do_sample=true \
-      actor_rollout_ref.rollout.temperature=0.8 \
-      actor_rollout_ref.rollout.top_p=0.9 \
-      actor_rollout_ref.rollout.n=4 \
+      actor_rollout_ref.rollout.temperature=1.0 \
+      actor_rollout_ref.rollout.top_p=1.0 \
+      actor_rollout_ref.rollout.top_k=-1 \
+      actor_rollout_ref.rollout.n=8 \
       actor_rollout_ref.rollout.val_kwargs.do_sample=true \
-      actor_rollout_ref.rollout.val_kwargs.temperature=0.8 \
-      actor_rollout_ref.rollout.val_kwargs.top_p=0.9 \
+      actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \
+      actor_rollout_ref.rollout.val_kwargs.top_p=1.0 \
+      actor_rollout_ref.rollout.val_kwargs.top_k=-1 \
       reward_model.reward_manager=prime \
       custom_reward_function.path=reward_tts.py \
       custom_reward_function.name=compute_score \
       trainer.project_name='llasa_tts_grpo' \
-      trainer.experiment_name='aishell3_reward_tts_prime_test' \
+      trainer.experiment_name=${exp_name} \
       trainer.logger=['console','wandb'] \
       trainer.n_gpus_per_node=$n_gpus_per_node \
       trainer.nnodes=1 \
-      trainer.save_freq=100 \
-      trainer.test_freq=400 \
+      trainer.save_freq=50 \
+      trainer.test_freq=50 \
       trainer.resume_mode='auto' \
       trainer.total_epochs=1 \
-      trainer.val_before_train=False
+      trainer.val_before_train=False \
+      algorithm.norm_adv_by_std_in_grpo=False
 fi
 
-step=2100
-llm_path=./checkpoints/llasa_tts_grpo/aishell3_reward_tts_prime/global_step_${step}
+steps=(1300 600 700 800)
+
+for step in ${steps[@]}; do
+llm_path=./checkpoints/llasa_tts_grpo/${exp_name}/global_step_${step}
 if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
   log "stage 3: merge the model"
   python -m verl.model_merger merge \
       --backend fsdp \
       --local_dir $llm_path/actor \
       --target_dir $llm_path/merged_hf_model || exit 1
 
-fi 
+fi
+
+token2wav_path=/workspace/CosyVoice2-0.5B
+model_path=$llm_path/merged_hf_model
 if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
   log "stage 4: Test the model"
-  dataset=zero_shot_zh
-  output_dir=./outputs_rl_aishell3_step${step}_${dataset}_jit_trt_fp16_reward_tts
-  token2wav_path=/workspace/CosyVoice2-0.5B
-  model_path=$llm_path/merged_hf_model
-
+  datasets=(zero_shot_zh test_zh)
+  for dataset in ${datasets[@]}; do
+  output_dir=./outputs_rl_${exp_name}_step${step}_${dataset}
   CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
-  torchrun --nproc_per_node=8 \
+  torchrun --nproc_per_node=$n_gpus \
       infer_dataset.py \
         --output-dir $output_dir \
         --llm-model-name-or-path $model_path \
         --token2wav-path $token2wav_path \
         --split-name ${dataset} || exit 1
-
   bash scripts/compute_wer.sh $output_dir ${dataset}
+  done
 fi
+done
 
 if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then
   log "stage 5: Infer with single case"
@@ -163,4 +177,4 @@ if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then
     --prompt-speech-path ./assets/prompt_audio.wav \
     --model-path $llm_path/merged_hf_model \
     --input-text "扁担长，板凳宽，扁担绑在板凳上。吃葡萄不吐葡萄皮，不吃葡萄倒吐葡萄皮。"
-fi
+fi