[FEATURE]Horovod support for training transformer (PART 2)

scripts/datasets/machine_translation/wmt2014_ende.sh

@@ -12,8 +12,8 @@ nlp_data prepare_wmt \
 # We use sacrebleu to fetch the dev set (newstest2013) and test set (newstest2014)
 sacrebleu -t wmt13 -l ${SRC}-${TGT} --echo src > ${SAVE_PATH}/dev.raw.${SRC}
 sacrebleu -t wmt13 -l ${SRC}-${TGT} --echo ref > ${SAVE_PATH}/dev.raw.${TGT}
-sacrebleu -t wmt14 -l ${SRC}-${TGT} --echo src > ${SAVE_PATH}/test.raw.${SRC}
-sacrebleu -t wmt14 -l ${SRC}-${TGT} --echo ref > ${SAVE_PATH}/test.raw.${TGT}
+sacrebleu -t wmt14/full -l ${SRC}-${TGT} --echo src > ${SAVE_PATH}/test.raw.${SRC}
+sacrebleu -t wmt14/full -l ${SRC}-${TGT} --echo ref > ${SAVE_PATH}/test.raw.${TGT}
 
 
 # Clean and tokenize the training + dev corpus
@@ -34,6 +34,7 @@ nlp_preprocess clean_tok_para_corpus --src-lang ${SRC} \
                       --tgt-corpus dev.raw.${TGT} \
                       --min-num-words 1 \
                       --max-num-words 100 \
+                      --max-ratio 1.5 \
                       --src-save-path dev.tok.${SRC} \
                       --tgt-save-path dev.tok.${TGT}
 

scripts/machine_translation/README.md

@@ -30,9 +30,36 @@ python3 train_transformer.py \
     --save_dir transformer_base_wmt2014_en_de_${SUBWORD_ALGO} \
     --cfg transformer_base \
     --lr 0.002 \
-    --batch_size 2700 \
-    --num_averages 5 \
-    --warmup_steps 4000 \
+    --sampler BoundedBudgetSampler \
+    --max_num_tokens 2700 \
+    --max_update 15000 \
+    --save_interval_update 500 \
+    --warmup_steps 6000 \
+    --warmup_init_lr 0.0 \
+    --seed 123 \
+    --gpus 0,1,2,3
+```
+
+Or training via horovod
+```
+horovodrun -np 4 -H localhost:4 python3 train_transformer.py \
+    --comm_backend horovod \
+    --train_src_corpus ${datapath}/wmt2014_ende/train.tok.${SUBWORD_ALGO}.${SRC} \
+    --train_tgt_corpus ${datapath}/wmt2014_ende/train.tok.${SUBWORD_ALGO}.${TGT} \
+    --dev_src_corpus ${datapath}/wmt2014_ende/dev.tok.${SUBWORD_ALGO}.${SRC} \
+    --dev_tgt_corpus ${datapath}/wmt2014_ende/dev.tok.${SUBWORD_ALGO}.${TGT} \
+    --src_subword_model_path ${datapath}/wmt2014_ende/${SUBWORD_ALGO}.model \
+    --src_vocab_path ${datapath}/wmt2014_ende/${SUBWORD_ALGO}.vocab \
+    --tgt_subword_model_path ${datapath}/wmt2014_ende/${SUBWORD_ALGO}.model \
+    --tgt_vocab_path ${datapath}/wmt2014_ende/${SUBWORD_ALGO}.vocab \
+    --save_dir transformer_base_wmt2014_en_de_${SUBWORD_ALGO} \
+    --cfg transformer_base \
+    --lr 0.002 \
+    --sampler BoundedBudgetSampler \
+    --max_num_tokens 2700 \
+    --max_update 15000 \
+    --save_interval_update 500 \
+    --warmup_steps 6000 \
     --warmup_init_lr 0.0 \
     --seed 123 \
     --gpus 0,1,2,3
@@ -42,18 +69,16 @@ Use the average_checkpoint cli to average the last 10 checkpoints
 
 ```bash
 gluon_average_checkpoint --checkpoints transformer_base_wmt2014_en_de_${SUBWORD_ALGO}/epoch*.params \
-    --begin 21 \
-    --end 30 \
-    --save-path transformer_base_wmt2014_en_de_${SUBWORD_ALGO}/avg_21_30.params
+    --begin 30 \
+    --end 39 \
+    --save-path transformer_base_wmt2014_en_de_${SUBWORD_ALGO}/epoch_avg_30_39.params
 ```
 
-
 Use the following command to inference/evaluate the Transformer model:
 
 ```bash
-SUBWORD_MODEL=yttm
 python3 evaluate_transformer.py \
-    --param_path transformer_base_wmt2014_en_de_${SUBWORD_MODEL}/average_21_30.params \
+    --param_path transformer_base_wmt2014_en_de_${SUBWORD_MODEL}/epoch_avg_30_39.params \
     --src_lang en \
     --tgt_lang de \
     --cfg transformer_base_wmt2014_en_de_${SUBWORD_MODEL}/config.yml \
@@ -110,7 +135,6 @@ gluon_average_checkpoint --checkpoints transformer_big_wmt2014_en_de_${SUBWORD_A
 Use the following command to inference/evaluate the Transformer model:
 
 ```bash
-SUBWORD_MODEL=yttm
 python3 evaluate_transformer.py \
     --param_path transformer_big_wmt2014_en_de_${SUBWORD_MODEL}/average_21_30.params \
     --src_lang en \
@@ -131,16 +155,18 @@ Test BLEU score with 3 seeds (evaluated via sacre BLEU):
 
 - transformer_base
 
+(test bleu / valid bleu)
 | Subword Model | #Params    | Seed = 123  | Seed = 1234 | Seed = 12345 |  Mean±std   |
 |---------------|------------|-------------|-------------|--------------|-------------|
-| yttm          |            |  -          | -           |  -           |  -          |
+| yttm          |            | 26.50/26.29 | -           |  -           |  -          |
 | hf_bpe        |            |  -          | -           |  -           |  -          |
 | spm           |            |  -          | -           |  -           |  -          |
 
 - transformer_wmt_en_de_big
 
+(test bleu / valid bleu)
 | Subword Model | #Params    | Seed = 123  | Seed = 1234 | Seed = 12345 |  Mean±std   |
 |---------------|------------|-------------|-------------|--------------|-------------|
-| yttm          |            |  27.99      | -           |  -           |  -          |
+| yttm          |            | 27.93/26.82 | -           |  -           |  -          |
 | hf_bpe        |            |  -          | -           |  -           |  -          |
 | spm           |            |  -          | -           |  -           |  -          |

scripts/machine_translation/evaluate_transformer.py

@@ -247,10 +247,17 @@ def evaluate(args):
             of.write('\n'.join(pred_sentences))
             of.write('\n')
 
-        sacrebleu_out = sacrebleu.corpus_bleu(sys_stream=pred_sentences, ref_streams=[all_tgt_lines])
-        logging.info('Time Spent: {}, #Sent={}, SacreBlEU={} Avg NLL={}, Perplexity={}'
+        sacrebleu_out = sacrebleu.corpus_bleu(sys_stream=pred_sentences, ref_streams=[all_tgt_lines])        
+        logging.info('Time Spent: {}, #Sent={}, SacreBlEU={} '
+                     '({:2.1f} {:2.1f} {:2.1f} {:2.1f}) '
+                     '(BP={:.3f}, ratio={:.3f}, syslen={}, reflen={}), '
+                     'Avg NLL={}, Perplexity={}'
                      .format(end_eval_time - start_eval_time, len(all_tgt_lines),
-                             sacrebleu_out.score, avg_nll_loss, np.exp(avg_nll_loss)))
+                             sacrebleu_out.score,
+                             *sacrebleu_out.precisions,
+                             sacrebleu_out.bp, sacrebleu_out.sys_len / sacrebleu_out.ref_len,
+                             sacrebleu_out.sys_len, sacrebleu_out.ref_len,
+                             avg_nll_loss, np.exp(avg_nll_loss)))
     # inference only
     else:
         with open(os.path.join(args.save_dir, 'pred_sentences.txt'), 'w', encoding='utf-8') as of:

scripts/machine_translation/train_transformer.py

@@ -50,7 +50,8 @@
     LinearWidthBucket,
     ExpWidthBucket,
     FixedBucketSampler,
-    BoundedBudgetSampler
+    BoundedBudgetSampler,
+    ShardedIterator
 )
 import gluonnlp.data.batchify as bf
 from gluonnlp.data import Vocab
@@ -179,6 +180,7 @@ def parse_args():
     logging.info(args)
     return args
 
+
 def validation(model, data_loader, ctx_l):
     """Validate the model on the dataset
 
@@ -231,14 +233,16 @@ def load_dataset_with_cache(src_corpus_path: str,
                             tgt_corpus_path: str,
                             src_tokenizer: BaseTokenizerWithVocab,
                             tgt_tokenizer: BaseTokenizerWithVocab,
-                            overwrite_cache: bool):
+                            overwrite_cache: bool,
+                            local_rank: int):
     # TODO online h5py multi processing encode (Tao)
     src_md5sum = md5sum(src_corpus_path)
     tgt_md5sum = md5sum(tgt_corpus_path)
     cache_filepath = os.path.join(CACHE_PATH,
                                   '{}_{}.cache.npz'.format(src_md5sum[:6], tgt_md5sum[:6]))
     if os.path.exists(cache_filepath) and not overwrite_cache:
-        logging.info('Load cache from {}'.format(cache_filepath))
+        if local_rank == 0:
+            logging.info('Load cache from {}'.format(cache_filepath))
         npz_data = np.load(cache_filepath, allow_pickle=True)
         src_data, tgt_data = npz_data['src_data'][:], npz_data['tgt_data'][:]
     else:
@@ -288,7 +292,7 @@ def create_tokenizer(tokenizer_type, model_path, vocab_path):
 
 
 def train(args):
-    store, num_parts, rank, local_rank, is_master_node, ctx_l = init_comm(
+    _, num_parts, rank, local_rank, _, ctx_l = init_comm(
         args.comm_backend, args.gpus)
     src_tokenizer = create_tokenizer(args.src_tokenizer,
                                      args.src_subword_model_path,
@@ -302,12 +306,14 @@ def train(args):
                                                              args.train_tgt_corpus,
                                                              src_tokenizer,
                                                              tgt_tokenizer,
-                                                             args.overwrite_cache)
+                                                             args.overwrite_cache,
+                                                             local_rank)
     dev_src_data, dev_tgt_data = load_dataset_with_cache(args.dev_src_corpus,
                                                          args.dev_tgt_corpus,
                                                          src_tokenizer,
                                                          tgt_tokenizer,
-                                                         args.overwrite_cache)
+                                                         args.overwrite_cache,
+                                                         local_rank)
     data_train = gluon.data.SimpleDataset(
         [(src_tokens, tgt_tokens, len(src_tokens), len(tgt_tokens), i)
          for i, (src_tokens, tgt_tokens) in enumerate(zip(train_src_data, train_tgt_data))])
@@ -363,9 +369,9 @@ def train(args):
         train_batch_sampler = BoundedBudgetSampler(lengths=[(ele[2], ele[3]) for ele in data_train],
                                                      max_num_tokens=args.max_num_tokens,
                                                      max_num_sentences=args.max_num_sentences,
-                                                     seed=args.seed,
-                                                     num_parts=num_parts,
-                                                     part_index=rank)
+                                                     seed=args.seed)
+        if num_parts > 1:
+            train_batch_sampler = ShardedIterator(train_batch_sampler, num_parts=num_parts, part_index=rank)
     elif args.sampler == 'FixedBucketSampler':
         if args.comm_backend == 'horovod':
             raise NotImplementedError('FixedBucketSampler does not support horovod at present')
@@ -390,8 +396,7 @@ def train(args):
     else:
         raise NotImplementedError
 
-    if local_rank == 0:
-        logging.info(train_batch_sampler)
+    logging.info(train_batch_sampler)
 
     batchify_fn = bf.Tuple(bf.Pad(), bf.Pad(), bf.Stack(), bf.Stack(), bf.Stack())
     train_data_loader = gluon.data.DataLoader(data_train,
@@ -483,27 +488,31 @@ def train(args):
                     log_avg_loss = (log_avg_loss / log_loss_denom).asnumpy()
                     logging.info('[Epoch {} Batch {}/{}] loss={:.4f}, ppl={:.4f}, '
                                  'throughput={:.2f}K wps, wc={:.2f}K, LR={}'
-                                 .format(epoch_id, processed_batch_num * num_parts, len(train_data_loader),
-                                         log_avg_loss, np.exp(log_avg_loss),
+                                 .format(epoch_id, processed_batch_num * num_parts,
+                                         len(train_data_loader), log_avg_loss, np.exp(log_avg_loss),
                                          wps / 1000, log_wc / 1000, trainer.learning_rate))
                     log_start_time = time.time()
                     log_avg_loss = 0
                     log_loss_denom = 0
                     log_wc = 0
                 if local_rank == 0 and \
                    (args.max_update > 0 and n_train_iters % args.save_interval_update == 0):
+                    n_update = n_train_iters // args.save_interval_update
                     model.save_parameters(os.path.join(args.save_dir,
-                                                       'update{:d}.params'.format(n_train_iters // args.save_interval_update)),
+                                                       'update{:d}.params'.format(n_update)),
                                           deduplicate=True)
+                    avg_valid_loss = validation(model, val_data_loader, ctx_l)
+                    logging.info('[Update {}] validation loss/ppl={:.4f}/{:.4f}'
+                                 .format(n_update, avg_valid_loss, np.exp(avg_valid_loss)))
                 if args.max_update > 0 and n_train_iters >= args.max_update:
                     break
-        if local_rank == 0 and args.epochs > 0:
+        if local_rank == 0:
             model.save_parameters(os.path.join(args.save_dir,
                                                'epoch{:d}.params'.format(epoch_id)),
                                   deduplicate=True)
-        avg_valid_loss = validation(model, val_data_loader, ctx_l)
-        logging.info('[Epoch {}] validation loss/ppl={:.4f}/{:.4f}'
-                     .format(epoch_id, avg_valid_loss, np.exp(avg_valid_loss)))
+            avg_valid_loss = validation(model, val_data_loader, ctx_l)
+            logging.info('[Epoch {}] validation loss/ppl={:.4f}/{:.4f}'
+                         .format(epoch_id, avg_valid_loss, np.exp(avg_valid_loss)))
 
         if args.max_update > 0 and n_train_iters >= args.max_update:
             break

scripts/question_answering/README.md

@@ -74,7 +74,7 @@ We could speed up multi-GPU training via horovod.
 Compared to KVStore, training RoBERTa Large model on SQuAD 2.0 with 3 epochs will save roughly 1/4 training resources (8.48 vs 11.32 hours). Results may vary depending on the training instances.
 
 ```bash
-mpirun -np 4 -H localhost:4 python3 run_squad.py \
+horovodrun -np 4 -H localhost:4 python3 run_squad.py \
     --comm_backend horovod \
     ...
 ```