Open tfogal opened 6 months ago
We need to add support for this. But in the meantime you should be able to get past the above error by jit
ting the nn.Module
only, which I believe for your example becomes:
model = MegatronImagen(cfg.model, trainer)
model.model = thunder.jit(model.model)
But in the meantime you should be able to get past the above error by
jit
ting thenn.Module
only
Just coming back to say: thanks! This worked like a charm.
We need to add support for this.
Agreed! Leaving this issue open to track it, but I have a workaround for now 😄 .
Description & Motivation
I'm trying to get the NeMo multimodel imagen example to use thunder, but Lightning itself does not support Thunder:
The patch to NeMo that hits this is:
Pitch
It would be great if Lightning could be updated to interoperate with Thunder such that we could pass the entire model to Thunder and Lightning could use / make sense of the returned
ThunderModule
.Alternatives
Only use Thunder on smaller pieces of the model.
Additional context
https://github.com/NVIDIA/NeMo/blob/23baa48e441ecb6cc6b49c23bf8cfc076db38bdc/examples/multimodal/text_to_image/imagen/imagen_training.py#L26 is the source for the model.
cc @borda