Lightning-AI / Lightning-AI/pytorch-lightning

Shark backend integration

Open
#11,884 2 comments 7 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

🚀 Feature

Integrate Shark as an accelerator backend for Lightning.

Motivation

Lightning users would benefit from a "performance first" backend option for running their models. Shark users would benefit from Lightning's clean UX-focused API and clean HW abstraction. Integration seems like a win-win.

Pitch

I'd like to integrate shark as an accelerator similar to how IPUs are implemented. A WIP PR should be incoming soon.

Additional context

Warning- not user friendly (hence the motivation for this issue)

SHARK repo: https://github.com/NodLabs/SHARK/tree/PerformanceDevBranch
Benchmarks: https://github.com/powderluv/transformer-benchmarks
SHARK examples: https://github.com/NodLabs/shark-samples/tree/main/examples

Examples of using torch mlir as a backend for pytorch can be found at:
https://github.com/llvm/torch-mlir/blob/main/examples
resnet jupyter notebook
plain python
Other models like bert can be generated using the heavydep testing instructions at top-level torch-mlir.

cc @borda @akihironitta @rohitgr7

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 reviewing the SHARK repository and its examples, then compare the torch-mlir examples linked in the issue, including the ResNet notebook and plain Python example. The issue is complete when SHARK is integrated as a Lightning accelerator backend with a user-friendly path and demonstrated model usage.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.