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

default transforms for thunder.jit #1299

Open kshitij12345 opened 1 month ago

kshitij12345 commented 1 month ago

Similar to default executors, we should have a good default for transforms to thunder.jit. This issue is to discuss how an API for the same would look like.

One option is that it could look similar to get_default_executors() -> List[Executor]. It makes more sense for executors as the higher priority executor will pick-up an operator (and lower priority would be ignored). In case of transforms, the transforms will applied in order and order of transform application maybe important. However, note that this is how thunder.jit takes in transforms currently (see docs), and applies them in the order they appear.

Context - We may want to make the ConstantFolding transform a default and this issue is to facilitate the discussion around the same. See https://github.com/Lightning-AI/lightning-thunder/issues/1250 and https://github.com/Lightning-AI/lightning-thunder/pull/1273

cc @apaz-cli

tfogal commented 1 month ago

@kshitij12345 is the constant folding as currently committed applied by default on at least the ThunderFX path, today? Or would we need this or manual intervention to apply that pass at all?

kshitij12345 commented 1 month ago

is the constant folding as currently committed applied by default on at least the ThunderFX path, today?

It is off by default even on ThunderFX path. For now, we can manually specify the pass with

from thunder.transforms.constant_folding import ConstantFolding
backend = ThunderCompiler(transforms=[ConstantFolding(),])
t-vi commented 1 month ago

So I think it would be really good if we could use recipes #1188 here to enable transforms by default. Last I knew, we had an eta of 1-2 weeks. (cc @lantiga )

tfogal commented 1 month ago

ack, appreciate the fast response! but also, take a break, it's late on a Friday ;-)