Lightning-AI / Lightning-AI/lightning-thunder

Hang using thunder.jit with tokenizer in NeMo Stable Diffusion

Open
#462 7 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug nemo
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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.