TPU Training fails with --evaluate_during_training

🐛 Bug

TPU Trainer does not seem to support --evaluate_during_training. When the training loop goes into logging part, the whole process just hangs up stalling training. The same code/dataset with a multi-gpu setup works well.

I am trying to move my company to Huggingface so want to train models on TPUs on our dataset which hung during the logging step. I was able to replicate the behavior with run_langugage_modelling.py, and the steps to replicate this are shown below.

Other observations are - I felt that multiprocessing way of doing TPU training wastes a lot of CPU memory because with large datasets one has to use a machine with 100s of GBs of RAM because the features are being replicated 8 times in memory. Another bug is that with TPU training there are 8 WandB runs generated and it creates a lot of clutter. Suggestions to fix this would be to only do wandb logging from a single process. If its unavoidable to generate 8 wandb runs, tag all the runs to belong to a single ‘group’ that leads to better organization of the runs. (https://docs.wandb.com/library/advanced/grouping)

Information

Model I am using (Bert, XLNet …): Roberta with run_language_modelling.py to replicate, T5 with our internal data.

Language I am using the model on (English, Chinese …): English

The problem arises when using:

• the official example scripts: (give details below)
• my own modified scripts: (give details below)

The tasks I am working on is:

• an official GLUE/SQUaD task: (give the name)
• my own task or dataset: (give details below)

To reproduce

Steps to reproduce the behavior:

1. Create a new n1-highmem-32 machine with debian-9-torch-xla OS image in us-central1-c zone
2. conda activate torch-xla-nightly and start a v2-8 TPU in us-central1-c zone. Set the TPU env vars
3. Use the master branch of transformers
4. Download Wikitext 103 raw char level data from https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-103-raw-v1.zip (according to examples for run_language_modelling). Extract it
5. Run the example script

export TRAIN_FILE=/path/to/dataset/wiki.train.raw
export TEST_FILE=/path/to/dataset/wiki.test.raw

python xla_spawn.py --num_cores 8 language_modelling/run_language_modeling.py \
    --output_dir=output \
    --model_type=roberta \
    --model_name_or_path=roberta-base \
    --do_train \
    --train_data_file=$TRAIN_FILE \
    --do_eval \
    --eval_data_file=$TEST_FILE \
    --mlm
    --evaluate_during_training
    --per_device_train_batch_size=4
    --per_device_eval_batch_size=4

When it hangs, the tqdm counter is stuck at step 499 (with 500 as the logging interval) and nothing happens. When I do a Keyboard Interrupt, I get this stack trace.

main()
  File "../../../vendor/transformers/examples/xla_spawn.py", line 68, in main
    xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores)
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 296, in spawn
    start_method=start_method)
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 158, in start_processes
    while not context.join():
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 78, in join
    timeout=timeout,
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/multiprocessing/connection.py", line 911, in wait
    ready = selector.select(timeout)
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/selectors.py", line 376, in select
    fd_event_list = self._poll.poll(timeout)
KeyboardInterrupt

Expected behavior

Being able to log validation set loss during training

Environment info

transformers version: 2.11.0

• Platform: Linux-4.9.0-12-amd64-x86_64-with-debian-9.12
• Python version: 3.6.10
• PyTorch version (GPU?): 1.6.0a0+03eca38 (False)
• Tensorflow version (GPU?): not installed (NA)
• Using GPU in script?: no
• Using distributed or parallel set-up in script?: Yes, 8 core parallelism with xla_spawn.py

[vanpelt]

I’m one of the founders of wandb. We’re digging into the root cause of this now. We’re planning to issue a new release ASAP to ensure users can never get into this hung state. I’ll update the thread here. For anyone finding this thread online and hitting the issue, you can add the following code to disable the gradient monitoring in wandb with huggingface.

import os
os.environ["WANDB_WATCH"] = "false"

Or if you’re shelling out to a python script:

export WANDB_WATCH=false
python your_script.py

[misrasaurabh1]

I am not sure if wandb supports logging of gradients with Pytorch/XLA. I reached out to Wandb to ask about this, should get a reply by tomorrow. It is possible that Pytorch/XLA does not support gradient logging as well. I looked at the XLA github repo and couldn’t find a mention of gradients logging with TPUs. I am unfamiliar with XLA interface with wandb and not keen on digging deeper into this. Hopefully wandb offers more clarity soon.