Lightning-AI / Lightning-AI/pytorch-lightning
`TQDMProgressBar` refresh forces TPU to recompile compute graph
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
TQDMProgressBar refresh forces TPU to recompile compute graph. This causes slow execute time.
### How to reproduce the bug
```python
import re
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import pytorch_lightning as pl
import torch_xla.debug.metrics as met
from pytorch_lightning.callbacks import TQDMProgressBar
from torch.utils.data import DataLoader, Dataset
from pytorch_lightning.callbacks import Callback
import torch_xla.core.xla_model as xm
class dummyDataset(Dataset):
def __getitem__(self, index):
return torch.rand(512), torch.rand(1)
def __len__(self):
return 10_000
class dummyModel(pl.LightningModule):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(512, 512),
nn.Linear(512, 1),
)
def training_step(self, batch):
x, y = batch
logits = self.net(x)
time.sleep(1)
return {'loss': F.cross_entropy(logits, y)}
def configure_optimizers(self):
return torch.optim.AdamW(self.parameters())
class TPUMetricCallback(Callback):
def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
if xm.is_master_ordinal():
report = met.metrics_report()
xrt_compile_count = re.search('Metric: XrtCompile\s+TotalSamples: (\d+)', report).group(1)
print(f'XrtCompile: {xrt_compile_count}, batch_idx = {batch_idx}')
def main():
refresh_rate = 1
ds = dummyDataset()
dl = DataLoader(ds, batch_size=16)
model = dummyModel()
metrics_callback = TPUMetricCallback()
tqdm_callback = TQDMProgressBar(refresh_rate=refresh_rate)
trainer = pl.Trainer(max_epochs=10,
accelerator='tpu',
devices=1,
callbacks=[tqdm_callback, metrics_callback])
trainer.fit(model=model,
train_dataloaders=dl
)
if __name__ == '__main__':
main()
```
### Error messages and logs
Refresh_rate = 1
```
XrtCompile: 3, batch_idx = 0
Epoch 0: 0%|▎ | 2/625 [00:02<12:42, 1.22s/it, loss=0, v_num=0]
XrtCompile: 4, batch_idx = 1
Epoch 0: 0%|▍ | 3/625 [00:03<12:41, 1.22s/it, loss=0, v_num=0]
XrtCompile: 5, batch_idx = 2
Epoch 0: 1%|▋ | 4/625 [00:04<12:55, 1.25s/it, loss=0, v_num=0]
XrtCompile: 6, batch_idx = 3
Epoch 0: 1%|▊ | 5/625 [00:06<13:22, 1.29s/it, loss=0, v_num=0]
XrtCompile: 7, batch_idx = 4
Epoch 0: 1%|▉ | 6/625 [00:08<13:50, 1.34s/it, loss=0, v_num=0]
XrtCompile: 8, batch_idx = 5
Epoch 0: 1%|█ | 7/625 [00:09<14:26, 1.40s/it, loss=0, v_num=0]
XrtCompile: 9, batch_idx = 6
Epoch 0: 1%|█▎ | 8/625 [00:11<15:02, 1.46s/it, loss=0, v_num=0]
XrtCompile: 10, batch_idx = 7
Epoch 0: 1%|█▍
```
Refresh_rate = 5
```
Epoch 0: 0%| | 0/625 [00:00=0.15.0
tensorboardX
protobuf==3.19.5
cc @carmocca @JackCaoG @steventk-g @Liyang90 @awaelchli
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the TQDMProgressBar callback and reproduce the report on a TPU using the provided dummy model, TPUMetricCallback, and refresh_rate values of 1 and 5. Compare XrtCompile counts from torch_xla metrics; done means refreshing the progress bar no longer causes a compile on every refresh while training behavior remains unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100