Lightning-AI / Lightning-AI/pytorch-lightning

Add something like `use_compile` parameter for Trainer

Open
#20,242 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature torch.compile
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Description & Motivation

For below example, model is being compiled, `DDPStrategy` is passed to Trainer, then during fit method `DDPStrategy` is being applied, so forward is compiled but `_pre_forward/_post_forward` in DDP class is not.
Due to this in DDP `_pre_forward/_post_forward` cpp_reducer is not being disabled later on causing problem with queueing callback.
When DDP is also compiled cpp_reducer is disabled as expected.

```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
import pytorch_lightning as pl
from pytorch_lightning.strategies import DDPStrategy
import functools
from torch._dynamo import compiled_autograd

torch._dynamo.config.optimize_ddp = "python_reducer"

class SomeModel(pl.LightningModule):
def __init__(self):
super().__init__()
self.layer = nn.Linear(10, 1)

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

def training_step(self, batch, batch_idx):
x, y = batch
return torch.mean((self(x) - y) ** 2)

def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=0.02)

def create_dataset(num_samples=1000):
x = torch.randn(num_samples, 10)
y = torch.sum(x, dim=1, keepdim=True)
return TensorDataset(x, y)

def run_training():
dataset = create_dataset()
train_loader = DataLoader(dataset, batch_size=32)

model = SomeModel()

# First compile whole model
model = torch.compile(model)

# static_graph has to be true causing _DDPSink.backward method queueing callback
ddp_strategy = DDPStrategy(static_graph=True)

trainer = pl.Trainer(
max_epochs=2,
accelerator='cpu',
devices=1,
strategy=ddp_strategy
)

# DDP will be applied inside fit method, so DDP pre/post forward won't be compiled while forward is
with compiled_autograd.enable(torch.compile()):
trainer.fit(model, train_loader)

if __name__ == "__main__":
run_training()
```

Expected repro:
```
[rank0]: Variable._execution_engine.queue_callback( # type: ignore[call-arg,misc]
[rank0]: RuntimeError: Final callbacks can only be installed during backward pass.
```

### Pitch

It seems useful to compile after applying strategy, so my suggestion is to add something like bool `use_compile` parameter for Trainer which would help for example in this situation, and also be cleaner to use.
Looks like it should be more advanced than just bool to setup specific backend and other optional compile parameters.

### Alternatives

src/lightning/pytorch/trainer/trainer.py
```python
class Trainer:
def _run(
self, model: "pl.LightningModule", ckpt_path: Optional[_PATH] = None
) -> Optional[Union[_EVALUATE_OUTPUT, _PREDICT_OUTPUT]]:
(...)
# ----------------------------
# SET UP THE TRAINER
# ----------------------------
(...)
self.strategy.setup(self)
(...)

------> ### Pseudo proposition
------> if self.use_compile
------> self.model = torch.compile(self.model)

# ----------------------------
# RUN THE TRAINER
# ----------------------------
results = self._run_stage()

(...)
```

### Additional context

File: torch/nn/parallel/distributed.py

```python
class DistributedDataParallel(Module, Joinable):
def _should_disable_cpp_reducer(self) -> bool:
return self._use_python_reducer and (
torch._utils.is_compiling() or self._force_to_disable_cpp_reducer
)

def _pre_forward(self, *inputs, **kwargs):
if self._should_disable_cpp_reducer():
return inputs, kwargs
(...)

def _post_forward(self, output):
if self._should_disable_cpp_reducer():
return output
(...)
```

Dynamo replaces output of `torch._utils.is_compiling()` to True when code is compiled, False otherwise.

cc @lantiga @borda

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in src/lightning/pytorch/trainer/trainer.py around Trainer._run and the strategy.setup(self) call, then review the reported interaction with torch/nn/parallel/distributed.py. Determine how compilation options should be represented and applied after strategy setup, including backend and other parameters. Done means the reproduced DDP compiled-autograd failure is resolved and targeted tests cover the new Trainer behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.