fix MegatronLayerPolicy to be compatible with the newest ParallelTransformerLayer

[dc3671]

Problem

Currently in DeepSpeed, will try to get client_module.attention while the actual ParallelTransformerLayer has no attention but self_attention according to this file https://github.com/microsoft/Megatron-DeepSpeed/blob/main/megatron/model/transformer.py#L927

The whole structure is:

Solution

I followed existed way by modifying version attribute to make it compatible with the newest ParallelTransformerLayer, and change attention to self_attention.

But I'm not sure whether ParallelTransformerLayer is changed or it's a new module, which means changing MegatronLayerPolicy.version is correct.

[RezaYazdaniAminabadi]

LGTM. Thanks @dc3671