lightning.Fabric meets deadlock when loading nn.Module

[forestbat]

Bug description

When I try to use lightning.Fabric.setup() to load torch.nn.Module under multi-process, the program will meet deadlock and stuck in lightning/fabric/strategies/launchers/subprocess_script.py.

I doubt this problem comes from popen start method of process, but I have not more evidence.

What version are you seeing the problem on?

v2.4, v2.5

How to reproduce the bug

I try to reproduce this bug with smaller demo with my project, but I failed, this bug seems only in my actual project:

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset
import lightning
fabric = lightning.Fabric(devices=[1, 2], num_nodes=1, strategy='ddp')


class MyEvaluator:
    def __init__(self):
        fabric.launch()

    def eval_model(self, dataset, crit):
        model = LinearModel()
        model = fabric.setup_module(module=model)
        model.eval()
        # I have changed type of model to custom class, but still can't reproduce this problem
        # model = Seq2Seq_Min_LSTM_GNN(en_input_size=30, de_input_size=18, output_size=2, hidden_size=256, forecast_history=168, forecast_length=56, graph=dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 2]))))
        test_loader = fabric.setup_dataloaders(DataLoader(dataset, batch_size=10, shuffle=True,
                                                           num_workers=1, multiprocessing_context='spawn'))
        for x, y in test_loader:
            output = model(x)
            loss = crit(output, y)
            yield loss

class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 2)

    def forward(self, x):
        return self.linear(x)


if __name__ == '__main__':
    x = torch.randn(100, 10)
    y = torch.rand(100, 2)
    dataset = TensorDataset(x, y)
    crit = nn.MSELoss()
    evaluator = MyEvaluator()
    for loss in evaluator.eval_model(dataset, crit):
        print(loss)

Error messages and logs

This is no error messages and logs when deadlock occurs. What should I do to know what happened in my program and give you enough messages?

Environment

Current environment

• CUDA:
  - GPU:
  - NVIDIA RTX 5000 Ada Generation
  - NVIDIA A40
  - NVIDIA A40
  - available: True
  - version: 12.1
• Lightning:
  - lightning: 2.4.0
  - lightning-utilities: 0.11.9
  - pytorch-lightning: 2.4.0
  - torch: 2.2.2
  - torchaudio: 2.2.2
  - torchdata: 0.7.1
  - torchmetrics: 1.6.0
  - torchvision: 0.17.2
• Packages:
  - absl-py: 2.1.0
  - affine: 2.4.0
  - aiobotocore: 2.13.2
  - aiodns: 3.2.0
  - aiohappyeyeballs: 2.3.7
  - aiohttp: 3.10.4
  - aiohttp-client-cache: 0.11.1
  - aioitertools: 0.11.0
  - aiosignal: 1.3.1
  - aiosqlite: 0.20.0
  - annotated-types: 0.7.0
  - appdirs: 1.4.4
  - argon2-cffi: 23.1.0
  - argon2-cffi-bindings: 21.2.0
  - asciitree: 0.3.3
  - async-retriever: 0.17.0
  - attrs: 24.2.0
  - autocommand: 2.2.2
  - backports.tarfile: 1.2.0
  - black: 24.8.0
  - bleach: 6.1.0
  - bokeh: 3.5.1
  - boto3: 1.34.131
  - botocore: 1.34.131
  - branca: 0.7.2
  - brotli: 1.1.0
  - bump2version: 1.0.1
  - cachetools: 5.5.0
  - cartopy: 0.23.0
  - cattrs: 23.2.3
  - certifi: 2024.8.30
  - cffi: 1.17.0
  - cfgrib: 0.9.14.0
  - cftime: 1.6.4
  - chardet: 5.2.0
  - charset-normalizer: 3.3.2
  - click: 8.1.7
  - click-plugins: 1.1.1
  - cligj: 0.7.2
  - cloudpickle: 3.0.0
  - codetiming: 1.4.0
  - colorama: 0.4.6
  - contourpy: 1.2.1
  - cryptography: 43.0.0
  - cupy: 13.3.0
  - cycler: 0.12.1
  - cytoolz: 0.12.3
  - dask: 2024.8.1
  - dask-expr: 1.1.11
  - dataretrieval: 1.0.10
  - deepspeed: 0.16.1
  - defusedxml: 0.7.1
  - dgl: 2.2.1+cu121
  - distributed: 2024.8.1
  - docutils: 0.21.2
  - eccodes: 1.7.1
  - einops: 0.8.0
  - et-xmlfile: 1.1.0
  - exceptiongroup: 1.2.2
  - fasteners: 0.19
  - fastrlock: 0.8.2
  - filelock: 3.15.4
  - findlibs: 0.0.5
  - flake8: 7.1.1
  - flexcache: 0.3
  - flexparser: 0.3.1
  - folium: 0.17.0
  - fonttools: 4.53.1
  - frozenlist: 1.4.1
  - fsspec: 2024.6.1
  - geopandas: 1.0.1
  - gmpy2: 2.1.5
  - greenlet: 3.0.3
  - grpcio: 1.62.2
  - h2: 4.1.0
  - h5netcdf: 1.3.0
  - h5py: 3.11.0
  - hjson: 3.1.0
  - hpack: 4.0.0
  - hydrodataset: 0.1.13
  - hydrodatasource: 0.0.8
  - hydroerr: 1.24
  - hydrosignatures: 0.17.0
  - hydrotopo: 0.0.6
  - hydroutils: 0.0.12
  - hyperframe: 6.0.1
  - idna: 3.7
  - igraph: 0.11.6
  - importlib-metadata: 8.2.0
  - importlib-resources: 6.4.0
  - inflect: 7.3.1
  - iniconfig: 2.0.0
  - intake: 2.0.6
  - itsdangerous: 2.2.0
  - jaraco.classes: 3.4.0
  - jaraco.context: 5.3.0
  - jaraco.functools: 4.0.2
  - jaraco.text: 3.12.1
  - jeepney: 0.8.0
  - jinja2: 3.1.4
  - jmespath: 1.0.1
  - joblib: 1.4.2
  - kaggle: 1.6.17
  - kerchunk: 0.2.6
  - keyring: 25.3.0
  - kiwisolver: 1.4.5
  - lightning: 2.4.0
  - lightning-utilities: 0.11.9
  - llvmlite: 0.43.0
  - locket: 1.0.0
  - loguru: 0.7.2
  - lxml: 5.3.0
  - lz4: 4.3.3
  - markdown: 3.6
  - markdown-it-py: 3.0.0
  - markupsafe: 2.1.5
  - matplotlib: 3.9.2
  - mccabe: 0.7.0
  - mdurl: 0.1.2
  - minio: 7.2.8
  - more-itertools: 10.4.0
  - mpmath: 1.3.0
  - msgpack: 1.0.8
  - multidict: 6.0.5
  - mypy-extensions: 1.0.0
  - netcdf4: 1.7.1.post2
  - networkx: 3.3
  - nh3: 0.2.18
  - ninja: 1.11.1.3
  - nuitka: 2.4.7
  - numba: 0.60.0
  - numcodecs: 0.13.0
  - numpy: 1.26.4
  - nvidia-ml-py: 12.535.161
  - nvitop: 1.3.2
  - openpyxl: 3.1.5
  - ordered-set: 4.1.0
  - owslib: 0.31.0
  - packaging: 24.1
  - pandas: 2.2.2
  - partd: 1.4.2
  - pathspec: 0.12.1
  - pillow: 10.4.0
  - pint: 0.24.3
  - pint-pandas: 0.6.2
  - pint-xarray: 0.4
  - pip: 24.2
  - pkginfo: 1.10.0
  - platformdirs: 4.2.2
  - pluggy: 1.5.0
  - polars: 1.17.1
  - protobuf: 4.25.3
  - psutil: 6.0.0
  - psycopg2-binary: 2.9.9
  - py-cpuinfo: 9.0.0
  - pyarrow: 17.0.0
  - pyarrow-hotfix: 0.6
  - pycairo: 1.27.0
  - pycares: 4.4.0
  - pycodestyle: 2.12.1
  - pycparser: 2.22
  - pycryptodome: 3.20.0
  - pydantic: 2.8.2
  - pydantic-core: 2.20.1
  - pyflakes: 3.2.0
  - pygeohydro: 0.17.0
  - pygeoogc: 0.17.0
  - pygeoutils: 0.17.0
  - pygments: 2.18.0
  - pykalman: 0.9.7
  - pynhd: 0.17.0
  - pyogrio: 0.9.0
  - pyparsing: 3.1.2
  - pyproj: 3.6.1
  - pyshp: 2.3.1
  - pysocks: 1.7.1
  - pytest: 8.3.2
  - python-dateutil: 2.9.0
  - python-slugify: 8.0.4
  - pytorch-lightning: 2.4.0
  - pytz: 2024.1
  - pyyaml: 6.0.2
  - rasterio: 1.3.10
  - readme-renderer: 44.0
  - requests: 2.32.3
  - requests-cache: 1.2.1
  - requests-toolbelt: 1.0.0
  - rfc3986: 2.0.0
  - rich: 13.7.1
  - rioxarray: 0.17.0
  - s3fs: 2024.6.1
  - s3transfer: 0.10.2
  - scikit-learn: 1.5.1
  - scipy: 1.14.0
  - seaborn: 0.13.2
  - secretstorage: 3.3.3
  - setuptools: 72.2.0
  - shap: 0.45.1
  - shapely: 2.0.1
  - six: 1.16.0
  - slicer: 0.0.8
  - snuggs: 1.4.7
  - sortedcontainers: 2.4.0
  - sqlalchemy: 2.0.32
  - sympy: 1.13.2
  - tblib: 3.0.0
  - tbparse: 0.0.9
  - tensorboard: 2.17.1
  - tensorboard-data-server: 0.7.0
  - termcolor: 2.5.0
  - text-unidecode: 1.3
  - texttable: 1.7.0
  - threadpoolctl: 3.5.0
  - tomli: 2.0.1
  - toolz: 0.12.1
  - torch: 2.2.2
  - torchaudio: 2.2.2
  - torchdata: 0.7.1
  - torchmetrics: 1.6.0
  - torchvision: 0.17.2
  - tornado: 6.4.1
  - tqdm: 4.66.5
  - triton: 2.2.0
  - twine: 5.1.1
  - typeguard: 4.3.0
  - typing-extensions: 4.12.2
  - tzdata: 2024.1
  - tzfpy: 0.15.5
  - ujson: 5.10.0
  - url-normalize: 1.4.3
  - urllib3: 2.2.2
  - webencodings: 0.5.1
  - werkzeug: 3.0.3
  - wget: 3.2
  - wheel: 0.44.0
  - wrapt: 1.16.0
  - xarray: 2024.7.0
  - xlrd: 2.0.1
  - xyzservices: 2024.6.0
  - yarl: 1.9.4
  - zarr: 2.18.2
  - zict: 3.0.0
  - zipp: 3.20.0
  - zstandard: 0.23.0
• System:
  - OS: Linux
  - architecture:
  - 64bit
  - ELF
  - processor: x86_64
  - python: 3.11.9
  - release: 5.4.0-195-generic
  - version: Demos #215-Ubuntu SMP Fri Aug 2 18:28:05 UTC 2024

More info

In my project, this error occurs in training and evaluating.

Hope the file can give you more help and tell me how to reproduce or solve it correctly.

[forestbat]

@lantiga 5 days passed, so do you or your team know any probable reason of this problem?

[lantiga]

hey @forestbat I'm going to reproduce this now on my end

[lantiga]

@forestbat unfortunately I can't reproduce:

Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/2
Initializing distributed: GLOBAL_RANK: 1, MEMBER: 2/2
----------------------------------------------------------------------------------------------------
distributed_backend=nccl
All distributed processes registered. Starting with 2 processes
----------------------------------------------------------------------------------------------------

tensor(0.5245, device='cuda:1', grad_fn=<MseLossBackward0>)
tensor(0.5649, device='cuda:2', grad_fn=<MseLossBackward0>)
tensor(0.3020, device='cuda:2', grad_fn=<MseLossBackward0>)
tensor(0.7710, device='cuda:1', grad_fn=<MseLossBackward0>)
tensor(0.7531, device='cuda:1', grad_fn=<MseLossBackward0>)
tensor(0.8569, device='cuda:2', grad_fn=<MseLossBackward0>)
tensor(0.6621, device='cuda:2', grad_fn=<MseLossBackward0>)
tensor(0.3022, device='cuda:1', grad_fn=<MseLossBackward0>)
tensor(0.3494, device='cuda:2', grad_fn=<MseLossBackward0>)
tensor(0.7401, device='cuda:1', grad_fn=<MseLossBackward0>)

this is what I get on a 4xT4 (on Lightning Studio), with the latest version of lightning (2.5.0.post0). I tried numerous variations, both with PyTorch 2.2.1 and 2.5.1, I can't get a deadlock.

In the absence of a reproduction it's really hard to guess what's going on.
Deadlocks in ddp usually arise from processes going out of sync, i.e. not executing exactly the same programs, so that one process may reach a barrier that another may not be able to reach. This sometimes happens when dataloaders in different processes have amounts of items, for instance.

Since you're only running the model in eval, there's no synchronization point related to synchronizing gradients priori to the optimizer step, so the opportunity for deadlocks are ultimately very small and probably related to spinning up processes and distributed sampling.

I suggest you work backwards on your codebase and comment out parts until you stop seeing a deadlock.

[forestbat]

I try to reproduce problem in code below，and it crashed with this report:

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset
import lightning
from tqdm import tqdm
import torch.optim as optim
fabric = lightning.Fabric(devices=[0, 1], num_nodes=1, strategy='ddp')


class MyEvaluator:
    def __init__(self):
        fabric.launch()
        x = torch.randn(100, 10)
        y = torch.rand(100, 2)
        self.train_dataset = self.make_dataset(x, y)
        self.test_dataset = self.make_dataset(x, y)

    def make_dataset(self, x, y):
        if fabric.global_rank == 0:
            # In actual program, there is some I/O so I used barrier
            dataset = TensorDataset(x, y)
            return dataset
        fabric.barrier()

    def train_model(self, crit, epoch):
        model = LinearModel()
        model = fabric.setup_module(module=model)
        train_loader = fabric.setup_dataloaders(DataLoader(self.train_dataset, batch_size=10, shuffle=True,
                                                          num_workers=1, multiprocessing_context='spawn'))
        optimizer = fabric.setup_optimizers(optim.Adam(model.parameters(), lr=0.01))
        model.train()
        for epoch in range(0, epoch):
            for _, (src, trg) in enumerate(tqdm(train_loader)):
                output = model(src)
                loss = crit(output, trg)
                fabric.backward(loss)
                optimizer.step()
                optimizer.zero_grad()
        return model

    def eval_model(self, model):
        model.eval()
        test_loader = fabric.setup_dataloaders(DataLoader(self.test_dataset, batch_size=10, shuffle=True,
                                                           num_workers=1, multiprocessing_context='spawn'))
        for _, (src, trg) in enumerate(tqdm(test_loader)):
            output = model(src)
            loss = crit(output, trg)
            yield loss

class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 2)

    def forward(self, x):
        return self.linear(x)

if __name__ == '__main__':
    evaluator = MyEvaluator()
    # dataset = evaluator.make_dataset(x, y)
    crit = nn.MSELoss()
    model = evaluator.train_model(crit, 10)
    for loss in evaluator.eval_model(model):
        print(loss)

Although this is a crash not deadlock, I find something same as deadlock:

  File "/home/username/.conda/envs/torchhydro1/lib/python3.11/site-packages/lightning/fabric/strategies/ddp.py", line 129, in setup_module
    return DistributedDataParallel(module=module, device_ids=device_ids, **self._ddp_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/username/.conda/envs/torchhydro1/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 798, in __init__
    _verify_param_shape_across_processes(self.process_group, parameters)
  File "/home/username/.conda/envs/torchhydro1/lib/python3.11/site-packages/torch/distributed/utils.py", line 263, in _verify_param_shape_across_processes
    return dist._verify_params_across_processes(process_group, tensors, logger)

When deadlock occurs, the program stucks in DistributedDataParallel and dist._verify_params_across_processes too.

So how to use fabric.barrier correctly?

And tqdm is also a reason which probably leads to deadlock, I will try removing it later.

@lantiga

[lantiga]

The issue arises because the two ranks diverge, specifically because make_dataset returns None on rank 1

    def make_dataset(self, x, y):
        if fabric.global_rank == 0:
            # In actual program, there is some I/O so I used barrier
            dataset = TensorDataset(x, y)
            return dataset
        fabric.barrier()  # <- this will never be reached on rank 0, also the make_dataset function will return None on rank 1 

Using if global_rank == 0 and barrier is good when you need to download data to disk, for instance. But then you still need to create a valid dataset object on all ranks if you want to iterate through it.

In this case if you remove the conditional and the barrier, training proceeds correctly

    def make_dataset(self, x, y):
        # In actual program, there is some I/O so I used barrier
        dataset = TensorDataset(x, y)
        return dataset

if you need to actually do I/O, this is the pattern to follow

    def make_dataset(self, x, y):
        if fabric.global_rank == 0:
            # In actual program, there is some I/O so I used barrier
            # do some I/O on one rank only
        fabric.barrier()  # make sure no one runs ahead
        # create the dataset on all ranks based on whatever the I/O did
        dataset = TensorDataset(x, y)
        # return on all ranks
        return dataset

[forestbat]

Sorry, I still can't reproduce this problem easily, but I can try to provide some logs.

I put breakpoints in the two positions:
https://github.com/iHeadWater/torchhydro/blob/dev-gnn/torchhydro/trainers/train_utils.py#L25

# First breakpoint
config_data = create_config_fabric()
total_fab = lightning.Fabric(devices=config_data['training_cfgs']['device'], num_nodes=1, strategy=config_data['training_cfgs']['strategy'])

https://github.com/iHeadWater/torchhydro/blob/dev-gnn/torchhydro/trainers/deep_hydro.py#L144

from torchhydro.trainers.train_utils import total_fab
def __init__(self,  cfgs: Dict,  pre_model=None):
    # Second breakpoint
    total_fab.launch()

When I run my program, log before the first breakpoint is:

import sys; print('Python %s on %s' % (sys.version, sys.platform))
/home/username/.conda/envs/torchhydro1/bin/python3.11 -X pycache_prefix=/home/username/.cache/JetBrains/IntelliJIdea2024.3/cpython-cache /home/username/.local/share/JetBrains/IntelliJIdea2024.3/python-ce/helpers/pydev/pydevd.py --multiprocess --qt-support=auto --port 29781 --file /home/username/torchhydro/experiments/train_with_era5land_gnn_ddp.py
Console output is saving to: /home/username/torchhydro/experiments/results/train_gnn_ddp.txt
[15:25:40] DEBUG    No module named 'forge'                     signatures.py:43
           DEBUG    No module named 'forge'                     signatures.py:43
[15:25:43] DEBUG    CACHEDIR=/home/username/.cache/matplotlib   __init__.py:341
           DEBUG    Using fontManager instance from         font_manager.py:1580
                    /home/username/.cache/matplotlib/fontl
                    ist-v390.json
[15:25:45] DEBUG    Using selector: EpollSelector          selector_events.py:54
           DEBUG    Changing event name from                        hooks.py:482
                    creating-client-class.iot-data to
                    creating-client-class.iot-data-plane
           DEBUG    Changing event name from before-call.apigateway hooks.py:482
                    to before-call.api-gateway
           DEBUG    Changing event name from                        hooks.py:482
                    request-created.machinelearning.Predict to
                    request-created.machine-learning.Predict
           DEBUG    Changing event name from                        hooks.py:482
                    before-parameter-build.autoscaling.CreateLaunch
                    Configuration to
                    before-parameter-build.auto-scaling.CreateLaunc
                    hConfiguration
           DEBUG    Changing event name from                        hooks.py:482
                    before-parameter-build.route53 to
                    before-parameter-build.route-53
           DEBUG    Changing event name from                        hooks.py:482
                    request-created.cloudsearchdomain.Search to
                    request-created.cloudsearch-domain.Search
           DEBUG    Changing event name from                        hooks.py:482
                    docs.*.autoscaling.CreateLaunchConfiguration.co
                    mplete-section to
                    docs.*.auto-scaling.CreateLaunchConfiguration.c
                    omplete-section
           DEBUG    Changing event name from                        hooks.py:482
                    before-parameter-build.logs.CreateExportTask to
                    before-parameter-build.cloudwatch-logs.CreateEx
                    portTask
           DEBUG    Changing event name from                        hooks.py:482
                    docs.*.logs.CreateExportTask.complete-section
                    to
                    docs.*.cloudwatch-logs.CreateExportTask.complet
                    e-section
           DEBUG    Changing event name from                        hooks.py:482
                    before-parameter-build.cloudsearchdomain.Search
                    to
                    before-parameter-build.cloudsearch-domain.Searc
                    h
           DEBUG    Changing event name from                        hooks.py:482
                    docs.*.cloudsearchdomain.Search.complete-sectio
                    n to
                    docs.*.cloudsearch-domain.Search.complete-secti
                    on
           DEBUG    Loading JSON file:                            loaders.py:180
                    /home/username/.conda/envs/torchhydro1/lib/p
                    ython3.11/site-packages/botocore/data/endpoin
                    ts.json
           DEBUG    Loading JSON file:                            loaders.py:180
                    /home/username/.conda/envs/torchhydro1/lib/p
                    ython3.11/site-packages/botocore/data/sdk-def
                    ault-configuration.json
           DEBUG    Event choose-service-name: calling handler      hooks.py:238
                    <function handle_service_name_alias at
                    0x7ff724f4f880>
           DEBUG    Loading JSON file:                            loaders.py:180
                    /home/username/.conda/envs/torchhydro1/lib/p
                    ython3.11/site-packages/botocore/data/s3/2006
                    -03-01/service-2.json.gz
           DEBUG    Loading JSON file:                            loaders.py:180
                    /home/username/.conda/envs/torchhydro1/lib/p
                    ython3.11/site-packages/botocore/data/s3/2006
                    -03-01/endpoint-rule-set-1.json.gz
           DEBUG    Loading JSON file:                            loaders.py:180
                    /home/username/.conda/envs/torchhydro1/lib/p
                    ython3.11/site-packages/botocore/data/partiti
                    ons.json
           DEBUG    Event creating-client-class.s3: calling handler hooks.py:238
                    <function add_generate_presigned_post at
                    0x7ff7250aba60>
           DEBUG    Event creating-client-class.s3: calling handler hooks.py:238
                    <function lazy_call.<locals>._handler at
                    0x7ff724aba5c0>
           DEBUG    Event creating-client-class.s3: calling handler hooks.py:238
                    <function add_generate_presigned_url at
                    0x7ff7250ab7e0>
           DEBUG    Setting s3 timeout as (60, 60)               endpoint.py:408
           DEBUG    Loading JSON file:                            loaders.py:180
                    /home/username/.conda/envs/torchhydro1/lib/p
                    ython3.11/site-packages/botocore/data/_retry.
                    json
           DEBUG    Registering retry handlers for service: s3     client.py:285
           DEBUG    Registering S3 region redirector handler       utils.py:1719
           DEBUG    Registering S3Express Identity Resolver        utils.py:1667
[15:25:47] DEBUG    Registering codec 'zlib'                      registry.py:71
           DEBUG    Registering codec 'gzip'                      registry.py:71
           DEBUG    Registering codec 'bz2'                       registry.py:71
           DEBUG    Registering codec 'lzma'                      registry.py:71
           DEBUG    Registering codec 'blosc'                     registry.py:71
           DEBUG    Registering codec 'zstd'                      registry.py:71
           DEBUG    Registering codec 'lz4'                       registry.py:71
           DEBUG    Registering codec 'astype'                    registry.py:71
           DEBUG    Registering codec 'delta'                     registry.py:71
           DEBUG    Registering codec 'quantize'                  registry.py:71
           DEBUG    Registering codec 'fixedscaleoffset'          registry.py:71
           DEBUG    Registering codec 'packbits'                  registry.py:71
           DEBUG    Registering codec 'categorize'                registry.py:71
           DEBUG    Registering codec 'pickle'                    registry.py:71
           DEBUG    Registering codec 'base64'                    registry.py:71
           DEBUG    Registering codec 'shuffle'                   registry.py:71
           DEBUG    Registering codec 'bitround'                  registry.py:71
           DEBUG    Registering codec 'msgpack2'                  registry.py:71
           DEBUG    Registering codec 'crc32'                     registry.py:71
           DEBUG    Registering codec 'adler32'                   registry.py:71
           DEBUG    Registering codec 'jenkins_lookup3'           registry.py:71
           DEBUG    Registering codec 'json2'                     registry.py:71
           DEBUG    Registering codec 'vlen-utf8'                 registry.py:71
           DEBUG    Registering codec 'vlen-bytes'                registry.py:71
           DEBUG    Registering codec 'vlen-array'                registry.py:71
           DEBUG    Registering codec 'fletcher32'                registry.py:71
           DEBUG    Registering codec 'pcodec'                    registry.py:71
           DEBUG    Registering codec 'n5_wrapper'                registry.py:71
           DEBUG    Registering codec 'fill_hdf_strings'          registry.py:71
           DEBUG    Registering codec 'grib'                      registry.py:71
           DEBUG    Registering codec 'FITSAscii'                 registry.py:71
           DEBUG    Registering codec 'FITSVarBintable'           registry.py:71
           DEBUG    Creating converter from 7 to 5                __init__.py:47
           DEBUG    Creating converter from 5 to 7                __init__.py:47
           DEBUG    Creating converter from 7 to 5                __init__.py:47
           DEBUG    Creating converter from 5 to 7                __init__.py:47
update config file
Notice, you have not set model_hyperparams!
Notice, you have not set batch_size!

And logs between the 1st and the 2nd breakpoint is:

update config file
Notice, you have not set model_hyperparams!
Notice, you have not set batch_size!
update config file
!!!!!!NOTE!!!!!!!!
-------Please make sure the PRECIPITATION variable is in the 1st location in var_t setting!!---------
If you have POTENTIAL_EVAPOTRANSPIRATION, please set it the 2nd!!!-
!!!!!!NOTE!!!!!!!!
-------Please make sure the STREAMFLOW variable is in the 1st location in var_out setting!!---------

When I continue executing my program from the 2nd breakpoint, it come back to the 1st breakpoint, and log in this period is:

INFO: You are using a CUDA device ('NVIDIA RTX 5000 Ada Generation') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision
[15:28:16] INFO     You are using a CUDA device ('NVIDIA RTX     rank_zero.py:63
                    5000 Ada Generation') that has Tensor Cores.
                    To properly utilize them, you should set
                    `torch.set_float32_matmul_precision('medium'
                    | 'high')` which will trade-off precision
                    for performance. For more details, read
                    https://pytorch.org/docs/stable/generated/to
                    rch.set_float32_matmul_precision.html#torch.
                    set_float32_matmul_precision
           INFO     Initializing distributed: GLOBAL_RANK: 0, distributed.py:296
                    MEMBER: 1/2
INFO: Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/2
[15:28:20] DEBUG    No module named 'forge'                     signatures.py:43
           DEBUG    No module named 'forge'                     signatures.py:43
[15:28:23] DEBUG    CACHEDIR=/home/username/.cache/matplotlib   __init__.py:341
           DEBUG    Using fontManager instance from         font_manager.py:1580
                    /home/username/.cache/matplotlib/fontl
                    ist-v390.json
[15:28:24] DEBUG    Using selector: EpollSelector          selector_events.py:54
           DEBUG    Changing event name from                        hooks.py:482
                    creating-client-class.iot-data to
                    creating-client-class.iot-data-plane
           DEBUG    Changing event name from before-call.apigateway hooks.py:482
                    to before-call.api-gateway
           DEBUG    Changing event name from                        hooks.py:482
                    request-created.machinelearning.Predict to
                    request-created.machine-learning.Predict
           DEBUG    Changing event name from                        hooks.py:482
                    before-parameter-build.autoscaling.CreateLaunch
                    Configuration to
                    before-parameter-build.auto-scaling.CreateLaunc
                    hConfiguration
           DEBUG    Changing event name from                        hooks.py:482
                    before-parameter-build.route53 to
                    before-parameter-build.route-53
           DEBUG    Changing event name from                        hooks.py:482
                    request-created.cloudsearchdomain.Search to
                    request-created.cloudsearch-domain.Search
           DEBUG    Changing event name from                        hooks.py:482
                    docs.*.autoscaling.CreateLaunchConfiguration.co
                    mplete-section to
                    docs.*.auto-scaling.CreateLaunchConfiguration.c
                    omplete-section
           DEBUG    Changing event name from                        hooks.py:482
                    before-parameter-build.logs.CreateExportTask to
                    before-parameter-build.cloudwatch-logs.CreateEx
                    portTask
           DEBUG    Changing event name from                        hooks.py:482
                    docs.*.logs.CreateExportTask.complete-section
                    to
                    docs.*.cloudwatch-logs.CreateExportTask.complet
                    e-section
           DEBUG    Changing event name from                        hooks.py:482
                    before-parameter-build.cloudsearchdomain.Search
                    to
                    before-parameter-build.cloudsearch-domain.Searc
                    h
           DEBUG    Changing event name from                        hooks.py:482
                    docs.*.cloudsearchdomain.Search.complete-sectio
                    n to
                    docs.*.cloudsearch-domain.Search.complete-secti
                    on
           DEBUG    Loading JSON file:                            loaders.py:180
                    /home/username/.conda/envs/torchhydro1/lib/p
                    ython3.11/site-packages/botocore/data/endpoin
                    ts.json
           DEBUG    Loading JSON file:                            loaders.py:180
                    /home/username/.conda/envs/torchhydro1/lib/p
                    ython3.11/site-packages/botocore/data/sdk-def
                    ault-configuration.json
           DEBUG    Event choose-service-name: calling handler      hooks.py:238
                    <function handle_service_name_alias at
                    0x7f968c69fe20>
           DEBUG    Loading JSON file:                            loaders.py:180
                    /home/username/.conda/envs/torchhydro1/lib/p
                    ython3.11/site-packages/botocore/data/s3/2006
                    -03-01/service-2.json.gz
           DEBUG    Loading JSON file:                            loaders.py:180
                    /home/username/.conda/envs/torchhydro1/lib/p
                    ython3.11/site-packages/botocore/data/s3/2006
                    -03-01/endpoint-rule-set-1.json.gz
[15:28:25] DEBUG    Loading JSON file:                            loaders.py:180
                    /home/username/.conda/envs/torchhydro1/lib/p
                    ython3.11/site-packages/botocore/data/partiti
                    ons.json
           DEBUG    Event creating-client-class.s3: calling handler hooks.py:238
                    <function add_generate_presigned_post at
                    0x7f968c621580>
           DEBUG    Event creating-client-class.s3: calling handler hooks.py:238
                    <function lazy_call.<locals>._handler at
                    0x7f968c1c3ce0>
           DEBUG    Event creating-client-class.s3: calling handler hooks.py:238
                    <function add_generate_presigned_url at
                    0x7f968c621300>
           DEBUG    Setting s3 timeout as (60, 60)               endpoint.py:408
           DEBUG    Loading JSON file:                            loaders.py:180
                    /home/username/.conda/envs/torchhydro1/lib/p
                    ython3.11/site-packages/botocore/data/_retry.
                    json
           DEBUG    Registering retry handlers for service: s3     client.py:285
           DEBUG    Registering S3 region redirector handler       utils.py:1719
           DEBUG    Registering S3Express Identity Resolver        utils.py:1667
[15:28:27] DEBUG    Registering codec 'zlib'                      registry.py:71
           DEBUG    Registering codec 'gzip'                      registry.py:71
           DEBUG    Registering codec 'bz2'                       registry.py:71
           DEBUG    Registering codec 'lzma'                      registry.py:71
           DEBUG    Registering codec 'blosc'                     registry.py:71
           DEBUG    Registering codec 'zstd'                      registry.py:71
           DEBUG    Registering codec 'lz4'                       registry.py:71
           DEBUG    Registering codec 'astype'                    registry.py:71
           DEBUG    Registering codec 'delta'                     registry.py:71
           DEBUG    Registering codec 'quantize'                  registry.py:71
           DEBUG    Registering codec 'fixedscaleoffset'          registry.py:71
           DEBUG    Registering codec 'packbits'                  registry.py:71
           DEBUG    Registering codec 'categorize'                registry.py:71
           DEBUG    Registering codec 'pickle'                    registry.py:71
           DEBUG    Registering codec 'base64'                    registry.py:71
           DEBUG    Registering codec 'shuffle'                   registry.py:71
           DEBUG    Registering codec 'bitround'                  registry.py:71
           DEBUG    Registering codec 'msgpack2'                  registry.py:71
           DEBUG    Registering codec 'crc32'                     registry.py:71
           DEBUG    Registering codec 'adler32'                   registry.py:71
           DEBUG    Registering codec 'jenkins_lookup3'           registry.py:71
           DEBUG    Registering codec 'json2'                     registry.py:71
           DEBUG    Registering codec 'vlen-utf8'                 registry.py:71
           DEBUG    Registering codec 'vlen-bytes'                registry.py:71
           DEBUG    Registering codec 'vlen-array'                registry.py:71
           DEBUG    Registering codec 'fletcher32'                registry.py:71
           DEBUG    Registering codec 'pcodec'                    registry.py:71
           DEBUG    Registering codec 'n5_wrapper'                registry.py:71
           DEBUG    Registering codec 'fill_hdf_strings'          registry.py:71
           DEBUG    Registering codec 'grib'                      registry.py:71
           DEBUG    Registering codec 'FITSAscii'                 registry.py:71
           DEBUG    Registering codec 'FITSVarBintable'           registry.py:71
           DEBUG    Creating converter from 7 to 5                __init__.py:47
           DEBUG    Creating converter from 5 to 7                __init__.py:47
           DEBUG    Creating converter from 7 to 5                __init__.py:47
           DEBUG    Creating converter from 5 to 7                __init__.py:47
update config file
Notice, you have not set model_hyperparams!
Notice, you have not set batch_size!

Then the program comes to the 2nd breakpoint again, and log in this period is:

Notice, you have not set model_hyperparams!
Notice, you have not set batch_size!
update config file
!!!!!!NOTE!!!!!!!!
-------Please make sure the PRECIPITATION variable is in the 1st location in var_t setting!!---------
If you have POTENTIAL_EVAPOTRANSPIRATION, please set it the 2nd!!!-
!!!!!!NOTE!!!!!!!!
-------Please make sure the STREAMFLOW variable is in the 1st location in var_out setting!!---------
INFO: Initializing distributed: GLOBAL_RANK: 1, MEMBER: 2/2
[2025-01-16 15:29:31,882] torch.distributed.distributed_c10d: [INFO] Using backend config: {'cuda': 'nccl'}
[W CUDAAllocatorConfig.h:30] Warning: expandable_segments not supported on this platform (function operator())
[2025-01-16 15:29:31,891] torch.distributed.distributed_c10d: [INFO] Using backend config: {'cuda': 'nccl'}
[W CUDAAllocatorConfig.h:30] Warning: expandable_segments not supported on this platform (function operator())
INFO: ----------------------------------------------------------------------------------------------------
distributed_backend=nccl
All distributed processes registered. Starting with 2 processes
----------------------------------------------------------------------------------------------------
[15:29:31] INFO     Initializing distributed: GLOBAL_RANK: 1, distributed.py:296
                    MEMBER: 2/2
[15:29:31] INFO     -------------------------------------------- rank_zero.py:63
                    --------------------------------------------
                    ------------
                    distributed_backend=nccl
                    All distributed processes registered.
                    Starting with 2 processes
                    --------------------------------------------
                    --------------------------------------------
                    ------------
# Then the process stucks in total.fab.setup_module()

Deadlocks in ddp usually arise from processes going out of sync, i.e. not executing exactly the same programs, so that one process may reach a barrier that another may not be able to reach.

What's my mistake when initializing process groups?

And in my test code, 2 processes listened 6 ports like this:

But when my program stuck, 2 processes are only listening 5 ports.

What probably happened in this situation? Hope for your reply.

@lantiga

[lantiga]

Quite hard to tell, the library is pretty extensive. I'm ok installing torchhydro, but I really need to be able to reproduce on my end or even better a Lightning Studio you can share with me.

[forestbat]

After moving fabric.launch() from custom trainer's __init__ method

def __init__(self, cfgs: Dict, pre_model=None):
        self.device_num = cfgs["training_cfgs"]["device"]
        total_fab.launch()

to main method in actual project()

if __name__ == "__main__":
    total_fab.launch()
    test_run_model()

the problem disappeared.

Thank for @lantiga spending time on this problem, I will close it and open a new issue with better demo or lightning studio project when I meet similar issue.