ChipNeMo: Domain-Adapted LLMs for Chip Design

In plain words

General-purpose chatbots know a little about everything, but a chip engineer asks about their company's own designs, internal tools and bug history, which no public model has seen. The authors' answer is to keep a capable open model, teach it the company's domain text by continuing its training on that text, and then teach it to follow instructions and to look up internal documents when it answers.

The idea has four parts. First, the model's vocabulary gets extra word pieces for hardware jargon, so long technical names cost fewer tokens. Second, the model reads a large amount of internal and public chip material as if it were ordinary text. Third, it is trained on a small set of question-and-answer examples plus general chat data, so it answers in a useful form. Fourth, a search index of internal documents, searched with a retriever that was itself tuned on the company's data, supplies the facts the model would otherwise have to guess.

The authors report that the largest version beats GPT-4 on two of their three test tasks, and is close on the third. These are the authors' own tests on internal data, which outsiders cannot rerun.

TL;DR

• Starts from LLaMA2 (7B, 13B, 70B) and runs domain-adaptive pretraining (DAPT) on about 24B tokens of chip-design text, roughly one epoch §2.2 app. Training Data.
• Extends the LLaMA2 tokenizer with about 9K domain tokens, initialised from the original tokenizer §2.1 §3.1.
• Aligns the model with SFT or SteerLM on general chat data plus 1.4K domain examples §2.3 app. Model Alignment.
• Fine-tunes the e5-small retriever on 3,000 synthetic query-passage samples and uses it for RAG §2.4 app. Domain-Adaptive Retrieval Model.
• Author-reported: ChipNeMo-70B beats GPT-4 on the engineering chatbot and on EDA script generation, and is competitive on bug summarization §1 §3.7 §3.8.

Problem

Off-the-shelf LLMs lack the internal APIs, tool scripts and bug history of a given chip company, and sending proprietary design data to a third-party API is a security concern §1. Training a domain model from scratch would need millions of GPU hours, so the authors adapt an existing foundation model instead §1.

Three applications are targeted: an engineering assistant chatbot, EDA script generation (Python and Tcl), and bug summarization and analysis §1. The paper assumes the company holds enough internal text for continued pretraining, and that the public general-data mix preserves general ability [app. Training Data].

Method

The pipeline has four stages, shown in fig. 1 (described here in words: a base model goes through DAPT, then alignment; RAG is added at inference).

Stage 1, domain-adaptive tokenizer §2.1. The authors train a new tokenizer on domain data and keep only tokens that are absent from the base vocabulary and rare in general text. They add these to the base tokenizer. Each new token's input embedding is the average of the embeddings of the pieces the original tokenizer produces for it. The output-layer weights of new tokens are set to zero. The threshold for "rare" is not stated.

Stage 2, DAPT §2.2. Each LLaMA2 model is initialised from its base weights and trained with the standard next-token objective on the domain corpus. The corpus mixes proprietary chip material (design source, documentation, verification, bug summaries, other) with a sampled slice of public Wikipedia and GitHub text. Public data is about 9.2% of training tokens app. Training Data. Code is downsampled and documentation upsampled across 2 to 4 epochs of the blend; the public part is seen for one epoch [app. Data Blend]. Training uses NeMo with tensor parallelism and FlashAttention §2.2.

Stage 3, alignment §2.3. Models are fine-tuned on a chat template with system, user and assistant turns, and loss is applied only to the assistant tokens app. Model Alignment. Two variants are used:

• SFT: one epoch over 128K general chat samples plus the 1.4K domain samples, shuffled together §2.3.
• SteerLM: an attribute model (initialised from LLaMA2-13B) is trained on HelpSteer and OASST and labels the domain data. The attribute-conditioned model is then trained on 56K OASST samples plus the 1.4K domain samples for two epochs app. Model Alignment.

Repeating the domain SFT set for more than one epoch led to overfitting, and DAPT applied to an already chat-aligned model badly degraded alignment §2.3.

Stage 4, RAG with a domain retriever §2.4 [app. Domain-Adaptive Retrieval Model]. The retriever is e5-small-unsupervised fine-tuned with contrastive learning on 3,000 generated samples, using the Tevatron framework. Each sample needs positives and hard negatives. They are generated as follows:

# Build one contrastive training sample  [app. Sampling Procedure]
# Input:  document corpus (passages), query generator LM, retriever
# Output: (query, positives, hard negatives)

def make_sample(corpus, query_model, retriever, top_n):
    # Pick a random passage from the corpus  [app. Sampling Procedure]
    passage = random_choice(corpus)
    # Ask an LM for a query this passage answers  [app. Sampling Procedure]
    query = query_model.generate_query(passage)
    # Fetch the top-N passages as candidate hard negatives
    # [app. Sampling Procedure]
    candidates = retriever.top_n(query, corpus, top_n)
    # Drop candidates the same LM judges to answer the query
    # [app. Sampling Procedure]
    negatives = [
        c for c in candidates if not query_model.answers(query, c)
    ]
    # Top up with random passages if too few negatives remain
    # [app. Sampling Procedure]
    while len(negatives) < top_n:
        negatives.append(random_choice(corpus))
    return query, passage, negatives

At answer time, the query is embedded, the top 8 passages are fetched from a store of about 67K passages (each about 512 characters from 1.8K documents), and they are placed in the prompt before the question §3.6. The prompt layout is not stated.

Hyperparameters

Results

Author-reported. Setup: internal chip-design multiple-choice benchmarks (AutoEval, 94 to 227 questions each), public academic benchmarks, a 88-question RAG chatbot benchmark rated by human experts on a 7-point scale, EDA script benchmarks (automatic and human), and a 30-bug summarization set §3 app. Domain Evaluation Benchmarks.

Domain multiple-choice (AutoEval), accuracy app. Auto Evaluation Results:

General academic benchmarks, 70B app. Auto Evaluation Results: MMLU 68.6 (LLaMA2) vs 69.4 (ChipNeMo); HellaSwag 83.6 vs 85.1; ARC-e 79.6 vs 80.9; RACE-H 48.9 vs 48.9; Winogrande 77.6 vs 78.0. The 7B and 13B models lose a little on MMLU (44.6 vs 45.7; 53.4 vs 55.4). The authors describe this as a slight degradation on open-domain tasks §3.2.

Code app. Coding Evaluation Results: pass rate with greedy decoding. HumanEval, 70B: 28.0 (LLaMA2) vs 30.5 (ChipNeMo). VerilogEval-Human, 70B: 30.8 vs 27.6. VerilogEval-Machine, 70B: 51.0 vs 53.8. At 13B, VerilogEval-Human rises from 9.0 to 22.4.

Engineering chatbot, human rating on a 7-point Likert scale [app. Engineering Assistant Chatbot Human Evaluation]:

"(G)" marks the variant aligned on general chat data only. The retriever (fine-tuned) beats e5-small-unsupervised by about 2x and Sentence Transformer by about 30% on top-8 hit accuracy §3.6; the per-category values are in a figure whose numbers are not stated in the text.

EDA script generation, Pass@5 automatic and 0 to 10 human score app. EDA Script Generation Evaluation:

Benchmark sizes: Python Easy 146, Medium 28, Hard 25; Tcl Easy 708, Medium 27, Hard 25 §3.7 tab. EDA Script Generation Evaluation Benchmarks.

Bug summarization §3.8: rated on a 7-point scale on a 30-bug set. GPT-4 is best on all three sub-tasks; ChipNeMo-70B-Steer improves on LLaMA2-70B-Chat. Per-task values are shown only in a figure, and the text does not give them as numbers §3.8.

Training cost §3.5 tab. Training cost of LLaMA2 models in A100 GPU hours: DAPT costs 2,620 (7B), 4,940 (13B) and 20,500 (70B) A100-hours, versus the published pretraining cost of 184,320, 368,640 and 1,720,320. Alignment costs 90, 160 and 840 A100-hours. The authors state DAPT is under 1.5% of pretraining cost.

Paper vs code

No official code. The paper links no repository or project page. Three GitHub repository searches (chipnemo, chipnemo org:NVIDIA, and chip design llm domain adapted) returned no results, so no code was read and no code_commit is recorded. All statements above come from the paper.

Pitfalls

• Intro headline numbers do not all match the tables. The intro gives a "score of 6.0" for the chatbot §1. The table average for ChipNeMo-70B-Steer(G) with RAG is 5.68; 6.02 is the retrieval-hit cell. Likewise, "more than 70% correctness" on simple EDA scripts §1 matches the RAG rows (77% and 84%), not the no-RAG rows (19% to 49%).
• Text and table label mismatch on RAG gains §3.6 [app. Engineering Assistant Chatbot Human Evaluation]. The text says RAG raises ChipNeMo-70B-Steer by 0.56 and says it beats LLaMA2-70B-Chat by 3.31 (no RAG) and 1.81 (RAG). In the table those differences hold for the (G) row (5.12 - 1.81 = 3.31; 5.68 - 3.86 = 1.82). For the non-(G) row the RAG gain is 0.46. Verify which variant a number refers to.
• Token counts disagree §2.2 app. Training Data. The text says the corpus has 23.1 billion tokens. The table totals are 25.3 B (raw) and 24.1 B (training blend). The 9.2% public share matches neither total exactly.
• Tokenizer details. The text says the tokenizer is extended by "about 9K" tokens §3.1. The adaptation steps are listed as four §2.1. The "rare in general data" threshold and the number of tokens actually kept are not stated. The tokenizer-efficiency gain is 1.6% to 3.3% §3.1, and the intro says "up to 3.3%" §1.
• Loss spikes are unexplained. DAPT shows spikes at similar steps across model sizes. The authors attribute them to bad data and did not remove them §2.2. A larger learning rate (3e-4) showed large loss spikes early in training, and degraded language and domain benchmarks except coding app. Learning Rate §3.4.
• Training data is not released. Most of the corpus is internal and cannot be rebuilt. The public share is Wikipedia plus GitHub code (the paper cites the Stack dataset), about 2.2B tokens in the training blend app. Training Data. Reproduction requires your own domain corpus.
• Tiny evaluation sets. Hard EDA sets have 25 items each, the bug set has 30 bugs, and the chatbot set has 88 questions §3.6 §3.7 §3.8. Human scores are single-rater 0 to 10 or 7-point ratings with no reported agreement. Treat differences of a few tenths as noise unless replicated.
• Benchmark availability. The AutoEval, EDA and chatbot benchmarks are internal. The paper does not release them. The Python EDA tool is proprietary §3.7.
• Retriever training is under-specified. Learning rate, epochs, batch size, negatives per query and the query/passage prompt prefixes for e5 are not stated. The sampling procedure uses Vicuna and Sentence Transformer; the authors say LLaMA2 and BM25 could replace them for commercial use, but they did not report that variant app. Dataset Sampling Procedure.
• Alignment data license. The general chat data is described as commercially usable app. Alignment Data. Each dataset's own terms still govern use, and the paper does not list them.
• Chat template. SFT uses <extra_id_0> / <extra_id_1> role tags, and loss is masked on system and user tokens app. Model Alignment. Using a different template changes the result.
• Overfitting. Repeating the 1.4K domain samples for more than one epoch caused the model to repeat unrelated training answers §2.3.
• Domain DAPT on a chat model fails. Running DAPT on LLaMA2-Chat badly degraded alignment §2.3. DAPT must start from base weights.
• Costs the abstract hides. RAG adds retrieval latency and an index build over about 67K passages. DAPT at 70B needs 20,500 A100-hours §3.5. The authors also note that LoRA (26.4M and 211.2M trainable parameters) was clearly worse than full DAPT on in-domain tasks §3.4.

Lineage

• Builds on: DAPT §2.2, e5-small-unsupervised §2.4, SteerLM and HelpSteer for alignment §2.3, Tevatron for retriever training §2.4, LLaMA2 base models §2.2.
• Compared against: LLaMA2-70B-Chat, GPT-4 and GPT-3.5 (academic scores cited from OpenAI) §3 app. Auto Evaluation Results.
• Superseded by: not stated in this card's wiki content.