search

Lightning-AI / lightning-thunder

Make PyTorch models up to 40% faster! Thunder is a source to source compiler for PyTorch. It enables using different hardware executors at once; across one or thousands of GPUs.
Apache License 2.0
1.2k stars 80 forks source link

pickling failure for auto registered ops fix #1104

Closed k223kim closed 2 months ago

k223kim commented 2 months ago
Before submitting - [ ] Was this discussed/approved via a Github issue? (no need for typos and docs improvements) - [ ] Did you read the [contributor guideline](https://github.com/Lightning-AI/pytorch-lightning/blob/main/.github/CONTRIBUTING.md), Pull Request section? - [ ] Did you make sure to update the docs? - [ ] Did you write any new necessary tests?

What does this PR do?

Fixes #993 .

PR review

Anyone in the community is free to review the PR once the tests have passed. If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

mattteochen commented 2 months ago

Should this PR solve pickling traces coming from nn.Module?

class M(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = torch.nn.Linear(1,1)

    def forward(self, x):
        return self.linear(x)
t-vi commented 2 months ago

There are a number of things that are not currently expected to be pickleable, including traces containing fusions and traces that have been transformed for execution. What the PR achieves is to be able to pickle (many) traces containing auto-registered symbols, e.g. initial compute traces that use the auto-registered things. Eventually, it would be cool to check all symbols for discoverability.