Lightning-AI / Lightning-AI/pytorch-lightning

AWS Neuron support

Open
#15,172 10 comments 3 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

accelerator fabric feature pl strategy
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

## 🚀 Feature

https://awsdocs-neuron.readthedocs-hosted.com/en/latest/index.html
https://aws.amazon.com/machine-learning/neuron/

### Motivation

https://aws.amazon.com/about-aws/whats-new/2022/10/ec2-trn1-instances-high-performance-cost-effective-deep-learning-training/

### Pitch

Neuron is XLA based, so it'll probably be an accelerator and strategy. Their marketing materials also advertise that it supports FSDP, and Megatron-LM. It's unclear at this point whether we would have specialized Neuron-based implementations for those.

We would also need access to hardware for testing.

______________________________________________________________________

#### If you enjoy Lightning, check out our other projects! ⚡

- [**Metrics**](https://github.com/Lightning-AI/metrics): Machine learning metrics for distributed, scalable PyTorch applications.

- [**Lite**](https://pytorch-lightning.readthedocs.io/en/latest/starter/lightning_lite.html): enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.

- [**Flash**](https://github.com/Lightning-AI/lightning-flash): The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.

- [**Bolts**](https://github.com/Lightning-AI/lightning-bolts): Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.

- [**Lightning Transformers**](https://github.com/Lightning-AI/lightning-transformers): Flexible interface for high-performance research using SOTA Transformers leveraging PyTorch Lightning, Transformers, and Hydra.

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 reading the linked AWS Neuron documentation and the Neuron hardware requirements. Determine the integration points for an XLA-based accelerator or strategy, including whether FSDP or Megatron-LM need specialized implementations; completion would require an agreed design, implementation, and access to Neuron hardware for testing.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, pytorch
Domain
cloud, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
20/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.