P-tuning refactor Part 2/N

Jenkinsfile

@@ -3361,7 +3361,8 @@ assert_frame_equal(training_curve, gt_curve, rtol=1e-3, atol=1e-3)"'''
                 trainer.max_epochs=null \
                 model.data.num_workers=1 \
                 model.tensor_model_parallel_size=1 \
-                model.virtual_prompt_style='prompt-tuning' \
+                model.virtual_prompt_style='p-tuning' \
+                model.p_tuning.encoder_type='embedding' \
                 model.language_model_path='/home/TestData/nlp/megatron_gpt/tiny/megatron_14m_gpt_tp1_pp1.nemo' \
                 model.existing_tasks=[] \
                 model.new_tasks=['rte'] \

nemo/collections/nlp/data/language_modeling/megatron/gpt_prompt_learning_dataset.py

@@ -34,7 +34,7 @@ class GPTPromptLearningDataset(Dataset):
     Args:
         data (list[strings], list[dicts]): (1) paths to .jsonl or .json files, (2) dict objects corresponding to each input example
         tokenizer (tokenizer): Tokenizer from frozen language model
-        virtual_prompt_source (Enum): Either VirtualPromptSource.PROMPT_TABLE or VirtualPromptSource.PROMPT_ENCODER
+        virtual_prompt_source (Enum): Either VirtualPromptSource.NO_PROMPTS or VirtualPromptSource.PROMPT_ENCODER
         task_templates (dict): Dictionary containing all task template information needed to format prompts. Created in the GPTPromptLearningModel class.
         pseudo_tokens (list[strings]): A list of virtual prompt token placeholders e.g [<prompt_1>, <prompt_2>, ...] up to max num virtual tokens
         pad_token_id (int): ID of pad token from tokenizer
@@ -179,10 +179,6 @@ def load_data(self, dataset):
             if self.min_seq_length <= len(input_ids) <= self.max_seq_length:
                 if self.virtual_prompt_source == VirtualPromptSource.PROMPT_ENCODER:
                     taskname_id = self.tokenizer.text_to_ids(taskname)
-
-                elif self.virtual_prompt_source == VirtualPromptSource.PROMPT_TABLE:
-                    taskname_id = self.task_templates[taskname]["task_id_num"]
-
                 elif self.virtual_prompt_source == VirtualPromptSource.NO_PROMPT:
                     taskname_id = -1
                 else:
@@ -342,7 +338,7 @@ def collate_fn(self, batch, tp_workers=0):
             taskname_ids = torch.tensor(taskname_ids)
 
         # Task ids are just used for a look up embeddings for prompt-table
-        elif self.virtual_prompt_source in [VirtualPromptSource.PROMPT_TABLE, VirtualPromptSource.NO_PROMPT]:
+        elif self.virtual_prompt_source == VirtualPromptSource.NO_PROMPT:
             taskname_ids = torch.tensor(taskname_ids)
 
         # Get max sequence length of batch

nemo/collections/nlp/models/language_modeling/megatron_base_prompt_learning_model.py

@@ -237,7 +237,10 @@ def add_ptuned_prompts_to_prompt_table(self):
                 device = next(self.word_embeddings.parameters()).device
                 tokenized_taskname = torch.tensor(self.tokenizer.text_to_ids(taskname)).to(device)
                 taskname_embeddings = self.word_embeddings(tokenized_taskname).unsqueeze(0)
-                virtual_prompt_embeddings = self.prompt_encoder(taskname_embeddings=taskname_embeddings).squeeze(0)
+                batch_size = taskname_embeddings.shape[0]
+                virtual_prompt_embeddings = self.prompt_encoder(batch_size=batch_size, use_cached_reps=False).squeeze(
+                    0
+                )
                 total_virtual_tokens = self.task_templates[taskname]["total_virtual_tokens"]
                 self.prompt_table.add_prompt_from_p_tuning_encoder(
                     taskname, virtual_prompt_embeddings, total_virtual_tokens
@@ -343,8 +346,11 @@ def embed_input_train(self, input_ids: Tensor, taskname_ids: Tensor):
             virtual_token_embeds = torch.stack(virtual_token_embeds)
 
         elif self.virtual_prompt_source == VirtualPromptSource.PROMPT_ENCODER:
-            taskname_embeddings = self.word_embeddings(taskname_ids)
-            virtual_token_embeds = self.prompt_encoder(taskname_embeddings=taskname_embeddings)
+            # taskname_embeddings = self.word_embeddings(taskname_ids)
+            batch_size, _ = taskname_ids.size()
+            virtual_token_embeds = self.prompt_encoder(batch_size=batch_size, use_cached_reps=False)
+        else:
+            raise RuntimeError("invalid VirtualPromptSource..")
 
         # Create index template specifying where virtual token embeddings should be placed
         batch_size, _, embedding_size = discrete_token_embeds.shape