Repository: Lightning-AI/pytorch-lightning

Issue:
`num_tpu_cores=8` does not work on kaggle
## 🐛 Bug

When I try to train a model on Kaggle TPU's with `num_tpu_cores` set to 8, I receive an error `Exception: process 2 terminated with exit code 1` . Would be great if this worked on kaggle.

### To Reproduce

Steps to reproduce the behavior:
1. Run this notebook: 
https://www.kaggle.com/lezwon/pytorch-on-tpu-with-pytorch-lightning

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
```
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-9-9251330963d1> in <module>
      3 # most basic trainer, uses good defaults (1 TPU)
      4 trainer = pl.Trainer(num_tpu_cores=8)
----> 5 trainer.fit(mnist_model)

/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, test_dataloaders)
    714 
    715             # train
--> 716             xmp.spawn(self.tpu_train, args=(model,), nprocs=self.num_tpu_cores, start_method=start_method)
    717 
    718             # load weights if not interrupted

/opt/conda/lib/python3.6/site-packages/torch_xla/distributed/xla_multiprocessing.py in spawn(fn, args, nprocs, join, daemon, start_method)
    180         join=join,
    181         daemon=daemon,
--> 182         start_method=start_method)

/opt/conda/lib/python3.6/site-packages/torch/multiprocessing/spawn.py in start_processes(fn, args, nprocs, join, daemon, start_method)
    156 
    157     # Loop on join until it returns True or raises an exception.
--> 158     while not context.join():
    159         pass
    160 

/opt/conda/lib/python3.6/site-packages/torch/multiprocessing/spawn.py in join(self, timeout)
    111                 raise Exception(
    112                     "process %d terminated with exit code %d" %
--> 113                     (error_index, exitcode)
    114                 )
    115 

Exception: process 3 terminated with exit code 1
```

#### Code sample
`trainer = pl.Trainer(num_tpu_cores=8, precision=16) `

### Expected behavior
Run the model utilizing all 8 TPU cores.

### Environment
```
cuda:
	GPU:
	available:           False
	version:             None
packages:
	numpy:               1.18.2
	pyTorch_debug:       False
	pyTorch_version:     1.6.0a0+30e7055
	pytorch-lightning:   0.7.3
	tensorboard:         2.1.1
	tqdm:                4.42.0
system:
	OS:                  Linux
	architecture:
		64bit
		
	processor:           
	python:              3.6.6
	version:             #1 SMP Sat Apr 4 00:12:45 PDT 2020
```





Discussion:
I think this is a kaggle problem?
@dlibenzi any ideas?
It prolly needs this on top:
```
!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
```

those lines are already at the top:
https://www.kaggle.com/pytorchlightning/pytorch-on-tpu-with-pytorch-lightning

![image](https://user-images.githubusercontent.com/3640001/79809189-eb2c1580-833d-11ea-80d2-4954e5ffca0d.png)

I bet the issue is here:

https://github.com/PyTorchLightning/pytorch-lightning/blob/bd168819f2e6433a8b0f0d497eda785199cf65cf/pytorch_lightning/trainer/trainer.py#L762

ah... yes. good catch.
know of something more general that we can check? i assume the only two options are kaggle and colab?
@lezwon want to find an environment variable we can check to know if on kaggle and submit a PR?
Honestly, pytorch does not like `fork` because of CUDA, but I would make that the default, with ability to change via some environment variable in cases someone have issues.

on GCP it would still be fork? 
when would it not be fork with TPUs?
Fork is an issue with pytorch/CUDA mostly.
But for safety, I would just add a Kaggle check as well in your code, and leave `spawn` as default.
Fork also helps Colab and Kaggle because, being them low memory VMs, one can reduce the memory consumption by creating the model (on default pytorch/cpu) at global scope, and then doing to(xla_device) from within the xmp.spawn() target functions.
This avoids creating pytorch/cpu models in each of the processes (one per core).
You can see a few tricks to fit models on Colab here:

https://colab.research.google.com/drive/1IvCxIg-Q_DlI7UNJuajpl4UZXNiW5jMg

Like create model at global scope, and serialize the to(xla_device) calls to avoid all 8 processes rushing into allocation host memory at the same time.




Resolving patch:
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -755,7 +755,7 @@ def fit(
             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'
+            start_method = 'fork' if os.getenv('COLAB_GPU') or os.getenv('KAGGLE_URL_BASE') else 'spawn'
 
             # track for predict
             self.model = model


Test patch: