NeMo Megatron-GPT 1.3B

| | |

Model Description

Megatron-GPT 1.3B 是一种基于 Transformer 的语言模型。GPT 是指一类类似于 GPT-2 和 GPT-3 的仅包含解码器的 Transformer 模型，而 1.3B 是指总的可训练参数数量为 13 亿 [1, 2]。它具有 Tensor 并行性 (TP) 为 1，Pipeline 并行性 (PP) 为 1，可以适配单个 NVIDIA GPU。

该模型是通过 NeMo Megatron 进行训练的。

入门指南

步骤 1：安装 NeMo 和相关依赖

您需要安装 NVIDIA Apex 和 NeMo。

git clone https://github.com/ericharper/apex.git
cd apex
git checkout nm_v1.11.0
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_layer_norm" --global-option="--distributed_adam" --global-option="--deprecated_fused_adam" ./

pip install nemo_toolkit['nlp']==1.11.0

或者，您可以使用预先安装了所有依赖项的 NeMo Megatron 训练 Docker 容器。

步骤 2：启动评估服务器

注意。该模型的 Tensor 并行性 (TP) 为 1，Pipeline 并行性 (PP) 为 1，可以适配单个 NVIDIA GPU。

git clone https://github.com/NVIDIA/NeMo.git 
cd NeMo/examples/nlp/language_modeling
git checkout v1.11.0
python megatron_gpt_eval.py gpt_model_file=nemo_gpt1.3B_fp16.nemo server=True tensor_model_parallel_size=1 trainer.devices=1

步骤 3：向模型发送提示！

import json
import requests

port_num = 5555
headers = {"Content-Type": "application/json"}

def request_data(data):
    resp = requests.put('http://localhost:{}/generate'.format(port_num),
                        data=json.dumps(data),
                        headers=headers)
    sentences = resp.json()['sentences']
    return sentences


data = {
    "sentences": ["Tell me an interesting fact about space travel."]*1,
    "tokens_to_generate": 50,
    "temperature": 1.0,
    "add_BOS": True,
    "top_k": 0,
    "top_p": 0.9,
    "greedy": False,
    "all_probs": False,
    "repetition_penalty": 1.2,
    "min_tokens_to_generate": 2,
}

sentences = request_data(data)
print(sentences)