megatron.model.fused_layer_norm

Description

This code is copied fron NVIDIA apex:

NVIDIA/apex

with some changes.

Classes

FusedLayerNormAffineFunction(*args, **kwargs)
MixedFusedLayerNorm(normalized_shape[, eps, ...])
RMSNorm(dim[, eps, sequence_parallel])