Model trains on 1 node 8xA100 but hits OOM in 4 nodes 8xA100
Author: aqred1
Pull request: no

Hello,

I am attempting to train a model on 32 A100s split into 4 nodes. The queueing system is SLURM. When training on a single node, I can use all 8 GPUs with a batch size of 8:

`python -m torch.distributed.launch --nproc_per_node=8  ./script.py  --model_name_or_path gpt2-large --train_file sample.txt --tokenizer_name embeddings --do_train --do_eval --output_dir output/ --evaluation_strategy steps --eval_steps 5000 --num_train_epochs 120 --per_device_train_batch_size 8 --cache_dir .cache2/ --bf16 --gradient_checkpointing True --save_total_limit 2 --learning_rate 1e-05`

But when switching to multi-node, it hits OOM even with a batch size of 1. Here's the script:

```
#SBATCH --time=23:59:00
#SBATCH --gres=gpu:a100:8
#SBATCH --export=NONE
#SBATCH --partition=a100
#SBATCH --nodes=4

export MASTER_PORT=12340

echo "NODELIST="${SLURM_NODELIST}
master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_ADDR=$master_addr

export RANK_NODE=$SLURM_NODEID
SLAVES=`scontrol show hostnames $SLURM_JOB_NODELIST | grep -v $MASTER_ADDR`
HOSTLIST="$MASTER_ADDR $SLAVES"

unset SLURM_EXPORT_ENV

source /path/to/py37/bin/activate


RANK=0
for node in $HOSTLIST; do
	ssh -q $node; 
	python -m torch.distributed.launch --nproc_per_node=8 --nnodes=4 --node_rank=$RANK --master_addr=$MASTER_ADDR   
        --master_port=$MASTER_PORT ./5.run_clm-post.py  --model_name_or_path gpt2-large --train_file sample.txt
        --tokenizer_name embeddings --do_train --do_eval --output_dir output/ --evaluation_strategy steps --eval_steps 5000
        --num_train_epochs 120 --per_device_train_batch_size 1 --cache_dir .cache2/ --bf16 --gradient_checkpointing True
        --save_total_limit 2 --learning_rate 1e-05 & RANK=$((RANK+1));
done
wait
```
Funnily enough, I get the CUDA OOM error 32 times, instead of 4, so I guess whether it's treating each GPU independently.

Note: I had to switch the backend to `gloo` because got nccl errors. 

I will try now with deepspeed + bf16,  I've been working on this setup for a couple of days and don't seem to get anywhere.
Any help is appreciated.

Best




Comments:
This may be of interest to @stas00 as I believe he has quite a bit of experience with this kind of setups.

I haven't encountered this specific issue, so first let's try to understand what is going on here.

> Note: I had to switch the backend to gloo because got nccl errors.

and did you test that one node was still working OK with `gloo` after switching to it?

What nccl errors were you getting?

> Funnily enough, I get the CUDA OOM error 32 times, instead of 4

each process gets OOM in your case, why would you expect 4 errors and not 32?

------------

unrelated, have a look at how we launch similar jobs on BigScience:
https://github.com/bigscience-workshop/bigscience/blob/master/train/tr8-104B-wide/tr8-104B.slurm

you can let `srun` handle the ssh'ing for you. But perhaps your SLURM env is slightly different. Ours is JeanZay HPC.

it probably doesn't matter, just thought I'd share how I launch these.

Thanks for the prompt response and the `srun` example, I really appreciate it! I tried with `srun` but it didn't work, so I had to ssh. But the cluster I am using is still under beta version, so I will get in contact with them and try again. 

But first. I am now directly logging into the nodes to rule out SLURM errors. 
Using `nccl`, the command above runs on a single node. But when running on multiple nodes, it hangs indefinitely. It stops here: 

```
  FutureWarning,
WARNING:torch.distributed.run:
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
*****************************************
```

And after some time, I get the following error:

```
  File "/path/to/py37/lib/python3.7/site-packages/torch/distributed/elastic/rendezvous/static_tcp_rendezvous.py", line 61, in next_rendezvous
    multi_tenant=True,
RuntimeError: connect() timed out. Original timeout was 900000 ms.
```
When I send it through SLURM, I get a RuntimeError as well. This is why I switched to `gloo`...

Sounds like a firewall issue then. The slave nodes try to connect to the master node and fail to.

You can install `py-spy` and run:

```
py-spy dump --pid PID
```

and it will dump the stack trace of the hanging process PID (on some systems `sudo` is required)

but it's almost certain you have a firewall issue there.

I have the same issues if I try to connect to an external world from a gpu node, e.g. to fetch `transformers` or `datasets` files.

So we basically devised this solution for the latter:

1. download all that is needed and cache by running the code on a node with internet access 
2. then run the code on a node after setting:

```
export HF_DATASETS_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
```

you will most certainly want this for your normal work.

but your issue is inter-node firewall.

Can you ssh from one gpu worker node to another? This is most likely where your problem is. If you can ssh then the port elastic is trying to sync on is firewalled. Ask your sysadmins which port is safe to use and set it explicitly in your launcher command line.

> Thanks for the prompt response and the srun example, I really appreciate it! I tried with srun but it didn't work, so I had to ssh. > But the cluster I am using is still under beta version, so I will get in contact with them and try again.

It'll make things much simpler if you could make `srun` to work. That's the whole point of SLURM - those things are designed to remove the complexity.

Here is a simple test script for you to work with rather than a big application.
```

# skeleton script for .launch
# test.py
import torch.distributed as dist
import argparse
import torch
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)

dist.init_process_group("nccl")
dist.all_reduce(torch.ones(1).to(device), op=dist.ReduceOp.SUM)
dist.barrier()


# to run
python -m torch.distributed.launch --nproc_per_node=2 test.py
```

adjust `--nproc_per_node` to your setup. and if needed bolt `srun` to it. otherwise this is just for a single node as is.

it doesn't print anything, we just want it not to fail ;)

Here is the same for the new `torch.distributed.run` API instead, which requires a slightly different setup

```
# skeleton script for .run
import torch.distributed as dist
import torch
import os
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
dist.init_process_group("nccl")
dist.barrier()
dist.get_world_size()
dist.is_available()
dist.get_rank()


# to run
python -m torch.distributed.run --nproc_per_node=2 test.py
```

but `launch` should be just fine.

Thanks a lot for the great response!

I installed `py-spy` and had a look at the stack trace of the command, but I am not sure if I am gaining much information. Here is the output.
```

Thread 0x1544E69E5740 (active+gil): "MainThread"
    next_rendezvous (torch/distributed/elastic/rendezvous/static_tcp_rendezvous.py:61)
    _rendezvous (torch/distributed/elastic/agent/server/api.py:538)
    wrapper (torch/distributed/elastic/metrics/api.py:125)
    _initialize_workers (torch/distributed/elastic/agent/server/api.py:678)
    wrapper (torch/distributed/elastic/metrics/api.py:125)
    _invoke_run (torch/distributed/elastic/agent/server/api.py:837)
    run (torch/distributed/elastic/agent/server/api.py:709)
    wrapper (torch/distributed/elastic/metrics/api.py:125)
    launch_agent (torch/distributed/launcher/api.py:252)
    __call__ (torch/distributed/launcher/api.py:131)
    run (torch/distributed/run.py:713)
    launch (torch/distributed/launch.py:174)
    main (torch/distributed/launch.py:189)
    <module> (torch/distributed/launch.py:193)
    _run_code (runpy.py:85)
    _run_module_as_main (runpy.py:193)
```
It must then be a firewall issue. I already downloaded the dataset and transformer on a node with access and cached everything for the `slurm` nodes. I didn't know I had to export those variables, thank you! I've asked the system administrators about a safe port to use. I'm however able to run deepspeed on several nodes (and benchmarking that at the moment); perhaps because pdsh is able to find a port?

Thanks for the `srun` example, it will come in handy!




Yes,  `HF_DATASETS_OFFLINE=1 TRANSFORMERS_OFFLINE=1` simply allows you to continue using the online specification for the models and datasets instead of the local paths, as it instructs the software to use the cache and not look anything up, but your way is just as fine.

The `deepspeed` launcher actually doesn't work on JeanZay, that's why I'm using `srun` + pytorch launcher. Perhaps your situation is reversed.

BTW,  the `deepspeed` launcher can be used as a drop in replacement for the pytorch launcher - you don't have to use Deepspeed in the rest of your code. And it has quite a few goodies in it. So yes, perhaps `pdsh` solves the issue. I have never used it, other than being aware of it, so can't add further commentary on it.

And additionally the new `elastic` launcher that replaced the old one in pt-1.9 also has a variety of connection methods, you might want to experiment with those as well. https://pytorch.org/docs/stable/distributed.elastic.html

Thanks for your response.

I'm going to switch to `torchrun` and give it a try at the connection options in your link.

Since this seems a very particular issue of my own, I'll wait to see what the systems administrators report and write here the solution once I get it running so that maybe it's helpful to someone - then after that close the issue. But of course feel free to close already, since this does not seem to be an issue _per se_.

But we also want to check that the OOM issue goes away once nccl is restored. Unless you have already validated it to be so by using the `deepspeed` launcher, then yes, it'd be safe to close.

and yes sharing your solution would definitely be helpful to others, @aqred1

It was indeed a problem of ports not being opened. In the end I'm using deepspeed so I haven't had the chance to encounter the OOM issue again. Thanks for your great help, I'm closing the issue.