Z Loss in CE

[Fr0do]

🚀 The feature, motivation and pitch

Often used in pretraining of LMs for stabilization, i.e. the recent Chameleon & PaLM.

Alternatives

flash-attn has implementations of abovementioned features, however, does not support fusing with linear head.

Additional context

No response

[ByronHsu]

Legit ask! We have tracked smooth label at https://github.com/linkedin/Liger-Kernel/issues/81. I modify the title for only Z loss to prevent duplication.

[Tcc0403]

@ByronHsu #take To support z loss, I just need a little add-ons to #198. I'll work on it after merging label_smoothing PR.