Allow user to select individual TPU core to train on (#1729)

* added tpu_id added tpu_id to mixins * train on individual tpu * parallel loader if tpu_id is None * removed progress_bar_refresh_rate * chlog * replaced num_tpu_cores with tpu_cores * set tpu_id to None if int * changed num_tpu_cores to tpu_cores in docs * updated docs * updated __init__.py removed self.tpu_id for ParallelLoader * Update pytorch_lightning/trainer/__init__.py * check if tpu_cores is a list Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com> * xla device conditional * num_tpu_cores deprecation * removed duplicate warning * fixed pep8 error * Revert "removed duplicate warning" This reverts commit 8adb0a9b * deprecated api update * fixed recursion error * fixed tests * fixed flake errors * removed current_tpu_index * Update CHANGELOG.md * Update trainer.py Co-authored-by: Jirka <jirka.borovec@seznam.cz> Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com> Co-authored-by: William Falcon <waf2107@columbia.edu>
2026-09-09 11:32:07 +08:00 · 2020-05-17 16:30:54 -04:00
parent 1a797bdad5
commit 7c7e50ca47
14 changed files with 104 additions and 126 deletions
@@ -8,8 +8,12 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).

 ### Added

+- Added type hints in `Trainer.fit()` and `Trainer.test()` to reflect that also a list of dataloaders can be passed in ([#1723](https://github.com/PyTorchLightning/pytorch-lightning/pull/1723)).
+
 ### Changed

+- Allow user to select individual TPU core to train on ([#1729](https://github.com/PyTorchLightning/pytorch-lightning/pull/1729))
+
 ### Deprecated

 ### Removed
@@ -287,7 +287,11 @@ trainer = Trainer(max_epochs=1, gpus=8, num_nodes=32)

 Or TPUs
 ```python
-trainer = Trainer(num_tpu_cores=8)
+# Distributes TPU core training
+trainer = Trainer(tpu_cores=8)
+
+# Single TPU core training
+trainer = Trainer(tpu_cores=[1])
 ```

 When you're done training, run the test accuracy
@@ -58,7 +58,7 @@ TPU 16-bit
 .. testcode::

    # DEFAULT
-    trainer = Trainer(num_tpu_cores=8, precision=32)
+    trainer = Trainer(tpu_cores=8, precision=32)

    # turn on 16-bit
-    trainer = Trainer(num_tpu_cores=8, precision=16)
+    trainer = Trainer(tpu_cores=8, precision=16)
@@ -185,7 +185,7 @@ EXACTLY the same as you would a PyTorch Module.

 Out:

- .. code-block:: none
+ .. code-block:: python

    torch.Size([1, 10])

@@ -519,50 +519,8 @@ First, change the runtime to TPU (and reinstall lightning).

 Next, install the required xla library (adds support for PyTorch on TPUs)

-.. code-block:: python
-
-    import collections
-    from datetime import datetime, timedelta
-    import os
-    import requests
-    import threading
-
-    _VersionConfig = collections.namedtuple('_VersionConfig', 'wheels,server')
-    VERSION = "torch_xla==nightly"  #@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()
-
-.. code-block::
-
-    # 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()
+    !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
+    !python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev

 In distributed training (multiple GPUs and multiple TPU cores) each GPU or TPU core will run a copy
 of this program. This means that without taking any care you will download the dataset N times which
@@ -609,7 +567,7 @@ Now we can train the LightningModule on a TPU without doing anything else!
 .. code-block:: python

    model = LitMNIST()
-    trainer = Trainer(num_tpu_cores=8)
+    trainer = Trainer(tpu_cores=8)
    trainer.fit(model)

 You'll now see the TPU cores booting up.
@@ -696,7 +654,7 @@ while checking the validation set.
    from pytorch_lightning import Trainer

    model = LitMNIST()
-    trainer = Trainer(num_tpu_cores=8)
+    trainer = Trainer(tpu_cores=8)
    trainer.fit(model)

 You may have noticed the words `Validation sanity check` logged. This is because Lightning runs 5 batches
@@ -747,7 +705,7 @@ Once you train your model simply call `.test()`.
    from pytorch_lightning import Trainer

    model = LitMNIST()
-    trainer = Trainer(num_tpu_cores=8)
+    trainer = Trainer(tpu_cores=8)
    trainer.fit(model)

    # run test set
@@ -769,7 +727,7 @@ You can also run the test from a saved lightning model
 .. code-block:: python

    model = LitMNIST.load_from_checkpoint(PATH)
-    trainer = Trainer(num_tpu_cores=8)
+    trainer = Trainer(tpu_cores=8)
    trainer.test(model)

 .. note:: Lightning disables gradients, puts model in eval mode and does everything needed for testing.
@@ -130,7 +130,7 @@ Lightning allows multiple ways of training
 - DistributedDataParallel (`distributed_backend='ddp'`) (multiple-gpus across many machines).
 - DistributedDataParallel2 (`distributed_backend='ddp2'`) (dp in a machine, ddp across machines).
 - Horovod (`distributed_backend='horovod'`) (multi-machine, multi-gpu, configured at runtime)
- TPUs (`num_tpu_cores=8|x`) (tpu or TPU pod)
+- TPUs (`tpu_cores=8|x`) (tpu or TPU pod)

 .. note:: If you request multiple GPUs without setting a mode, ddp will be automatically used.

@@ -189,7 +189,7 @@ However, this time you need to specifically call test (this is done so you don't
    # OPTION 2:
    # test after loading weights
    model = LitModel.load_from_checkpoint(PATH)
-    trainer = Trainer(num_tpu_cores=1)
+    trainer = Trainer(tpu_cores=1)
    trainer.test()

 Again, under the hood, lightning does the following in (pseudocode):
@@ -236,7 +236,7 @@ Without changing a SINGLE line of your code, you can now do the following with t
    # train on TPUs using 16 bit precision with early stopping
    # using only half the training data and checking validation every quarter of a training epoch
    trainer = Trainer(
-        nb_tpu_cores=8,
+        tpu_cores=8,
        precision=16,
        early_stop_checkpoint=True,
        train_percent_check=0.5,
@@ -1,8 +1,8 @@
 TPU support
 ===========

-Lightning supports running on TPUs. At this moment, TPUs are only available
-on Google Cloud (GCP). For more information on TPUs
+Lightning supports running on TPUs. At this moment, TPUs are available
+on Google Cloud (GCP), Google Colab and Kaggle Environments. For more information on TPUs
 `watch this video <https://www.youtube.com/watch?v=kPMpmcl_Pyw>`_.

 ---------------
@@ -31,6 +31,7 @@ To access TPUs there are two main ways.

 1. Using google colab.
 2. Using Google Cloud (GCP).
+3. Using Kaggle.

 ---------------

@@ -51,50 +52,10 @@ To get a TPU on colab, follow these steps:
 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()
-
   .. code-block::

-        # 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()
+    !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
+    !python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev

 5. Once the above is done, install PyTorch Lightning (v 0.7.0+).

@@ -156,13 +117,23 @@ To use a full TPU pod skip to the TPU pod section.
    import pytorch_lightning as pl

    my_model = MyLightningModule()
-    trainer = pl.Trainer(num_tpu_cores=8)
+    trainer = pl.Trainer(tpu_cores=8)
    trainer.fit(my_model)

 That's it! Your model will train on all 8 TPU cores.

 ---------------

+Single TPU core training
+----------------------------
+Lightning supports training on a single TPU core. Just pass the TPU core ID [1-8] in a list.
+
+.. code-block:: python
+
+    trainer = pl.Trainer(tpu_cores=[1])
+
+---------------
+
 Distributed Backend with TPU
 ----------------------------
 The ```distributed_backend``` option used for GPUs does not apply to TPUs.
@@ -195,7 +166,7 @@ set the 16-bit flag.
    import pytorch_lightning as pl

    my_model = MyLightningModule()
-    trainer = pl.Trainer(num_tpu_cores=8, precision=16)
+    trainer = pl.Trainer(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>`_.
@@ -598,7 +598,22 @@ nb_sanity_val_steps:

 num_tpu_cores
 ^^^^^^^^^^^^^
-How many TPU cores to train on (1 or 8).
+.. warning:: .. deprecated:: 0.7.6
+
+    Use `tpu_cores` instead. Will remove 0.9.0.
+
+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
+
+tpu_cores
+^^^^^^^^^
+- How many TPU cores to train on (1 or 8).
+- Which TPU core to train on [1-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
@@ -615,21 +630,21 @@ Example::
    # your_trainer_file.py

    # default used by the Trainer (ie: train on CPU)
-    trainer = Trainer(num_tpu_cores=None)
+    trainer = Trainer(tpu_cores=None)

    # int: train on a single core
-    trainer = Trainer(num_tpu_cores=1)
+    trainer = Trainer(tpu_cores=1)
+
+    # list: train on a single selected core
+    trainer = Trainer(tpu_cores=[2])

    # int: train on all cores few cores
-    trainer = Trainer(num_tpu_cores=8)
+    trainer = Trainer(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)
+    trainer = Trainer(tpu_cores=8)

 To train on more than 8 cores (ie: a POD),
 submit this script using the xla_dist script.
@@ -135,3 +135,9 @@ class TrainerDeprecatedAPITillVer0_9(ABC):
        rank_zero_warn("`training_tqdm_dict` was renamed to `progress_bar_dict` in v0.7.3"
                       " and this method will be removed in v0.9.0", DeprecationWarning)
        return self.progress_bar_dict
+
+    @property
+    def num_tpu_cores(self):
+        """Back compatibility, will be removed in v0.9.0"""
+        rank_zero_warn("Argument `num_tpu_cores` is now set by `tpu_cores` since v0.7.6"
+                       " and this argument will be removed in v0.9.0", DeprecationWarning)
@@ -389,7 +389,6 @@ class TrainerDPMixin(ABC):
    root_gpu: ...
    amp_level: str
    precision: ...
-    current_tpu_idx: ...
    proc_rank: int
    tpu_local_core_rank: int
    tpu_global_core_rank: int
@@ -398,6 +397,7 @@ class TrainerDPMixin(ABC):
    data_parallel_device_ids: ...
    logger: Union[LightningLoggerBase, bool]
    progress_bar_callback: ...
+    tpu_id: int

    @property
    @abstractmethod
@@ -442,7 +442,8 @@ class TrainerDPMixin(ABC):
        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())
+                xla_device = xm.xla_device(self.tpu_id) if self.tpu_id is not None else xm.xla_device()
+                return batch.to(xla_device)

        if device == 'gpu':
            # base case: object can be directly moved using `cuda` or `to`
@@ -501,7 +502,8 @@ class TrainerDPMixin(ABC):

    def tpu_train(self, tpu_core_idx, model):
        # put model on tpu
-        model.to(xm.xla_device())
+        self._device = xm.xla_device(self.tpu_id) if self.tpu_id is not None else xm.xla_device()
+        model.to(self._device)

        # get the appropriate tpu ranks
        self.tpu_local_core_rank = xm.get_local_ordinal()
@@ -511,8 +513,6 @@ class TrainerDPMixin(ABC):
        if self.tpu_global_core_rank != 0 and self.progress_bar_callback is not None:
            self.progress_bar_callback.disable()

-        # track current tpu
-        self.current_tpu_idx = tpu_core_idx
        self.proc_rank = self.tpu_local_core_rank
        rank_zero_only.rank = self.proc_rank

@@ -174,6 +174,7 @@ class TrainerEvaluationLoopMixin(ABC):
    val_dataloaders: DataLoader
    use_tpu: bool
    reload_dataloaders_every_epoch: ...
+    tpu_id: int

    # Callback system
    on_validation_batch_start: Callable
@@ -248,7 +249,7 @@ class TrainerEvaluationLoopMixin(ABC):
            dl_outputs = []

            # on TPU we have to wrap it under the ParallelLoader
-            if self.use_tpu:
+            if self.use_tpu and self.tpu_id is None:
                device = xm.xla_device()
                dataloader = xla_pl.ParallelLoader(dataloader, [device])
                dataloader = dataloader.per_device_loader(device)
@@ -35,7 +35,6 @@ from pytorch_lightning.trainer.lr_finder import TrainerLRFinderMixin
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities import rank_zero_warn, parsing

-
 try:
    from apex import amp
 except ImportError:
@@ -82,7 +81,7 @@ class Trainer(
        'gradient_clip', 'nb_gpu_nodes', 'max_nb_epochs', 'min_nb_epochs',
        'add_row_log_interval', 'nb_sanity_val_steps', 'tng_tqdm_dic',
    )
-    DEPRECATED_IN_0_9 = ('use_amp', 'show_progress_bar', 'training_tqdm_dict')
+    DEPRECATED_IN_0_9 = ('use_amp', 'show_progress_bar', 'training_tqdm_dict', 'num_tpu_cores')

    def __init__(
            self,
@@ -97,7 +96,7 @@ class Trainer(
            num_processes: int = 1,
            gpus: Optional[Union[List[int], str, int]] = None,
            auto_select_gpus: bool = False,
-            num_tpu_cores: Optional[int] = None,
+            tpu_cores: Optional[Union[List[int], int]] = None,
            log_gpu_memory: Optional[str] = None,
            progress_bar_refresh_rate: int = 1,
            overfit_pct: float = 0.0,
@@ -133,6 +132,7 @@ class Trainer(
            progress_bar_callback: Optional[Union[ProgressBarBase, bool]] = True,
            terminate_on_nan: bool = False,
            auto_scale_batch_size: Union[str, bool] = False,
+            num_tpu_cores: Optional[int] = None,  # backward compatible, todo: remove in v0.9.0
            amp_level: str = 'O1',  # backward compatible, todo: remove in v0.8.0
            default_save_path=None,  # backward compatible, todo: remove in v0.8.0
            gradient_clip=None,  # backward compatible, todo: remove in v0.8.0
@@ -188,7 +188,10 @@ class Trainer(
                GPUs are configured to be in "exclusive mode", such
                that only one process at a time can access them.

-            num_tpu_cores: How many TPU cores to train on (1 or 8).
+            tpu_cores: How many TPU cores to train on (1 or 8) / Single TPU to train on [1]
+
+            num_tpu_cores: How many TPU cores to train on (1 or 8)
+                .. warning:: .. deprecated:: 0.7.6. Will remove 0.9.0.

            log_gpu_memory: None, 'min_max', 'all'. Might slow performance

@@ -342,9 +345,19 @@ class Trainer(
        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'
+        if num_tpu_cores is not None:
+            rank_zero_warn("Argument `num_tpu_cores` is now set by `tpu_cores` since v0.7.6"
+                           " and this argument will be removed in v0.9.0", DeprecationWarning)
+
+        if tpu_cores is None:
+            tpu_cores = num_tpu_cores
+        self.on_tpu = tpu_cores is not None
+        self.tpu_cores = tpu_cores
+        assert self.tpu_cores in (1, 8, None) or (
+            isinstance(self.tpu_cores, (list, tuple, set)) and len(self.tpu_cores) == 1
+        ), '`tpu_cores` can only be 1, 8 or [<1-8>]'
+
+        self.tpu_id = tpu_cores[0] if isinstance(tpu_cores, list) else None

        if num_processes != 1 and distributed_backend != "ddp_cpu":
            rank_zero_warn("num_processes is only used for distributed_backend=\"ddp_cpu\". Ignoring it.")
@@ -477,7 +490,6 @@ class Trainer(
        # 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
@@ -858,7 +870,7 @@ class Trainer(
            self.single_gpu_train(model)

        elif self.use_tpu:  # pragma: no-cover
-            log.info(f'training on {self.num_tpu_cores} TPU cores')
+            log.info(f'training on {self.tpu_cores} TPU cores')

            #  COLAB_GPU is an env var available by default in Colab environments.
            start_method = 'fork' if self.on_colab_kaggle else 'spawn'
@@ -867,7 +879,10 @@ class Trainer(
            self.model = model

            # train
-            xmp.spawn(self.tpu_train, args=(model,), nprocs=self.num_tpu_cores, start_method=start_method)
+            if self.tpu_id is not None:
+                self.tpu_train(self.tpu_id, model)
+            else:
+                xmp.spawn(self.tpu_train, args=(model,), nprocs=self.tpu_cores, start_method=start_method)

            # load weights if not interrupted
            self.load_spawn_weights(model)
@@ -231,6 +231,7 @@ class TrainerTrainLoopMixin(ABC):
    total_batch_idx: int
    checkpoint_callback: ...
    terminate_on_nan: bool
+    tpu_id: int

    # Callback system
    callbacks: List[Callback]
@@ -393,7 +394,7 @@ class TrainerTrainLoopMixin(ABC):
        train_dataloader = self.train_dataloader

        # on TPU we have to wrap it under the ParallelLoader
-        if self.use_tpu:
+        if self.use_tpu and self.tpu_id is None:
            device = xm.xla_device()
            train_dataloader = xla_pl.ParallelLoader(train_dataloader, [device])
            train_dataloader = train_dataloader.per_device_loader(device)
@@ -96,6 +96,9 @@ def test_tbd_remove_in_v0_9_0_trainer():
        trainer = Trainer(progress_bar_refresh_rate=50, show_progress_bar=False)
    assert getattr(trainer, 'show_progress_bar')

+    with pytest.deprecated_call(match='v0.9.0'):
+        _ = Trainer(num_tpu_cores=8)
+

 def test_tbd_remove_in_v0_9_0_module_imports():
    _soft_unimport_module("pytorch_lightning.core.decorators")