detokenization parallelization

README.md

@@ -10,11 +10,11 @@ Below shows the generation speed gain by using FastSeq.
 | Model | W/O FastSeq (in samples/s) | W/ FastSeq (in samples/s) | Speedup |
 |------------------|:--------------------------:|:-------------------------:|:-----:|
 | [ProphetNet](examples/prophetnet/README.md) | 2.7 | 10.3 | 3.8x |
-| [Bart (`fs`)](examples/bart/README.md) | 2.7 | 14.5 | 5.4x |
-| [Bart (`hf`)](examples/bart/README.md#speedup-bart-huggingface-transformers-version-by-using-fastseq) | 3.4 | 6.4 | 1.9x |
-| [DistilBart (`hf`)](examples/distilbart/README.md) | 4.0 | 6.5 | 1.6x |
-| [T5 (`hf`)](examples/t5/README.md) | 4.8 | 7.5 | 1.6x |
-| [WMT16 En-De (`fs`)](examples/wmt/README.md) | 84.0 | 135.0 | 1.6x |
+| [Bart (`fs`)](examples/bart/README.md) | 2.7 | 13.3 | 5x |
+| [Bart (`hf`)](examples/bart/README.md#speedup-bart-huggingface-transformers-version-by-using-fastseq) | 3.4 | 11.0 | 3.2x |
+| [DistilBart (`hf`)](examples/distilbart/README.md) | 4.0 | 13.5 | 3.4x |
+| [T5 (`hf`)](examples/t5/README.md) | 4.8 | 17.0 | 3.5x |
+| [WMT16 En-De (`fs`)](examples/wmt/README.md) | 84.0 | 124.0 | 1.5x |
 
 - All the following benchmarking experiments run on NVIDIA-V100-16GB with [docker](docker/Dockerfile). Highest speed recorded for each model by tuning batch size. For parameter setting details, click link of corresponding model.
 - `fs` stands for [Fairseq](https://github.com/pytorch/fairseq) 0.9.0 version, `hf` stands for [Huggingface Transformers](https://github.com/huggingface/transformers) 3.0.2 version.

benchmarks/models/hf_bart.sh

@@ -20,9 +20,9 @@ source utils.sh
 grep "facebook/bart-large-cnn cnn_dm.1k/raw val " perf | awk '{print $9}' | awk -F'|' '{if($1!="NA"){c+=1;s+=$1}}END{print s/c}' | bash range.sh 34.8 35
 # Speed on V100 16GB 250W
 grep -E "transformers_v3.0.2 facebook/bart-large-cnn cnn_dm.1k/raw val 32 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 3.2 3.4
-grep -E "transformers_v3.0.2\+fastseq_v.* facebook/bart-large-cnn cnn_dm.1k/raw val 32 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 5.2 100
-grep -E "transformers_v3.0.2\+fastseq_v.* facebook/bart-large-cnn cnn_dm.1k/raw val 64 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 6.2 100
-grep -E "transformers_v3.0.2\+fastseq_v.* facebook/bart-large-cnn cnn_dm.1k/raw val 128 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 6.4 100
+grep -E "transformers_v3.0.2\+fastseq_v.* facebook/bart-large-cnn cnn_dm.1k/raw val 32 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 7.9 100
+grep -E "transformers_v3.0.2\+fastseq_v.* facebook/bart-large-cnn cnn_dm.1k/raw val 64 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 10.7 100
+grep -E "transformers_v3.0.2\+fastseq_v.* facebook/bart-large-cnn cnn_dm.1k/raw val 128 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 11.0 100
 
 ## Accuracy
 #grep "facebook/bart-large-cnn cnn_dm/raw val " perf | awk '{print $9}' | awk -F'|' '{if($1!="NA"){c+=1;s+=$1}}END{print s/c}' | bash range.sh 44.78 44.82

benchmarks/models/hf_distibart.sh

@@ -20,9 +20,9 @@ source utils.sh
 grep "sshleifer/distilbart-cnn-12-6 cnn_dm.1k/raw val " perf | awk '{print $9}' | awk -F'|' '{if($1!="NA"){c+=1;s+=$1}}END{print s/c}' | bash range.sh 35.1 35.3
 # Speed on V100 16GB 250W
 grep -E "transformers_v3.0.2 sshleifer/distilbart-cnn-12-6 cnn_dm.1k/raw val 64 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 3.9 4.2
-grep -E "transformers_v3.0.2\+fastseq_v.* sshleifer/distilbart-cnn-12-6 cnn_dm.1k/raw val 64 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 6.5 100
+grep -E "transformers_v3.0.2\+fastseq_v.* sshleifer/distilbart-cnn-12-6 cnn_dm.1k/raw val 64 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 13.3 100
 # todo: bigger bs doesn't increase speed
-grep -E "transformers_v3.0.2\+fastseq_v.* sshleifer/distilbart-cnn-12-6 cnn_dm.1k/raw val 128 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 6.5 100
+grep -E "transformers_v3.0.2\+fastseq_v.* sshleifer/distilbart-cnn-12-6 cnn_dm.1k/raw val 128 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 13.5 100
 
 ## Accuracy
 #grep "sshleifer/distilbart-cnn-12-6 cnn_dm/raw val " perf | awk '{print $9}' | awk -F'|' '{if($1!="NA"){c+=1;s+=$1}}END{print s/c}' | bash range.sh 45 45.1

benchmarks/models/hf_mbart.sh

@@ -14,5 +14,5 @@ source utils.sh
 # Accuracy
 grep "facebook/mbart-large-en-ro wmt_en_ro/raw val " perf | awk '{if($8!="NA"){c+=1;s+=$8}}END{print s/c}' | bash range.sh 27.79 27.95
 # Speed on V100 16GB 250W
-grep -E "transformers_v3.0.2 facebook/mbart-large-en-ro wmt_en_ro/raw val 64 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 5.8 6.2
+grep -E "transformers_v3.0.2 facebook/mbart-large-en-ro wmt_en_ro/raw val 64 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 7.6 7.7
 grep -E "transformers_v3.0.2\+fastseq_v.* facebook/mbart-large-en-ro wmt_en_ro/raw val 64 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 6.0 100

benchmarks/models/hf_t5.sh

@@ -14,6 +14,5 @@ source utils.sh
 grep "t5-base wmt_en_ro/raw val " perf | awk '{if($8!="NA"){c+=1;s+=$8}}END{print s/c}' | bash range.sh 27.42 27.44
 # Speed on V100 16GB 250W
 grep -E "transformers_v3.0.2 t5-base wmt_en_ro/raw val 64 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 4.6 5.2
-grep -E "transformers_v3.0.2\+fastseq_v.* t5-base wmt_en_ro/raw val 64 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 6.8 7.3
-grep -E "transformers_v3.0.2\+fastseq_v.* t5-base wmt_en_ro/raw val 128 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 7.5 8.0
-
+grep -E "transformers_v3.0.2\+fastseq_v.* t5-base wmt_en_ro/raw val 64 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 14.4 14.8
+grep -E "transformers_v3.0.2\+fastseq_v.* t5-base wmt_en_ro/raw val 128 " perf | awk '{s+=$13}END{print s/NR}' | bash range.sh 16.8 17.0

fastseq_cli/transformers_generate.py

@@ -2,21 +2,115 @@
 import argparse
 import json
 from pathlib import Path
-
-import torch
+from multiprocessing import Process, Queue
 from tqdm import tqdm
-
-from fastseq_cli.transformers_utils import use_task_specific_params, trim_batch, calculate_rouge, calculate_bleu_score
+import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from fastseq_cli.transformers_utils import use_task_specific_params, trim_batch, calculate_rouge, calculate_bleu_score
 
 DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+GENERATE_FINISHED = 'done'
+POSTPROCESS_FINISHED = None
+
+class TokenizeDataset(torch.utils.data.Dataset):
+ """Characterizes a dataset for PyTorch"""
+ def __init__(self, examples, tokenizer, model_name, prefix):
+ """Multiprocess Dataloader.
+ Args:
+ examples (List(str)): a list of input sentences.
+ tokenizer (AutoTokenizer): instance of AutoTokenizer.
+ model_name (string): model name.
+ prefix (string): input example prefix if any. 
+ """ 
+ self.examples = examples
+ self.tokenizer= tokenizer
+ self.model_name = model_name
+ self.prefix = prefix
+ self.return_tensors="pt"
+ self.truncation=True
+ self.padding="max_length"
+
+ def __len__(self):
+ return len(self.examples)
+
+ def __getitem__(self, index):
+ batch = self.examples[index]
+ if "t5" in self.model_name:
+ batch = self.prefix + batch
+ batch = self.tokenizer(batch,
+ return_tensors=self.return_tensors,
+ truncation=self.truncation,
+ padding=self.padding)
+ return batch['input_ids'], batch['attention_mask']
+
+class IOProcess (Process):
+ """ Write detokenized output to file in order."""
+ def __init__(self, msg_queue, fout):
+ super(IOProcess, self).__init__()
+ self.msg_queue = msg_queue
+ self.fout = fout
+ self.waiting_for=0
+ self.dec_buf = {}
+
+ def process_dec(self, dec):
+ for hypothesis in dec:
+ self.fout.write(hypothesis + "\n")
+ self.fout.flush()
+
+ def process_buffer(self):
+ while self.waiting_for in self.dec_buf:
+ self.process_dec(self.dec_buf[self.waiting_for])
+ del self.dec_buf[self.waiting_for]
+ self.waiting_for+=1
 
-def chunks(lst, n):
- """Yield successive n-sized chunks from lst."""
- for i in range(0, len(lst), n):
- yield lst[i:i + n]
+ def run(self):
+ while True:
+ ind, dec = self.msg_queue.get()
+ if dec == GENERATE_FINISHED:
+ break
+ elif ind != self.waiting_for:
+ self.dec_buf[ind] = dec
+ else:
+ self.process_dec(dec)
+ self.waiting_for+=1
+ self.process_buffer()
+ self.process_buffer()
+ assert not self.dec_buf, "IO Buffer not empty"
+ self.msg_queue.close()
+ self.msg_queue.join_thread()
 
+class PostProcess(Process):
+ """ Parallel detokenization """
+ def __init__(self, tokenizer, data_queue, msg_queue,
+ skip_special_tokens, clean_up_tokenization_spaces):
+ super(PostProcess, self).__init__()
+ self.data_queue = data_queue
+ self.msg_queue = msg_queue
+ self.tokenizer = tokenizer
+ self.clean_up_tokenization_spaces = clean_up_tokenization_spaces
+ self.skip_special_tokens = skip_special_tokens
+
+ def run(self):
+ while True:
+ ind, summaries = self.data_queue.get()
+ if summaries == GENERATE_FINISHED:
+ self.data_queue.put((-1, POSTPROCESS_FINISHED))
+ break
+ elif summaries == POSTPROCESS_FINISHED:
+ self.data_queue.put((-1, POSTPROCESS_FINISHED))
+ break
+ else:
+ dec = self.tokenizer.batch_decode(summaries,
+ skip_special_tokens = self.skip_special_tokens,
+ clean_up_tokenization_spaces =
+ self.clean_up_tokenization_spaces)
+ self.msg_queue.put((ind, dec))
+
+ self.data_queue.close()
+ self.data_queue.join_thread()
+ self.msg_queue.close()
+ self.msg_queue.join_thread()
 
 def generate_summaries_or_translations(
  examples: list,
@@ -29,6 +123,10 @@ def generate_summaries_or_translations(
  decoder_start_token_id=None,
  fastseq_opt=True,
  no_repeat_ngram_size=None,
+ skip_special_tokens=True,
+ clean_up_tokenization_spaces=False,
+ preprocess_cpu_num=2,
+ postprocess_cpu_num=2,
  **gen_kwargs,
 ) -> None:
  """Run generation"""
@@ -46,30 +144,44 @@ def generate_summaries_or_translations(
 
  # update config with summarization specific params
  use_task_specific_params(model, task)
+ data_queue = Queue()
+ msg_queue = Queue()
+ p_list = []
+
+ for i in range(postprocess_cpu_num):
+ p = PostProcess(tokenizer, data_queue, msg_queue,
+ skip_special_tokens, clean_up_tokenization_spaces)
+ p_list.append(p)
+ p.start()
 
- for batch in tqdm(list(chunks(examples, batch_size))):
- if "t5" in model_name:
- batch = [model.config.prefix + text for text in batch]
- batch = tokenizer(batch,
- return_tensors="pt",
- truncation=True,
- padding="max_length").to(device)
+ io_process = IOProcess( msg_queue, fout)
+ io_process.start()
+ dataset = TokenizeDataset(examples, tokenizer, model_name,
+ model.config.prefix)
+ training_generator = torch.utils.data.DataLoader(dataset,
+ batch_size=batch_size, num_workers = preprocess_cpu_num,
+ drop_last=True)
+ for ind, batch in tqdm(enumerate(training_generator)):
+ input_ids, attention_mask = batch
+ input_ids = input_ids.view(batch_size, -1).to(device)
+ attention_mask = attention_mask.view(batch_size, -1).to(device)
  input_ids, attention_mask = trim_batch(
-**batch, pad_token_id=tokenizer.pad_token_id)
+input_ids, tokenizer.pad_token_id, attention_mask)
  summaries = model.generate(
  input_ids=input_ids,
  attention_mask=attention_mask,
  decoder_start_token_id=decoder_start_token_id,
  no_repeat_ngram_size=no_repeat_ngram_size,
  **gen_kwargs,
  )
- dec = tokenizer.batch_decode(summaries,
- skip_special_tokens=True,
- clean_up_tokenization_spaces=False)
- for hypothesis in dec:
- fout.write(hypothesis + "\n")
- fout.flush()
-
+ summaries_cpu = summaries.cpu()
+ data_queue.put((ind, summaries_cpu))
+ data_queue.put((-1, GENERATE_FINISHED))
+ for p in p_list:
+ p.join()
+ msg_queue.put((-1, GENERATE_FINISHED))
+ io_process.join()
+ fout.close()
 
 def run_generate():
  """Entrance is here."""
@@ -118,6 +230,19 @@ def run_generate():
  parser.add_argument("--without_fastseq_opt", action="store_true")
  parser.add_argument("--no_repeat_ngram_size", type=int, default=None,
  required=False, help="size of no repeat ngram")
+ parser.add_argument("--include_special_tokens", action="store_true")
+ parser.add_argument("--clean_up_tokenization_spaces", action="store_true")
+ parser.add_argument("--preprocess_cpu_num",
+ type=int,
+ default=2,
+ required=False,
+ help="pre-processing worker threads")
+ parser.add_argument("--postprocess_cpu_num",
+ type=int,
+ default=2,
+ required=False,
+ help="post-processing worker threads")
+
  args = parser.parse_args()
  examples = [
  " " + x.rstrip() if "t5" in args.model_name else x.rstrip()
@@ -137,7 +262,11 @@ def run_generate():
  decoder_start_token_id=args.decoder_start_token_id,
  fastseq_opt=not args.without_fastseq_opt,
  no_repeat_ngram_size=args.no_repeat_ngram_size,
- )
+ skip_special_tokens=not args.include_special_tokens,
+ clean_up_tokenization_spaces=args.clean_up_tokenization_spaces,
+ preprocess_cpu_num=args.preprocess_cpu_num,
+ postprocess_cpu_num=args.postprocess_cpu_num,
+ )
  if args.reference_path is None:
  return
  # Compute scores