Papers
Topics
Authors
Recent
Gemini 2.5 Flash
Gemini 2.5 Flash
119 tokens/sec
GPT-4o
56 tokens/sec
Gemini 2.5 Pro Pro
43 tokens/sec
o3 Pro
6 tokens/sec
GPT-4.1 Pro
47 tokens/sec
DeepSeek R1 via Azure Pro
28 tokens/sec
2000 character limit reached

HLAT: High-quality Large Language Model Pre-trained on AWS Trainium (2404.10630v2)

Published 16 Apr 2024 in cs.CL and cs.LG

Abstract: Getting LLMs to perform well on the downstream tasks requires pre-training over trillions of tokens. This typically demands a large number of powerful computational devices in addition to a stable distributed training framework to accelerate the training. The growing number of applications leveraging AI/ML led to a scarcity of the expensive conventional accelerators (such as GPUs), which emphasizes the need for the alternative specialized-accelerators that are scalable and cost-efficient. AWS Trainium is the second-generation machine learning accelerator purposely built for training large deep learning models. However, training LLMs with billions of parameters on AWS Trainium is challenging due to its relatively nascent software ecosystem. In this paper, we showcase HLAT: a family of 7B and 70B decoder-only LLMs pre-trained using 4096 AWS Trainium accelerators over 1.8 trillion tokens. The performance of HLAT is benchmarked against popular open source models including LLaMA and OpenLLaMA, which have been trained on NVIDIA GPUs and Google TPUs, respectively. On various evaluation tasks, we show that HLAT achieves model quality on par with the baselines of similar model size. We also open-source all the training scripts and configurations of HLAT (https://github.com/awslabs/HLAT) and share the best practice of using the NeuronX Distributed Training (NxDT), a customized distributed training library for AWS Trainium. Our work demonstrates that AWS Trainium powered by NxDT is able to successfully pre-train state-of-the-art LLM models with high performance and cost-effectiveness.

Definition Search Book Streamline Icon: https://streamlinehq.com
References (59)
  1. TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems. https://www.tensorflow.org/ Software available from tensorflow.org.
  2. The Falcon Series of Open Language Models. arXiv:2311.16867 [cs.CL]
  3. PaLM 2 Technical Report. arXiv:2305.10403 [cs.CL]
  4. Apache Arrow. 2020. Apache Arrow, a crosslanguage development platform for in-memory analytics. https://arrow.apache.org/.
  5. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, Vol. 34. 7432–7439.
  6. Distributed Inference and Fine-tuning of Large Language Models Over The Internet. Advances in Neural Information Processing Systems 36 (2024).
  7. Language Models are Few-Shot Learners. In Advances in Neural Information Processing Systems, H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (Eds.), Vol. 33. Curran Associates, Inc., 1877–1901. https://proceedings.neurips.cc/paper_files/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf
  8. A survey on evaluation of large language models. ACM Transactions on Intelligent Systems and Technology (2023).
  9. Evaluating Large Language Models Trained on Code. (2021). arXiv:2107.03374 [cs.LG]
  10. Training Deep Nets with Sublinear Memory Cost. arXiv:1604.06174 [cs.LG]
  11. BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions. In NAACL.
  12. Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge. ArXiv abs/1803.05457 (2018). https://api.semanticscholar.org/CorpusID:3922816
  13. Training Verifiers to Solve Math Word Problems. arXiv preprint arXiv:2110.14168 (2021).
  14. Together Computer. 2023. RedPajama-Data: An Open Source Recipe to Reproduce LLaMA training dataset. https://github.com/togethercomputer/RedPajama-Data
  15. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), Jill Burstein, Christy Doran, and Thamar Solorio (Eds.). Association for Computational Linguistics, Minneapolis, Minnesota, 4171–4186. https://doi.org/10.18653/v1/N19-1423
  16. FairScale authors. 2021. FairScale: A general purpose modular PyTorch library for high performance and large scale training. https://github.com/facebookresearch/fairscale.
  17. A framework for few-shot language model evaluation. Version v0. 0.1. Sept (2021).
  18. Xinyang Geng and Hao Liu. 2023. OpenLLaMA: An Open Reproduction of LLaMA. https://github.com/openlm-research/open_llama
  19. OLMo: Accelerating the Science of Language Models. arXiv:2402.00838 [cs.CL]
  20. Deep learning with limited numerical precision. In Proceedings of the 32nd International Conference on International Conference on Machine Learning - Volume 37 (Lille, France) (ICML’15). JMLR.org, 1737–1746.
  21. NeMo: a toolkit for Conversational AI and Large Language Models. https://github.com/NVIDIA/NeMo
  22. Aligning AI With Shared Human Values. Proceedings of the International Conference on Learning Representations (ICLR) (2021).
  23. Measuring Massive Multitask Language Understanding. Proceedings of the International Conference on Learning Representations (ICLR) (2021).
  24. How good are gpt models at machine translation? a comprehensive evaluation. arXiv preprint arXiv:2302.09210 (2023).
  25. Training Compute-Optimal Large Language Models. arXiv:2203.15556 [cs.CL]
  26. Drew A. Hudson and Christopher D. Manning. 2019. GQA: A New Dataset for Real-World Visual Reasoning and Compositional Question Answering. arXiv:1902.09506 [cs.CL]
  27. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint arXiv:1705.03551 (2017).
  28. Reducing Activation Recomputation in Large Transformer Models. arXiv:2205.05198 [cs.LG]
  29. Natural Questions: a Benchmark for Question Answering Research. Transactions of the Association of Computational Linguistics (2019).
  30. BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault (Eds.). Association for Computational Linguistics, Online, 7871–7880. https://doi.org/10.18653/v1/2020.acl-main.703
  31. Pretrained language models for text generation: A survey. arXiv preprint arXiv:2201.05273 (2022).
  32. StarCoder: may the source be with you! (2023). arXiv:2305.06161 [cs.CL]
  33. TruthfulQA: Measuring How Models Mimic Human Falsehoods. arXiv:2109.07958 [cs.CL]
  34. RoBERTa: A Robustly Optimized BERT Pretraining Approach. arXiv:1907.11692 [cs.CL]
  35. Ilya Loshchilov and Frank Hutter. 2019. Decoupled Weight Decay Regularization. In International Conference on Learning Representations. https://openreview.net/forum?id=Bkg6RiCqY7
  36. Mixed precision training. arXiv preprint arXiv:1710.03740 (2017).
  37. N.A. [n. d.]. Paper under review.
  38. CrowS-Pairs: A Challenge Dataset for Measuring Social Biases in Masked Language Models. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, Online.
  39. OpenAI. 2023. GPT-4 Technical Report. ArXiv abs/2303.08774 (2023). https://arxiv.org/abs/2303.08774
  40. The RefinedWeb dataset for Falcon LLM: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.01116 (2023). arXiv:2306.01116 https://arxiv.org/abs/2306.01116
  41. Language models are unsupervised multitask learners. OpenAI blog 1, 8 (2019), 9.
  42. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. ArXiv. https://www.microsoft.com/en-us/research/publication/zero-memory-optimizations-toward-training-trillion-parameter-models/
  43. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. 3505–3506.
  44. Code llama: Open foundation models for code. arXiv preprint arXiv:2308.12950 (2023).
  45. Winogrande: An adversarial winograd schema challenge at scale. Commun. ACM 64, 9 (2021), 99–106.
  46. Noam Shazeer. 2020. Glu variants improve transformer. arXiv preprint arXiv:2002.05202 (2020).
  47. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053 [cs.CL]
  48. Beyond the Imitation Game: Quantifying and extrapolating the capabilities of language models. arXiv preprint arXiv:2206.04615 (2022).
  49. Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them. arXiv preprint arXiv:2210.09261 (2022).
  50. Spike No More: Stabilizing the Pre-training of Large Language Models. arXiv:2312.16903 [cs.CL]
  51. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 (2023).
  52. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 (2023).
  53. Attention is all you need. In Advances in Neural Information Processing Systems. 5998–6008.
  54. Transformers: State-of-the-Art Natural Language Processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations. Association for Computational Linguistics, Online, 38–45. https://www.aclweb.org/anthology/2020.emnlp-demos.6
  55. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830 (2019).
  56. Siren’s Song in the AI Ocean: A Survey on Hallucination in Large Language Models. arXiv preprint arXiv:2309.01219 (2023).
  57. A Survey of Large Language Models. arXiv preprint arXiv:2303.18223 (2023). http://arxiv.org/abs/2303.18223
  58. Pytorch FSDP: experiences on scaling fully sharded data parallel. arXiv preprint arXiv:2304.11277 (2023).
  59. Large language models for information retrieval: A survey. arXiv preprint arXiv:2308.07107 (2023).
User Edit Pencil Streamline Icon: https://streamlinehq.com
Authors (9)
  1. Haozheng Fan (4 papers)
  2. Hao Zhou (351 papers)
  3. Guangtai Huang (3 papers)
  4. Parameswaran Raman (11 papers)
  5. Xinwei Fu (6 papers)
  6. Gaurav Gupta (44 papers)
  7. Dhananjay Ram (10 papers)
  8. Yida Wang (62 papers)
  9. Jun Huan (31 papers)
Citations (1)
View on Semantic Scholar

Summary

We haven't generated a summary for this paper yet.

Ai Sparkles Streamline Icon: https://streamlinehq.com Summarize Now
X Twitter Logo Streamline Icon: https://streamlinehq.com

Tweets