TQDMProgressBar bug (?) causes DDP to hang

[tchittesh]

pytorch-lightning/src/lightning/pytorch/callbacks/progress/tqdm_progress.py

Lines 285 to 286 in 918a1a6

	if not self.train_progress_bar.disable:
	self.train_progress_bar.set_postfix(self.get_metrics(trainer, pl_module))

Is this^ a safe implementation for DDP?

Based on my understanding, this will lead to the following stack trace, ultimately leading to a sync operation on ResultMetrics if certain conditions are met. However, self.train_progress_bar.disable is only False (i.e. enabled) on the rank-zero process, so the metrics computations will only succeed if there happens to be another codepath on non-rank-zero devices that's requesting self._logger_connector.metrics at the same time.

Stack trace

pytorch-lightning/src/lightning/pytorch/callbacks/progress/progress_bar.py

Lines 180 to 201 in 918a1a6

	def get_metrics(
	self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
	) -> dict[str, Union[int, str, float, dict[str, float]]]:
	r"""Combines progress bar metrics collected from the trainer with standard metrics from get_standard_metrics.
	Implement this to override the items displayed in the progress bar.
	Here is an example of how to override the defaults:
	.. code-block:: python
	def get_metrics(self, trainer, model):
	# don't show the version number
	items = super().get_metrics(trainer, model)
	items.pop("v_num", None)
	return items
	Return:
	Dictionary with the items to be displayed in the progress bar.
	"""
	standard_metrics = get_standard_metrics(trainer)
	pbar_metrics = trainer.progress_bar_metrics

pytorch-lightning/src/lightning/pytorch/trainer/trainer.py

Lines 1675 to 1683 in 918a1a6

	@property
	def progress_bar_metrics(self) -> _PBAR_DICT:
	"""The metrics sent to the progress bar.
	This includes metrics logged via :meth:`~lightning.pytorch.core.LightningModule.log` with the
	:paramref:`~lightning.pytorch.core.LightningModule.log.prog_bar` argument set.
	"""
	return self._logger_connector.progress_bar_metrics

pytorch-lightning/src/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py

Lines 253 to 258 in 918a1a6

	@property
	def progress_bar_metrics(self) -> _PBAR_DICT:
	if self.trainer._results:
	metrics = self.metrics["pbar"]
	self._progress_bar_metrics.update(metrics)
	return self._progress_bar_metrics

pytorch-lightning/src/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py

Lines 232 to 237 in 918a1a6

	@property
	def metrics(self) -> _METRICS:
	"""This function returns either batch or epoch metrics."""
	on_step = self._first_loop_iter is not None
	assert self.trainer._results is not None
	return self.trainer._results.metrics(on_step)

pytorch-lightning/src/lightning/pytorch/trainer/connectors/logger_connector/result.py

Lines 471 to 476 in 918a1a6

	def metrics(self, on_step: bool) -> _METRICS:
	metrics = _METRICS(callback={}, log={}, pbar={})
	for _, result_metric in self.valid_items():
	# extract forward_cache or computed from the _ResultMetric
	value = self._get_cache(result_metric, on_step)

pytorch-lightning/src/lightning/pytorch/trainer/connectors/logger_connector/result.py

Lines 425 to 440 in 918a1a6

	@staticmethod
	def _get_cache(result_metric: _ResultMetric, on_step: bool) -> Optional[Tensor]:
	cache = None
	if on_step and result_metric.meta.on_step:
	cache = result_metric._forward_cache
	elif not on_step and result_metric.meta.on_epoch:
	if result_metric._computed is None:
	should = result_metric.meta.sync.should
	if not should and result_metric.is_tensor and _distributed_is_initialized():
	warning_cache.warn(
	f"It is recommended to use `self.log({result_metric.meta.name!r}, ..., sync_dist=True)`"
	" when logging on epoch level in distributed setting to accumulate the metric across"
	" devices.",
	category=PossibleUserWarning,
	)
	result_metric.compute()

I ran into this when debugging an issue in my training script where DDP was hanging at the end of the first epoch. Changing the original code chunk as such resolved the deadlock for me.

metrics = self.get_metrics(trainer, pl_module)
if not self.train_progress_bar.disable: 
    self.train_progress_bar.set_postfix(metrics) 

To provide a little more detail, I have a CustomCallback that calls trainer.strategy.barrier() during on_train_epoch_end. At the end of the first training epoch, the rank-zero process hangs on TQDMProgressBar.on_train_epoch_end and the non-rank-zero processes hang on CustomCallback.on_train_epoch_end.

cc @justusschock @lantiga

[deependujha]

Hi @tchittesh , I tried running a multi-GPU DDP setup with a similar callback and TQDMProgressBar, but I wasn’t able to reproduce the hanging issue.

Here’s my code:

import lightning as L
from lightning.pytorch.callbacks import TQDMProgressBar
from lightning.pytorch.demos.boring_classes import BoringModel

class MyModel(BoringModel):
    def __init__(self):
        super().__init__()
        self.latest_loss = 0

    def training_step(self, batch, batch_idx):
        loss = super().training_step(batch, batch_idx)["loss"]
        self.log("epoch_loss", loss, on_step=False, on_epoch=True, sync_dist=True)
        self.latest_loss = loss.item()
        return {"loss": loss}

class Bar(TQDMProgressBar):
    def on_train_epoch_end(self, trainer, pl_module):
        # barrier before progress bar update
        trainer.strategy.barrier()
        super().on_train_epoch_end(trainer, pl_module)

    def get_metrics(self, trainer, model):
        items = super().get_metrics(trainer, model)
        items["loss"] = model.latest_loss
        items.pop("v_num", None)
        return items

if __name__ == "__main__":
    bar = Bar()
    model = MyModel()
    trainer = L.Trainer(
        strategy="ddp",
        accelerator="gpu",
        devices=4,
        max_epochs=10,
        callbacks=[bar],
        log_every_n_steps=1,
    )
    trainer.fit(model)

• output on T4 machine with 4 gpus:

⚡ ~ python main.py
💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/4
Initializing distributed: GLOBAL_RANK: 3, MEMBER: 4/4
Initializing distributed: GLOBAL_RANK: 2, MEMBER: 3/4
Initializing distributed: GLOBAL_RANK: 1, MEMBER: 2/4
----------------------------------------------------------------------------------------------------
distributed_backend=nccl
All distributed processes registered. Starting with 4 processes
----------------------------------------------------------------------------------------------------

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]
LOCAL_RANK: 2 - CUDA_VISIBLE_DEVICES: [0,1,2,3]
LOCAL_RANK: 3 - CUDA_VISIBLE_DEVICES: [0,1,2,3]
LOCAL_RANK: 1 - CUDA_VISIBLE_DEVICES: [0,1,2,3]

  | Name  | Type   | Params | Mode 
-----------------------------------------
0 | layer | Linear | 66     | train
-----------------------------------------
66        Trainable params
0         Non-trainable params
66        Total params
0.000     Total estimated model params size (MB)
1         Modules in train mode
0         Modules in eval mode
Sanity Checking: |                                                                                                                                                                      | 0/? [00:00<?, ?it/s]/home/zeus/miniconda3/envs/cloudspace/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:433: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=7` in the `DataLoader` to improve performance.
/home/zeus/miniconda3/envs/cloudspace/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:433: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=7` in the `DataLoader` to improve performance.
Epoch 9: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 16/16 [00:00<00:00, 214.59it/s, loss=0.0919]`Trainer.fit` stopped: `max_epochs=10` reached.                                                                                                                                                               
Epoch 9: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 16/16 [00:00<00:00, 203.04it/s, loss=0.0919]
⚡ ~ 

Am I missing something in the setup that triggers the deadlock?

[tchittesh]

Here's a minimal reproduction of the deadlock (and the fix for it).

Some notes:

• It's important to instantiate CustomCallback in LightningModule.configure_callbacks() instead of directly passing it to the trainer.
  • If we pass it into the trainer, the order of callbacks is CustomCallback -> TQDMProgressBar -> ModelCheckpoint. In this case, the CustomCallback will execute correctly. Then, rank 0 will try to sync/compute ResultMetrics for the TQDMProgressBar. Rank 1 will skip past TQDMProgressBar, but it'll get lucky because ModelCheckpoint will also request to sync/compute the same ResultMetrics.
  • Passing it into LightningModule.configure_callbacks() will change the order of callbacks to TQDMProgressBar -> CustomCallback -> ModelCheckpoint. In this case, rank 0 will try to compute ResultMetrics for the TQDMProgressBar, but rank 1 will skip past the TQDMProgressBar and get stuck at the dist.gather_object sync point in CustomCallback.

import torch.distributed as dist
import pytorch_lightning as pl
import pytorch_lightning.callbacks
from pytorch_lightning.demos.boring_classes import BoringModel

def on_train_epoch_end_monkey_patch(self, trainer, pl_module):
    metrics = self.get_metrics(trainer, pl_module)
    if not self.train_progress_bar.disable:
        self.train_progress_bar.set_postfix(metrics)
    if self._leave:
        self.train_progress_bar.close()

# Uncomment this monkey patch to fix the deadlock:
# pytorch_lightning.callbacks.TQDMProgressBar.on_train_epoch_end = on_train_epoch_end_monkey_patch

class MyModel(BoringModel):
    def training_step(self, batch, batch_idx):
        loss = super().training_step(batch, batch_idx)["loss"]
        self.log("loss", loss, on_step=False, on_epoch=True, sync_dist=True)
        return {"loss": loss}

    def configure_callbacks(self):
        return [CustomCallback()]

class CustomCallback(pl.Callback):
    def on_train_epoch_end(self, trainer, pl_module):
        obj = 10
        objs = [None] * trainer.world_size
        dist.gather_object(
            obj,
            objs if pl_module.global_rank == 0 else None,
            dst=0,
        )

if __name__ == "__main__":
    model = MyModel()
    trainer = pl.Trainer(
        strategy="ddp",
        accelerator="gpu",
        devices=2,
        max_epochs=10,
    )
    trainer.fit(model)

[deependujha]

Hi @tchittesh, thanks for the detailed repro! Another way to fix this DDP hang is to adjust _reorder_callbacks():

tuner_callbacks -> other_callbacks -> progress_bar_callbacks -> checkpoint_callbacks.

pytorch-lightning/src/lightning/pytorch/trainer/connectors/callback_connector.py

Lines 214 to 239 in 4dc1dbc

	def _reorder_callbacks(callbacks: list[Callback]) -> list[Callback]:
	"""Moves all the tuner specific callbacks at the beginning of the list and all the `ModelCheckpoint` callbacks
	to the end of the list. The sequential order within the group of checkpoint callbacks is preserved, as well as
	the order of all other callbacks.
	Args:
	callbacks: A list of callbacks.
	Return:
	A new list in which the first elements are tuner specific callbacks and last elements are ModelCheckpoints
	if there were any present in the input.
	"""
	tuner_callbacks: list[Callback] = []
	other_callbacks: list[Callback] = []
	checkpoint_callbacks: list[Callback] = []
	for cb in callbacks:
	if isinstance(cb, (BatchSizeFinder, LearningRateFinder)):
	tuner_callbacks.append(cb)
	elif isinstance(cb, Checkpoint):
	checkpoint_callbacks.append(cb)
	else:
	other_callbacks.append(cb)
	return tuner_callbacks + other_callbacks + checkpoint_callbacks

to be something like:

@staticmethod
    def _reorder_callbacks(callbacks: list[Callback]) -> list[Callback]:
        
        tuner_callbacks: list[Callback] = []
        other_callbacks: list[Callback] = []
        progress_bar_callbacks: list[Callback] = []
        checkpoint_callbacks: list[Callback] = []

        for cb in callbacks:
            if isinstance(cb, (BatchSizeFinder, LearningRateFinder)):
                tuner_callbacks.append(cb)
            elif isinstance(cb, Checkpoint):
                checkpoint_callbacks.append(cb)
            elif isinstance(cb, ProgressBar):
                progress_bar_callbacks.append(cb)
            else:
                other_callbacks.append(cb)

        return tuner_callbacks + other_callbacks + progress_bar_callbacks + checkpoint_callbacks

This preserves the relative order within each group and avoids rank 0 being blocked while syncing metrics.
This approach also fixes the bug. Any concerns or edge cases I might be missing?

cc: @SkafteNicki

[SkafteNicki]

@deependujha yeah, i think just changing the the _reorder_callbacks method is a good strategy for solving this issue. I cannot really see a edge case where this change will be a problem. Please go ahead and send a PR.

[tchittesh]

I don't understand. Isn't this just hiding the bug in TQDMProgressBar? If someone later comes in and disables the ModelCheckpoint callback, won't they hit the same deadlock? Could you provide some more color on why reordering callbacks is a better solution than computing metrics on all ranks?

[deependujha]

works fine even with enable_checkpointing=False.

[SkafteNicki]

@tchittesh @deependujha lets start out with a PR that implement the appropriate test e.g. something that would fail with the current codebase and then implement the fix to make sure that it indeed fixes the error

[tchittesh]

Agreed that we should add a test here.

However, I strongly believe that reordering callbacks is not a good fix. Using some pseudo-code for simplicity, there's a

if global_rank == 0:
    sync_and_compute_result_metrics()

block in TQDMProgressBar, which will deadlock unless non-rank 0 processes also call sync_and_compute_result_metrics() somewhere else. This means that in isolation, TQDMProgressBar.on_train_epoch_end() will always hang with DDP!

We typically don't observe this on master, because the ModelCheckpoint callback tries to sync_and_compute_result_metrics() as well, on non-rank 0 processes. Similarly, with the proposed reordering fix, the deadlock "goes away" because LoggerConnector also runs sync_and_compute_result_metrics() on non-rank 0 processes. However, the correctness of TQDMProgressBar ideally shouldn't depend on the presence of other hooks that get executed afterwards.

Let me know if I'm overlooking something -- I hope we can reach some consensus on this issue.