Add triton implementation of layer norm

[epwalsh]

Addresses #258.

Adds a triton implementation of layer norm. At the moment this builds on #274, so either one can be merged.

I've included a benchmark that you can run with:

python tests/triton/layer_norm_bench.py

I've also added a script for building wheels from AMD's fork of triton. As far as I know the script can be ran on any Linux machine (such as a Beaker interactive session on cirrascale):

./scripts/build_triton.sh

This will build and upload the wheel to s3://ai2-llm/wheels/. I've already built one and tested it on LUMI.

You can run the tests on LUMI like this:

sbatch scripts/triton-test-on-lumi.sh

At the moment the low-precision tests are failing, but at least the fp32 ones pass. I've opened an issue here with a minimal example: ROCm/triton#323

[epwalsh]

Unfortunately this did not work well on LUMI, even when running the LN in full precision. The purple run here is using AMDLayerNorm and the beige is using this triton implementation.

[epwalsh]

I'm marking this blocked until we figure out what's wrong with triton on LUMI.