Is it possible to use the Layer Normalization optimization out of the transformer layer?

[bpmsilva]

Hello DeepSpeed Team,

I see you use two invertible operations for single GPU optimizations. One is related to the softmax, accessible through deepspeed.ops.sparse_attention.Softmax. The other is related to Layer Normalization. However, it seems there is no way to access the latter optimization. I want to use DeepSpeed Layer Normalization optimization in a project I am working on. Would that be possible? The code seems to be in a CUDA file (DeepSpeed/csrc/transformer/normalize_kernels.cu) and not in Python.

I really appreciate any help you can provide.

[RezaYazdaniAminabadi]

Hi @bpmsilva

Thanks for the interest in accessing these kernels. You are right we have some of the fine-grained modules accessible on PyTorch in python land. We can certainly bring both of the invertible functions you mentioned as separate individual modules. For that purpose, we need to just create the higher level APIs both on the C++ and PyTorch levels. I can help with designing such APIs to give proper access to such kernels, and if possible I will need your help to verify them and adding some unit tests for them.

Thanks,
Reza