Lightning-AI / Lightning-AI/lightning-thunder

Allow user to specify the process group for TEv2

Open
#2,437 0 comments 0 reactions 0 assignees View on GitHub
distributed enhancement TransformerEngine
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

## 🚀 Feature

Allow users to specify pg when using TEv2 executor. This issue stems from conversation in [#2404 (comment)](https://github.com/Lightning-AI/lightning-thunder/pull/2404#discussion_r2266348534)

This could be achieved with a compiler flag or, arguably better, with an executor option/argument.

This issue could be paired with giving the recipe as an option to the executor too to get rid of TEs autocast and global state manager #2080

cc @crcrpar @ali-alshaar7

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the TEv2 executor implementation and the discussion in PR #2404, then trace how executor options and compiler flags are currently handled. Define the supported way for users to provide a process group, and verify that TEv2 uses it without relying on the existing global state; the related recipe option in #2080 is a separate consideration.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers, distributed-systems
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.