LowPrecisionLayerNorm

class composer.algorithms.LowPrecisionLayerNorm(apply_at=Event.INIT)

Replaces all instances of torch.nn.LayerNorm with class:.LPLayerNorm.

LPLayerNorm is a thin wrapper around torch.nn.LayerNorm which forces the layer to run in lower precision (torch.float16 or torch.bfloat16) if autocast is enabled. This algorithm has no effect in FP32 or DeepSpeed FP16 mode, where autocast is disabled.

This algorithm is intended to be used instead of Fused LayerNorm. They have similar behavior and performance.

Parameters

apply_at (Event) – Event where algorithm is applied.