Lightning-AI / Lightning-AI/lightning-thunder
TypeError: Missing a required argument with thunder.jit in NeMo SD ResBlock
@athitten is already working on this.
Since Jun 10, 2024.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
### 🐛 Bug
Adding `thunder.jit` to `ResBlock` in the UNet stage of NeMo SD is raising an error. From looking at the ResBlock call in NeMo code, the class is called correctly with right arguments. In-spite of that its unsure why thunder is raising this error.
`Encountered exception TypeError: missing a required argument: 'emb' while tracing ResBlock`
Stack trace of the error can be found here:
[resblock_error.log](https://github.com/user-attachments/files/15685760/resblock_error.log)
### To Reproduce
Steps to reproduce the behavior:
1. Pull the appropriate NeMo docker image
2. Apply the git patch:
[resblock.patch](https://github.com/user-attachments/files/15685773/resblock.patch)
3. Run Stable Diffusion with the command:
```python
python examples/multimodal/text_to_image/stable_diffusion/sd_train.py trainer.precision=16 trainer.num_nodes=1 trainer.devices=1 ++exp_manager.max_time_per_run=00:00:03:00 trainer.max_steps=20 model.micro_batch_size=1 model.global_batch_size=1 model.data.synthetic_data=True exp_manager.exp_dir=/workspace/TestData/multimodal/stable_diffusion_train model.inductor=False model.cond_stage_config._target_=nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.FrozenCLIPEmbedder ++model.cond_stage_config.version=openai/clip-vit-large-patch14 ++model.cond_stage_config.max_length=77 ~model.cond_stage_config.restore_from_path ~model.cond_stage_config.freeze ~model.cond_stage_config.layer model.unet_config.from_pretrained=null model.first_stage_config.from_pretrained=null model.unet_config.use_flash_attention=False model.unet_config.attention_resolutions=\[1\] model.unet_config.channel_mult=\[1\]`
```
cc: @tfogal
cc @tfogal
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.