Merge 65865df300 into 72b9064f9d

remove leftover mid references (#491 )
fix bug in speedrun script, batch size that doesn't OOM on 8XH100 for d24 is 16
2026-07-08 22:09:16 +00:00 · 2026-02-03 13:57:19 +01:00 · 2026-02-02 08:33:46 -08:00 · 2026-02-02 15:50:14 +00:00 · 2026-01-15 20:47:26 +01:00 · 2025-12-30 15:13:42 +01:00
9 changed files with 202 additions and 39 deletions
--- a/nanochat/checkpoint_manager.py
+++ b/nanochat/checkpoint_manager.py
@ -164,7 +164,6 @@ def load_model_from_dir(checkpoints_dir, device, phase, model_tag=None, step=Non
 def load_model(source, *args, **kwargs):
    model_dir = {
        "base": "base_checkpoints",
-        "mid": "mid_checkpoints",
        "sft": "chatsft_checkpoints",
        "rl": "chatrl_checkpoints",
    }[source]
--- a/nanochat/report.py
+++ b/nanochat/report.py
@ -211,8 +211,6 @@ EXPECTED_FILES = [
    "base-model-training.md",
    "base-model-loss.md",
    "base-model-evaluation.md",
-    "midtraining.md",
-    "chat-evaluation-mid.md",
    "chat-sft.md",
    "chat-evaluation-sft.md",
    "chat-rl.md",
@ -316,8 +314,6 @@ class Report:
                # extract the most important metrics from the sections
                if file_name == "base-model-evaluation.md":
                    final_metrics["base"] = extract(section, "CORE")
-                if file_name == "chat-evaluation-mid.md":
-                    final_metrics["mid"] = extract(section, chat_metrics)
                if file_name == "chat-evaluation-sft.md":
                    final_metrics["sft"] = extract(section, chat_metrics)
                if file_name == "chat-evaluation-rl.md":
@ -337,7 +333,7 @@ class Report:
            # Custom ordering: CORE first, ChatCORE last, rest in middle
            all_metrics = sorted(all_metrics, key=lambda x: (x != "CORE", x == "ChatCORE", x))
            # Fixed column widths
-            stages = ["base", "mid", "sft", "rl"]
+            stages = ["base", "sft", "rl"]
            metric_width = 15
            value_width = 8
            # Write table header
--- a/nanochat/ui.html
+++ b/nanochat/ui.html
@ -1,10 +1,13 @@
 <!DOCTYPE html>
 <html lang="en">
+
 <head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0, viewport-fit=cover">
    <title>NanoChat</title>
    <link rel="icon" type="image/svg+xml" href="/logo.svg">
+    <script src="https://cdn.jsdelivr.net/npm/marked@9.1.6/marked.min.js"></script>
+    <script src="https://cdn.jsdelivr.net/npm/dompurify@3.0.5/dist/purify.min.js"></script>
    <style>
        :root {
            color-scheme: light;
@ -104,11 +107,14 @@
        }

        .message-content {
-            white-space: pre-wrap;
            line-height: 1.6;
            max-width: 100%;
        }

+        .message.user .message-content {
+            white-space: pre-wrap;
+        }
+
        .message.assistant .message-content {
            background: transparent;
            border: none;
@ -224,8 +230,16 @@
        }

        @keyframes typing {
-            0%, 60%, 100% { opacity: 0.2; }
-            30% { opacity: 1; }
+
+            0%,
+            60%,
+            100% {
+                opacity: 0.2;
+            }
+
+            30% {
+                opacity: 1;
+            }
        }

        .error-message {
@ -236,13 +250,109 @@
            border-radius: 0.75rem;
            margin-top: 0.5rem;
        }
+
+        /* Markdown styling */
+        .message-content table {
+            border-collapse: collapse;
+            margin: 0.25rem 0;
+            width: 100%;
+        }
+
+        .message-content th,
+        .message-content td {
+            border: 1px solid #e5e7eb;
+            padding: 0.5rem;
+            text-align: left;
+        }
+
+        .message-content th {
+            background-color: #f9fafb;
+            font-weight: 600;
+        }
+
+        .message-content code {
+            background-color: #f3f4f6;
+            padding: 0.125rem 0.25rem;
+            border-radius: 0.25rem;
+            font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', 'Consolas', 'Courier New', monospace;
+            font-size: 0.875em;
+        }
+
+        .message-content pre {
+            background-color: #f3f4f6;
+            padding: 0.75rem;
+            border-radius: 0.5rem;
+            overflow-x: auto;
+            margin: 0.25rem 0;
+            position: relative;
+        }
+
+        .message-content pre code {
+            background: none;
+            padding: 0;
+        }
+
+        .code-language-label {
+            position: absolute;
+            top: 0.5rem;
+            right: 0.5rem;
+            background-color: #e5e7eb;
+            color: #6b7280;
+            padding: 0.125rem 0.375rem;
+            border-radius: 0.25rem;
+            font-size: 0.75rem;
+            font-family: ui-sans-serif, -apple-system, system-ui, "Segoe UI", Helvetica, Arial, sans-serif;
+            text-transform: lowercase;
+            font-weight: 500;
+        }
+
+        .message-content h1,
+        .message-content h2,
+        .message-content h3,
+        .message-content h4,
+        .message-content h5,
+        .message-content h6 {
+            margin: 0.75rem 0 0.25rem 0;
+            font-weight: 600;
+        }
+
+        .message-content h1 {
+            font-size: 1.5em;
+        }
+
+        .message-content h2 {
+            font-size: 1.3em;
+        }
+
+        .message-content h3 {
+            font-size: 1.1em;
+        }
+
+        .message-content ul,
+        .message-content ol {
+            margin: 0.25rem 0;
+            padding-left: 1.5rem;
+        }
+
+        .message-content li {
+            margin: 0.1rem 0;
+        }
+
+        .message-content blockquote {
+            border-left: 4px solid #e5e7eb;
+            padding-left: 1rem;
+            margin: 0.5rem 0;
+            color: #6b7280;
+        }
    </style>
 </head>
+
 <body>
    <div class="header">
        <div class="header-left">
            <button class="new-conversation-btn" onclick="newConversation()" title="New Conversation (Ctrl+Shift+N)">
-                <svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
+                <svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"
+                    stroke-linecap="round" stroke-linejoin="round">
                    <path d="M12 5v14"></path>
                    <path d="M5 12h14"></path>
                </svg>
@ -259,15 +369,11 @@

    <div class="input-container">
        <div class="input-wrapper">
-            <textarea
-                id="chatInput"
-                class="chat-input"
-                placeholder="Ask anything"
-                rows="1"
-                onkeydown="handleKeyDown(event)"
-            ></textarea>
+            <textarea id="chatInput" class="chat-input" placeholder="Ask anything" rows="1"
+                onkeydown="handleKeyDown(event)"></textarea>
            <button id="sendButton" class="send-button" onclick="sendMessage()" disabled>
-                <svg width="22" height="22" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
+                <svg width="22" height="22" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"
+                    stroke-linecap="round" stroke-linejoin="round">
                    <path d="M22 2L11 13"></path>
                    <path d="M22 2l-7 20-4-9-9-4 20-7z"></path>
                </svg>
@ -287,7 +393,65 @@
        let currentTemperature = 0.8;
        let currentTopK = 50;

-        chatInput.addEventListener('input', function() {
+        // Configure marked.js for better rendering
+        if (typeof marked !== 'undefined') {
+            marked.setOptions({
+                breaks: true,        // Convert line breaks to <br>
+                gfm: true,          // GitHub Flavored Markdown
+                sanitize: false,    // We'll use DOMPurify instead
+                smartypants: false  // Disable smart quotes for simplicity
+            });
+        }
+
+        // Markdown rendering function
+        function renderMarkdown(content) {
+            try {
+                if (typeof marked !== 'undefined' && typeof DOMPurify !== 'undefined') {
+                    const rawHtml = marked.parse(content);
+                    const sanitizedHtml = DOMPurify.sanitize(rawHtml);
+                    return addLanguageLabels(sanitizedHtml);
+                } else {
+                    console.warn('Markdown libraries not loaded, falling back to plain text');
+                    return escapeHtml(content);
+                }
+            } catch (error) {
+                console.error('Markdown rendering failed:', error);
+                return escapeHtml(content);
+            }
+        }
+
+        // Add language labels to code blocks
+        function addLanguageLabels(html) {
+            const tempDiv = document.createElement('div');
+            tempDiv.innerHTML = html;
+            
+            const codeBlocks = tempDiv.querySelectorAll('pre code');
+            codeBlocks.forEach(codeBlock => {
+                const pre = codeBlock.parentElement;
+                const className = codeBlock.className;
+                
+                // Extract language from class (e.g., "language-javascript" -> "javascript")
+                const languageMatch = className.match(/language-(\w+)/);
+                if (languageMatch) {
+                    const language = languageMatch[1];
+                    const label = document.createElement('span');
+                    label.className = 'code-language-label';
+                    label.textContent = language;
+                    pre.appendChild(label);
+                }
+            });
+            
+            return tempDiv.innerHTML;
+        }
+
+        // HTML escape fallback function
+        function escapeHtml(text) {
+            const div = document.createElement('div');
+            div.textContent = text;
+            return div.innerHTML;
+        }
+
+        chatInput.addEventListener('input', function () {
            this.style.height = 'auto';
            this.style.height = Math.min(this.scrollHeight, 200) + 'px';
            sendButton.disabled = !this.value.trim() || isGenerating;
@ -300,7 +464,7 @@
            }
        }

-        document.addEventListener('keydown', function(event) {
+        document.addEventListener('keydown', function (event) {
            // Ctrl+Shift+N for new conversation
            if (event.ctrlKey && event.shiftKey && event.key === 'N') {
                event.preventDefault();
@ -326,13 +490,19 @@

            const contentDiv = document.createElement('div');
            contentDiv.className = 'message-content';
-            contentDiv.textContent = content;
+
+            // Render markdown for assistant messages, plain text for others
+            if (role === 'assistant' && content) {
+                contentDiv.innerHTML = renderMarkdown(content);
+            } else {
+                contentDiv.textContent = content;
+            }

            // Add click handler for user messages to enable editing
            if (role === 'user' && messageIndex !== null) {
                contentDiv.setAttribute('data-message-index', messageIndex);
                contentDiv.setAttribute('title', 'Click to edit and restart from here');
-                contentDiv.addEventListener('click', function() {
+                contentDiv.addEventListener('click', function () {
                    if (!isGenerating) {
                        editMessage(messageIndex);
                    }
@ -343,7 +513,7 @@
            if (role === 'assistant' && messageIndex !== null) {
                contentDiv.setAttribute('data-message-index', messageIndex);
                contentDiv.setAttribute('title', 'Click to regenerate this response');
-                contentDiv.addEventListener('click', function() {
+                contentDiv.addEventListener('click', function () {
                    if (!isGenerating) {
                        regenerateMessage(messageIndex);
                    }
@ -426,7 +596,8 @@
                                const data = JSON.parse(line.slice(6));
                                if (data.token) {
                                    fullResponse += data.token;
-                                    assistantContent.textContent = fullResponse;
+                                    // Render markdown in real-time as tokens arrive
+                                    assistantContent.innerHTML = renderMarkdown(fullResponse);
                                    chatContainer.scrollTop = chatContainer.scrollHeight;
                                }
                            } catch (e) {
@ -441,7 +612,7 @@
                // Add click handler to regenerate this assistant message
                assistantContent.setAttribute('data-message-index', assistantMessageIndex);
                assistantContent.setAttribute('title', 'Click to regenerate this response');
-                assistantContent.addEventListener('click', function() {
+                assistantContent.addEventListener('click', function () {
                    if (!isGenerating) {
                        regenerateMessage(assistantMessageIndex);
                    }
@ -563,4 +734,5 @@
            });
    </script>
 </body>
-</html>
+
+</html>
--- a/runs/speedrun.sh
+++ b/runs/speedrun.sh
@ -69,13 +69,10 @@ python -m scripts.tok_eval
 echo "Waiting for dataset download to complete..."
 wait $DATASET_DOWNLOAD_PID

-# Number of processes/GPUs to use
-NPROC_PER_NODE=8
-
 # d24 model (slightly overtrained is enough to beat GPT-2 => increase data:params ratio from compute optimal 10.5 (default) to 12)
-torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_train -- --depth=24 --target-param-data-ratio=12 --run=$WANDB_RUN
+torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- --depth=24 --target-param-data-ratio=12 --device-batch-size=16 --run=$WANDB_RUN
 # evaluate the model: CORE metric, BPB on train/val, and draw samples
-torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_eval
+torchrun --standalone --nproc_per_node=8 -m scripts.base_eval -- --device-batch-size=16

 # -----------------------------------------------------------------------------
 # SFT (teach the model conversation special tokens, tool use, multiple choice)
@ -85,8 +82,8 @@ torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_eval
 curl -L -o $NANOCHAT_BASE_DIR/identity_conversations.jsonl https://karpathy-public.s3.us-west-2.amazonaws.com/identity_conversations.jsonl

 # run SFT and eval the model
-torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.chat_sft -- --run=$WANDB_RUN
-torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.chat_eval -- -i sft
+torchrun --standalone --nproc_per_node=8 -m scripts.chat_sft -- --device-batch-size=16 --run=$WANDB_RUN
+torchrun --standalone --nproc_per_node=8 -m scripts.chat_eval -- -i sft

 # chat with the model over CLI! Leave out the -p to chat interactively
 # python -m scripts.chat_cli -p "Why is the sky blue?"
--- a/scripts/chat_cli.py
+++ b/scripts/chat_cli.py
@ -12,7 +12,7 @@ from nanochat.engine import Engine
 from nanochat.checkpoint_manager import load_model

 parser = argparse.ArgumentParser(description='Chat with the model')
-parser.add_argument('-i', '--source', type=str, default="sft", help="Source of the model: sft|mid|rl")
+parser.add_argument('-i', '--source', type=str, default="sft", help="Source of the model: sft|rl")
 parser.add_argument('-g', '--model-tag', type=str, default=None, help='Model tag to load')
 parser.add_argument('-s', '--step', type=int, default=None, help='Step to load')
 parser.add_argument('-p', '--prompt', type=str, default='', help='Prompt the model, get a single response back')
--- a/scripts/chat_eval.py
+++ b/scripts/chat_eval.py
@ -183,7 +183,7 @@ if __name__ == "__main__":

    # Parse command-line arguments
    parser = argparse.ArgumentParser()
-    parser.add_argument('-i', '--source', type=str, required=True, help="Source of the model: sft|mid|rl")
+    parser.add_argument('-i', '--source', type=str, required=True, help="Source of the model: sft|rl")
    parser.add_argument('-a', '--task-name', type=str, default=None, help="Task name. Default = all tasks. Use | to split multiple tasks.")
    parser.add_argument('-d', '--dtype', type=str, default='bfloat16', choices=['float32', 'bfloat16'])
    parser.add_argument('-t', '--temperature', type=float, default=0.0)
--- a/scripts/chat_rl.py
+++ b/scripts/chat_rl.py
@ -38,7 +38,6 @@ parser.add_argument("--run", type=str, default="dummy", help="wandb run name ('d
 parser.add_argument("--device-type", type=str, default="", help="cuda|cpu|mps (empty = autodetect)")
 parser.add_argument("--dtype", type=str, default="bfloat16", help="float32|bfloat16")
 # Model loading
-parser.add_argument("--source", type=str, default="sft", help="mid|sft - which checkpoint to load from")
 parser.add_argument("--model-tag", type=str, default=None, help="model tag to load from")
 parser.add_argument("--model-step", type=int, default=None, help="model step to load from")
 # Training horizon
@ -77,7 +76,7 @@ use_dummy_wandb = args.run == "dummy" or not master_process
 wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat-rl", name=args.run, config=user_config)

 # Init model and tokenizer
-model, tokenizer, meta = load_model(args.source, device, phase="eval", model_tag=args.model_tag, step=args.model_step)
+model, tokenizer, meta = load_model("sft", device, phase="eval", model_tag=args.model_tag, step=args.model_step)
 engine = Engine(model, tokenizer) # for sampling rollouts

 # -----------------------------------------------------------------------------
--- a/scripts/chat_web.py
+++ b/scripts/chat_web.py
@ -62,7 +62,7 @@ MAX_MAX_TOKENS = 4096

 parser = argparse.ArgumentParser(description='NanoChat Web Server')
 parser.add_argument('-n', '--num-gpus', type=int, default=1, help='Number of GPUs to use (default: 1)')
-parser.add_argument('-i', '--source', type=str, default="sft", help="Source of the model: sft|mid|rl")
+parser.add_argument('-i', '--source', type=str, default="sft", help="Source of the model: sft|rl")
 parser.add_argument('-t', '--temperature', type=float, default=0.8, help='Default temperature for generation')
 parser.add_argument('-k', '--top-k', type=int, default=50, help='Default top-k sampling parameter')
 parser.add_argument('-m', '--max-tokens', type=int, default=512, help='Default max tokens for generation')
--- a/tasks/spellingbee.py
+++ b/tasks/spellingbee.py
@ -20,7 +20,7 @@ LLM because it has to learn how every token (a little semantic chunk/atom)
 maps to the sequence of individual characters that make it up. Larger models
 learn this eventually on their own, but if we want this capability to exist
 in smaller models, we have to actively encourage it by over-representing it
-in the training data. Midtraining is a good place to do this.
+in the training data. SFT is a good place to do this.

 To preview a few example conversations, run:
 python -m tasks.spellingbee
Author	SHA1	Message	Date
Gaurav	30da19b19e	Merge `65865df300` into `72b9064f9d`	2026-02-03 13:57:19 +01:00
Sofie Van Landeghem	72b9064f9d	remove leftover mid references (#491 )	2026-02-02 08:33:46 -08:00
Andrej Karpathy	b19b4f3e49	fix bug in speedrun script, batch size that doesn't OOM on 8XH100 for d24 is 16	2026-02-02 15:50:14 +00:00
svlandeg	65865df300	Merge branch 'master' into master_goderr	2026-01-15 20:47:26 +01:00
svlandeg	c79559674b	Merge branch 'master' into master_goderr	2025-12-30 15:13:42 +01:00
Goderr	a3b701af8d	change in code block display	2025-10-21 20:17:03 +05:30
Goderr	e9c5e911f6	change in linespacing of response	2025-10-21 20:14:15 +05:30
Goderr	07348bb1d3	forgot to change the max_tokens	2025-10-21 18:01:41 +05:30
Goderr	5ec267fabc	Change in UI	2025-10-21 17:49:08 +05:30