README.md · nvidia/Nemotron-4-340B-Instruct at 950e64c0d974f946cac9720e932893312574c2b2

metadata

license: other
license_name: nvidia-open-model-license
license_link: LICENSE

Nemotron-4-340B-Instruct

License

NVIDIA Open Model License

Model Overview

Nemotron-4-340B-Instruct is a large language model (LLM) which is a fine-tuned version of the Nemotron-4-340B-Base base model, optimized for English single and multi-turn chat use-cases. The base model was pre-trained on a corpus of 8 trillion tokens consisting of a diverse assortment of English based texts, 40+ coding languages, and 50+ natural languages.

Subsequently the Nemotron-4-340B-Instruct model went through additional alignment steps including:

Supervised Fine-tuning (SFT)
Direct Policy Optimization (DPO)
Additional in-house alignment techniques (Publication work in progress)

This results in a final model that is aligned for human chat preferences, improvements in mathematical reasoning, coding and instruction following.

This model is ready for commercial use.

Model Developer: NVIDIA

Model Input: Text Input Format: String Input Parameters: One-Dimensional (1D)

Model Output: Text Output Format: String Output Parameters: 1D

Model Dates: Nemotron-4-340B-Instruct was trained between December 2023 and May 2024

Data Freshness: The pretraining data has a cutoff of June 2023

Required Hardware

BF16 Inference:

8x H200 (1x H200 Node)
16x H100 (2x H100 Nodes)
16x A100 (2x A100 Nodes)

FP8 Inference:

8x H100 (1x H100 Node)

Model Architecture:

The base model, Nemotron-4-340B, was trained with a global batch-size of 2304, a sequence length of 4096 tokens, uses Grouped-Query Attention (GQA), and RoPE positional embeddings.

Architecture Type: Transformer Decoder (auto-regressive language model)

Software Integration

Supported Hardware Architecture Compatibility: NVIDIA H100, A100 80GB, A100 40GB

Usage

We will spin up an inference server and then call the inference server in a python script. Let’s first define the python script call_server.py

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

def text_generation(data, ip='localhost', port=None):
    resp = requests.put(f'http://{ip}:{port}/generate', data=json.dumps(data), headers=headers)
    return resp.json()


def get_generation(prompt, greedy, add_BOS, token_to_gen, min_tokens, temp, top_p, top_k, repetition, batch=False):
    data = {
        "sentences": [prompt] if not batch else prompt,
        "tokens_to_generate": int(token_to_gen),
        "temperature": temp,
        "add_BOS": add_BOS,
        "top_k": top_k,
        "top_p": top_p,
        "greedy": greedy,
        "all_probs": False,
        "repetition_penalty": repetition,
        "min_tokens_to_generate": int(min_tokens),
        "end_strings": ["<|endoftext|>", "<extra_id_1>", "\x11", "<extra_id_1>User"],
    }
    sentences = text_generation(data, port=1424)['sentences']
    return sentences[0] if not batch else sentences

PROMPT_TEMPLATE = """<extra_id_0>System

<extra_id_1>User
{prompt}
<extra_id_1>Assistant
"""

question = "Write a poem on NVIDIA in the style of Shakespeare"
prompt = PROMPT_TEMPLATE.format(prompt=question)
print(prompt)

response = get_generation(prompt, greedy=True, add_BOS=False, token_to_gen=1024, min_tokens=1, temp=1.0, top_p=1.0, top_k=0, repetition=1.0, batch=False)
print(response)

Given this python script, we will create a bash script, which spins up the inference server within the NeMo container and calls the python script call_server.py. The bash script nemo_inference.sh is as follows,

WEB_PORT=1424

depends_on () { HOST=$1 PORT=$2 STATUS=$(curl -X PUT http://$HOST:$PORT >/dev/null 2>/dev/null; echo $?) while [ $STATUS -ne 0 ] do echo "waiting for server ($HOST:$PORT) to be up" sleep 10 STATUS=$(curl -X PUT http://$HOST:$PORT >/dev/null 2>/dev/null; echo $?) done echo "server ($HOST:$PORT) is up running" }

echo "output filename: $OUTPUT_FILENAME"

/usr/bin/python3 /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py
gpt_model_file=$NEMO_FILE
pipeline_model_parallel_split_rank=0
server=True tensor_model_parallel_size=8
trainer.precision=bf16 pipeline_model_parallel_size=4
trainer.devices=8
trainer.num_nodes=4
web_server=False
port=${WEB_PORT} & SERVER_PID=$!

readonly local_rank="${LOCAL_RANK:=${SLURM_LOCALID:=${OMPI_COMM_WORLD_LOCAL_RANK:-}}}"
if [ $SLURM_NODEID -eq 0 ] && [ $local_rank -eq 0 ]; then
    depends_on "0.0.0.0" ${WEB_PORT}

    echo "start get json"
    sleep 5

    echo "SLURM_NODEID: $SLURM_NODEID"
    echo "local_rank: $local_rank"
    /usr/bin/python3 call_server.py
    echo "clean up dameons: $$"
    kill -9 $SERVER_PID
    pkill python
fi
wait

3, We can launch the nemo_inferece.sh with a slurm script defined like below, which starts a 4-node job for the model inference.

#!/bin/bash #SBATCH -A SLURM-ACCOUNT #SBATCH -p SLURM-PARITION #SBATCH -N 4 # number of nodes #SBATCH -J generation
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8 set -x

read -r -d '' cmd <<EOF bash nemo_inference.sh EOF

srun -o $OUTFILE -e $ERRFILE --container-image="$CONTAINER" $MOUNTS bash -c "${cmd}"

Intended use

Nemotron-4-340B-Instruct is a chat model intended for use in over 50+ natural and coding languages. For best performance on a given task, users are encouraged to customize the chat model using the NeMo Framework suite of customization tools including Parameter-Efficient Fine-Tuning (P-tuning, Adapters, LoRA), and SFT/Steer-LM/RLHF.

Red Teaming:

TO BE UPDATED BASED ON RED TEAMING PICs + LEGAL REVIEW

Evaluation Results

MT-Bench (GPT-4-Turbo)

Evaluated using select datasets from the Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena

total	writing	roleplay	extraction	stem	humanities	reasoning	math	coding	turn 1	turn 2
8.22	8.70	8.70	9.20	8.75	8.95	6.40	8.40	6.70	8.61	7.84

IFEval

Evaluated using the Instruction Following Eval (IFEval) introduced in Instruction-Following Evaluation for Large Language Models.

Prompt-Strict Acc	Instruction-Strict Acc
79.9	86.1

MMLU

Evaluated using the Multi-task Language Understanding benchmarks as introduced in Measuring Massive Multitask Language Understanding

MMLU 0-shot
78.7

GSM8K

Evaluated using the Grade School Math 8K (GSM8K) bechmark as introduced in Training Verifiers to Solve Math Word Problems.

GSM8K 0-shot
92.3

HumanEval

Evaluated using the HumanEval benchmark as introduced in Evaluating Large Language Models Trained on Code.

HumanEval 0-shot
73.2

Arena Hard

Evaluated using the Arena-Hard Pipeline from the LMSys Org.

Arena Hard
54.2

AlpacaEval 2.0 LC

Evaluated using the AlpacaEval 2.0 LC (Length Controlled) as introduced in the paper: Length-Controlled AlpacaEval: A Simple Way to Debias Automatic Evaluators

AlpacaEval
54.2

MBPP

Evaluated using the MBPP Dataset as introduced in the Program Synthesis with Large Language Models paper.

MBPP
75.4

TFEval

Evaluated using the CantTalkAboutThis Dataset as introduced in the CantTalkAboutThis: Aligning Language Models to Stay on Topic in Dialogues paper.

Distractor F1	On-topic F1
81.7	97.7

Limitations

The base model was trained on data that contains toxic language and societal biases originally crawled from the internet. Therefore, the model may amplify those biases and return toxic responses especially when prompted with toxic prompts. The model may generate answers that may be inaccurate, omit key information, or include irrelevant or redundant text producing socially unacceptable or undesirable text, even if the prompt itself does not include anything