generalized chat sft prompt (NVIDIA#7655)

* fix dataset issues Signed-off-by: Yi Dong <yidong@nvidia.com> * working version Signed-off-by: Yi Dong <yidong@nvidia.com> * all passed Signed-off-by: Yi Dong <yidong@nvidia.com> * refactor tests Signed-off-by: Yi Dong <yidong@nvidia.com> * all pass Signed-off-by: Yi Dong <yidong@nvidia.com> * working version Signed-off-by: Yi Dong <yidong@nvidia.com> * use end name signal for labels Signed-off-by: Yi Dong <yidong@nvidia.com> * all fixed Signed-off-by: Yi Dong <yidong@nvidia.com> * update doc Signed-off-by: Yi Dong <yidong@nvidia.com> * style fix Signed-off-by: Yi Dong <yidong@nvidia.com> * remove unused imports Signed-off-by: Yi Dong <yidong@nvidia.com> * make sure nccl not timing out Signed-off-by: Yi Dong <yidong@nvidia.com> * style fix Signed-off-by: Yi Dong <yidong@nvidia.com> * generate example template Signed-off-by: Yi Dong <yidong@nvidia.com> * generic end of name token Signed-off-by: Yi Dong <yidong@nvidia.com> * style fix Signed-off-by: Yi Dong <yidong@nvidia.com> * add the chat prompt format into the config Signed-off-by: Yi Dong <yidong@nvidia.com> * make sure sft working Signed-off-by: Yi Dong <yidong@nvidia.com> * address reviewer comment Signed-off-by: Yi Dong <yidong@nvidia.com> * fix non Signed-off-by: Yi Dong <yidong@nvidia.com> * try openAI prompt Signed-off-by: Yi Dong <yidong@nvidia.com> * remove unused imports Signed-off-by: Yi Dong <yidong@nvidia.com> * remove human labels from the data Signed-off-by: Yi Dong <yidong@nvidia.com> * use hf dataset to clean Signed-off-by: Yi Dong <yidong@nvidia.com> * reviewer comments Signed-off-by: Yi Dong <yidong@nvidia.com> --------- Signed-off-by: Yi Dong <yidong@nvidia.com>
Kipok · Oct 17, 2023 · d507414 · d507414
1 parent 7c0f94e
commit d507414
Show file tree

Hide file tree

Showing 9 changed files with 488 additions and 237 deletions.
diff --git a/examples/nlp/language_modeling/megatron_gpt_eval.py b/examples/nlp/language_modeling/megatron_gpt_eval.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import datetime
 import os
 import threading
 from functools import partial
@@ -167,7 +168,11 @@ def remove_padded_prompts(response, nb_paddings):
 def main(cfg) -> None:
 
     # trainer required for restoring model parallel models
-    trainer = Trainer(strategy=NLPDDPStrategy(), **cfg.trainer, callbacks=[CustomProgressBar()])
+    trainer = Trainer(
+        strategy=NLPDDPStrategy(timeout=datetime.timedelta(seconds=18000)),
+        **cfg.trainer,
+        callbacks=[CustomProgressBar()],
+    )
 
     if cfg.gpt_model_file is not None:
         if (

diff --git a/examples/nlp/language_modeling/tuning/conf/megatron_gpt_sft.yaml b/examples/nlp/language_modeling/tuning/conf/megatron_gpt_sft.yaml
@@ -71,6 +71,12 @@ model:
 
   data:
     chat: False # whether use chatbot data or not
+    chat_prompt_tokens:  # special tokens for the chat prompts, a dictionary of {token_type: token}. note that some tokenizer may combine the characters at the junction between {end_of_turn}{turn_start}. e.g. '<im end><im start>', the '><' sometimes is merged to be a single token. This is not supported, try to avoid
+      system_turn_start: '<extra_id_0>'
+      turn_start: '<extra_id_1>'
+      label_start: '<extra_id_2>'
+      end_of_turn: "\x0A"  # \0x0A is '\n'
+      end_of_name: "\x0A"  # \0x0A is '\n'
     train_ds:
       # Example of how to specify paths to multiple datasets
       # file_names:

diff --git a/examples/nlp/language_modeling/tuning/megatron_gpt_sft.py b/examples/nlp/language_modeling/tuning/megatron_gpt_sft.py
@@ -15,11 +15,12 @@
 import os
 import tempfile
 
+import torch.multiprocessing as mp
 from omegaconf.omegaconf import OmegaConf, open_dict
 from pytorch_lightning import Trainer
 from pytorch_lightning.plugins.environments import TorchElasticEnvironment
-from pytorch_lightning.trainer.connectors.checkpoint_connector import _CheckpointConnector
 
+from nemo.collections.nlp.data.language_modeling.megatron.gpt_sft_chat_dataset import get_prompt_template_example
 from nemo.collections.nlp.models.language_modeling.megatron_gpt_sft_model import MegatronGPTSFTModel
 from nemo.collections.nlp.modules.common.megatron.megatron_init import fake_initialize_model_parallel
 from nemo.collections.nlp.parts.nlp_overrides import (
@@ -35,6 +36,8 @@
 from nemo.utils.exp_manager import exp_manager
 from nemo.utils.model_utils import inject_model_parallel_rank
 
+mp.set_start_method("spawn", force=True)
+
 
 def _modify_config(gpt_cfg, cfg, add_cfg_to_tree=False):
     """
@@ -67,6 +70,13 @@ def _modify_config(gpt_cfg, cfg, add_cfg_to_tree=False):
         gpt_cfg.ffn_dropout = cfg.model.ffn_dropout
         gpt_cfg.use_flash_attention = cfg.model.get('use_flash_attention', False)
 
+        if cfg.model.data.get('chat', False):
+            # chat model, overwrite the prompt template
+            prompt_template = get_prompt_template_example(cfg.model.data.chat_prompt_tokens)
+            gpt_cfg.data.train_ds.prompt_template = prompt_template
+            gpt_cfg.data.validation_ds.prompt_template = prompt_template
+            gpt_cfg.data.test_ds.prompt_template = prompt_template
+
         sft_cls = MegatronGPTSFTModel
         gpt_cfg.target = f"{sft_cls.__module__}.{sft_cls.__name__}"