Lightning-AI / Lightning-AI/pytorch-lightning

Add a docs peformance section

Open
#12,398 3 comments 3 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

docs performance
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

## 📚 Documentation

Have a page where we go through all the tips to improve performance and their caveats.

For example, some are suggested in https://william-falcon.medium.com/pytorch-lightning-vs-deepspeed-vs-fsdp-vs-ffcv-vs-e0d6b2a95719

A speed checklist could be:

- Tune `DataLoader(batch_size)`
- Tune `DataLoader(num_workers)`
- Tune `DataLoader(persistent_workers)`
- Check training will happen on device: `Trainer(accelerator="gpu")`
- (If desired) Check script uses the DDP strategy and an apt # of devices
- (If DDP) Use `Trainer(strategy="ddp_find_unused_parameters_false")`
- Check anomaly detection being turned off `Trainer(detect_anomaly=False)` (default)
- Maybe set `Trainer(benchmark=True)` (default) or `Trainer(benchmark=False)`. Whether this is helpful or disadvantageous depends on the data and operations used in your model.
- Disable loggers: `Trainer(logger=False)`. It will use `TensorBoard` by default.
- Disable checkpointing if this is a pure benchmark where we don’t really care about the weights: `Trainer(enable_checkpointing=False)`. Or just try to minimize how often you save
- Reduce progress bar printing: `Trainer(callbacks=ProgressBar(refresh_rate=N))` or disable it completely: `Trainer(enable_progress_bar=False)`
- During logging, avoid `self.log(on_step=True)` and `self.log(sync_dist=True)`
- Make sure `setup` is idempotent so that you don't do the same work more than once accidentally
- A better process group backend: [`DDPStrategy(pg_backend="fairring")`](https://github.com/facebookresearch/fairring)
- (If DDP) Use `DDP(static_graph=True)` (there's an automatic message suggesting this)
- (if torch<2.0) Override `optimizer_zero_grad` to call `zero_grad(set_to_none=True)`
- If using an Ampere+ GPU: set `torch.set_float32_matmul_precision('medium' | 'high')` (there's an automatic message suggesting this)
- Try to `torch.compile(your_lightningmodule)` before passing it to `trainer.fit()`

Then for memory:
- Beware of overriding `training_epoch_end`
- Model memory profiling: https://zdevito.github.io/2022/08/16/memory-snapshots.html

DataLoading:
- There are faster DataLoader implementations to PyTorch's such as https://ffcv.io or https://docs.nvidia.com/deeplearning/dali/user-guide/docs

cc @borda @rohitgr7 @akihironitta

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 by locating the project’s documentation structure and checking the current documentation for the referenced Trainer, DataLoader, DDP, logging, checkpointing, and profiling options. Verify each recommendation and caveat against current PyTorch Lightning behavior and the linked resources. Done means a maintainable performance and memory guide with a practical speed checklist and accurate examples.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
documentation, machine-learning, performance
Issue type
Documentation
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
28/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.