Enable AMP for xla:gpu device in trainer class

[ymwangg]

What does this PR do?

This PR enables AMP in trainer class for xla:gpu device.

Discussion

It looks like the torch_xla support in trainer class is primarily for xla:tpu device.
I found the following features may be useful but not essential and I can include them in this PR if necessary:

1. Rename tpu to xla in the codebase.
2. Currently xla device is always turned on when torch_xla is installed. It may be useful to allow users to optionally turn it off without uninstalling torch_xla.
3. Currently users need to set GPU_NUM_DEVICES manually when using xla:gpu device. It may be useful to set a default value for it when torch_xla and cuda devices are available.

[LysandreJik]

Oh, interesting! Thanks for your contribution, pinging @sgugger on the issue.

[sgugger]

I'm not entirely sure about this PR in the sense that PyTorch XLA support is mainly for TPU, and I don't know if traditional mixed precision training with the gradient scaler will work on TPUs.

So we should probably split the test to detect if we have GPUs available or TPUs. Some of the logic will stay common between the two, but the mixed precision part might only work for XLA GPUs?

[sgugger]

This part is done by the dataloader (which is wrapped in a ParallelLoader), so it shouldn't be here.

[ymwangg]

This is intentional since loss will be materialized in self._nested_gather(loss.repeat(batch_size)) and adding a mark_step here can significantly improve the speed. For example, the evaluation time of bert-base-uncased using run_mlm.py will be reduced from 32.53s to 18.73s by adding this mark_step.

[sgugger]

Will this be false on TPU? The test is there for that purpose since mixed precision training does not work on TPU.

[ymwangg]

You are right that this check won't filter out TPU. We can change it to not (self.device.type == "xla" and "GPU_NUM_DEVICES" in os.environ).

[ymwangg]

I'm not entirely sure about this PR in the sense that PyTorch XLA support is mainly for TPU, and I don't know if traditional mixed precision training with the gradient scaler will work on TPUs.

So we should probably split the test to detect if we have GPUs available or TPUs. Some of the logic will stay common between the two, but the mixed precision part might only work for XLA GPUs?

Right, XLA:TPU does not support AMP and only XLA:GPU support it.

[sgugger]

Right, XLA:TPU does not support AMP and only XLA:GPU support it.

So as I said in my previous comment, could you add a new test is_gpu_xla_available and use this one for the part where you add grad scalers? Otherwise the changes will make the Trainer stop working on TPU.

[ymwangg]

@sgugger Maybe I'm missing something. Could you elaborate why the changes will make the Trainer stop working on TPU? The code inside self.do_grad_scaling is unreachable if running on TPU since either --fp16 or --bf16 option will raise error on TPU. I tested the following training script on TPU:

python run_mlm.py \
    --model_name_or_path bert-base-uncased \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --overwrite_output_dir true \
    --output_dir /tmp/test-mlm \
    --per_device_train_batch_size 10 \
    --do_eval \
    --do_train

With master branch:

WARNING:root:TPU has started up successfully with version pytorch-1.9
WARNING:__main__:Process rank: -1, device: xla:1, n_gpu: 0distributed training: False, 16-bits training: False
...
***** train metrics *****
  epoch                    =        3.0
  train_loss               =     1.7568
  train_runtime            = 0:12:23.47
  train_samples            =       4627
  train_samples_per_second =      18.67
  train_steps_per_second   =      1.868

With this PR:

WARNING:root:TPU has started up successfully with version pytorch-1.9
WARNING:__main__:Process rank: -1, device: xla:1, n_gpu: 0distributed training: False, 16-bits training: False
...
***** train metrics *****
  epoch                    =        3.0
  train_loss               =     1.7577
  train_runtime            = 0:10:19.70
  train_samples            =       4627
  train_samples_per_second =     22.399
  train_steps_per_second   =      2.241

[sgugger]

Ah yes, you're right. Thanks for testing!