Sourced from pytorch-lightning's releases.

Lightning v2.3: Tensor Parallelism and 2D Parallelism

Lightning AI is excited to announce the release of Lightning 2.3 ⚡

Did you know? The Lightning philosophy extends beyond a boilerplate-free deep learning framework: We've been hard at work bringing you Lightning Studio. Code together, prototype, train, deploy, host AI web apps. All from your browser, with zero setup.

This release introduces experimental support for Tensor Parallelism and 2D Parallelism, PyTorch 2.3 support, and several bugfixes and stability improvements.

• Highlights
  • Tensor Parallelism (beta)
  • 2D Parallelism (beta)
  • Training Mode in Model Summary
  • Special Forward Methods in Fabric
• Notable Changes
• Full Changelog
  • PyTorch Lightning
  • Lightning Fabric
• Contributors

Highlights

Tensor Parallelism (beta)

Tensor parallelism (TP) is a technique that splits up the computation of selected layers across GPUs to save memory and speed up distributed models. To enable TP as well as other forms of parallelism, we introduce a ModelParallelStrategy for both Lightning Trainer and Fabric. Under the hood, TP is enabled through new experimental PyTorch APIs like DTensor and torch.distributed.tensor.parallel.

PyTorch Lightning

Enabling TP in a model with PyTorch Lightning requires you to implement the LightningModule.configure_model() method where you convert selected layers of a model to paralellized layers. This is an advanced feature, because it requires a deep understanding of the model architecture. Open the tutorial Studio to learn the basics of Tensor Parallelism.

 

import lightning as L
from lightning.pytorch.strategies import ModelParallelStrategy
from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
from torch.distributed.tensor.parallel import parallelize_module
1. Implement the configure_model() method in LightningModule
class LitModel(L.LightningModule):
def init(self):
super().init()
self.model = FeedForward(8192, 8192)
</tr></table>

... (truncated)

• a42484c Fix failing app tests (#19971)
• f6fd046 Release 2.3.0 (#19954)
• a97814a Update README.md
• fa5da26 Update README.md (#19968)
• 06ea3a0 Fix resetting epoch loop restarting flag in LearningRateFinder (#19819)
• 5fa32d9 Ignore parameters causing ValueError when dumping to YAML (#19804)
• 4f96c83 Sanitize argument-free object params before logging (#19771)
• a611de0 Removing numpy requirement from all files in examples/pytorch/domain_template...
• 812ffde Fix save_last type annotation for ModelCheckpoint (#19808)
• 7668a6b Flexible and easy to use HSDP setting (#19504)
• Additional commits viewable in compare view

You can trigger Dependabot actions by commenting on this PR:

• @dependabot rebase will rebase this PR
• @dependabot recreate will recreate this PR, overwriting any edits that have been made to it
• @dependabot merge will merge this PR after your CI passes on it
• @dependabot squash and merge will squash and merge this PR after your CI passes on it
• @dependabot cancel merge will cancel a previously requested merge and block automerging
• @dependabot reopen will reopen this PR if it is closed
• @dependabot close will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
• @dependabot show <dependency name> ignore conditions will show all of the ignore conditions of the specified dependency
• @dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)