Lightning-AI / Lightning-AI/pytorch-lightning

Should Fabric have a `.mark_step()`?

Open
#17,622 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

discussion fabric feature won't fix
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Description & Motivation

In [#265](https://github.com/Lightning-AI/lit-llama/pull/265), the inference code gets sped up by inserting the XLA mark_step() calls. Should Fabric provide an accelerator/strategy-agnostic method for that?

### Pitch

It could be implemented directly like this in the Fabric class:

```py
def mark_step(self):
if self.device.type == "xla":
import torch_xla.core.xla_model as xm

xm.mark_step()
# for all other accelerators, this is a no-op
```

The convenience method enables you to call

`fabric.mark_step()` directly in the loop and no code changes are required once switching to another accelerator or strategy.

### Alternatives

The above proposes to add the logic directly to the Fabric object since it is very simple and concerns only the XLA strategy and to keep the base strategy interface lean. If this turns out to generalize and be required for other accelerators in the future, we could then decide to add the interface to the strategy base class.

### Additional context

_No response_

cc @borda @carmocca @justusschock @awaelchli

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 with the Fabric class and the accelerator or strategy handling described in the issue. Check the proposed torch_xla mark_step() behavior and determine how the method should behave for non-XLA accelerators; done means the Fabric-level call works for XLA and is a no-op elsewhere.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.