diff --git a/README.rst b/README.rst index 3135bdbfabdd..ea35f9bc96ab 100644 --- a/README.rst +++ b/README.rst @@ -57,10 +57,10 @@ such as FSDP, Mixture-of-Experts, and RLHF with TensorRT-LLM to provide speedups Introduction ------------ -NVIDIA NeMo Framework is a generative AI framework built for researchers and pytorch developers +NVIDIA NeMo Framework is a generative AI framework built for researchers and pytorch developers working on large language models (LLMs), multimodal models (MM), automatic speech recognition (ASR), and text-to-speech synthesis (TTS). -The primary objective of NeMo is to provide a scalable framework for researchers and developers from industry and academia +The primary objective of NeMo is to provide a scalable framework for researchers and developers from industry and academia to more easily implement and design new generative AI models by being able to leverage existing code and pretrained models. For technical documentation, please see the `NeMo Framework User Guide `_. @@ -68,8 +68,8 @@ For technical documentation, please see the `NeMo Framework User Guide `_ and training is automatically scalable to 1000s of GPUs. -When applicable, NeMo models take advantage of the latest possible distributed training techniques, -including parallelism strategies such as +When applicable, NeMo models take advantage of the latest possible distributed training techniques, +including parallelism strategies such as * data parallelism * tensor parallelism @@ -84,7 +84,7 @@ and mixed precision training recipes with bfloat16 and FP8 training. NeMo's Transformer based LLM and Multimodal models leverage `NVIDIA Transformer Engine `_ for FP8 training on NVIDIA Hopper GPUs and leverages `NVIDIA Megatron Core `_ for scaling transformer model training. -NeMo LLMs can be aligned with state of the art methods such as SteerLM, DPO and Reinforcement Learning from Human Feedback (RLHF), +NeMo LLMs can be aligned with state of the art methods such as SteerLM, DPO and Reinforcement Learning from Human Feedback (RLHF), see `NVIDIA NeMo Aligner `_ for more details. NeMo LLM and Multimodal models can be deployed and optimized with `NVIDIA Inference Microservices (Early Access) `_. @@ -93,7 +93,7 @@ NeMo ASR and TTS models can be optimized for inference and deployed for producti For scaling NeMo LLM and Multimodal training on Slurm clusters or public clouds, please see the `NVIDIA Framework Launcher `_. The NeMo Framework launcher has extensive recipes, scripts, utilities, and documentation for training NeMo LLMs and Multimodal models and also has an `Autoconfigurator `_ -which can be used to find the optimal model parallel configuration for training on a specific cluster. +which can be used to find the optimal model parallel configuration for training on a specific cluster. To get started quickly with the NeMo Framework Launcher, please see the `NeMo Framework Playbooks `_ The NeMo Framework Launcher does not currently support ASR and TTS training but will soon. diff --git a/docs/source/nlp/nemo_megatron/peft/landing_page.rst b/docs/source/nlp/nemo_megatron/peft/landing_page.rst index c90dcdfff1c5..4461feffb5db 100644 --- a/docs/source/nlp/nemo_megatron/peft/landing_page.rst +++ b/docs/source/nlp/nemo_megatron/peft/landing_page.rst @@ -12,14 +12,14 @@ fraction of the computational and storage costs. NeMo supports four PEFT methods which can be used with various transformer-based models. -==================== ===== ===== ========= == -\ GPT 3 NvGPT LLaMa 1/2 T5 -==================== ===== ===== ========= == -Adapters (Canonical) ✅ ✅ ✅ ✅ -LoRA ✅ ✅ ✅ ✅ -IA3 ✅ ✅ ✅ ✅ -P-Tuning ✅ ✅ ✅ ✅ -==================== ===== ===== ========= == +==================== ===== ======== ========= ====== == +\ GPT 3 Nemotron LLaMa 1/2 Falcon T5 +==================== ===== ======== ========= ====== == +LoRA ✅ ✅ ✅ ✅ ✅ +P-Tuning ✅ ✅ ✅ ✅ ✅ +Adapters (Canonical) ✅ ✅ ✅ ✅ +IA3 ✅ ✅ ✅ ✅ +==================== ===== ======== ========= ====== == Learn more about PEFT in NeMo with the :ref:`peftquickstart` which provides an overview on how PEFT works in NeMo. Read about the supported PEFT methods diff --git a/docs/source/nlp/nemo_megatron/peft/quick_start.rst b/docs/source/nlp/nemo_megatron/peft/quick_start.rst index 000e242b9508..fd46444eee54 100644 --- a/docs/source/nlp/nemo_megatron/peft/quick_start.rst +++ b/docs/source/nlp/nemo_megatron/peft/quick_start.rst @@ -62,7 +62,7 @@ Base model classes PEFT in NeMo is built with a mix-in class that does not belong to any model in particular. This means that the same interface is available to different NeMo models. Currently, NeMo supports PEFT for GPT-style -models such as GPT 3, NvGPT, LLaMa 1/2 (``MegatronGPTSFTModel``), as +models such as GPT 3, Nemotron, LLaMa 1/2 (``MegatronGPTSFTModel``), as well as T5 (``MegatronT5SFTModel``). Full finetuning vs PEFT @@ -78,11 +78,13 @@ PEFT. trainer = MegatronTrainerBuilder(config).create_trainer() model_cfg = MegatronGPTSFTModel.merge_cfg_with(config.model.restore_from_path, config) + ### Training API ### model = MegatronGPTSFTModel.restore_from(restore_path, model_cfg, trainer) # restore from pretrained ckpt - + peft_cfg = LoRAPEFTConfig(model_cfg) + + peft_cfg = LoraPEFTConfig(model_cfg) + model.add_adapter(peft_cfg) trainer.fit(model) # saves adapter weights only + ### Inference API ### # Restore from base then load adapter API model = MegatronGPTSFTModel.restore_from(restore_path, trainer, model_cfg) + model.load_adapters(adapter_save_path, peft_cfg) diff --git a/docs/source/nlp/nemo_megatron/prompt_learning.rst b/docs/source/nlp/nemo_megatron/prompt_learning.rst deleted file mode 100644 index 8fe481019a6f..000000000000 --- a/docs/source/nlp/nemo_megatron/prompt_learning.rst +++ /dev/null @@ -1,390 +0,0 @@ -.. _promptlearning: - -Prompt Learning ---------------- - -Within NeMo we refer to **p-tuning** and **prompt tuning** methods collectively as prompt learning. Both methods are parameter efficient alternatives to fine-tuning pretrained language models. Our NeMo implementation makes it possible to use one pretrained GPT model on many downstream tasks without needing to tune the model's full set of parameters. It also allows for adding new tasks to your model without overwriting or disrupting previous tasks for which the model has already been p-tuned/prompt-tuned. Because the original model parameters are frozen and never altered by either method, p-tuning/prompt-tuning also avoids catastrophic forgetting issues often encountered when fine-tuning models. - -Instead of selecting discrete text prompts in a manual or automated fashion, prompt tuning and p-tuning utilize virtual prompt embeddings that can be optimized via gradient descent. The only difference between prompt tuning and p-tuning within NeMo-Megatron is the architecture used to tune the soft prompt tokens during training. - -- Our prompt tuning implementation is based off Lester et. al’s EMNLP 2021 paper "`The Power of Scale for Parameter-Efficient Prompt Tuning `_" -- Our p-tuning implementation is based off Liu et al's paper "`GPT Understands, Too `_" - -Our continuous learning capability for combined p-tuning and prompt tuning with GPT style models is a NeMo specific extension of the author's original work. - -Please also checkout our `prompt learning tutorial notebook. `_ - - -Terminology -^^^^^^^^^^^ -We will be using the terms ``continuous``, ``soft``, and ``virtual`` token interchangeably to refer to embeddings inserted into the model prompt that have no concrete mapping to strings or characters within the model’s vocabulary. These virtual token embeddings exist in contrast to the ``discrete``, ``hard``, or ``real`` tokens that do make up the model’s vocabulary. Virtual tokens are purely 1D vectors with dimensionality equal to that of each real token embedding, matching the ``hidden_size`` hyperparameter. In training and inference, continuous token embeddings are inserted among discrete token embeddings according to a template you provide in the model's config. We will demonstrate how to do this below. - -When referring to p-tuning and prompt tuning together, we will be using the phrase prompt learning for simplicity. - -Prompt Tuning -^^^^^^^^^^^^^ - -In prompt-tuning a pretrained GPT model, soft prompt embeddings are initialized as a 2D matrix of size ``total_virtual_tokens X hidden_size``. Each task the model is prompt-tuned to perform has its own 2D embedding matrix associated with it. Tasks do not share any parameters during training or inference. All GPT model parameters are frozen and only the embedding parameters for each task are updated during training. - -In prompt tuning you can specify how the embeddings are initialized for each task. You can either - -- Initialize embedding parameters according to some random distribution -- Initialize embedding parameters from existing vocabulary embeddings (recommended) - -If you choose to initialize virtual token embeddings from existing embedding weights, you can provide the string of words you want to use for initialization in the model's config. This string will be tokenized and tiled or truncated to match the specified number of virtual tokens you would like to use (``total_virtual_tokens``). Vocab embeddings are copied and used to initialize the soft prompt embedding matrix for each task. The vocab embeddings themselves are not updated or changed during prompt tuning. - -P-Tuning -^^^^^^^^ - -In p-tuning, an LSTM model is used to predict virtual token embeddings. We refer to this LSTM model as our ``prompt_encoder``. LSTM parameters are randomly initialized at the start of p-tuning. All GPT model parameters are frozen, and only the LSTM weights are updated at each training step. LSTM parameters are shared between all tasks that are p-tuned at the same time, but the LSTM model outputs unique virtual token embeddings for each task. The virtual tokens predicted by the LSTM are inserted among the discrete token input in the exact same manner as with prompt-tuning. You still specify the number of virtual tokens you want to use by setting ``total_virtual_tokens`` and each virtual token embedding is still a 1D vector of size ``hidden_size``. - -Using Both Prompt and P-Tuning -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -A single pretrained GPT model can use both p-tuning and prompt-tuning. While you must decide to use either p-tuning or prompt-tuning for each task you want your model to perform, you can p-tune your model on a set of tasks *A*, then prompt tune your same model on a different set of tasks *B*, then finally run inference on tasks from both *A* and *B* at the same time. During prompt-tuning or p-tuning, tasks tuned at the same time must use the same number of virtual tokens. During inference, tasks using differing amounts of virtual tokens can be run at the same time. - -When p-tuning completes, prompt tuned virtual tokens from the p-tuning ``prompt_encoder`` are automatically moved to the ``prompt_table`` where all prompt tuned and p-tuned soft prompts are stored. The LSTM ``prompt_encoder`` is then removed from the model. This allows us to preserve previously p-tuned soft prompts while still maintaining the ability to add new p-tuned or prompt-tuned soft prompts in the future. The ``prompt_table`` uses the ``taskname`` as a key to look up the correct virtual tokens for a specified task. The ``prompt_table``'s hash table data structure also makes it possible for each task to flexibly use a different number of virtual tokens. - -P-tuning usually requires fewer virtual tokens per task to achieve good results but uses a higher number of parameters compared to prompt-tuning. For example, if you prompt tune a 125M parameter GPT model (with hidden size 768) on two tasks, using 100 virtual tokens per task, the total parameters tuned during prompt tuning would equal 153k (~.1% of the pre-trained model size). If you p-tune the same 125M GPT model on 2 tasks, using an LSTM with two layers and 10 tokens per task, you will be tuning 8.3M parameters (~6.6% of the pre-trained model size). The increased number of parameters used during p-tuning is mitigated by our ``prompt_table``. When p-tuned soft prompts are placed in the prompt table, only the parameters for the predicted virtual tokens are saved. This allows us to keep the benefit of tuning a larger number of parameters during training, while also preserving the parameter efficiency of prompt-tuning during inference and storing of the model. - -Because p-tuning shares parameters between tasks during training, p-tuning your model on multiple tasks that are similar might allow your model to share insight between tasks. In the same vein, p-tuning on many very different tasks at once might perform worse than prompt tuning, which tunes a distinct set of parameters per task. **Generally we recommend using p-tuning over prompt tuning.** - -Users can also optionally tune the model's full parameters in addition to the soft prompt parameters. See ``model.lm_finetune`` in the Prompt Learning Config section for details on how to configure this. - -Dataset Preprocessing -^^^^^^^^^^^^^^^^^^^^^ - -The prompt learning dataset accepts a list of json/dictionary objects or a list of json file names where each json file contains a collection of json objects. Each json object must include the field ``taskname`` which is a string identifier for the task the data example corresponds to. They should also include one or more fields corresponding to different sections of the discrete text prompt. The input data might look like: - -.. code:: - - [ - {"taskname": "squad", "context": [CONTEXT_PARAGRAPH_TEXT1], "question": [QUESTION_TEXT1], "answer": [ANSWER_TEXT1]}, - {"taskname": "squad", "context": [CONTEXT_PARAGRAPH_TEXT2], "question": [QUESTION_TEXT2], "answer": [ANSWER_TEXT2]}, - {"taskname": "intent_and_slot", "utterance": [UTTERANCE_TEXT1], "label": [INTENT_TEXT1][SLOT_TEXT1]}, - {"taskname": "intent_and_slot", "utterance": [UTTERANCE_TEXT2], "label": [INTENT_TEXT2][SLOT_TEXT2]}, - {"taskname": "sentiment", "sentence": [SENTENCE_TEXT1], "label": [SENTIMENT_LABEL1]}, - {"taskname": "sentiment", "sentence": [SENTENCE_TEXT2], "label": [SENTIMENT_LABEL2]}, - ] - -These additional fields can be unlimited in number and will be used to help map different parts of the discrete text input to a prompt template that you define. We show how this mapping works and how to construct your prompt template in the Prompt Formatting section. Data examples for each dataset can all be passed to the dataset class in one file, or in separate ``.jsonl`` files in a list. - -.. _data-example-label: - -Prompt Formatting -^^^^^^^^^^^^^^^^^ - -To customize different prompts for different tasks, we simply need to specify the prompt task template in the config file at ``model.task_templates``. The virtual token markers ``<|VIRTUAL_PROMPT_#|>`` signify where you want virtual tokens to be placed in the template string. ``<|VIRTUAL_PROMPT_0|>``, ``<|VIRTUAL_PROMPT_1|>``, and ``<|VIRTUAL_PROMPT_2|>`` indicate where a number of virtual tokens matching the values given at ``virtual_token_splits[0]``, ``virtual_token_splits[1]`` and ``virtual_token_splits[2]`` will be placed. The other variable fields ``{var}`` refer to the fields in the data json. - -For example, given: - -- the data json ``{"sentence1": "And he said, Mama, I'm home.", "sentence2": "He didn't say a word."}`` -- virtual token splits set to ``virtual_token_splits = [3, 3, 3]`` -- a prompt template set to ``prompt_template = "<|VIRTUAL_PROMPT_0|> Hypothesis: [sentence1], <|VIRTUAL_PROMPT_1|> Premise: [sentence2] <|VIRTUAL_PROMPT_2|> Answer:"`` - -the input will be translated into ``VVV Hypothesis: And he said, Mama, I'm home. VVV Premise: He didn't say a word. VVV Answer:``, where ``VVV`` are three virtual tokens. - -**We recommend you first try prompt learning by placing all virtual tokens at the very beginning of your prompt template** like we do with the ``sentiment`` task example below. We've found this gives strong performance. -.. code:: - - config.model.task_templates = [ - { - "taskname": "sentiment", - "prompt_template": "<|VIRTUAL_PROMPT_0|> {sentence} sentiment: {label}", - "total_virtual_tokens": 10, - "virtual_token_splits": [10], - "truncate_field": "sentence", - "answer_only_loss": False, - }, - { - "taskname": "intent_and_slot", - "prompt_template": "<|VIRTUAL_PROMPT_0|> Predict intent and slot <|VIRTUAL_PROMPT_1|> :\n{utterance}{label}", - "total_virtual_tokens": 10, - "virtual_token_splits": [7, 3], - "truncate_field": None, - "answer_only_loss": True, - "answer_field": "label" - } - ] - -.. _prompt-formatting-label: - -``model.task_templates`` Config Parameters -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -.. list-table:: - :widths: 15 15 25 - :header-rows: 1 - - * - **Parameter** - - **Data type** - - **Description** - * - **taskname** - - string - - Short string denoting the task, used to lookup task specific virtual tokens from the ``prompt_table``. Refers to the same ``taskname`` in the dataset json objects. - * - **prompt_template** - - string - - a string showing the model where to place virtual tokens and how to map dataset json fields to where they belong in the model prompt - * - **total_virtual_tokens** - - int - - specifies the total number of virtual tokens that will be inserted into the model prompt - * - **virtual_token_splits** - - list of ints - - specifies the number of virtual tokens that belong at each ``<|VIRTUAL_PROMPT_#|>`` marker. ``virtual_token_splits`` values should add up to ``total_virtual_tokens``. The number of ``virtual_token_splits`` should match the number of ``<|VIRTUAL_PROMPT_#|>`` markers. - * - **answer_only_loss** - - bool - - Whether to limit loss calculation to only the answer portion of the prompt during tuning. Strongly recommended for long prompts. - * - **answer_field** - - string - - The field in the data json corresponding to the answer. The loss will only be calculated on this portion of the prompt if ``answer_only_loss`` is ``True``. The answer field must be at the end of the prompt template. - * - **truncate_field** - - string - - specifies which field in the data json to truncate if the length of the input exceeds the maximum sequence length of the model. If ``truncate_field`` is set to ``None``, examples that are too long are simply dropped from the dataset. - -Prompt Learning Specific Config Values -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -.. list-table:: - :widths: 15 15 25 - :header-rows: 1 - - * - **Parameter** - - **Data type** - - **Description** - * - **model.nemo_path** - - string - - Path to where you want to save your model after prompt tuning/p-tuning, must end in `.nemo` - * - **model.virtual_prompt_style** - - string - - one of 'prompt-tuning', 'p-tuning', or 'inference' - * - **model.language_model_path** - - string - - Path to the GPT language model .nemo file you want to use for prompt learning, not needed if ``restore_path`` is set - * - **model.restore_path** - - string - - Path to a .nemo file of existing ``MegatronGPTPromptLearningModel`` that has already been prompt tuned or p-tuned on at least one task. P-tuned or prompt tuned in this training session will be added to this model's `prompt_table`. Should be set to ``null`` if none. - * - **model.new_tasks** - - list of strings - - List of new tasknames to be prompt or p-tuned, - * - **model.existing_tasks** - - list of strings - - List of tasks the model has already been p-tuned/prompt-tuned for, needed when a restore path is given. Should be set to ``[]`` if None. - * - **model.task_templates** - - list - - See the ``model.task_templates`` Config Parameters Table above - * - **model.prompt_tuning.new_prompt_init_methods** - - list of strings - - List of 'text' or 'random', should correspond to the order of tasks listed in ``model.new_tasks``. Only needed if `virtual_prompt_style='prompt-tuning'` - * - **model.prompt_tuning.new_prompt_init_text** - - list of strings - - The text you want to use for soft prompt initalization if ``model.prompt_tuning.new_prompt_init_methods`` is set to 'text' for a task. Should correspond to the order of tasks listed in ``model.new_tasks``. The text is tokenized and clipped or tiled to match ``total_virtual_tokens`` in ``model.task_templates``. The vocab embeddings associated with each token are copied and use to initialize the soft prompts before tuning. - * - **model.p_tuning.dropout** - - float - - LSTM prompt encoder dropout prob - * - **model.p_tuning.num_layers** - - int - - Num layers in LSTM prompt encoder - * - **model.tensor_model_parallel_size** - - int - - intra-layer model parallelism, must match the ``tensor_model_parallel_size`` of the GPT model given at ``language_model_path`` - * - **model.batch_size** - - int - - global batch size - * - **model.data.train_ds** - - list of strings - - list of ``.json`` or ``.jsonl`` training dataset files with json ojects that have the dataset format described above - * - **model.data.validation_ds** - - list of strings - - list of ``.json`` or ``.jsonl`` validation dataset files with json ojects that have the dataset format described above - * - **model.data.add_eos** - - bool - - Whether to add an EOS token at the end of each training example (recommended). - -An example config file can be found at https://github.com/NVIDIA/NeMo/tree/stable/examples/nlp/language_modeling/conf/megatron_gpt_prompt_learning_config.yaml - -Setting New Tasks -^^^^^^^^^^^^^^^^^ - -After you p-tune or prompt-tune your model, you can always go back and p-tune or prompt-tune your model on more tasks without over writing the virtual prompts who've trained already. You can also use a different number of ``total_virtual_tokens`` between each training session as long as tasks ptuned or prompt tuned at the same time have the same number of ``total_virtual_tokens``. For this reason, when you ptune on a new task, you need to tell your model which of your tasks are new and which ones already exist (and thus you don't want to tune them). You do this by setting the ``new_tasks`` and ``existing_tasks`` values in the config file. - -Example Multi-Task Prompt Tuning Config and Command -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -First define a config called ``multitask-prompt-learning.yaml`` demonstrated below. **In the** ``exp_manager`` **portion of the config,** ``save_nemo_on_train_end`` **should be set to** ``False`` **to avoid unnecessarily saving the incorrect model weights.** This is already done in the example `megatron_gpt_prompt_learning_config.yaml config `_ that you should use as your starting point. The correct prompt learning model will be saved at the ``model.nemo_path`` you set. - -.. code:: - - name: multitask_prompt_tuning - trainer: ... - exp_manager: ... - model: - seed: 1234 - nemo_path: ${name}.nemo - virtual_prompt_style: "prompt-tuning" - encoder_seq_length: 2048 - tensor_model_parallel_size: 1 - pipeline_model_parallel_size: 1 - global_batch_size: 16 - micro_batch_size: 4 - - restore_path: null - language_model_path: models/megatron_125M_gpt.nemo - existing_tasks: [] - new_tasks: ["sentiment", "intent_and_slot"] - - task_templates: - - taskname: "sentiment" - prompt_template: "<|VIRTUAL_PROMPT_0|> {sentence} sentiment: {label}" - total_virtual_tokens: 100 - virtual_token_splits: [100] - truncate_field: null - answer_only_loss: False - - - taskname: "intent_and_slot" - prompt_template: "<|VIRTUAL_PROMPT_0|> Predict intent and slot <|VIRTUAL_PROMPT_1|> :\n{utterance}{label}" - total_virtual_tokens: 100 - virtual_token_splits: [80, 20] - truncate_field: null - answer_only_loss: True - answer_field: "label" - - prompt_tuning: - new_prompt_init_methods: ["text", "text"] - new_prompt_init_text: ["financial sentiment analysis postive neutral negative", "intent and slot classification virtual assistant task bot please"] - - data: - train_ds: ["data/financial_phrase_bank_train.jsonl", "data/assistent_train.jsonl"] - validation_ds: ["data/financial_phrase_bank_val.jsonl", "data/assistent_val.jsonl"] - add_eos: True - shuffle: True - num_workers: 1 - pin_memory: True - - optim: ... - -(See https://github.com/NVIDIA/NeMo/tree/stable/examples/nlp/language_modeling/conf/megatron_gpt_prompt_learning_config.yaml for what should go in the ``trainer``, ``exp_manager``, and ``optim`` sections.) - -Then run the command - -.. code:: - - python megatron_gpt_prompt_learning.py --config-name=multitask-prompt-learning.yaml - - -Example Multi-Task P-Tuning Config and Command After Prompt-Tuning -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -Update ``multitask-prompt-learning.yaml`` from the example above with p-tuning parameters for the new task. Be sure to update ``model.existing_tasks`` with the tasknames from previous prompt learning runs and to use the ``.nemo`` file saved at the end of your last prompt learning session. Values different from the config above have stars commented next to them. - -In this example, the SQuAD task includes the question context as part of the prompt. Because the context is long, we recommend setting ``answer_only_loss`` to ``True`` for this task, and any task where a significant portion of the prompt is not a part of the answer. ``answer_only_loss`` tells the model to only calculate the cross-entropy loss on the answer portion of the training example. Though we recommend placing all virtual tokens at the beginning of the prompt, we place them throughout the prompt in this example to demonstrate how to do so. - -.. code:: - - name: multitask_p_tuning # *** - trainer: ... - exp_manager: ... - model: - seed: 1234 - nemo_path: ${name}.nemo - virtual_prompt_style: "p-tuning" # *** - encoder_seq_length: 2048 - tensor_model_parallel_size: 1 - pipeline_model_parallel_size: 1 - global_batch_size: 16 - micro_batch_size: 4 - - restore_path: multitask_prompt_tuning.nemo # *** - language_model_path: models/megatron_125M_gpt.nemo - existing_tasks: ["sentiment", "intent_and_slot"] # *** - new_tasks: ["squad"] - - task_templates: - - taskname: "sentiment" - prompt_template: "<|VIRTUAL_PROMPT_0|> {sentence} sentiment: {label}" - total_virtual_tokens: 100 - virtual_token_splits: [100] - truncate_field: null - answer_only_loss: False - - - taskname: "intent_and_slot" - prompt_template: "<|VIRTUAL_PROMPT_0|> Predict intent and slot <|VIRTUAL_PROMPT_1|> :\n{utterance}{label}" - total_virtual_tokens: 100 - virtual_token_splits: [80, 20] - truncate_field: null - answer_only_loss: True - answer_field: "label" - - - taskname: "squad" # *** - prompt_template: "<|VIRTUAL_PROMPT_0|> Answer the question from the context {question} {context} Answer: {answer}" # *** - total_virtual_tokens: 9 # *** - virtual_token_splits: [9] # *** - truncate_field: context # *** - answer_only_loss: True # *** - answer_field: "answer" # *** - - p_tuning: # *** - dropout: 0.0 # *** - num_layers: 2 # *** - - data: - train_ds: ["data/squad_train.jsonl"] # *** - validation_ds: ["data/squad_val.jsonl"] # *** - add_eos: True - shuffle: True - num_workers: 1 - pin_memory: True - - optim: ... - -Then run the command again: - -.. code:: - - python megatron_gpt_prompt_learning.py --config-name=multitask-prompt-learning.yaml - - -Example Multi-Task Inference -^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -The inference file can contain a mix of prompts from all the tasks the model has been prompt tuned on. - -.. code:: - - python megatron_gpt_prompt_learning_eval.py \ - virtual_prompt_model_file=PATH_TO_NEMO_PROMPT_LEARNING_MODEL_FILE \ - gpt_model_file=PATH_TO_FROZEN_GPT_MODEL_FILE \ - inference.greedy=True \ - inference.add_BOS=False \ - trainer.devices=1 \ - trainer.num_nodes=1 \ - tensor_model_parallel_size=1 \ - pipeline_model_parallel_size=1 \ - prompts=[prompt1,prompt2] - -``virtual_prompt_model_file`` should be a path to a .nemo file saved after p-tuning/prompt tuning and ``model_file`` is still the path to the gpt model's .nemo file. - -prompts in this case should be a list of .json or .jsonl files containing json objects similar to the ones used during prompt learning. They should have keys that match the fields specified in the prompt template. Fields can be dropped from the prompt dict and their corresponding section of the prompt template will be automatically removed. - -For example, say the prompt template during p-tuning/prompt-tuning looked like: - -.. code:: - - '<|VIRTUAL_PROMPT_0|> Context: {context} Question: {question} Answer: {answer}' - -but you don't want to include the answer field during inference. Just don't include the answer field in the prompt dict like below: - -.. code:: - - {"taskname": "squad", "context": "some paragraph", "question": "question related to paragraph"} - {"taskname": "squad", "context": "another paragraph", "question": "a different question related to paragraph"} - - -And the dataset class will automatically format your input to have the form: - -.. code:: - - [ - '<|VIRTUAL_PROMPT_0|> Context: some paragraph Question: question related to paragraph Answer: ', - '<|VIRTUAL_PROMPT_0|> Context: another paragraph Question: a different question related to paragraph Answer: ' - ] - -Generally prompt learning inference is just like running inference with a GPT model. The only difference is you need to add ``virtual_prompt_model_file=PATH_TO_NEMO_PROMPT_LEARNING_MODEL_FILE`` to your command if you're using a p-tuned/prompt-tuned model. - -Example prompt learning script: `NeMo/examples/nlp/language_modeling/megatron_gpt_prompt_learning.py.py `__. - -Example prompt tuned inference script: `NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py `__. diff --git a/tutorials/nlp/Multitask_Prompt_and_PTuning.ipynb b/tutorials/nlp/Multitask_Prompt_and_PTuning.ipynb deleted file mode 100644 index 076a8ffad3df..000000000000 --- a/tutorials/nlp/Multitask_Prompt_and_PTuning.ipynb +++ /dev/null @@ -1,786 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "b7a434f4", - "metadata": {}, - "outputs": [], - "source": [ - "BRANCH='main'" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "developmental-gibraltar", - "metadata": {}, - "outputs": [], - "source": [ - "\"\"\"\n", - "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n", - "\n", - "Instructions for setting up Colab are as follows:\n", - "1. Open a new Python 3 notebook.\n", - "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n", - "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n", - "4. Run this cell to set up dependencies.\n", - "\"\"\"\n", - "# If you're using Google Colab and not running locally, run this cell\n", - "\n", - "# install NeMo\n", - "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[nlp]" - ] - }, - { - "cell_type": "markdown", - "id": "42daf8bf", - "metadata": {}, - "source": [ - "# Introduction\n", - "\n", - "In this notebook we demonstrate how to use p-tuning and prompt tuning within NeMo-Megatron. Both methods are parameter efficient alternatives to fine-tuning pretrained language models. Our NeMo implementation makes it possible to use one pretrained GPT model on many downstream tasks without needing to tune the model’s full set of parameters. It also allows for adding new tasks to your model without overwriting or disrupting previous tasks for which the model has already been p-tuned/prompt-tuned. Because the original model parameters are frozen and never altered by either method, p-tuning/prompt-tuning also avoid catastrophic forgetting issues often encountered when fine-tuning models.\n", - "\n", - "- Our prompt tuning implementation is based off Lester et. al’s EMNLP 2021 paper [The Power of Scale for Parameter-Efficient Prompt Tuning](https://arxiv.org/abs/2104.08691)\n", - "\n", - "- Our p-tuning implementation is based off Liu et al's paper [GPT Understands, Too](https://arxiv.org/abs/2103.10385).\n", - "\n", - "- Command line usage examples and API documentation can be found in [our user docs](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/nemo_megatron/prompt_learning.html). \n", - "\n", - "\"Prompt\n", - "\n", - "Our continuous learning capability for combined p-tuning and prompt tuning with GPT style models is a NeMo specific extension of the author’s original work.\n", - "\n", - "# The Plan\n", - "\n", - "We are going to show you how to:\n", - " \n", - " 1. P-Tune/Prompt Tune a model on multiple tasks at the same time\n", - " 2. Add a new task to a model that has already been P-Tuned/Prompt Tuned previously\n", - " \n", - "We will first p-tune a GPT model on sentiment analysis, and intent and slot classification tasks. Then we will show how to add the squad question answering task to the same model we already p-tuned once.\n", - "\n", - "\n", - "# Technical Overview\n", - "Instead of selecting discrete text prompts in a manual or automated fashion, prompt tuning and p-tuning utilize virtual prompt embeddings that can be optimized via gradient decent. The only difference between prompt tuning and p-tuning within NeMo-Megatron is the architecture used to tune the soft prompt tokens during training.\n", - "\n", - "### Terminology\n", - "We will be using the terms `continuous`, `soft`, and `virtual` token interchangeably to refer to embeddings inserted into the model prompt that have no concrete mapping to strings or characters within the model’s vocabulary. These virtual token embeddings exist in contrast to the `discrete`, `hard`, or `real` tokens that do make up the model’s vocabulary. Virtual tokens are purely 1D vectors with dimensionality equal to that of each real token embedding, matching the `hidden_size` hyperparameter. In training and inference, continuous token embeddings are inserted among discrete token embeddings according to a template you provide in the model’s config. We will demonstrate how to do this below.\n", - "\n", - "When referring to p-tuning and prompt tuning together, we will be using the phrase prompt learning for simplicity.\n", - "\n", - "### Prompt-Tuning\n", - "In prompt-tuning a pretrained GPT model, soft prompt embeddings are initialized as a 2D matrix of size `total_virtual_tokens X hidden_size`. Each task the model is prompt-tuned to perform has its own 2D embedding matrix associated with it. Tasks do not share any parameters during training or inference. All GPT model parameters are frozen and only the embedding parameters for each task are updated during training.\n", - "\n", - "In prompt tuning you can specify how the embeddings are initialized for each task. You can either\n", - "\n", - "1. Initialize embedding parameters according to some random distribution\n", - "2. Initialize embedding parameters from existing vocabulary embeddings (recommended)\n", - "\n", - "If you choose to initialize virtual token embeddings from existing embedding weights, you can provide the string of words you want to use for initialization in the model’s config. This string will be tokenized and tiled or truncated to match the specified number of virtual tokens you would like to use (`total_virtual_tokens`). Vocab embeddings are copied and used to initialize the soft prompt embedding matrix for each task. The vocab embeddings themselves are not updated or changed during prompt tuning.\n", - "\n", - "\n", - "### P-Tuning\n", - "In p-tuning, an LSTM model is used to predict virtual token embeddings. We refer to this LSTM model as our `prompt_encoder`. LSTM parameters are randomly initialized at the start of p-tuning. All GPT model parameters are frozen, and only the LSTM weights are updated at each training step. LSTM parameters are shared between all tasks that are p-tuned at the same time, but the LSTM model outputs unique virtual token embeddings for each task. The virtual tokens predicted by the LSTM are inserted among the discrete token input in the exact same manner as with prompt-tuning. You still specify the number of virtual tokens you want to use by setting `total_virtual_tokens` and each virtual token embedding is still a 1D vector of size `hidden_size`.\n", - "\n", - "\n", - "\n", - "# The Best of Both\n", - "A single pretrained GPT model can use both p-tuning and prompt-tuning. While you must decide to use either p-tuning or prompt-tuning for each task you want your model to perform, you can p-tune your model on a set of tasks A, then prompt tune your same model on a different set of tasks B, then finally run inference on tasks from both A and B at the same time. During prompt-tuning or p-tuning, tasks tuned at the same time must use the same number of virtual tokens. During inference, tasks using differing amounts of virtual tokens can be run at the same time.\n", - "\n", - "Please see our [docs for more comparisons between prompt and p-tuning](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/nemo_megatron/prompt_learning.html). \n", - "\n", - "With all that covered, let's get started!\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "31c27562", - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "import wget" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "0bfc7709", - "metadata": {}, - "source": [ - "# Tasks and Datasets\n", - "We will be using p-tuning to teach our GPT model to do **Question Answering**.\n", - "\n", - "We will be using the [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/) reading comprehension dataset, consisting of questions posed by crowd workers on a set of Wikipedia articles, where the answer to every question is a segment of text. More information on [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/) can be found on their website or in their paper by Rajpurkar et. al \"[Know What You Don’t Know: Unanswerable Questions for SQuAD](https://arxiv.org/pdf/1806.03822.pdf)\"." - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "e0b0072a", - "metadata": {}, - "source": [ - "# Data Preparation\n", - "\n", - "The prompt learning dataset loader accepts a list of json/dictionary objects or a list of json file names where each json file contains a collection of json objects. Each json object must include the field `taskname` which is a string identifier for the task the data example corresponds to. They should also include one or more fields corresponding to different sections of the discrete text prompt. The input data might look like:\n", - "\n", - "```\n", - "[\n", - " {\"taskname\": \"squad\", \"context\": [CONTEXT_PARAGRAPH_TEXT1], \"question\": [QUESTION_TEXT1], \"answer\": [ANSWER_TEXT1]},\n", - " {\"taskname\": \"squad\", \"context\": [CONTEXT_PARAGRAPH_TEXT2], \"question\": [QUESTION_TEXT2], \"answer\": [ANSWER_TEXT2]},\n", - "]\n", - "```\n", - "\n", - "These additional fields can be unlimited in number and will be used to help map different parts of the discrete text input to a prompt template that you define. We will show how this mapping works and how to construct your prompt template in the `Prompt Formatting` section. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "0dbd41fd", - "metadata": {}, - "outputs": [], - "source": [ - "# You can replace DATA_DIR and NEMO_DIR with your own locations\n", - "DATA_DIR = \"data\"\n", - "NEMO_DIR = '.'\n", - "\n", - "os.makedirs(DATA_DIR, exist_ok=True)" - ] - }, - { - "cell_type": "markdown", - "id": "504a7b40", - "metadata": {}, - "source": [ - "\n", - "For each dataset we have preprocessing scripts pre-written in NeMo's example directory located in `examples/nlp`. Let's download those now. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "e72a1dc1", - "metadata": {}, - "outputs": [], - "source": [ - "# download the preprocessing scripts from github for the purpose of this tutorial\n", - "wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/scripts/dataset_processing/nlp/squad/prompt_learning_squad_preprocessing.py', NEMO_DIR)" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "71813919", - "metadata": {}, - "source": [ - "Now let's down load and process the dataset." - ] - }, - { - "cell_type": "markdown", - "id": "816791de", - "metadata": {}, - "source": [ - "### SQuAD Dataset" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "fa16d8ac", - "metadata": {}, - "outputs": [], - "source": [ - "SQUAD_DIR = os.path.join(DATA_DIR, \"SQuAD\")\n", - "os.makedirs(SQUAD_DIR, exist_ok=True)\n", - "\n", - "# Download the SQuAD dataset\n", - "!wget https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json\n", - "!wget https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json\n", - "!mv train-v1.1.json {SQUAD_DIR}\n", - "!mv dev-v1.1.json {SQUAD_DIR}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "64e3e25b", - "metadata": {}, - "outputs": [], - "source": [ - "# Preprocess squad data\n", - "!python $NEMO_DIR/prompt_learning_squad_preprocessing.py --data-dir {SQUAD_DIR}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "b562d1de", - "metadata": {}, - "outputs": [], - "source": [ - "# What the squad dataset looks like after processing\n", - "!head -4 $SQUAD_DIR/squad_train.jsonl" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "a385d319", - "metadata": {}, - "source": [ - "We made a `.jsonl` file for each of the train, validation, and testing splits of the squad data. Every `.jsonl` file contains json objects with the fields `taskname`, `context`, `question`, and `answer`. The preprocessing script is called `prompt_learning_squad_preprocessing.py`. It should be in your `NEMO_DIR` and at `scripts/dataset_processing/nlp/squad/prompt_learning_squad_preprocessing.py` in the NeMo repo. \n", - "\n", - "The SQuAD dataset consists of various topics like `Beyoncé`, `IPod`, and `Symbiosis`. Each topic has several paragraphs associated with it, and each paragraph has several questions and answers related to it. When we separated the train/validation/test splits, we separated them on the topic level. For example, if the training set contains paragraphs and questions about the topic `Beyoncé`, neither the validation nor test sets will contain any questions on this topic. All questions about a certain topic are isolated to one split of the data. \n", - "\n", - "Like the Financial PhraseBank Dataset, we randomly selected 80% of the questions for training, 10% for validation, and 10% for test. This resulted in `69125` test examples, `8952` validation examples, and `8744` testing examples. The `answer` field was removed from test examples.\n", - "\n", - "Training on the full train split could take a lot of time, so we are going to clip the train split to 2k examples for the sake of this tutorial, and limit the validation dataset to 200 samples." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "0f1473ba", - "metadata": {}, - "outputs": [], - "source": [ - "! head -2000 $SQUAD_DIR/squad_train.jsonl > $SQUAD_DIR/squad_short_train.jsonl\n", - "! head -200 $SQUAD_DIR/squad_val.jsonl > $SQUAD_DIR/squad_short_val.jsonl\n" - ] - }, - { - "cell_type": "markdown", - "id": "2e19c8dc", - "metadata": {}, - "source": [ - "# P-Tuning Model Config Setup\n", - "\n", - "Now we will begin setting up the config file used for prompt/p-tuning our GPT models! GPT Prompt learning within NeMo uses a class called `MegatronGPTPromptLearningModel` which has its own config file. We will start by loading an example prompt learning config file, then make changes to it to fit our tasks and training plans. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5749c387", - "metadata": {}, - "outputs": [], - "source": [ - "from omegaconf import OmegaConf\n", - "\n", - "CONFIG_DIR = os.path.join(NEMO_DIR, \"conf\")\n", - "os.makedirs(CONFIG_DIR, exist_ok=True)\n", - "\n", - "# Download the example config file\n", - "wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/language_modeling/conf/megatron_gpt_prompt_learning_config.yaml', CONFIG_DIR)\n", - "\n", - "# Load the example config file so we can start editing it\n", - "CONFIG_PATH = os.path.join(CONFIG_DIR, \"megatron_gpt_prompt_learning_config.yaml\")\n", - "config = OmegaConf.load(CONFIG_PATH)" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "ce966bcf", - "metadata": {}, - "source": [ - "First let's set the datasets we've created in the config. We are going to start by p-tuning a GPT model on a small subset of the **Squad** task. We do this by setting the following config params below: " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "6bb1590f", - "metadata": {}, - "outputs": [], - "source": [ - "config.model.data.train_ds = [f\"{SQUAD_DIR}/squad_short_train.jsonl\"]\n", - "config.model.data.validation_ds = [f\"{SQUAD_DIR}/squad_short_val.jsonl\"]" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "4e021b24", - "metadata": {}, - "source": [ - "### Prompt Formatting\n", - "Now that we have our dataset, lets define what we want the prompt to look like. \n", - "\n", - "The squad dataset json files contain fields named \"context\", \"question\" and \"answer\". The prompt formatting template allows us to arrange these fields and decide where to insert virtual prompts. We can add the `<|VIRTUAL_PROMPT_0|>` token anywhere between the fields (although we recommend simply adding it in the leftmost position will be sufficient).\n", - "\n", - "For example, given a data jsonl file with examples like this: \n", - "\n", - "\n", - "**{\"taskname\": \"squad\", \"context\": \"Super Bowl 50 was an American football ga... numerals 50.\", \"question\": \"What does AFC stand for?\", \"answer\": \"American Football Conference\"}**. \n", - "\n", - "\n", - "We can create a prompt template set to `prompt_template = \"<|VIRTUAL_PROMPT_0|> Context: {context}\\n\\nquestion: {question}\\n\\nanswer: {answer}\"` other options are also possible, for example the `\\n` can be replaced with whitespace or the other of the context and question can be swapped. The answer however, should be at the end.\n", - "\n", - "Let's configure the prompt template for the task below:\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "f935b411", - "metadata": {}, - "outputs": [], - "source": [ - "config.model.task_templates = [\n", - " \n", - " {\n", - " \"taskname\": \"squad\",\n", - " \"prompt_template\": \"<|VIRTUAL_PROMPT_0|> Context: {context}\\n\\nQuestion: {question}\\n\\nAnswer:{answer}\",\n", - " \"total_virtual_tokens\": 15,\n", - " \"virtual_token_splits\": [15],\n", - " \"truncate_field\": \"context\",\n", - " \"answer_only_loss\": True,\n", - " \"answer_field\": \"answer\",\n", - " },\n", - " \n", - "]" - ] - }, - { - "cell_type": "markdown", - "id": "dcc438b5", - "metadata": {}, - "source": [ - "Note each `task_template` item has 5 fields. \n", - "\n", - "- **`prompt_template`** is a string showing the model where to place virtual tokens and how to map dataset json fields to where they belong in the model prompt. \n", - "\n", - "\n", - "- **`taskname`** refers to the same `taskname` in the dataset json objects. \n", - "\n", - "\n", - "- **`total_virtual_tokens`** specifies the total number of virtual tokens that will be inserted into the model prompt.\n", - "\n", - "\n", - "- **`virtual_token_splits`** specifies the number of virtual tokens that belong at each `<|VIRTUAL_PROMPT_#|>` marker. `virtual_token_splits` values should add up to `total_virtual_tokens`. The number of `virtual_token_splits` should match the number of `<|VIRTUAL_PROMPT_#|>` markers. \n", - "\n", - "\n", - "- **`truncate_field`** specifies which field in the data json to truncate if the length of the input exceeds the maximum sequence length of the model. If `truncate_field` is set to `None`, examples that are too long are simply dropped from the dataset.\n", - "\n", - "\n", - "- **`answer_only_loss`** Whether to limit loss calculation to only the answer portion of the prompt during tuning. `True` Strongly recommended for long prompts, but shorter prompts with single word answers seem to benefit from setting this to `False`. \n", - "\n", - "\n", - "- **`answer_field`** The field in the data json corresponding to the answer. The loss will only be calculated on this portion of the prompt if `answer_only_loss` is `True`. The answer field must be at the end of the prompt template.\n", - "\n", - "In the `task_templates` we set above, `squad` has a different number of virtual tokens than `sentiment` and `intent_and_slot`. This is because we will be p-tuning on `squad` after we p-tune on the other two tasks and **we do not need to use the same number of virtual tokens between sessions**. We also set the `truncate` field for squad because the context can sometimes be longer than the model's max sequence length, and we want that field to be truncated if the example is too long. Lastly, we set `answer_only_loss` to true for `squad` due to the longer prompt. We've found `answer_only_loss=True` to work significantly better for this task." - ] - }, - { - "cell_type": "markdown", - "id": "84579c7a", - "metadata": {}, - "source": [ - "### Setting New Tasks\n", - "After you p-tune your model this time, you can always go back and p-tune or prompt-tune your model on more tasks without over writing the virtual prompts who've trained this time. You can also use a different number of `total_virtual_tokens` between each training session as long as tasks p-tuned or prompt tuned at the same time have the same number of `total_virtual_tokens`. For this reason, when you p-tune on a new task, you need to tell your model which of your tasks are new and which ones already exist (and thus you don't want to tune them). \n", - "\n", - "You do this by setting the `new_tasks` and `existing_tasks` values in the config file. Because we are p-tuning a model with no existing tasks, you should set `existing_tasks=[]` and `new_tasks=[\"sentiment\", \"intent_and_slot\"]` as follows:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "57a73e01", - "metadata": {}, - "outputs": [], - "source": [ - "config.model.existing_tasks = []\n", - "config.model.new_tasks = [\"squad\"]" - ] - }, - { - "cell_type": "markdown", - "id": "3b77e88c", - "metadata": {}, - "source": [ - "After p-tuning and/or prompt tuning is complete, you can run inference on all tasks at the same time, regardless of their `total_virtual_tokens` value." - ] - }, - { - "cell_type": "markdown", - "id": "a0d5017e", - "metadata": {}, - "source": [ - "### Setting The Pre-Trained GPT Model\n", - "We still need to set which GPT model we want to p-tune/prompt tune. Prompt learning methods work best with large GPT language models (5B or above), but the purposes of this tutorial, we are going to download a 345M parameter GPT model from NVIDIA NGC." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "48cdf868", - "metadata": {}, - "outputs": [], - "source": [ - "# Check what GPT .nemo models we have available on NGC\n", - "from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel\n", - "MegatronGPTModel.list_available_models()" - ] - }, - { - "cell_type": "markdown", - "id": "ede350ed", - "metadata": {}, - "source": [ - "If we wanted to use the GPT model class directly, we could instantiate a trainer then download the model by calling running \n", - "`gpt_model = MegatronGPTModel.from_pretrained(model_name=\"megatron_gpt_345m\", trainer=trainer).cuda()`. But we just need the `.nemo` file in our working NeMo directory in this tutorial, so we will download it using `wget`. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "364439a1", - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "# Download the model from NGC\n", - "gpt_file_name = \"megatron_gpt_345m.nemo\"\n", - "!wget -nc --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/nemo/megatron_gpt_345m/versions/1/files/megatron_gpt_345m.nemo -O {NEMO_DIR}/{gpt_file_name}" - ] - }, - { - "cell_type": "markdown", - "id": "1d6a8a67", - "metadata": {}, - "source": [ - "Now that we have a `.nemo` GPT file to work with. We need to add its path in our prompt learning config. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2778a5fa", - "metadata": {}, - "outputs": [], - "source": [ - "# Set GPT model path on prompt learning config\n", - "config.model.language_model_path = gpt_file_name" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "943a9c83", - "metadata": {}, - "source": [ - "We can also set where we want the final prompt tuned model to be saved by setting `model.nemo_path`. By default the tuned prompt learning model will be saved in your current working directory to a `.nemo` file with the same name as your experiment (`config.name`). Let's change the save name to be `p_tuned_gpt.nemo`. **Your model path must end in `.nemo`.**" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "a278cbdf", - "metadata": {}, - "outputs": [], - "source": [ - "config.exp_manager.checkpoint_callback_params.save_nemo_on_train_end= True\n", - "config.exp_manager.checkpoint_callback_params.always_save_nemo= True\n", - "config.exp_manager.checkpoint_callback_params.save_best_model= True" - ] - }, - { - "cell_type": "markdown", - "id": "378a73e7", - "metadata": {}, - "source": [ - "### Setting P-Tuning Specific Params\n", - "Within the config file, p-tuning and prompt-tuning each have a couple of hyperparameters specific to them. We first need to tell the model that we want to do p-tuning, not prompt-tuning. To do this, we set the **`model.virtual_prompt_style`** hyperparameter like this:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "68763763", - "metadata": {}, - "outputs": [], - "source": [ - "from nemo.collections.nlp.modules.common import VirtualPromptStyle\n", - "config.model.virtual_prompt_style = VirtualPromptStyle.P_TUNING" - ] - }, - { - "cell_type": "markdown", - "id": "947dec63", - "metadata": {}, - "source": [ - "Then we can set the 2 p-tuning specific parameters. Reminder, p-tuning uses an LSTM prompt encoder to predict virtual tokens. \n", - "\n", - "- **`p_tuning.dropout`** the LSTM prompt encoder dropout probability \n", - "- **`p_tuning.num_layers`** the number of LSTM layers you want your p-tuning prompt encoder to have\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "03f893ef", - "metadata": {}, - "outputs": [], - "source": [ - "config.model.p_tuning.dropout = 0.0\n", - "config.model.p_tuning.num_layers = 2\n", - "config.model.global_batch_size = 2\n", - "config.model.micro_batch_size = 1" - ] - }, - { - "cell_type": "markdown", - "id": "a988d16e", - "metadata": {}, - "source": [ - "Let's have a look at all the values we've set in the model config. You can change any of these values in the same manner we've been using above. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "12a37ada", - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "# Final model config\n", - "print(OmegaConf.to_yaml(config.model))" - ] - }, - { - "cell_type": "markdown", - "id": "6b4bc7f3", - "metadata": {}, - "source": [ - "### Setting Prompt-Tuning Specific Params\n", - "\n", - "Though we are not using prompt tuning in this training session, let's go over the prompt tuning specific parameters we would use if we were. \n", - "\n", - "- **`prompt_tuning.new_prompt_init_methods`** Whether you want to initialize virtual token embeddings from the embeddings of existing parts of the model's vocabulary (either 'text' or 'random')\n", - "- **`prompt_tuning.new_prompt_init_text`** The text you want to use if you have 'text' in the list above, should be None otherwise. \n", - "\n", - "Each of the above hyperparameters are a list of strings. \n", - "\n", - "`new_prompt_init_methods` would look like `[\"text\", \"random\", \"text\", \"text\"]` if you were prompt tuning on 4 tasks at once, and you wanted the second task in `new_tasks` to use random initialization. \n", - "\n", - "`new_prompt_init_text` might look like `[\"some text I want to use\", None, \"some other text\", \"task text goes here\"]` for those four new tasks. \n", - "\n", - "The order of both should correspond to the order of the tasks you have listed in `model.new_tasks`. " - ] - }, - { - "cell_type": "markdown", - "id": "4c048852", - "metadata": {}, - "source": [ - "# Building the PyTorch Lightning Trainer\n", - "NeMo models are primarily PyTorch Lightning modules - and therefore are entirely compatible with the PyTorch Lightning ecosystem.\n", - "\n", - "Let's first instantiate a Trainer object" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "90f85b2a", - "metadata": {}, - "outputs": [], - "source": [ - "import torch\n", - "import pytorch_lightning as pl\n", - "from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategyNotebook\n", - "from pytorch_lightning.plugins.environments import TorchElasticEnvironment\n", - "\n", - "# let's modify some trainer configs\n", - "# check if we have GPU available and uses it\n", - "accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n", - "config.trainer.accelerator = accelerator\n", - "config.trainer.devices = 1\n", - "config.trainer.max_epochs = 4\n", - "config.trainer.val_check_interval = 1.0\n", - "\n", - "# for PyTorch Native AMP set precision=16\n", - "config.trainer.precision = 16 if torch.cuda.is_available() else 32\n", - "\n", - "# setup cluster environment parameters\"\n", - "# use torch elastic cluster environment so `create_process_externally` is True\n", - "# the launcher is set to None. It will not try to spawn new processes.\n", - "# It won't create the misconfiguration error because of the `interactive session`\n", - "os.environ[\"LOCAL_RANK\"] = '0'\n", - "os.environ[\"RANK\"] = '0'\n", - "os.environ[\"WORLD_SIZE\"] = '1'\n", - "\n", - "strategy = NLPDDPStrategyNotebook(find_unused_parameters=False, no_ddp_communication_hook=True)\n", - "plugins = [TorchElasticEnvironment()]\n", - "trainer = pl.Trainer(plugins= plugins, strategy=strategy, **config.trainer)\n", - "\n", - "print(\"Trainer config - \\n\")\n", - "print(OmegaConf.to_yaml(config.trainer))" - ] - }, - { - "cell_type": "markdown", - "id": "4d0124c1", - "metadata": {}, - "source": [ - "# Setting up a NeMo Experiment\n", - "\n", - "NeMo has an experiment manager that handles logging and checkpointing for us, so let's use it:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "f2c943ba", - "metadata": {}, - "outputs": [], - "source": [ - "from nemo.utils.exp_manager import exp_manager\n", - "\n", - "# Set name of the experiment \n", - "config.name = 'p_tuning'\n", - "config.exp_manager.resume_if_exists = False\n", - "\n", - "# Init the experiment manager and view the exp_dir\n", - "exp_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n", - "exp_dir = str(exp_dir)\n", - "print(exp_dir)" - ] - }, - { - "cell_type": "markdown", - "id": "5860bd90", - "metadata": {}, - "source": [ - "We can also set learning hyperparameters as follows:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4c4ec542", - "metadata": {}, - "outputs": [], - "source": [ - "# Set some of the learning parameters\n", - "config.model.optim.lr = 1e-4\n", - "config.model.precision = config.trainer.precision" - ] - }, - { - "cell_type": "markdown", - "id": "298b3dce", - "metadata": {}, - "source": [ - "# First P-Tuning Session\n", - "The only thing left to do is load up the model and begin p-tuning!" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "b4bda19b", - "metadata": {}, - "outputs": [], - "source": [ - "from nemo.collections.nlp.models.language_modeling.megatron_gpt_prompt_learning_model import MegatronGPTPromptLearningModel\n", - "\n", - "model = MegatronGPTPromptLearningModel(cfg=config.model, trainer=trainer)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2d99f433", - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "# Training set to 2 epochs by default in a cell above\n", - "# Each epoch will take around 1min 15sec, but training time can vary\n", - "trainer.fit(model)" - ] - }, - { - "attachments": {}, - "cell_type": "markdown", - "id": "6aab09d4", - "metadata": {}, - "source": [ - "# Inference After P-Tuning\n", - "One way to run inference after p-tuning or prompt-tuning your model is to call `model.generate()`. `model.generate()` takes in \n", - "\n", - "- `inputs` which can be either a list of dictionary objects or `.jsonl` files containing dictionary objects, \n", - "- `length_params`\n", - "- `sampling_params`\n", - "\n", - "as arguments. More information about the [text generation API can be found here](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/modules/common/transformer/text_generation.py).\n", - "\n", - "If `length_params` and `sampling_params` are set to `None`, the model generates output with a greedy decoding strategy and generates up to `30` new tokens. Most predictive downstream tasks (not text generation tasks), use greedy sampling. To see other ways to run inference with your prompt learning model and more details on how to define various inference parameters, visit `examples/nlp/language_modeling/megatron_gpt_eval.py`.\n", - "\n", - "Below are some randomly selected test examples from the sentiment classification and intent and slot classification test files. Notice that the `label` field is dropped from all test examples. The `MegatronPromptLearningDataset` called within `.generate()` automatically leaves fields in the prompt template empty when they are not provided in the data json. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "dc95e764", - "metadata": {}, - "outputs": [], - "source": [ - "test_examples = [\n", - " {\"taskname\": \"squad\", \"context\": \"The build was released for download later in the day in standard 32-bit and 64-bit versions, plus a special 64-bit version which included SDKs and developer tools (Visual Studio Express and Expression Blend) for developing Metro-style apps. The Windows Store was announced during the presentation, but was not available in this build. According to Microsoft, there were about 535,000 downloads of the developer preview within the first 12 hours of its release. Originally set to expire on March 11, 2012, in February 2012 the Developer Preview's expiry date was changed to January 15, 2013.\", \"question\": \"When was the Developer preview initially intended to expire?\"},\n", - " {\"taskname\": \"squad\", \"context\": \"The structures of most federal governments incorporate mechanisms to protect the rights of component states. One method, known as 'intrastate federalism', is to directly represent the governments of component states in federal political institutions. Where a federation has a bicameral legislature the upper house is often used to represent the component states while the lower house represents the people of the nation as a whole. A federal upper house may be based on a special scheme of apportionment, as is the case in the senates of the United States and Australia, where each state is represented by an equal number of senators irrespective of the size of its population.\", \"question\": \"What is a bicameral legislature?\"},\n", - " {\"taskname\": \"squad\", \"context\": \"Imported mystery religions, which offered initiates salvation in the afterlife, were a matter of personal choice for an individual, practiced in addition to carrying on one's family rites and participating in public religion. The mysteries, however, involved exclusive oaths and secrecy, conditions that conservative Romans viewed with suspicion as characteristic of \\\"magic\\\", conspiratorial (coniuratio), or subversive activity. Sporadic and sometimes brutal attempts were made to suppress religionists who seemed to threaten traditional morality and unity, as with the senate's efforts to restrict the Bacchanals in 186 BC.\", \"question\": \"What was the practice of religion to the Romans?\"}\n", - "]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "74a5a358", - "metadata": { - "scrolled": true - }, - "outputs": [], - "source": [ - "response = model.generate(inputs=test_examples, length_params=None)\n", - "\n", - "print('The prediction results of some sample queries with the trained model:')\n", - "for result in response['sentences']:\n", - " print(result)\n", - " print(\"-\" * 30)" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.8.16" - } - }, - "nbformat": 4, - "nbformat_minor": 5 - }