NVIDIA Megatron-LM Powers 172 Billion Parameter LLM for Japanese Language Proficiency

NVIDIA's Megatron-LM is at the forefront of a significant development in the field of Natural Language Processing (NLP) by powering a large language model (LLM) with 172 billion parameters, aimed at enhancing Japanese language processing capabilities. This initiative is part of the Generative AI Accelerator Challenge (GENIAC) project, as reported by NVIDIA.

Advancements in Generative AI

Generative AI has revolutionized content creation, outperforming traditional machine learning methods. Large language models have spearheaded this transformation, finding applications in customer support, voice assistance, text summarization, and translation. However, many existing models are predominantly trained on English data, leading to deficiencies in other languages, including Japanese. For instance, only 0.11% of the GPT-3 corpus is comprised of Japanese data, necessitating the development of more robust Japanese language models.

GENIAC and the LLM-jp Project

The Ministry of Economy, Trade and Industry (METI) launched the GENIAC initiative to bolster model development capabilities in Japan. This initiative supports companies and researchers by providing computational resources, fostering industry collaboration, and evaluating model performance. The LLM-jp project, developed under this initiative, focuses on creating an open model with 172 billion parameters, emphasizing Japanese language proficiency. This model, the largest of its kind in Japan during its development phase, aims to share its development insights broadly.

Training with NVIDIA Megatron-LM

NVIDIA Megatron-LM, a research-oriented framework, facilitates the training of LLMs at unprecedented speeds. It utilizes Megatron-Core, an open-source library optimized for GPU performance, supporting various parallelism techniques. This framework is compatible with NVIDIA Tensor Core GPUs and supports FP8 precision, enhancing training efficiency with the NVIDIA Hopper architecture.

The LLM-jp 172B model is being trained using 2.1 trillion tokens from a multilingual corpus, primarily Japanese and English. Training is conducted on NVIDIA H100 Tensor Core GPUs with Google Cloud A3 Instances, utilizing FP8 hybrid training. This approach incorporates techniques like z-loss and batch-skipping to stabilize training and utilizes flash attention for speed enhancements.

Performance and Results

Training of the LLM-jp 172B model is ongoing, with evaluations conducted periodically to ensure accuracy in Japanese and English tasks. The transition from BF16 to FP8 hybrid precision has significantly increased training throughput, demonstrating a 1.4x speed acceleration from 400 to 550 TFLOP/s. This improvement highlights FP8 hybrid training as a promising method for enhancing large-scale model pretraining efficiency.

The ongoing development of the LLM-jp 172B model not only aims to advance Japanese language capabilities but also sets a precedent for future multilingual AI models. The project underscores the importance of efficient training frameworks like Megatron-LM in accelerating generative AI research and development.