Lightning-AI / Lightning-AI/lightning-thunder

Do not perform DTensor isinstance check for all nvFuser inputs at runtime

Open
#2,651 0 comments 1 reaction 1 assignee View on GitHub

Nobody has claimed this yet.

DTensor nvfuser
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

## 🐛 Bug

Current nvFuser executor in Thunder checks every argument at runtime if it's DTensor or not. This check should be performed at trace construction time because the dispatch to DTensor-based execution or not should be done solely based on the types recorded in the trace.

Place where DTensor check is performed: https://github.com/Lightning-AI/lightning-thunder/blob/750620bdc3e956ffda287e3a24af6a20b9e06f5f/thunder/executors/nvfuserex_impl.py#L555

This would involve updating the construction of `FusionDefinitionWrapper` to include a new attribute "use_dtensor_execute" (or some better name) to control the dispatch.

cc @tfogal

Contributor guide

No contributing guide indexed for this repository

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.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.