nanochat: “ChatGPT tốt nhất mà $100 có thể mua”#

1
# Model architecture trong nanochat/model.py
2
class GPT(nn.Module):
3
    def __init__(self, config):
4
        super().__init__()
5
        self.config = config
6
        # Simple, hackable architecture
7

8
# Training loop trong scripts/base_train.py
9
for step in range(max_steps):
10
    # Clean training logic
11
    loss = model(x, y)
12
    loss.backward()

1
# 8XH100 (Recommended)
2
gpu: 8x H100 80GB
3
memory: 2TB RAM
4
storage: 1TB NVMe SSD
5

6
# Alternative: 8XA100 (slower)
7
gpu: 8x A100 80GB
8
memory: 1TB RAM
9
storage: 500GB SSD

1
# GPT-style transformer trong nanochat/model.py
2
class GPT(nn.Module):
3
    def __init__(self, config):
4
        # Embedding layers
5
        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
6
        self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd))
7

8
        # Transformer blocks
9
        self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
10

11
        # Output head
12
        self.ln_f = nn.LayerNorm(config.n_embd)
13
        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

1
# Base training - Language modeling
2
python -m scripts.base_train
3

4
# Mid training - Continued pretraining
5
python -m scripts.mid_train
6

7
# Supervised fine-tuning
8
python -m scripts.sft_train
9

10
# Reinforcement learning (optional)
11
python -m scripts.rl_train

1
# Download và prepare data
2
python -m nanochat.dataset -n 180  # 180 shards for $100 tier
3

4
# Automatic data processing:
5
# 1. Download FineWeb dataset
6
# 2. Tokenize với custom BPE
7
# 3. Create training shards
8
# 4. Shuffle và batch data

1
// rustbpe - High-performance tokenizer
2
use rustbpe::Tokenizer;
3

4
impl Tokenizer {
5
    pub fn encode(&self, text: &str) -> Vec<u32> {
6
        // Efficient BPE encoding
7
    }
8

9
    pub fn decode(&self, tokens: &[u32]) -> String {
10
        // Fast decoding với caching
11
    }
12
}

1
# Automatic mixed precision
2
@torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16)
3
def forward(self, x, targets=None):
4
    # Efficient forward pass
5

6
# Gradient accumulation for memory efficiency
7
effective_batch_size = device_batch_size * gradient_accumulation_steps * world_size

1
# Configurable training parameters
2
config = {
3
    'learning_rate': 6e-4,
4
    'batch_size': 32,
5
    'sequence_length': 1024,
6
    'warmup_steps': 2000,
7
    'weight_decay': 0.1,
8
    'beta1': 0.9,
9
    'beta2': 0.95,
10
}

1
# Progressive training stages
2
stages = [
3
    {'name': 'base', 'data': 'fineweb', 'steps': 50000},
4
    {'name': 'mid', 'data': 'fineweb', 'steps': 5000},
5
    {'name': 'sft', 'data': 'smoltalk', 'steps': 3000},
6
    {'name': 'rl', 'data': 'preferences', 'steps': 1000},
7
]

1
# Built-in evaluation tasks
2
tasks = [
3
    'core',           # Common sense
4
    'arc_challenge',  # Abstract reasoning
5
    'arc_easy',       # Basic reasoning
6
    'gsm8k',         # Math problems
7
    'humaneval',     # Code generation
8
    'mmlu',          # Knowledge
9
    'chatcore',      # Conversation
10
]

1
# Add your own evaluation tasks
2
class CustomTask(Task):
3
    def evaluate(self, model, tokenizer):
4
        # Custom evaluation logic
5
        return accuracy_score

1
# Students can modify every aspect:
2
# 1. Change model architecture
3
class CustomGPT(GPT):
4
    def __init__(self, config):
5
        # Custom modifications
6

7
# 2. Implement new training techniques
8
def custom_training_loop():
9
    # Novel optimization strategies
10

11
# 3. Add new evaluation tasks
12
def custom_benchmark():
13
    # Domain-specific evaluation