FSDP + TP Duplicate GPU detected : rank 0 and rank 1 both on CUDA device 1b000

pzanders: <p>Hello,</p>
<p>I am trying to do the FSDP + TP parallelism with a simple Autoencoder. When I do it on a single node with multiple GPUs it works fine (device_mesh = init_device_mesh(“cuda”, (2, 2), mesh_dim_names=(“dp”, “tp”))), but when i scale it to multiple nodes with multiple GPUs I get the following error. I have 4 GPUs for each node.</p>
<p>[rank0]: Traceback (most recent call last):<br>
[rank0]:   File “/gpfs/scratch/bsc23/MN4/bsc23/bsc23645/AE_tp/tp_test.py”, line 111, in <br>
[rank0]:     output = sharded_model(inp)<br>
[rank0]:              ^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/gpfs/apps/MN5/ACC/PYTORCH/2.4.0/torch/nn/modules/module.py”, line 1532, in _wrapped_call_impl<br>
[rank0]:     return self._call_impl(*args, **kwargs)<br>
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/gpfs/apps/MN5/ACC/PYTORCH/2.4.0/torch/nn/modules/module.py”, line 1541, in _call_impl<br>
[rank0]:     return forward_call(*args, **kwargs)<br>
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/gpfs/apps/MN5/ACC/PYTORCH/2.4.0/torch/distributed/fsdp/fully_sharded_data_parallel.py”, line 823, in forward<br>
[rank0]:     args, kwargs = _pre_forward(<br>
[rank0]:                    ^^^^^^^^^^^^^<br>
[rank0]:   File “/gpfs/apps/MN5/ACC/PYTORCH/2.4.0/torch/distributed/fsdp/_runtime_utils.py”, line 380, in _pre_forward<br>
[rank0]:     unshard_fn(state, handle)<br>
[rank0]:   File “/gpfs/apps/MN5/ACC/PYTORCH/2.4.0/torch/distributed/fsdp/_runtime_utils.py”, line 415, in _pre_forward_unshard<br>
[rank0]:     _unshard(state, handle, state._unshard_stream, state._pre_unshard_stream)<br>
[rank0]:   File “/gpfs/apps/MN5/ACC/PYTORCH/2.4.0/torch/distributed/fsdp/_runtime_utils.py”, line 299, in _unshard<br>
[rank0]:     handle.unshard()<br>
[rank0]:   File “/gpfs/apps/MN5/ACC/PYTORCH/2.4.0/torch/distributed/fsdp/_flat_param.py”, line 1308, in unshard<br>
[rank0]:     padded_unsharded_flat_param = self._all_gather_flat_param(unsharded_flat_param)<br>
[rank0]:                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/gpfs/apps/MN5/ACC/PYTORCH/2.4.0/torch/distributed/fsdp/_flat_param.py”, line 1399, in _all_gather_flat_param<br>
[rank0]:     dist.all_gather_into_tensor(<br>
[rank0]:   File “/gpfs/apps/MN5/ACC/PYTORCH/2.4.0/torch/distributed/c10d_logger.py”, line 78, in wrapper<br>
[rank0]:     return func(*args, **kwargs)<br>
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]:   File “/gpfs/apps/MN5/ACC/PYTORCH/2.4.0/torch/distributed/distributed_c10d.py”, line 2971, in all_gather_into_tensor<br>
[rank0]:     work = group._allgather_base(output_tensor, input_tensor, opts)<br>
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
[rank0]: torch.distributed.DistBackendError: NCCL error in: /gpfs/apps/MN5/ACC/PYTORCH/2.4.0/torch/csrc/distributed/c10d/NCCLUtils.hpp:275, invalid usage (run with NCCL_DEBUG=WARN for details), NCCL version 2.20.5<br>
[rank0]: ncclInvalidUsage: This usually reflects invalid usage of NCCL library.<br>
[rank0]: Last error:<br>
[rank0]: Duplicate GPU detected : rank 0 and rank 1 both on CUDA device 1b000</p>
<pre><code class="lang-auto">logger = get_logger()
world_size = int(os.environ["WORLD_SIZE"])
rank = int(os.environ["RANK"])
local_rank = int(os.environ["LOCAL_RANK"])

# Device Mesh
device_mesh = init_device_mesh("cuda", (4, 4), mesh_dim_names=("dp", "tp"))

tp_mesh = device_mesh["tp"]
dp_mesh = device_mesh["dp"]
dp_rank = dp_mesh.get_local_rank()

# Prepare Dataset
dataset = np.load("../ML/Simulations/NPY/WTx2_apo_CA_ChainsA_DRY_imagedFit.npy")
dataset = np.reshape(dataset, (len(dataset), -1))
num_frames, n_features = dataset.shape

sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True, seed=42, drop_last=False)
dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, drop_last=True)

model = AutoEncoder2(input_dim=n_features, nlayers=6, latent_dim=3)

tp_model = parallelize_module(
    model,
    tp_mesh,
    parallelize_plan=parallel_plan_ae_nl6)

# FSDP
sharded_model = FSDP(tp_model, device_mesh=dp_mesh, use_orig_params=True).to("cuda")

# Create an optimizer for the parallelized and sharded model.
lr = 1e-4
rank_log(rank, logger, f"Creating Adam optimizer with learning rate {lr}")
optimizer = torch.optim.AdamW(sharded_model.parameters(), lr=lr, foreach=True)

# Training loop:
num_epochs = 10
model.train()
for epoch in range(num_epochs):
   for i, inp in enumerate(dataloader):    
       torch.manual_seed(i + dp_rank)
       inp = inp.to(torch.cuda.current_device())
  
       output = sharded_model(inp)
       loss = torch.nn.functional.mse_loss(output, inp)
       train_losses.append(loss.item())

       optimizer.zero_grad()
       loss.sum().backward()
       optimizer.step()
   
   train_loss = torch.mean(torch.tensor(train_losses))

</code></pre>
<p>Here is my SLURM job:</p>
<pre><code class="lang-auto">#!/bin/bash

#SBATCH --job-name=multinode-fsdp
#SBATCH -o \%x.out
#SBATCH -e \%x.err
#SBATCH --nodes=4
#SBATCH --ntasks=80
#SBATCH --qos=acc_debug
#SBATCH -t 0-02:00:00
#SBATCH --gres=gpu:4

nodes=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) )
nodes_array=($nodes)
head_node=${nodes_array[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)

echo Node IP: $head_node_ip
export LOGLEVEL=INFO

#srun torchrun
srun torchrun \
--nnodes 4 \
--nproc_per_node 4 \
--rdzv_id $RANDOM \
--rdzv_backend c10d \
--rdzv_endpoint $head_node_ip:29501 \
tp_test.py &amp;&gt; tp-fsdp_output.log

</code></pre>
<p>What am I doing wrong?<br>
I’m new in distributed training, can somebody help?</p>
<p><a href="https://pytorch.org/tutorials/intermediate/TP_tutorial.html#combine-tensor-parallel-with-fully-sharded-data-parallel-together" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/tutorials/intermediate/TP_tutorial.html#combine-tensor-parallel-with-fully-sharded-data-parallel-together</a></p>

agu: <p>Just a wild guess: could you try <code>torch.cuda.set_device(local_rank)</code> before initing the device mesh?</p>