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.21k stars 80 forks source link

Support NeMo StableDiffusion network #266

Open athitten opened 7 months ago

athitten commented 7 months ago

🚀 Feature

NeMo's "MegatronLatentDiffusion" model implements Stable Diffusion

Initial examine:

Found 37 distinct operations, of which 30 (81.1%) are supported

Motivation

Pitch

Work items

cc @tfogal

jjsjann123 commented 7 months ago

Test memory use of NeMo model with out-of-place add instead of the in-place variant

Wondering how we are going to evaluate this one. i.e. if we are just using eager mode for the time being, an inplace version does save memory footprint, but that might not translate to a compiled program. Since we might have a chance to fuse out-of-place operation.

tfogal commented 7 months ago

Test memory use of NeMo model with out-of-place add instead of the in-place variant

Wondering how we are going to evaluate this one. i.e. if we are just using eager mode for the time being, an inplace version does save memory footprint, but that might not translate to a compiled program. Since we might have a chance to fuse out-of-place operation.

Yeah, it's definitely not an apples-to-apples comparison.

But actually that's not the comparison we're interested in here. Here we'd want to evaluate NeMo alone (i.e. w/o thunder) using in-place ops vs. NeMo alone using out-of-place ops.

While we can always modify any network to use Thunder (and this is the approach I would recommend to most users!), NeMo isn't going to adopt Thunder instantly. First we need to show some value to NeMo.

Until then, the NeMo team is happy to consider changes that make networks more amenable to Thunder, but only if they improve or at least do not regress non-Thunder-based use cases.

Does that make help/make sense?

jjsjann123 commented 7 months ago

Until then, the NeMo team is happy to consider changes that make networks more amenable to Thunder, but only if they improve or at least do not regress non-Thunder-based use cases.

Got'ya. That makes sense. I guess that's also why the last item doesn't have a corresponding thunder issue. Sorry for the noise.