Repository: facebookresearch/fairseq

Issue #2506: TPU + MT Command
Created: 2020-08-20T22:00:06Z
@myleott I have been struggling with the GCP TPU tutorial and getting either OOM or "connection to mesh master" failed from
https://cloud.google.com/tpu/docs/tutorials/transformer-pytorch
In the case that you guys are not maintaining that tutorial, is there a working MT command with --tpu checked in?
I tried the command below on master
with
torch==1.5.0a0+ab660ae
torch-xla==1.5

fairseq-train \
  $HOME/pytorch-tutorial-data/wmt18_en_de_bpej32k \
  --save-interval=1 \
  --arch=transformer_vaswani_wmt_en_de_big \
  --max-target-positions=64  --max-tokens=64 \
  --attention-dropout=0.1 \
  --no-progress-bar \
  --criterion=label_smoothed_cross_entropy \
  --source-lang=en \
  --lr-scheduler=inverse_sqrt \
  --min-lr 1e-09 \
  --skip-invalid-size-inputs-valid-test \
  --target-lang=de \
  --label-smoothing=0.1 \
  --update-freq=1 \
  --optimizer adam \
  --adam-betas '(0.9, 0.98)' \
  --warmup-init-lr 1e-07 \
  --lr 0.0005 \
  --warmup-updates 4000 \
  --share-all-embeddings \
  --dropout 0.3 \
  --weight-decay 0.0 \
  --valid-subset=valid \
  --max-epoch=25 \
  --tpu --bf16 --distributed-world-size 8
and got the following traceback (once for each device)
Traceback (most recent call last):
  File "/anaconda3/envs/torch-xla-1.5/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py", lin
e 119, in _start_fn
    fn(gindex, *args)
  File "/home/shleifer/fairseq/fairseq/distributed_utils.py", line 156, in distributed_main
    main(args, **kwargs)
  File "/home/shleifer/fairseq/fairseq_cli/train.py", line 119, in main
    valid_losses, should_stop = train(args, trainer, task, epoch_itr)
  File "/anaconda3/envs/torch-xla-1.5/lib/python3.6/contextlib.py", line 52, in inner
    return func(*args, **kwds)
  File "/home/shleifer/fairseq/fairseq_cli/train.py", line 199, in train
    log_output = trainer.train_step(samples)
  File "/anaconda3/envs/torch-xla-1.5/lib/python3.6/contextlib.py", line 52, in inner
    return func(*args, **kwds)
  File "/home/shleifer/fairseq/fairseq/trainer.py", line 491, in train_step
    logging_outputs, sample_size, ooms, train_time, ignore=is_dummy_batch,
  File "/home/shleifer/fairseq/fairseq/trainer.py", line 820, in _aggregate_logging_outputs
    logging_outputs, *extra_stats_to_sum, ignore=ignore
  File "/home/shleifer/fairseq/fairseq/trainer.py", line 883, in _fast_stat_sync_sum
    group=self.data_parallel_process_group
  File "/home/shleifer/fairseq/fairseq/distributed_utils.py", line 312, in all_reduce_dict
    cpu_data = _all_reduce_dict(cpu_data)
  File "/home/shleifer/fairseq/fairseq/distributed_utils.py", line 309, in _all_reduce_dict
    all_reduce(buf, group=group)
  File "/home/shleifer/fairseq/fairseq/distributed_utils.py", line 207, in all_reduce
    return xm.all_reduce('sum', [tensor], groups=group[1])
  File "/anaconda3/envs/torch-xla-1.5/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 360, in all_red
uce
    reduce_type, inputs, _get_all_reduce_token(), scale, groups)
TypeError: _xla_all_reduce(): incompatible function arguments. The following argument types are supported:
    1. (arg0: str, arg1: List[at::Tensor], arg2: _XLAC.IrValue, arg3: float, arg4: list) -> _XLAC.IrValue

Invoked with: 'sum', [tensor([51.0000,  0.0000, 40.2500, 51.0000,  3.0000, 51.0000],device='xla:0',
       dtype=torch.float64)], <_XLAC.IrValue object at 0x7ff4085356c0>, 1.0, None

Duplicate issue #3317: How to efficiently run translation task on TPUs?
Created: 2021-03-07T07:33:02Z
❓ Questions and Help
Before asking:

search the issues.
search the docs.


What is your question?
I'm running training on a custom translation dataset with 14M examples with vocab_size of 32K on a tpu-v2-8.
The current training time seems to be very slow(estimates of around 260 hrs per epoch) and has lot of WARNING | fairseq.trainer | XLA compilation detected on device #0; too many of these can lead to slow training, but we expect a few in the beginning.
To improve the training times, I used num-batch-buckets=8  :
# to use all 8 tpu cores
export MKL_SERVICE_FORCE_INTEL=1

fairseq-train {exp_dir}/final_bin \
        --arch transformer \
        --optimizer adam --adam-betas '(0.9, 0.98)' --clip-norm 1.0  \
        --lr 0.0005 --lr-scheduler inverse_sqrt --warmup-updates 4000 --warmup-init-lr 1e-07  \
        --dropout 0.2  \
        --criterion label_smoothed_cross_entropy --label-smoothing 0.1  \
        --max-tokens 8192  \
        --max-update 1000000  \
        --max-source-positions 200  \
        --max-target-positions 200  \
        --tensorboard-logdir {expdir}/tensorboard  \
        --save-dir {expdir}/model  \
        --save-interval 1  \
        --keep-last-epochs 5  \
        --patience 5  \
        --source-lang SRC  \
        --target-lang TGT  \
        --update-freq 1  \
        --tpu  \
        --distributed-world-size 8  \
        --num-batch-buckets 8  \
        --wandb-project <name of wandb project>
^ after a slow startup, the estimated training times for one epoch seems to be around 2 hrs 5 min.
But I noticed that I'm losing crucial 2M datapoints after bucketing and removing sentences greater than 200 (max_source_pos and max_target_pos). Without num_buckets=8, I only lose 20K samples

To avoid skipping a lot of crucial training data, I remove the 20K training examples where the len(source_tokens) or len(target_tokens)  exceed 200 before running fairseq-train. This results in no loss of data after running the same training command as above but the model training is more slower at 4 hr 40 min per epoch.
2021-03-07 05:35:25 | INFO | fairseq_cli.train | task: TranslationTask                                                                                                                                                                  
2021-03-07 05:35:25 | INFO | fairseq_cli.train | model: TransformerModel                                                                                                                                                                
2021-03-07 05:35:25 | INFO | fairseq_cli.train | criterion: LabelSmoothedCrossEntropyCriterion                                                                                                                                          
2021-03-07 05:35:25 | INFO | fairseq_cli.train | num. model params: 93,073,408 (num. trained: 93,073,408)                                                                                                                               
2021-03-07 05:35:25 | INFO | fairseq_cli.train | training on 8 devices (GPUs/TPUs)                                                                                                                                                      
2021-03-07 05:35:25 | INFO | fairseq_cli.train | max tokens per GPU = 8192 and batch size per GPU = None                                                                                                                                
2021-03-07 05:35:25 | INFO | fairseq.trainer | Preparing to load checkpoint ../experiments/all-v1_WAT/model-testing/checkpoint_last.pt                                                                                                  
2021-03-07 05:35:25 | INFO | fairseq.trainer | No existing checkpoint found ../experiments/all-v1_WAT/model-testing/checkpoint_last.pt                                                                                                  
2021-03-07 05:35:25 | INFO | fairseq.trainer | loading train data for epoch 1                                                                                                                                                           
2021-03-07 05:35:35 | INFO | fairseq.data.data_utils | loaded 14,470,625 examples from: ../experiments/all-v1_WAT/final_bin/train.SRC-TGT.SRC                                                                                           
2021-03-07 05:35:44 | INFO | fairseq.data.data_utils | loaded 14,470,625 examples from: ../experiments/all-v1_WAT/final_bin/train.SRC-TGT.TGT                                                                                           
2021-03-07 05:35:44 | INFO | fairseq.tasks.translation | ../experiments/all-v1_WAT/final_bin train SRC-TGT 14470625 examples                                                                                                            
2021-03-07 05:35:45 | INFO | fairseq.data.language_pair_dataset | bucketing source lengths: [11, 16, 21, 25, 31, 39, 52, 201]                                                                                                           
2021-03-07 05:35:46 | INFO | fairseq.data.language_pair_dataset | bucketing target lengths: [8, 13, 17, 22, 27, 34, 45, 201]                                                                                                            
epoch 001:   0%|                                                                                                                                                                                              | 0/12379 [00:00<?, ?it/s]
2021-03-07 05:36:25 | INFO | fairseq.trainer | begin training epoch 1                                                                                                                                                                   
2021-03-07 05:36:26 | INFO | fairseq_cli.train | Start iterating over samples 

<lots of XLA compilation warnings>
epoch 001:   2%| | 300/12379 [28:48<5:15:44,  1.57s/it, loss=10.73, nll_loss=10.105, ppl=1101.2, wps=239.7, ups=0.01, wpb=30865, bsz=1584, num_updates=300, lr=3.75925e-05, gnorm=1.25, clip=100, train_wall=49, gb_free=0.1, gb_total=8
2021-03-07 06:05:14 | WARNING | fairseq.trainer | XLA compilation detected on device #5; too many of these can lead to slow training, but we expect a few in the beginning
2021-03-07 06:05:14 | WARNING | fairseq.trainer | XLA compilation detected on device #6; too many of these can lead to slow training, but we expect a few in the beginning
2021-03-07 06:05:14 | WARNING | fairseq.trainer | XLA compilation detected on device #0; too many of these can lead to slow training, but we expect a few in the beginning
epoch 001:   2%| | 301/12379 [28:49<4:08:42,  1.24s/it, loss=10.73, nll_loss=10.105, ppl=1101.2, wps=239.7, ups=0.01, wpb=30865, bsz=1584, num_updates=300, lr=3.75925e-05, gnorm=1.25, clip=100, train_wall=49, gb_free=0.1, gb_total=8
2021-03-07 06:05:14 | WARNING | fairseq.trainer | XLA compilation detected on device #2; too many of these can lead to slow training, but we expect a few in the beginning
2021-03-07 06:05:14 | WARNING | fairseq.trainer | XLA compilation detected on device #4; too many of these can lead to slow training, but we expect a few in the beginning
2021-03-07 06:05:14 | WARNING | fairseq.trainer | XLA compilation detected on device #1; too many of these can lead to slow training, but we expect a few in the beginning
2021-03-07 06:05:14 | WARNING | fairseq.trainer | XLA compilation detected on device #7; too many of these can lead to slow training, but we expect a few in the beginning
2021-03-07 06:05:14 | WARNING | fairseq.trainer | XLA compilation detected on device #3; too many of these can lead to slow training, but we expect a few in the beginning
epoch 001:   2%| | 307/12379 [28:57<4:13:49,  1.26s/it, loss=10.73, nll_loss=10.105, ppl=1101.2, wps=239.7, ups=0.01, wpb=30865, bsz=1584, num_updates=300, lr=3.75925e-05, gnorm=1.25, clip=100, train_wall=49, gb_free=0.1, gb_total=8
2021-03-07 06:05:22 | WARNING | fairseq.trainer | XLA compilation detected on device #3; too many of these can lead to slow training, but we expect a few in the beginning
epoch 001:   3%| | 427/12379 [31:31<4:22:29,  1.32s/it, loss=10.602, nll_loss=9.916, ppl=965.77, wps=142.3, ups=0.01, wpb=18262, bsz=320, num_updates=400, lr=5.009e-05, gnorm=2.404, clip=100, train_wall=46, gb_free=0.1, gb_total=8,
2021-03-07 06:07:56 | WARNING | fairseq.trainer | XLA compilation detected on device #1; too many of these can lead to slow training, but we expect a few in the beginning
epoch 001:   5%| | 610/12379 [35:24<3:46:48,  1.16s/it, loss=9.964, nll_loss=9.172, ppl=576.71, wps=236.6, ups=0.01, wpb=30216, bsz=880, num_updates=600, lr=7.5085e-05, gnorm=1.446, clip=100, train_wall=49, gb_free=0.4, gb_total=8,
2021-03-07 06:11:50 | WARNING | fairseq.trainer | XLA compilation detected on device #6; too many of these can lead to slow training, but we expect a few in the beginning
epoch 001:   5%| | 644/12379 [36:08<4:07:56,  1.27s/it, loss=9.964, nll_loss=9.172, ppl=576.71, wps=236.6, ups=0.01, wpb=30216, bsz=880, num_updates=600, lr=7.5085e-05, gnorm=1.446, clip=100, train_wall=49, gb_free=0.4, gb_total=8,
2021-03-07 06:12:34 | WARNING | fairseq.trainer | XLA compilation detected on device #2; too many of these can lead to slow training, but we expect a few in the beginning
epoch 001:  28%|▎| 3410/12379 [1:35:12<3:14:40,  1.30s/it, loss=6.855, nll_loss=5.597, ppl=48.4, wps=208, ups=0.01, wpb=26565, bsz=1016, num_updates=3400, lr=0.000425015, gnorm=0.609, clip=0, train_wall=47, gb_free=0.1, gb_total=8,
2021-03-07 07:12:45 | WARNING | fairseq.trainer | XLA compilation detected on device #0; too many of these can lead to slow training, but we expect a few in the beginning
epoch 001:  28%|▎| 3505/12379 [1:38:50<3:02:15,  1.23s/it, loss=6.816, nll_loss=5.537, ppl=46.43, wps=116.1, ups=0, wpb=26016, bsz=688, num_updates=3500, lr=0.000437513, gnorm=0.71, clip=0, train_wall=116, gb_free=0.1, gb_total=8, $
epoch 001:  29%|▎| 3555/12379 [1:39:55<3:12:04,  1.31s/it, loss=6.816, nll_loss=5.537, ppl=46.43, wps=116.1, ups=0, wpb=26016, bsz=688, num_updates=3500, lr=0.000437513, gnorm=0.71, clip=0, train_wall=116, gb_free=0.1, gb_total=8, $                                                                                                                                                        
^ these are the training statistics at the moment. I also tried --bf16 without export XLA_USE_BF16=1 but that seems to be slower than the current training time.
Questions:

Is there some documentation or fairseq examples on how to run translation on tpu efficiently with the fairseq master branch?
Currently, bucketing(num-batch-buckets) on src and tgt lengths runs before restricting the sentences to max_source_pos and max_target_pos. Shouldn't this be the other way around?
How do we decide the ideal num-batch-buckets ? (I've tried 3,5,8, 16 and found 8 to have the fastest training time estimate so far)
Do you have any advice to speed up the current training times or things I could try to see if training time improves?

Please let me know if you need any more information.
Thanks
Code
   
What have you tried?
What's your environment?

fairseq Version (e.g., 1.0 or master): master (commit da9eaba)
PyTorch Version (e.g., 1.0) 1.7 (torch-xla-1.7 from here )
OS (e.g., Linux): Linux
How you installed fairseq (pip, source): source  (commit da9eaba)
Build command you used (if compiling from source): pip install -e .
Python version: 3.6.10
CUDA/cuDNN version: NA
GPU models and configuration: TPU v2-8
Any other relevant information:

Closing comment (2022-01-17T13:47:22Z): might be related to #2506