148 lines
3.8 KiB
Markdown
148 lines
3.8 KiB
Markdown
---
|
|
language:
|
|
- en
|
|
license: apache-2.0
|
|
tags:
|
|
- llama
|
|
- causal-lm
|
|
- from-scratch
|
|
- dpo
|
|
- chat
|
|
- text-generation
|
|
library_name: transformers
|
|
pipeline_tag: text-generation
|
|
model-index:
|
|
- name: Transformer-1B-Chat
|
|
results: []
|
|
---
|
|
|
|
# Transformer-1B-Chat
|
|
|
|
A **1.1 billion parameter** decoder-only language model trained **entirely from scratch** -- pretraining, supervised fine-tuning, and preference alignment -- on 8x NVIDIA H100 GPUs.
|
|
|
|
## Model Details
|
|
|
|
| Property | Value |
|
|
|---|---|
|
|
| Parameters | 1,105,827,840 (1.1B) |
|
|
| Architecture | LLaMA-style Decoder-only Transformer |
|
|
| Hidden Size | 2048 |
|
|
| Intermediate Size | 5504 (SwiGLU) |
|
|
| Layers | 22 |
|
|
| Attention Heads | 32 (Grouped Query Attention) |
|
|
| KV Heads | 8 |
|
|
| Head Dim | 64 |
|
|
| Max Sequence Length | 2048 |
|
|
| Vocab Size | 32,003 |
|
|
| Precision | BFloat16 |
|
|
|
|
### Architecture Highlights
|
|
|
|
- **RoPE** (Rotary Position Embeddings) with theta=10,000
|
|
- **Grouped Query Attention** (GQA) -- 4:1 query-to-KV head ratio for efficient inference
|
|
- **SwiGLU** Feed-Forward Network
|
|
- **RMSNorm** in a pre-norm configuration
|
|
- **Flash Attention 2** via PyTorch SDPA
|
|
|
|
## Training Pipeline
|
|
|
|
This model was built through a complete 3-stage training pipeline:
|
|
|
|
### Stage 1: Pretraining
|
|
|
|
| Detail | Value |
|
|
|---|---|
|
|
| Dataset | HuggingFaceFW/fineweb-edu (sample-10BT) |
|
|
| Tokens Trained | ~20B tokens |
|
|
| Steps | 19,070 |
|
|
| Duration | ~12.3 hours |
|
|
| Optimizer | AdamW (lr=3e-4, betas=0.9/0.95, wd=0.1) |
|
|
| Schedule | WSD (Warmup-Stable-Decay), warmup=1000 steps |
|
|
| Batch Size | 512 sequences (8 GPUs x 8 micro x 8 grad accum) |
|
|
| Final Loss | 2.43 |
|
|
| Throughput | ~338K tokens/sec |
|
|
|
|
### Stage 2: Supervised Fine-Tuning (SFT)
|
|
|
|
| Detail | Value |
|
|
|---|---|
|
|
| Dataset | HuggingFaceH4/ultrachat_200k (207,865 conversations) |
|
|
| Steps | 3,240 (2 epochs) |
|
|
| Duration | ~52 minutes |
|
|
| Optimizer | AdamW (lr=2e-5, cosine decay) |
|
|
| Batch Size | 256 sequences |
|
|
| Final Loss | 1.20 |
|
|
|
|
### Stage 3: Direct Preference Optimization (DPO)
|
|
|
|
| Detail | Value |
|
|
|---|---|
|
|
| Dataset | argilla/ultrafeedback-binarized-preferences-cleaned (60,917 pairs) |
|
|
| Steps | 952 (1 epoch) |
|
|
| Duration | ~14 minutes |
|
|
| Optimizer | AdamW (lr=5e-7, cosine decay) |
|
|
| Beta | 0.1 |
|
|
| Batch Size | 64 pairs |
|
|
| Final Loss | 0.49 |
|
|
| Final Accuracy | 72.5% (chosen preferred over rejected) |
|
|
| Final Reward Margin | 0.84 |
|
|
|
|
### Hardware
|
|
|
|
- **8x NVIDIA H100 80GB HBM3**
|
|
- **Distributed Strategy**: PyTorch DDP (DistributedDataParallel)
|
|
- **Communication**: NCCL
|
|
- **Mixed Precision**: BF16 autocast
|
|
- **Total Training Time**: ~13.5 hours (all 3 stages)
|
|
|
|
## Chat Template
|
|
|
|
The model uses a simple chat template with special tokens:
|
|
|
|
```
|
|
<|user|>
|
|
Your message here
|
|
<|end|>
|
|
<|assistant|>
|
|
Model response here
|
|
<|end|>
|
|
```
|
|
|
|
### Special Tokens
|
|
|
|
| Token | ID | Purpose |
|
|
|---|---|---|
|
|
| `<|user|>` | 32000 | Start of user turn |
|
|
| `<|assistant|>` | 32001 | Start of assistant turn |
|
|
| `<|end|>` | 32002 | End of turn |
|
|
|
|
## Limitations
|
|
|
|
- **1.1B parameters** -- smaller models have inherent limitations in reasoning depth and factual accuracy
|
|
- Trained on English data only
|
|
- May generate plausible-sounding but incorrect information
|
|
- The DPO alignment is single-epoch; additional iterations could improve quality
|
|
- Not safety-tuned beyond what the UltraFeedback dataset provides
|
|
|
|
## Training Code
|
|
|
|
The full training code is open-sourced alongside this model.
|
|
|
|
```
|
|
model/
|
|
config.py # Model and training hyperparameters
|
|
transformer.py # Full transformer implementation from scratch
|
|
data.py # Pretraining data pipeline (FineWeb-Edu)
|
|
sft_data.py # SFT data pipeline (UltraChat)
|
|
dpo_data.py # DPO data pipeline (UltraFeedback)
|
|
train.py # Pretraining script (DDP, 8-GPU)
|
|
train_sft.py # SFT script
|
|
train_dpo.py # DPO script
|
|
chat.py # Interactive chat interface
|
|
export_to_hf.py # Export to HuggingFace format
|
|
```
|
|
|
|
## License
|
|
|
|
Apache 2.0
|