| 1 | try: |
| 2 | from apex.normalization.fused_layer_norm import FusedLayerNorm |
| 3 | class LayerNorm(FusedLayerNorm): |
| 4 | def __init__(self, *args, pb_relax=False, **kwargs): |
| 5 | super().__init__(*args, **kwargs) |
| 6 | self.pb_relax = pb_relax |
| 7 | |
| 8 | def forward(self, x): |
| 9 | if not self.pb_relax: |
| 10 | return super().forward(x) |
| 11 | return super().forward(x / (x.abs().max().detach() / 8)) |
| 12 | except ModuleNotFoundError: |
| 13 | from sat.helpers import print_rank0 |
| 14 | print_rank0('Please install apex to use fused_layer_norm, fall back to torch.nn.LayerNorm', level='DEBUG') |