Enable multiline prompts by setting a custom prompt end (EleutherAI#754)

* add multiline prompt support to input-file mode

* reroll the input in the interactive mode

* add comments about customizable prompt ends

* reformat files modified in EleutherAI#754 with Black

* add comments about 'input' stripping '\n's

* Formatting and spelling

Co-authored-by: Quentin Anthony <qganthony@yahoo.com>

README.md

@@ -146,10 +146,10 @@ To reproduce our evaluation numbers on, for example, TriviaQA and PIQA use:
 
 You can add an arbitrary list of evaluation tasks here, for details of all tasks available, see [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness).
 
-For more details on each entry point, see the [Training and Finetuning](#training-and-finetuning), [Inference](#inference) and [Evaluation](#evaluation) 
+For more details on each entry point, see the [Training and Finetuning](#training-and-finetuning), [Inference](#inference) and [Evaluation](#evaluation)
 # Configuration
 
-GPT-NeoX parameters are defined in a YAML configuration file which is passed to the deepy.py launcher. We have provided some example .yaml files in [configs](./configs/), including one for GPT-NeoX-20B, and example configuration files for other model sizes. 
+GPT-NeoX parameters are defined in a YAML configuration file which is passed to the deepy.py launcher. We have provided some example .yaml files in [configs](./configs/), including one for GPT-NeoX-20B, and example configuration files for other model sizes.
 
 These files are generally complete, but non-optimal. For example, depending on your specific GPU configuration, you may need to change some settings such as `pipe-parallel-size`, `model-parallel-size` to increase or decrease the degree of parallelisation, `train_micro_batch_size_per_gpu` or `gradient-accumulation-steps` to modify batch size related settings, or the `zero_optimization` dict to modify how optimizer states are parallelised across workers.
 

configs/neox_arguments.md

@@ -958,7 +958,7 @@ Text Generation arguments
 
 - **eval_results_prefix**: str
 
-    Default = 
+    Default =
 
     prefix to which to save evaluation results - final fp will be {eval_results_prefix}_eval_results_yy-mm-dd-HH-MM.json
 
@@ -1155,10 +1155,10 @@ Training Arguments
 
     Acts as a multiplier on either the "log" or "linear" checkpoint spacing.
 
-    With `checkpoint-scale="linear"`, `checkpoint-factor=20`, and `train-iters=100`, checkpoints will be saved at 
+    With `checkpoint-scale="linear"`, `checkpoint-factor=20`, and `train-iters=100`, checkpoints will be saved at
     steps [20, 40, 60, 80, 100].
 
-    With `checkpoint-scale="log"`, `checkpoint-factor=2`, and `train-iters=100`, checkpoints will be saved at 
+    With `checkpoint-scale="log"`, `checkpoint-factor=2`, and `train-iters=100`, checkpoints will be saved at
     steps [1, 2, 4, 8, 16, 32, 64, 100].
 
     Note that the last checkpoint step is always saved.
@@ -1572,7 +1572,7 @@ Args for deepspeed config
 
     Default = None
 
-    
+
 
 
 
@@ -1706,4 +1706,3 @@ Args for deepspeed runner (deepspeed.launcher.runner).
     Default = None
 
     Adds a `--comment` to the DeepSpeed launch command. In DeeperSpeed this is passed on to the SlurmLauncher as well. Sometime necessary for cluster rules, or so I've heard.
-

configs/text_generation.yml

@@ -6,6 +6,7 @@
 
   # Params for all
   "maximum_tokens": 102,
+  "prompt_end": "\n",
   "temperature": 1.0,
   "top_p": 0.0,
   "top_k": 0,

generate.py

@@ -62,6 +62,7 @@ def main():
             input_file=neox_args.sample_input_file,
             output_file=neox_args.sample_output_file,
             maximum_tokens=neox_args.maximum_tokens,
+            prompt_end=neox_args.prompt_end,
             recompute=neox_args.recompute,
             temperature=neox_args.temperature,
             top_k=neox_args.top_k,
@@ -75,6 +76,7 @@ def main():
             recompute=neox_args.recompute,
             temperature=neox_args.temperature,
             maximum_tokens=neox_args.maximum_tokens,
+            prompt_end=neox_args.prompt_end,
             top_k=neox_args.top_k,
             top_p=neox_args.top_p,
         )

megatron/model/transformer.py

@@ -663,15 +663,15 @@ def forward(self, x, attention_mask, layer_past=None):
             # to save communication time (we can do a single allreduce after we add mlp / attn outputs).
             # due to a bug, the two layernorms are not tied in GPT-NeoX-20B. This is non-desirable, but
             # we preserve the functionality for backwards compatibility
-            
+
             residual = x
             # applies the correct normalization depending on if the norms are tied
             if self.gpt_j_tied:
                 x1, x2 = self.input_layernorm(x), self.post_attention_layernorm(x)
             else:
                 x = self.input_layernorm(x)
                 x1, x2 = x, x
-                
+
             # attention operator
             attention_output, attention_bias = self.attention(
                 x1, attention_mask, layer_past=layer_past
@@ -699,7 +699,7 @@ def forward(self, x, attention_mask, layer_past=None):
                 )
 
             # output = (x + attn(ln(x)) + mlp(ln(x))
-            output   = residual         + self.reduce(output)
+            output = residual + self.reduce(output)
         else:
             # pseudocode:
             # x = x + attn(ln1(x))

megatron/neox_arguments/arguments.py

@@ -741,18 +741,18 @@ def calculate_derived(self):
                 save_iters = set(self.extra_save_iters)
             else:
                 save_iters = set()
-            
-            step = self.checkpoint_factor # don't save step 0 or 1
+
+            step = self.checkpoint_factor  # don't save step 0 or 1
             while step < self.train_iters:
                 save_iters.add(step)
                 if self.checkpoint_scale == "log":
                     step *= self.checkpoint_factor
                 elif self.checkpoint_scale == "linear":
                     step += self.checkpoint_factor
-            
+
             save_iters = list(save_iters)
             save_iters.sort()
-        
+
         self.update_values(
             {
                 "save_iters": save_iters,
@@ -848,7 +848,7 @@ def calculate_derived(self):
         if self.sparsity_config is None:
             # Can't have a default value as an empty dict so need to set it here
             self.update_value("sparsity_config", {})
-        
+
         # Adding equal dataset weights if none are provided
         if self.train_data_paths and (self.train_data_weights is None):
             self.train_data_weights = [1.0] * len(self.train_data_paths)
@@ -947,7 +947,11 @@ def validate_values(self):
             raise ValueError(error_message)
             return False
 
-        if self.save is not None and self.checkpoint_factor is None and self.extra_save_iters is None:
+        if (
+            self.save is not None
+            and self.checkpoint_factor is None
+            and self.extra_save_iters is None
+        ):
             error_message = (
                 self.__class__.__name__
                 + ".validate_values() checkpoint_factor or extra_save_iters must be defined if save is defined"

megatron/neox_arguments/neox_args.py

@@ -332,7 +332,7 @@ class NeoXArgsModel(NeoXArgsTemplate):
       x = ln(x)
       x = x + attn(x) + mlp(x)
     """
-    
+
     gpt_j_tied: bool = False
     """
     If false, we use
@@ -785,10 +785,10 @@ class NeoXArgsTraining(NeoXArgsTemplate):
     """
     Acts as a multiplier on either the "log" or "linear" checkpoint spacing.
 
-    With `checkpoint-scale="linear"`, `checkpoint-factor=20`, and `train-iters=100`, checkpoints will be saved at 
+    With `checkpoint-scale="linear"`, `checkpoint-factor=20`, and `train-iters=100`, checkpoints will be saved at
     steps [20, 40, 60, 80, 100].
 
-    With `checkpoint-scale="log"`, `checkpoint-factor=2`, and `train-iters=100`, checkpoints will be saved at 
+    With `checkpoint-scale="log"`, `checkpoint-factor=2`, and `train-iters=100`, checkpoints will be saved at
     steps [1, 2, 4, 8, 16, 32, 64, 100].
 
     Note that the last checkpoint step is always saved.
@@ -1008,6 +1008,11 @@ class NeoXArgsTextgen(NeoXArgsTemplate):
     maximum number of tokens to be generated
     """
 
+    prompt_end: str = "\n"
+    """
+    a single prompt's end. Defaults to newline
+    """
+
     sample_input_file: str = None
     """
     Get input from file instead of interactive mode, each line is an input.

megatron/text_generation_utils.py

@@ -552,6 +552,7 @@ def generate_samples_input_from_file(
     output_file=None,
     eos_token_id: int = None,
     maximum_tokens: int = 64,
+    prompt_end: str = "\n",
     recompute: bool = False,
     temperature: float = 0.0,
     top_k: int = 0,
@@ -570,6 +571,7 @@ def generate_samples_input_from_file(
 
     eos_token_id: end of text token at which completion is terminated, even if max_tokes count has not been reached
     maximum_tokens: maximum number of tokens to be generated
+    prompt_end: end of a single input prompt. Defaults to newline character '\n'. Other prompt-end sequences may be useful when generating indent-aware completions (e.g. code)
 
     recompute: flag indicating whether a cache is used for already forwarded tokens (true) or whether all tokens are recomputed at every iteration (false)
 
@@ -592,8 +594,9 @@ def generate_samples_input_from_file(
     print_rank_0(
         "generate_samples_input_from_file() loading input from {}".format(input_file)
     )
-    with open(input_file, "r") as f:
-        prompts = f.readlines()
+    with open(input_file, "r", encoding="utf-8") as f:
+        prompts = f.read()
+        prompts = prompts.split(prompt_end)
     prompts = [p.strip() for p in prompts]
     prompts = [p for p in prompts if len(p) > 0]
     print_rank_0(
@@ -654,6 +657,7 @@ def generate_samples_unconditional(
 
     eos_token_id: end of text token at which completion is terminated, even if max_tokes count has not been reached
     maximum_tokens: maximum number of tokens to be generated
+    prompt_end: end of a single input prompt. Defaults to newline character '\n'. Other prompt-end sequences may be useful when generating indent-aware completions (e.g. code). The interactive mode will reroll the user-input request until the stop-char is met
 
     recompute: flag indicating whether a cache is used for already forwarded tokens (true) or whether all tokens are recomputed at every iteration (false)
 
@@ -699,6 +703,7 @@ def generate_samples_interactive(
     neox_args,
     model,
     maximum_tokens: int = 64,
+    prompt_end: str = "\n",
     eos_token_id: int = None,
     recompute: bool = False,
     temperature: float = 0.0,
@@ -738,7 +743,20 @@ def generate_samples_interactive(
 
         if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
             os.system("clear")
-            raw_text = input("Context prompt >>> ")
+            raw_text = ""
+            while True:
+                current_input = input("Context prompt >>> ")
+                if (
+                    prompt_end == "\n"
+                ):  # we need to handle '\n' case as 'input' strips it and leads to lines being squashed
+                    raw_text += current_input
+                    break
+                if prompt_end in current_input:
+                    raw_text += current_input.split(prompt_end)[0]
+                    break
+                raw_text += (
+                    current_input + "\n"
+                )  # re-add newline since we stripped it on input
             context_tokens = neox_args.tokenizer.tokenize(raw_text)
             if len(context_tokens) == 0:
                 context_tokens = [neox_args.tokenizer.eod]

megatron/training.py

@@ -112,7 +112,7 @@ def pretrain(neox_args):
                 optimizer=optimizer,
                 lr_scheduler=lr_scheduler,
             )
-            
+
         iteration = train(
             neox_args=neox_args,
             timers=timers,
@@ -615,10 +615,7 @@ def train(
         )
 
         # Checkpointing
-        if (
-            neox_args.save
-            and iteration in neox_args.save_iters
-        ):
+        if neox_args.save and iteration in neox_args.save_iters:
             save_checkpoint(
                 neox_args=neox_args,
                 iteration=iteration,

tools/convert_to_hf.py

@@ -101,16 +101,16 @@ def __init__(self, neox_config):
             1  # pad defaulting to 1. follows convention from GPT-NeoX-20b tokenizer
         )
 
-
     # TODO: change the default value here based on discussion regarding `gpt_j_tied` config parameter's default
-    use_tied_lns = get_key(neox_config, 'gpt-j-tied', False)
+    use_tied_lns = get_key(neox_config, "gpt-j-tied", False)
 
     if use_tied_lns:
         raise NotImplementedError(
-                """ERROR: Huggingface Transformers does not yet support a single shared layernorm 
+            """ERROR: Huggingface Transformers does not yet support a single shared layernorm
                 per transformer block for GPT-NeoX models trained  w/ GPT-J parallel residuals.
-                See https://github.com/EleutherAI/gpt-neox/pull/481 for further details.""")
-
+                See https://github.com/EleutherAI/gpt-neox/pull/481 for further details."""
+        )
+
     # set all config values.
     hf_config = GPTNeoXConfig(
         vocab_size=args.padded_vocab_size,