Comet-1-Preview / README.md
Bc-AI's picture
Add Transformers support (modeling + config code), tokenizer/generation configs and model card
9df61c9 verified
|
Raw History Blame Contribute Delete
11.8 kB
---
library_name: transformers
pipeline_tag: text-generation
language:
- en
tags:
- comet
- openmodel-1
- causal-lm
- base-model
- hybrid-architecture
- custom_code
datasets:
- HuggingFaceFW/fineweb-edu
- HuggingFaceFW/fineweb
- HuggingFaceTB/finemath
- Cyrile/dataset-the-stack-v2-dedup-sub
---
# Comet-1-Preview
Comet-1-Preview is a 7.0B-parameter causal language model from OpenModel-Labs, trained from scratch with the single-file pipeline in [`train.py`](./train.py) (internal name: OpenModel-1). Three out of every four blocks are grouped-query attention; the fourth replaces attention with a gated, dilated causal convolution.
**This checkpoint is a base model.** It continues text; it does not follow chat turns. Prompt it with the beginning of the document you want, or with a few-shot pattern. See [Prompting](#prompting).
It is a preview release: no benchmark results have been published for it, and no license has been declared for the weights yet.
| | |
|---|---|
| Parameters | 7,001,436,928 |
| Weights | `model.safetensors`, float32, 28 GB |
| Context | trained on 4,096-token sequences; position table covers 8,192 |
| Tokenizer | GPT-NeoX-20B BPE, 50,277 tokens |
| Language | English web text, with maths and source code |
## Quick start
The architecture is not built into Transformers, so loading needs `trust_remote_code=True`. That runs [`modeling_comet.py`](./modeling_comet.py) and [`configuration_comet.py`](./configuration_comet.py) from this repository; they are short and worth reading first.
```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "OpenModel-Labs/Comet-1-Preview"
tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
repo,
trust_remote_code=True,
dtype=torch.bfloat16, # on older 4.x releases that reject `dtype`, use torch_dtype=
device_map="auto", # needs `accelerate`
)
inputs = tokenizer("Photosynthesis is the process by which", return_tensors="pt").to(model.device)
output = model.generate(
**inputs,
max_new_tokens=128,
do_sample=True,
temperature=0.7,
top_p=0.9,
repetition_penalty=1.1,
)
print(tokenizer.decode(output[0], skip_special_tokens=True))
```
The sampling values above are ordinary starting points, not tuned settings.
Memory: the weights take about 14 GB in bfloat16 and 28 GB in float32. Training computed in bfloat16, so bfloat16 inference matches the precision the weights were trained under. The generation cache adds roughly 0.6 GB per sequence at 4,096 tokens in bfloat16.
Requirements: a recent PyTorch 2.x and `transformers` 4.46 or newer.
## Prompting
The model was trained on plain documents separated by `<|endoftext|>`. Write the start of a document and let it continue.
These are greedy continuations (`do_sample=False`, 20 new tokens) from the released weights. `\n` marks a line break.
| Prompt | Continuation |
|---|---|
| `The capital of Australia is` | ` Canberra.\n\nThe population of Australia is about 23 million.\n\nThe currency of` |
| `Photosynthesis is the process by which` | ` green plants convert carbon dioxide and water into glucose and oxygen. This process is essential for the survival of` |
| `def is_prime(n):\n """Return True if n is a prime number."""\n` | ` if n < 2:\n return False\n for i in range(2, int` |
| `Question: What is 12 + 30?\nAnswer:` | ` 12 + 30 = 42\n\nQuestion: What is 12 + 30 + 42?\nAnswer:` |
The last row shows the usual base-model behaviour: after answering, it writes the next question itself. Stop on a string of your choosing (for example `"\nQuestion:"`) or cap `max_new_tokens`.
Things to know:
- The tokenizer adds no special tokens. `<|endoftext|>` (id 0) is the document separator and the end-of-sequence token for generation. `<|padding|>` (id 1) is the pad token, and batches are padded on the left.
- No chat template is shipped. `train.py` defines later fine-tuning stages that use a `<|user|>` / `<|assistant|>` turn format, but this checkpoint does not respond to it. Given `<|user|>\nWhat is the capital of France?\n<|endoftext|><|assistant|>\n`, its greedy continuation was `\n# How to Calculate the Area of a Circle\n\nThe area of a circle is the amount`, which is the start of an unrelated document.
- For question answering or instructions, use a few-shot pattern such as the `Question:` / `Answer:` one above.
## Architecture
```
tokens -> embedding (50,277 x 4,096)
-> 36 blocks, repeating [global, global, global, recall]
-> RMSNorm -> logits through the same embedding matrix (tied)
block: x = x + mixer_scale * mixer(RMSNorm(x))
x = x + mlp_scale * SwiGLU(RMSNorm(x))
```
| Component | Detail |
|---|---|
| Blocks | 36 in total: 27 `global`, 9 `recall` (blocks 3, 7, 11, ..., 35) |
| Width | 4,096 |
| `global` mixer | Causal grouped-query attention: 32 query heads, 8 key/value heads, head size 128. Queries and keys pass through a per-head RMSNorm, then RoPE (base 1,000,000). |
| `recall` mixer | Project 4,096 to a gate and a value of 2,048 channels each; run a depthwise causal convolution over the value (kernel 7, dilation 8); multiply by the sigmoid of the gate; project back to 4,096. |
| Feed-forward | SwiGLU, inner width 12,288, in every block |
| Residual scaling | A learned per-channel scale on each branch, initialised to 0.5 |
| Normalisation | Pre-norm RMSNorm, epsilon 1e-6 |
| Biases | None |
A `recall` block sees seven positions: the current token and the tokens 8, 16, 24, 32, 40 and 48 places back. The config lists `memory_dilations: [1, 2, 4, 8]`, and a block picks its dilation by its index modulo four. Because `recall` blocks always sit at an index that is 3 modulo 4, all nine use dilation 8 in this checkpoint.
A `global` block has about 193M parameters and a `recall` block about 176M; the embedding matrix has 206M.
## Training
This section describes the settings in `train.py`. It is not a log of the run: the upload does not include the run metadata, so the number of tokens this checkpoint actually saw is not recorded here.
**Data.** Pretraining streams FineWeb-Edu and FineWeb, then gradually mixes in maths and code. The mixture moves linearly from the first column to the second over 12B tokens.
| Source | Start | After anneal | License note in `train.py` |
|---|---|---|---|
| [HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) (`sample-100BT`) | 70% | 55% | ODC-By, Common Crawl-derived |
| [HuggingFaceFW/fineweb](https://huggingface.co/datasets/HuggingFaceFW/fineweb) (`sample-100BT`) | 30% | 20% | ODC-By, Common Crawl-derived |
| [HuggingFaceTB/finemath](https://huggingface.co/datasets/HuggingFaceTB/finemath) (`finemath-4plus`) | 0% | 15% | ODC-By, Common Crawl-derived |
| [Cyrile/dataset-the-stack-v2-dedup-sub](https://huggingface.co/datasets/Cyrile/dataset-the-stack-v2-dedup-sub) (Python 6%, JavaScript 2%, Java 1%, C++ 0.75%, C 0.25%) | 0% | 10% | per-file license review required |
**Optimisation.**
| Setting | Value |
|---|---|
| Sequence length | 4,096 tokens; documents are packed end to end with the end-of-text token between them |
| Hardware | 8 GPUs, DDP |
| Batch | 1 sequence per GPU, 16 accumulation steps: 524,288 tokens per optimizer step |
| Optimizer | AdamW, betas (0.9, 0.95), weight decay 0.1 (not applied to norm gains or residual scales) |
| Learning rate | 3e-4 peak, linear warm-up over 600M tokens, cosine decay to 3e-5 |
| Gradient clipping | 1.0 |
| Precision | float32 parameters, bfloat16 autocast |
| Token budget | 22 tokens per parameter (about 154B tokens), with an 11-day wall-clock cap on the stage |
`train.py` also contains supervised fine-tuning and SimPO preference-tuning stages. As shown under [Prompting](#prompting), this checkpoint behaves as the output of the pretraining stage.
## Files
| File | Purpose |
|---|---|
| `model.safetensors` | Weights, exactly as written by `train.py` (flat keys: `embed.weight`, `blocks.N.*`, `norm.weight`) |
| `config.json` | Architecture hyperparameters and the `auto_map` that points Transformers at the code below |
| `configuration_comet.py` | `CometConfig` |
| `modeling_comet.py` | `CometModel` (hidden states) and `CometForCausalLM` (logits, loss, `generate`) |
| `generation_config.json` | End-of-sequence and padding ids |
| `tokenizer.json`, `tokenizer_config.json`, `special_tokens_map.json` | Tokenizer (same `tokenizer.json` content as `EleutherAI/gpt-neox-20b`) |
| `train.py` | The training pipeline, including the reference `OpenModel1` module |
## Implementation notes
`modeling_comet.py` keeps the parameter names and shapes of the training module, so the checkpoint loads with no key remapping. It adds what the training loop did not need:
- **Generation cache.** `use_cache=True` returns a standard `DynamicCache`. Attention blocks store keys and values in it. `recall` blocks store their convolution input in the key slot, so beam search, batch expansion and cache cropping work on them unchanged.
- **Padding.** An `attention_mask` hides padded keys in attention, zeroes padded positions in the convolution, and numbers real tokens from position 0. A left-padded batch therefore gives the same logits as running each row alone.
- **Loss.** Pass `labels=input_ids`; the shift happens inside the model and `-100` is ignored.
- **Gradient checkpointing.** `model.gradient_checkpointing_enable()`.
Not supported: `output_attentions` (attention runs through `scaled_dot_product_attention`), legacy tuple caches, and inputs longer than 8,192 positions, which raise an error.
## Verification
What was checked before these files were published:
- **Layout.** A full-size `CometForCausalLM` has the same 353 tensor names and shapes as the header of the released `model.safetensors`.
- **Equivalence.** With identical weights on small random configurations, logits, hidden states and loss match the `OpenModel1` module from `train.py` exactly in float32 (maximum absolute difference 0).
- **Cache and padding.** Token-by-token and chunked decoding match a full forward pass to about 1e-6. Greedy `generate` with the cache matches recomputing from scratch, beam search matches with and without the cache, and left-padded batches match unbatched runs.
- **Released weights.** A float32 CPU forward pass over the full checkpoint gave perplexity 4.5 on a 68-token encyclopedic paragraph and 1.7 on a short Python function, put 98.8% on ` Tokyo` after a list of capitals ending in `The capital of Japan is`, and gave perplexity above 400,000 on random tokens. These are smoke tests on a handful of sentences, not benchmarks.
- **Versions.** The test suite passes on `transformers` 4.46.3, 4.57.3 and 5.22.2 with PyTorch 2.15.
Not covered: the full checkpoint has not been run on a GPU or in bfloat16 through this code, and the continuations in the Prompting table were produced on CPU with weights rounded through float16. Multi-GPU `device_map` placement was only exercised on a small model.
## Limitations
- It is a base model with no instruction or safety tuning. It can produce incorrect, biased or offensive text, and it will state false things fluently.
- Training data is English web text, maths and code. Expect weak results in other languages.
- Sequences longer than 4,096 tokens were never seen in training. Up to 8,192 the model runs, but quality there has not been measured.
- No evaluation results are available yet.
## Fine-tuning for chat
To continue with the fine-tuning stage that `train.py` defines, format each message as
```
<|{role}|>\n{content}\n<|endoftext|>
```
and compute the loss on assistant messages only, including their closing `<|endoftext|>`. The role markers are ordinary text, not special tokens. A generation prompt ends with `<|assistant|>\n`.