Architektur im Überblick

Quelldateien

Diese Seite wurde aus den folgenden Quelldateien erstellt:

• nanochat/gpt.py
• nanochat/engine.py
• nanochat/dataloader.py
• nanochat/checkpoint_manager.py
• nanochat/optim.py
• nanochat/common.py
• nanochat/report.py
• nanochat/execution.py
• nanochat/tokenizer.py
• nanochat/flash_attention.py
• pyproject.toml
• tasks/arc.py
• tasks/mmlu.py
• tasks/gsm8k.py
• nanochat/fp8.py
• tasks/common.py
• tasks/smoltalk.py
• scripts/chat_rl.py
• tasks/humaneval.py
• nanochat/dataset.py

NanoChat ist ein minimalistisches, aber leistungsstarkes Framework zum Training und zur Feinabstimmung von GPT-ähnlichen Sprachmodellen. Die Architektur zeichnet sich durch eine klare Trennung der Verantwortlichkeiten aus, wobei jedes Modul eine spezifische Aufgabe im Trainingspipeline übernimmt. Das System ist für verteiltes Training (DDP) konzipiert und nutzt fortschrittliche Optimierungstechniken wie den Muon-Optimizer für 2D-Matrixparameter.

Systemarchitektur Übersicht

Die folgende Grafik zeigt die Hauptmodule von NanoChat und deren Abhängigkeiten:

正在加载图表渲染器...

Erklärung der Architekturkomponenten:

1. Datenpipeline (links): Der DataLoader implementiert einen Best-Fit Cropping Algorithmus, der Dokumente aus Parquet-Dateien lädt und tokenisiert. Die 100% GPU-Auslastung ohne Padding wird durch intelligente Dokumentenauswahl erreicht.

2. Modell-Kern (Mitte): Das GPT-Modell nutzt Transformer-Blöcke mit Sliding-Window-Attention und einer präzisen FLOPs-Berechnung für das Training.

3. Training-Engine (rechts): Der MuonAdamW-Optimizer kombiniert Muon für 2D-Matrizen mit AdamW für Embeddings und Skalare.

4. Infrastruktur (unten): Die automatische DType-Erkennung wählt basierend auf der GPU-Capability zwischen BF16 und FP32.

---

Modellarchitektur und FLOPs-Berechnung

Verantwortlichkeiten und Grenzen

Das GPT-Modul (nanochat/gpt.py) ist das Herzstück der Architektur und verantwortlich für:

• Definition der Transformer-Architektur mit konfigurierbaren Hyperparametern
• FLOPs-Schätzung für Trainingseffizienz-Analysen
• Optimizer-Setup mit parametrisierten Lernraten für verschiedene Parametertypen

Nicht im Verantwortungsbereich: Datenladen, Tokenisierung, Checkpoint-Persistenz (diese werden von separaten Modulen übernommen).

Eintrittspunkte und Kerna APIs

Die Hauptklasse GPT bietet folgende zentrale Methoden:

FLOPs-Berechnung im Detail

Die FLOPs-Schätzung basiert auf der Chinchilla-Paper-Methodik mit folgenden Komponenten:

python
1# Kernformel (vereinfacht)
2num_flops_per_token = 6 * (nparams - nparams_exclude) + attn_flops

Berechnungsfaktoren:

1. Matmul-Parameter: Jeder Gewichtungsparameter trägt 6 FLOPs bei (2 Forward + 4 Backward)
2. Ausgeschlossene Parameter: Embeddings und Skalare werden nicht als Matmul-FLOPs gezählt
3. Attention-FLOPs: 12 * h * q * effective_seq pro Layer, wobei effective_seq durch Sliding Windows begrenzt wird

python
1# Sliding Window Berücksichtigung
2for window_size in self.window_sizes:
3    window = window_size[0]  # (left, right) tuple
4    effective_seq = t if window < 0 else min(window, t)
5    attn_flops += 12 * h * q * effective_seq

Quelle: nanochat/gpt.py:331-336

Parametergruppen und Optimizer-Konfiguration

Die Optimizer-Setup-Methode unterteilt alle Modellparameter in sieben Gruppen mit unterschiedlichen Lernraten:

LR-Skalierung: Die Lernrate für AdamW-Parameter wird mit (model_dim / 768) ** -0.5 skaliert, um für verschiedene Modellgrößen konsistentes Training zu gewährleisten.

Quelle: nanochat/gpt.py:383-385

Fehlerbehandlung und Randbedingungen

• Assertion bei Parameteranzahl: Zeile 381 verifiziert, dass alle Parameter erfasst werden
• Distributed Detection: get_dist_info() entscheidet zwischen DistMuonAdamW und MuonAdamW
• Initial LR Setting: Alle Parametergruppen erhalten initial_lr für LR-Scheduling

---

Optimierungsstrategien Muon und AdamW

Verantwortlichkeiten und Design-Philosophie

Der MuonAdamW-Optimizer implementiert eine hybride Optimierungsstrategie:

• Muon für 2D-Matrixparameter (Attention, MLP-Gewichte)
• AdamW für Embeddings, LM-Head und 0D/1D-Parameter

Design-Begründung: Muon orthogonalisiert Updates durch Newton-Schulz-Iteration, was für 2D-Matrizen effektiv ist, aber für Embeddings ungeeignet.

Fused Kernels für Effizienz

Beide Optimizer-Teile nutzen fused Kernels, um Python-Overhead zu eliminieren:

AdamW Fused Step

Der adamw_step_fused-Kernel führt in einem kompilierten Graphen aus:

python
1# Reihenfolge der Operationen
21. Weight Decay: p.mul_(1 - lr_t * wd_t)
32. Momentum Update: exp_avg.lerp_(grad, 1 - beta1_t)
43. Bias Correction: bias1 = 1 - beta1_t ** step_t
54. Parameter Update: p.add_(exp_avg / denom, alpha=-step_size)

0-D CPU Tensoren: Alle Hyperparameter werden als 0-D CPU-Tensoren gespeichert, um Rekompilierung bei Wertänderungen zu vermeiden.

Quelle: nanochat/optim.py:36-37

Muon Fused Step

Der muon_step_fused-Kernel implementiert eine vierstufige Pipeline:

正在加载图表渲染器...

Polar Express Details:

python
1# Iterative Orthogonalisierung (5 Iterationen)
2X = g.bfloat16()
3X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01 + 1e-6)
4for a, b, c in polar_express_coeffs[:ns_steps]:
5    if tall_matrix:
6        A = X.mT @ X
7        B = b * A + c * (A @ A)
8        X = a * X + X @ B
9    else:  # wide matrix
10        A = X @ X.mT
11        X = a * X + B @ X

Quelle: nanochat/optim.py:115-126

Zustandsverwaltung

Der Muon-Optimizer speichert zwei Puffer pro Parametergruppe:

Factored Second Moment: Die Form hängt vom Seitenverhältnis der Matrix ab – bei "tall" Matrizen wird zeilenweise, bei "wide" Matrizen spaltenweise reduziert.

Quelle: nanochat/optim.py:250-254

Cautious Update Mechanismus

Der Muon-Update nutzt eine Maske, um Updates zu unterdrücken, die das Vorzeichen des Parameters ändern würden:

python
1mask = (g * stacked_params) >= 0
2stacked_params.sub_(lr * g + lr * wd * stacked_params * mask)

Dies verhindert destabilisierende Vorzeichenwechsel während des Trainings.

Quelle: nanochat/optim.py:145-146

---

Datenpipeline und Tokenisierung

Verantwortlichkeiten des DataLoaders

Der tokenizing_distributed_data_loader_with_state_bos_bestfit ist verantwortlich für:

• Effiziente Batch-Erstellung mit 100% Token-Utilisierung (kein Padding)
• DDP-Sharding für verteiltes Training
• Resume-Funktionalität mit exakter Positionswiederherstellung

Nicht im Verantwortungsbereich: Tokenizer-Training, Parquet-Dateierstellung.

Best-Fit Cropping Algorithmus

Der Algorithmus maximiert die GPU-Auslastung durch intelligente Dokumentenauswahl:

正在加载图表渲染器...

Algorithmus-Schritte (pro Zeile):

1. Buffer-Check: Sicherstellen, dass ≥buffer_size Dokumente verfügbar sind
2. Best-Fit-Suche: Größtes Dokument finden, das in verbleibenden Platz passt
3. Fallback-Crop: Wenn nichts passt, kürzestes Dokument cropen

Quelle: nanochat/dataloader.py:130-151

DDP-Sharding und Resume-Logik

Der interne _document_batches-Generator handhabt verteiltes Training:

python
1# DDP-Sharding: Jeder Rank verarbeitet jeden ddp_world_size-ten Row Group
2rg_idx = ddp_rank
3while rg_idx < pf.num_row_groups:
4    rg = pf.read_row_group(rg_idx)
5    # ...
6    rg_idx += ddp_world_size

Resume-Mechanismus:

python
1# Bei Resume: Startposition berechnen
2if first_pass and resume_rg_idx is not None and pq_idx == resume_pq_idx:
3    base_idx = resume_rg_idx // ddp_world_size
4    base_idx += 1  # Nicht wiederholen!
5    rg_idx = base_idx * ddp_world_size + ddp_rank

Quelle: nanochat/dataloader.py:53-60

Speichereffiziente Buffer-Verwaltung

Der DataLoader nutzt pre-allokierte Buffer für minimale Speicherfragmentierung:

Single HtoD Transfer: Nach Batch-Konstruktion wird ein einziger asynchroner Transfer ausgeführt:

python
1gpu_buffer.copy_(cpu_buffer, non_blocking=use_cuda)
2yield inputs, targets, state_dict

Quelle: nanochat/dataloader.py:160-161

---

Tokenizer-Implementierung

Zwei-Wege-Strategie

NanoChat bietet zwei Tokenizer-Implementierungen:

HuggingFace Tokenizer

Der HuggingFaceTokenizer nutzt GPT-4-Style Pre-Tokenisierung:

python
1# Pre-Tokenizer Konfiguration
2tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
3    pre_tokenizers.Split(pattern=gpt4_split_regex, behavior="isolated"),
4    pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False)
5])

Regex-Pattern Anpassung: Das Pattern wurde von \p{N}{1,3} zu \p{N}{1,2} geändert, um Token-Verschwendung bei kleinen Vokabularen zu reduzieren.

Quelle: nanochat/tokenizer.py:71-78

RustBPE/Tiktoken für Effizienz

Der RustBPETokenizer kombiniert:

1. Training mit rustbpe: Schnelles BPE-Training
2. Inferenz mit tiktoken: Optimierte Produktionstokenisierung

python
1# Tiktoken Encoding Konstruktion
2enc = tiktoken.Encoding(
3    name="rustbpe",
4    pat_str=pattern,
5    mergeable_ranks=mergeable_ranks,
6    special_tokens=special_tokens,
7)

Multi-Threading: Die encode-Methode unterstützt num_threads für parallele Kodierung.

Quelle: nanochat/tokenizer.py:184-189

BOS-Token-Erkennung

Die get_bos_token_id-Methode implementiert eine Fallback-Kaskade:

1. Suche nach <|bos|>
2. Fallback auf `` (GPT-2 kompatibel)
3. Assertion-Fehler wenn keines gefunden

---

Infrastruktur-Komponenten

Automatische DType-Erkennung

Die _detect_compute_dtype-Funktion wählt automatisch den optimalen Datentyp:

python
1# Entscheidunglogik
2if capability >= (8, 0):  # Ampere+
3    return torch.bfloat16  # Native BF16 Support
4else:  # Pre-Ampere (V100, T4)
5    return torch.float32   # FP16 benötigt GradScaler

Umgebungsvariable-Override: NANOCHAT_DTYPE kann die automatische Erkennung überschreiben.

Quelle: nanochat/common.py:18-29

Checkpoint-Management

Der load_optimizer_state-Loader unterstützt verteiltes Training:

python
1# Rank-spezifischer Checkpoint-Pfad
2optimizer_path = os.path.join(
3    checkpoint_dir, 
4    f"optim_{step:06d}_rank{rank:d}.pt"
5)

Checkpoint-Typen:

Quelle: nanochat/checkpoint_manager.py:176-180

Kostenabschätzung

Die estimate_cost-Funktion berechnet Trainingskosten basierend auf GPU-Typ:

python
1gpu_hourly_rates = {
2    "H100": 3.00,
3    "A100": 1.79,
4    "V100": 0.55,
5}

Fallback: Unbekannte GPUs werden mit $2.00/GPU/Stunde geschätzt.

Quelle: nanochat/report.py:94-99

---

Datenfluss und Aufrufkette

End-to-End Training Flow

正在加载图表渲染器...

Erklärung der Sequenz:

1. Datenladen: Der DataLoader liest Row Groups aus Parquet-Dateien mit DDP-Sharding
2. Tokenisierung: Der Tokenizer kodiert Text-Batches mit BOS-Prepending
3. Best-Fit Cropping: Der Algorithmus füllt jede Zeile optimal ohne Padding
4. Forward Pass: Das GPT-Modell berechnet Loss und Gradients
5. Optimizer Step: MuonAdamW führt getrennte Updates für verschiedene Parametertypen aus
6. Checkpointing: Periodisches Speichern von Modell- und Optimizer-State

Kritische Aufrufkette

Die folgende Tabelle zeigt die wichtigsten Funktionsaufrufe im Training-Loop:

---

Modulabhängigkeiten

正在加载图表渲染器...

Abhängigkeitsrichtungen:

• Core → Common: Modell und Optimizer nutzen DType-Erkennung und Logging
• Data → Common: DataLoader nutzt verteilte Utilities
• Data → Tokenizer: DataLoader ist vom Tokenizer abhängig
• Infra → Common: Alle Infrastrukturmodule nutzen gemeinsame Utilities

---

Kerndesign-Entscheidungen

1. Hybrid-Optimizer (Muon + AdamW)

Entscheidung: Verwendung von Muon für 2D-Matrizen und AdamW für alle anderen Parameter.

Begründung: Muon orthogonalisiert Updates durch Newton-Schulz-Iteration, was für Gewichtsmatrizen effektiv ist. Für Embeddings und Skalare ist diese Orthogonalisierung nicht sinnvoll.

Trade-off: Erhöhte Komplexität durch zwei Optimizer-Pfade, aber bessere Konvergenz für Matrixparameter.

Quelle: nanochat/optim.py:152-177

2. Best-Fit Cropping statt Padding

Entscheidung: Kein Padding, sondern intelligentes Cropping für 100% Token-Utilisierung.

Begründung: Padding verschwendet Rechenkapazität. Best-Fit Cropping erreicht ~35% Token-Discard-Rate bei voller GPU-Auslastung.

Trade-off: Komplexere Batch-Erstellung, aber signifikant effizienteres Training.

Quelle: nanochat/dataloader.py:80-94

3. Fused Kernels mit 0-D CPU Tensoren

Entscheidung: Hyperparameter als 0-D CPU-Tensoren speichern.

Begründung: Vermeidet torch.compile Rekompilierung bei Hyperparameter-Änderungen während des Trainings.

Trade-off: Leicht erhöhter Speicherverbrauch, aber eliminiert Rekompilierungs-Overhead.

Quelle: nanochat/optim.py:180-192

4. Automatische DType-Erkennung

Entscheidung: BF16 für Ampere+ GPUs, FP32 für ältere Hardware.

Begründung: BF16 erfordert keine GradScaler und ist auf SM 80+ nativ unterstützt. Ältere GPUs (V100, T4) haben nur FP16 Tensor Cores.

Trade-off: Kein FP16-Support ohne manuellen Override, aber vereinfachtes Training auf moderner Hardware.

Quelle: nanochat/common.py:21-29

5. Parquet-basierte Datenspeicherung

Entscheidung: Datensätze als Parquet-Dateien mit Row Groups.

Begründung: Ermöglicht effizientes DDP-Sharding auf Row-Group-Ebene und Resume ohne Datenwiederholung.

Trade-off: Erfordert Vorverarbeitung der Daten, aber effiziente Speicherung und schnelle Random Access.

Quelle: nanochat/dataloader.py:36-38

---

Technologie-Stack

---

Konfiguration und Startprozess

Umgebungsvariablen

Trainingskonfiguration

Die wichtigsten Konfigurationsparameter im Optimizer-Setup:

python
1# Beispielkonfiguration
2setup_optimizer(
3    unembedding_lr=0.004,  # LM-Head Lernrate
4    embedding_lr=0.2,      # Embedding Lernrate
5    matrix_lr=0.02,        # Matrix-Lernrate (Muon)
6    weight_decay=0.0,      # Weight Decay für Muon
7    scalar_lr=0.5          # Skalar-Lernrate
8)

Checkpoint-Struktur

{base_dir}/
├── base_checkpoints/
│   └── {model_tag}/
│       ├── model_{step:06d}.pt
│       └── optim_{step:06d}_rank{rank:d}.pt
├── chatsft_checkpoints/
└── chatrl_checkpoints/

---

Bekannte Einschränkungen

1. Kein FP16-Support ohne GradScaler: Pre-Ampere GPUs fallen auf FP32 zurück
2. Muon nicht für Embeddings geeignet: Muon ist nur für 2D-Matrizen konzipiert
3. ~35% Token-Discard-Rate: Best-Fit Cropping verwirft einen Teil der Tokens
4. Keine FSDP-Unterstützung: Nur DDP für verteiltes Training
5. Parquet-Voraussetzung: Datensätze müssen vorab in Parquet konvertiert werden