{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "481c7ee8",
   "metadata": {},
   "source": [
    "### Start by making sure you have your API key setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a0d70ed1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "if os.getenv(\"A3_API_KEY\", \"\") == \"\":\n",
    "    raise ValueError(\n",
    "        \"A3_API_KEY is not set. Please set it in your environment variables or use `os.environ['A3_API_KEY'] = 'your-8-digit-student-id'`\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97b64b44",
   "metadata": {},
   "source": [
    "### Do imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "10118323",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/herman/dev/cs336-teach/cs336-monorepo/assignments/scaling/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import rich\n",
    "\n",
    "from cs336_scaling.client import (\n",
    "    get_budget,\n",
    "    get_experiment,\n",
    "    list_experiments,\n",
    "    save_final_submission,\n",
    "    submit_experiment,\n",
    ")\n",
    "from cs336_scaling.training.model.config import BasicTransformerConfig\n",
    "from cs336_scaling.training.optimizer import AdamWConfig, WarmupCosineDecay\n",
    "from cs336_scaling.training.training_config import TrainingConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d98fe94",
   "metadata": {},
   "source": [
    "### 1. Start by checking your budget. It should be empty for now"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4b7aa14f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BudgetSummary(used_seconds=0.0, remaining_seconds=43200.0, total_budget_seconds=43200.0)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_budget()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b53793a",
   "metadata": {},
   "source": [
    "### 2. Submit an experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "76be1f80",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TrainingConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">architecture_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">BasicTransformerConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">attention_bias</span>=<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">head_dim</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">64</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">hidden_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">448</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">intermediate_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1280</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">num_attention_heads</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">num_hidden_layers</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">num_key_value_heads</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">rms_norm_eps</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-06</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">rope_theta</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1000000</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">tie_word_embeddings</span>=<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">dtype</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'bfloat16'</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">vocab_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32000</span>\n",
       "    <span style=\"font-weight: bold\">)</span>,\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">optimizer_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">AdamWConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">lr_scheduler</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">WarmupCosineDecay</span><span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">peak_value</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0003</span>, <span style=\"color: #808000; text-decoration-color: #808000\">final_lr_frac</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1</span>, <span style=\"color: #808000; text-decoration-color: #808000\">warmup_frac</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.05</span>, <span style=\"color: #808000; text-decoration-color: #808000\">init_value</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0</span><span style=\"font-weight: bold\">)</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">weight_decay</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.01</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">beta1</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">beta2</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.95</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">eps</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-08</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">eps_root</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-08</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">grad_clip_norm</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>\n",
       "    <span style=\"font-weight: bold\">)</span>,\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">train_batch_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span>,\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">val_batch_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32</span>,\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">n_evals</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">16</span>,\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">total_train_tokens</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8388608</span>,\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">max_runtime_seconds</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">60.0</span>,\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">model_seed</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>\n",
       "<span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mTrainingConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "    \u001b[33marchitecture_config\u001b[0m=\u001b[1;35mBasicTransformerConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "        \u001b[33mattention_bias\u001b[0m=\u001b[3;91mFalse\u001b[0m,\n",
       "        \u001b[33mhead_dim\u001b[0m=\u001b[1;36m64\u001b[0m,\n",
       "        \u001b[33mhidden_size\u001b[0m=\u001b[1;36m448\u001b[0m,\n",
       "        \u001b[33mintermediate_size\u001b[0m=\u001b[1;36m1280\u001b[0m,\n",
       "        \u001b[33mnum_attention_heads\u001b[0m=\u001b[1;36m7\u001b[0m,\n",
       "        \u001b[33mnum_hidden_layers\u001b[0m=\u001b[1;36m9\u001b[0m,\n",
       "        \u001b[33mnum_key_value_heads\u001b[0m=\u001b[1;36m7\u001b[0m,\n",
       "        \u001b[33mrms_norm_eps\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-06\u001b[0m,\n",
       "        \u001b[33mrope_theta\u001b[0m=\u001b[1;36m1000000\u001b[0m,\n",
       "        \u001b[33mtie_word_embeddings\u001b[0m=\u001b[3;91mFalse\u001b[0m,\n",
       "        \u001b[33mdtype\u001b[0m=\u001b[32m'bfloat16'\u001b[0m,\n",
       "        \u001b[33mvocab_size\u001b[0m=\u001b[1;36m32000\u001b[0m\n",
       "    \u001b[1m)\u001b[0m,\n",
       "    \u001b[33moptimizer_config\u001b[0m=\u001b[1;35mAdamWConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "        \u001b[33mlr_scheduler\u001b[0m=\u001b[1;35mWarmupCosineDecay\u001b[0m\u001b[1m(\u001b[0m\u001b[33mpeak_value\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.0003\u001b[0m, \u001b[33mfinal_lr_frac\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.1\u001b[0m, \u001b[33mwarmup_frac\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.05\u001b[0m, \u001b[33minit_value\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.0\u001b[0m\u001b[1m)\u001b[0m,\n",
       "        \u001b[33mweight_decay\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.01\u001b[0m,\n",
       "        \u001b[33mbeta1\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.9\u001b[0m,\n",
       "        \u001b[33mbeta2\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.95\u001b[0m,\n",
       "        \u001b[33meps\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-08\u001b[0m,\n",
       "        \u001b[33meps_root\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-08\u001b[0m,\n",
       "        \u001b[33mgrad_clip_norm\u001b[0m=\u001b[1;36m1\u001b[0m\u001b[1;36m.0\u001b[0m\n",
       "    \u001b[1m)\u001b[0m,\n",
       "    \u001b[33mtrain_batch_size\u001b[0m=\u001b[1;36m128\u001b[0m,\n",
       "    \u001b[33mval_batch_size\u001b[0m=\u001b[1;36m32\u001b[0m,\n",
       "    \u001b[33mn_evals\u001b[0m=\u001b[1;36m16\u001b[0m,\n",
       "    \u001b[33mtotal_train_tokens\u001b[0m=\u001b[1;36m8388608\u001b[0m,\n",
       "    \u001b[33mmax_runtime_seconds\u001b[0m=\u001b[1;36m60\u001b[0m\u001b[1;36m.0\u001b[0m,\n",
       "    \u001b[33mmodel_seed\u001b[0m=\u001b[1;36m0\u001b[0m\n",
       "\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# define our experiment config\n",
    "training_config = TrainingConfig(\n",
    "    architecture_config=BasicTransformerConfig(\n",
    "        attention_bias=False,\n",
    "        head_dim=64,\n",
    "        hidden_size=448,\n",
    "        intermediate_size=1280,\n",
    "        num_attention_heads=7,\n",
    "        num_hidden_layers=9,\n",
    "        num_key_value_heads=7,\n",
    "        rms_norm_eps=1e-6,\n",
    "        rope_theta=1_000_000,\n",
    "        tie_word_embeddings=False,\n",
    "        dtype=\"bfloat16\",\n",
    "        vocab_size=32_000,\n",
    "    ),\n",
    "    optimizer_config=AdamWConfig(\n",
    "        lr_scheduler=WarmupCosineDecay(\n",
    "            peak_value=3e-4,\n",
    "            final_lr_frac=0.1,\n",
    "            warmup_frac=0.05,\n",
    "            init_value=0.0,\n",
    "        ),\n",
    "        weight_decay=1e-2,\n",
    "        beta1=0.9,\n",
    "        beta2=0.95,\n",
    "        eps=1e-8,\n",
    "        eps_root=1e-8,\n",
    "        grad_clip_norm=1.0,\n",
    "    ),\n",
    "    train_batch_size=128,\n",
    "    val_batch_size=32,\n",
    "    n_evals=16,\n",
    "    total_train_tokens=8_388_608,\n",
    "    max_runtime_seconds=60.0,\n",
    "    model_seed=0,\n",
    ")\n",
    "rich.print(training_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "18cd1d09",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SubmitResponse(experiment_id=14, budget_summary=BudgetSummary(used_seconds=60.0, remaining_seconds=43140.0, total_budget_seconds=43200.0))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# submit the experiment\n",
    "submit_result = submit_experiment(training_config)\n",
    "submit_result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee3eec4e",
   "metadata": {},
   "source": [
    "#### You can also submit your config as a JSON-string\n",
    "\n",
    "The code below will crash, since you've already submitted the same experiment above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "64f224cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Experiment already exists for this training config (expected, continuing).\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    submit_result_json = submit_experiment(\n",
    "        {\n",
    "            \"architecture_config\": {\n",
    "                \"attention_bias\": False,\n",
    "                \"head_dim\": 64,\n",
    "                \"hidden_size\": 448,\n",
    "                \"intermediate_size\": 1280,\n",
    "                \"num_attention_heads\": 7,\n",
    "                \"num_hidden_layers\": 9,\n",
    "                \"num_key_value_heads\": 7,\n",
    "                \"rms_norm_eps\": 1e-06,\n",
    "                \"rope_theta\": 1000000,\n",
    "                \"tie_word_embeddings\": False,\n",
    "                \"dtype\": \"bfloat16\",\n",
    "                \"vocab_size\": 32000,\n",
    "            },\n",
    "            \"optimizer_config\": {\n",
    "                \"lr_scheduler\": {\n",
    "                    \"peak_value\": 0.0003,\n",
    "                    \"final_lr_frac\": 0.1,\n",
    "                    \"warmup_frac\": 0.05,\n",
    "                    \"init_value\": 0.0,\n",
    "                },\n",
    "                \"weight_decay\": 0.01,\n",
    "                \"beta1\": 0.9,\n",
    "                \"beta2\": 0.95,\n",
    "                \"eps\": 1e-08,\n",
    "                \"eps_root\": 1e-08,\n",
    "                \"grad_clip_norm\": 1.0,\n",
    "            },\n",
    "            \"train_batch_size\": 128,\n",
    "            \"val_batch_size\": 32,\n",
    "            \"n_evals\": 16,\n",
    "            \"total_train_tokens\": 8388608,\n",
    "            \"max_runtime_seconds\": 60.0,\n",
    "            \"model_seed\": 0,\n",
    "        }\n",
    "    )\n",
    "except RuntimeError as e:\n",
    "    if \"409 Client Error: Conflict\" in str(e) and \"experiment already exists\" in str(e):\n",
    "        print(\n",
    "            \"Experiment already exists for this training config (expected, continuing).\"\n",
    "        )\n",
    "        submit_result_json = None\n",
    "    else:\n",
    "        raise"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb3babe0",
   "metadata": {},
   "source": [
    "### Query experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "702d8273",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">[</span>\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">ExperimentResponse</span><span style=\"font-weight: bold\">(</span>\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">experiment_id</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">14</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">training_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TrainingConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">architecture_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">BasicTransformerConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">attention_bias</span>=<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">head_dim</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">64</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">hidden_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">448</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">intermediate_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1280</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">num_attention_heads</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">num_hidden_layers</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">num_key_value_heads</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">rms_norm_eps</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-06</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">rope_theta</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1000000</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">tie_word_embeddings</span>=<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">dtype</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'bfloat16'</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">vocab_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32000</span>\n",
       "            <span style=\"font-weight: bold\">)</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">optimizer_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">AdamWConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">lr_scheduler</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">WarmupCosineDecay</span><span style=\"font-weight: bold\">(</span>\n",
       "                    <span style=\"color: #808000; text-decoration-color: #808000\">peak_value</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0003</span>,\n",
       "                    <span style=\"color: #808000; text-decoration-color: #808000\">final_lr_frac</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1</span>,\n",
       "                    <span style=\"color: #808000; text-decoration-color: #808000\">warmup_frac</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.05</span>,\n",
       "                    <span style=\"color: #808000; text-decoration-color: #808000\">init_value</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0</span>\n",
       "                <span style=\"font-weight: bold\">)</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">weight_decay</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.01</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">beta1</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">beta2</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.95</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">eps</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-08</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">eps_root</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-08</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">grad_clip_norm</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>\n",
       "            <span style=\"font-weight: bold\">)</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">train_batch_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">val_batch_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">n_evals</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">16</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">total_train_tokens</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8388608</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">max_runtime_seconds</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">60.0</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">model_seed</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>\n",
       "        <span style=\"font-weight: bold\">)</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">status</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">QueuedExperimentStatus</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">queued_at</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">datetime</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">.datetime</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2026</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">29</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">23</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">46</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">36</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">972380</span>, <span style=\"color: #808000; text-decoration-color: #808000\">tzinfo</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TzInfo</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span><span style=\"font-weight: bold\">))</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">status_type</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'queued'</span>\n",
       "        <span style=\"font-weight: bold\">)</span>\n",
       "    <span style=\"font-weight: bold\">)</span>\n",
       "<span style=\"font-weight: bold\">]</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m[\u001b[0m\n",
       "    \u001b[1;35mExperimentResponse\u001b[0m\u001b[1m(\u001b[0m\n",
       "        \u001b[33mexperiment_id\u001b[0m=\u001b[1;36m14\u001b[0m,\n",
       "        \u001b[33mtraining_config\u001b[0m=\u001b[1;35mTrainingConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33marchitecture_config\u001b[0m=\u001b[1;35mBasicTransformerConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "                \u001b[33mattention_bias\u001b[0m=\u001b[3;91mFalse\u001b[0m,\n",
       "                \u001b[33mhead_dim\u001b[0m=\u001b[1;36m64\u001b[0m,\n",
       "                \u001b[33mhidden_size\u001b[0m=\u001b[1;36m448\u001b[0m,\n",
       "                \u001b[33mintermediate_size\u001b[0m=\u001b[1;36m1280\u001b[0m,\n",
       "                \u001b[33mnum_attention_heads\u001b[0m=\u001b[1;36m7\u001b[0m,\n",
       "                \u001b[33mnum_hidden_layers\u001b[0m=\u001b[1;36m9\u001b[0m,\n",
       "                \u001b[33mnum_key_value_heads\u001b[0m=\u001b[1;36m7\u001b[0m,\n",
       "                \u001b[33mrms_norm_eps\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-06\u001b[0m,\n",
       "                \u001b[33mrope_theta\u001b[0m=\u001b[1;36m1000000\u001b[0m,\n",
       "                \u001b[33mtie_word_embeddings\u001b[0m=\u001b[3;91mFalse\u001b[0m,\n",
       "                \u001b[33mdtype\u001b[0m=\u001b[32m'bfloat16'\u001b[0m,\n",
       "                \u001b[33mvocab_size\u001b[0m=\u001b[1;36m32000\u001b[0m\n",
       "            \u001b[1m)\u001b[0m,\n",
       "            \u001b[33moptimizer_config\u001b[0m=\u001b[1;35mAdamWConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "                \u001b[33mlr_scheduler\u001b[0m=\u001b[1;35mWarmupCosineDecay\u001b[0m\u001b[1m(\u001b[0m\n",
       "                    \u001b[33mpeak_value\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.0003\u001b[0m,\n",
       "                    \u001b[33mfinal_lr_frac\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.1\u001b[0m,\n",
       "                    \u001b[33mwarmup_frac\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.05\u001b[0m,\n",
       "                    \u001b[33minit_value\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.0\u001b[0m\n",
       "                \u001b[1m)\u001b[0m,\n",
       "                \u001b[33mweight_decay\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.01\u001b[0m,\n",
       "                \u001b[33mbeta1\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.9\u001b[0m,\n",
       "                \u001b[33mbeta2\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.95\u001b[0m,\n",
       "                \u001b[33meps\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-08\u001b[0m,\n",
       "                \u001b[33meps_root\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-08\u001b[0m,\n",
       "                \u001b[33mgrad_clip_norm\u001b[0m=\u001b[1;36m1\u001b[0m\u001b[1;36m.0\u001b[0m\n",
       "            \u001b[1m)\u001b[0m,\n",
       "            \u001b[33mtrain_batch_size\u001b[0m=\u001b[1;36m128\u001b[0m,\n",
       "            \u001b[33mval_batch_size\u001b[0m=\u001b[1;36m32\u001b[0m,\n",
       "            \u001b[33mn_evals\u001b[0m=\u001b[1;36m16\u001b[0m,\n",
       "            \u001b[33mtotal_train_tokens\u001b[0m=\u001b[1;36m8388608\u001b[0m,\n",
       "            \u001b[33mmax_runtime_seconds\u001b[0m=\u001b[1;36m60\u001b[0m\u001b[1;36m.0\u001b[0m,\n",
       "            \u001b[33mmodel_seed\u001b[0m=\u001b[1;36m0\u001b[0m\n",
       "        \u001b[1m)\u001b[0m,\n",
       "        \u001b[33mstatus\u001b[0m=\u001b[1;35mQueuedExperimentStatus\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33mqueued_at\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2026\u001b[0m, \u001b[1;36m4\u001b[0m, \u001b[1;36m29\u001b[0m, \u001b[1;36m23\u001b[0m, \u001b[1;36m46\u001b[0m, \u001b[1;36m36\u001b[0m, \u001b[1;36m972380\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
       "            \u001b[33mstatus_type\u001b[0m=\u001b[32m'queued'\u001b[0m\n",
       "        \u001b[1m)\u001b[0m\n",
       "    \u001b[1m)\u001b[0m\n",
       "\u001b[1m]\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# the experiment is in queued state for now\n",
    "rich.print(list_experiments())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5392619f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">[</span>\n",
       "    <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">ExperimentResponse</span><span style=\"font-weight: bold\">(</span>\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">experiment_id</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">14</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">training_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TrainingConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">architecture_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">BasicTransformerConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">attention_bias</span>=<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">head_dim</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">64</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">hidden_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">448</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">intermediate_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1280</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">num_attention_heads</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">num_hidden_layers</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">num_key_value_heads</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">rms_norm_eps</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-06</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">rope_theta</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1000000</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">tie_word_embeddings</span>=<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">dtype</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'bfloat16'</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">vocab_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32000</span>\n",
       "            <span style=\"font-weight: bold\">)</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">optimizer_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">AdamWConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">lr_scheduler</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">WarmupCosineDecay</span><span style=\"font-weight: bold\">(</span>\n",
       "                    <span style=\"color: #808000; text-decoration-color: #808000\">peak_value</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0003</span>,\n",
       "                    <span style=\"color: #808000; text-decoration-color: #808000\">final_lr_frac</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1</span>,\n",
       "                    <span style=\"color: #808000; text-decoration-color: #808000\">warmup_frac</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.05</span>,\n",
       "                    <span style=\"color: #808000; text-decoration-color: #808000\">init_value</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0</span>\n",
       "                <span style=\"font-weight: bold\">)</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">weight_decay</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.01</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">beta1</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">beta2</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.95</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">eps</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-08</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">eps_root</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-08</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">grad_clip_norm</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>\n",
       "            <span style=\"font-weight: bold\">)</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">train_batch_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">val_batch_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">n_evals</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">16</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">total_train_tokens</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8388608</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">max_runtime_seconds</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">60.0</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">model_seed</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>\n",
       "        <span style=\"font-weight: bold\">)</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">status</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">RunningExperimentStatus</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">queued_at</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">datetime</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">.datetime</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2026</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">29</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">23</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">46</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">36</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">972380</span>, <span style=\"color: #808000; text-decoration-color: #808000\">tzinfo</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TzInfo</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span><span style=\"font-weight: bold\">))</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">dispatched_at</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">datetime</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">.datetime</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2026</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">29</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">23</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">46</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">38</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">244101</span>, <span style=\"color: #808000; text-decoration-color: #808000\">tzinfo</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TzInfo</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span><span style=\"font-weight: bold\">))</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">run_id</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'fc-01KQDT6BXQFMPGVWR3GTS3PVJ3'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">val_losses</span>=<span style=\"font-weight: bold\">[]</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">status_type</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'running'</span>\n",
       "        <span style=\"font-weight: bold\">)</span>\n",
       "    <span style=\"font-weight: bold\">)</span>\n",
       "<span style=\"font-weight: bold\">]</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m[\u001b[0m\n",
       "    \u001b[1;35mExperimentResponse\u001b[0m\u001b[1m(\u001b[0m\n",
       "        \u001b[33mexperiment_id\u001b[0m=\u001b[1;36m14\u001b[0m,\n",
       "        \u001b[33mtraining_config\u001b[0m=\u001b[1;35mTrainingConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33marchitecture_config\u001b[0m=\u001b[1;35mBasicTransformerConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "                \u001b[33mattention_bias\u001b[0m=\u001b[3;91mFalse\u001b[0m,\n",
       "                \u001b[33mhead_dim\u001b[0m=\u001b[1;36m64\u001b[0m,\n",
       "                \u001b[33mhidden_size\u001b[0m=\u001b[1;36m448\u001b[0m,\n",
       "                \u001b[33mintermediate_size\u001b[0m=\u001b[1;36m1280\u001b[0m,\n",
       "                \u001b[33mnum_attention_heads\u001b[0m=\u001b[1;36m7\u001b[0m,\n",
       "                \u001b[33mnum_hidden_layers\u001b[0m=\u001b[1;36m9\u001b[0m,\n",
       "                \u001b[33mnum_key_value_heads\u001b[0m=\u001b[1;36m7\u001b[0m,\n",
       "                \u001b[33mrms_norm_eps\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-06\u001b[0m,\n",
       "                \u001b[33mrope_theta\u001b[0m=\u001b[1;36m1000000\u001b[0m,\n",
       "                \u001b[33mtie_word_embeddings\u001b[0m=\u001b[3;91mFalse\u001b[0m,\n",
       "                \u001b[33mdtype\u001b[0m=\u001b[32m'bfloat16'\u001b[0m,\n",
       "                \u001b[33mvocab_size\u001b[0m=\u001b[1;36m32000\u001b[0m\n",
       "            \u001b[1m)\u001b[0m,\n",
       "            \u001b[33moptimizer_config\u001b[0m=\u001b[1;35mAdamWConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "                \u001b[33mlr_scheduler\u001b[0m=\u001b[1;35mWarmupCosineDecay\u001b[0m\u001b[1m(\u001b[0m\n",
       "                    \u001b[33mpeak_value\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.0003\u001b[0m,\n",
       "                    \u001b[33mfinal_lr_frac\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.1\u001b[0m,\n",
       "                    \u001b[33mwarmup_frac\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.05\u001b[0m,\n",
       "                    \u001b[33minit_value\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.0\u001b[0m\n",
       "                \u001b[1m)\u001b[0m,\n",
       "                \u001b[33mweight_decay\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.01\u001b[0m,\n",
       "                \u001b[33mbeta1\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.9\u001b[0m,\n",
       "                \u001b[33mbeta2\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.95\u001b[0m,\n",
       "                \u001b[33meps\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-08\u001b[0m,\n",
       "                \u001b[33meps_root\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-08\u001b[0m,\n",
       "                \u001b[33mgrad_clip_norm\u001b[0m=\u001b[1;36m1\u001b[0m\u001b[1;36m.0\u001b[0m\n",
       "            \u001b[1m)\u001b[0m,\n",
       "            \u001b[33mtrain_batch_size\u001b[0m=\u001b[1;36m128\u001b[0m,\n",
       "            \u001b[33mval_batch_size\u001b[0m=\u001b[1;36m32\u001b[0m,\n",
       "            \u001b[33mn_evals\u001b[0m=\u001b[1;36m16\u001b[0m,\n",
       "            \u001b[33mtotal_train_tokens\u001b[0m=\u001b[1;36m8388608\u001b[0m,\n",
       "            \u001b[33mmax_runtime_seconds\u001b[0m=\u001b[1;36m60\u001b[0m\u001b[1;36m.0\u001b[0m,\n",
       "            \u001b[33mmodel_seed\u001b[0m=\u001b[1;36m0\u001b[0m\n",
       "        \u001b[1m)\u001b[0m,\n",
       "        \u001b[33mstatus\u001b[0m=\u001b[1;35mRunningExperimentStatus\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33mqueued_at\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2026\u001b[0m, \u001b[1;36m4\u001b[0m, \u001b[1;36m29\u001b[0m, \u001b[1;36m23\u001b[0m, \u001b[1;36m46\u001b[0m, \u001b[1;36m36\u001b[0m, \u001b[1;36m972380\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
       "            \u001b[33mdispatched_at\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2026\u001b[0m, \u001b[1;36m4\u001b[0m, \u001b[1;36m29\u001b[0m, \u001b[1;36m23\u001b[0m, \u001b[1;36m46\u001b[0m, \u001b[1;36m38\u001b[0m, \u001b[1;36m244101\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
       "            \u001b[33mrun_id\u001b[0m=\u001b[32m'fc-01KQDT6BXQFMPGVWR3GTS3PVJ3'\u001b[0m,\n",
       "            \u001b[33mval_losses\u001b[0m=\u001b[1m[\u001b[0m\u001b[1m]\u001b[0m,\n",
       "            \u001b[33mstatus_type\u001b[0m=\u001b[32m'running'\u001b[0m\n",
       "        \u001b[1m)\u001b[0m\n",
       "    \u001b[1m)\u001b[0m\n",
       "\u001b[1m]\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# after a few seconds, it will be in running state, as long as the queue is not full\n",
    "time.sleep(10)\n",
    "rich.print(list_experiments())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "78ce7bec",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">ExperimentResponse</span><span style=\"font-weight: bold\">(</span>\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">experiment_id</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">14</span>,\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">training_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TrainingConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">architecture_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">BasicTransformerConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">attention_bias</span>=<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">head_dim</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">64</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">hidden_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">448</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">intermediate_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1280</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">num_attention_heads</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">num_hidden_layers</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">num_key_value_heads</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">rms_norm_eps</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-06</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">rope_theta</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1000000</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">tie_word_embeddings</span>=<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">dtype</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'bfloat16'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">vocab_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32000</span>\n",
       "        <span style=\"font-weight: bold\">)</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">optimizer_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">AdamWConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">lr_scheduler</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">WarmupCosineDecay</span><span style=\"font-weight: bold\">(</span>\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">peak_value</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0003</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">final_lr_frac</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">warmup_frac</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.05</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">init_value</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0</span>\n",
       "            <span style=\"font-weight: bold\">)</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">weight_decay</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.01</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">beta1</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">beta2</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.95</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">eps</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-08</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">eps_root</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-08</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">grad_clip_norm</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>\n",
       "        <span style=\"font-weight: bold\">)</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">train_batch_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">val_batch_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">n_evals</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">16</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">total_train_tokens</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8388608</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">max_runtime_seconds</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">60.0</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">model_seed</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>\n",
       "    <span style=\"font-weight: bold\">)</span>,\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">status</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">CompletedExperimentStatus</span><span style=\"font-weight: bold\">(</span>\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">queued_at</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">datetime</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">.datetime</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2026</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">29</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">23</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">46</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">36</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">972380</span>, <span style=\"color: #808000; text-decoration-color: #808000\">tzinfo</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TzInfo</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span><span style=\"font-weight: bold\">))</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">dispatched_at</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">datetime</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">.datetime</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2026</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">29</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">23</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">46</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">38</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">244101</span>, <span style=\"color: #808000; text-decoration-color: #808000\">tzinfo</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TzInfo</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span><span style=\"font-weight: bold\">))</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">run_id</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'fc-01KQDT6BXQFMPGVWR3GTS3PVJ3'</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">used_runtime_seconds</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">10.18408550300228</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">val_losses</span>=<span style=\"font-weight: bold\">[</span>\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.625</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8.92578125</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8.37890625</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.97265625</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.607421875</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.4453125</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.328125</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.291015625</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.25</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.22265625</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.208984375</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.19921875</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.193359375</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.189453125</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.189453125</span>,\n",
       "            <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.189453125</span>\n",
       "        <span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">completed_at</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">datetime</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">.datetime</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2026</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">29</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">23</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">47</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">25</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">805998</span>, <span style=\"color: #808000; text-decoration-color: #808000\">tzinfo</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TzInfo</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span><span style=\"font-weight: bold\">))</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">status_type</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'completed'</span>\n",
       "    <span style=\"font-weight: bold\">)</span>\n",
       "<span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mExperimentResponse\u001b[0m\u001b[1m(\u001b[0m\n",
       "    \u001b[33mexperiment_id\u001b[0m=\u001b[1;36m14\u001b[0m,\n",
       "    \u001b[33mtraining_config\u001b[0m=\u001b[1;35mTrainingConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "        \u001b[33marchitecture_config\u001b[0m=\u001b[1;35mBasicTransformerConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33mattention_bias\u001b[0m=\u001b[3;91mFalse\u001b[0m,\n",
       "            \u001b[33mhead_dim\u001b[0m=\u001b[1;36m64\u001b[0m,\n",
       "            \u001b[33mhidden_size\u001b[0m=\u001b[1;36m448\u001b[0m,\n",
       "            \u001b[33mintermediate_size\u001b[0m=\u001b[1;36m1280\u001b[0m,\n",
       "            \u001b[33mnum_attention_heads\u001b[0m=\u001b[1;36m7\u001b[0m,\n",
       "            \u001b[33mnum_hidden_layers\u001b[0m=\u001b[1;36m9\u001b[0m,\n",
       "            \u001b[33mnum_key_value_heads\u001b[0m=\u001b[1;36m7\u001b[0m,\n",
       "            \u001b[33mrms_norm_eps\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-06\u001b[0m,\n",
       "            \u001b[33mrope_theta\u001b[0m=\u001b[1;36m1000000\u001b[0m,\n",
       "            \u001b[33mtie_word_embeddings\u001b[0m=\u001b[3;91mFalse\u001b[0m,\n",
       "            \u001b[33mdtype\u001b[0m=\u001b[32m'bfloat16'\u001b[0m,\n",
       "            \u001b[33mvocab_size\u001b[0m=\u001b[1;36m32000\u001b[0m\n",
       "        \u001b[1m)\u001b[0m,\n",
       "        \u001b[33moptimizer_config\u001b[0m=\u001b[1;35mAdamWConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33mlr_scheduler\u001b[0m=\u001b[1;35mWarmupCosineDecay\u001b[0m\u001b[1m(\u001b[0m\n",
       "                \u001b[33mpeak_value\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.0003\u001b[0m,\n",
       "                \u001b[33mfinal_lr_frac\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.1\u001b[0m,\n",
       "                \u001b[33mwarmup_frac\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.05\u001b[0m,\n",
       "                \u001b[33minit_value\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.0\u001b[0m\n",
       "            \u001b[1m)\u001b[0m,\n",
       "            \u001b[33mweight_decay\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.01\u001b[0m,\n",
       "            \u001b[33mbeta1\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.9\u001b[0m,\n",
       "            \u001b[33mbeta2\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.95\u001b[0m,\n",
       "            \u001b[33meps\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-08\u001b[0m,\n",
       "            \u001b[33meps_root\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-08\u001b[0m,\n",
       "            \u001b[33mgrad_clip_norm\u001b[0m=\u001b[1;36m1\u001b[0m\u001b[1;36m.0\u001b[0m\n",
       "        \u001b[1m)\u001b[0m,\n",
       "        \u001b[33mtrain_batch_size\u001b[0m=\u001b[1;36m128\u001b[0m,\n",
       "        \u001b[33mval_batch_size\u001b[0m=\u001b[1;36m32\u001b[0m,\n",
       "        \u001b[33mn_evals\u001b[0m=\u001b[1;36m16\u001b[0m,\n",
       "        \u001b[33mtotal_train_tokens\u001b[0m=\u001b[1;36m8388608\u001b[0m,\n",
       "        \u001b[33mmax_runtime_seconds\u001b[0m=\u001b[1;36m60\u001b[0m\u001b[1;36m.0\u001b[0m,\n",
       "        \u001b[33mmodel_seed\u001b[0m=\u001b[1;36m0\u001b[0m\n",
       "    \u001b[1m)\u001b[0m,\n",
       "    \u001b[33mstatus\u001b[0m=\u001b[1;35mCompletedExperimentStatus\u001b[0m\u001b[1m(\u001b[0m\n",
       "        \u001b[33mqueued_at\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2026\u001b[0m, \u001b[1;36m4\u001b[0m, \u001b[1;36m29\u001b[0m, \u001b[1;36m23\u001b[0m, \u001b[1;36m46\u001b[0m, \u001b[1;36m36\u001b[0m, \u001b[1;36m972380\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
       "        \u001b[33mdispatched_at\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2026\u001b[0m, \u001b[1;36m4\u001b[0m, \u001b[1;36m29\u001b[0m, \u001b[1;36m23\u001b[0m, \u001b[1;36m46\u001b[0m, \u001b[1;36m38\u001b[0m, \u001b[1;36m244101\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
       "        \u001b[33mrun_id\u001b[0m=\u001b[32m'fc-01KQDT6BXQFMPGVWR3GTS3PVJ3'\u001b[0m,\n",
       "        \u001b[33mused_runtime_seconds\u001b[0m=\u001b[1;36m10\u001b[0m\u001b[1;36m.18408550300228\u001b[0m,\n",
       "        \u001b[33mval_losses\u001b[0m=\u001b[1m[\u001b[0m\n",
       "            \u001b[1;36m9.625\u001b[0m,\n",
       "            \u001b[1;36m8.92578125\u001b[0m,\n",
       "            \u001b[1;36m8.37890625\u001b[0m,\n",
       "            \u001b[1;36m7.97265625\u001b[0m,\n",
       "            \u001b[1;36m7.607421875\u001b[0m,\n",
       "            \u001b[1;36m7.4453125\u001b[0m,\n",
       "            \u001b[1;36m7.328125\u001b[0m,\n",
       "            \u001b[1;36m7.291015625\u001b[0m,\n",
       "            \u001b[1;36m7.25\u001b[0m,\n",
       "            \u001b[1;36m7.22265625\u001b[0m,\n",
       "            \u001b[1;36m7.208984375\u001b[0m,\n",
       "            \u001b[1;36m7.19921875\u001b[0m,\n",
       "            \u001b[1;36m7.193359375\u001b[0m,\n",
       "            \u001b[1;36m7.189453125\u001b[0m,\n",
       "            \u001b[1;36m7.189453125\u001b[0m,\n",
       "            \u001b[1;36m7.189453125\u001b[0m\n",
       "        \u001b[1m]\u001b[0m,\n",
       "        \u001b[33mcompleted_at\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2026\u001b[0m, \u001b[1;36m4\u001b[0m, \u001b[1;36m29\u001b[0m, \u001b[1;36m23\u001b[0m, \u001b[1;36m47\u001b[0m, \u001b[1;36m25\u001b[0m, \u001b[1;36m805998\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
       "        \u001b[33mstatus_type\u001b[0m=\u001b[32m'completed'\u001b[0m\n",
       "    \u001b[1m)\u001b[0m\n",
       "\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# finally, it will turn into finished state\n",
    "time.sleep(60)\n",
    "rich.print(get_experiment(submit_result.experiment_id))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f10110c",
   "metadata": {},
   "source": [
    "### We are refunded the difference between the budget and the cost of the experiment once it finishes\n",
    "\n",
    "We reserved 60 seconds for our experiment, but "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f8e5d42c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BudgetSummary(used_seconds=10.18408550300228, remaining_seconds=43189.815914497, total_budget_seconds=43200.0)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_budget()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5020138f",
   "metadata": {},
   "source": [
    "### Submit your final validation. This will run for 48 B200-hours"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9ab4c6a0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">FinalSubmissionResponse</span><span style=\"font-weight: bold\">(</span>\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">training_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TrainingConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">architecture_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">BasicTransformerConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">attention_bias</span>=<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">head_dim</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">64</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">hidden_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">448</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">intermediate_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1280</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">num_attention_heads</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">num_hidden_layers</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">num_key_value_heads</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">rms_norm_eps</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-06</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">rope_theta</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1000000</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">tie_word_embeddings</span>=<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">dtype</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'bfloat16'</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">vocab_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32000</span>\n",
       "        <span style=\"font-weight: bold\">)</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">optimizer_config</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">AdamWConfig</span><span style=\"font-weight: bold\">(</span>\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">lr_scheduler</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">WarmupCosineDecay</span><span style=\"font-weight: bold\">(</span>\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">peak_value</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0003</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">final_lr_frac</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">warmup_frac</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.05</span>,\n",
       "                <span style=\"color: #808000; text-decoration-color: #808000\">init_value</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0</span>\n",
       "            <span style=\"font-weight: bold\">)</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">weight_decay</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.01</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">beta1</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">beta2</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.95</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">eps</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-08</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">eps_root</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1e-08</span>,\n",
       "            <span style=\"color: #808000; text-decoration-color: #808000\">grad_clip_norm</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0</span>\n",
       "        <span style=\"font-weight: bold\">)</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">train_batch_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">val_batch_size</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">n_evals</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">16</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">total_train_tokens</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8388608</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">max_runtime_seconds</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">60.0</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">model_seed</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>\n",
       "    <span style=\"font-weight: bold\">)</span>,\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">predicted_final_loss</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.2</span>,\n",
       "    <span style=\"color: #808000; text-decoration-color: #808000\">submitted_at</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">datetime</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">.datetime</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2026</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">29</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">23</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">47</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">47</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">218799</span>, <span style=\"color: #808000; text-decoration-color: #808000\">tzinfo</span>=<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">TzInfo</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span><span style=\"font-weight: bold\">))</span>\n",
       "<span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mFinalSubmissionResponse\u001b[0m\u001b[1m(\u001b[0m\n",
       "    \u001b[33mtraining_config\u001b[0m=\u001b[1;35mTrainingConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "        \u001b[33marchitecture_config\u001b[0m=\u001b[1;35mBasicTransformerConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33mattention_bias\u001b[0m=\u001b[3;91mFalse\u001b[0m,\n",
       "            \u001b[33mhead_dim\u001b[0m=\u001b[1;36m64\u001b[0m,\n",
       "            \u001b[33mhidden_size\u001b[0m=\u001b[1;36m448\u001b[0m,\n",
       "            \u001b[33mintermediate_size\u001b[0m=\u001b[1;36m1280\u001b[0m,\n",
       "            \u001b[33mnum_attention_heads\u001b[0m=\u001b[1;36m7\u001b[0m,\n",
       "            \u001b[33mnum_hidden_layers\u001b[0m=\u001b[1;36m9\u001b[0m,\n",
       "            \u001b[33mnum_key_value_heads\u001b[0m=\u001b[1;36m7\u001b[0m,\n",
       "            \u001b[33mrms_norm_eps\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-06\u001b[0m,\n",
       "            \u001b[33mrope_theta\u001b[0m=\u001b[1;36m1000000\u001b[0m,\n",
       "            \u001b[33mtie_word_embeddings\u001b[0m=\u001b[3;91mFalse\u001b[0m,\n",
       "            \u001b[33mdtype\u001b[0m=\u001b[32m'bfloat16'\u001b[0m,\n",
       "            \u001b[33mvocab_size\u001b[0m=\u001b[1;36m32000\u001b[0m\n",
       "        \u001b[1m)\u001b[0m,\n",
       "        \u001b[33moptimizer_config\u001b[0m=\u001b[1;35mAdamWConfig\u001b[0m\u001b[1m(\u001b[0m\n",
       "            \u001b[33mlr_scheduler\u001b[0m=\u001b[1;35mWarmupCosineDecay\u001b[0m\u001b[1m(\u001b[0m\n",
       "                \u001b[33mpeak_value\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.0003\u001b[0m,\n",
       "                \u001b[33mfinal_lr_frac\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.1\u001b[0m,\n",
       "                \u001b[33mwarmup_frac\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.05\u001b[0m,\n",
       "                \u001b[33minit_value\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.0\u001b[0m\n",
       "            \u001b[1m)\u001b[0m,\n",
       "            \u001b[33mweight_decay\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.01\u001b[0m,\n",
       "            \u001b[33mbeta1\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.9\u001b[0m,\n",
       "            \u001b[33mbeta2\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.95\u001b[0m,\n",
       "            \u001b[33meps\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-08\u001b[0m,\n",
       "            \u001b[33meps_root\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-08\u001b[0m,\n",
       "            \u001b[33mgrad_clip_norm\u001b[0m=\u001b[1;36m1\u001b[0m\u001b[1;36m.0\u001b[0m\n",
       "        \u001b[1m)\u001b[0m,\n",
       "        \u001b[33mtrain_batch_size\u001b[0m=\u001b[1;36m128\u001b[0m,\n",
       "        \u001b[33mval_batch_size\u001b[0m=\u001b[1;36m32\u001b[0m,\n",
       "        \u001b[33mn_evals\u001b[0m=\u001b[1;36m16\u001b[0m,\n",
       "        \u001b[33mtotal_train_tokens\u001b[0m=\u001b[1;36m8388608\u001b[0m,\n",
       "        \u001b[33mmax_runtime_seconds\u001b[0m=\u001b[1;36m60\u001b[0m\u001b[1;36m.0\u001b[0m,\n",
       "        \u001b[33mmodel_seed\u001b[0m=\u001b[1;36m0\u001b[0m\n",
       "    \u001b[1m)\u001b[0m,\n",
       "    \u001b[33mpredicted_final_loss\u001b[0m=\u001b[1;36m7\u001b[0m\u001b[1;36m.2\u001b[0m,\n",
       "    \u001b[33msubmitted_at\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2026\u001b[0m, \u001b[1;36m4\u001b[0m, \u001b[1;36m29\u001b[0m, \u001b[1;36m23\u001b[0m, \u001b[1;36m47\u001b[0m, \u001b[1;36m47\u001b[0m, \u001b[1;36m218799\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\n",
       "\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "rich.print(save_final_submission(training_config, 7.2))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
