diff --git a/docs/source/apex.rst b/docs/source/apex.rst index d4dbc4cd..3c6b86c7 100644 --- a/docs/source/apex.rst +++ b/docs/source/apex.rst @@ -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`. \ No newline at end of file diff --git a/docs/source/index.rst b/docs/source/index.rst index ebd0fb9c..3232e74f 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -55,6 +55,7 @@ PyTorch-Lightning Documentation single_gpu sequences training_tricks + tpu test_set optimizers profiler diff --git a/docs/source/new-project.rst b/docs/source/new-project.rst index 35834518..c269c819 100644 --- a/docs/source/new-project.rst +++ b/docs/source/new-project.rst @@ -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) diff --git a/docs/source/tpu.rst b/docs/source/tpu.rst new file mode 100644 index 00000000..e4dd0c33 --- /dev/null +++ b/docs/source/tpu.rst @@ -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 `_. + +Live demo +---------- +Check out this `Google Colab `_ 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 `_. + + +About XLA +---------- +XLA is the library that interfaces PyTorch with the TPUs. +For more information check out `XLA `_. diff --git a/pytorch_lightning/core/hooks.py b/pytorch_lightning/core/hooks.py index 6c5f74b2..dca3b36c 100644 --- a/pytorch_lightning/core/hooks.py +++ b/pytorch_lightning/core/hooks.py @@ -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() diff --git a/pytorch_lightning/trainer/auto_mix_precision.py b/pytorch_lightning/trainer/auto_mix_precision.py index bd6edd93..135a0bce 100644 --- a/pytorch_lightning/trainer/auto_mix_precision.py +++ b/pytorch_lightning/trainer/auto_mix_precision.py @@ -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: diff --git a/pytorch_lightning/trainer/data_loading.py b/pytorch_lightning/trainer/data_loading.py index ee6c341c..496debbf 100644 --- a/pytorch_lightning/trainer/data_loading.py +++ b/pytorch_lightning/trainer/data_loading.py @@ -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)) diff --git a/pytorch_lightning/trainer/distrib_data_parallel.py b/pytorch_lightning/trainer/distrib_data_parallel.py index 96adeaf0..01bcdb1d 100644 --- a/pytorch_lightning/trainer/distrib_data_parallel.py +++ b/pytorch_lightning/trainer/distrib_data_parallel.py @@ -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: diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py index d1417d2c..8cec776a 100644 --- a/pytorch_lightning/trainer/distrib_parts.py +++ b/pytorch_lightning/trainer/distrib_parts.py @@ -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 diff --git a/pytorch_lightning/trainer/evaluation_loop.py b/pytorch_lightning/trainer/evaluation_loop.py index d974292c..f5d2b932 100644 --- a/pytorch_lightning/trainer/evaluation_loop.py +++ b/pytorch_lightning/trainer/evaluation_loop.py @@ -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) diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index c7da1111..d2395ee5 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -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 diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py index 0047f7e7..d8bdb3f4 100644 --- a/pytorch_lightning/trainer/training_loop.py +++ b/pytorch_lightning/trainer/training_loop.py @@ -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) diff --git a/tests/test_amp.py b/tests/test_amp.py index 0ea77991..5a349062 100644 --- a/tests/test_amp.py +++ b/tests/test_amp.py @@ -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) diff --git a/tests/test_gpu_models.py b/tests/test_gpu_models.py index 51f2541b..e982d2f2 100644 --- a/tests/test_gpu_models.py +++ b/tests/test_gpu_models.py @@ -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):