Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Tpu logging #2230

Merged
merged 5 commits into from
Jun 18, 2020
Merged

Tpu logging #2230

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
48948e4
add tpu view
williamFalcon Jun 18, 2020
c040095
add tpu view
williamFalcon Jun 18, 2020
2f3b1c2
add tpu view
williamFalcon Jun 18, 2020
ea529e4
add tpu view
williamFalcon Jun 18, 2020
e28bfb9
add tpu view
williamFalcon Jun 18, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion pytorch_lightning/trainer/distrib_data_parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,16 @@ def train_fx(trial_hparams, cluster_manager, _):
HYDRA_AVAILABLE = True


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


class TrainerDDPMixin(ABC):

# this is just a summary on variables used in this abstract class,
Expand All @@ -172,6 +182,7 @@ class TrainerDDPMixin(ABC):
num_processes: int
num_nodes: int
node_rank: int
tpu_cores: int

@property
def is_global_zero(self) -> int:
Expand Down Expand Up @@ -277,6 +288,8 @@ def set_distributed_mode(self, distributed_backend):
)

rank_zero_info(f'GPU available: {torch.cuda.is_available()}, used: {self.on_gpu}')
num_cores = self.tpu_cores if self.tpu_cores is not None else 0
rank_zero_info(f'TPU available: {XLA_AVAILABLE}, using: {num_cores} TPU cores')

def configure_slurm_ddp(self, num_gpu_nodes):
self.is_slurm_managing_tasks = False
Expand Down Expand Up @@ -329,7 +342,6 @@ def determine_ddp_node_rank(self):
node_ids = [(k, os.environ.get(k, None)) for k in env_vars]
node_ids = [(k, v) for k, v in node_ids if v is not None]
if len(node_ids) == 0:
log.warning("No environment variable for node rank defined. Set as 0.")
return 0
if len(node_ids) > 1:
log.warning(f"Multiple environment variables ({node_ids}) defined for node rank. "
Expand Down