Investigate RMSNorm as an alternative to LayerNorm about olmo HOT 11 CLOSED

I tried this out on the branch Torch2-RMSNorm using torch.autocast() to manually control precision. It was much slower.

However when I stopped messing with autocast it was faster:

from olmo.

Final results: https://wandb.ai/ai2-llm/petew-benchmarks-2/reports/LayerNorm-Benchmarks--VmlldzozOTI1NTA4

These results tell us that to optimize throughput we should use --model.layer_norm_type=rms when compiling and --model.layer_norm_type=low_precision when not compiling.

from olmo.

@epwalsh at this point we should also look at end task performance than absolute speedup results

Some of these decisions in bloom were made based on downstream eval and not just throughput: https://arxiv.org/pdf/2210.15424.pdf

from olmo.

like I think SwiGLU and alibi are non-negotiable based on end task performance

from olmo.

Absolutely

from olmo.

Only problem with ALiBi is that it incurs a significant performance hit since it doesn't work with the current Flash Attention implementation.

from olmo.

Is low_precision stable enough for use? I thought doing LN in 32-bits was a major stability hack in BLOOM.

from olmo.

Is low_precision stable enough for use? I thought doing LN in 32-bits was a major stability hack in BLOOM.

That remains to be seen.

from olmo.

What about Alibi vs. Rope?

from olmo.

performance or throughput?

from olmo.

from olmo.