Updates the export for onnx script to work with HF models

[dakinggg]

• updates the onnx export script to work directly with HF models
• fixes some onnx incompatibilities in mosaicgpt
• adds an onnx export unit test

1b training on this pr: https://wandb.ai/mosaic-ml/daniel-debug/runs/3rnsjn09?workspace=user-danielking
pytest passes: ========================================================== 145 passed, 3 skipped, 31 xfailed, 352 warnings in 311.60s (0:05:11) ===========================================================

[dskhudia]

Would the onnx model work with kv caching? It's not returning past_key_values.

[dakinggg]

No, I don't think this works with kv caching right now. I think HF has some support for doing this with https://huggingface.co/docs/transformers/v4.27.2/en/main_classes/onnx#transformers.onnx.OnnxConfigWithPast, but saw the point of this PR more as making sure that our model ops are able to be exported.