global process count incorrect with elastic, fault tolerant training

[srib]

🐛 Bug

Problem

Count of the total number of processes incorrectly set.

Context

I am trying to run elastic training with torchelastic. I have tried with both gloo and nccl backends.

Error message

Error coming from gloo backend:

Traceback (most recent call last):
  File "train_hydra.py", line 20, in hydra_main
    train(cfg)
  File "/bdata/bdata1/sribkain/learnseis/learnseis/training.py", line 39, in train
    t.fit(module, data_module)
  File "/ldata/Code/salt-identification/SRIBKAIN_ENVS/pl_env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 496, in fit
    self.pre_dispatch()
  File "/ldata/Code/salt-identification/SRIBKAIN_ENVS/pl_env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 525, in pre_dispatch
    self.accelerator.pre_dispatch()
  File "/ldata/Code/salt-identification/SRIBKAIN_ENVS/pl_env/lib/python3.7/site-packages/pytorch_lightning/accelerators/accelerator.py", line 83, in pre_dispatch
    self.training_type_plugin.pre_dispatch()
  File "/ldata/Code/salt-identification/SRIBKAIN_ENVS/pl_env/lib/python3.7/site-packages/pytorch_lightning/plugins/training_type/ddp.py", line 243, in pre_dispatch
    self.init_ddp_connection(self.global_rank, self.world_size)
  File "/ldata/Code/salt-identification/SRIBKAIN_ENVS/pl_env/lib/python3.7/site-packages/pytorch_lightning/plugins/training_type/ddp.py", line 226, in init_ddp_connection
    torch_distrib.init_process_group(self.torch_distributed_backend, rank=global_rank, world_size=world_size)
  File "/ldata/Code/salt-identification/SRIBKAIN_ENVS/pl_env/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 432, in init_process_group
    timeout=timeout)
  File "/ldata/Code/salt-identification/SRIBKAIN_ENVS/pl_env/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 503, in _new_process_group_helper
    timeout=timeout)
RuntimeError: [enforce fail at /pytorch/third_party/gloo/gloo/context.cc:27] rank < size. 13 vs 8

NCCL backend gives this error: pytorch/pytorch#20313

Please reproduce using the BoringModel

I am running imagenet example from pl using torchvision.models.resnet34. Happy to reproduce with BoringModel if needed.

Before launching, I have exported the variable GLOO_SOCKET_IFNAME and set it to the appropriate interface name.

On node 0:

PL_TORCH_DISTRIBUTED_BACKEND=gloo python -m torchelastic.distributed.launch --nnodes=1:5 --rdzv_id='nodockertestelasticlaunch7' --rdzv_backend=etcd --rdzv_endpoint=10.18.0.15:2379 train_hydra.py +experiment=elastic_config.yaml

On node 1:

PL_TORCH_DISTRIBUTED_BACKEND=gloo python -m torchelastic.distributed.launch --nnodes=1:5 --rdzv_id='nodockertestelasticlaunch7' --rdzv_backend=etcd --rdzv_endpoint=10.18.0.15:2379 train_hydra.py +experiment=elastic_config.yaml

To Reproduce

Use following BoringModel and post here

Expected behavior

To be able to run distributed fault tolerant training :)

Environment

Note: Bugs with code are solved faster ! Colab Notebook should be made public !

• IDE: Please, use our python bug_report_model.py template.

• Colab Notebook: Please copy and paste the output from our environment collection script (or fill out the checklist below manually).

Output of collect_env_details.py:

* CUDA:
        - GPU:
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
                - GeForce RTX 2080 Ti
        - available:         True
        - version:           10.2
* Packages:
        - numpy:             1.19.2
        - pyTorch_debug:     False
        - pyTorch_version:   1.6.0
        - pytorch-lightning: 1.2.6
        - tqdm:              4.48.2
        - torchelastic:    0.2.0
* System:
        - OS:                Linux
        - architecture:
                - 64bit
                -
        - processor:         x86_64
        - python:            3.7.7
        - version:           #88-Ubuntu SMP Tue Feb 11 20:11:34 UTC 2020

Additional context

[srib]

Related issues: #6797, #6527

[srib]

Shouldn't torch_distrib.init_process_group use world_size from os.environ['WORLD_SIZE'] instead of world_size?

    def init_ddp_connection(self, global_rank: int, world_size: int) -> None:
        os.environ["MASTER_ADDR"] = str(self.cluster_environment.master_address())
        os.environ["MASTER_PORT"] = str(self.cluster_environment.master_port())
        os.environ["WORLD_SIZE"] = str(self.cluster_environment.world_size())

        if not torch.distributed.is_initialized():
            log.info(f"initializing ddp: GLOBAL_RANK: {global_rank}, MEMBER: {global_rank + 1}/{world_size}")
            torch_distrib.init_process_group(self.torch_distributed_backend, rank=global_rank, world_size=world_size)

From: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/plugins/training_type/ddp.py#L234

[ananthsub]

@srib what does your os.environ look like before init_ddp_connection is called? could you print this out? Lightning currently doesn't handle the case where min_nodes != max_nodes in torchelastic. The trainer needs to listen for when the world size changes and rebalance accordingly, which currently isn't implemented.

[srib]

@ananthsub Thank you for the reply.

All the following cases were run with the following command:

PL_TORCH_DISTRIBUTED_BACKEND=gloo python -m torchelastic.distributed.launch --nnodes=1:5 --rdzv_id='nodockertestelasticlaunch7' --rdzv_backend=etcd --rdzv_endpoint=10.18.0.15:2379 train_hydra.py +experiment=elastic_config.yaml

It seems to work fine when min_nodes != max_nodes. I ran the following tests and here is what I have to report.

1. On a single node with 8 GPUs, here is what my os.environ looks like before init_ddp_connection is called:

world_size: 8
os.environ["WORLD_SIZE"]: 8

2. On two nodes with 8 GPUs each:

world_size: 8
os.environ["WORLD_SIZE"]: 16

3. Interrupt one node:

world_size: 8
os.environ["WORLD_SIZE"]: 8

4. Add a new node again:

world_size: 8
os.environ["WORLD_SIZE"]: 16

I even got training to run with the following change:

Before:

torch_distrib.init_process_group(self.torch_distributed_backend, rank=global_rank, world_size=world_size)

After:

torch_distrib.init_process_group(self.torch_distributed_backend, rank=global_rank, world_size=int(os.environ["WORLD_SIZE"]))

Am I missing something here?

[srib]

@ananthsub Can I submit a PR for this or are there any other implications that I am missing here?

[ananthsub]

@awaelchli what do you think of making global_rank a method on the ClusterEnvironment. world_size is already there, and we know whether the cluster environment spawns processes based on creates_children.

Given this state is already available from the env variables, this could be used to set the rank_zero_only.rank on trainer init as well, and we can reset this in the DDP/parallel plugin whenever we do spawn. We can add a setter for world size and global rank as well. For Slurm and Torchelastic, the setter should be a no-op: we can always set the fields to be whatever the environment is and ignore the passed in argument

Then the cluster environment becomes the source of truth for these fields, which can be used in the training type plugin, which is used in the accelerator, which is used in the trainer.

[awaelchli]

Sounds good yes. Let's try it. This will affect all major plugins, so we need to be careful :)

[awaelchli]

@srib ok I think I understand now the bug you are experiencing. The world size is changing in the middle of training, and we are keeping the old value saved. And so we should read the world size always from the cluster environment (which will read it from the environment)

[awaelchli]

Found a related bug in master:

Running
python pl_examples/basic_examples/simple_image_classifier.py --accelerator ddp --gpus 2
then on rank 0 we have LightningEnvironment but on rank 1 it is TorchElasticEnvironment
arghh

[ananthsub]

@awaelchli is it because the lightning environment uses the same environment variables as torchelastic for local rank?