| --- |
| 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 |
| |
