Enable TPU support (#868)

* added tpu docs * added tpu flags * add tpu docs + init training call * amp * amp * amp * amp * optimizer step * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * fix test pkg create (#873) * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added test return and print * added test return and print * added test return and print * added test return and print * added test return and print * Update pytorch_lightning/trainer/trainer.py Co-Authored-By: Luis Capelo <luiscape@gmail.com> * Fix segmentation example (#876) * removed torchvision model and added custom model * minor fix * Fixed relative imports issue * Fix/typo (#880) * Update greetings.yml * Update greetings.yml * Changelog (#869) * Create CHANGELOG.md * Update CHANGELOG.md * Update CHANGELOG.md * Update PULL_REQUEST_TEMPLATE.md * Update PULL_REQUEST_TEMPLATE.md * Add PR links to Version 0.6.0 in CHANGELOG.md * Add PR links for Unreleased in CHANGELOG.md * Update PULL_REQUEST_TEMPLATE.md * Fixing Function Signatures (#871) * added tpu docs * added tpu flags * add tpu docs + init training call * amp * amp * amp * amp * optimizer step * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added auto data transfer to TPU * added test return and print * added test return and print * added test return and print * added test return and print * added test return and print * added test return and print * added test return and print * added test return and print Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com> Co-authored-by: Luis Capelo <luiscape@gmail.com> Co-authored-by: Akshay Kulkarni <akshayk.vnit@gmail.com> Co-authored-by: Ethan Harris <ewah1g13@soton.ac.uk> Co-authored-by: Shikhar Chauhan <xssChauhan@users.noreply.github.com>
2026-09-09 11:32:07 +08:00 · 2020-02-17 16:01:20 -05:00
parent e38b18e9eb
commit d4a31f02e0
14 changed files with 489 additions and 48 deletions
@@ -1,13 +1,18 @@
 16-bit training
 =================
+Lightning offers 16-bit training for CPUs, GPUs and TPUs.
+
+GPU 16-bit
+-----------
 Lightning uses NVIDIA apex to handle 16-bit precision training.

 To use 16-bit precision, do two things:
+
 1. Install Apex
-2. Set the amp trainer flag.
+2. Set the "precision" trainer flag.

 Install apex
----------------------------------------------
+^^^^^^^^^^^^
 .. code-block:: bash

    $ git clone https://github.com/NVIDIA/apex
@@ -31,12 +36,25 @@ Install apex


 Enable 16-bit
--------------
+^^^^^^^^^^^^^
+
+.. code-block:: python
+
+    # turn on 16-bit
+    trainer = Trainer(amp_level='O1', precision=16)
+
+If you need to configure the apex init for your particular use case or want to use a different way of doing
+16-bit training, override   :meth:`pytorch_lightning.core.LightningModule.configure_apex`.
+
+TPU 16-bit
+----------
+16-bit on TPus is much simpler. To use 16-bit with TPUs set precision to 16 when using the tpu flag

 .. code-block:: python

    # DEFAULT
-    trainer = Trainer(amp_level='O1', use_amp=False)
+    trainer = Trainer(num_tpu_cores=8, precision=32)
+
+    # turn on 16-bit
+    trainer = Trainer(num_tpu_cores=8, precision=16)

-If you need to configure the apex init for your particular use case or want to use a different way of doing
-16-bit training, override   :meth:`pytorch_lightning.core.LightningModule.configure_apex`.
@@ -55,6 +55,7 @@ PyTorch-Lightning Documentation
   single_gpu
   sequences
   training_tricks
+   tpu
   test_set
   optimizers
   profiler
@@ -60,7 +60,7 @@ Then you could do rapid research by switching between these two and using the sa
    else:
        model = CoolerNotBERT()

-    trainer = Trainer(gpus=4, use_amp=True)
+    trainer = Trainer(gpus=4, precision=16)
    trainer.fit(model)


@@ -0,0 +1,175 @@
+TPU support
+===========
+
+Lightning supports running on TPUs. At this moment, TPUs are only available
+on Google Cloud (GCP). For more information on TPUs
+`watch this video <https://www.youtube.com/watch?v=kPMpmcl_Pyw>`_.
+
+Live demo
+----------
+Check out this `Google Colab <https://colab.research.google.com/drive/1-_LKx4HwAxl5M6xPJmqAAu444LTDQoa3>`_ to see how to train MNIST on TPUs.
+
+TPU Terminology
+---------------
+A TPU is a Tensor processing unit. Each TPU has 8 cores where each
+core is optimized for 128x128 matrix multiplies. In general, a single
+TPU is about as fast as 5 V100 GPUs!
+
+A TPU pod hosts many TPUs on it. Currently, TPU pod v2 has 2048 cores!
+You can request a full pod from Google cloud or a "slice" which gives you
+some subset of those 2048 cores.
+
+How to access TPUs
+-------------------
+To access TPUs there are two main ways.
+
+1. Using google colab.
+2. Using Google Cloud (GCP).
+
+Colab TPUs
+-----------
+Colab is like a jupyter notebook with a free GPU or TPU
+hosted on GCP.
+
+To get a TPU on colab, follow these steps:
+
+1. Go to https://colab.research.google.com/.
+
+2. Click "new notebook" (bottom right of pop-up).
+
+3. Click runtime > change runtime settings. Select Python 3,
+and hardware accelerator "TPU". This will give you a TPU with 8 cores.
+
+4. Next, insert this code into the first cell and execute. This
+will install the xla library that interfaces between PyTorch and
+the TPU.
+
+.. code-block:: python
+
+    import collections
+    from datetime import datetime, timedelta
+    import os
+    import requests
+    import threading
+
+    _VersionConfig = collections.namedtuple('_VersionConfig', 'wheels,server')
+    VERSION = "xrt==1.15.0"  #@param ["xrt==1.15.0", "torch_xla==nightly"]
+    CONFIG = {
+        'xrt==1.15.0': _VersionConfig('1.15', '1.15.0'),
+        'torch_xla==nightly': _VersionConfig('nightly', 'XRT-dev{}'.format(
+            (datetime.today() - timedelta(1)).strftime('%Y%m%d'))),
+    }[VERSION]
+    DIST_BUCKET = 'gs://tpu-pytorch/wheels'
+    TORCH_WHEEL = 'torch-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)
+    TORCH_XLA_WHEEL = 'torch_xla-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)
+    TORCHVISION_WHEEL = 'torchvision-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)
+
+    # Update TPU XRT version
+    def update_server_xrt():
+      print('Updating server-side XRT to {} ...'.format(CONFIG.server))
+      url = 'http://{TPU_ADDRESS}:8475/requestversion/{XRT_VERSION}'.format(
+          TPU_ADDRESS=os.environ['COLAB_TPU_ADDR'].split(':')[0],
+          XRT_VERSION=CONFIG.server,
+      )
+      print('Done updating server-side XRT: {}'.format(requests.post(url)))
+
+    update = threading.Thread(target=update_server_xrt)
+    update.start()
+
+    # Install Colab TPU compat PyTorch/TPU wheels and dependencies
+    !pip uninstall -y torch torchvision
+    !gsutil cp "$DIST_BUCKET/$TORCH_WHEEL" .
+    !gsutil cp "$DIST_BUCKET/$TORCH_XLA_WHEEL" .
+    !gsutil cp "$DIST_BUCKET/$TORCHVISION_WHEEL" .
+    !pip install "$TORCH_WHEEL"
+    !pip install "$TORCH_XLA_WHEEL"
+    !pip install "$TORCHVISION_WHEEL"
+    !sudo apt-get install libomp5
+    update.join()
+5. Once the above is done, install PyTorch Lightning (v 0.6.1+).
+
+.. code-block::
+
+    ! pip install pytorch-lightning
+
+6. Then set up your LightningModule as normal.
+
+7. TPUs require a DistributedSampler. That means you should change your
+train_dataloader (and val, train) code as follows.
+
+.. code-block:: python
+
+    import torch_xla.core.xla_model as xm
+
+    @pl.data_loader
+    def train_dataloader(self):
+        dataset = MNIST(
+            os.getcwd(),
+            train=True,
+            download=True,
+            transform=transforms.ToTensor()
+        )
+
+        # required for TPU support
+        sampler = None
+        if use_tpu:
+            sampler = torch.utils.data.distributed.DistributedSampler(
+                dataset,
+                num_replicas=xm.xrt_world_size(),
+                rank=xm.get_ordinal(),
+                shuffle=True
+            )
+
+        loader = DataLoader(
+            dataset,
+            sampler=sampler,
+            batch_size=32
+        )
+
+        return loader
+
+8. Configure the number of TPU cores in the trainer. You can only choose
+1 or 8. To use a full TPU pod skip to the TPU pod section.
+
+.. code-block:: python
+
+    import pytorch_lightning as pl
+
+    my_model = MyLightningModule()
+    trainer = pl.Trainer(num_tpu_cores=8)
+    trainer.fit(my_model)
+
+That's it! Your model will train on all 8 TPU cores.
+
+TPU Pod
+--------
+To train on more than 8 cores, your code actually doesn't change!
+All you need to do is submit the following command:
+
+.. code-block:: bash
+    $ python -m torch_xla.distributed.xla_dist
+    --tpu=$TPU_POD_NAME
+    --conda-env=torch-xla-nightly
+    -- python /usr/share/torch-xla-0.5/pytorch/xla/test/test_train_imagenet.py --fake_data
+
+16 bit precision
+-----------------
+Lightning also supports training in 16-bit precision with TPUs.
+By default, TPU training will use 32-bit precision. To enable 16-bit, also
+set the 16-bit flag.
+
+.. code-block:: python
+
+    import pytorch_lightning as pl
+
+    my_model = MyLightningModule()
+    trainer = pl.Trainer(num_tpu_cores=8, precision=16)
+    trainer.fit(my_model)
+
+Under the hood the xla library will use the `bfloat16 type <https://en.wikipedia.org/wiki/Bfloat16_floating-point_format>`_.
+
+
+About XLA
+----------
+XLA is the library that interfaces PyTorch with the TPUs.
+For more information check out `XLA <https://github.com/pytorch/xla>`_.
@@ -113,10 +113,10 @@ class ModelHooks(torch.nn.Module):

        """

-    def backward(self, use_amp, loss, optimizer, optimizer_idx):
+    def backward(self, trainer, loss, optimizer, optimizer_idx):
        """Override backward with your own implementation if you need to

-        :param use_amp: Whether amp was requested or not
+        :param trainer: Pointer to the trainer
        :param loss: Loss is already scaled by accumulated grads
        :param optimizer: Current optimizer being used
        :param optimizer_idx: Index of the current optimizer being used
@@ -137,8 +137,11 @@ class ModelHooks(torch.nn.Module):
                    loss.backward()

        """
-        if use_amp:
-            with amp.scale_loss(loss, optimizer) as scaled_loss:
-                scaled_loss.backward()
+        if trainer.precision == 16:
+
+            # .backward is not special on 16-bit with TPUs
+            if not trainer.on_tpu:
+                with amp.scale_loss(loss, optimizer) as scaled_loss:
+                    scaled_loss.backward()
        else:
            loss.backward()
@@ -12,6 +12,9 @@ import logging as log

 class TrainerAMPMixin(ABC):

+    def __init__(self):
+        self.use_amp = None
+
    def init_amp(self, use_amp):
        self.use_amp = use_amp and APEX_AVAILABLE
        if self.use_amp:
@@ -36,6 +36,8 @@ class TrainerDataLoadingMixin(ABC):
        self.use_ddp2 = None
        self.shown_warnings = None
        self.val_check_interval = None
+        self.use_tpu = None
+        self.tpu_local_core_rank = None

    def _percent_range_check(self, name):
        value = getattr(self, name)
@@ -80,9 +82,10 @@ class TrainerDataLoadingMixin(ABC):
            self.val_check_batch = max(1, self.val_check_batch)

        on_ddp = self.use_ddp or self.use_ddp2
-        if on_ddp and not isinstance(self.get_train_dataloader().sampler, DistributedSampler):
+        needs_sampler = on_ddp or self.use_tpu
+        if needs_sampler and not isinstance(self.get_train_dataloader().sampler, DistributedSampler):
            msg = """
-            You're using multiple gpus and multiple nodes without using a DistributedSampler
+            You're using multiple gpus and multiple nodes, or TPUs without using a
            to assign a subset of your data to each process. To silence this warning, pass a
            DistributedSampler to your DataLoader.

@@ -119,13 +122,14 @@ class TrainerDataLoadingMixin(ABC):
            self.num_val_batches = int(self.num_val_batches * self.val_percent_check)

        on_ddp = self.use_ddp or self.use_ddp2
-        if on_ddp and self.get_val_dataloaders() is not None:
+        needs_sampler = on_ddp or self.use_tpu
+        if needs_sampler and self.get_val_dataloaders() is not None:
            for dataloader in self.get_val_dataloaders():
                if not isinstance(dataloader.sampler, DistributedSampler):
                    msg = """
                    Your val_dataloader(s) don't use DistributedSampler.

-                    You're using multiple gpus and multiple nodes without using a
+                    You're using multiple gpus and multiple nodes, or TPUs without using a
                    DistributedSampler to assign a subset of your data to each process.
                    To silence this warning, pass a DistributedSampler to your DataLoader.

@@ -162,13 +166,14 @@ class TrainerDataLoadingMixin(ABC):
            self.num_test_batches = int(self.num_test_batches * self.test_percent_check)

        on_ddp = self.use_ddp or self.use_ddp2
-        if on_ddp and self.get_test_dataloaders() is not None:
+        needs_sampler = on_ddp or self.use_tpu
+        if needs_sampler and self.get_test_dataloaders() is not None:
            for dataloader in self.get_test_dataloaders():
                if not isinstance(dataloader.sampler, DistributedSampler):
                    msg = """
                    Your `test_dataloader(s)` don't use DistributedSampler.

-                    You're using multiple gpus and multiple nodes without using a
+                    You're using multiple gpus and multiple nodes, or TPUs without using a
                    DistributedSampler to assign a subset of your data to each process.
                    To silence this warning, pass a DistributedSampler to your DataLoader.

@@ -210,6 +215,14 @@ class TrainerDataLoadingMixin(ABC):
            self.get_test_dataloaders()
            self.get_val_dataloaders()

+        # on TPUs load each dataloader only on process 0
+        # this will trigger the data downloads
+        if self.use_tpu:
+            if self.tpu_local_core_rank == 0:
+                self.get_train_dataloader()
+                self.get_test_dataloaders()
+                self.get_val_dataloaders()
+
        # support IterableDataset for train data
        self.is_iterable_train_dataloader = (
            EXIST_ITER_DATASET and isinstance(self.get_train_dataloader().dataset, IterableDataset))
@@ -144,6 +144,7 @@ class TrainerDDPMixin(ABC):
        self.distributed_backend = None
        self.use_amp = None
        self.amp_level = None
+        self.use_tpu = None

    @abstractmethod
    def copy_trainer_model_properties(self, model):
@@ -160,6 +161,13 @@ class TrainerDDPMixin(ABC):
        # this is just empty shell for code from other class
        pass

+    def init_tpu(self):
+        # turn off all the GPU stuff
+        self.distributed_backend = None
+
+        # enable tpu
+        self.use_tpu = True
+
    def set_distributed_mode(self, distributed_backend, num_gpu_nodes):
        # skip for CPU
        if self.num_gpus == 0:
@@ -335,6 +335,8 @@ Here lightning distributes parts of your module across available GPUs to optimiz
 """

 from abc import ABC, abstractmethod
+import logging as log
+import os

 import torch

@@ -351,6 +353,13 @@ try:
 except ImportError:
    APEX_AVAILABLE = False

+try:
+    import torch_xla.core.xla_model as xm
+    XLA_AVAILABLE = True
+
+except ImportError:
+    XLA_AVAILABLE = False
+

 class TrainerDPMixin(ABC):

@@ -366,6 +375,12 @@ class TrainerDPMixin(ABC):
        self.single_gpu = None
        self.root_gpu = None
        self.amp_level = None
+        self.precision = None
+        self.current_tpu_idx = None
+        self.proc_rank = None
+        self.tpu_local_core_rank = None
+        self.tpu_global_core_rank = None
+        self.use_tpu = None

    @abstractmethod
    def run_pretrain_routine(self, model):
@@ -394,32 +409,47 @@ class TrainerDPMixin(ABC):
            m.use_amp = self.use_amp
            m.testing = self.testing
            m.single_gpu = self.single_gpu
+            m.use_tpu = self.use_tpu
+            m.tpu_local_core_rank = self.tpu_local_core_rank
+            m.tpu_global_core_rank = self.tpu_global_core_rank
+
+    def transfer_batch_to_tpu(self, batch):
+        return self.__transfer_data_to_device(batch, device='tpu')

    def transfer_batch_to_gpu(self, batch, gpu_id):
-        # base case: object can be directly moved using `cuda` or `to`
-        if callable(getattr(batch, 'cuda', None)):
-            return batch.cuda(gpu_id)
+        return self.__transfer_data_to_device(batch, device='gpu', gpu_id=gpu_id)

-        if callable(getattr(batch, 'to', None)):
-            return batch.to(torch.device('cuda', gpu_id))
+    def __transfer_data_to_device(self, batch, device, gpu_id=None):
+        if device == 'tpu' and XLA_AVAILABLE:
+            # base case: object can be directly moved using `to`
+            if callable(getattr(batch, 'to', None)):
+                return batch.to(xm.xla_device())
+
+        if device == 'gpu':
+            # base case: object can be directly moved using `cuda` or `to`
+            if callable(getattr(batch, 'cuda', None)):
+                return batch.cuda(gpu_id)
+
+            if callable(getattr(batch, 'to', None)):
+                return batch.to(torch.device('cuda', gpu_id))

        # when list
        if isinstance(batch, list):
            for i, x in enumerate(batch):
-                batch[i] = self.transfer_batch_to_gpu(x, gpu_id)
+                batch[i] = self.transfer_batch_to_tpu(x)
            return batch

        # when tuple
        if isinstance(batch, tuple):
            batch = list(batch)
            for i, x in enumerate(batch):
-                batch[i] = self.transfer_batch_to_gpu(x, gpu_id)
+                batch[i] = self.transfer_batch_to_tpu(x)
            return tuple(batch)

        # when dict
        if isinstance(batch, dict):
            for k, v in batch.items():
-                batch[k] = self.transfer_batch_to_gpu(v, gpu_id)
+                batch[k] = self.transfer_batch_to_tpu(v)

            return batch

@@ -440,6 +470,34 @@ class TrainerDPMixin(ABC):

        self.run_pretrain_routine(model)

+    def tpu_train(self, tpu_core_idx, model):
+        # put model on tpu
+        model.to(xm.xla_device())
+
+        # get the appropriate tpu ranks
+        self.tpu_local_core_rank = xm.get_local_ordinal()
+        self.tpu_global_core_rank = xm.get_ordinal()
+
+        # avoid duplicating progress bar
+        self.show_progress_bar = self.show_progress_bar and self.tpu_global_core_rank == 0
+
+        # track current tpu
+        self.current_tpu_idx = tpu_core_idx
+        self.proc_rank = self.tpu_local_core_rank
+
+        # CHOOSE OPTIMIZER
+        # allow for lr schedulers as well
+        self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
+
+        # init 16 bit for TPU
+        if self.precision == 16:
+            os.environ['XLA_USE_BF16'] = 1
+
+        m = f'INIT TPU local core: {self.tpu_local_core_rank}, ' \
+            f'global rank: {self.tpu_global_core_rank}'
+        log.info(m)
+        self.run_pretrain_routine(model)
+
    def dp_train(self, model):

        # CHOOSE OPTIMIZER
@@ -131,6 +131,14 @@ from tqdm.auto import tqdm

 from pytorch_lightning.utilities.debugging import MisconfigurationException

+try:
+    import torch_xla.distributed.parallel_loader as xla_pl
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+except ImportError:
+    XLA_AVAILABLE = False
+

 class TrainerEvaluationLoopMixin(ABC):

@@ -159,6 +167,7 @@ class TrainerEvaluationLoopMixin(ABC):
        self.callback_metrics = None
        self.get_test_dataloaders = None
        self.get_val_dataloaders = None
+        self.use_tpu = None

    @abstractmethod
    def copy_trainer_model_properties(self, model):
@@ -175,6 +184,11 @@ class TrainerEvaluationLoopMixin(ABC):
        # this is just empty shell for code from other class
        pass

+    @abstractmethod
+    def transfer_batch_to_tpu(self, batch):
+        # this is just empty shell for code from other class
+        pass
+
    @abstractmethod
    def transfer_batch_to_gpu(self, batch, gpu):
        # this is just empty shell for code from other class
@@ -215,6 +229,13 @@ class TrainerEvaluationLoopMixin(ABC):
        # run validation
        for dataloader_idx, dataloader in enumerate(dataloaders):
            dl_outputs = []
+
+            # on TPU we have to wrap it under the ParallelLoader
+            if self.use_tpu:
+                device = xm.xla_device()
+                dataloader = xla_pl.ParallelLoader(dataloader, [device])
+                dataloader = dataloader.per_device_loader(device)
+
            for batch_idx, batch in enumerate(dataloader):

                if batch is None:  # pragma: no cover
@@ -356,6 +377,11 @@ class TrainerEvaluationLoopMixin(ABC):
            batch = self.transfer_batch_to_gpu(batch, root_gpu)
            args[0] = batch

+        # TPU
+        if self.use_tpu:
+            batch = self.transfer_batch_to_tpu(batch)
+            args[0] = batch
+
        # CPU
        if test:
            output = model.test_step(*args)
@@ -36,6 +36,15 @@ try:
 except ImportError:
    APEX_AVAILABLE = False

+try:
+    import torch_xla
+    import torch_xla.core.xla_model as xm
+    import torch_xla.distributed.xla_multiprocessing as xmp
+
+    XLA_AVAILABLE = True
+except ImportError:
+    XLA_AVAILABLE = False
+

 class Trainer(TrainerIOMixin,
              TrainerDPMixin,
@@ -62,6 +71,7 @@ class Trainer(TrainerIOMixin,
            nb_gpu_nodes=None,  # backward compatible, todo: remove in v0.8.0
            num_nodes=1,
            gpus=None,
+            num_tpu_cores=None,
            log_gpu_memory=None,
            show_progress_bar=True,
            overfit_pct=0.0,
@@ -81,7 +91,8 @@ class Trainer(TrainerIOMixin,
            row_log_interval=10,
            add_row_log_interval=None,  # backward compatible, todo: remove in v0.8.0
            distributed_backend=None,
-            use_amp=False,
+            use_amp=False,  # backward compatible, todo: remove in v0.8.0
+            precision=32,
            print_nan_grads=False,
            weights_summary='full',
            weights_save_path=None,
@@ -163,7 +174,7 @@ class Trainer(TrainerIOMixin,
                    trainer = Trainer(gradient_clip_val=0.0)

            gradient_clip (int):
-                .. deprecated:: 0.5.0
+                .. warning: .. deprecated:: 0.5.0
                    Use `gradient_clip_val` instead. Will remove 0.8.0.

            process_position (int): orders the tqdm bar when running multiple models on same machine.
@@ -182,7 +193,7 @@ class Trainer(TrainerIOMixin,
                    trainer = Trainer(num_nodes=8)

            nb_gpu_nodes (int):
-                .. deprecated:: 0.5.0
+                ..warning:: .. deprecated:: 0.5.0
                    Use `num_nodes` instead. Will remove 0.8.0.

            gpus (list|str|int): Which GPUs to train on.
@@ -205,6 +216,48 @@ class Trainer(TrainerIOMixin,
                    # combine with num_nodes to train on multiple GPUs across nodes
                    trainer = Trainer(gpus=2, num_nodes=4) # uses 8 gpus in total

+            num_tpu_cores (int): How many TPU cores to train on (1 or 8).
+                A single TPU v2 or v3 has 8 cores. A TPU pod has
+                up to 2048 cores. A slice of a POD means you get as many cores
+                as you request.
+
+                You MUST use DistributedDataSampler with your dataloader for this
+                to work. Your effective batch size is batch_size * total tpu cores.
+
+                This parameter can be either 1 or 8.
+
+                Example::
+
+                    # your_trainer_file.py
+
+                    # default used by the Trainer (ie: train on CPU)
+                    trainer = Trainer(num_tpu_cores=None)
+
+                    # int: train on a single core
+                    trainer = Trainer(num_tpu_cores=1)
+
+                    # int: train on all cores few cores
+                    trainer = Trainer(num_tpu_cores=8)
+
+                    # for 8+ cores must submit via xla script with
+                    # a max of 8 cores specified. The XLA script
+                    # will duplicate script onto each TPU in the POD
+                    trainer = Trainer(num_tpu_cores=8)
+
+                    # -1: train on all available TPUs
+                    trainer = Trainer(num_tpu_cores=-1)
+
+            To train on more than 8 cores (ie: a POD),
+            submit this script using the xla_dist script.
+
+            Example::
+
+                $ python -m torch_xla.distributed.xla_dist
+                --tpu=$TPU_POD_NAME
+                --conda-env=torch-xla-nightly
+                --env=XLA_USE_BF16=1
+                -- python your_trainer_file.py
+
            log_gpu_memory (str): None, 'min_max', 'all'. Might slow performance
                because it uses the output of nvidia-smi.
                Example::
@@ -279,7 +332,7 @@ class Trainer(TrainerIOMixin,
                    trainer = Trainer(max_epochs=1000)

            max_nb_epochs (int):
-                .. deprecated:: 0.5.0
+                .. warning:: .. deprecated:: 0.5.0
                    Use `max_epochs` instead. Will remove 0.8.0.

            min_epochs (int): Force training for at least these many epochs
@@ -289,7 +342,7 @@ class Trainer(TrainerIOMixin,
                    trainer = Trainer(min_epochs=1)

            min_nb_epochs (int):
-                .. deprecated:: 0.5.0
+                .. warning:: .. deprecated:: 0.5.0
                    Use `min_nb_epochs` instead. Will remove 0.8.0.

            train_percent_check (int): How much of training dataset to check.
@@ -350,7 +403,7 @@ class Trainer(TrainerIOMixin,
                    trainer = Trainer(row_log_interval=10)

            add_row_log_interval (int):
-                .. deprecated:: 0.5.0
+                .. warning:: .. deprecated:: 0.5.0
                    Use `row_log_interval` instead. Will remove 0.8.0.

            distributed_backend (str): The distributed backend to use.
@@ -374,11 +427,26 @@ class Trainer(TrainerIOMixin,
                    # useful for things like increasing the number of negative samples
                    trainer = Trainer(gpus=2, num_nodes=2, distributed_backend='ddp2')

-            use_amp (bool): If true uses apex for 16bit precision
+            use_amp (bool):
+                .. warning:: .. deprecated:: 0.6.1
+                    Use `precision` instead. Will remove 0.8.0.
+
+            precision (int): Full precision (32), half precision (16).
+                Can be used on CPU, GPU or TPUs.
+
+                If used on TPU will use torch.bfloat16 but tensor printing
+                will still show torch.float32.
+
                Example::

                    # default used by the Trainer
-                    trainer = Trainer(use_amp=False)
+                    trainer = Trainer(precision=32)
+
+                    # 16-bit precision
+                    trainer = Trainer(precision=16)
+
+                    # one day
+                    trainer = Trainer(precision=8|4|2)

            print_nan_grads (bool): Prints gradients with nan values
                Example::
@@ -435,7 +503,7 @@ class Trainer(TrainerIOMixin,
                    trainer = Trainer(num_sanity_val_steps=0)

            nb_sanity_val_steps (int):
-                .. deprecated:: 0.5.0
+                .. warning:: .. deprecated:: 0.5.0
                    Use `num_sanity_val_steps` instead. Will remove 0.8.0.

            truncated_bptt_steps (int): Truncated back prop breaks performs backprop every k steps of
@@ -517,6 +585,12 @@ class Trainer(TrainerIOMixin,
        self.check_val_every_n_epoch = check_val_every_n_epoch
        self.track_grad_norm = track_grad_norm
        self.on_gpu = True if (gpus and torch.cuda.is_available()) else False
+
+        # tpu config
+        self.on_tpu = num_tpu_cores is not None
+        self.num_tpu_cores = num_tpu_cores
+        assert num_tpu_cores in [1, 8, None], 'num_tpu_cores can only be 1 or 8'
+
        self.process_position = process_position
        self.weights_summary = weights_summary

@@ -614,6 +688,11 @@ class Trainer(TrainerIOMixin,
        self.data_parallel_device_ids = parse_gpu_ids(gpus)
        self.root_gpu = determine_root_gpu_device(self.data_parallel_device_ids)

+        # tpu state flags
+        self.use_tpu = False
+        self.tpu_local_core_rank = None
+        self.tpu_global_core_rank = None
+
        # distributed backend choice
        self.use_ddp = False
        self.use_ddp2 = False
@@ -622,6 +701,11 @@ class Trainer(TrainerIOMixin,
        self.distributed_backend = distributed_backend
        self.set_distributed_mode(distributed_backend, num_nodes)

+        # override dist backend when using tpus
+        if self.on_tpu:
+            self.init_tpu()
+            self.current_tpu_idx = None
+
        # init flags for SLURM+ddp to work
        self.proc_rank = 0
        self.world_size = 1
@@ -653,6 +737,9 @@ class Trainer(TrainerIOMixin,

        # 16 bit mixed precision training using apex
        self.amp_level = amp_level
+        self.precision = precision
+        if self.precision == 16:
+            use_amp = True
        self.init_amp(use_amp)

    @property
@@ -724,7 +811,7 @@ class Trainer(TrainerIOMixin,

        :return: dictionary

-        .. deprecated:: 0.5.0
+        .. warning:: .. deprecated:: 0.5.0
                    Use `training_tqdm_dict` instead. Will remove 0.8.0.
        """
        warnings.warn("`tng_tqdm_dic` has renamed to `training_tqdm_dict` since v0.5.0"
@@ -765,6 +852,13 @@ class Trainer(TrainerIOMixin,
        elif self.single_gpu:
            self.single_gpu_train(model)

+        elif self.use_tpu:
+            log.info(f'training on {self.num_tpu_cores} TPU cores')
+
+            #  COLAB_GPU is an env var available by default in Colab environments.
+            start_method = 'fork' if os.getenv('COLAB_GPU') else 'spawn'
+            xmp.spawn(self.tpu_train, args=(model,), nprocs=self.num_tpu_cores, start_method=start_method)
+
        # ON CPU
        else:
            # run through amp wrapper
@@ -167,6 +167,21 @@ try:
 except ImportError:
    APEX_AVAILABLE = False

+try:
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+except ImportError:
+    XLA_AVAILABLE = False
+
+try:
+    import torch_xla.distributed.parallel_loader as xla_pl
+
+    XLA_AVAILABLE = True
+
+except ImportError:
+    XLA_AVAILABLE = False
+

 class TrainerTrainLoopMixin(ABC):

@@ -179,6 +194,7 @@ class TrainerTrainLoopMixin(ABC):
        self.use_dp = None
        self.use_ddp2 = None
        self.single_gpu = None
+        self.use_tpu = None
        self.data_parallel_device_ids = None
        self.check_val_every_n_epoch = None
        self.num_training_batches = None
@@ -212,6 +228,8 @@ class TrainerTrainLoopMixin(ABC):
        self.get_train_dataloader = None
        self.reduce_lr_on_plateau_scheduler = None
        self.profiler = None
+        self.batch_idx = None
+        self.precision = None

    @property
    def max_nb_epochs(self):
@@ -251,6 +269,11 @@ class TrainerTrainLoopMixin(ABC):
        # this is just empty shell for code from other class
        pass

+    @abstractmethod
+    def transfer_batch_to_tpu(self, batch):
+        # this is just empty shell for code from other class
+        pass
+
    @abstractmethod
    def clip_gradients(self):
        # this is just empty shell for code from other class
@@ -288,7 +311,8 @@ class TrainerTrainLoopMixin(ABC):
        # run all epochs
        for epoch in range(self.current_epoch, self.max_epochs):
            # set seed for distributed sampler (enables shuffling for each epoch)
-            if self.use_ddp and hasattr(self.get_train_dataloader().sampler, 'set_epoch'):
+            if (self.use_ddp or self.use_tpu) \
+                    and hasattr(self.get_train_dataloader().sampler, 'set_epoch'):
                self.get_train_dataloader().sampler.set_epoch(epoch)

            # get model
@@ -376,9 +400,18 @@ class TrainerTrainLoopMixin(ABC):
            with self.profiler.profile('on_epoch_start'):
                model.on_epoch_start()

+        # request the dataloader
+        train_dataloader = self.get_train_dataloader()
+
+        # on TPU we have to wrap it under the ParallelLoader
+        if self.use_tpu:
+            device = xm.xla_device()
+            train_dataloader = xla_pl.ParallelLoader(train_dataloader, [device])
+            train_dataloader = train_dataloader.per_device_loader(device)
+
        # run epoch
        for batch_idx, batch in self.profiler.profile_iterable(
-            enumerate(self.get_train_dataloader()), "get_train_batch"
+            enumerate(train_dataloader), "get_train_batch"
        ):
            # stop epoch if we limited the number of training batches
            if batch_idx >= self.num_training_batches:
@@ -505,7 +538,7 @@ class TrainerTrainLoopMixin(ABC):
                    # backward pass
                    model_ref = self.get_model()
                    with self.profiler.profile('model_backward'):
-                        model_ref.backward(self.use_amp, closure_loss, optimizer, opt_idx)
+                        model_ref.backward(self, closure_loss, optimizer, opt_idx)

                    # track metrics for callbacks
                    all_callback_metrics.append(callback_metrics)
@@ -549,8 +582,11 @@ class TrainerTrainLoopMixin(ABC):
                    # override function to modify this behavior
                    model = self.get_model()
                    with self.profiler.profile('optimizer_step'):
-                        model.optimizer_step(self.current_epoch, batch_idx,
-                                             optimizer, opt_idx, optimizer_closure)
+                        if self.use_tpu:
+                            xm.optimizer_step(optimizer, barrier=True)
+                        else:
+                            model.optimizer_step(self.current_epoch, batch_idx,
+                                                 optimizer, opt_idx, optimizer_closure)

                    # calculate running loss for display
                    self.running_loss.append(self.batch_loss_value)
@@ -614,6 +650,12 @@ class TrainerTrainLoopMixin(ABC):
            args[0] = batch
            output = self.model.training_step(*args)

+        # TPU support
+        elif self.use_tpu:
+            batch = self.transfer_batch_to_tpu(copy.copy(batch))
+            args[0] = batch
+            output = self.model.training_step(*args)
+
        # CPU forward
        else:
            output = self.model.training_step(*args)
@@ -26,7 +26,7 @@ def test_amp_single_gpu(tmpdir):
        max_epochs=1,
        gpus=1,
        distributed_backend='ddp',
-        use_amp=True
+        precision=16
    )

    tutils.run_model_test(trainer_options, model)
@@ -49,7 +49,7 @@ def test_no_amp_single_gpu(tmpdir):
        max_epochs=1,
        gpus=1,
        distributed_backend='dp',
-        use_amp=True
+        precision=16
    )

    trainer = Trainer(**trainer_options)
@@ -75,7 +75,7 @@ def test_amp_gpu_ddp(tmpdir):
        max_epochs=1,
        gpus=2,
        distributed_backend='ddp',
-        use_amp=True
+        precision=16
    )

    tutils.run_model_test(trainer_options, model)
@@ -101,7 +101,7 @@ def test_amp_gpu_ddp_slurm_managed(tmpdir):
        max_epochs=1,
        gpus=[0],
        distributed_backend='ddp',
-        use_amp=True
+        precision=16
    )

    # exp file to get meta
@@ -140,7 +140,7 @@ def test_cpu_model_with_amp(tmpdir):
        max_epochs=1,
        train_percent_check=0.4,
        val_percent_check=0.4,
-        use_amp=True
+        precision=16
    )

    model, hparams = tutils.get_model()
@@ -163,7 +163,7 @@ def test_amp_gpu_dp(tmpdir):
        max_epochs=1,
        gpus='0, 1',  # test init with gpu string
        distributed_backend='dp',
-        use_amp=True
+        precision=16
    )

    trainer = Trainer(**trainer_options)
@@ -230,7 +230,7 @@ def test_ddp_sampler_error(tmpdir):
        max_epochs=1,
        gpus=[0, 1],
        distributed_backend='ddp',
-        use_amp=True
+        precision=16
    )

    with pytest.warns(UserWarning):