Removing qk layernorms for DiT diffusion attention modules as qk layernorm is not supported for CrossAttention
State: open | Pull request: yes

* Docker Image: `nvcr.io/nvidia/nemo:24.12`

* Command: `torchrun --nproc-per-node 8 nemo/collections/diffusion/train.py --yes --factory pretrain_xl`

* Error: `CrossAttentionSubmodules.__init__() got an unexpected keyword argument q_layernorm`

* Reason: `megatron.core.transformer.attention.CrossAttentionSubmodules` does not have argument `q_layernorm` and `k_layernorm`.



<!-- pylint-output-strict-mode-false -->

beep boop 🤖: 🙏 The following files have warnings. In case you are familiar with these, please try helping us to improve the code base.

---

Your code was analyzed with PyLint. The following annotations have been identified:

```
************* Module nemo.collections.diffusion.models.dit.dit_layer_spec
nemo/collections/diffusion/models/dit/dit_layer_spec.py:194:0: C0301: Line too long (138/119) (line-too-long)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:234:0: C0301: Line too long (132/119) (line-too-long)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:249:0: C0301: Line too long (124/119) (line-too-long)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:268:0: C0301: Line too long (124/119) (line-too-long)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:288:0: C0301: Line too long (124/119) (line-too-long)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:306:0: C0301: Line too long (124/119) (line-too-long)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:362:0: C0301: Line too long (138/119) (line-too-long)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:396:0: C0301: Line too long (129/119) (line-too-long)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:416:0: C0301: Line too long (129/119) (line-too-long)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:549:0: C0301: Line too long (127/119) (line-too-long)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:552:0: C0301: Line too long (138/119) (line-too-long)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:53:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:59:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:65:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:74:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:108:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:112:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:116:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:120:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:132:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:138:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:157:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:221:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:383:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:482:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:565:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:623:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:643:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:675:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:734:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:771:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:796:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:822:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:841:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:21:0: W0611: Unused rearrange imported from einops (unused-import)
nemo/collections/diffusion/models/dit/dit_layer_spec.py:22:0: W0611: Unused jit_fuser imported from megatron.core.jit (unused-import)

-----------------------------------
Your code has been rated at 8.46/10
```

---

Mitigation guide:

* Add sensible and useful docstrings to functions and methods
* For trivial methods like getter/setters, consider adding `# pylint: disable=C0116` inside the function itself
* To disable multiple functions/methods at once, put a `# pylint: disable=C0116` before the first and a `# pylint: enable=C0116` after the last.

By applying these rules, we reduce the occurance of this message in future.

Thank you for improving NeMo's documentation!

hi @Sy-Zhang we have a change to add qk LN to cross attention. it hasn't been merged to megatron-lm yet. you can try the container nvcr.io/nvidia/nemo:cosmos.01