Lightning-AI / Lightning-AI/lightning-thunder
Hang using thunder.jit with tokenizer in NeMo Stable Diffusion
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
### 🐛 Bug
NeMo's Stable Diffusion uses [CLIPTokenizer](https://github.com/NVIDIA/NeMo/blob/v1.23.0/nemo/collections/multimodal/modules/stable_diffusion/encoders/modules.py#L264) from HuggingFace. Adding thunder.jit to the tokenizer is causing a hang.
### To Reproduce
[tokenizer.patch](https://github.com/Lightning-AI/lightning-thunder/files/15445697/tokenizer.patch)
Steps to reproduce the behavior:
1. Apply the attached git diff to NeMo
2. Run NeMo using the command below:
```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.
Research direction
Apply the attached tokenizer.patch to NeMo and run the provided sd_train.py command to reproduce the hang around CLIPTokenizer with thunder.jit. Start from the linked NeMo modules.py tokenizer entry point and trace the execution until the hang; done means the same Stable Diffusion training command completes without hanging.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- huggingface, python, pytorch
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100