Lightning-AI/pytorch-lightning

Returning None from training_step with multi GPU DDP training

Open

#5,243 opened on Dec 23, 2020

View on GitHub
 (26 comments) (4 reactions) (1 assignee)Python (3,233 forks)batch import
distributedfeaturehelp wantedpriority: 1

Repository metrics

Stars
 (26,687 stars)
PR merge metrics
 (Avg merge 9d 15h) (3 merged PRs in 30d)

Description

🐛 Bug

Returning None from training_step with multi GPU DDP training freezes the training without exception

To Reproduce

Starting multi-gpu training with a None-returning training_step function.

Example training_step function:

    def training_step(self, batch, batch_idx):
        data, target = batch
        model_outputs = self.forward(images)
        loss = calc_loss(model_outputs, target)

        if torch.isnan(loss) or random.random() < .05:
            return None

        return loss

Example trainer:

 trainer = Trainer(
    gpus=2,
    distributed_backend="ddp"
)

Expected behavior

To continue training with skipping the current batch as pointed out at here.

Environment

No specific environment is needed to reproduce this bug.

Additional context

This issue was mentioned here: #4956 but not with specifics.

Note: By the time this issue being investigated, a help for a workaround would be great!

Contributor guide

Research direction
Investigate the distributed data parallel (DDP) synchronization mechanism when training step returns None. Trace how the Lightning module handles None outputs and why it causes a deadlock. Look at the gradient synchronization in the DDP plugin and see if None triggers an incomplete all reduce. Check related issues and PRs for potential fixes.
Tech stack
pythonpytorch
Domain
backendmachine learning
Issue type
Bug
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
3
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
1-3 hours
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
Active
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
Mostly clear
Prerequisites
PythonPyTorchDistributed Training
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
65