Megatron-LM、NVIDIA NeMo、MegaMolBART 、model_optim_rng.pt 文件是什么?

本文涉及以下几个概念，分别是：

Megatron和Megatron-LM-v1.1.5-3D_parallelism

NVIDIA NeMo

Megatron和Megatron-LM-v1.1.5-3D_parallelism是什么？

Megatron是由NVIDIA开发的一种用于训练大规模语言模型的开源框架。它旨在提供高效的分布式训练方案，使研究人员和开发者能够训练更大规模的模型来解决自然语言处理等领域的挑战。Megatron是基于PyTorch深度学习框架构建的，并且专注于利用多个GPU或超级计算机上的分布式训练进行高性能计算。它提供了一系列高级的优化和并行化策略，以实现对超大规模模型和数据集的有效训练。

Megatron-LM-v1.1.5-3D_parallelism是一个用于大规模语言模型训练的开源软件框架。它是由NVIDIA开发的，旨在通过并行计算来加速训练过程。它是Megatron下面的一部分。

NVIDIA NeMo

1）NVIDIA NeMo 是什么？

Official 地址：GitHub - NVIDIA/NeMo: NeMo: a toolkit for conversational AI

NVIDIA NeMo Megatron是在Megatron的基础上发展起来的开源项目，由NVIDIA研究人员主导，是一个框架，用于借助简单的 Python 界面构建、训练和微调 GPU 加速的语音和自然语言理解 (NLU) 模型。使用 NeMo，开发者可以创建新的模型架构，并通过易于使用的应用编程接口 (API)，在 NVIDIA GPU 中的 Tensor Core 上使用混合精度计算对其进行训练。

NeMo Megatron 是此框架的一部分，提供并行化技术（例如 Megatron-LM 研究项目中的工作流和张量并行化），用于训练大规模语言模型。

2）使用 NeMo Megatron 进行大规模语言建模

基于 Transformer 的大规模语言模型正应用于广泛的自然语言任务，例如文本生成、摘要和聊天机器人。NeMo Megatron 能够管理训练数据，并使用多达数万亿个参数训练大规模模型。它会执行数据管理任务，例如格式化、过滤、重复数据删除和混合，使用其他方式完成这些任务可能需要花费数月时间。它还使用张量并行性在节点内扩展模型，并利用数据和工作流并行性跨节点扩展数据和模型。

NeMo Megatron 可依托数千个 GPU 提供高训练效率，使企业训练大规模 NLP 变得可行。

NeMo Megatron 可以将经过训练的模型导出到 NVIDIA Triton™ 推理服务器，以在多个 GPU 和多个节点上运行大规模 NLP 模型。

NVIDIA NeMo | NVIDIA 开发者

MegaMolBART 是什么？

代码地址：GitHub - NVIDIA/MegaMolBART: A deep learning model for small molecule drug discovery and cheminformatics based on SMILES

它是根据Chemformer改造的：GitHub - MolecularAI/Chemformer

MegaMolBART是一个基于SMILES的小分子药物发现和化学信息学的深度学习模型。MegaMolBART使用NVIDIA的Megatron 框架，旨在开发大型transformer 模型。

MegaMolBART relies on NeMo. NeMo provides a robust environment for developing and deploying deep learning models, including Megatron models. NeMo provides enhancements to PyTorch Lightning, such as hyperparameter configurabilityconfiguarbility with YAML files and checkpoint management. It also enables the development and training of large transformer models using NVIDIA’s Megatron framework, which makes multi-GPU, multi-node training with data parallelism, model parallelism, and mixed precision.

The ZINC-15 database is used for pre-training. Approximately 1.45 Billion molecules (SMILES strings) 从满足下面条件中选出:

1. molecular weight <= 500 Daltons,
2. LogP <= 5,
3. reactivity level was “reactive”,
4. purchasability was “annotated”.

During pre-processing, the compounds are filtered to ensure a maximum length of 512 characters. Train, validation, and test splits are randomly split using a seed as 99% / 0.5% / 0.5%. Data canonicalization and augmentation during training are performed using RDKIT via masking and SMILES randomization, as described previously.

Overview - NVIDIA Docs

首先我在github上看到两个issue：

Convert megatron lm ckpt to nemo：Convert megatron lm ckpt to nemo · Issue #5517 · NVIDIA/NeMo · GitHub

Cannot Convert Megatron GPT checkpoint ：Cannot Convert Megatron GPT checkpoint · Issue #14939 · huggingface/transformers · GitHub