Final submission: code, predictions, training

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+training/data/star3-thoughts.jsonl filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,76 @@
 ---
+language:
+- si
 license: apache-2.0
+base_model: google/gemma-4-E4B-it
+tags:
+- sinhala
+- mmlu
+- multiple-choice
+- lora
+- star
 ---
+
+# Team SinhalaNuwana - Sinhala MMLU Shared Task System
+
+Codabench macro score 63.5755 (Easy 64.3302, Medium 56.7057, Hard 69.6907).
+
+## System summary
+- **Base model(s):** [google/gemma-4-E4B-it](https://huggingface.co/google/gemma-4-E4B-it), revision `ee0ef6023621cff504d758262d4e04895a5af4a2`, text decoder only
+- **What we did:** LoRA fine-tuning with STaR (self-taught reasoner): the base model thought about 8,894 Sinhala training questions, and we trained a LoRA adapter on its 4,073 correct first-try thoughts. At inference the model thinks in Sinhala before it answers (thinking mode, greedy). Each question is answered with the choices in several orders, and the option probabilities are averaged. Easy and Hard questions also get a longer thought of up to 2,048 tokens, and Medium questions 4 more choice orders. One model, no retrieval.
+- **Total inference parameters:** 7.50B (7,501,679,616: 7,463,013,376 base and 38,666,240 adapter)
+- **Extra data used for training or retrieval:** Our own Sinhala multiple-choice questions, generated by Claude models from Sri Lankan school textbooks and teacher guides. They are in `training/data/` of this repository. No retrieval.
+
+## Models and resources in this submission
+| Repo | What it is |
+|---|---|
+| [`SinhalaMMLU-SharedTask/team-sinhalanuwana`](https://huggingface.co/SinhalaMMLU-SharedTask/team-sinhalanuwana) | code, inference script, final predictions, training code and training data (this repo) |
+| [`SinhalaMMLU-SharedTask/team-sinhalanuwana-gemma4-e4b-star3-lora`](https://huggingface.co/SinhalaMMLU-SharedTask/team-sinhalanuwana-gemma4-e4b-star3-lora) | LoRA adapter on gemma-4-E4B-it: weights, config, tokenizer files, training log |
+| [`google/gemma-4-E4B-it`](https://huggingface.co/google/gemma-4-E4B-it) | public base model (Apache 2.0) |
+
+## How to run
+```bash
+pip install -r requirements.txt
+python inference.py
+```
+- **GPU used:** Kaggle T4 x2 (16 GB each). One GPU with 24 GB or more also holds the model.
+- **Approximate run time on the test set:** about 16 hours for the 2,238 questions on T4 x2 (six passes of 4.3, 4.2, 1.7, 1.9, 2.1 and 2.1 hours)
+- **Anything else the organizers need to know:**
+  - Python 3.13. If `torchao` is installed (Kaggle has it), run `pip uninstall -y torchao` first, because `peft` imports it and rejects that version.
+  - `inference.py` calls `predict.py` and `merge_votes.py` from the same folder, so run it inside a clone of this repository.
+  - It writes `predictions.jsonl` in the current folder, which replaces our submitted copy in a clone. The submitted file stays in the repository history.
+  - Only `DATA_REPO` needs to change. No token is needed if the repositories are public.
+
+## Files
+- `inference.py` - the exact system that produced our final Codabench submission
+- `predict.py`, `merge_votes.py` - one answering pass and the vote over the passes, called by `inference.py`
+- `predictions.jsonl` - the exact file we submitted to the final leaderboard
+- `requirements.txt` - package versions we used
+- `manifest.yaml`, `model_manifest.json` - models, adapter SHA-256, parameter count, closed-API use
+- `training/` - training and data-preparation code, Kaggle notebooks, settings and training data (see `training/README.md`)
+
+## Team details
+
+| Name | Affiliation | Role |
+| --- | --- | --- |
+| Senali Dilumika | UCSC | Corresponding member |
+
+Team name on Codabench: SinhalaNuwana.
+
+## Verification
+
+- Fresh Kaggle session (2 x T4, 28 October 2026), the same code, the same pinned versions and the same adapter, on the first 3 questions of each of the 41 test files (123 questions): 122 answers are the same as in `predictions.jsonl`. All Medium and Hard answers are the same. The one changed answer is an Easy question whose 5 votes split 2-2-1. Float16 GPU arithmetic is not fully deterministic, so a full rerun can change a few close votes.
+- `merge_votes.py`, run on the six scored passes, rebuilds `predictions.jsonl` byte for byte.
+
+## Training (details in `training/README.md`)
+
+`training/README.md` gives the full pipeline, the settings and the data. In short:
+
+1. Claude models generated Sinhala multiple-choice questions from Sri Lankan school textbooks, NIE teacher guides and resource books, A/L learning outcome notes and Wikibooks. A stronger Claude model checked the answer keys of 10,000 of them.
+2. `review_overlap.py` removed the 1,106 questions that resemble a dev or test question by wording or by meaning (LaBSE cosine 0.70 or more). 8,894 questions remained.
+3. `star.py` let the base model (no adapter) think about each question for up to 512 tokens. It kept the 4,073 thoughts that gave the correct answer on the first try.
+4. `train_lora.py` trained a new LoRA adapter on those thoughts: 2 epochs, AdamW, learning rate 1e-4 with a cosine schedule and 3% warmup, 16 examples per step, sequences of up to 1,024 tokens. It took 4.3 hours on 2 x T4.
+
+## Licence and citation
+
+The code and the adapter are released under the Apache 2.0 licence, the licence of Gemma 4.

inference.py

@@ -0,0 +1,104 @@
+"""
+Sinhala MMLU Shared Task - Final inference script
+Team: SinhalaNuwana
+
+    pip install -r requirements.txt
+    python inference.py
+
+Downloads the questions from DATA_REPO, the base model and our LoRA adapter from Hugging Face,
+runs predict.py six times and merge_votes.py once (both next to this file), and writes
+predictions.jsonl and predictions.zip in the current folder. Greedy decoding, no sampling.
+About 16 hours for the 2,238 test questions on 2 x T4.
+"""
+
+# ======================= ORGANIZERS CHANGE ONLY THIS LINE =======================
+DATA_REPO = "SinhalaMMLU-SharedTask/Sinhala-MMLU-SharedTask-Test-Dataset"
+# ================================================================================
+
+import argparse
+import hashlib
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+MODEL = "google/gemma-4-E4B-it"
+REVISION = "ee0ef6023621cff504d758262d4e04895a5af4a2"
+ADAPTER = "SinhalaMMLU-SharedTask/team-sinhalanuwana-gemma4-e4b-star3-lora"
+ADAPTER_SHA256 = "b82696e2541e0133d24bf14721caa20485dd70ab8ef78a7e40c1d2a77db1ba71"
+ROOT = Path(__file__).resolve().parent
+
+PASSES = [
+    ("orders-02", ["--think", "512", "--rotations", "0", "2"]),
+    ("orders-13", ["--think", "512", "--rotations", "1", "3"]),
+    ("hard-2048", ["--level", "hard", "--think", "2048", "--score_batch_size", "8"]),
+    ("easy-2048", ["--level", "easy", "--think", "2048", "--score_batch_size", "8"]),
+    ("medium-rev02", ["--level", "medium", "--reverse", "--think", "512", "--rotations", "0", "2"]),
+    ("medium-rev13", ["--level", "medium", "--reverse", "--think", "512", "--rotations", "1", "3"]),
+]
+
+
+def slice_questions(data_dir, out_dir, per_file):
+    for path in sorted(data_dir.rglob("*.json")):
+        if any(part.startswith(".") for part in path.relative_to(data_dir).parts):
+            continue
+        target = out_dir / path.relative_to(data_dir)
+        target.parent.mkdir(parents=True, exist_ok=True)
+        questions = json.loads(path.read_text(encoding="utf-8"))[:per_file]
+        target.write_text(json.dumps(questions, ensure_ascii=False), encoding="utf-8")
+    return out_dir
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Answer the SinhalaMMLU questions with Gemma 4 E4B-it and the round 3 STaR LoRA adapter, "
+        "and write OUT_DIR/predictions.jsonl and OUT_DIR/predictions.zip"
+    )
+    parser.add_argument("--data_dir", type=Path, help=f"local folder of the per-subject .json files; {DATA_REPO} when not given")
+    parser.add_argument("--out_dir", type=Path, default=Path("."), help="where predictions.jsonl goes; each pass writes to OUT_DIR/runs/")
+    parser.add_argument("--adapter", help=f"local adapter folder; {ADAPTER} from the Hub when not given")
+    parser.add_argument("--per_file", type=int, help="answer only the first N questions of each file, for a quick check")
+    args = parser.parse_args()
+
+    snapshot_download(MODEL, revision=REVISION)
+    adapter = Path(args.adapter or snapshot_download(ADAPTER))
+    digest = hashlib.sha256((adapter / "adapter_model.safetensors").read_bytes()).hexdigest()
+    if digest != ADAPTER_SHA256:
+        raise SystemExit(f"error: {adapter}/adapter_model.safetensors is {digest}, expected {ADAPTER_SHA256}")
+
+    data_dir = args.data_dir or Path(snapshot_download(DATA_REPO, repo_type="dataset"))
+    runs = args.out_dir / "runs"
+    if args.per_file:
+        data_dir = slice_questions(data_dir, runs / "questions", args.per_file)
+
+    env = {**os.environ, "HF_HUB_OFFLINE": "1"}
+    long_env = {**env, "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}
+    common = [
+        "--model", MODEL, "--revision", REVISION, "--dtype", "float16", "--batch_size", "32",
+        "--adapter", str(adapter), "--data_dir", str(data_dir),
+    ]
+    for name, options in PASSES:
+        print(f"== pass {name}", flush=True)
+        subprocess.run(
+            [sys.executable, str(ROOT / "predict.py"), *common, "--out_dir", str(runs / name), *options],
+            check=True,
+            env=long_env if "2048" in options else env,
+        )
+
+    run = {name: str(runs / name) for name, _ in PASSES}
+    subprocess.run(
+        [
+            sys.executable, str(ROOT / "merge_votes.py"), run["orders-02"], run["orders-13"],
+            "--override", run["hard-2048"], "--add", run["easy-2048"],
+            "--add", run["medium-rev02"], "--add", run["medium-rev13"], "--out_dir", str(args.out_dir),
+        ],
+        check=True,
+        env=env,
+    )
+
+
+if __name__ == "__main__":
+    main()

manifest.yaml

@@ -0,0 +1,54 @@
+system: SinhalaNuwana
+team: SinhalaNuwana
+models:
+  - id: google/gemma-4-E4B-it
+    revision: ee0ef6023621cff504d758262d4e04895a5af4a2
+    params: 7463013376
+    role: answer selection
+  - id: adapter/
+    sha256: b82696e2541e0133d24bf14721caa20485dd70ab8ef78a7e40c1d2a77db1ba71
+    params: 38666240
+    role: adapter on the answer-selection model
+total_params: 7501679616
+closed_api_used_for_data: true
+closed_api_details: >-
+  Anthropic claude-haiku-5-5 generated the 18,398 training questions in data/train/mcq.jsonl from
+  Sri Lankan school textbooks, NIE teacher guides and resource books, A/L learning outcome notes
+  and Wikibooks. Anthropic claude-sonnet-5-5 (9,228) and claude-opus-5-5 (772) checked the answer
+  keys of 10,000 of them (data/train/verified-mcq.jsonl). review_overlap.py removed the 1,106 that
+  match a dev or test question by wording or by meaning (LaBSE cosine 0.70 or more). The model
+  without any adapter thinks about the other 8,894, and the adapter trains from the base model on
+  its 4,073 correct first-try thoughts (STaR round 3, star.py --no_hint_retry, 2 epochs). No closed
+  API is used at inference, and no dev or test question is used for training.
+retrieval_or_internet_at_inference: false
+decoding:
+  method: >-
+    zero-shot; Sinhala prompt in the chat template with thinking mode on; a greedy thought, then the
+    next-token logits after the closed thought and the answer cue "පිළිතුර:", softmax over the
+    option labels "1".."4". Every question is answered in all 4 choice orders (rotated by 0, 1, 2
+    and 3) with a thought of up to 512 tokens; the probability vectors are mapped back to the
+    original order and averaged. Hard questions (metadata difficulty "hard") are answered once more
+    in the given order with a thought of up to 2,048 tokens, and that answer replaces the 4-order
+    vote for them, because most hard thoughts do not finish within 512 tokens. Easy questions are
+    also answered once more in the given order with a thought of up to 2,048 tokens, which joins
+    their 4 orders as a fifth vote. Medium questions are also answered in the 4 reversed orders
+    (predict.py --reverse: the choices as 4, 3, 2, 1, rotated by 0, 1, 2 and 3) with a thought of up
+    to 512 tokens, so they get 8 orders. merge_votes.py combines the passes with equal weight per
+    thought; the highest average probability is the answer
+  reasoning_tokens: {easy: [512, 2048], medium: 512, hard: 2048}
+  choice_orders: {easy: [0, 1, 2, 3, "0 at 2048 tokens"], medium: [0, 1, 2, 3, r0, r1, r2, r3], hard: [0]}
+  sampling: none
+  batch_size: 32
+  score_batch_size: {easy: "32 (512 tokens), 8 (2048 tokens)", medium: 32, hard: 8}
+  batch_tokens: 36864
+  dtype: float16
+  device_map: auto
+  prompt: predict.py (user_prompt, render)
+environment:
+  requirements: requirements.txt
+  python: "3.13"
+  gpu: >-
+    NVIDIA; verified on two Tesla T4 (16 GB each) on Kaggle; one GPU of 24 GB or more also holds
+    the model
+  dev_accuracy: null
+  test_accuracy_macro: 0.635755

merge_votes.py

@@ -0,0 +1,108 @@
+import argparse
+import json
+import zipfile
+from pathlib import Path
+
+LABELS = ["1", "2", "3", "4"]
+
+
+def key(row):
+    return (row["subject"], row["level"], row["q_no"])
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Average the answer probabilities of every thought in several predict.py runs on the same "
+        "questions, with equal weight per thought, and write one predictions.zip"
+    )
+    parser.add_argument("run_dirs", nargs="+", type=Path, help="predict.py output folders with thoughts.jsonl")
+    parser.add_argument("--out_dir", type=Path, required=True)
+    parser.add_argument(
+        "--override", type=Path, action="append", default=[],
+        help="a run on a subset of the questions whose thoughts replace the others for those questions; repeat it for more subsets",
+    )
+    parser.add_argument(
+        "--add", type=Path, action="append", default=[],
+        help="a run on a subset of the questions whose thoughts join the others for those questions, with the same weight per thought; repeat it for more subsets",
+    )
+    parser.add_argument(
+        "--tiebreak", type=Path,
+        help="a run, usually without a thought, whose average probabilities pick between the answers that got the same "
+        "most votes; questions without a tie or not in this run keep the highest average probability",
+    )
+    args = parser.parse_args()
+
+    probs = {}
+    for run_dir in args.run_dirs:
+        for line in (run_dir / "thoughts.jsonl").open(encoding="utf-8"):
+            row = json.loads(line)
+            probs.setdefault(key(row), []).append(row["probs"])
+
+    counts = {len(group) for group in probs.values()}
+    if len(counts) != 1:
+        raise SystemExit(f"error: questions have different numbers of thoughts: {sorted(counts)}")
+    print(f"Averaging {counts.pop()} thoughts for each of {len(probs)} questions")
+
+    for override in args.override:
+        replaced = {}
+        for line in (override / "thoughts.jsonl").open(encoding="utf-8"):
+            row = json.loads(line)
+            if key(row) not in probs:
+                raise SystemExit(f"error: {key(row)} from {override} is not in the merged runs")
+            replaced.setdefault(key(row), []).append(row["probs"])
+        probs.update(replaced)
+        print(f"Replaced the thoughts of {len(replaced)} questions with those of {override}")
+
+    for extra in args.add:
+        added = set()
+        for line in (extra / "thoughts.jsonl").open(encoding="utf-8"):
+            row = json.loads(line)
+            if key(row) not in probs:
+                raise SystemExit(f"error: {key(row)} from {extra} is not in the merged runs")
+            probs[key(row)].append(row["probs"])
+            added.add(key(row))
+        print(f"Added the thoughts of {extra} to {len(added)} questions")
+
+    tiebreak = {}
+    if args.tiebreak:
+        for line in (args.tiebreak / "thoughts.jsonl").open(encoding="utf-8"):
+            row = json.loads(line)
+            if key(row) not in probs:
+                raise SystemExit(f"error: {key(row)} from {args.tiebreak} is not in the merged runs")
+            tiebreak.setdefault(key(row), []).append(row["probs"])
+
+    first = [json.loads(line) for line in (args.run_dirs[0] / "predictions.jsonl").open(encoding="utf-8")]
+    if len(first) != len(probs):
+        raise SystemExit(f"error: {len(first)} predictions but {len(probs)} questions with thoughts")
+    predictions = []
+    broken = 0
+    for row in first:
+        group = probs[key(row)]
+        mean = [sum(p) / len(group) for p in zip(*group)]
+        pick = max(range(len(mean)), key=mean.__getitem__)
+        votes = [max(range(len(p)), key=p.__getitem__) for p in group]
+        top = max(votes.count(c) for c in set(votes))
+        tied = [c for c in sorted(set(votes)) if votes.count(c) == top]
+        if len(tied) > 1 and key(row) in tiebreak:
+            extra = tiebreak[key(row)]
+            judge = [sum(p) / len(extra) for p in zip(*extra)]
+            pick = max(tied, key=judge.__getitem__)
+            broken += 1
+        predictions.append({**row, "answer": LABELS[pick]})
+    if args.tiebreak:
+        print(f"Broke {broken} tied votes with {args.tiebreak}")
+
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+    jsonl_path = args.out_dir / "predictions.jsonl"
+    jsonl_path.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in predictions), encoding="utf-8")
+    with zipfile.ZipFile(args.out_dir / "predictions.zip", "w", zipfile.ZIP_DEFLATED) as zf:
+        zf.write(jsonl_path, "predictions.jsonl")
+    for run_dir in args.run_dirs:
+        before = {key(r): r["answer"] for r in map(json.loads, (run_dir / "predictions.jsonl").open(encoding="utf-8"))}
+        changed = sum(before[key(r)] != r["answer"] for r in predictions)
+        print(f"{changed} answers differ from {run_dir}")
+    print(f"Wrote {len(predictions)} predictions to {args.out_dir / 'predictions.zip'}")
+
+
+if __name__ == "__main__":
+    main()

model_manifest.json

@@ -0,0 +1,35 @@
+{
+  "system": "Gemma 4 E4B instruction-tuned model (text-only) in thinking mode with a LoRA adapter trained by STaR on its own thoughts; each question is answered in all 4 choice orders with a thought of up to 512 tokens, easy questions once more with a thought of up to 2,048 tokens as a fifth vote, medium questions also in the 4 reversed choice orders, hard questions once more with a thought of up to 2,048 tokens that replaces their vote, and the answer is the highest average probability among the option numbers after the thought",
+  "components": [
+    {
+      "role": "generator",
+      "model": "google/gemma-4-E4B-it",
+      "url": "https://huggingface.co/google/gemma-4-E4B-it",
+      "license": "Apache-2.0",
+      "released": "2026-03-02",
+      "loaded_parts": "language model only (vision and audio encoders are not loaded)",
+      "variant": "instruction-tuned (-it), thinking mode on with a greedy thought of up to 512 tokens on every question in 4 choice orders (predict.py --rotations 0 2 and --rotations 1 3) and on medium questions in the 4 reversed orders (predict.py --level medium --reverse), and of up to 2,048 tokens on easy and hard questions in the given order (predict.py --level easy|hard --think 2048)",
+      "parameters_loaded": 7463013376,
+      "parameters_full_checkpoint": 7996156490,
+      "precision": "float16 (run.sh sets it on every GPU; predict.py alone picks bfloat16 on GPUs with compute capability 8 or more)",
+      "fine_tuned": false
+    },
+    {
+      "role": "adapter",
+      "model": "LoRA adapter made by train_lora.py",
+      "method": "LoRA r=16, alpha 32, dropout 0.05 on every linear layer of the 42 decoder layers (342 modules), merged into the base weights at load",
+      "training": "STaR round 3: the model without any adapter thinks about each of 8,894 verified training questions; star.py --no_hint_retry keeps the thoughts that finish within 512 tokens and give the verified answer on the first try (4,073); a new adapter is trained from the base model on them, 2 epochs. Nothing from the round 1 or round 2 adapters is used",
+      "parameters": 38666240,
+      "optional": true
+    }
+  ],
+  "total_inference_parameters": {
+    "without_adapter": 7463013376,
+    "with_adapter": 7501679616
+  },
+  "parameter_limit": 8000000000,
+  "retrieval": "none",
+  "internet_at_inference": false,
+  "training_data": "data/train/verified-mcq.jsonl: 10,000 of the 18,398 Sinhala multiple-choice questions in data/train/mcq.jsonl (generated from Sri Lankan school textbooks, NIE teacher guides and resource books, A/L learning outcome notes and Wikibooks), with answer keys checked by a stronger model. review_overlap.py removes the 1,106 that match a dev or test question by wording or meaning, and the adapter trains only on the first-try thoughts of the model without an adapter that star.py keeps for the other 8,894. The dev and test sets are not used for training.",
+  "closed_api_usage": "Anthropic claude-haiku-5-5 generated the training questions in data/train/mcq.jsonl, and claude-sonnet-5-5 and claude-opus-5-5 checked the answer keys in data/train/verified-mcq.jsonl. They are used for training data only. No closed API is used at inference."
+}

predict.py

@@ -0,0 +1,335 @@
+import argparse
+import json
+import zipfile
+from pathlib import Path
+
+import torch
+from tqdm import tqdm
+from transformers import AutoTokenizer, Gemma4ForCausalLM
+
+MAX_PARAMS = 8_000_000_000
+MAX_BATCH_TOKENS = 36864
+LABELS = ["1", "2", "3", "4"]
+ANSWER_CUE = "පිළිතුර:"
+THOUGHT_OPEN = "<|channel>thought\n"
+THOUGHT_CLOSE = "\n<channel|>"
+STOP_TOKENS = ["<channel|>", "<eos>", "<turn|>"]
+
+
+def normalise(question):
+    metadata = question.get("metadata", {})
+    question = {
+        **question,
+        "subject_original": metadata.get("subject_original", question.get("subject_original")),
+        "level": metadata.get("difficulty", question.get("level")),
+    }
+    answer = question.get("answer")
+    if answer is not None and not (isinstance(answer, int) and 1 <= answer <= len(question["choices"])):
+        texts = [choice.strip() for choice in question["choices"]]
+        question["answer"] = texts.index(str(answer).strip()) + 1
+    return question
+
+
+def load_questions(data_dir):
+    questions = []
+    for path in sorted(Path(data_dir).rglob("*.json")):
+        if any(part.startswith(".") for part in path.relative_to(data_dir).parts):
+            continue
+        questions.extend(normalise(q) for q in json.loads(path.read_text(encoding="utf-8")))
+    return questions
+
+
+def user_prompt(question):
+    lines = [
+        "පහත බහුවරණ ප්‍රශ්නයට නිවැරදි පිළිතුරේ අංකය පමණක් ලියන්න.",
+        "",
+        f"විෂයය: {question['subject_original']}",
+        f"ප්‍රශ්නය: {question['question'].strip()}",
+        *(f"{i}. {choice}" for i, choice in enumerate(question["choices"], 1)),
+        ANSWER_CUE,
+    ]
+    return "\n".join(lines)
+
+
+def render(tokenizer, prompt, think):
+    if tokenizer.chat_template is None:
+        return tokenizer.bos_token + prompt + " "
+    messages = [{"role": "user", "content": prompt}]
+    if think:
+        text = tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True, enable_thinking=True
+        )
+        return text + THOUGHT_OPEN
+    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+
+def thinking_text(tokenizer, prompt, thought):
+    return render(tokenizer, prompt, True) + thought + THOUGHT_CLOSE + ANSWER_CUE
+
+
+DTYPES = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}
+
+
+def placement(dtype_name="auto"):
+    if not torch.cuda.is_available():
+        return torch.float32, "cpu"
+    if dtype_name != "auto":
+        return DTYPES[dtype_name], "auto"
+    return (torch.float16 if torch.cuda.get_device_capability()[0] < 8 else torch.bfloat16), "auto"
+
+
+def load_model(name, adapter=None, revision=None, dtype_name="auto"):
+    tokenizer = AutoTokenizer.from_pretrained(name, revision=revision, padding_side="left")
+    dtype, device_map = placement(dtype_name)
+    model, info = Gemma4ForCausalLM.from_pretrained(
+        name,
+        revision=revision,
+        dtype=dtype,
+        device_map=device_map,
+        key_mapping={r"^model\.language_model\.": "model."},
+        output_loading_info=True,
+    )
+    assert not info["missing_keys"], f"weights missing from checkpoint: {sorted(info['missing_keys'])[:5]}"
+    n_adapter = 0
+    if adapter is not None:
+        from peft import PeftModel
+
+        model = PeftModel.from_pretrained(model, adapter)
+        n_adapter = sum(p.numel() for name, p in model.named_parameters() if "lora_" in name)
+        assert n_adapter > 0, f"adapter {adapter} has no LoRA parameters"
+        model = model.merge_and_unload()
+    model.eval()
+    n_params = sum(p.numel() for p in model.parameters()) + n_adapter
+    print(
+        f"Loaded {name} (text-only) on {model.device} in {model.dtype}: {n_params:,} parameters"
+        f" ({n_adapter:,} adapter, merged)"
+    )
+    assert n_params <= MAX_PARAMS, f"{n_params:,} parameters exceeds the 8B limit"
+    return tokenizer, model
+
+
+def label_ids(tokenizer):
+    ids = [tokenizer.encode(label, add_special_tokens=False) for label in LABELS]
+    assert all(len(i) == 1 for i in ids), f"answer labels are not single tokens: {ids}"
+    return [i[0] for i in ids]
+
+
+def batches(lengths, budget, batch_size):
+    batch = []
+    for i in sorted(range(len(lengths)), key=lambda i: -lengths[i]):
+        if batch and (
+            len(batch) == batch_size or (len(batch) + 1) * (lengths[batch[0]] + budget) > MAX_BATCH_TOKENS
+        ):
+            yield batch
+            batch = []
+        batch.append(i)
+    if batch:
+        yield batch
+
+
+def encode(tokenizer, texts, device):
+    return tokenizer(texts, return_tensors="pt", padding=True, add_special_tokens=False).to(device)
+
+
+def score(model, tokenizer, texts, n_choices, batch_size):
+    labels = label_ids(tokenizer)
+    lengths = [len(tokenizer.encode(t, add_special_tokens=False)) for t in texts]
+    probs = [None] * len(texts)
+    for batch in tqdm(list(batches(lengths, 0, batch_size)), desc="score"):
+        inputs = encode(tokenizer, [texts[i] for i in batch], model.device)
+        position_ids = (inputs["attention_mask"].cumsum(-1) - 1).clamp(min=0)
+        with torch.inference_mode():
+            logits = model(**inputs, position_ids=position_ids, logits_to_keep=1).logits[:, -1, labels].float()
+        if not torch.isfinite(logits).all():
+            raise RuntimeError("non-finite answer logits")
+        counts = torch.tensor([n_choices[i] for i in batch], device=logits.device)
+        valid = torch.arange(len(labels), device=logits.device)[None, :] < counts[:, None]
+        for i, row in zip(batch, logits.masked_fill(~valid, float("-inf")).softmax(-1).tolist()):
+            probs[i] = row[: n_choices[i]]
+    return probs
+
+
+def think(model, tokenizer, contexts, budget, batch_size):
+    stops = tokenizer.convert_tokens_to_ids(STOP_TOKENS)
+    lengths = [len(tokenizer.encode(t, add_special_tokens=False)) for t in contexts]
+    thoughts = [None] * len(contexts)
+    for batch in tqdm(list(batches(lengths, budget, batch_size)), desc="think"):
+        inputs = encode(tokenizer, [contexts[i] for i in batch], model.device)
+        with torch.inference_mode():
+            out = model.generate(
+                **inputs,
+                max_new_tokens=budget,
+                do_sample=False,
+                eos_token_id=stops,
+                pad_token_id=tokenizer.pad_token_id,
+            )
+        for i, row in zip(batch, out[:, inputs["input_ids"].shape[1]:].tolist()):
+            end = next((k for k, token in enumerate(row) if token in stops), len(row))
+            thoughts[i] = {
+                "thought": tokenizer.decode(row[:end], skip_special_tokens=True).strip(),
+                "finished": end < len(row) and row[end] == stops[0],
+                "thought_tokens": end,
+            }
+    return thoughts
+
+
+def argmax(values):
+    return max(range(len(values)), key=values.__getitem__)
+
+
+def answer(model, tokenizer, prompts, n_choices, budget, batch_size, score_batch_size=None):
+    if budget == 0:
+        texts = [render(tokenizer, p, False) for p in prompts]
+        thoughts = [{} for _ in prompts]
+    else:
+        contexts = [render(tokenizer, p, True) for p in prompts]
+        thoughts = think(model, tokenizer, contexts, budget, batch_size)
+        texts = [thinking_text(tokenizer, p, t["thought"]) for p, t in zip(prompts, thoughts)]
+    probs = score(model, tokenizer, texts, n_choices, score_batch_size or batch_size)
+    return [{**t, "probs": p, "pred": argmax(p) + 1} for t, p in zip(thoughts, probs)]
+
+
+def order(n, rotation, reverse):
+    positions = list(range(n))[::-1] if reverse else list(range(n))
+    k = rotation % n
+    return positions[k:] + positions[:k]
+
+
+def rotate(question, rotation, reverse=False):
+    return {**question, "choices": [question["choices"][i] for i in order(len(question["choices"]), rotation, reverse)]}
+
+
+def unrotate(probs, rotation, reverse=False):
+    original = [0.0] * len(probs)
+    for i, p in zip(order(len(probs), rotation, reverse), probs):
+        original[i] = p
+    return original
+
+
+def vote(model, tokenizer, questions, rotations, budget, batch_size, score_batch_size=None, reverse=False):
+    prompts = [user_prompt(rotate(q, r, reverse)) for q in questions for r in rotations]
+    n_choices = [len(q["choices"]) for q in questions for _ in rotations]
+    runs = answer(model, tokenizer, prompts, n_choices, budget, batch_size, score_batch_size)
+    results = []
+    for start in range(0, len(runs), len(rotations)):
+        group = [
+            {**run, "rotation": r, "reverse": reverse, "probs": unrotate(run["probs"], r, reverse)}
+            for run, r in zip(runs[start : start + len(rotations)], rotations)
+        ]
+        group = [{**run, "pred": argmax(run["probs"]) + 1} for run in group]
+        probs = [sum(p) / len(group) for p in zip(*(run["probs"] for run in group))]
+        results.append({"pred": argmax(probs) + 1, "probs": probs, "runs": group})
+    return results
+
+
+def metrics(questions, results):
+    levels = {}
+    for question, result in zip(questions, results):
+        counts = levels.setdefault(question["level"], [0, 0])
+        counts[0] += result["pred"] == question["answer"]
+        counts[1] += 1
+    correct = sum(c for c, _ in levels.values())
+    return {
+        "accuracy": correct / len(questions),
+        "correct": correct,
+        "total": len(questions),
+        "levels": {k: {"accuracy": c / t, "correct": c, "total": t} for k, (c, t) in sorted(levels.items())},
+    }
+
+
+def write_jsonl(path, rows):
+    path.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows), encoding="utf-8")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", default="google/gemma-4-E4B-it")
+    parser.add_argument("--revision", help="model commit to load; the latest one when not given")
+    parser.add_argument("--dtype", choices=["auto", *DTYPES], default="auto")
+    parser.add_argument("--data_dir", default="data/test")
+    parser.add_argument("--out_dir", default="outputs")
+    parser.add_argument("--think", type=int, default=512)
+    parser.add_argument("--adapter")
+    parser.add_argument("--limit", type=int)
+    parser.add_argument("--level", help="answer only the questions of this level, for example hard")
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument(
+        "--score_batch_size", type=int,
+        help="batch size for scoring the answer after the thought; --batch_size when not given, lower it for long thoughts",
+    )
+    parser.add_argument(
+        "--rotations", type=int, nargs="+", default=[0],
+        help="answer once per rotation of the choices and average the answer probabilities; 0 2 votes over two orders",
+    )
+    parser.add_argument(
+        "--reverse", action="store_true",
+        help="reverse the choices before rotating them, which gives 4 more orders of 4 choices",
+    )
+    args = parser.parse_args()
+
+    questions = load_questions(args.data_dir)
+    if args.level:
+        questions = [q for q in questions if q["level"] == args.level]
+    questions = questions[: args.limit]
+    print(f"Loaded {len(questions)} questions from {args.data_dir}")
+    out_dir = Path(args.out_dir)
+    if not questions:
+        out_dir.mkdir(parents=True, exist_ok=True)
+        write_jsonl(out_dir / "predictions.jsonl", [])
+        write_jsonl(out_dir / "thoughts.jsonl", [])
+        return
+
+    tokenizer, model = load_model(args.model, args.adapter, args.revision, args.dtype)
+    if args.think and tokenizer.chat_template is None:
+        parser.error(f"{args.model} has no chat template, so it has no thinking mode; pass --think 0")
+
+    results = vote(
+        model, tokenizer, questions, args.rotations, args.think, args.batch_size, args.score_batch_size, args.reverse
+    )
+    predictions = [
+        {
+            "q_no": question["q_no"],
+            "subject": question["subject"],
+            "category": question["category"],
+            "level": question["level"],
+            "answer": LABELS[result["pred"] - 1],
+        }
+        for question, result in zip(questions, results)
+    ]
+
+    out_dir.mkdir(parents=True, exist_ok=True)
+    jsonl_path = out_dir / "predictions.jsonl"
+    write_jsonl(jsonl_path, predictions)
+    with zipfile.ZipFile(out_dir / "predictions.zip", "w", zipfile.ZIP_DEFLATED) as zf:
+        zf.write(jsonl_path, "predictions.jsonl")
+    print(f"Wrote {len(predictions)} predictions to {out_dir / 'predictions.zip'}")
+
+    runs = [(p, run) for p, r in zip(predictions, results) for run in r["runs"]]
+    thoughts = [
+        {"q_no": p["q_no"], "subject": p["subject"], "level": p["level"]}
+        | {k: run[k] for k in ("rotation", "reverse", "pred", "probs", "finished", "thought_tokens", "thought") if k in run}
+        for p, run in runs
+    ]
+    write_jsonl(out_dir / "thoughts.jsonl", thoughts)
+    if args.think:
+        finished = sum(run["finished"] for _, run in runs)
+        print(f"Thoughts finished within {args.think} tokens: {finished}/{len(runs)}")
+
+    if questions and all(q.get("answer") is not None for q in questions):
+        report = metrics(questions, results) | {"think": args.think, "adapter": args.adapter}
+        if len(args.rotations) > 1:
+            report["rotations"] = {
+                str(r): metrics(questions, [res["runs"][i] for res in results])["accuracy"]
+                for i, r in enumerate(args.rotations)
+            }
+        if args.think:
+            report["finished"] = finished
+        (out_dir / "metrics.json").write_text(json.dumps(report, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
+        levels = ", ".join(f"{k} {v['accuracy']:.4f}" for k, v in report["levels"].items())
+        print(f"Accuracy {report['accuracy']:.4f} ({report['correct']}/{report['total']}); {levels}")
+        if "rotations" in report:
+            print("Accuracy of each rotation alone:", report["rotations"])
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,9 @@
+--extra-index-url https://download.pytorch.org/whl/cu128
+torch==2.11.0+cu128
+transformers==5.20.1
+accelerate==1.15.0
+huggingface_hub==2.2.0
+tokenizers==0.23.3
+safetensors==0.8.0
+peft==0.21.2
+tqdm==4.67.3

training/README.md

@@ -0,0 +1,96 @@
+# Training: STaR round 3 LoRA adapter
+
+This folder holds the code, settings and data that made the adapter in [SinhalaMMLU-SharedTask/team-sinhalanuwana-gemma4-e4b-star3-lora](https://huggingface.co/SinhalaMMLU-SharedTask/team-sinhalanuwana-gemma4-e4b-star3-lora). The adapter trains from the base model `google/gemma-4-E4B-it` (revision `ee0ef6023621cff504d758262d4e04895a5af4a2`). It uses no earlier adapter.
+
+## Files
+
+| File | Purpose |
+| --- | --- |
+| `review_overlap.py` | Compares each training question with every dev and test question, by wording (character 6-gram Jaccard and containment, same choices) and by meaning (LaBSE cosine). Lists the matches for removal. |
+| `prepare_star_data.py` | Removes the listed questions and splits the rest into 2 shards with seed 0 |
+| `star.py` | Lets the model think about each training question and keeps the thoughts that finish and give the correct answer |
+| `train_lora.py` | Trains the LoRA adapter on the kept thoughts |
+| `kaggle/star3-gen-0`, `kaggle/star3-gen-1` | The Kaggle notebooks that ran `star.py` on the two shards (2 x T4 each, about 4.2 hours each) |
+| `kaggle/star3-train` | The Kaggle notebook that joined their kept thoughts and ran `train_lora.py` (2 x T4, 4.3 hours) |
+| `train_log.json` | Settings, parameter counts and the loss of every step of the training run |
+| `data/shard-0.jsonl`, `data/shard-1.jsonl` | The 8,894 training questions (4,447 each), the input of `star.py` |
+| `data/excluded_ids.txt` | The ids of the 1,106 questions that `review_overlap.py` removed |
+| `data/overlap-summary.json` | The counts of the overlap review, by rule |
+| `data/star3-thoughts.jsonl` | The 4,073 training examples: prompt, thought and answer. This is the exact file the adapter trained on. |
+
+## Training data
+
+Anthropic claude-haiku-5-5 generated Sinhala multiple-choice questions from Sri Lankan school textbooks, NIE teacher guides and resource books, A/L learning outcome notes and Wikibooks. Anthropic claude-sonnet-5-5 (9,228 questions) and claude-opus-5-5 (772 questions) checked the answer keys of 10,000 of them. The shared task allows closed models for training data.
+
+Each question in `data/shard-*.jsonl` has its `id`, `subject`, `level`, `grade`, `source`, `question`, the 4 `choices`, the checked `answer` (1 to 4) and an `explanation`.
+
+We removed the questions that resemble an organiser question before any training. `review_overlap.py` compared the 10,000 questions with the 1,851 dev and 2,238 test questions, and with an earlier copy of the test set (2,146 questions), and removed 1,106 by these rules. A question can match more than one rule, so the counts add up to more than 1,106:
+
+| Rule | Removed |
+| --- | --- |
+| Same question text | 0 |
+| Question Jaccard of 0.5 or more | 17 |
+| Question containment of 0.8 or more | 48 |
+| 3 or more same choices and containment of 0.3 or more | 11 |
+| All 4 choices the same | 43 |
+| LaBSE cosine of 0.70 or more | 1,033 |
+
+The removed questions are not in this repository. The dev and test questions were not used for training in any form.
+
+## Steps
+
+Run the commands from the root of this repository. The training scripts import `predict.py` from the root, so set `PYTHONPATH=.`.
+
+1. Remove the questions that resemble a dev or test question, and make the shards. These steps need the full 10,000-question file and the organiser dev and test sets, which are not in this repository. Their output is in `data/`.
+
+   ```bash
+   PYTHONPATH=. python training/review_overlap.py --questions verified-mcq.jsonl \
+     --eval dev=DEV_DIR test=TEST_DIR --out_dir overlap
+   PYTHONPATH=. python training/prepare_star_data.py --questions verified-mcq.jsonl \
+     --exclude overlap/excluded_ids.txt --out_dir training/data --shards 2 --seed 0
+   ```
+
+2. Let the base model think about each question, with no adapter and no hinted retry. One run for each shard, a thought of up to 512 tokens, greedy, batches of 32:
+
+   ```bash
+   PYTHONPATH=. python training/star.py --train training/data/shard-0.jsonl --no_hint_retry \
+     --revision ee0ef6023621cff504d758262d4e04895a5af4a2 --out_dir star-0
+   PYTHONPATH=. python training/star.py --train training/data/shard-1.jsonl --no_hint_retry \
+     --revision ee0ef6023621cff504d758262d4e04895a5af4a2 --out_dir star-1
+   ```
+
+   `star.py` keeps a thought only if it ends within 512 tokens and the answer after it is the checked answer. It kept 4,073 of the 8,894.
+
+3. Join the kept thoughts of the two runs:
+
+   ```bash
+   cat star-0/star.jsonl star-1/star.jsonl > star3.jsonl
+   ```
+
+   The result is byte for byte `data/star3-thoughts.jsonl`. `train_lora.py` shuffles the examples with seed 0.
+
+4. Train the adapter:
+
+   ```bash
+   PYTHONPATH=. python training/train_lora.py --data star3.jsonl \
+     --revision ee0ef6023621cff504d758262d4e04895a5af4a2 --out_dir adapter
+   ```
+
+   To train on the published examples directly, pass `--data training/data/star3-thoughts.jsonl`.
+
+## Settings
+
+| Setting | Value |
+| --- | --- |
+| LoRA | Rank 16, alpha 32, dropout 0.05, on all 342 linear layers of the 42 decoder layers (`q_proj`, `k_proj`, `v_proj`, `o_proj`, `gate_proj`, `up_proj`, `down_proj`, `per_layer_input_gate`, `per_layer_projection`) |
+| Trainable parameters | 38,666,240 |
+| Loss | Cross-entropy on the thought, the answer cue and the answer number, not on the prompt |
+| Optimiser | AdamW, learning rate 1e-4, cosine schedule with 3% warmup |
+| Batch | 4 examples, 4 accumulation steps, so 16 examples per step |
+| Epochs | 2 (510 steps) |
+| Maximum length | 1,024 tokens (no example was dropped) |
+| Precision | float16 with gradient scaling (T4 GPUs have no bfloat16) |
+| Seed | 0 |
+| Hardware and time | 2 x Tesla T4 on Kaggle, 15,552 seconds (4.3 hours) |
+
+The Kaggle notebooks used the same scripts from a snapshot of our code with all scripts in one folder, and the shards from a private Kaggle dataset. Their `run.py` installs `requirements.txt`, removes `torchao`, downloads the model and runs the command above.

training/data/excluded_ids.txt

@@ -0,0 +1,1106 @@
+mcq_915451748_058_4
+mcq_915451748_007_0
+mcq_902667986_007_2
+mcq_915451748_051_0
+mcq_902667986_006_0
+mcq_902667986_009_2
+mcq_915451748_007_4
+mcq_915451748_052_1
+mcq_915451748_044_3
+mcq_902667986_014_4
+mcq_902667986_028_1
+mcq_915451748_057_1
+mcq_902667986_025_1
+mcq_915451748_031_4
+mcq_915451748_053_0
+mcq_915451748_053_1
+mcq_902667986_024_1
+mcq_902667986_008_4
+mcq_915451748_026_2
+mcq_902667986_013_0
+mcq_902667986_012_3
+mcq_915451748_050_0
+mcq_902667986_007_1
+mcq_902667986_013_1
+mcq_915451748_017_2
+mcq_902667986_028_4
+mcq_902667986_016_4
+mcq_902667986_014_2
+mcq_902667986_027_2
+mcq_915451748_050_1
+mcq_902667986_016_0
+mcq_23255_002_4
+mcq_23243_000_0
+mcq_22040_001_0
+mcq_22026_003_0
+mcq_23241_000_0
+mcq_22056_001_0
+mcq_19860_000_0
+mcq_23267_000_2
+mcq_22063_000_0
+mcq_23237_001_0
+mcq_23254_003_0
+mcq_23237_001_4
+mcq_22053_001_0
+mcq_23237_003_4
+mcq_23239_000_0
+mcq_22047_003_2
+mcq_22049_000_2
+mcq_23248_001_2
+mcq_23241_000_3
+mcq_23245_001_2
+mcq_22034_002_2
+mcq_22066_000_2
+mcq_19838_004_0
+mcq_19832_001_3
+mcq_23240_002_1
+mcq_23270_000_1
+mcq_22037_000_3
+mcq_23045_000_2
+mcq_23032_000_0
+mcq_20578_001_0
+mcq_20591_000_0
+mcq_20576_001_2
+mcq_20580_001_0
+mcq_23034_001_2
+mcq_23052_000_0
+mcq_23042_000_0
+mcq_23420_000_0
+mcq_23047_001_2
+mcq_23432_001_2
+mcq_20579_000_0
+mcq_23044_000_0
+mcq_20601_000_2
+mcq_23062_000_2
+mcq_23054_001_2
+mcq_20581_001_0
+mcq_23416_001_0
+mcq_23058_001_0
+mcq_20595_001_0
+mcq_23031_000_0
+mcq_20580_000_0
+mcq_23047_001_0
+mcq_23416_000_2
+mcq_20583_000_0
+mcq_23040_001_2
+mcq_23436_000_2
+mcq_23333_000_0
+mcq_23320_001_1
+mcq_23333_001_0
+mcq_23321_001_0
+mcq_23317_002_0
+mcq_19983_000_0
+mcq_19982_002_2
+mcq_22758_001_2
+mcq_22745_000_0
+mcq_22748_002_0
+mcq_22757_001_0
+mcq_19976_000_1
+mcq_22749_000_0
+mcq_19972_002_0
+mcq_22741_003_0
+mcq_22746_001_0
+mcq_22755_000_0
+mcq_23317_000_0
+mcq_19959_001_0
+mcq_22745_000_2
+mcq_23332_000_0
+mcq_22746_000_1
+mcq_23324_000_1
+mcq_23318_000_0
+mcq_19966_001_0
+mcq_19959_002_0
+mcq_22752_002_2
+mcq_21720_003_1
+mcq_19423_018_0
+mcq_19423_000_0
+mcq_22921_014_1
+mcq_22924_003_1
+mcq_22921_014_2
+mcq_21718_002_2
+mcq_22921_012_1
+mcq_21718_003_1
+mcq_21718_001_1
+mcq_21723_003_2
+mcq_21718_002_1
+mcq_22924_008_0
+mcq_19423_020_1
+mcq_21723_003_0
+mcq_22926_012_0
+mcq_19423_014_2
+mcq_22924_009_0
+mcq_19423_014_0
+mcq_21718_006_1
+mcq_19432_001_2
+mcq_19423_007_0
+mcq_22926_007_0
+mcq_22926_015_2
+mcq_19432_000_1
+mcq_21714_003_0
+mcq_21718_009_1
+mcq_21718_006_0
+mcq_19423_002_0
+mcq_22926_004_2
+mcq_19423_002_1
+mcq_22921_011_0
+mcq_22921_006_0
+mcq_19423_008_1
+mcq_22924_001_0
+mcq_21718_001_2
+mcq_21718_003_2
+mcq_21718_001_0
+mcq_21718_000_0
+mcq_19423_020_0
+mcq_21712_016_0
+mcq_22924_006_0
+mcq_903137782_011_2
+mcq_906302875_010_0
+mcq_903137782_019_3
+mcq_903137782_011_0
+mcq_906302875_029_0
+mcq_906302875_009_0
+mcq_906302875_018_6
+mcq_906302875_020_0
+mcq_906302875_010_6
+mcq_903137782_012_0
+mcq_903137782_012_1
+mcq_903137782_020_0
+mcq_906302875_020_2
+mcq_906302875_022_5
+mcq_903137782_011_7
+mcq_906302875_025_3
+mcq_906302875_009_2
+mcq_903137782_009_0
+mcq_903137782_019_2
+mcq_903137782_021_7
+mcq_903137782_008_7
+mcq_903137782_020_1
+mcq_906302875_015_0
+mcq_903137782_012_6
+mcq_903137782_010_3
+mcq_906302875_033_1
+mcq_906302875_027_0
+mcq_906302875_033_0
+mcq_903137782_025_6
+mcq_903137782_011_6
+mcq_903137782_018_2
+mcq_906302875_011_3
+mcq_903137782_009_1
+mcq_906302875_011_1
+mcq_903137782_010_6
+mcq_903137782_011_3
+mcq_903137782_021_2
+mcq_903137782_024_0
+mcq_903137782_010_0
+mcq_903137782_025_2
+mcq_906302875_010_3
+mcq_910693558_025_2
+mcq_910693558_034_0
+mcq_903779140_044_2
+mcq_903779140_009_0
+mcq_910693558_024_2
+mcq_903779140_016_0
+mcq_903779140_006_1
+mcq_903779140_019_0
+mcq_910693558_007_2
+mcq_903779140_009_4
+mcq_910693558_009_1
+mcq_904719733_010_4
+mcq_910693558_010_0
+mcq_903779140_015_5
+mcq_903779140_015_3
+mcq_910693558_023_5
+mcq_903779140_015_1
+mcq_903779140_009_1
+mcq_903779140_007_2
+mcq_903779140_007_0
+mcq_910693558_009_5
+mcq_910693558_014_0
+mcq_916520645_017_0
+mcq_912050769_018_0
+mcq_903568982_005_1
+mcq_916520645_037_4
+mcq_916520645_036_2
+mcq_916520645_020_0
+mcq_916520645_013_1
+mcq_903568982_017_1
+mcq_903568982_015_5
+mcq_916520645_023_0
+mcq_912050769_010_0
+mcq_903568982_042_4
+mcq_903568982_031_4
+mcq_912050769_009_3
+mcq_916520645_023_5
+mcq_903568982_015_0
+mcq_903568982_009_0
+mcq_903568982_011_1
+mcq_903568982_029_0
+mcq_903568982_044_0
+mcq_903568982_034_5
+mcq_903568982_009_1
+mcq_916520645_018_0
+mcq_916520645_046_0
+mcq_916520645_017_4
+mcq_916520645_029_0
+mcq_916520645_029_5
+mcq_903568982_008_4
+mcq_916520645_037_1
+mcq_903568982_031_0
+mcq_903568982_031_1
+mcq_903568982_023_0
+mcq_903568982_008_0
+mcq_916520645_009_1
+mcq_916520645_036_3
+mcq_916520645_033_1
+mcq_912050769_009_1
+mcq_903568982_024_1
+mcq_916520645_007_5
+mcq_916520645_011_4
+mcq_903568982_042_3
+mcq_916520645_041_1
+mcq_916520645_040_5
+mcq_903568982_009_3
+mcq_912050769_005_5
+mcq_916520645_020_1
+mcq_912050769_033_3
+mcq_916520645_013_2
+mcq_916520645_005_4
+mcq_912050769_019_1
+mcq_903568982_038_1
+mcq_912050769_007_2
+mcq_916520645_046_2
+mcq_916520645_011_0
+mcq_916520645_037_0
+mcq_916520645_024_5
+mcq_912050769_024_0
+mcq_916520645_008_0
+mcq_903568982_011_5
+mcq_903568982_016_5
+mcq_916520645_045_5
+mcq_916520645_037_5
+mcq_912050769_014_3
+mcq_903568982_023_1
+mcq_916520645_040_1
+mcq_916520645_036_1
+mcq_903568982_008_5
+mcq_19377_002_3
+mcq_19388_000_0
+mcq_21399_007_1
+mcq_19370_000_0
+mcq_21399_000_0
+mcq_22349_003_0
+mcq_21388_004_3
+mcq_19388_001_3
+mcq_21388_005_3
+mcq_21395_005_0
+mcq_22349_001_0
+mcq_19370_007_0
+mcq_21388_005_1
+mcq_21388_000_2
+mcq_19370_002_1
+mcq_19388_001_0
+mcq_21388_006_0
+mcq_19370_003_0
+mcq_21399_008_2
+mcq_19370_005_0
+mcq_22349_003_1
+mcq_19370_003_1
+mcq_22349_000_0
+mcq_22349_004_0
+mcq_21388_003_0
+mcq_21395_001_2
+mcq_21399_005_3
+mcq_19370_002_0
+mcq_21388_003_2
+mcq_21388_004_0
+mcq_21388_004_1
+mcq_21399_001_0
+mcq_19382_001_1
+mcq_19382_002_0
+mcq_24610_001_0
+mcq_24613_001_0
+mcq_24612_000_0
+mcq_23868_000_2
+mcq_21135_000_0
+mcq_23888_000_2
+mcq_23909_003_2
+mcq_23877_001_0
+mcq_24611_001_0
+mcq_21118_001_0
+mcq_23917_005_1
+mcq_23909_001_0
+mcq_24610_005_1
+mcq_23906_002_0
+mcq_23874_003_0
+mcq_23865_001_1
+mcq_23865_000_0
+mcq_21135_003_0
+mcq_23913_001_2
+mcq_24608_000_0
+mcq_23868_003_0
+mcq_23898_001_0
+mcq_23928_005_2
+mcq_24613_003_2
+mcq_21122_000_2
+mcq_24608_003_0
+mcq_23870_000_0
+mcq_24606_003_0
+mcq_23880_001_0
+mcq_21118_018_0
+mcq_23917_001_0
+mcq_23874_001_1
+mcq_21122_000_0
+mcq_24597_003_0
+mcq_24604_001_1
+mcq_23921_001_1
+mcq_21125_000_0
+mcq_21118_007_0
+mcq_21929_003_3
+mcq_21926_002_0
+mcq_21926_012_0
+mcq_19156_003_1
+mcq_21926_004_0
+mcq_19174_005_0
+mcq_21928_008_1
+mcq_21183_006_0
+mcq_21174_003_0
+mcq_19174_007_0
+mcq_21928_002_1
+mcq_21926_007_1
+mcq_19174_009_3
+mcq_21174_004_1
+mcq_21179_004_0
+mcq_21174_002_0
+mcq_19151_004_0
+mcq_19174_000_0
+mcq_21171_002_3
+mcq_21171_007_0
+mcq_19156_005_0
+mcq_19174_000_1
+mcq_21934_002_3
+mcq_21924_005_0
+mcq_21185_000_0
+mcq_21174_009_4
+mcq_19174_008_1
+mcq_21928_003_0
+mcq_19174_004_3
+mcq_19163_001_0
+mcq_21171_009_0
+mcq_19174_006_2
+mcq_19171_000_0
+mcq_19174_002_0
+mcq_19163_004_4
+mcq_21171_020_2
+mcq_21171_006_0
+mcq_19156_006_4
+mcq_17191_004_2
+mcq_17177_004_1
+mcq_17144_008_0
+mcq_18534_012_1
+mcq_17105_000_2
+mcq_17207_000_0
+mcq_18595_000_2
+mcq_17137_005_0
+mcq_17130_001_1
+mcq_19265_004_0
+mcq_18584_003_2
+mcq_19292_003_0
+mcq_17137_000_0
+mcq_18512_002_0
+mcq_18589_009_0
+mcq_17130_002_1
+mcq_18589_010_0
+mcq_17207_004_0
+mcq_19272_000_1
+mcq_18584_008_1
+mcq_18553_005_1
+mcq_17177_000_1
+mcq_19255_002_2
+mcq_19237_006_2
+mcq_18562_009_0
+mcq_19298_002_1
+mcq_18534_012_2
+mcq_19284_001_2
+mcq_19265_000_1
+mcq_19265_002_0
+mcq_18557_003_1
+mcq_19315_001_0
+mcq_19237_003_2
+mcq_17199_000_1
+mcq_18557_002_1
+mcq_19288_003_0
+mcq_18534_008_1
+mcq_19284_004_0
+mcq_18534_005_0
+mcq_18567_011_0
+mcq_19284_005_1
+mcq_19284_004_2
+mcq_17207_004_2
+mcq_18584_008_0
+mcq_18547_003_0
+mcq_18567_010_0
+mcq_18528_000_1
+mcq_19292_000_0
+mcq_19284_000_0
+mcq_17130_001_0
+mcq_19237_004_0
+mcq_18567_006_1
+mcq_19284_003_1
+mcq_18562_008_0
+mcq_18553_005_0
+mcq_17169_001_0
+mcq_19255_001_0
+mcq_17144_004_2
+mcq_18534_003_2
+mcq_19308_005_2
+mcq_19288_000_0
+mcq_18553_001_2
+mcq_18547_005_2
+mcq_19284_000_2
+mcq_17130_005_0
+mcq_19255_002_0
+mcq_18519_004_2
+mcq_19298_003_2
+mcq_19308_003_1
+mcq_18553_001_0
+mcq_19298_003_1
+mcq_19320_008_1
+mcq_17177_007_0
+mcq_18534_007_0
+mcq_19269_002_1
+mcq_18562_009_2
+mcq_19272_007_0
+mcq_19247_001_0
+mcq_19237_003_0
+mcq_19261_001_2
+mcq_17130_001_2
+mcq_18528_003_0
+mcq_18578_004_2
+mcq_19269_004_2
+mcq_18567_014_0
+mcq_19261_001_1
+mcq_18528_001_1
+mcq_19269_004_0
+mcq_19284_005_0
+mcq_19265_003_0
+mcq_18589_001_0
+mcq_18589_016_2
+mcq_19292_002_0
+mcq_18567_005_1
+mcq_17144_005_1
+mcq_19251_004_1
+mcq_17191_006_0
+mcq_19265_007_0
+mcq_19303_001_0
+mcq_19320_008_0
+mcq_18547_002_1
+mcq_18567_005_0
+mcq_18519_002_2
+mcq_18589_015_2
+mcq_18553_002_0
+mcq_19237_009_0
+mcq_19288_001_2
+mcq_18512_000_2
+mcq_17207_001_0
+mcq_19744_001_3
+mcq_18424_001_0
+mcq_18412_004_2
+mcq_18397_001_3
+mcq_19767_006_2
+mcq_19746_005_3
+mcq_19767_002_1
+mcq_18380_005_2
+mcq_18412_000_0
+mcq_18457_000_0
+mcq_19761_002_2
+mcq_19752_001_0
+mcq_18416_000_0
+mcq_18397_003_1
+mcq_19741_002_1
+mcq_20288_000_0
+mcq_20269_000_1
+mcq_20292_000_0
+mcq_23194_000_0
+mcq_23206_001_0
+mcq_38534_008_1
+mcq_38542_011_2
+mcq_28806_000_0
+mcq_38554_009_1
+mcq_28809_033_0
+mcq_28809_019_0
+mcq_38472_009_2
+mcq_38523_000_2
+mcq_28815_003_0
+mcq_28809_010_0
+mcq_28817_000_2
+mcq_40041_002_0
+mcq_25632_008_2
+mcq_40041_002_4
+mcq_40139_001_1
+mcq_40062_003_0
+mcq_40041_002_3
+mcq_25681_001_0
+mcq_40062_004_0
+mcq_25712_001_2
+mcq_25704_002_0
+mcq_40116_003_0
+mcq_40020_001_4
+mcq_40062_001_4
+mcq_25695_001_4
+mcq_25719_003_2
+mcq_40155_001_0
+mcq_25726_002_1
+mcq_40099_002_4
+mcq_40139_000_2
+mcq_40103_002_3
+mcq_25732_002_3
+mcq_40093_001_2
+mcq_40168_002_2
+mcq_40093_003_4
+mcq_25643_002_0
+mcq_40116_004_3
+mcq_40168_001_0
+mcq_40030_003_2
+mcq_40146_001_4
+mcq_25638_009_0
+mcq_25632_003_3
+mcq_40020_003_4
+mcq_40093_000_1
+mcq_40139_003_4
+mcq_25643_000_3
+mcq_40111_002_1
+mcq_40076_003_2
+mcq_40057_002_2
+mcq_40041_003_1
+mcq_40093_004_1
+mcq_25712_001_4
+mcq_25697_002_0
+mcq_40111_005_2
+mcq_40076_000_4
+mcq_25686_000_4
+mcq_25625_002_2
+mcq_40155_004_0
+mcq_40146_000_1
+mcq_25700_001_0
+mcq_40089_004_0
+mcq_40036_003_3
+mcq_40036_001_2
+mcq_25638_005_4
+mcq_40146_000_2
+mcq_40076_002_3
+mcq_25732_001_2
+mcq_25655_002_2
+mcq_40139_004_0
+mcq_25676_001_2
+mcq_40062_003_1
+mcq_40111_004_0
+mcq_25686_000_2
+mcq_40076_002_1
+mcq_25663_001_2
+mcq_25719_001_3
+mcq_25681_002_2
+mcq_40076_003_0
+mcq_40111_005_0
+mcq_25861_000_0
+mcq_21781_002_0
+mcq_21778_001_0
+mcq_25879_000_0
+mcq_27135_001_0
+mcq_25892_001_2
+mcq_27149_003_2
+mcq_21779_000_2
+mcq_27144_000_2
+mcq_21778_003_2
+mcq_25873_000_0
+mcq_25876_000_1
+mcq_22977_000_0
+mcq_25833_001_0
+mcq_22976_000_0
+mcq_22968_021_0
+mcq_22986_002_2
+mcq_22985_001_2
+mcq_25799_000_0
+mcq_22968_045_2
+mcq_22968_041_0
+mcq_22972_002_2
+mcq_25833_000_0
+mcq_22982_004_1
+mcq_22986_001_2
+mcq_25816_001_2
+mcq_22969_002_2
+mcq_26405_013_1
+mcq_28651_005_2
+mcq_28651_004_0
+mcq_28583_011_1
+mcq_26377_001_1
+mcq_28583_009_1
+mcq_28651_003_1
+mcq_28651_000_1
+mcq_20629_014_1
+mcq_26405_009_0
+mcq_28696_001_1
+mcq_28583_014_0
+mcq_28583_014_1
+mcq_26377_006_2
+mcq_28696_002_1
+mcq_26365_008_1
+mcq_28583_006_1
+mcq_26390_024_0
+mcq_20633_005_0
+mcq_20632_001_0
+mcq_28651_004_1
+mcq_26425_010_1
+mcq_20625_013_0
+mcq_20629_016_2
+mcq_26405_017_1
+mcq_28651_001_0
+mcq_28665_009_0
+mcq_28583_013_2
+mcq_28583_008_1
+mcq_20633_003_1
+mcq_26377_000_2
+mcq_28623_009_1
+mcq_26390_023_0
+mcq_28665_013_1
+mcq_26405_013_0
+mcq_26377_000_1
+mcq_28583_013_1
+mcq_20633_000_1
+mcq_28583_015_1
+mcq_28665_014_0
+mcq_28665_014_2
+mcq_26425_008_0
+mcq_26365_007_0
+mcq_28651_005_0
+mcq_28583_019_0
+mcq_26405_015_0
+mcq_20632_004_0
+mcq_26405_018_0
+mcq_28665_014_1
+mcq_20634_006_1
+mcq_28651_004_2
+mcq_20625_000_0
+mcq_26390_022_2
+mcq_20632_004_2
+mcq_26377_007_0
+mcq_20633_002_0
+mcq_20633_002_2
+mcq_28665_003_0
+mcq_26390_003_0
+mcq_26390_021_0
+mcq_20632_001_2
+mcq_20633_011_0
+mcq_20632_005_0
+mcq_26405_008_2
+mcq_28623_006_0
+mcq_20633_001_0
+mcq_26390_022_0
+mcq_28623_009_0
+mcq_26390_012_2
+mcq_26405_014_1
+mcq_28583_018_1
+mcq_28651_000_0
+mcq_28583_011_0
+mcq_20633_004_0
+mcq_20632_002_0
+mcq_28623_012_1
+mcq_28583_013_0
+mcq_28623_008_1
+mcq_20632_007_2
+mcq_28583_007_0
+mcq_26390_003_1
+mcq_26390_017_0
+mcq_29458_002_2
+mcq_38216_020_0
+mcq_38206_005_0
+mcq_29454_013_0
+mcq_38206_013_2
+mcq_29460_016_1
+mcq_38209_006_2
+mcq_29454_005_2
+mcq_29454_000_1
+mcq_29458_006_0
+mcq_29456_001_1
+mcq_29458_003_0
+mcq_29460_017_1
+mcq_29460_003_2
+mcq_38187_023_0
+mcq_912552759_005_4
+mcq_912552759_006_5
+mcq_900325340_011_2
+mcq_907248739_030_5
+mcq_900325340_011_4
+mcq_912552759_006_4
+mcq_912552759_005_1
+mcq_912552759_006_7
+mcq_907248739_010_4
+mcq_907248739_010_5
+mcq_900325340_015_0
+mcq_907248739_030_2
+mcq_907248739_018_7
+mcq_907248739_009_7
+mcq_907248739_010_1
+mcq_907248739_011_1
+mcq_900325340_016_6
+mcq_907248739_021_3
+mcq_900325340_017_0
+mcq_907248739_008_0
+mcq_907248739_017_3
+mcq_900325340_010_2
+mcq_907248739_009_3
+mcq_907248739_013_5
+mcq_907248739_021_5
+mcq_900325340_011_7
+mcq_900325340_017_7
+mcq_907248739_022_2
+mcq_907248739_022_5
+mcq_907248739_026_6
+mcq_907248739_016_2
+mcq_900325340_013_1
+mcq_900325340_026_0
+mcq_907248739_030_0
+mcq_907248739_009_5
+mcq_907248739_020_4
+mcq_907248739_016_0
+mcq_907248739_020_6
+mcq_907248739_019_0
+mcq_900325340_009_3
+mcq_38260_006_1
+mcq_29487_001_2
+mcq_38249_003_1
+mcq_38262_000_0
+mcq_29485_004_2
+mcq_29486_003_2
+mcq_38260_000_0
+mcq_29488_002_0
+mcq_38263_000_1
+mcq_902928789_006_5
+mcq_902928789_043_0
+mcq_902928789_038_5
+mcq_902928789_011_0
+mcq_902928789_026_2
+mcq_902928789_007_0
+mcq_902928789_034_2
+mcq_902928789_011_2
+mcq_902928789_031_4
+mcq_902928789_027_1
+mcq_908755184_050_0
+mcq_902928789_005_1
+mcq_902928789_029_2
+mcq_902928789_019_2
+mcq_902928789_036_5
+mcq_902928789_035_1
+mcq_902928789_029_1
+mcq_908755184_043_0
+mcq_907562714_007_0
+mcq_908755184_031_2
+mcq_908755184_026_3
+mcq_26233_001_0
+mcq_20322_002_0
+mcq_28019_000_1
+mcq_26208_002_0
+mcq_20322_005_0
+mcq_28004_001_0
+mcq_20327_002_0
+mcq_26186_005_1
+mcq_26226_006_0
+mcq_26216_000_1
+mcq_27955_000_2
+mcq_20319_003_1
+mcq_26216_001_3
+mcq_26216_001_0
+mcq_27975_001_0
+mcq_26226_003_1
+mcq_20322_005_1
+mcq_26186_004_0
+mcq_26220_000_1
+mcq_26186_011_0
+mcq_22172_000_0
+mcq_29340_005_1
+mcq_22190_002_2
+mcq_29363_005_0
+mcq_29376_004_0
+mcq_22194_001_0
+mcq_22175_003_0
+mcq_29340_000_0
+mcq_29363_001_1
+mcq_38855_014_0
+mcq_38838_004_1
+mcq_22189_001_0
+mcq_38879_004_2
+mcq_22186_002_1
+mcq_29340_005_0
+mcq_29368_007_0
+mcq_22192_005_0
+mcq_29363_008_1
+mcq_38843_001_0
+mcq_22194_005_2
+mcq_22192_000_1
+mcq_22177_001_1
+mcq_38883_005_1
+mcq_38829_000_0
+mcq_29368_012_0
+mcq_38861_007_1
+mcq_38883_010_1
+mcq_22194_009_1
+mcq_38829_002_2
+mcq_38861_000_0
+mcq_39084_001_0
+mcq_38976_002_2
+mcq_39060_013_0
+mcq_39084_003_4
+mcq_39095_005_0
+mcq_39084_000_1
+mcq_39072_003_3
+mcq_39023_003_1
+mcq_38976_003_3
+mcq_39072_004_0
+mcq_39007_013_0
+mcq_25538_003_2
+mcq_38948_008_1
+mcq_39072_013_2
+mcq_39023_000_0
+mcq_38948_008_0
+mcq_38976_016_4
+mcq_39060_013_2
+mcq_25538_000_4
+mcq_25538_007_2
+mcq_38976_005_1
+mcq_25549_017_1
+mcq_25476_014_2
+mcq_25459_013_2
+mcq_39023_010_0
+mcq_39095_003_1
+mcq_25555_014_2
+mcq_39007_007_0
+mcq_38976_011_4
+mcq_39060_004_0
+mcq_39084_000_0
+mcq_39023_002_1
+mcq_38948_003_1
+mcq_39007_003_0
+mcq_25549_007_4
+mcq_24866_012_1
+mcq_18673_005_2
+mcq_24683_005_3
+mcq_32977_000_0
+mcq_24866_012_0
+mcq_24649_002_0
+mcq_24649_004_0
+mcq_24763_003_3
+mcq_24843_007_2
+mcq_24702_003_4
+mcq_18633_000_2
+mcq_25114_019_2
+mcq_24755_003_0
+mcq_24866_012_4
+mcq_18637_001_4
+mcq_24773_001_4
+mcq_25176_001_1
+mcq_24702_005_4
+mcq_32977_007_2
+mcq_25167_004_3
+mcq_18641_001_2
+mcq_18658_006_4
+mcq_25066_005_4
+mcq_24792_000_2
+mcq_24803_000_1
+mcq_24812_001_3
+mcq_903217400_152_2
+mcq_903217400_169_3
+mcq_907792485_012_1
+mcq_903217400_212_2
+mcq_907792485_014_3
+mcq_907792485_090_2
+mcq_907792485_014_0
+mcq_906597809_020_2
+mcq_907792485_066_2
+mcq_907792485_089_2
+mcq_913873011_051_1
+mcq_907792485_063_0
+mcq_907792485_026_0
+mcq_907792485_052_2
+mcq_907792485_063_2
+mcq_31810_000_1
+mcq_31763_000_2
+mcq_905145333_054_1
+mcq_905145333_006_1
+mcq_34917_000_0
+mcq_34910_000_3
+mcq_32068_000_2
+mcq_905884636_021_1
+mcq_915703041_031_2
+mcq_915703041_001_0
+mcq_32024_000_2
+mcq_32062_000_5
+mcq_32072_000_4
+mcq_32025_000_1
+mcq_32082_000_2
+mcq_906744851_001_3
+mcq_32084_000_2
+mcq_32075_000_2
+mcq_905719217_187_2
+mcq_915134069_013_3
+mcq_905719217_016_2
+mcq_905719217_290_2
+mcq_912464987_007_1
+mcq_905719217_079_2
+mcq_912464987_015_2
+mcq_905719217_171_1
+mcq_905719217_012_2
+mcq_915134069_014_1
+mcq_905719217_217_2
+mcq_915134069_008_0
+mcq_913593445_028_2
+mcq_930002845_034_3
+mcq_911625854_040_0
+mcq_911625854_038_2
+mcq_911625854_042_3
+mcq_911625854_026_2
+mcq_911625854_023_3
+mcq_930002930_019_2
+mcq_911625854_048_2
+mcq_911625854_026_0
+mcq_913593445_041_2
+mcq_911625854_023_2
+mcq_930002461_029_0
+mcq_29866_000_1
+mcq_29857_000_1
+mcq_29808_000_0
+mcq_29851_000_0
+mcq_29876_000_3
+mcq_901772665_005_1
+mcq_29814_000_2
+mcq_29855_000_2
+mcq_905630786_021_0
+mcq_31039_000_5
+mcq_900181973_021_3
+mcq_900181973_021_2
+mcq_30988_000_3
+mcq_905906820_021_1
+mcq_909740322_060_0
+mcq_909740322_019_2
+mcq_916659363_049_3
+mcq_909740322_031_2
+mcq_909740322_056_2
+mcq_909740322_005_0
+mcq_916659363_003_2
+mcq_909740322_043_1
+mcq_909740322_023_2
+mcq_916659363_057_2
+mcq_916659363_039_2
+mcq_909740322_053_2
+mcq_909740322_022_0
+mcq_909740322_008_2
+mcq_909740322_037_2
+mcq_909740322_057_3
+mcq_909740322_024_2
+mcq_909740322_008_0
+mcq_909740322_037_0
+mcq_916659363_021_2
+mcq_916659363_033_1
+mcq_916659363_045_2
+mcq_909740322_013_2
+mcq_903208054_020_2
+mcq_902849618_031_2
+mcq_903208054_045_3
+mcq_902849618_010_0
+mcq_903208054_036_1
+mcq_902849618_033_3
+mcq_903208054_033_0
+mcq_902849618_031_0
+mcq_903208054_036_0
+mcq_902849618_022_0
+mcq_903208054_026_0
+mcq_903208054_026_1
+mcq_34960_000_1
+mcq_34961_000_0
+mcq_903208054_024_2
+mcq_902849618_009_1
+mcq_910273911_060_3
+mcq_910273911_053_3
+mcq_912636831_090_1
+mcq_912636831_059_0
+mcq_910273911_026_3
+mcq_910273911_050_3
+mcq_910273911_047_2
+mcq_910273911_031_1
+mcq_910273911_104_0
+mcq_910273911_011_1
+mcq_910273911_117_1
+mcq_910273911_011_3
+mcq_910273911_041_1
+mcq_910273911_120_2
+mcq_912636831_016_1
+mcq_910273911_044_3
+mcq_910273911_025_1
+mcq_32675_000_0
+mcq_32675_000_3
+mcq_32697_000_1
+mcq_910273911_069_3
+mcq_910273911_046_3
+mcq_32675_000_1
+mcq_910273911_101_0
+mcq_912636831_072_1
+mcq_912636831_084_0
+mcq_912636831_033_2
+mcq_907144258_025_3
+mcq_909937432_065_2
+mcq_909937432_049_0
+mcq_909937432_027_2
+mcq_909937432_097_0
+mcq_907144258_043_0
+mcq_907144258_006_1
+mcq_907144258_055_0
+mcq_909937432_065_0
+mcq_909937432_065_1
+mcq_909937432_007_2
+mcq_909937432_090_3
+mcq_909937432_019_0
+mcq_909937432_122_3
+mcq_907144258_045_0
+mcq_909937432_128_1
+mcq_909937432_019_3
+mcq_909937432_011_0
+mcq_909937432_085_2
+mcq_907144258_025_0
+mcq_909937432_044_0
+mcq_903906712_080_0
+mcq_908494993_151_2
+mcq_908494993_037_0
+mcq_28067_000_0
+mcq_903906712_067_1
+mcq_908494993_040_2
+mcq_28087_000_2
+mcq_905057208_125_1
+mcq_30561_000_2
+mcq_30642_000_4
+mcq_30658_000_2
+mcq_30608_000_2
+mcq_30645_000_0
+mcq_901474957_002_2
+mcq_30607_000_5
+mcq_30669_000_3
+mcq_916454267_028_3
+mcq_30584_000_3
+mcq_30590_000_4
+mcq_30664_000_4
+mcq_914440602_024_1
+mcq_914440602_042_1
+mcq_914440602_031_1
+mcq_915437942_041_0
+mcq_915437942_029_0
+mcq_914440602_045_1
+mcq_914440602_016_1
+mcq_914440602_012_1
+mcq_914440602_012_0
+mcq_914440602_031_0
+mcq_915437942_014_1
+mcq_914440602_016_3
+mcq_915437942_003_1
+mcq_911005812_011_0
+mcq_911005812_010_1
+mcq_901202257_046_2
+mcq_901202257_025_0
+mcq_930002764_034_0
+mcq_930002818_024_2

training/data/overlap-summary.json

@@ -0,0 +1,24 @@
+{
+ "training_questions": 10000,
+ "organiser_questions": {
+  "dev=data/dev-set": 1851,
+  "test=data/test-set": 2238,
+  "test-old=../Dataset/Test": 2146
+ },
+ "excluded": 1106,
+ "excluded_by_rule": {
+  "same question": 0,
+  "question jaccard >= 0.5": 17,
+  "question containment >= 0.8": 48,
+  "3+ same choices and containment >= 0.3": 11,
+  "all 4 choices the same": 43,
+  "meaning cosine >= 0.70": 1033
+ },
+ "nearest_cosine_percentiles": {
+  "50": 0.598,
+  "90": 0.702,
+  "99": 0.805,
+  "99.9": 0.899
+ },
+ "reported_below_threshold": 1598
+}

training/data/star3-thoughts.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3eddcbfa7d45b5a155f593f0128469a28740c76310f4c7feb42a004abe1feafb
+size 11102486

training/kaggle/star3-gen-0/kernel-metadata.json

@@ -0,0 +1,16 @@
+{
+  "id": "{owner}/sinhalanuwanacode-star3-gen-0",
+  "title": "sinhalanuwanacode-star3-gen-0",
+  "code_file": "run.py",
+  "language": "python",
+  "kernel_type": "script",
+  "is_private": true,
+  "enable_gpu": true,
+  "enable_tpu": false,
+  "enable_internet": true,
+  "machine_shape": "NvidiaTeslaT4",
+  "dataset_sources": ["{owner}/sinhalanuwanacode", "{owner}/sinhalanuwanacode-data"],
+  "competition_sources": [],
+  "kernel_sources": [],
+  "model_sources": []
+}

training/kaggle/star3-gen-0/run.py

@@ -0,0 +1,33 @@
+import glob
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+REVISION = "ee0ef6023621cff504d758262d4e04895a5af4a2"
+
+
+def find(name):
+    hits = glob.glob(f"/kaggle/input/**/{name}", recursive=True)
+    if len(hits) != 1:
+        raise SystemExit(f"error: expected one {name} under /kaggle/input: {hits}")
+    return Path(hits[0])
+
+
+code = find("predict.py").parent
+questions = find("star2/shard-0.jsonl")
+print("code snapshot:", code, (code / "COMMIT").read_text().strip())
+print("questions:", questions)
+
+env = {**os.environ, "HF_HOME": "/kaggle/temp/hf"}
+offline = {**env, "HF_HUB_OFFLINE": "1"}
+subprocess.run([sys.executable, "-m", "pip", "install", "-q", "-r", str(code / "requirements.txt")], check=True)
+subprocess.run([sys.executable, "-m", "pip", "uninstall", "-y", "-q", "torchao"], check=True)
+subprocess.run(["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv"], check=True)
+subprocess.run(["hf", "download", "google/gemma-4-E4B-it", "--revision", REVISION], check=True, env=env)
+subprocess.run(
+    [sys.executable, str(code / "star.py"), "--train", str(questions), "--no_hint_retry",
+     "--revision", REVISION, "--out_dir", "/kaggle/working/star"],
+    check=True,
+    env=offline,
+)

training/kaggle/star3-gen-1/kernel-metadata.json

@@ -0,0 +1,16 @@
+{
+  "id": "{owner}/sinhalanuwanacode-star3-gen-1",
+  "title": "sinhalanuwanacode-star3-gen-1",
+  "code_file": "run.py",
+  "language": "python",
+  "kernel_type": "script",
+  "is_private": true,
+  "enable_gpu": true,
+  "enable_tpu": false,
+  "enable_internet": true,
+  "machine_shape": "NvidiaTeslaT4",
+  "dataset_sources": ["{owner}/sinhalanuwanacode", "{owner}/sinhalanuwanacode-data"],
+  "competition_sources": [],
+  "kernel_sources": [],
+  "model_sources": []
+}

training/kaggle/star3-gen-1/run.py

@@ -0,0 +1,33 @@
+import glob
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+REVISION = "ee0ef6023621cff504d758262d4e04895a5af4a2"
+
+
+def find(name):
+    hits = glob.glob(f"/kaggle/input/**/{name}", recursive=True)
+    if len(hits) != 1:
+        raise SystemExit(f"error: expected one {name} under /kaggle/input: {hits}")
+    return Path(hits[0])
+
+
+code = find("predict.py").parent
+questions = find("star2/shard-1.jsonl")
+print("code snapshot:", code, (code / "COMMIT").read_text().strip())
+print("questions:", questions)
+
+env = {**os.environ, "HF_HOME": "/kaggle/temp/hf"}
+offline = {**env, "HF_HUB_OFFLINE": "1"}
+subprocess.run([sys.executable, "-m", "pip", "install", "-q", "-r", str(code / "requirements.txt")], check=True)
+subprocess.run([sys.executable, "-m", "pip", "uninstall", "-y", "-q", "torchao"], check=True)
+subprocess.run(["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv"], check=True)
+subprocess.run(["hf", "download", "google/gemma-4-E4B-it", "--revision", REVISION], check=True, env=env)
+subprocess.run(
+    [sys.executable, str(code / "star.py"), "--train", str(questions), "--no_hint_retry",
+     "--revision", REVISION, "--out_dir", "/kaggle/working/star"],
+    check=True,
+    env=offline,
+)

training/kaggle/star3-train/kernel-metadata.json

@@ -0,0 +1,16 @@
+{
+  "id": "{owner}/sinhalanuwanacode-star3-train",
+  "title": "sinhalanuwanacode-star3-train",
+  "code_file": "run.py",
+  "language": "python",
+  "kernel_type": "script",
+  "is_private": true,
+  "enable_gpu": true,
+  "enable_tpu": false,
+  "enable_internet": true,
+  "machine_shape": "NvidiaTeslaT4",
+  "dataset_sources": ["{owner}/sinhalanuwanacode"],
+  "competition_sources": [],
+  "kernel_sources": ["{owner}/sinhalanuwanacode-star3-gen-0", "{owner}/sinhalanuwanacode-star3-gen-1"],
+  "model_sources": []
+}

training/kaggle/star3-train/run.py

@@ -0,0 +1,40 @@
+import glob
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+REVISION = "ee0ef6023621cff504d758262d4e04895a5af4a2"
+
+
+def find(name):
+    hits = glob.glob(f"/kaggle/input/**/{name}", recursive=True)
+    if len(hits) != 1:
+        raise SystemExit(f"error: expected one {name} under /kaggle/input: {hits}")
+    return Path(hits[0])
+
+
+code = find("predict.py").parent
+print("code snapshot:", code, (code / "COMMIT").read_text().strip())
+parts = sorted(glob.glob("/kaggle/input/**/star/star.jsonl", recursive=True))
+if len(parts) != 2:
+    raise SystemExit(f"error: expected the star.jsonl of both star3-gen notebooks: {parts}")
+lines = [line for p in parts for line in Path(p).read_text(encoding="utf-8").splitlines()]
+direct = [line for line in lines if json.loads(line)["source"] == "direct"]
+data = Path("/kaggle/working/star3.jsonl")
+data.write_text("".join(line + "\n" for line in direct), encoding="utf-8")
+print(f"examples: {len(direct)} first-try thoughts written by the model without an adapter")
+
+env = {**os.environ, "HF_HOME": "/kaggle/temp/hf"}
+offline = {**env, "HF_HUB_OFFLINE": "1"}
+subprocess.run([sys.executable, "-m", "pip", "install", "-q", "-r", str(code / "requirements.txt")], check=True)
+subprocess.run([sys.executable, "-m", "pip", "uninstall", "-y", "-q", "torchao"], check=True)
+subprocess.run(["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv"], check=True)
+subprocess.run(["hf", "download", "google/gemma-4-E4B-it", "--revision", REVISION], check=True, env=env)
+subprocess.run(
+    [sys.executable, str(code / "train_lora.py"), "--data", str(data), "--revision", REVISION,
+     "--out_dir", "/kaggle/working/adapter"],
+    check=True,
+    env=offline,
+)

training/prepare_star_data.py

@@ -0,0 +1,33 @@
+import argparse
+import json
+import random
+from pathlib import Path
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--questions", default="data/train/verified-mcq.jsonl")
+    parser.add_argument("--exclude", default="data/train/star2/overlap/excluded_ids.txt",
+                        help="ids that review_overlap.py found similar to a dev or test question")
+    parser.add_argument("--out_dir", default="data/train/star2")
+    parser.add_argument("--shards", type=int, default=2)
+    parser.add_argument("--seed", type=int, default=0)
+    args = parser.parse_args()
+
+    lines = Path(args.questions).read_text(encoding="utf-8").splitlines()
+    excluded = set(Path(args.exclude).read_text(encoding="utf-8").split())
+    kept = [line for line in lines if json.loads(line)["id"] not in excluded]
+    assert len(lines) - len(kept) == len(excluded), "an excluded id is not in the question file"
+
+    random.Random(args.seed).shuffle(kept)
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    for n in range(args.shards):
+        shard = kept[n :: args.shards]
+        (out_dir / f"shard-{n}.jsonl").write_text("".join(line + "\n" for line in shard), encoding="utf-8")
+        print(f"shard-{n}.jsonl: {len(shard)} questions")
+    print(f"Kept {len(kept)} of {len(lines)} questions; excluded {len(excluded)} listed in {args.exclude}")
+
+
+if __name__ == "__main__":
+    main()

training/review_overlap.py

@@ -0,0 +1,151 @@
+import argparse
+import json
+import re
+import unicodedata
+from pathlib import Path
+
+import numpy as np
+import torch
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+from transformers import AutoModel, AutoTokenizer
+
+from predict import load_questions, normalise
+
+NGRAM = 6
+EMBEDDER = "sentence-transformers/LaBSE"
+RULES = {
+    "same question": lambda m: m["question_key_equal"],
+    "question jaccard >= 0.5": lambda m: m["jaccard"] >= 0.5,
+    "question containment >= 0.8": lambda m: m["containment"] >= 0.8 and m["min_shingles"] >= 8,
+    "3+ same choices and containment >= 0.3": lambda m: m["same_choices"] >= 3 and m["containment"] >= 0.3,
+    "all 4 choices the same": lambda m: m["same_choices"] >= 4,
+    "meaning cosine >= 0.70": lambda m: m["cosine"] >= 0.70,
+}
+
+
+def key(text):
+    return re.sub(r"[\s​-‍\W_]+", "", unicodedata.normalize("NFC", str(text))).lower()
+
+
+def shingles(text):
+    k = key(text)
+    return {k[i : i + NGRAM] for i in range(max(1, len(k) - NGRAM + 1))}
+
+
+def as_text(q):
+    return q["question"].strip() + "\n" + "\n".join(str(c) for c in q["choices"])
+
+
+def embed(texts, batch_size=64):
+    path = Path(snapshot_download(EMBEDDER, allow_patterns=[
+        "config.json", "model.safetensors", "tokenizer*", "vocab.txt", "special_tokens_map.json", "2_Dense/model.safetensors",
+    ]))
+    tokenizer = AutoTokenizer.from_pretrained(path)
+    model = AutoModel.from_pretrained(path).eval()
+    dense = load_file(path / "2_Dense" / "model.safetensors")
+    weight, bias = dense["linear.weight"], dense["linear.bias"]
+    out = []
+    with torch.inference_mode():
+        for i in range(0, len(texts), batch_size):
+            batch = tokenizer(texts[i : i + batch_size], padding=True, truncation=True, max_length=256,
+                              return_tensors="pt")
+            cls = model(**batch).last_hidden_state[:, 0]
+            vec = torch.tanh(cls @ weight.T + bias)
+            out.append(torch.nn.functional.normalize(vec, dim=-1).numpy())
+    return np.concatenate(out)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--questions", default="data/train/verified-mcq.jsonl")
+    parser.add_argument("--eval", nargs="+", default=["dev=data/dev-set", "test=data/test-set", "test-old=../Dataset/Test"],
+                        help="name=folder of organiser question files")
+    parser.add_argument("--out_dir", default="outputs/overlap")
+    parser.add_argument("--report_cosine", type=float, default=0.65)
+    args = parser.parse_args()
+
+    train = [normalise(json.loads(line)) for line in Path(args.questions).read_text(encoding="utf-8").splitlines()]
+    evals = []
+    for spec in args.eval:
+        name, folder = spec.split("=", 1)
+        evals.extend({**q, "set": name} for q in load_questions(folder))
+    print(f"{len(train)} training questions against {len(evals)} organiser questions from {args.eval}")
+
+    eval_keys = [key(q["question"]) for q in evals]
+    eval_shingles = [shingles(q["question"]) for q in evals]
+    eval_choices = [{key(c) for c in q["choices"]} for q in evals]
+    index, choice_index = {}, {}
+    for i, s in enumerate(eval_shingles):
+        for g in s:
+            index.setdefault(g, []).append(i)
+        for c in eval_choices[i]:
+            choice_index.setdefault(c, []).append(i)
+
+    vectors = embed([as_text(q) for q in train] + [as_text(q) for q in evals])
+    cosine = vectors[: len(train)] @ vectors[len(train) :].T
+
+    rows, excluded = [], []
+    for t, q in enumerate(train):
+        s = shingles(q["question"])
+        shared = {}
+        for g in s:
+            for i in index.get(g, ()):
+                shared[i] = shared.get(i, 0) + 1
+        choices = {key(c) for c in q["choices"]}
+        shared_choices = {}
+        for c in choices:
+            for i in choice_index.get(c, ()):
+                shared_choices[i] = shared_choices.get(i, 0) + 1
+        candidates = (
+            {i for i, n in shared.items()
+             if n / (len(s) + len(eval_shingles[i]) - n) >= 0.5 or n / min(len(s), len(eval_shingles[i])) >= 0.8}
+            | {i for i, n in shared_choices.items() if n >= 3}
+            | set(np.argsort(-cosine[t])[:5].tolist())
+        )
+        matches = []
+        for i in sorted(candidates):
+            n = shared.get(i, 0)
+            m = {
+                "question_key_equal": key(q["question"]) == eval_keys[i],
+                "jaccard": n / (len(s) + len(eval_shingles[i]) - n),
+                "containment": n / min(len(s), len(eval_shingles[i])),
+                "min_shingles": min(len(s), len(eval_shingles[i])),
+                "same_choices": len(choices & eval_choices[i]),
+                "cosine": float(cosine[t, i]),
+            }
+            m["rules"] = [name for name, rule in RULES.items() if rule(m)]
+            matches.append((i, m))
+        i, m = max(matches, key=lambda x: (len(x[1]["rules"]), x[1]["cosine"]))
+        if m["rules"] or m["cosine"] >= args.report_cosine:
+            rows.append({
+                "id": q["id"], "flagged": bool(m["rules"]), "rules": m["rules"],
+                "cosine": round(m["cosine"], 3), "jaccard": round(m["jaccard"], 3),
+                "containment": round(m["containment"], 3), "same_choices": m["same_choices"],
+                "train_question": q["question"], "train_choices": q["choices"],
+                "eval_set": evals[i]["set"], "eval_q_no": evals[i].get("q_no"), "eval_subject": evals[i].get("subject"),
+                "eval_question": evals[i]["question"], "eval_choices": evals[i]["choices"],
+            })
+        if m["rules"]:
+            excluded.append(q["id"])
+
+    rows.sort(key=lambda r: (not r["flagged"], -r["cosine"]))
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    (out_dir / "pairs.jsonl").write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows), encoding="utf-8")
+    (out_dir / "excluded_ids.txt").write_text("".join(i + "\n" for i in excluded), encoding="utf-8")
+    best = cosine.max(axis=1)
+    summary = {
+        "training_questions": len(train),
+        "organiser_questions": {spec: sum(q["set"] == spec.split("=")[0] for q in evals) for spec in args.eval},
+        "excluded": len(excluded),
+        "excluded_by_rule": {name: sum(name in r["rules"] for r in rows) for name in RULES},
+        "nearest_cosine_percentiles": {p: round(float(np.percentile(best, p)), 3) for p in (50, 90, 99, 99.9)},
+        "reported_below_threshold": sum(not r["flagged"] for r in rows),
+    }
+    (out_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
+    print(json.dumps(summary, ensure_ascii=False, indent=1))
+
+
+if __name__ == "__main__":
+    main()

training/star.py

@@ -0,0 +1,112 @@
+import argparse
+import json
+import random
+import re
+from pathlib import Path
+
+from predict import ANSWER_CUE, answer, argmax, load_model, normalise, score, thinking_text, user_prompt, write_jsonl
+
+HINT_MENTION = re.compile(r"hint|ඉඟි|(?:given|provided|stated|suggested) answer", re.IGNORECASE)
+SENTENCE_BREAK = re.compile(r"((?<=[.!?])\s+|\n+)")
+MIN_KEPT_FRACTION = 0.5
+
+
+def hinted_prompt(question):
+    gold = question["answer"]
+    hint = f"Hint: the correct answer is {gold}. {question['choices'][gold - 1]}"
+    return user_prompt(question).removesuffix(ANSWER_CUE) + hint + "\n" + ANSWER_CUE
+
+
+def strip_hint(thought):
+    pieces = SENTENCE_BREAK.split(thought)
+    sentences = zip(pieces[0::2], pieces[1::2] + [""])
+    return "".join(s + sep for s, sep in sentences if not HINT_MENTION.search(s)).strip()
+
+
+def attempt(model, tokenizer, questions, think, batch_size):
+    prompts = [user_prompt(q) for q in questions]
+    results = answer(model, tokenizer, prompts, [len(q["choices"]) for q in questions], think, batch_size)
+    attempts = [
+        {"id": q["id"], "gold": q["answer"], "pred": r["pred"], "finished": r["finished"],
+         "thought_tokens": r["thought_tokens"], "hinted": False, "thought": r["thought"]}
+        for q, r in zip(questions, results)
+    ]
+    kept = [
+        {"id": q["id"], "prompt": p, "thought": r["thought"], "answer": q["answer"], "source": "direct"}
+        for q, p, r in zip(questions, prompts, results)
+        if r["finished"] and r["pred"] == q["answer"]
+    ]
+    return attempts, kept
+
+
+def rationalize(model, tokenizer, questions, think, batch_size):
+    n_choices = [len(q["choices"]) for q in questions]
+    prompts = [hinted_prompt(q) for q in questions]
+    results = answer(model, tokenizer, prompts, n_choices, think, batch_size)
+    cleaned = [strip_hint(r["thought"]) for r in results]
+    texts = [thinking_text(tokenizer, user_prompt(q), c) for q, c in zip(questions, cleaned)]
+    rescored = [argmax(p) + 1 for p in score(model, tokenizer, texts, n_choices, batch_size)]
+    attempts, kept = [], []
+    for q, r, c, pred in zip(questions, results, cleaned, rescored):
+        attempts.append(
+            {"id": q["id"], "gold": q["answer"], "pred": r["pred"], "finished": r["finished"],
+             "thought_tokens": r["thought_tokens"], "hinted": True, "thought": r["thought"],
+             "cleaned": c, "pred_without_hint": pred}
+        )
+        if (
+            r["finished"]
+            and r["pred"] == q["answer"]
+            and pred == q["answer"]
+            and len(c) >= MIN_KEPT_FRACTION * len(r["thought"])
+        ):
+            kept.append({"id": q["id"], "prompt": user_prompt(q), "thought": c, "answer": q["answer"],
+                         "source": "rationalized"})
+    return attempts, kept
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", default="google/gemma-4-E4B-it")
+    parser.add_argument("--revision", help="model commit to load; the latest one when not given")
+    parser.add_argument("--adapter", help="LoRA adapter that writes the thoughts, merged into the model")
+    parser.add_argument("--train", default="data/train/mcq.jsonl")
+    parser.add_argument("--n", type=int, help="sample this many questions with --seed; all of them when not given")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--out_dir", default="star")
+    parser.add_argument("--think", type=int, default=512)
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--no_hint_retry", action="store_true", help="keep only first-try thoughts and skip the hinted retry")
+    args = parser.parse_args()
+    if args.think <= 0:
+        parser.error("--think must be positive, STaR keeps thoughts")
+
+    lines = Path(args.train).read_text(encoding="utf-8").splitlines()
+    questions = [normalise(json.loads(line)) for line in lines]
+    if args.n is not None:
+        questions = random.Random(args.seed).sample(questions, args.n)
+        print(f"Sampled {len(questions)} of {len(lines)} questions from {args.train} with seed {args.seed}")
+    else:
+        print(f"Using all {len(questions)} questions from {args.train}")
+
+    tokenizer, model = load_model(args.model, args.adapter, args.revision)
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    attempts, kept = attempt(model, tokenizer, questions, args.think, args.batch_size)
+    write_jsonl(out_dir / "attempts.jsonl", attempts)
+    write_jsonl(out_dir / "star.jsonl", kept)
+    print(f"Direct: kept {len(kept)}/{len(questions)}")
+    if args.no_hint_retry:
+        return
+
+    done = {row["id"] for row in kept}
+    missed = [q for q in questions if q["id"] not in done]
+    hinted_attempts, rationalized = rationalize(model, tokenizer, missed, args.think, args.batch_size)
+    write_jsonl(out_dir / "attempts.jsonl", attempts + hinted_attempts)
+    write_jsonl(out_dir / "star.jsonl", kept + rationalized)
+    print(f"Rationalized: kept {len(rationalized)}/{len(missed)}")
+    print(f"Wrote {len(kept) + len(rationalized)} examples to {out_dir / 'star.jsonl'}")
+
+
+if __name__ == "__main__":
+    main()

training/train_log.json

@@ -0,0 +1,4626 @@
+{
+ "config": {
+  "model": "google/gemma-4-E4B-it",
+  "revision": "ee0ef6023621cff504d758262d4e04895a5af4a2",
+  "data": "/kaggle/working/star3.jsonl",
+  "out_dir": "/kaggle/working/adapter",
+  "epochs": 2,
+  "lr": 0.0001,
+  "batch_size": 4,
+  "grad_accum": 4,
+  "max_length": 1024,
+  "rank": 16,
+  "alpha": 32,
+  "dropout": 0.05,
+  "seed": 0
+ },
+ "dtype": "torch.float16",
+ "grad_scaling": true,
+ "target_modules": 342,
+ "base_params": 7463013376,
+ "adapter_params": 38666240,
+ "total_params": 7501679616,
+ "examples": 4073,
+ "trained_examples": 4073,
+ "dropped_too_long": 0,
+ "dropped_misaligned": 0,
+ "skipped_steps": 0,
+ "seconds": 15551.6,
+ "tokens_per_s": 317.6,
+ "peak_memory_gb": [
+  10.56,
+  6.99
+ ],
+ "steps": [
+  {
+   "step": 1,
+   "epoch": 0,
+   "loss": 0.19640883523944777,
+   "lr": 0.0,
+   "grad_norm": 0.4852641522884369,
+   "skipped": false,
+   "seconds": 29.4
+  },
+  {
+   "step": 2,
+   "epoch": 0,
+   "loss": 0.2024406308274182,
+   "lr": 6.666666666666667e-06,
+   "grad_norm": 0.48907169699668884,
+   "skipped": false,
+   "seconds": 58.9
+  },
+  {
+   "step": 3,
+   "epoch": 0,
+   "loss": 0.20496504929589732,
+   "lr": 1.3333333333333333e-05,
+   "grad_norm": 0.5109589695930481,
+   "skipped": false,
+   "seconds": 89.4
+  },
+  {
+   "step": 4,
+   "epoch": 0,
+   "loss": 0.20524678683676226,
+   "lr": 2e-05,
+   "grad_norm": 0.5412611365318298,
+   "skipped": false,
+   "seconds": 120.8
+  },
+  {
+   "step": 5,
+   "epoch": 0,
+   "loss": 0.16709267570479677,
+   "lr": 2.6666666666666667e-05,
+   "grad_norm": 0.5452609658241272,
+   "skipped": false,
+   "seconds": 151.0
+  },
+  {
+   "step": 6,
+   "epoch": 0,
+   "loss": 0.15597860216284942,
+   "lr": 3.3333333333333335e-05,
+   "grad_norm": 0.4582180082798004,
+   "skipped": false,
+   "seconds": 180.9
+  },
+  {
+   "step": 7,
+   "epoch": 0,
+   "loss": 0.14760536539304112,
+   "lr": 4e-05,
+   "grad_norm": 0.3433503210544586,
+   "skipped": false,
+   "seconds": 210.8
+  },
+  {
+   "step": 8,
+   "epoch": 0,
+   "loss": 0.14063837973937082,
+   "lr": 4.666666666666667e-05,
+   "grad_norm": 0.304032564163208,
+   "skipped": false,
+   "seconds": 241.3
+  },
+  {
+   "step": 9,
+   "epoch": 0,
+   "loss": 0.15517175885577025,
+   "lr": 5.333333333333333e-05,
+   "grad_norm": 0.2853395938873291,
+   "skipped": false,
+   "seconds": 271.0
+  },
+  {
+   "step": 10,
+   "epoch": 0,
+   "loss": 0.12460355233487874,
+   "lr": 6e-05,
+   "grad_norm": 0.26190483570098877,
+   "skipped": false,
+   "seconds": 301.3
+  },
+  {
+   "step": 11,
+   "epoch": 0,
+   "loss": 0.13429050899623296,
+   "lr": 6.666666666666667e-05,
+   "grad_norm": 0.2822328805923462,
+   "skipped": false,
+   "seconds": 332.0
+  },
+  {
+   "step": 12,
+   "epoch": 0,
+   "loss": 0.10786084480626937,
+   "lr": 7.333333333333333e-05,
+   "grad_norm": 0.2527726888656616,
+   "skipped": false,
+   "seconds": 360.5
+  },
+  {
+   "step": 13,
+   "epoch": 0,
+   "loss": 0.09978185194823579,
+   "lr": 8e-05,
+   "grad_norm": 0.23101957142353058,
+   "skipped": false,
+   "seconds": 391.4
+  },
+  {
+   "step": 14,
+   "epoch": 0,
+   "loss": 0.10522525868517288,
+   "lr": 8.666666666666667e-05,
+   "grad_norm": 0.2543775737285614,
+   "skipped": false,
+   "seconds": 422.5
+  },
+  {
+   "step": 15,
+   "epoch": 0,
+   "loss": 0.08935808227175758,
+   "lr": 9.333333333333334e-05,
+   "grad_norm": 0.24526648223400116,
+   "skipped": false,
+   "seconds": 452.1
+  },
+  {
+   "step": 16,
+   "epoch": 0,
+   "loss": 0.10548639101148338,
+   "lr": 0.0001,
+   "grad_norm": 0.30216461420059204,
+   "skipped": false,
+   "seconds": 484.5
+  },
+  {
+   "step": 17,
+   "epoch": 0,
+   "loss": 0.0861717793149071,
+   "lr": 9.999899300364532e-05,
+   "grad_norm": 0.26806312799453735,
+   "skipped": false,
+   "seconds": 515.7
+  },
+  {
+   "step": 18,
+   "epoch": 0,
+   "loss": 0.08647315657249598,
+   "lr": 9.999597205514297e-05,
+   "grad_norm": 0.3897492587566376,
+   "skipped": false,
+   "seconds": 546.5
+  },
+  {
+   "step": 19,
+   "epoch": 0,
+   "loss": 0.07331878050296768,
+   "lr": 9.99909372761763e-05,
+   "grad_norm": 0.2865033745765686,
+   "skipped": false,
+   "seconds": 575.8
+  },
+  {
+   "step": 20,
+   "epoch": 0,
+   "loss": 0.0808455167275755,
+   "lr": 9.998388886954547e-05,
+   "grad_norm": 0.3762819170951843,
+   "skipped": false,
+   "seconds": 604.1
+  },
+  {
+   "step": 21,
+   "epoch": 0,
+   "loss": 0.08293077495736136,
+   "lr": 9.997482711915927e-05,
+   "grad_norm": 0.3262479901313782,
+   "skipped": false,
+   "seconds": 635.3
+  },
+  {
+   "step": 22,
+   "epoch": 0,
+   "loss": 0.0765426305353788,
+   "lr": 9.996375239002369e-05,
+   "grad_norm": 0.36088672280311584,
+   "skipped": false,
+   "seconds": 664.5
+  },
+  {
+   "step": 23,
+   "epoch": 0,
+   "loss": 0.08666027548914276,
+   "lr": 9.99506651282272e-05,
+   "grad_norm": 0.5469640493392944,
+   "skipped": false,
+   "seconds": 695.4
+  },
+  {
+   "step": 24,
+   "epoch": 0,
+   "loss": 0.08630259219662725,
+   "lr": 9.993556586092281e-05,
+   "grad_norm": 0.456892192363739,
+   "skipped": false,
+   "seconds": 728.2
+  },
+  {
+   "step": 25,
+   "epoch": 0,
+   "loss": 0.0724076038978558,
+   "lr": 9.991845519630678e-05,
+   "grad_norm": 0.39262059330940247,
+   "skipped": false,
+   "seconds": 758.8
+  },
+  {
+   "step": 26,
+   "epoch": 0,
+   "loss": 0.07770107819389818,
+   "lr": 9.989933382359422e-05,
+   "grad_norm": 0.4262242317199707,
+   "skipped": false,
+   "seconds": 791.2
+  },
+  {
+   "step": 27,
+   "epoch": 0,
+   "loss": 0.06933413065568625,
+   "lr": 9.987820251299122e-05,
+   "grad_norm": 0.46644866466522217,
+   "skipped": false,
+   "seconds": 821.9
+  },
+  {
+   "step": 28,
+   "epoch": 0,
+   "loss": 0.065214456355704,
+   "lr": 9.985506211566388e-05,
+   "grad_norm": 0.9025358557701111,
+   "skipped": false,
+   "seconds": 852.3
+  },
+  {
+   "step": 29,
+   "epoch": 0,
+   "loss": 0.06811053849340527,
+   "lr": 9.982991356370404e-05,
+   "grad_norm": 0.4470258355140686,
+   "skipped": false,
+   "seconds": 881.8
+  },
+  {
+   "step": 30,
+   "epoch": 0,
+   "loss": 0.053144117122032736,
+   "lr": 9.98027578700917e-05,
+   "grad_norm": 0.38273853063583374,
+   "skipped": false,
+   "seconds": 912.0
+  },
+  {
+   "step": 31,
+   "epoch": 0,
+   "loss": 0.05634953098012355,
+   "lr": 9.977359612865423e-05,
+   "grad_norm": 0.4882364273071289,
+   "skipped": false,
+   "seconds": 941.3
+  },
+  {
+   "step": 32,
+   "epoch": 0,
+   "loss": 0.055631592902682075,
+   "lr": 9.974242951402235e-05,
+   "grad_norm": 0.3918277621269226,
+   "skipped": false,
+   "seconds": 971.1
+  },
+  {
+   "step": 33,
+   "epoch": 0,
+   "loss": 0.07527670451454374,
+   "lr": 9.970925928158274e-05,
+   "grad_norm": 0.5305290222167969,
+   "skipped": false,
+   "seconds": 1002.1
+  },
+  {
+   "step": 34,
+   "epoch": 0,
+   "loss": 0.05458790223464191,
+   "lr": 9.967408676742751e-05,
+   "grad_norm": 0.4513584077358246,
+   "skipped": false,
+   "seconds": 1031.2
+  },
+  {
+   "step": 35,
+   "epoch": 0,
+   "loss": 0.050869956201693954,
+   "lr": 9.963691338830044e-05,
+   "grad_norm": 0.36274057626724243,
+   "skipped": false,
+   "seconds": 1061.9
+  },
+  {
+   "step": 36,
+   "epoch": 0,
+   "loss": 0.0570102071433521,
+   "lr": 9.959774064153977e-05,
+   "grad_norm": 0.42373621463775635,
+   "skipped": false,
+   "seconds": 1091.4
+  },
+  {
+   "step": 37,
+   "epoch": 0,
+   "loss": 0.05033422113724285,
+   "lr": 9.955657010501806e-05,
+   "grad_norm": 0.46532633900642395,
+   "skipped": false,
+   "seconds": 1121.0
+  },
+  {
+   "step": 38,
+   "epoch": 0,
+   "loss": 0.054250309157325254,
+   "lr": 9.951340343707852e-05,
+   "grad_norm": 0.5128265023231506,
+   "skipped": false,
+   "seconds": 1150.8
+  },
+  {
+   "step": 39,
+   "epoch": 0,
+   "loss": 0.051204504607117715,
+   "lr": 9.946824237646824e-05,
+   "grad_norm": 0.3985678553581238,
+   "skipped": false,
+   "seconds": 1181.8
+  },
+  {
+   "step": 40,
+   "epoch": 0,
+   "loss": 0.06049295101113949,
+   "lr": 9.942108874226811e-05,
+   "grad_norm": 0.525611400604248,
+   "skipped": false,
+   "seconds": 1212.5
+  },
+  {
+   "step": 41,
+   "epoch": 0,
+   "loss": 0.05473621184691105,
+   "lr": 9.937194443381972e-05,
+   "grad_norm": 0.503436803817749,
+   "skipped": false,
+   "seconds": 1242.8
+  },
+  {
+   "step": 42,
+   "epoch": 0,
+   "loss": 0.05886410636377241,
+   "lr": 9.93208114306486e-05,
+   "grad_norm": 0.47664114832878113,
+   "skipped": false,
+   "seconds": 1274.4
+  },
+  {
+   "step": 43,
+   "epoch": 0,
+   "loss": 0.05740216232730541,
+   "lr": 9.926769179238466e-05,
+   "grad_norm": 0.994827389717102,
+   "skipped": false,
+   "seconds": 1304.9
+  },
+  {
+   "step": 44,
+   "epoch": 0,
+   "loss": 0.05837026187973685,
+   "lr": 9.921258765867919e-05,
+   "grad_norm": 0.47819700837135315,
+   "skipped": false,
+   "seconds": 1335.1
+  },
+  {
+   "step": 45,
+   "epoch": 0,
+   "loss": 0.048231410801350474,
+   "lr": 9.915550124911866e-05,
+   "grad_norm": 0.41990071535110474,
+   "skipped": false,
+   "seconds": 1366.3
+  },
+  {
+   "step": 46,
+   "epoch": 0,
+   "loss": 0.04562145920400366,
+   "lr": 9.909643486313533e-05,
+   "grad_norm": 0.3867731988430023,
+   "skipped": false,
+   "seconds": 1396.0
+  },
+  {
+   "step": 47,
+   "epoch": 0,
+   "loss": 0.04700247855921044,
+   "lr": 9.903539087991462e-05,
+   "grad_norm": 0.4377094805240631,
+   "skipped": false,
+   "seconds": 1427.3
+  },
+  {
+   "step": 48,
+   "epoch": 0,
+   "loss": 0.04771211157444897,
+   "lr": 9.897237175829926e-05,
+   "grad_norm": 0.4685591757297516,
+   "skipped": false,
+   "seconds": 1457.2
+  },
+  {
+   "step": 49,
+   "epoch": 0,
+   "loss": 0.0489195060784648,
+   "lr": 9.890738003669029e-05,
+   "grad_norm": 0.4277944564819336,
+   "skipped": false,
+   "seconds": 1487.9
+  },
+  {
+   "step": 50,
+   "epoch": 0,
+   "loss": 0.05864074655654856,
+   "lr": 9.884041833294476e-05,
+   "grad_norm": 0.45892971754074097,
+   "skipped": false,
+   "seconds": 1519.0
+  },
+  {
+   "step": 51,
+   "epoch": 0,
+   "loss": 0.04818263103261309,
+   "lr": 9.877148934427037e-05,
+   "grad_norm": 0.39819756150245667,
+   "skipped": false,
+   "seconds": 1552.3
+  },
+  {
+   "step": 52,
+   "epoch": 0,
+   "loss": 0.05366593056003958,
+   "lr": 9.870059584711668e-05,
+   "grad_norm": 0.6030884385108948,
+   "skipped": false,
+   "seconds": 1581.7
+  },
+  {
+   "step": 53,
+   "epoch": 0,
+   "loss": 0.043071207878591276,
+   "lr": 9.862774069706346e-05,
+   "grad_norm": 0.37357521057128906,
+   "skipped": false,
+   "seconds": 1614.2
+  },
+  {
+   "step": 54,
+   "epoch": 0,
+   "loss": 0.04875012851107735,
+   "lr": 9.855292682870551e-05,
+   "grad_norm": 0.4176650047302246,
+   "skipped": false,
+   "seconds": 1645.7
+  },
+  {
+   "step": 55,
+   "epoch": 0,
+   "loss": 0.05916158294450872,
+   "lr": 9.847615725553456e-05,
+   "grad_norm": 0.5978649258613586,
+   "skipped": false,
+   "seconds": 1675.6
+  },
+  {
+   "step": 56,
+   "epoch": 0,
+   "loss": 0.058098166422967945,
+   "lr": 9.839743506981782e-05,
+   "grad_norm": 0.42823487520217896,
+   "skipped": false,
+   "seconds": 1706.4
+  },
+  {
+   "step": 57,
+   "epoch": 0,
+   "loss": 0.045172476261220076,
+   "lr": 9.831676344247342e-05,
+   "grad_norm": 0.4181194603443146,
+   "skipped": false,
+   "seconds": 1737.6
+  },
+  {
+   "step": 58,
+   "epoch": 0,
+   "loss": 0.05060952837004283,
+   "lr": 9.82341456229428e-05,
+   "grad_norm": 0.5255125761032104,
+   "skipped": false,
+   "seconds": 1765.9
+  },
+  {
+   "step": 59,
+   "epoch": 0,
+   "loss": 0.05047130558174233,
+   "lr": 9.814958493905963e-05,
+   "grad_norm": 0.43823546171188354,
+   "skipped": false,
+   "seconds": 1797.2
+  },
+  {
+   "step": 60,
+   "epoch": 0,
+   "loss": 0.0491023870761557,
+   "lr": 9.806308479691595e-05,
+   "grad_norm": 0.4243999719619751,
+   "skipped": false,
+   "seconds": 1828.1
+  },
+  {
+   "step": 61,
+   "epoch": 0,
+   "loss": 0.043288892164662245,
+   "lr": 9.797464868072488e-05,
+   "grad_norm": 0.4273908734321594,
+   "skipped": false,
+   "seconds": 1857.9
+  },
+  {
+   "step": 62,
+   "epoch": 0,
+   "loss": 0.04162379607413579,
+   "lr": 9.788428015268027e-05,
+   "grad_norm": 0.45563948154449463,
+   "skipped": false,
+   "seconds": 1887.6
+  },
+  {
+   "step": 63,
+   "epoch": 0,
+   "loss": 0.046959113067771105,
+   "lr": 9.779198285281325e-05,
+   "grad_norm": 0.5101006031036377,
+   "skipped": false,
+   "seconds": 1919.9
+  },
+  {
+   "step": 64,
+   "epoch": 0,
+   "loss": 0.04868613549499087,
+   "lr": 9.769776049884563e-05,
+   "grad_norm": 0.4730517864227295,
+   "skipped": false,
+   "seconds": 1951.2
+  },
+  {
+   "step": 65,
+   "epoch": 0,
+   "loss": 0.052818452819319794,
+   "lr": 9.760161688604008e-05,
+   "grad_norm": 0.4125804007053375,
+   "skipped": false,
+   "seconds": 1982.8
+  },
+  {
+   "step": 66,
+   "epoch": 0,
+   "loss": 0.04659271240234375,
+   "lr": 9.750355588704727e-05,
+   "grad_norm": 0.45149803161621094,
+   "skipped": false,
+   "seconds": 2013.5
+  },
+  {
+   "step": 67,
+   "epoch": 0,
+   "loss": 0.04500601429651194,
+   "lr": 9.740358145174998e-05,
+   "grad_norm": 0.44188907742500305,
+   "skipped": false,
+   "seconds": 2045.2
+  },
+  {
+   "step": 68,
+   "epoch": 0,
+   "loss": 0.046970408207588896,
+   "lr": 9.730169760710386e-05,
+   "grad_norm": 0.5863640904426575,
+   "skipped": false,
+   "seconds": 2076.5
+  },
+  {
+   "step": 69,
+   "epoch": 0,
+   "loss": 0.044865242302458606,
+   "lr": 9.719790845697533e-05,
+   "grad_norm": 0.39309433102607727,
+   "skipped": false,
+   "seconds": 2108.1
+  },
+  {
+   "step": 70,
+   "epoch": 0,
+   "loss": 0.04593456075913328,
+   "lr": 9.709221818197624e-05,
+   "grad_norm": 0.44956448674201965,
+   "skipped": false,
+   "seconds": 2139.3
+  },
+  {
+   "step": 71,
+   "epoch": 0,
+   "loss": 0.045809862771893126,
+   "lr": 9.698463103929542e-05,
+   "grad_norm": 0.48698508739471436,
+   "skipped": false,
+   "seconds": 2170.5
+  },
+  {
+   "step": 72,
+   "epoch": 0,
+   "loss": 0.044335218049706415,
+   "lr": 9.687515136252731e-05,
+   "grad_norm": 0.536542534828186,
+   "skipped": false,
+   "seconds": 2201.8
+  },
+  {
+   "step": 73,
+   "epoch": 0,
+   "loss": 0.034864149286479816,
+   "lr": 9.676378356149734e-05,
+   "grad_norm": 0.36436447501182556,
+   "skipped": false,
+   "seconds": 2231.0
+  },
+  {
+   "step": 74,
+   "epoch": 0,
+   "loss": 0.043896083493069775,
+   "lr": 9.665053212208426e-05,
+   "grad_norm": 0.4838595688343048,
+   "skipped": false,
+   "seconds": 2260.5
+  },
+  {
+   "step": 75,
+   "epoch": 0,
+   "loss": 0.04309926780932544,
+   "lr": 9.653540160603956e-05,
+   "grad_norm": 0.4336928725242615,
+   "skipped": false,
+   "seconds": 2291.4
+  },
+  {
+   "step": 76,
+   "epoch": 0,
+   "loss": 0.049297388306396996,
+   "lr": 9.641839665080363e-05,
+   "grad_norm": 0.4917004406452179,
+   "skipped": false,
+   "seconds": 2320.4
+  },
+  {
+   "step": 77,
+   "epoch": 0,
+   "loss": 0.04662362143164262,
+   "lr": 9.629952196931901e-05,
+   "grad_norm": 0.7901037335395813,
+   "skipped": false,
+   "seconds": 2350.6
+  },
+  {
+   "step": 78,
+   "epoch": 0,
+   "loss": 0.04559790509483134,
+   "lr": 9.617878234984055e-05,
+   "grad_norm": 0.48896324634552,
+   "skipped": false,
+   "seconds": 2381.8
+  },
+  {
+   "step": 79,
+   "epoch": 0,
+   "loss": 0.041544652383531604,
+   "lr": 9.60561826557425e-05,
+   "grad_norm": 0.43355318903923035,
+   "skipped": false,
+   "seconds": 2411.3
+  },
+  {
+   "step": 80,
+   "epoch": 0,
+   "loss": 0.04375480673779016,
+   "lr": 9.593172782532268e-05,
+   "grad_norm": 0.4469282031059265,
+   "skipped": false,
+   "seconds": 2442.5
+  },
+  {
+   "step": 81,
+   "epoch": 0,
+   "loss": 0.043051442354047106,
+   "lr": 9.580542287160348e-05,
+   "grad_norm": 0.4262138903141022,
+   "skipped": false,
+   "seconds": 2474.1
+  },
+  {
+   "step": 82,
+   "epoch": 0,
+   "loss": 0.05389878683908257,
+   "lr": 9.567727288213005e-05,
+   "grad_norm": 0.533173680305481,
+   "skipped": false,
+   "seconds": 2505.5
+  },
+  {
+   "step": 83,
+   "epoch": 0,
+   "loss": 0.04077400034217979,
+   "lr": 9.554728301876526e-05,
+   "grad_norm": 0.4008263349533081,
+   "skipped": false,
+   "seconds": 2538.1
+  },
+  {
+   "step": 84,
+   "epoch": 0,
+   "loss": 0.042045998732131294,
+   "lr": 9.541545851748186e-05,
+   "grad_norm": 0.4874799847602844,
+   "skipped": false,
+   "seconds": 2568.2
+  },
+  {
+   "step": 85,
+   "epoch": 0,
+   "loss": 0.042924036080179126,
+   "lr": 9.528180468815155e-05,
+   "grad_norm": 0.4657498896121979,
+   "skipped": false,
+   "seconds": 2599.4
+  },
+  {
+   "step": 86,
+   "epoch": 0,
+   "loss": 0.04021313511834994,
+   "lr": 9.514632691433107e-05,
+   "grad_norm": 0.4809572398662567,
+   "skipped": false,
+   "seconds": 2628.6
+  },
+  {
+   "step": 87,
+   "epoch": 0,
+   "loss": 0.04092615930167263,
+   "lr": 9.50090306530454e-05,
+   "grad_norm": 0.48094093799591064,
+   "skipped": false,
+   "seconds": 2660.8
+  },
+  {
+   "step": 88,
+   "epoch": 0,
+   "loss": 0.0385408574271613,
+   "lr": 9.486992143456792e-05,
+   "grad_norm": 0.47969570755958557,
+   "skipped": false,
+   "seconds": 2690.8
+  },
+  {
+   "step": 89,
+   "epoch": 0,
+   "loss": 0.040967103461144676,
+   "lr": 9.472900486219769e-05,
+   "grad_norm": 0.4366934597492218,
+   "skipped": false,
+   "seconds": 2721.1
+  },
+  {
+   "step": 90,
+   "epoch": 0,
+   "loss": 0.047251623422032964,
+   "lr": 9.458628661203367e-05,
+   "grad_norm": 0.48443788290023804,
+   "skipped": false,
+   "seconds": 2752.0
+  },
+  {
+   "step": 91,
+   "epoch": 0,
+   "loss": 0.0382828146632823,
+   "lr": 9.444177243274618e-05,
+   "grad_norm": 0.4706588685512543,
+   "skipped": false,
+   "seconds": 2783.3
+  },
+  {
+   "step": 92,
+   "epoch": 0,
+   "loss": 0.04608729774082383,
+   "lr": 9.429546814534529e-05,
+   "grad_norm": 0.44139933586120605,
+   "skipped": false,
+   "seconds": 2813.4
+  },
+  {
+   "step": 93,
+   "epoch": 0,
+   "loss": 0.04839884853088036,
+   "lr": 9.414737964294636e-05,
+   "grad_norm": 0.4453069567680359,
+   "skipped": false,
+   "seconds": 2843.4
+  },
+  {
+   "step": 94,
+   "epoch": 0,
+   "loss": 0.0385095516794599,
+   "lr": 9.399751289053267e-05,
+   "grad_norm": 0.42764177918434143,
+   "skipped": false,
+   "seconds": 2874.1
+  },
+  {
+   "step": 95,
+   "epoch": 0,
+   "loss": 0.03839392269232086,
+   "lr": 9.384587392471515e-05,
+   "grad_norm": 0.4124580919742584,
+   "skipped": false,
+   "seconds": 2905.1
+  },
+  {
+   "step": 96,
+   "epoch": 0,
+   "loss": 0.043852357471378486,
+   "lr": 9.369246885348926e-05,
+   "grad_norm": 0.47256457805633545,
+   "skipped": false,
+   "seconds": 2936.3
+  },
+  {
+   "step": 97,
+   "epoch": 0,
+   "loss": 0.041886468532681916,
+   "lr": 9.353730385598887e-05,
+   "grad_norm": 0.4709848165512085,
+   "skipped": false,
+   "seconds": 2965.5
+  },
+  {
+   "step": 98,
+   "epoch": 0,
+   "loss": 0.044607980541999256,
+   "lr": 9.338038518223747e-05,
+   "grad_norm": 0.42527738213539124,
+   "skipped": false,
+   "seconds": 2995.5
+  },
+  {
+   "step": 99,
+   "epoch": 0,
+   "loss": 0.04058560680755121,
+   "lr": 9.322171915289635e-05,
+   "grad_norm": 0.41362443566322327,
+   "skipped": false,
+   "seconds": 3026.1
+  },
+  {
+   "step": 100,
+   "epoch": 0,
+   "loss": 0.036925773563923486,
+   "lr": 9.306131215901003e-05,
+   "grad_norm": 0.39879411458969116,
+   "skipped": false,
+   "seconds": 3056.0
+  },
+  {
+   "step": 101,
+   "epoch": 0,
+   "loss": 0.04319861298706494,
+   "lr": 9.289917066174886e-05,
+   "grad_norm": 0.45430848002433777,
+   "skipped": false,
+   "seconds": 3085.9
+  },
+  {
+   "step": 102,
+   "epoch": 0,
+   "loss": 0.038103842340568954,
+   "lr": 9.273530119214868e-05,
+   "grad_norm": 1.6617456674575806,
+   "skipped": false,
+   "seconds": 3117.3
+  },
+  {
+   "step": 103,
+   "epoch": 0,
+   "loss": 0.04705844420699579,
+   "lr": 9.256971035084785e-05,
+   "grad_norm": 0.5413803458213806,
+   "skipped": false,
+   "seconds": 3148.0
+  },
+  {
+   "step": 104,
+   "epoch": 0,
+   "loss": 0.04588349056332997,
+   "lr": 9.24024048078213e-05,
+   "grad_norm": 0.48643723130226135,
+   "skipped": false,
+   "seconds": 3176.6
+  },
+  {
+   "step": 105,
+   "epoch": 0,
+   "loss": 0.03868896050454584,
+   "lr": 9.223339130211192e-05,
+   "grad_norm": 0.4337347447872162,
+   "skipped": false,
+   "seconds": 3207.2
+  },
+  {
+   "step": 106,
+   "epoch": 0,
+   "loss": 0.036370536171722645,
+   "lr": 9.206267664155907e-05,
+   "grad_norm": 0.41438397765159607,
+   "skipped": false,
+   "seconds": 3236.2
+  },
+  {
+   "step": 107,
+   "epoch": 0,
+   "loss": 0.045510898442028866,
+   "lr": 9.189026770252436e-05,
+   "grad_norm": 0.4909036159515381,
+   "skipped": false,
+   "seconds": 3266.4
+  },
+  {
+   "step": 108,
+   "epoch": 0,
+   "loss": 0.041277409169551436,
+   "lr": 9.171617142961477e-05,
+   "grad_norm": 0.506044864654541,
+   "skipped": false,
+   "seconds": 3295.5
+  },
+  {
+   "step": 109,
+   "epoch": 0,
+   "loss": 0.048108986375681596,
+   "lr": 9.154039483540273e-05,
+   "grad_norm": 0.8060802221298218,
+   "skipped": false,
+   "seconds": 3325.8
+  },
+  {
+   "step": 110,
+   "epoch": 0,
+   "loss": 0.04355858804628586,
+   "lr": 9.136294500014386e-05,
+   "grad_norm": 0.4543282687664032,
+   "skipped": false,
+   "seconds": 3356.2
+  },
+  {
+   "step": 111,
+   "epoch": 0,
+   "loss": 0.038672522803189745,
+   "lr": 9.118382907149165e-05,
+   "grad_norm": 0.40678146481513977,
+   "skipped": false,
+   "seconds": 3387.2
+  },
+  {
+   "step": 112,
+   "epoch": 0,
+   "loss": 0.042898832788448915,
+   "lr": 9.100305426420956e-05,
+   "grad_norm": 0.46925193071365356,
+   "skipped": false,
+   "seconds": 3415.2
+  },
+  {
+   "step": 113,
+   "epoch": 0,
+   "loss": 0.044555553183160804,
+   "lr": 9.082062785988049e-05,
+   "grad_norm": 0.4316609501838684,
+   "skipped": false,
+   "seconds": 3445.5
+  },
+  {
+   "step": 114,
+   "epoch": 0,
+   "loss": 0.05310061437274338,
+   "lr": 9.06365572066134e-05,
+   "grad_norm": 0.47825491428375244,
+   "skipped": false,
+   "seconds": 3476.2
+  },
+  {
+   "step": 115,
+   "epoch": 0,
+   "loss": 0.039852327125268344,
+   "lr": 9.045084971874738e-05,
+   "grad_norm": 0.4318072199821472,
+   "skipped": false,
+   "seconds": 3507.8
+  },
+  {
+   "step": 116,
+   "epoch": 0,
+   "loss": 0.04183180106903781,
+   "lr": 9.026351287655294e-05,
+   "grad_norm": 0.41331470012664795,
+   "skipped": false,
+   "seconds": 3540.0
+  },
+  {
+   "step": 117,
+   "epoch": 0,
+   "loss": 0.03974040917733554,
+   "lr": 9.007455422593077e-05,
+   "grad_norm": 0.5103293657302856,
+   "skipped": false,
+   "seconds": 3570.9
+  },
+  {
+   "step": 118,
+   "epoch": 0,
+   "loss": 0.044967052919180975,
+   "lr": 8.988398137810777e-05,
+   "grad_norm": 0.4268462359905243,
+   "skipped": false,
+   "seconds": 3602.0
+  },
+  {
+   "step": 119,
+   "epoch": 0,
+   "loss": 0.036238125160800834,
+   "lr": 8.969180200933047e-05,
+   "grad_norm": 0.3959253132343292,
+   "skipped": false,
+   "seconds": 3634.4
+  },
+  {
+   "step": 120,
+   "epoch": 0,
+   "loss": 0.03730900039784556,
+   "lr": 8.949802386055581e-05,
+   "grad_norm": 0.4061117172241211,
+   "skipped": false,
+   "seconds": 3665.0
+  },
+  {
+   "step": 121,
+   "epoch": 0,
+   "loss": 0.04315282394073202,
+   "lr": 8.930265473713938e-05,
+   "grad_norm": 0.43387264013290405,
+   "skipped": false,
+   "seconds": 3695.2
+  },
+  {
+   "step": 122,
+   "epoch": 0,
+   "loss": 0.04045547115192259,
+   "lr": 8.910570250852097e-05,
+   "grad_norm": 0.43794116377830505,
+   "skipped": false,
+   "seconds": 3725.1
+  },
+  {
+   "step": 123,
+   "epoch": 0,
+   "loss": 0.039255543758994656,
+   "lr": 8.890717510790763e-05,
+   "grad_norm": 0.421598881483078,
+   "skipped": false,
+   "seconds": 3753.7
+  },
+  {
+   "step": 124,
+   "epoch": 0,
+   "loss": 0.04119531671735658,
+   "lr": 8.870708053195413e-05,
+   "grad_norm": 0.41828998923301697,
+   "skipped": false,
+   "seconds": 3782.7
+  },
+  {
+   "step": 125,
+   "epoch": 0,
+   "loss": 0.03619634128686488,
+   "lr": 8.850542684044078e-05,
+   "grad_norm": 0.3894548714160919,
+   "skipped": false,
+   "seconds": 3813.8
+  },
+  {
+   "step": 126,
+   "epoch": 0,
+   "loss": 0.04429675767294329,
+   "lr": 8.83022221559489e-05,
+   "grad_norm": 0.4620448052883148,
+   "skipped": false,
+   "seconds": 3843.6
+  },
+  {
+   "step": 127,
+   "epoch": 0,
+   "loss": 0.03996556843091452,
+   "lr": 8.809747466353356e-05,
+   "grad_norm": 0.4065825343132019,
+   "skipped": false,
+   "seconds": 3874.3
+  },
+  {
+   "step": 128,
+   "epoch": 0,
+   "loss": 0.0439553558475418,
+   "lr": 8.789119261039385e-05,
+   "grad_norm": 0.4792153239250183,
+   "skipped": false,
+   "seconds": 3904.4
+  },
+  {
+   "step": 129,
+   "epoch": 0,
+   "loss": 0.038623366088057345,
+   "lr": 8.768338430554082e-05,
+   "grad_norm": 0.38725003600120544,
+   "skipped": false,
+   "seconds": 3933.6
+  },
+  {
+   "step": 130,
+   "epoch": 0,
+   "loss": 0.034875345331959726,
+   "lr": 8.74740581194627e-05,
+   "grad_norm": 0.46025949716567993,
+   "skipped": false,
+   "seconds": 3962.5
+  },
+  {
+   "step": 131,
+   "epoch": 0,
+   "loss": 0.03991871696805529,
+   "lr": 8.726322248378775e-05,
+   "grad_norm": 0.5329133868217468,
+   "skipped": false,
+   "seconds": 3993.1
+  },
+  {
+   "step": 132,
+   "epoch": 0,
+   "loss": 0.04528188935060945,
+   "lr": 8.705088589094459e-05,
+   "grad_norm": 0.6367639899253845,
+   "skipped": false,
+   "seconds": 4023.4
+  },
+  {
+   "step": 133,
+   "epoch": 0,
+   "loss": 0.04366885924049116,
+   "lr": 8.683705689382024e-05,
+   "grad_norm": 0.8130128979682922,
+   "skipped": false,
+   "seconds": 4054.7
+  },
+  {
+   "step": 134,
+   "epoch": 0,
+   "loss": 0.04263949807086629,
+   "lr": 8.662174410541555e-05,
+   "grad_norm": 0.47220155596733093,
+   "skipped": false,
+   "seconds": 4086.5
+  },
+  {
+   "step": 135,
+   "epoch": 0,
+   "loss": 0.03748772638578097,
+   "lr": 8.640495619849821e-05,
+   "grad_norm": 0.7511334419250488,
+   "skipped": false,
+   "seconds": 4115.9
+  },
+  {
+   "step": 136,
+   "epoch": 0,
+   "loss": 0.03817583855450704,
+   "lr": 8.618670190525352e-05,
+   "grad_norm": 0.45687562227249146,
+   "skipped": false,
+   "seconds": 4146.9
+  },
+  {
+   "step": 137,
+   "epoch": 0,
+   "loss": 0.045329121330418734,
+   "lr": 8.596699001693255e-05,
+   "grad_norm": 0.5111539959907532,
+   "skipped": false,
+   "seconds": 4177.5
+  },
+  {
+   "step": 138,
+   "epoch": 0,
+   "loss": 0.04316695398120789,
+   "lr": 8.574582938349817e-05,
+   "grad_norm": 0.47795966267585754,
+   "skipped": false,
+   "seconds": 4207.5
+  },
+  {
+   "step": 139,
+   "epoch": 0,
+   "loss": 0.03760606574186284,
+   "lr": 8.552322891326846e-05,
+   "grad_norm": 0.4370393753051758,
+   "skipped": false,
+   "seconds": 4238.5
+  },
+  {
+   "step": 140,
+   "epoch": 0,
+   "loss": 0.03411327835019874,
+   "lr": 8.529919757255783e-05,
+   "grad_norm": 0.5971311926841736,
+   "skipped": false,
+   "seconds": 4269.4
+  },
+  {
+   "step": 141,
+   "epoch": 0,
+   "loss": 0.040823135821758484,
+   "lr": 8.507374438531607e-05,
+   "grad_norm": 0.4045528471469879,
+   "skipped": false,
+   "seconds": 4300.7
+  },
+  {
+   "step": 142,
+   "epoch": 0,
+   "loss": 0.03689725307248674,
+   "lr": 8.484687843276469e-05,
+   "grad_norm": 0.503023624420166,
+   "skipped": false,
+   "seconds": 4330.5
+  },
+  {
+   "step": 143,
+   "epoch": 0,
+   "loss": 0.04255549059662948,
+   "lr": 8.461860885303114e-05,
+   "grad_norm": 1.0986597537994385,
+   "skipped": false,
+   "seconds": 4362.6
+  },
+  {
+   "step": 144,
+   "epoch": 0,
+   "loss": 0.03933713789797076,
+   "lr": 8.438894484078086e-05,
+   "grad_norm": 0.384954571723938,
+   "skipped": false,
+   "seconds": 4394.5
+  },
+  {
+   "step": 145,
+   "epoch": 0,
+   "loss": 0.04379340127772898,
+   "lr": 8.415789564684673e-05,
+   "grad_norm": 0.42970529198646545,
+   "skipped": false,
+   "seconds": 4424.8
+  },
+  {
+   "step": 146,
+   "epoch": 0,
+   "loss": 0.0348121738856441,
+   "lr": 8.392547057785661e-05,
+   "grad_norm": 0.3794669210910797,
+   "skipped": false,
+   "seconds": 4454.9
+  },
+  {
+   "step": 147,
+   "epoch": 0,
+   "loss": 0.03421579952458653,
+   "lr": 8.369167899585841e-05,
+   "grad_norm": 0.3601197600364685,
+   "skipped": false,
+   "seconds": 4485.3
+  },
+  {
+   "step": 148,
+   "epoch": 0,
+   "loss": 0.042821301594670964,
+   "lr": 8.345653031794292e-05,
+   "grad_norm": 0.4644930958747864,
+   "skipped": false,
+   "seconds": 4516.7
+  },
+  {
+   "step": 149,
+   "epoch": 0,
+   "loss": 0.044358414786973915,
+   "lr": 8.322003401586462e-05,
+   "grad_norm": 0.4486781060695648,
+   "skipped": false,
+   "seconds": 4547.8
+  },
+  {
+   "step": 150,
+   "epoch": 0,
+   "loss": 0.04154949754934255,
+   "lr": 8.298219961566009e-05,
+   "grad_norm": 0.4688260853290558,
+   "skipped": false,
+   "seconds": 4579.2
+  },
+  {
+   "step": 151,
+   "epoch": 0,
+   "loss": 0.04206281148280588,
+   "lr": 8.274303669726426e-05,
+   "grad_norm": 0.4345685541629791,
+   "skipped": false,
+   "seconds": 4609.6
+  },
+  {
+   "step": 152,
+   "epoch": 0,
+   "loss": 0.03806614825690483,
+   "lr": 8.250255489412463e-05,
+   "grad_norm": 0.41471055150032043,
+   "skipped": false,
+   "seconds": 4639.3
+  },
+  {
+   "step": 153,
+   "epoch": 0,
+   "loss": 0.036419154975548464,
+   "lr": 8.226076389281316e-05,
+   "grad_norm": 0.39894402027130127,
+   "skipped": false,
+   "seconds": 4670.0
+  },
+  {
+   "step": 154,
+   "epoch": 0,
+   "loss": 0.0412287198855671,
+   "lr": 8.201767343263612e-05,
+   "grad_norm": 0.5558825135231018,
+   "skipped": false,
+   "seconds": 4699.3
+  },
+  {
+   "step": 155,
+   "epoch": 0,
+   "loss": 0.03543429750900332,
+   "lr": 8.177329330524182e-05,
+   "grad_norm": 0.3999188542366028,
+   "skipped": false,
+   "seconds": 4731.7
+  },
+  {
+   "step": 156,
+   "epoch": 0,
+   "loss": 0.0313229477893021,
+   "lr": 8.152763335422613e-05,
+   "grad_norm": 0.36880818009376526,
+   "skipped": false,
+   "seconds": 4762.2
+  },
+  {
+   "step": 157,
+   "epoch": 0,
+   "loss": 0.04267634258650647,
+   "lr": 8.128070347473609e-05,
+   "grad_norm": 0.476824015378952,
+   "skipped": false,
+   "seconds": 4793.5
+  },
+  {
+   "step": 158,
+   "epoch": 0,
+   "loss": 0.036534420304552194,
+   "lr": 8.103251361307119e-05,
+   "grad_norm": 0.3953181505203247,
+   "skipped": false,
+   "seconds": 4824.2
+  },
+  {
+   "step": 159,
+   "epoch": 0,
+   "loss": 0.04185770834502996,
+   "lr": 8.07830737662829e-05,
+   "grad_norm": 0.4075761139392853,
+   "skipped": false,
+   "seconds": 4854.8
+  },
+  {
+   "step": 160,
+   "epoch": 0,
+   "loss": 0.03086744183635943,
+   "lr": 8.053239398177191e-05,
+   "grad_norm": 0.3988531529903412,
+   "skipped": false,
+   "seconds": 4885.7
+  },
+  {
+   "step": 161,
+   "epoch": 0,
+   "loss": 0.030656032216276878,
+   "lr": 8.028048435688333e-05,
+   "grad_norm": 0.41813963651657104,
+   "skipped": false,
+   "seconds": 4915.8
+  },
+  {
+   "step": 162,
+   "epoch": 0,
+   "loss": 0.03457538713973674,
+   "lr": 8.002735503850016e-05,
+   "grad_norm": 0.5292803049087524,
+   "skipped": false,
+   "seconds": 4946.6
+  },
+  {
+   "step": 163,
+   "epoch": 0,
+   "loss": 0.04431626313934932,
+   "lr": 7.97730162226344e-05,
+   "grad_norm": 0.46250706911087036,
+   "skipped": false,
+   "seconds": 4977.1
+  },
+  {
+   "step": 164,
+   "epoch": 0,
+   "loss": 0.03613297641050265,
+   "lr": 7.95174781540165e-05,
+   "grad_norm": 0.4473329484462738,
+   "skipped": false,
+   "seconds": 5006.5
+  },
+  {
+   "step": 165,
+   "epoch": 0,
+   "loss": 0.04166648178634685,
+   "lr": 7.926075112568259e-05,
+   "grad_norm": 0.45099830627441406,
+   "skipped": false,
+   "seconds": 5039.0
+  },
+  {
+   "step": 166,
+   "epoch": 0,
+   "loss": 0.0461718571295432,
+   "lr": 7.900284547855991e-05,
+   "grad_norm": 0.4615466594696045,
+   "skipped": false,
+   "seconds": 5070.3
+  },
+  {
+   "step": 167,
+   "epoch": 0,
+   "loss": 0.03558595718876008,
+   "lr": 7.874377160105036e-05,
+   "grad_norm": 0.40478748083114624,
+   "skipped": false,
+   "seconds": 5101.0
+  },
+  {
+   "step": 168,
+   "epoch": 0,
+   "loss": 0.03655152665563376,
+   "lr": 7.848353992861195e-05,
+   "grad_norm": 0.8066860437393188,
+   "skipped": false,
+   "seconds": 5131.0
+  },
+  {
+   "step": 169,
+   "epoch": 0,
+   "loss": 0.04193836329213726,
+   "lr": 7.822216094333847e-05,
+   "grad_norm": 0.47176116704940796,
+   "skipped": false,
+   "seconds": 5161.5
+  },
+  {
+   "step": 170,
+   "epoch": 0,
+   "loss": 0.04063318507788144,
+   "lr": 7.795964517353735e-05,
+   "grad_norm": 0.6494160294532776,
+   "skipped": false,
+   "seconds": 5190.9
+  },
+  {
+   "step": 171,
+   "epoch": 0,
+   "loss": 0.031576382511724656,
+   "lr": 7.769600319330552e-05,
+   "grad_norm": 0.33635416626930237,
+   "skipped": false,
+   "seconds": 5222.1
+  },
+  {
+   "step": 172,
+   "epoch": 0,
+   "loss": 0.03734357056854235,
+   "lr": 7.74312456221035e-05,
+   "grad_norm": 0.4150267243385315,
+   "skipped": false,
+   "seconds": 5253.9
+  },
+  {
+   "step": 173,
+   "epoch": 0,
+   "loss": 0.032456102336556145,
+   "lr": 7.716538312432766e-05,
+   "grad_norm": 0.4122450649738312,
+   "skipped": false,
+   "seconds": 5285.7
+  },
+  {
+   "step": 174,
+   "epoch": 0,
+   "loss": 0.04108603528774564,
+   "lr": 7.689842640888063e-05,
+   "grad_norm": 0.4406798183917999,
+   "skipped": false,
+   "seconds": 5316.2
+  },
+  {
+   "step": 175,
+   "epoch": 0,
+   "loss": 0.034947429970073884,
+   "lr": 7.663038622873999e-05,
+   "grad_norm": 0.388076514005661,
+   "skipped": false,
+   "seconds": 5345.7
+  },
+  {
+   "step": 176,
+   "epoch": 0,
+   "loss": 0.03984182050221951,
+   "lr": 7.636127338052512e-05,
+   "grad_norm": 0.4097675085067749,
+   "skipped": false,
+   "seconds": 5376.6
+  },
+  {
+   "step": 177,
+   "epoch": 0,
+   "loss": 0.03554444531306821,
+   "lr": 7.60910987040623e-05,
+   "grad_norm": 0.4486690163612366,
+   "skipped": false,
+   "seconds": 5404.5
+  },
+  {
+   "step": 178,
+   "epoch": 0,
+   "loss": 0.03406732369196045,
+   "lr": 7.58198730819481e-05,
+   "grad_norm": 0.36147722601890564,
+   "skipped": false,
+   "seconds": 5437.3
+  },
+  {
+   "step": 179,
+   "epoch": 0,
+   "loss": 0.03996114220770357,
+   "lr": 7.554760743911103e-05,
+   "grad_norm": 0.6578200459480286,
+   "skipped": false,
+   "seconds": 5468.4
+  },
+  {
+   "step": 180,
+   "epoch": 0,
+   "loss": 0.03974208406398927,
+   "lr": 7.52743127423715e-05,
+   "grad_norm": 0.5207664370536804,
+   "skipped": false,
+   "seconds": 5498.2
+  },
+  {
+   "step": 181,
+   "epoch": 0,
+   "loss": 0.04600800008996591,
+   "lr": 7.500000000000001e-05,
+   "grad_norm": 0.5798813104629517,
+   "skipped": false,
+   "seconds": 5528.0
+  },
+  {
+   "step": 182,
+   "epoch": 0,
+   "loss": 0.039908104562122335,
+   "lr": 7.472468026127385e-05,
+   "grad_norm": 0.4419390559196472,
+   "skipped": false,
+   "seconds": 5559.1
+  },
+  {
+   "step": 183,
+   "epoch": 0,
+   "loss": 0.03292939452479223,
+   "lr": 7.444836461603195e-05,
+   "grad_norm": 0.3788493871688843,
+   "skipped": false,
+   "seconds": 5588.1
+  },
+  {
+   "step": 184,
+   "epoch": 0,
+   "loss": 0.04538562315394896,
+   "lr": 7.417106419422819e-05,
+   "grad_norm": 0.48531362414360046,
+   "skipped": false,
+   "seconds": 5617.4
+  },
+  {
+   "step": 185,
+   "epoch": 0,
+   "loss": 0.03726213373671817,
+   "lr": 7.389279016548316e-05,
+   "grad_norm": 0.4027515649795532,
+   "skipped": false,
+   "seconds": 5650.0
+  },
+  {
+   "step": 186,
+   "epoch": 0,
+   "loss": 0.03393715986206677,
+   "lr": 7.361355373863414e-05,
+   "grad_norm": 0.3956986963748932,
+   "skipped": false,
+   "seconds": 5680.1
+  },
+  {
+   "step": 187,
+   "epoch": 0,
+   "loss": 0.04360706535456188,
+   "lr": 7.333336616128369e-05,
+   "grad_norm": 0.4776063561439514,
+   "skipped": false,
+   "seconds": 5709.6
+  },
+  {
+   "step": 188,
+   "epoch": 0,
+   "loss": 0.05108193350914854,
+   "lr": 7.305223871934657e-05,
+   "grad_norm": 0.4673639237880707,
+   "skipped": false,
+   "seconds": 5739.7
+  },
+  {
+   "step": 189,
+   "epoch": 0,
+   "loss": 0.03513687010553615,
+   "lr": 7.277018273659517e-05,
+   "grad_norm": 0.3786344826221466,
+   "skipped": false,
+   "seconds": 5769.9
+  },
+  {
+   "step": 190,
+   "epoch": 0,
+   "loss": 0.037264012151926776,
+   "lr": 7.24872095742033e-05,
+   "grad_norm": 0.4696302115917206,
+   "skipped": false,
+   "seconds": 5801.1
+  },
+  {
+   "step": 191,
+   "epoch": 0,
+   "loss": 0.03839773619213788,
+   "lr": 7.220333063028872e-05,
+   "grad_norm": 1.4208773374557495,
+   "skipped": false,
+   "seconds": 5832.1
+  },
+  {
+   "step": 192,
+   "epoch": 0,
+   "loss": 0.037638320253307794,
+   "lr": 7.191855733945387e-05,
+   "grad_norm": 0.4687727987766266,
+   "skipped": false,
+   "seconds": 5861.9
+  },
+  {
+   "step": 193,
+   "epoch": 0,
+   "loss": 0.033059018984484906,
+   "lr": 7.163290117232542e-05,
+   "grad_norm": 0.46348440647125244,
+   "skipped": false,
+   "seconds": 5892.3
+  },
+  {
+   "step": 194,
+   "epoch": 0,
+   "loss": 0.037771729388868955,
+   "lr": 7.13463736350921e-05,
+   "grad_norm": 0.39239150285720825,
+   "skipped": false,
+   "seconds": 5921.5
+  },
+  {
+   "step": 195,
+   "epoch": 0,
+   "loss": 0.04014919540012298,
+   "lr": 7.105898626904134e-05,
+   "grad_norm": 0.4191938638687134,
+   "skipped": false,
+   "seconds": 5952.6
+  },
+  {
+   "step": 196,
+   "epoch": 0,
+   "loss": 0.032287506898942726,
+   "lr": 7.077075065009433e-05,
+   "grad_norm": 0.3714689612388611,
+   "skipped": false,
+   "seconds": 5983.1
+  },
+  {
+   "step": 197,
+   "epoch": 0,
+   "loss": 0.03435342616117841,
+   "lr": 7.048167838833977e-05,
+   "grad_norm": 0.43851959705352783,
+   "skipped": false,
+   "seconds": 6014.3
+  },
+  {
+   "step": 198,
+   "epoch": 0,
+   "loss": 0.031339017286041354,
+   "lr": 7.019178112756624e-05,
+   "grad_norm": 1.0957211256027222,
+   "skipped": false,
+   "seconds": 6045.5
+  },
+  {
+   "step": 199,
+   "epoch": 0,
+   "loss": 0.03213614967368697,
+   "lr": 6.990107054479312e-05,
+   "grad_norm": 0.4667567014694214,
+   "skipped": false,
+   "seconds": 6075.9
+  },
+  {
+   "step": 200,
+   "epoch": 0,
+   "loss": 0.03776222146967406,
+   "lr": 6.960955834980028e-05,
+   "grad_norm": 0.49867287278175354,
+   "skipped": false,
+   "seconds": 6106.6
+  },
+  {
+   "step": 201,
+   "epoch": 0,
+   "loss": 0.03299907082185107,
+   "lr": 6.931725628465643e-05,
+   "grad_norm": 0.4388504922389984,
+   "skipped": false,
+   "seconds": 6137.2
+  },
+  {
+   "step": 202,
+   "epoch": 0,
+   "loss": 0.02845420716160882,
+   "lr": 6.902417612324615e-05,
+   "grad_norm": 0.4300897717475891,
+   "skipped": false,
+   "seconds": 6167.6
+  },
+  {
+   "step": 203,
+   "epoch": 0,
+   "loss": 0.048140803469887264,
+   "lr": 6.873032967079561e-05,
+   "grad_norm": 0.4720727205276489,
+   "skipped": false,
+   "seconds": 6196.9
+  },
+  {
+   "step": 204,
+   "epoch": 0,
+   "loss": 0.04197649929110564,
+   "lr": 6.843572876339705e-05,
+   "grad_norm": 0.5733126997947693,
+   "skipped": false,
+   "seconds": 6228.0
+  },
+  {
+   "step": 205,
+   "epoch": 0,
+   "loss": 0.0407797022241052,
+   "lr": 6.814038526753205e-05,
+   "grad_norm": 0.4141710102558136,
+   "skipped": false,
+   "seconds": 6258.9
+  },
+  {
+   "step": 206,
+   "epoch": 0,
+   "loss": 0.03560397631227188,
+   "lr": 6.784431107959359e-05,
+   "grad_norm": 0.4068869650363922,
+   "skipped": false,
+   "seconds": 6289.4
+  },
+  {
+   "step": 207,
+   "epoch": 0,
+   "loss": 0.03401413873482035,
+   "lr": 6.754751812540679e-05,
+   "grad_norm": 0.3947623372077942,
+   "skipped": false,
+   "seconds": 6319.6
+  },
+  {
+   "step": 208,
+   "epoch": 0,
+   "loss": 0.04135789566635663,
+   "lr": 6.725001835974853e-05,
+   "grad_norm": 0.4253414571285248,
+   "skipped": false,
+   "seconds": 6351.6
+  },
+  {
+   "step": 209,
+   "epoch": 0,
+   "loss": 0.03787880008467034,
+   "lr": 6.695182376586603e-05,
+   "grad_norm": 0.4344177544116974,
+   "skipped": false,
+   "seconds": 6382.8
+  },
+  {
+   "step": 210,
+   "epoch": 0,
+   "loss": 0.03571759499921796,
+   "lr": 6.665294635499404e-05,
+   "grad_norm": 0.3693549633026123,
+   "skipped": false,
+   "seconds": 6413.8
+  },
+  {
+   "step": 211,
+   "epoch": 0,
+   "loss": 0.037497748775431376,
+   "lr": 6.635339816587109e-05,
+   "grad_norm": 0.4440297484397888,
+   "skipped": false,
+   "seconds": 6443.4
+  },
+  {
+   "step": 212,
+   "epoch": 0,
+   "loss": 0.03740905688192839,
+   "lr": 6.605319126425454e-05,
+   "grad_norm": 0.4463842213153839,
+   "skipped": false,
+   "seconds": 6474.5
+  },
+  {
+   "step": 213,
+   "epoch": 0,
+   "loss": 0.04070149968421992,
+   "lr": 6.575233774243465e-05,
+   "grad_norm": 0.6535020470619202,
+   "skipped": false,
+   "seconds": 6505.5
+  },
+  {
+   "step": 214,
+   "epoch": 0,
+   "loss": 0.03551791987161815,
+   "lr": 6.545084971874738e-05,
+   "grad_norm": 0.4679766297340393,
+   "skipped": false,
+   "seconds": 6536.2
+  },
+  {
+   "step": 215,
+   "epoch": 0,
+   "loss": 0.0383298000410759,
+   "lr": 6.514873933708638e-05,
+   "grad_norm": 0.3855707347393036,
+   "skipped": false,
+   "seconds": 6566.2
+  },
+  {
+   "step": 216,
+   "epoch": 0,
+   "loss": 0.037831665922547046,
+   "lr": 6.484601876641375e-05,
+   "grad_norm": 1.139644742012024,
+   "skipped": false,
+   "seconds": 6595.9
+  },
+  {
+   "step": 217,
+   "epoch": 0,
+   "loss": 0.03648110361960469,
+   "lr": 6.454270020026995e-05,
+   "grad_norm": 0.4173310101032257,
+   "skipped": false,
+   "seconds": 6626.4
+  },
+  {
+   "step": 218,
+   "epoch": 0,
+   "loss": 0.033973121468026765,
+   "lr": 6.423879585628261e-05,
+   "grad_norm": 0.4316014349460602,
+   "skipped": false,
+   "seconds": 6655.2
+  },
+  {
+   "step": 219,
+   "epoch": 0,
+   "loss": 0.04896280768634715,
+   "lr": 6.39343179756744e-05,
+   "grad_norm": 1.1312209367752075,
+   "skipped": false,
+   "seconds": 6685.7
+  },
+  {
+   "step": 220,
+   "epoch": 0,
+   "loss": 0.03579341232938403,
+   "lr": 6.36292788227699e-05,
+   "grad_norm": 0.40495115518569946,
+   "skipped": false,
+   "seconds": 6716.8
+  },
+  {
+   "step": 221,
+   "epoch": 0,
+   "loss": 0.05005555049648453,
+   "lr": 6.332369068450174e-05,
+   "grad_norm": 0.47920170426368713,
+   "skipped": false,
+   "seconds": 6749.8
+  },
+  {
+   "step": 222,
+   "epoch": 0,
+   "loss": 0.03518659511842654,
+   "lr": 6.30175658699156e-05,
+   "grad_norm": 0.47350093722343445,
+   "skipped": false,
+   "seconds": 6780.5
+  },
+  {
+   "step": 223,
+   "epoch": 0,
+   "loss": 0.03665763939553687,
+   "lr": 6.271091670967436e-05,
+   "grad_norm": 0.42811599373817444,
+   "skipped": false,
+   "seconds": 6810.9
+  },
+  {
+   "step": 224,
+   "epoch": 0,
+   "loss": 0.07557717694559939,
+   "lr": 6.240375555556145e-05,
+   "grad_norm": 0.5362834930419922,
+   "skipped": false,
+   "seconds": 6841.8
+  },
+  {
+   "step": 225,
+   "epoch": 0,
+   "loss": 0.034370320694503485,
+   "lr": 6.209609477998338e-05,
+   "grad_norm": 0.3875359296798706,
+   "skipped": false,
+   "seconds": 6873.3
+  },
+  {
+   "step": 226,
+   "epoch": 0,
+   "loss": 0.03586724312368848,
+   "lr": 6.178794677547137e-05,
+   "grad_norm": 0.4375789761543274,
+   "skipped": false,
+   "seconds": 6903.7
+  },
+  {
+   "step": 227,
+   "epoch": 0,
+   "loss": 0.03204754625441711,
+   "lr": 6.147932395418205e-05,
+   "grad_norm": 0.37311556935310364,
+   "skipped": false,
+   "seconds": 6934.6
+  },
+  {
+   "step": 228,
+   "epoch": 0,
+   "loss": 0.029651823812095246,
+   "lr": 6.117023874739772e-05,
+   "grad_norm": 0.3873976469039917,
+   "skipped": false,
+   "seconds": 6965.5
+  },
+  {
+   "step": 229,
+   "epoch": 0,
+   "loss": 0.03313081313206076,
+   "lr": 6.0860703605025395e-05,
+   "grad_norm": 0.40196579694747925,
+   "skipped": false,
+   "seconds": 6997.5
+  },
+  {
+   "step": 230,
+   "epoch": 0,
+   "loss": 0.032857811465442976,
+   "lr": 6.05507309950955e-05,
+   "grad_norm": 0.3450320065021515,
+   "skipped": false,
+   "seconds": 7025.7
+  },
+  {
+   "step": 231,
+   "epoch": 0,
+   "loss": 0.03128418165351189,
+   "lr": 6.024033340325954e-05,
+   "grad_norm": 0.33902105689048767,
+   "skipped": false,
+   "seconds": 7058.4
+  },
+  {
+   "step": 232,
+   "epoch": 0,
+   "loss": 0.033792713284052735,
+   "lr": 5.992952333228728e-05,
+   "grad_norm": 0.3558861017227173,
+   "skipped": false,
+   "seconds": 7088.5
+  },
+  {
+   "step": 233,
+   "epoch": 0,
+   "loss": 0.03406227171537694,
+   "lr": 5.9618313301563055e-05,
+   "grad_norm": 0.4204740524291992,
+   "skipped": false,
+   "seconds": 7120.3
+  },
+  {
+   "step": 234,
+   "epoch": 0,
+   "loss": 0.03475309157182153,
+   "lr": 5.9306715846581506e-05,
+   "grad_norm": 0.38286036252975464,
+   "skipped": false,
+   "seconds": 7149.4
+  },
+  {
+   "step": 235,
+   "epoch": 0,
+   "loss": 0.032176283458823195,
+   "lr": 5.8994743518442694e-05,
+   "grad_norm": 0.40629762411117554,
+   "skipped": false,
+   "seconds": 7180.5
+  },
+  {
+   "step": 236,
+   "epoch": 0,
+   "loss": 0.031434473539711294,
+   "lr": 5.868240888334653e-05,
+   "grad_norm": 0.3440655469894409,
+   "skipped": false,
+   "seconds": 7211.0
+  },
+  {
+   "step": 237,
+   "epoch": 0,
+   "loss": 0.036813643134599484,
+   "lr": 5.836972452208654e-05,
+   "grad_norm": 0.3786338269710541,
+   "skipped": false,
+   "seconds": 7242.7
+  },
+  {
+   "step": 238,
+   "epoch": 0,
+   "loss": 0.028255267931254515,
+   "lr": 5.805670302954321e-05,
+   "grad_norm": 0.3729824125766754,
+   "skipped": false,
+   "seconds": 7274.4
+  },
+  {
+   "step": 239,
+   "epoch": 0,
+   "loss": 0.041683754948011256,
+   "lr": 5.7743357014176624e-05,
+   "grad_norm": 0.475638210773468,
+   "skipped": false,
+   "seconds": 7306.1
+  },
+  {
+   "step": 240,
+   "epoch": 0,
+   "loss": 0.03886621624397832,
+   "lr": 5.7429699097518585e-05,
+   "grad_norm": 0.44749245047569275,
+   "skipped": false,
+   "seconds": 7337.5
+  },
+  {
+   "step": 241,
+   "epoch": 0,
+   "loss": 0.0370755632110252,
+   "lr": 5.7115741913664264e-05,
+   "grad_norm": 0.4899304211139679,
+   "skipped": false,
+   "seconds": 7368.2
+  },
+  {
+   "step": 242,
+   "epoch": 0,
+   "loss": 0.032659575216887886,
+   "lr": 5.680149810876322e-05,
+   "grad_norm": 0.3824728727340698,
+   "skipped": false,
+   "seconds": 7398.7
+  },
+  {
+   "step": 243,
+   "epoch": 0,
+   "loss": 0.041655277951663454,
+   "lr": 5.6486980340510086e-05,
+   "grad_norm": 0.5157487988471985,
+   "skipped": false,
+   "seconds": 7428.5
+  },
+  {
+   "step": 244,
+   "epoch": 0,
+   "loss": 0.03493482644935842,
+   "lr": 5.617220127763474e-05,
+   "grad_norm": 0.8254408836364746,
+   "skipped": false,
+   "seconds": 7458.7
+  },
+  {
+   "step": 245,
+   "epoch": 0,
+   "loss": 0.03271396995609642,
+   "lr": 5.585717359939192e-05,
+   "grad_norm": 0.4264238178730011,
+   "skipped": false,
+   "seconds": 7489.6
+  },
+  {
+   "step": 246,
+   "epoch": 0,
+   "loss": 0.035557805931201635,
+   "lr": 5.5541909995050554e-05,
+   "grad_norm": 0.7259626388549805,
+   "skipped": false,
+   "seconds": 7520.5
+  },
+  {
+   "step": 247,
+   "epoch": 0,
+   "loss": 0.03830919612537731,
+   "lr": 5.522642316338268e-05,
+   "grad_norm": 0.5846589803695679,
+   "skipped": false,
+   "seconds": 7550.8
+  },
+  {
+   "step": 248,
+   "epoch": 0,
+   "loss": 0.030078849312401398,
+   "lr": 5.4910725812151864e-05,
+   "grad_norm": 0.3776968717575073,
+   "skipped": false,
+   "seconds": 7581.3
+  },
+  {
+   "step": 249,
+   "epoch": 0,
+   "loss": 0.03155149267491511,
+   "lr": 5.4594830657601384e-05,
+   "grad_norm": 0.3725808262825012,
+   "skipped": false,
+   "seconds": 7610.1
+  },
+  {
+   "step": 250,
+   "epoch": 0,
+   "loss": 0.028693306861097875,
+   "lr": 5.427875042394199e-05,
+   "grad_norm": 0.4124157130718231,
+   "skipped": false,
+   "seconds": 7642.6
+  },
+  {
+   "step": 251,
+   "epoch": 0,
+   "loss": 0.03217434105471025,
+   "lr": 5.396249784283942e-05,
+   "grad_norm": 0.5312711000442505,
+   "skipped": false,
+   "seconds": 7673.0
+  },
+  {
+   "step": 252,
+   "epoch": 0,
+   "loss": 0.03104523542030881,
+   "lr": 5.364608565290155e-05,
+   "grad_norm": 0.4328680634498596,
+   "skipped": false,
+   "seconds": 7704.1
+  },
+  {
+   "step": 253,
+   "epoch": 0,
+   "loss": 0.038518297395887566,
+   "lr": 5.3329526599165204e-05,
+   "grad_norm": 0.4788975417613983,
+   "skipped": false,
+   "seconds": 7735.2
+  },
+  {
+   "step": 254,
+   "epoch": 0,
+   "loss": 0.03644856197733275,
+   "lr": 5.301283343258293e-05,
+   "grad_norm": 0.3760280907154083,
+   "skipped": false,
+   "seconds": 7765.5
+  },
+  {
+   "step": 255,
+   "epoch": 0,
+   "loss": 0.040502291837035805,
+   "lr": 5.2696018909509306e-05,
+   "grad_norm": 0.609524667263031,
+   "skipped": false,
+   "seconds": 7781.9
+  },
+  {
+   "step": 256,
+   "epoch": 1,
+   "loss": 0.024327973316301866,
+   "lr": 5.2379095791187124e-05,
+   "grad_norm": 0.34370267391204834,
+   "skipped": false,
+   "seconds": 7813.3
+  },
+  {
+   "step": 257,
+   "epoch": 1,
+   "loss": 0.023225568969183718,
+   "lr": 5.2062076843233366e-05,
+   "grad_norm": 0.2865596413612366,
+   "skipped": false,
+   "seconds": 7844.7
+  },
+  {
+   "step": 258,
+   "epoch": 1,
+   "loss": 0.020964424002418502,
+   "lr": 5.174497483512506e-05,
+   "grad_norm": 0.29307377338409424,
+   "skipped": false,
+   "seconds": 7874.0
+  },
+  {
+   "step": 259,
+   "epoch": 1,
+   "loss": 0.02225420688061004,
+   "lr": 5.142780253968481e-05,
+   "grad_norm": 0.25329098105430603,
+   "skipped": false,
+   "seconds": 7905.2
+  },
+  {
+   "step": 260,
+   "epoch": 1,
+   "loss": 0.019704535258939154,
+   "lr": 5.1110572732566475e-05,
+   "grad_norm": 0.2547796964645386,
+   "skipped": false,
+   "seconds": 7936.0
+  },
+  {
+   "step": 261,
+   "epoch": 1,
+   "loss": 0.029142963339943442,
+   "lr": 5.0793298191740404e-05,
+   "grad_norm": 0.457135945558548,
+   "skipped": false,
+   "seconds": 7966.4
+  },
+  {
+   "step": 262,
+   "epoch": 1,
+   "loss": 0.024066689716361174,
+   "lr": 5.047599169697884e-05,
+   "grad_norm": 0.2987236976623535,
+   "skipped": false,
+   "seconds": 7998.0
+  },
+  {
+   "step": 263,
+   "epoch": 1,
+   "loss": 0.021392749322089965,
+   "lr": 5.015866602934112e-05,
+   "grad_norm": 0.2809486985206604,
+   "skipped": false,
+   "seconds": 8028.5
+  },
+  {
+   "step": 264,
+   "epoch": 1,
+   "loss": 0.018585208267671207,
+   "lr": 4.984133397065889e-05,
+   "grad_norm": 0.2966468930244446,
+   "skipped": false,
+   "seconds": 8060.1
+  },
+  {
+   "step": 265,
+   "epoch": 1,
+   "loss": 0.023124013230306148,
+   "lr": 4.952400830302117e-05,
+   "grad_norm": 0.4006018340587616,
+   "skipped": false,
+   "seconds": 8091.2
+  },
+  {
+   "step": 266,
+   "epoch": 1,
+   "loss": 0.019659400110981634,
+   "lr": 4.92067018082596e-05,
+   "grad_norm": 0.29323676228523254,
+   "skipped": false,
+   "seconds": 8120.9
+  },
+  {
+   "step": 267,
+   "epoch": 1,
+   "loss": 0.024482026068742894,
+   "lr": 4.888942726743353e-05,
+   "grad_norm": 0.3428184390068054,
+   "skipped": false,
+   "seconds": 8151.8
+  },
+  {
+   "step": 268,
+   "epoch": 1,
+   "loss": 0.02003022017938944,
+   "lr": 4.85721974603152e-05,
+   "grad_norm": 0.3119227886199951,
+   "skipped": false,
+   "seconds": 8182.3
+  },
+  {
+   "step": 269,
+   "epoch": 1,
+   "loss": 0.026439435787854695,
+   "lr": 4.825502516487497e-05,
+   "grad_norm": 0.3782959580421448,
+   "skipped": false,
+   "seconds": 8212.2
+  },
+  {
+   "step": 270,
+   "epoch": 1,
+   "loss": 0.014993694335972417,
+   "lr": 4.7937923156766646e-05,
+   "grad_norm": 0.23553837835788727,
+   "skipped": false,
+   "seconds": 8242.5
+  },
+  {
+   "step": 271,
+   "epoch": 1,
+   "loss": 0.01803465460008746,
+   "lr": 4.762090420881289e-05,
+   "grad_norm": 0.2514285147190094,
+   "skipped": false,
+   "seconds": 8273.6
+  },
+  {
+   "step": 272,
+   "epoch": 1,
+   "loss": 0.019902860205701138,
+   "lr": 4.730398109049071e-05,
+   "grad_norm": 0.2984468340873718,
+   "skipped": false,
+   "seconds": 8304.0
+  },
+  {
+   "step": 273,
+   "epoch": 1,
+   "loss": 0.01947339402151121,
+   "lr": 4.698716656741708e-05,
+   "grad_norm": 0.2736494541168213,
+   "skipped": false,
+   "seconds": 8335.5
+  },
+  {
+   "step": 274,
+   "epoch": 1,
+   "loss": 0.01885466807551386,
+   "lr": 4.667047340083481e-05,
+   "grad_norm": 0.25517645478248596,
+   "skipped": false,
+   "seconds": 8364.6
+  },
+  {
+   "step": 275,
+   "epoch": 1,
+   "loss": 0.01877950233943932,
+   "lr": 4.635391434709847e-05,
+   "grad_norm": 0.2774704396724701,
+   "skipped": false,
+   "seconds": 8394.3
+  },
+  {
+   "step": 276,
+   "epoch": 1,
+   "loss": 0.018360153023932645,
+   "lr": 4.603750215716057e-05,
+   "grad_norm": 0.32435178756713867,
+   "skipped": false,
+   "seconds": 8424.5
+  },
+  {
+   "step": 277,
+   "epoch": 1,
+   "loss": 0.01531967046614941,
+   "lr": 4.5721249576058027e-05,
+   "grad_norm": 0.3591376543045044,
+   "skipped": false,
+   "seconds": 8456.0
+  },
+  {
+   "step": 278,
+   "epoch": 1,
+   "loss": 0.016068661741435234,
+   "lr": 4.5405169342398634e-05,
+   "grad_norm": 0.32163316011428833,
+   "skipped": false,
+   "seconds": 8484.3
+  },
+  {
+   "step": 279,
+   "epoch": 1,
+   "loss": 0.017871897473814482,
+   "lr": 4.508927418784815e-05,
+   "grad_norm": 0.2604014277458191,
+   "skipped": false,
+   "seconds": 8515.2
+  },
+  {
+   "step": 280,
+   "epoch": 1,
+   "loss": 0.021977489380508728,
+   "lr": 4.477357683661734e-05,
+   "grad_norm": 0.3197914958000183,
+   "skipped": false,
+   "seconds": 8547.2
+  },
+  {
+   "step": 281,
+   "epoch": 1,
+   "loss": 0.02219422790948942,
+   "lr": 4.445809000494946e-05,
+   "grad_norm": 0.36903974413871765,
+   "skipped": false,
+   "seconds": 8577.6
+  },
+  {
+   "step": 282,
+   "epoch": 1,
+   "loss": 0.017691280488369144,
+   "lr": 4.4142826400608086e-05,
+   "grad_norm": 0.29918819665908813,
+   "skipped": false,
+   "seconds": 8607.6
+  },
+  {
+   "step": 283,
+   "epoch": 1,
+   "loss": 0.016365259626193267,
+   "lr": 4.3827798722365264e-05,
+   "grad_norm": 0.34842997789382935,
+   "skipped": false,
+   "seconds": 8638.8
+  },
+  {
+   "step": 284,
+   "epoch": 1,
+   "loss": 0.020046620939416723,
+   "lr": 4.351301965948991e-05,
+   "grad_norm": 0.2710837423801422,
+   "skipped": false,
+   "seconds": 8671.2
+  },
+  {
+   "step": 285,
+   "epoch": 1,
+   "loss": 0.01676117178713448,
+   "lr": 4.3198501891236804e-05,
+   "grad_norm": 0.27728036046028137,
+   "skipped": false,
+   "seconds": 8700.4
+  },
+  {
+   "step": 286,
+   "epoch": 1,
+   "loss": 0.017474436593507493,
+   "lr": 4.288425808633575e-05,
+   "grad_norm": 0.2791229486465454,
+   "skipped": false,
+   "seconds": 8731.9
+  },
+  {
+   "step": 287,
+   "epoch": 1,
+   "loss": 0.01543151501258758,
+   "lr": 4.2570300902481426e-05,
+   "grad_norm": 0.25070998072624207,
+   "skipped": false,
+   "seconds": 8762.9
+  },
+  {
+   "step": 288,
+   "epoch": 1,
+   "loss": 0.017339263243646895,
+   "lr": 4.2256642985823395e-05,
+   "grad_norm": 0.2656727433204651,
+   "skipped": false,
+   "seconds": 8793.3
+  },
+  {
+   "step": 289,
+   "epoch": 1,
+   "loss": 0.022872012997628945,
+   "lr": 4.19432969704568e-05,
+   "grad_norm": 0.36633551120758057,
+   "skipped": false,
+   "seconds": 8823.6
+  },
+  {
+   "step": 290,
+   "epoch": 1,
+   "loss": 0.020003974834431417,
+   "lr": 4.163027547791347e-05,
+   "grad_norm": 0.26786261796951294,
+   "skipped": false,
+   "seconds": 8853.5
+  },
+  {
+   "step": 291,
+   "epoch": 1,
+   "loss": 0.016302781465769377,
+   "lr": 4.131759111665349e-05,
+   "grad_norm": 0.2734525501728058,
+   "skipped": false,
+   "seconds": 8884.3
+  },
+  {
+   "step": 292,
+   "epoch": 1,
+   "loss": 0.01887829475529337,
+   "lr": 4.100525648155731e-05,
+   "grad_norm": 0.31800055503845215,
+   "skipped": false,
+   "seconds": 8913.4
+  },
+  {
+   "step": 293,
+   "epoch": 1,
+   "loss": 0.020399836793905158,
+   "lr": 4.06932841534185e-05,
+   "grad_norm": 0.3321515619754791,
+   "skipped": false,
+   "seconds": 8946.0
+  },
+  {
+   "step": 294,
+   "epoch": 1,
+   "loss": 0.021608482857536256,
+   "lr": 4.038168669843697e-05,
+   "grad_norm": 0.3299846351146698,
+   "skipped": false,
+   "seconds": 8976.4
+  },
+  {
+   "step": 295,
+   "epoch": 1,
+   "loss": 0.015938243765422702,
+   "lr": 4.007047666771274e-05,
+   "grad_norm": 0.32036611437797546,
+   "skipped": false,
+   "seconds": 9007.7
+  },
+  {
+   "step": 296,
+   "epoch": 1,
+   "loss": 0.016782717475817,
+   "lr": 3.9759666596740476e-05,
+   "grad_norm": 0.2806603014469147,
+   "skipped": false,
+   "seconds": 9039.7
+  },
+  {
+   "step": 297,
+   "epoch": 1,
+   "loss": 0.014644961975055359,
+   "lr": 3.944926900490452e-05,
+   "grad_norm": 0.2742427885532379,
+   "skipped": false,
+   "seconds": 9071.9
+  },
+  {
+   "step": 298,
+   "epoch": 1,
+   "loss": 0.016398670697453047,
+   "lr": 3.913929639497462e-05,
+   "grad_norm": 0.28872784972190857,
+   "skipped": false,
+   "seconds": 9104.5
+  },
+  {
+   "step": 299,
+   "epoch": 1,
+   "loss": 0.017263746750174284,
+   "lr": 3.882976125260229e-05,
+   "grad_norm": 0.41232573986053467,
+   "skipped": false,
+   "seconds": 9133.6
+  },
+  {
+   "step": 300,
+   "epoch": 1,
+   "loss": 0.019409321586866885,
+   "lr": 3.852067604581794e-05,
+   "grad_norm": 0.36655980348587036,
+   "skipped": false,
+   "seconds": 9165.4
+  },
+  {
+   "step": 301,
+   "epoch": 1,
+   "loss": 0.013817376942734192,
+   "lr": 3.821205322452863e-05,
+   "grad_norm": 0.4269910156726837,
+   "skipped": false,
+   "seconds": 9197.3
+  },
+  {
+   "step": 302,
+   "epoch": 1,
+   "loss": 0.02174159104628318,
+   "lr": 3.790390522001662e-05,
+   "grad_norm": 0.3511791229248047,
+   "skipped": false,
+   "seconds": 9226.2
+  },
+  {
+   "step": 303,
+   "epoch": 1,
+   "loss": 0.016236488907425493,
+   "lr": 3.759624444443858e-05,
+   "grad_norm": 0.23098163306713104,
+   "skipped": false,
+   "seconds": 9255.1
+  },
+  {
+   "step": 304,
+   "epoch": 1,
+   "loss": 0.017032122276003547,
+   "lr": 3.728908329032567e-05,
+   "grad_norm": 0.3353688716888428,
+   "skipped": false,
+   "seconds": 9285.3
+  },
+  {
+   "step": 305,
+   "epoch": 1,
+   "loss": 0.01921281299615016,
+   "lr": 3.69824341300844e-05,
+   "grad_norm": 0.2812039256095886,
+   "skipped": false,
+   "seconds": 9317.4
+  },
+  {
+   "step": 306,
+   "epoch": 1,
+   "loss": 0.016250054041544598,
+   "lr": 3.6676309315498256e-05,
+   "grad_norm": 0.28371351957321167,
+   "skipped": false,
+   "seconds": 9346.9
+  },
+  {
+   "step": 307,
+   "epoch": 1,
+   "loss": 0.017268654154290946,
+   "lr": 3.6370721177230116e-05,
+   "grad_norm": 0.30808404088020325,
+   "skipped": false,
+   "seconds": 9377.3
+  },
+  {
+   "step": 308,
+   "epoch": 1,
+   "loss": 0.018877130984260353,
+   "lr": 3.606568202432562e-05,
+   "grad_norm": 0.2958217263221741,
+   "skipped": false,
+   "seconds": 9408.6
+  },
+  {
+   "step": 309,
+   "epoch": 1,
+   "loss": 0.018510348657640053,
+   "lr": 3.5761204143717385e-05,
+   "grad_norm": 0.2983313500881195,
+   "skipped": false,
+   "seconds": 9438.1
+  },
+  {
+   "step": 310,
+   "epoch": 1,
+   "loss": 0.027753260579261186,
+   "lr": 3.545729979973005e-05,
+   "grad_norm": 0.49328893423080444,
+   "skipped": false,
+   "seconds": 9468.4
+  },
+  {
+   "step": 311,
+   "epoch": 1,
+   "loss": 0.021194045424342375,
+   "lr": 3.515398123358627e-05,
+   "grad_norm": 0.29437315464019775,
+   "skipped": false,
+   "seconds": 9498.7
+  },
+  {
+   "step": 312,
+   "epoch": 1,
+   "loss": 0.01686973282576468,
+   "lr": 3.485126066291364e-05,
+   "grad_norm": 0.24846428632736206,
+   "skipped": false,
+   "seconds": 9529.9
+  },
+  {
+   "step": 313,
+   "epoch": 1,
+   "loss": 0.020775053623928477,
+   "lr": 3.4549150281252636e-05,
+   "grad_norm": 0.32179757952690125,
+   "skipped": false,
+   "seconds": 9561.7
+  },
+  {
+   "step": 314,
+   "epoch": 1,
+   "loss": 0.01746008980630427,
+   "lr": 3.424766225756537e-05,
+   "grad_norm": 0.2708369493484497,
+   "skipped": false,
+   "seconds": 9592.7
+  },
+  {
+   "step": 315,
+   "epoch": 1,
+   "loss": 0.02235273347283802,
+   "lr": 3.394680873574546e-05,
+   "grad_norm": 0.30049774050712585,
+   "skipped": false,
+   "seconds": 9623.0
+  },
+  {
+   "step": 316,
+   "epoch": 1,
+   "loss": 0.020062926840855224,
+   "lr": 3.364660183412892e-05,
+   "grad_norm": 0.27229195833206177,
+   "skipped": false,
+   "seconds": 9654.1
+  },
+  {
+   "step": 317,
+   "epoch": 1,
+   "loss": 0.017258836476076313,
+   "lr": 3.334705364500596e-05,
+   "grad_norm": 0.26787665486335754,
+   "skipped": false,
+   "seconds": 9684.8
+  },
+  {
+   "step": 318,
+   "epoch": 1,
+   "loss": 0.01641855517474037,
+   "lr": 3.304817623413397e-05,
+   "grad_norm": 0.2533664405345917,
+   "skipped": false,
+   "seconds": 9715.7
+  },
+  {
+   "step": 319,
+   "epoch": 1,
+   "loss": 0.019590947524596114,
+   "lr": 3.274998164025148e-05,
+   "grad_norm": 0.2662172317504883,
+   "skipped": false,
+   "seconds": 9746.9
+  },
+  {
+   "step": 320,
+   "epoch": 1,
+   "loss": 0.015173372780071455,
+   "lr": 3.2452481874593234e-05,
+   "grad_norm": 0.2734486162662506,
+   "skipped": false,
+   "seconds": 9777.2
+  },
+  {
+   "step": 321,
+   "epoch": 1,
+   "loss": 0.01690472215391683,
+   "lr": 3.215568892040641e-05,
+   "grad_norm": 0.2521533966064453,
+   "skipped": false,
+   "seconds": 9806.3
+  },
+  {
+   "step": 322,
+   "epoch": 1,
+   "loss": 0.02087894162760005,
+   "lr": 3.1859614732467954e-05,
+   "grad_norm": 0.31441473960876465,
+   "skipped": false,
+   "seconds": 9836.5
+  },
+  {
+   "step": 323,
+   "epoch": 1,
+   "loss": 0.013367866929058109,
+   "lr": 3.156427123660297e-05,
+   "grad_norm": 0.2140335887670517,
+   "skipped": false,
+   "seconds": 9867.4
+  },
+  {
+   "step": 324,
+   "epoch": 1,
+   "loss": 0.014857198143227932,
+   "lr": 3.12696703292044e-05,
+   "grad_norm": 0.25735756754875183,
+   "skipped": false,
+   "seconds": 9898.8
+  },
+  {
+   "step": 325,
+   "epoch": 1,
+   "loss": 0.017185792865311567,
+   "lr": 3.097582387675385e-05,
+   "grad_norm": 0.311068594455719,
+   "skipped": false,
+   "seconds": 9929.0
+  },
+  {
+   "step": 326,
+   "epoch": 1,
+   "loss": 0.02158736824714009,
+   "lr": 3.0682743715343564e-05,
+   "grad_norm": 0.40494808554649353,
+   "skipped": false,
+   "seconds": 9960.3
+  },
+  {
+   "step": 327,
+   "epoch": 1,
+   "loss": 0.017084697947514934,
+   "lr": 3.0390441650199724e-05,
+   "grad_norm": 0.40005937218666077,
+   "skipped": false,
+   "seconds": 9990.1
+  },
+  {
+   "step": 328,
+   "epoch": 1,
+   "loss": 0.018331635892169955,
+   "lr": 3.0098929455206904e-05,
+   "grad_norm": 0.29285380244255066,
+   "skipped": false,
+   "seconds": 10020.1
+  },
+  {
+   "step": 329,
+   "epoch": 1,
+   "loss": 0.01609054411061065,
+   "lr": 2.9808218872433767e-05,
+   "grad_norm": 0.3758953809738159,
+   "skipped": false,
+   "seconds": 10048.8
+  },
+  {
+   "step": 330,
+   "epoch": 1,
+   "loss": 0.022970987022469896,
+   "lr": 2.9518321611660237e-05,
+   "grad_norm": 0.3137178421020508,
+   "skipped": false,
+   "seconds": 10079.0
+  },
+  {
+   "step": 331,
+   "epoch": 1,
+   "loss": 0.017783200439661022,
+   "lr": 2.9229249349905684e-05,
+   "grad_norm": 0.29773151874542236,
+   "skipped": false,
+   "seconds": 10109.6
+  },
+  {
+   "step": 332,
+   "epoch": 1,
+   "loss": 0.016127112212563518,
+   "lr": 2.894101373095867e-05,
+   "grad_norm": 0.2919289171695709,
+   "skipped": false,
+   "seconds": 10137.5
+  },
+  {
+   "step": 333,
+   "epoch": 1,
+   "loss": 0.01922946286664721,
+   "lr": 2.8653626364907917e-05,
+   "grad_norm": 0.2525309920310974,
+   "skipped": false,
+   "seconds": 10168.2
+  },
+  {
+   "step": 334,
+   "epoch": 1,
+   "loss": 0.013304160924842517,
+   "lr": 2.8367098827674578e-05,
+   "grad_norm": 0.23909994959831238,
+   "skipped": false,
+   "seconds": 10199.6
+  },
+  {
+   "step": 335,
+   "epoch": 1,
+   "loss": 0.016354744904398068,
+   "lr": 2.8081442660546125e-05,
+   "grad_norm": 0.22283875942230225,
+   "skipped": false,
+   "seconds": 10228.6
+  },
+  {
+   "step": 336,
+   "epoch": 1,
+   "loss": 0.015653531660566027,
+   "lr": 2.7796669369711294e-05,
+   "grad_norm": 0.2464396357536316,
+   "skipped": false,
+   "seconds": 10258.2
+  },
+  {
+   "step": 337,
+   "epoch": 1,
+   "loss": 0.016884816086275165,
+   "lr": 2.7512790425796718e-05,
+   "grad_norm": 0.25359708070755005,
+   "skipped": false,
+   "seconds": 10288.8
+  },
+  {
+   "step": 338,
+   "epoch": 1,
+   "loss": 0.02054701702025414,
+   "lr": 2.7229817263404866e-05,
+   "grad_norm": 0.22650858759880066,
+   "skipped": false,
+   "seconds": 10320.9
+  },
+  {
+   "step": 339,
+   "epoch": 1,
+   "loss": 0.01796677281747861,
+   "lr": 2.694776128065345e-05,
+   "grad_norm": 0.3973900377750397,
+   "skipped": false,
+   "seconds": 10350.4
+  },
+  {
+   "step": 340,
+   "epoch": 1,
+   "loss": 0.01574439181162856,
+   "lr": 2.6666633838716314e-05,
+   "grad_norm": 0.2303837239742279,
+   "skipped": false,
+   "seconds": 10380.1
+  },
+  {
+   "step": 341,
+   "epoch": 1,
+   "loss": 0.013421567281087239,
+   "lr": 2.638644626136587e-05,
+   "grad_norm": 0.256214439868927,
+   "skipped": false,
+   "seconds": 10411.5
+  },
+  {
+   "step": 342,
+   "epoch": 1,
+   "loss": 0.01560402047647239,
+   "lr": 2.6107209834516854e-05,
+   "grad_norm": 0.2719871401786804,
+   "skipped": false,
+   "seconds": 10444.1
+  },
+  {
+   "step": 343,
+   "epoch": 1,
+   "loss": 0.02437654976241827,
+   "lr": 2.5828935805771802e-05,
+   "grad_norm": 0.5778554081916809,
+   "skipped": false,
+   "seconds": 10474.6
+  },
+  {
+   "step": 344,
+   "epoch": 1,
+   "loss": 0.025893028253831666,
+   "lr": 2.5551635383968065e-05,
+   "grad_norm": 1.2740740776062012,
+   "skipped": false,
+   "seconds": 10505.2
+  },
+  {
+   "step": 345,
+   "epoch": 1,
+   "loss": 0.01596813173179167,
+   "lr": 2.5275319738726165e-05,
+   "grad_norm": 0.3198592960834503,
+   "skipped": false,
+   "seconds": 10535.3
+  },
+  {
+   "step": 346,
+   "epoch": 1,
+   "loss": 0.02277532178302144,
+   "lr": 2.500000000000001e-05,
+   "grad_norm": 0.36020344495773315,
+   "skipped": false,
+   "seconds": 10565.9
+  },
+  {
+   "step": 347,
+   "epoch": 1,
+   "loss": 0.03668016211487209,
+   "lr": 2.4725687257628534e-05,
+   "grad_norm": 0.3347547948360443,
+   "skipped": false,
+   "seconds": 10596.5
+  },
+  {
+   "step": 348,
+   "epoch": 1,
+   "loss": 0.018362691837820064,
+   "lr": 2.4452392560888976e-05,
+   "grad_norm": 0.5395991802215576,
+   "skipped": false,
+   "seconds": 10627.8
+  },
+  {
+   "step": 349,
+   "epoch": 1,
+   "loss": 0.021723086045103975,
+   "lr": 2.418012691805191e-05,
+   "grad_norm": 0.30525851249694824,
+   "skipped": false,
+   "seconds": 10655.7
+  },
+  {
+   "step": 350,
+   "epoch": 1,
+   "loss": 0.013355130647918101,
+   "lr": 2.3908901295937713e-05,
+   "grad_norm": 0.2056380659341812,
+   "skipped": false,
+   "seconds": 10687.6
+  },
+  {
+   "step": 351,
+   "epoch": 1,
+   "loss": 0.01658166001536888,
+   "lr": 2.363872661947488e-05,
+   "grad_norm": 0.29792430996894836,
+   "skipped": false,
+   "seconds": 10717.7
+  },
+  {
+   "step": 352,
+   "epoch": 1,
+   "loss": 0.01986786977065489,
+   "lr": 2.336961377126001e-05,
+   "grad_norm": 0.33303678035736084,
+   "skipped": false,
+   "seconds": 10749.8
+  },
+  {
+   "step": 353,
+   "epoch": 1,
+   "loss": 0.019453472974544632,
+   "lr": 2.310157359111938e-05,
+   "grad_norm": 0.35702478885650635,
+   "skipped": false,
+   "seconds": 10779.0
+  },
+  {
+   "step": 354,
+   "epoch": 1,
+   "loss": 0.015191067429245598,
+   "lr": 2.283461687567236e-05,
+   "grad_norm": 0.2647939622402191,
+   "skipped": false,
+   "seconds": 10808.7
+  },
+  {
+   "step": 355,
+   "epoch": 1,
+   "loss": 0.015050849344550908,
+   "lr": 2.2568754377896516e-05,
+   "grad_norm": 0.2411239743232727,
+   "skipped": false,
+   "seconds": 10840.6
+  },
+  {
+   "step": 356,
+   "epoch": 1,
+   "loss": 0.012920646420356383,
+   "lr": 2.2303996806694488e-05,
+   "grad_norm": 0.22789867222309113,
+   "skipped": false,
+   "seconds": 10871.0
+  },
+  {
+   "step": 357,
+   "epoch": 1,
+   "loss": 0.014633365573149149,
+   "lr": 2.2040354826462668e-05,
+   "grad_norm": 0.2780497670173645,
+   "skipped": false,
+   "seconds": 10902.1
+  },
+  {
+   "step": 358,
+   "epoch": 1,
+   "loss": 0.015241143214731,
+   "lr": 2.1777839056661554e-05,
+   "grad_norm": 0.2793889045715332,
+   "skipped": false,
+   "seconds": 10933.8
+  },
+  {
+   "step": 359,
+   "epoch": 1,
+   "loss": 0.017730983503526435,
+   "lr": 2.1516460071388062e-05,
+   "grad_norm": 0.2675914466381073,
+   "skipped": false,
+   "seconds": 10964.3
+  },
+  {
+   "step": 360,
+   "epoch": 1,
+   "loss": 0.02158524816624618,
+   "lr": 2.125622839894964e-05,
+   "grad_norm": 0.29603734612464905,
+   "skipped": false,
+   "seconds": 10995.1
+  },
+  {
+   "step": 361,
+   "epoch": 1,
+   "loss": 0.016972364705656273,
+   "lr": 2.09971545214401e-05,
+   "grad_norm": 0.3738466203212738,
+   "skipped": false,
+   "seconds": 11026.7
+  },
+  {
+   "step": 362,
+   "epoch": 1,
+   "loss": 0.021522874766122458,
+   "lr": 2.073924887431744e-05,
+   "grad_norm": 0.31328442692756653,
+   "skipped": false,
+   "seconds": 11058.1
+  },
+  {
+   "step": 363,
+   "epoch": 1,
+   "loss": 0.017776748608779275,
+   "lr": 2.048252184598352e-05,
+   "grad_norm": 0.2673688232898712,
+   "skipped": false,
+   "seconds": 11089.1
+  },
+  {
+   "step": 364,
+   "epoch": 1,
+   "loss": 0.020682575531182877,
+   "lr": 2.0226983777365604e-05,
+   "grad_norm": 0.3776637017726898,
+   "skipped": false,
+   "seconds": 11120.0
+  },
+  {
+   "step": 365,
+   "epoch": 1,
+   "loss": 0.01417801892561934,
+   "lr": 1.9972644961499854e-05,
+   "grad_norm": 0.20646420121192932,
+   "skipped": false,
+   "seconds": 11150.7
+  },
+  {
+   "step": 366,
+   "epoch": 1,
+   "loss": 0.019781235268653725,
+   "lr": 1.9719515643116674e-05,
+   "grad_norm": 0.31171584129333496,
+   "skipped": false,
+   "seconds": 11181.4
+  },
+  {
+   "step": 367,
+   "epoch": 1,
+   "loss": 0.016028509728283503,
+   "lr": 1.946760601822809e-05,
+   "grad_norm": 0.271849662065506,
+   "skipped": false,
+   "seconds": 11211.2
+  },
+  {
+   "step": 368,
+   "epoch": 1,
+   "loss": 0.01597410992545235,
+   "lr": 1.9216926233717085e-05,
+   "grad_norm": 0.26115888357162476,
+   "skipped": false,
+   "seconds": 11241.2
+  },
+  {
+   "step": 369,
+   "epoch": 1,
+   "loss": 0.01464600051658741,
+   "lr": 1.8967486386928817e-05,
+   "grad_norm": 0.20717978477478027,
+   "skipped": false,
+   "seconds": 11271.8
+  },
+  {
+   "step": 370,
+   "epoch": 1,
+   "loss": 0.014690513336154173,
+   "lr": 1.8719296525263922e-05,
+   "grad_norm": 0.2489292472600937,
+   "skipped": false,
+   "seconds": 11300.7
+  },
+  {
+   "step": 371,
+   "epoch": 1,
+   "loss": 0.019532598908972135,
+   "lr": 1.847236664577389e-05,
+   "grad_norm": 0.2948475182056427,
+   "skipped": false,
+   "seconds": 11332.0
+  },
+  {
+   "step": 372,
+   "epoch": 1,
+   "loss": 0.019013768308514138,
+   "lr": 1.8226706694758195e-05,
+   "grad_norm": 0.24077652394771576,
+   "skipped": false,
+   "seconds": 11363.0
+  },
+  {
+   "step": 373,
+   "epoch": 1,
+   "loss": 0.015933554452930872,
+   "lr": 1.7982326567363888e-05,
+   "grad_norm": 0.287951797246933,
+   "skipped": false,
+   "seconds": 11394.9
+  },
+  {
+   "step": 374,
+   "epoch": 1,
+   "loss": 0.01682713304488164,
+   "lr": 1.773923610718686e-05,
+   "grad_norm": 0.22983764111995697,
+   "skipped": false,
+   "seconds": 11424.4
+  },
+  {
+   "step": 375,
+   "epoch": 1,
+   "loss": 0.017295656077654046,
+   "lr": 1.7497445105875377e-05,
+   "grad_norm": 0.2799907922744751,
+   "skipped": false,
+   "seconds": 11455.1
+  },
+  {
+   "step": 376,
+   "epoch": 1,
+   "loss": 0.014463379727276486,
+   "lr": 1.725696330273575e-05,
+   "grad_norm": 0.2402082234621048,
+   "skipped": false,
+   "seconds": 11483.6
+  },
+  {
+   "step": 377,
+   "epoch": 1,
+   "loss": 0.01813724382882879,
+   "lr": 1.7017800384339928e-05,
+   "grad_norm": 0.39935919642448425,
+   "skipped": false,
+   "seconds": 11513.5
+  },
+  {
+   "step": 378,
+   "epoch": 1,
+   "loss": 0.02063750761723156,
+   "lr": 1.6779965984135377e-05,
+   "grad_norm": 0.26636335253715515,
+   "skipped": false,
+   "seconds": 11542.8
+  },
+  {
+   "step": 379,
+   "epoch": 1,
+   "loss": 0.01720981912825147,
+   "lr": 1.6543469682057106e-05,
+   "grad_norm": 0.25584572553634644,
+   "skipped": false,
+   "seconds": 11573.4
+  },
+  {
+   "step": 380,
+   "epoch": 1,
+   "loss": 0.016045881793675845,
+   "lr": 1.6308321004141607e-05,
+   "grad_norm": 0.24034883081912994,
+   "skipped": false,
+   "seconds": 11604.8
+  },
+  {
+   "step": 381,
+   "epoch": 1,
+   "loss": 0.016837453895027516,
+   "lr": 1.60745294221434e-05,
+   "grad_norm": 0.27986347675323486,
+   "skipped": false,
+   "seconds": 11636.5
+  },
+  {
+   "step": 382,
+   "epoch": 1,
+   "loss": 0.020221882823811536,
+   "lr": 1.5842104353153287e-05,
+   "grad_norm": 27.728328704833984,
+   "skipped": false,
+   "seconds": 11665.6
+  },
+  {
+   "step": 383,
+   "epoch": 1,
+   "loss": 0.01467958791457942,
+   "lr": 1.5611055159219152e-05,
+   "grad_norm": 0.27296140789985657,
+   "skipped": false,
+   "seconds": 11696.4
+  },
+  {
+   "step": 384,
+   "epoch": 1,
+   "loss": 0.014055224498078696,
+   "lr": 1.5381391146968866e-05,
+   "grad_norm": 0.22324781119823456,
+   "skipped": false,
+   "seconds": 11727.0
+  },
+  {
+   "step": 385,
+   "epoch": 1,
+   "loss": 0.016929681425198358,
+   "lr": 1.5153121567235335e-05,
+   "grad_norm": 0.2702989876270294,
+   "skipped": false,
+   "seconds": 11757.3
+  },
+  {
+   "step": 386,
+   "epoch": 1,
+   "loss": 0.015159464784758233,
+   "lr": 1.4926255614683932e-05,
+   "grad_norm": 0.2476302534341812,
+   "skipped": false,
+   "seconds": 11788.4
+  },
+  {
+   "step": 387,
+   "epoch": 1,
+   "loss": 0.018204662800171433,
+   "lr": 1.4700802427442179e-05,
+   "grad_norm": 1.1057723760604858,
+   "skipped": false,
+   "seconds": 11819.5
+  },
+  {
+   "step": 388,
+   "epoch": 1,
+   "loss": 0.01382485911035508,
+   "lr": 1.4476771086731567e-05,
+   "grad_norm": 0.2819043695926666,
+   "skipped": false,
+   "seconds": 11848.7
+  },
+  {
+   "step": 389,
+   "epoch": 1,
+   "loss": 0.013176220352306017,
+   "lr": 1.4254170616501827e-05,
+   "grad_norm": 0.2870631217956543,
+   "skipped": false,
+   "seconds": 11879.2
+  },
+  {
+   "step": 390,
+   "epoch": 1,
+   "loss": 0.015560504752145687,
+   "lr": 1.4033009983067452e-05,
+   "grad_norm": 0.2705937325954437,
+   "skipped": false,
+   "seconds": 11909.6
+  },
+  {
+   "step": 391,
+   "epoch": 1,
+   "loss": 0.0172401643150863,
+   "lr": 1.3813298094746491e-05,
+   "grad_norm": 0.3437158465385437,
+   "skipped": false,
+   "seconds": 11940.3
+  },
+  {
+   "step": 392,
+   "epoch": 1,
+   "loss": 0.021647427999498762,
+   "lr": 1.3595043801501794e-05,
+   "grad_norm": 0.2943425476551056,
+   "skipped": false,
+   "seconds": 11970.0
+  },
+  {
+   "step": 393,
+   "epoch": 1,
+   "loss": 0.02094827180419021,
+   "lr": 1.3378255894584463e-05,
+   "grad_norm": 0.2669637203216553,
+   "skipped": false,
+   "seconds": 11999.7
+  },
+  {
+   "step": 394,
+   "epoch": 1,
+   "loss": 0.016528638708385404,
+   "lr": 1.3162943106179749e-05,
+   "grad_norm": 0.29634103178977966,
+   "skipped": false,
+   "seconds": 12031.1
+  },
+  {
+   "step": 395,
+   "epoch": 1,
+   "loss": 0.013307068416540571,
+   "lr": 1.2949114109055415e-05,
+   "grad_norm": 0.2215854972600937,
+   "skipped": false,
+   "seconds": 12060.2
+  },
+  {
+   "step": 396,
+   "epoch": 1,
+   "loss": 0.015899094069978637,
+   "lr": 1.2736777516212266e-05,
+   "grad_norm": 0.2678273022174835,
+   "skipped": false,
+   "seconds": 12089.3
+  },
+  {
+   "step": 397,
+   "epoch": 1,
+   "loss": 0.011314437034100134,
+   "lr": 1.2525941880537307e-05,
+   "grad_norm": 0.2707388997077942,
+   "skipped": false,
+   "seconds": 12119.0
+  },
+  {
+   "step": 398,
+   "epoch": 1,
+   "loss": 0.01665331870056188,
+   "lr": 1.2316615694459189e-05,
+   "grad_norm": 0.4135851263999939,
+   "skipped": false,
+   "seconds": 12149.1
+  },
+  {
+   "step": 399,
+   "epoch": 1,
+   "loss": 0.022952738952416948,
+   "lr": 1.2108807389606158e-05,
+   "grad_norm": 0.3070237338542938,
+   "skipped": false,
+   "seconds": 12179.5
+  },
+  {
+   "step": 400,
+   "epoch": 1,
+   "loss": 0.017071752871094466,
+   "lr": 1.1902525336466464e-05,
+   "grad_norm": 0.23360101878643036,
+   "skipped": false,
+   "seconds": 12209.9
+  },
+  {
+   "step": 401,
+   "epoch": 1,
+   "loss": 0.02012303668381288,
+   "lr": 1.1697777844051105e-05,
+   "grad_norm": 0.5425054430961609,
+   "skipped": false,
+   "seconds": 12240.2
+  },
+  {
+   "step": 402,
+   "epoch": 1,
+   "loss": 0.013822645244752747,
+   "lr": 1.1494573159559213e-05,
+   "grad_norm": 0.24790161848068237,
+   "skipped": false,
+   "seconds": 12269.8
+  },
+  {
+   "step": 403,
+   "epoch": 1,
+   "loss": 0.01809105012161628,
+   "lr": 1.1292919468045877e-05,
+   "grad_norm": 0.2774483263492584,
+   "skipped": false,
+   "seconds": 12299.4
+  },
+  {
+   "step": 404,
+   "epoch": 1,
+   "loss": 0.014330840732265222,
+   "lr": 1.1092824892092373e-05,
+   "grad_norm": 0.30954477190971375,
+   "skipped": false,
+   "seconds": 12330.3
+  },
+  {
+   "step": 405,
+   "epoch": 1,
+   "loss": 0.018357951065589642,
+   "lr": 1.0894297491479045e-05,
+   "grad_norm": 1.0409938097000122,
+   "skipped": false,
+   "seconds": 12360.4
+  },
+  {
+   "step": 406,
+   "epoch": 1,
+   "loss": 0.016694805807263444,
+   "lr": 1.0697345262860636e-05,
+   "grad_norm": 0.3196686804294586,
+   "skipped": false,
+   "seconds": 12390.2
+  },
+  {
+   "step": 407,
+   "epoch": 1,
+   "loss": 0.017687978299173884,
+   "lr": 1.0501976139444191e-05,
+   "grad_norm": 0.27741143107414246,
+   "skipped": false,
+   "seconds": 12420.6
+  },
+  {
+   "step": 408,
+   "epoch": 1,
+   "loss": 0.02315771045021266,
+   "lr": 1.0308197990669538e-05,
+   "grad_norm": 0.33277827501296997,
+   "skipped": false,
+   "seconds": 12450.5
+  },
+  {
+   "step": 409,
+   "epoch": 1,
+   "loss": 0.014765257162961296,
+   "lr": 1.0116018621892237e-05,
+   "grad_norm": 0.26430729031562805,
+   "skipped": false,
+   "seconds": 12481.0
+  },
+  {
+   "step": 410,
+   "epoch": 1,
+   "loss": 0.021278894573238723,
+   "lr": 9.92544577406923e-06,
+   "grad_norm": 0.2686549127101898,
+   "skipped": false,
+   "seconds": 12511.7
+  },
+  {
+   "step": 411,
+   "epoch": 1,
+   "loss": 0.014673768498664539,
+   "lr": 9.73648712344707e-06,
+   "grad_norm": 0.22932101786136627,
+   "skipped": false,
+   "seconds": 12540.4
+  },
+  {
+   "step": 412,
+   "epoch": 1,
+   "loss": 0.012732859466220641,
+   "lr": 9.549150281252633e-06,
+   "grad_norm": 0.21991147100925446,
+   "skipped": false,
+   "seconds": 12571.3
+  },
+  {
+   "step": 413,
+   "epoch": 1,
+   "loss": 0.01871246826772906,
+   "lr": 9.363442793386606e-06,
+   "grad_norm": 0.24232839047908783,
+   "skipped": false,
+   "seconds": 12603.8
+  },
+  {
+   "step": 414,
+   "epoch": 1,
+   "loss": 0.013039972662259746,
+   "lr": 9.179372140119525e-06,
+   "grad_norm": 0.21381565928459167,
+   "skipped": false,
+   "seconds": 12636.0
+  },
+  {
+   "step": 415,
+   "epoch": 1,
+   "loss": 0.01748069505959961,
+   "lr": 8.996945735790447e-06,
+   "grad_norm": 0.2568955421447754,
+   "skipped": false,
+   "seconds": 12664.7
+  },
+  {
+   "step": 416,
+   "epoch": 1,
+   "loss": 0.014440955766563963,
+   "lr": 8.816170928508365e-06,
+   "grad_norm": 0.2171037197113037,
+   "skipped": false,
+   "seconds": 12695.5
+  },
+  {
+   "step": 417,
+   "epoch": 1,
+   "loss": 0.013840633916014087,
+   "lr": 8.637054999856148e-06,
+   "grad_norm": 0.265576034784317,
+   "skipped": false,
+   "seconds": 12723.2
+  },
+  {
+   "step": 418,
+   "epoch": 1,
+   "loss": 0.014396657270685956,
+   "lr": 8.459605164597267e-06,
+   "grad_norm": 0.26587891578674316,
+   "skipped": false,
+   "seconds": 12754.4
+  },
+  {
+   "step": 419,
+   "epoch": 1,
+   "loss": 0.01791663248030866,
+   "lr": 8.283828570385238e-06,
+   "grad_norm": 0.3830026388168335,
+   "skipped": false,
+   "seconds": 12785.5
+  },
+  {
+   "step": 420,
+   "epoch": 1,
+   "loss": 0.016166413976127142,
+   "lr": 8.109732297475635e-06,
+   "grad_norm": 0.26750504970550537,
+   "skipped": false,
+   "seconds": 12816.0
+  },
+  {
+   "step": 421,
+   "epoch": 1,
+   "loss": 0.012980075526736159,
+   "lr": 7.937323358440935e-06,
+   "grad_norm": 0.23061294853687286,
+   "skipped": false,
+   "seconds": 12847.4
+  },
+  {
+   "step": 422,
+   "epoch": 1,
+   "loss": 0.01326038453287219,
+   "lr": 7.766608697888095e-06,
+   "grad_norm": 0.2267688363790512,
+   "skipped": false,
+   "seconds": 12877.5
+  },
+  {
+   "step": 423,
+   "epoch": 1,
+   "loss": 0.01464304382652022,
+   "lr": 7.597595192178702e-06,
+   "grad_norm": 0.26651710271835327,
+   "skipped": false,
+   "seconds": 12906.8
+  },
+  {
+   "step": 424,
+   "epoch": 1,
+   "loss": 0.014991654364183522,
+   "lr": 7.430289649152156e-06,
+   "grad_norm": 3.3915822505950928,
+   "skipped": false,
+   "seconds": 12938.5
+  },
+  {
+   "step": 425,
+   "epoch": 1,
+   "loss": 0.01039733584913624,
+   "lr": 7.264698807851328e-06,
+   "grad_norm": 0.1901535987854004,
+   "skipped": false,
+   "seconds": 12967.8
+  },
+  {
+   "step": 426,
+   "epoch": 1,
+   "loss": 0.01814713527780788,
+   "lr": 7.100829338251147e-06,
+   "grad_norm": 0.3595639765262604,
+   "skipped": false,
+   "seconds": 12997.6
+  },
+  {
+   "step": 427,
+   "epoch": 1,
+   "loss": 0.017600897006762102,
+   "lr": 6.9386878409899715e-06,
+   "grad_norm": 0.26794108748435974,
+   "skipped": false,
+   "seconds": 13028.5
+  },
+  {
+   "step": 428,
+   "epoch": 1,
+   "loss": 0.013089698404604005,
+   "lr": 6.778280847103669e-06,
+   "grad_norm": 0.22527790069580078,
+   "skipped": false,
+   "seconds": 13058.3
+  },
+  {
+   "step": 429,
+   "epoch": 1,
+   "loss": 0.014918052926552072,
+   "lr": 6.619614817762537e-06,
+   "grad_norm": 0.23448950052261353,
+   "skipped": false,
+   "seconds": 13090.1
+  },
+  {
+   "step": 430,
+   "epoch": 1,
+   "loss": 0.01596330983663676,
+   "lr": 6.462696144011149e-06,
+   "grad_norm": 0.2590503692626953,
+   "skipped": false,
+   "seconds": 13121.9
+  },
+  {
+   "step": 431,
+   "epoch": 1,
+   "loss": 0.0151835978799376,
+   "lr": 6.3075311465107535e-06,
+   "grad_norm": 0.20222195982933044,
+   "skipped": false,
+   "seconds": 13152.0
+  },
+  {
+   "step": 432,
+   "epoch": 1,
+   "loss": 0.016157260446313156,
+   "lr": 6.154126075284855e-06,
+   "grad_norm": 0.26201319694519043,
+   "skipped": false,
+   "seconds": 13182.5
+  },
+  {
+   "step": 433,
+   "epoch": 1,
+   "loss": 0.017083022302860837,
+   "lr": 6.002487109467347e-06,
+   "grad_norm": 0.3449777364730835,
+   "skipped": false,
+   "seconds": 13212.7
+  },
+  {
+   "step": 434,
+   "epoch": 1,
+   "loss": 0.01500254542848523,
+   "lr": 5.852620357053651e-06,
+   "grad_norm": 0.2756975293159485,
+   "skipped": false,
+   "seconds": 13243.5
+  },
+  {
+   "step": 435,
+   "epoch": 1,
+   "loss": 0.02222851620663891,
+   "lr": 5.704531854654721e-06,
+   "grad_norm": 0.3507115840911865,
+   "skipped": false,
+   "seconds": 13273.1
+  },
+  {
+   "step": 436,
+   "epoch": 1,
+   "loss": 0.019562456892998556,
+   "lr": 5.558227567253832e-06,
+   "grad_norm": 0.49055540561676025,
+   "skipped": false,
+   "seconds": 13303.6
+  },
+  {
+   "step": 437,
+   "epoch": 1,
+   "loss": 0.016707392397116234,
+   "lr": 5.413713387966329e-06,
+   "grad_norm": 0.30712446570396423,
+   "skipped": false,
+   "seconds": 13331.3
+  },
+  {
+   "step": 438,
+   "epoch": 1,
+   "loss": 0.015860420129522206,
+   "lr": 5.270995137802315e-06,
+   "grad_norm": 0.2504942715167999,
+   "skipped": false,
+   "seconds": 13361.0
+  },
+  {
+   "step": 439,
+   "epoch": 1,
+   "loss": 0.017581775125712695,
+   "lr": 5.13007856543209e-06,
+   "grad_norm": 0.24654711782932281,
+   "skipped": false,
+   "seconds": 13393.6
+  },
+  {
+   "step": 440,
+   "epoch": 1,
+   "loss": 0.017193995594407954,
+   "lr": 4.99096934695461e-06,
+   "grad_norm": 0.2912333011627197,
+   "skipped": false,
+   "seconds": 13423.5
+  },
+  {
+   "step": 441,
+   "epoch": 1,
+   "loss": 0.016572819199672965,
+   "lr": 4.853673085668947e-06,
+   "grad_norm": 0.24467350542545319,
+   "skipped": false,
+   "seconds": 13452.7
+  },
+  {
+   "step": 442,
+   "epoch": 1,
+   "loss": 0.0161877468859019,
+   "lr": 4.7181953118484556e-06,
+   "grad_norm": 0.3681254982948303,
+   "skipped": false,
+   "seconds": 13482.4
+  },
+  {
+   "step": 443,
+   "epoch": 1,
+   "loss": 0.017729474168756372,
+   "lr": 4.58454148251814e-06,
+   "grad_norm": 0.25886455178260803,
+   "skipped": false,
+   "seconds": 13513.4
+  },
+  {
+   "step": 444,
+   "epoch": 1,
+   "loss": 0.017672039854985017,
+   "lr": 4.452716981234744e-06,
+   "grad_norm": 0.3544117510318756,
+   "skipped": false,
+   "seconds": 13543.3
+  },
+  {
+   "step": 445,
+   "epoch": 1,
+   "loss": 0.014902976864204671,
+   "lr": 4.322727117869951e-06,
+   "grad_norm": 0.2283497452735901,
+   "skipped": false,
+   "seconds": 13573.1
+  },
+  {
+   "step": 446,
+   "epoch": 1,
+   "loss": 0.017319539780133623,
+   "lr": 4.19457712839652e-06,
+   "grad_norm": 0.36809778213500977,
+   "skipped": false,
+   "seconds": 13602.4
+  },
+  {
+   "step": 447,
+   "epoch": 1,
+   "loss": 0.021381757385359655,
+   "lr": 4.068272174677335e-06,
+   "grad_norm": 0.3147641718387604,
+   "skipped": false,
+   "seconds": 13632.3
+  },
+  {
+   "step": 448,
+   "epoch": 1,
+   "loss": 0.01543112478605562,
+   "lr": 3.9438173442575e-06,
+   "grad_norm": 0.22462071478366852,
+   "skipped": false,
+   "seconds": 13662.9
+  },
+  {
+   "step": 449,
+   "epoch": 1,
+   "loss": 0.01753268618065157,
+   "lr": 3.821217650159453e-06,
+   "grad_norm": 0.22539229691028595,
+   "skipped": false,
+   "seconds": 13694.8
+  },
+  {
+   "step": 450,
+   "epoch": 1,
+   "loss": 0.01648473962055351,
+   "lr": 3.700478030680987e-06,
+   "grad_norm": 0.22293080389499664,
+   "skipped": false,
+   "seconds": 13725.2
+  },
+  {
+   "step": 451,
+   "epoch": 1,
+   "loss": 0.01837570977896805,
+   "lr": 3.581603349196372e-06,
+   "grad_norm": 0.24224314093589783,
+   "skipped": false,
+   "seconds": 13755.3
+  },
+  {
+   "step": 452,
+   "epoch": 1,
+   "loss": 0.011402834787101922,
+   "lr": 3.4645983939604496e-06,
+   "grad_norm": 0.21129898726940155,
+   "skipped": false,
+   "seconds": 13784.0
+  },
+  {
+   "step": 453,
+   "epoch": 1,
+   "loss": 0.013870194191312664,
+   "lr": 3.349467877915746e-06,
+   "grad_norm": 0.21010389924049377,
+   "skipped": false,
+   "seconds": 13813.6
+  },
+  {
+   "step": 454,
+   "epoch": 1,
+   "loss": 0.014568589164859322,
+   "lr": 3.2362164385026706e-06,
+   "grad_norm": 0.2321111261844635,
+   "skipped": false,
+   "seconds": 13845.4
+  },
+  {
+   "step": 455,
+   "epoch": 1,
+   "loss": 0.015580425932430933,
+   "lr": 3.1248486374726883e-06,
+   "grad_norm": 0.2638421356678009,
+   "skipped": false,
+   "seconds": 13877.1
+  },
+  {
+   "step": 456,
+   "epoch": 1,
+   "loss": 0.013652178491824374,
+   "lr": 3.0153689607045845e-06,
+   "grad_norm": 0.24279172718524933,
+   "skipped": false,
+   "seconds": 13907.7
+  },
+  {
+   "step": 457,
+   "epoch": 1,
+   "loss": 0.01109541549945247,
+   "lr": 2.9077818180237693e-06,
+   "grad_norm": 0.22842855751514435,
+   "skipped": false,
+   "seconds": 13939.2
+  },
+  {
+   "step": 458,
+   "epoch": 1,
+   "loss": 0.016054675009398334,
+   "lr": 2.802091543024671e-06,
+   "grad_norm": 0.3169231414794922,
+   "skipped": false,
+   "seconds": 13968.8
+  },
+  {
+   "step": 459,
+   "epoch": 1,
+   "loss": 0.021520703001592777,
+   "lr": 2.6983023928961404e-06,
+   "grad_norm": 0.4069320857524872,
+   "skipped": false,
+   "seconds": 14000.2
+  },
+  {
+   "step": 460,
+   "epoch": 1,
+   "loss": 0.013510759170049744,
+   "lr": 2.596418548250029e-06,
+   "grad_norm": 0.21013274788856506,
+   "skipped": false,
+   "seconds": 14031.5
+  },
+  {
+   "step": 461,
+   "epoch": 1,
+   "loss": 0.013634231950015235,
+   "lr": 2.496444112952734e-06,
+   "grad_norm": 0.29028260707855225,
+   "skipped": false,
+   "seconds": 14063.1
+  },
+  {
+   "step": 462,
+   "epoch": 1,
+   "loss": 0.01630439492373671,
+   "lr": 2.3983831139599287e-06,
+   "grad_norm": 0.23812183737754822,
+   "skipped": false,
+   "seconds": 14094.4
+  },
+  {
+   "step": 463,
+   "epoch": 1,
+   "loss": 0.015388567299092912,
+   "lr": 2.3022395011543686e-06,
+   "grad_norm": 0.22293654084205627,
+   "skipped": false,
+   "seconds": 14123.4
+  },
+  {
+   "step": 464,
+   "epoch": 1,
+   "loss": 0.015927969555081974,
+   "lr": 2.208017147186736e-06,
+   "grad_norm": 0.28787466883659363,
+   "skipped": false,
+   "seconds": 14154.5
+  },
+  {
+   "step": 465,
+   "epoch": 1,
+   "loss": 0.014266629292836635,
+   "lr": 2.1157198473197414e-06,
+   "grad_norm": 0.23060643672943115,
+   "skipped": false,
+   "seconds": 14185.9
+  },
+  {
+   "step": 466,
+   "epoch": 1,
+   "loss": 0.019002444870969852,
+   "lr": 2.0253513192751373e-06,
+   "grad_norm": 0.2352025806903839,
+   "skipped": false,
+   "seconds": 14216.0
+  },
+  {
+   "step": 467,
+   "epoch": 1,
+   "loss": 0.016235436245141923,
+   "lr": 1.9369152030840556e-06,
+   "grad_norm": 0.3779015839099884,
+   "skipped": false,
+   "seconds": 14246.5
+  },
+  {
+   "step": 468,
+   "epoch": 1,
+   "loss": 0.02078640302022298,
+   "lr": 1.8504150609403858e-06,
+   "grad_norm": 0.3310161232948303,
+   "skipped": false,
+   "seconds": 14277.5
+  },
+  {
+   "step": 469,
+   "epoch": 1,
+   "loss": 0.012792520942559773,
+   "lr": 1.7658543770572189e-06,
+   "grad_norm": 0.21761828660964966,
+   "skipped": false,
+   "seconds": 14307.9
+  },
+  {
+   "step": 470,
+   "epoch": 1,
+   "loss": 0.016218211290812422,
+   "lr": 1.6832365575265741e-06,
+   "grad_norm": 0.20596861839294434,
+   "skipped": false,
+   "seconds": 14338.0
+  },
+  {
+   "step": 471,
+   "epoch": 1,
+   "loss": 0.019815962166179115,
+   "lr": 1.6025649301821876e-06,
+   "grad_norm": 0.28267160058021545,
+   "skipped": false,
+   "seconds": 14369.4
+  },
+  {
+   "step": 472,
+   "epoch": 1,
+   "loss": 0.013912331480348034,
+   "lr": 1.523842744465437e-06,
+   "grad_norm": 0.23686163127422333,
+   "skipped": false,
+   "seconds": 14399.0
+  },
+  {
+   "step": 473,
+   "epoch": 1,
+   "loss": 0.013959775194577858,
+   "lr": 1.4470731712944884e-06,
+   "grad_norm": 0.2434159815311432,
+   "skipped": false,
+   "seconds": 14429.9
+  },
+  {
+   "step": 474,
+   "epoch": 1,
+   "loss": 0.01357290552109755,
+   "lr": 1.372259302936546e-06,
+   "grad_norm": 0.2746292054653168,
+   "skipped": false,
+   "seconds": 14461.3
+  },
+  {
+   "step": 475,
+   "epoch": 1,
+   "loss": 0.014042890674668011,
+   "lr": 1.2994041528833266e-06,
+   "grad_norm": 0.2632994055747986,
+   "skipped": false,
+   "seconds": 14492.8
+  },
+  {
+   "step": 476,
+   "epoch": 1,
+   "loss": 0.013474665587544012,
+   "lr": 1.2285106557296477e-06,
+   "grad_norm": 0.2416546642780304,
+   "skipped": false,
+   "seconds": 14523.3
+  },
+  {
+   "step": 477,
+   "epoch": 1,
+   "loss": 0.01666516979245374,
+   "lr": 1.1595816670552428e-06,
+   "grad_norm": 0.23275969922542572,
+   "skipped": false,
+   "seconds": 14552.4
+  },
+  {
+   "step": 478,
+   "epoch": 1,
+   "loss": 0.011933755860480118,
+   "lr": 1.0926199633097157e-06,
+   "grad_norm": 0.22021210193634033,
+   "skipped": false,
+   "seconds": 14582.8
+  },
+  {
+   "step": 479,
+   "epoch": 1,
+   "loss": 0.01824787427528745,
+   "lr": 1.02762824170074e-06,
+   "grad_norm": 0.2591199576854706,
+   "skipped": false,
+   "seconds": 14612.6
+  },
+  {
+   "step": 480,
+   "epoch": 1,
+   "loss": 0.01649331980774414,
+   "lr": 9.646091200853802e-07,
+   "grad_norm": 0.3194705545902252,
+   "skipped": false,
+   "seconds": 14645.0
+  },
+  {
+   "step": 481,
+   "epoch": 1,
+   "loss": 0.02036778127466428,
+   "lr": 9.035651368646648e-07,
+   "grad_norm": 0.43635284900665283,
+   "skipped": false,
+   "seconds": 14675.4
+  },
+  {
+   "step": 482,
+   "epoch": 1,
+   "loss": 0.016031105149885316,
+   "lr": 8.444987508813451e-07,
+   "grad_norm": 0.22120577096939087,
+   "skipped": false,
+   "seconds": 14706.1
+  },
+  {
+   "step": 483,
+   "epoch": 1,
+   "loss": 0.015851634324078795,
+   "lr": 7.874123413208145e-07,
+   "grad_norm": 0.3351321816444397,
+   "skipped": false,
+   "seconds": 14737.7
+  },
+  {
+   "step": 484,
+   "epoch": 1,
+   "loss": 0.017026997200027875,
+   "lr": 7.323082076153509e-07,
+   "grad_norm": 0.28247079253196716,
+   "skipped": false,
+   "seconds": 14768.3
+  },
+  {
+   "step": 485,
+   "epoch": 1,
+   "loss": 0.021605440543634866,
+   "lr": 6.791885693514133e-07,
+   "grad_norm": 0.27851662039756775,
+   "skipped": false,
+   "seconds": 14799.0
+  },
+  {
+   "step": 486,
+   "epoch": 1,
+   "loss": 0.017235946980944803,
+   "lr": 6.280555661802856e-07,
+   "grad_norm": 0.2936553657054901,
+   "skipped": false,
+   "seconds": 14829.3
+  },
+  {
+   "step": 487,
+   "epoch": 1,
+   "loss": 0.01512704289275211,
+   "lr": 5.78911257731879e-07,
+   "grad_norm": 0.21169595420360565,
+   "skipped": false,
+   "seconds": 14860.4
+  },
+  {
+   "step": 488,
+   "epoch": 1,
+   "loss": 0.016688139252046503,
+   "lr": 5.317576235317756e-07,
+   "grad_norm": 0.2827041745185852,
+   "skipped": false,
+   "seconds": 14891.6
+  },
+  {
+   "step": 489,
+   "epoch": 1,
+   "loss": 0.014433720455782647,
+   "lr": 4.865965629214819e-07,
+   "grad_norm": 0.24454833567142487,
+   "skipped": false,
+   "seconds": 14922.8
+  },
+  {
+   "step": 490,
+   "epoch": 1,
+   "loss": 0.01710877140949955,
+   "lr": 4.434298949819449e-07,
+   "grad_norm": 0.2576327621936798,
+   "skipped": false,
+   "seconds": 14954.1
+  },
+  {
+   "step": 491,
+   "epoch": 1,
+   "loss": 0.024297605621693766,
+   "lr": 4.02259358460233e-07,
+   "grad_norm": 0.46404898166656494,
+   "skipped": false,
+   "seconds": 14983.7
+  },
+  {
+   "step": 492,
+   "epoch": 1,
+   "loss": 0.01663580822664791,
+   "lr": 3.630866116995757e-07,
+   "grad_norm": 0.27184104919433594,
+   "skipped": false,
+   "seconds": 15014.6
+  },
+  {
+   "step": 493,
+   "epoch": 1,
+   "loss": 0.020265022200381948,
+   "lr": 3.2591323257248893e-07,
+   "grad_norm": 0.24386714398860931,
+   "skipped": false,
+   "seconds": 15045.4
+  },
+  {
+   "step": 494,
+   "epoch": 1,
+   "loss": 0.016951551928618313,
+   "lr": 2.907407184172706e-07,
+   "grad_norm": 0.27427002787590027,
+   "skipped": false,
+   "seconds": 15075.4
+  },
+  {
+   "step": 495,
+   "epoch": 1,
+   "loss": 0.016049816498071592,
+   "lr": 2.5757048597765396e-07,
+   "grad_norm": 0.22172033786773682,
+   "skipped": false,
+   "seconds": 15106.9
+  },
+  {
+   "step": 496,
+   "epoch": 1,
+   "loss": 0.01609070582001692,
+   "lr": 2.2640387134577058e-07,
+   "grad_norm": 0.2710008919239044,
+   "skipped": false,
+   "seconds": 15136.7
+  },
+  {
+   "step": 497,
+   "epoch": 1,
+   "loss": 0.015595577054981822,
+   "lr": 1.9724212990830938e-07,
+   "grad_norm": 0.25194889307022095,
+   "skipped": false,
+   "seconds": 15166.6
+  },
+  {
+   "step": 498,
+   "epoch": 1,
+   "loss": 0.011678507451703827,
+   "lr": 1.7008643629596866e-07,
+   "grad_norm": 2.243579626083374,
+   "skipped": false,
+   "seconds": 15197.3
+  },
+  {
+   "step": 499,
+   "epoch": 1,
+   "loss": 0.014588745646268718,
+   "lr": 1.449378843361271e-07,
+   "grad_norm": 0.2866562604904175,
+   "skipped": false,
+   "seconds": 15228.6
+  },
+  {
+   "step": 500,
+   "epoch": 1,
+   "loss": 0.013277890113232332,
+   "lr": 1.2179748700879012e-07,
+   "grad_norm": 0.2638900876045227,
+   "skipped": false,
+   "seconds": 15259.5
+  },
+  {
+   "step": 501,
+   "epoch": 1,
+   "loss": 0.013807131787219783,
+   "lr": 1.0066617640578368e-07,
+   "grad_norm": 0.23005108535289764,
+   "skipped": false,
+   "seconds": 15290.9
+  },
+  {
+   "step": 502,
+   "epoch": 1,
+   "loss": 0.02514329022550263,
+   "lr": 8.15448036932176e-08,
+   "grad_norm": 0.24505768716335297,
+   "skipped": false,
+   "seconds": 15323.3
+  },
+  {
+   "step": 503,
+   "epoch": 1,
+   "loss": 0.019522010527541444,
+   "lr": 6.443413907720186e-08,
+   "grad_norm": 0.3740701973438263,
+   "skipped": false,
+   "seconds": 15353.2
+  },
+  {
+   "step": 504,
+   "epoch": 1,
+   "loss": 0.015291050259849599,
+   "lr": 4.933487177280482e-08,
+   "grad_norm": 0.29071199893951416,
+   "skipped": false,
+   "seconds": 15383.7
+  },
+  {
+   "step": 505,
+   "epoch": 1,
+   "loss": 0.014405335253931761,
+   "lr": 3.624760997631982e-08,
+   "grad_norm": 0.2308592051267624,
+   "skipped": false,
+   "seconds": 15414.3
+  },
+  {
+   "step": 506,
+   "epoch": 1,
+   "loss": 0.012780588496198009,
+   "lr": 2.5172880840745873e-08,
+   "grad_norm": 0.25191834568977356,
+   "skipped": false,
+   "seconds": 15445.2
+  },
+  {
+   "step": 507,
+   "epoch": 1,
+   "loss": 0.01951780338693646,
+   "lr": 1.6111130454543598e-08,
+   "grad_norm": 0.47091758251190186,
+   "skipped": false,
+   "seconds": 15474.9
+  },
+  {
+   "step": 508,
+   "epoch": 1,
+   "loss": 0.014618728624938905,
+   "lr": 9.06272382371065e-09,
+   "grad_norm": 0.2414289265871048,
+   "skipped": false,
+   "seconds": 15506.0
+  },
+  {
+   "step": 509,
+   "epoch": 1,
+   "loss": 0.014318742624718845,
+   "lr": 4.0279448570323954e-09,
+   "grad_norm": 0.24327567219734192,
+   "skipped": false,
+   "seconds": 15534.2
+  },
+  {
+   "step": 510,
+   "epoch": 1,
+   "loss": 0.010492502323318565,
+   "lr": 1.0069963546743832e-09,
+   "grad_norm": 0.2578527629375458,
+   "skipped": false,
+   "seconds": 15551.6
+  }
+ ]
+}

training/train_lora.py

@@ -0,0 +1,205 @@
+import argparse
+import json
+import math
+import random
+import time
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from peft import LoraConfig, get_peft_model
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+from transformers import get_cosine_schedule_with_warmup
+
+from predict import LABELS, label_ids, load_model, render, thinking_text
+
+LOSS_CHUNK = 256
+WARMUP_FRACTION = 0.03
+MAX_GRAD_NORM = 1.0
+MAX_SKIPPED_IN_A_ROW = 10
+
+
+def target_modules(model):
+    layers = model.model.layers.named_modules(prefix="model.layers")
+    return [name for name, module in layers if isinstance(module, nn.Linear)]
+
+
+def encode_example(tokenizer, labels, row):
+    context = tokenizer.encode(render(tokenizer, row["prompt"], True), add_special_tokens=False)
+    text = thinking_text(tokenizer, row["prompt"], row["thought"]) + LABELS[row["answer"] - 1]
+    ids = tokenizer.encode(text, add_special_tokens=False)
+    if ids[: len(context)] != context or ids[-1] != labels[row["answer"] - 1]:
+        return None
+    return ids, len(context)
+
+
+def collate(batch, pad_id, device):
+    width = max(len(ids) for ids, _ in batch)
+    input_ids = torch.full((len(batch), width), pad_id)
+    mask = torch.zeros((len(batch), width), dtype=torch.long)
+    targets = torch.full((len(batch), width), -100)
+    for row, (ids, context) in enumerate(batch):
+        start = width - len(ids)
+        input_ids[row, start:] = torch.tensor(ids)
+        mask[row, start:] = 1
+        targets[row, start + context :] = torch.tensor(ids[context:])
+    return input_ids.to(device), mask.to(device), targets
+
+
+def chunk_loss(base, hidden, targets):
+    logits = base.lm_head(hidden).float()
+    cap = base.config.final_logit_softcapping
+    if cap is not None:
+        logits = torch.tanh(logits / cap) * cap
+    return F.cross_entropy(logits, targets.to(logits.device), reduction="sum")
+
+
+def loss_sum(model, input_ids, mask, targets):
+    base = model.get_base_model()
+    position_ids = (mask.cumsum(-1) - 1).clamp(min=0)
+    hidden = base.model(
+        input_ids=input_ids, attention_mask=mask, position_ids=position_ids, use_cache=False
+    ).last_hidden_state
+    shifted = targets[:, 1:].to(hidden.device)
+    select = shifted != -100
+    hidden, shifted = hidden[:, :-1][select], shifted[select]
+    chunks = range(0, len(shifted), LOSS_CHUNK)
+    return sum(
+        checkpoint(chunk_loss, base, hidden[i : i + LOSS_CHUNK], shifted[i : i + LOSS_CHUNK], use_reentrant=False)
+        for i in chunks
+    )
+
+
+def peak_memory():
+    if not torch.cuda.is_available():
+        return []
+    return [round(torch.cuda.max_memory_allocated(i) / 2**30, 2) for i in range(torch.cuda.device_count())]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", default="google/gemma-4-E4B-it")
+    parser.add_argument("--revision", help="model commit to load; the latest one when not given")
+    parser.add_argument("--data", default="star/star.jsonl")
+    parser.add_argument("--out_dir", default="adapter")
+    parser.add_argument("--epochs", type=int, default=2)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--grad_accum", type=int, default=4)
+    parser.add_argument("--max_length", type=int, default=1024)
+    parser.add_argument("--rank", type=int, default=16)
+    parser.add_argument("--alpha", type=int, default=32)
+    parser.add_argument("--dropout", type=float, default=0.05)
+    parser.add_argument("--seed", type=int, default=0)
+    args = parser.parse_args()
+
+    tokenizer, model = load_model(args.model, revision=args.revision)
+    base_params = sum(p.numel() for p in model.parameters())
+    labels = label_ids(tokenizer)
+    rows = [json.loads(line) for line in Path(args.data).read_text(encoding="utf-8").splitlines()]
+    encoded = [encode_example(tokenizer, labels, row) for row in rows]
+    misaligned = sum(e is None for e in encoded)
+    examples = [e for e in encoded if e is not None and len(e[0]) <= args.max_length]
+    too_long = len(encoded) - misaligned - len(examples)
+    print(f"Training on {len(examples)} of {len(rows)} examples: {too_long} longer than {args.max_length} tokens"
+          f" and {misaligned} that do not tokenize as context + target were dropped")
+    if not examples:
+        raise SystemExit("error: no training examples left")
+
+    dtype = model.dtype
+    torch.manual_seed(args.seed)
+    modules = target_modules(model)
+    config = LoraConfig(
+        r=args.rank, lora_alpha=args.alpha, lora_dropout=args.dropout, bias="none", target_modules=modules
+    )
+    model = get_peft_model(model, config)
+    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    model.enable_input_require_grads()
+    model.train()
+    params = [p for p in model.parameters() if p.requires_grad]
+    adapter_params = sum(p.numel() for p in params)
+    print(f"LoRA r={args.rank} on {len(modules)} linear modules: {adapter_params:,} adapter parameters,"
+          f" {base_params + adapter_params:,} in total")
+
+    longest = sorted(examples, key=lambda e: -len(e[0]))[: args.batch_size]
+    loss_sum(model, *collate(longest, tokenizer.pad_token_id, model.device)).backward()
+    model.zero_grad(set_to_none=True)
+    print(f"Longest micro-batch ({len(longest[0][0])} tokens) fits: peak GPU memory {peak_memory()} GB")
+
+    step_size = args.batch_size * args.grad_accum
+    total_steps = args.epochs * math.ceil(len(examples) / step_size)
+    optimizer = torch.optim.AdamW(params, lr=args.lr, weight_decay=0.0)
+    scheduler = get_cosine_schedule_with_warmup(optimizer, round(WARMUP_FRACTION * total_steps), total_steps)
+    scaler = torch.amp.GradScaler(model.device.type, init_scale=2.0**10, enabled=dtype == torch.float16)
+    rng = random.Random(args.seed)
+    steps, tokens, skipped_in_a_row = [], 0, 0
+    start_time = time.time()
+    for epoch in range(args.epochs):
+        order = examples[:]
+        rng.shuffle(order)
+        for start in range(0, len(order), step_size):
+            step_examples = sorted(order[start : start + step_size], key=lambda e: len(e[0]))
+            n_targets = sum(len(ids) - context for ids, context in step_examples)
+            step_loss = 0.0
+            for i in range(0, len(step_examples), args.batch_size):
+                micro = step_examples[i : i + args.batch_size]
+                input_ids, mask, targets = collate(micro, tokenizer.pad_token_id, model.device)
+                loss = loss_sum(model, input_ids, mask, targets)
+                if not math.isfinite(loss.item()):
+                    raise RuntimeError(f"non-finite loss {loss.item()} at step {len(steps) + 1}")
+                scaler.scale(loss / n_targets).backward()
+                step_loss += loss.item()
+                tokens += int(mask.sum())
+            scaler.unscale_(optimizer)
+            grad_norm = float(nn.utils.clip_grad_norm_(params, MAX_GRAD_NORM))
+            scale = scaler.get_scale()
+            scaler.step(optimizer)
+            scaler.update()
+            skipped = scaler.get_scale() < scale
+            if not math.isfinite(grad_norm) and not skipped:
+                raise RuntimeError(f"non-finite gradients at step {len(steps) + 1}")
+            skipped_in_a_row = skipped_in_a_row + 1 if skipped else 0
+            if skipped_in_a_row == MAX_SKIPPED_IN_A_ROW:
+                raise RuntimeError(f"{MAX_SKIPPED_IN_A_ROW} steps in a row overflowed in fp16")
+            row = {
+                "step": len(steps) + 1,
+                "epoch": epoch,
+                "loss": step_loss / n_targets,
+                "lr": scheduler.get_last_lr()[0],
+                "grad_norm": grad_norm,
+                "skipped": skipped,
+                "seconds": round(time.time() - start_time, 1),
+            }
+            scheduler.step()
+            optimizer.zero_grad(set_to_none=True)
+            steps.append(row)
+            print(json.dumps(row), flush=True)
+    seconds = time.time() - start_time
+
+    out_dir = Path(args.out_dir)
+    model.save_pretrained(out_dir)
+    log = {
+        "config": vars(args),
+        "dtype": str(dtype),
+        "grad_scaling": scaler.is_enabled(),
+        "target_modules": len(modules),
+        "base_params": base_params,
+        "adapter_params": adapter_params,
+        "total_params": base_params + adapter_params,
+        "examples": len(rows),
+        "trained_examples": len(examples),
+        "dropped_too_long": too_long,
+        "dropped_misaligned": misaligned,
+        "skipped_steps": sum(row["skipped"] for row in steps),
+        "seconds": round(seconds, 1),
+        "tokens_per_s": round(tokens / seconds, 1),
+        "peak_memory_gb": peak_memory(),
+        "steps": steps,
+    }
+    (out_dir / "train_log.json").write_text(json.dumps(log, indent=1) + "\n", encoding="utf-8")
+    print(f"Saved the adapter to {out_dir}: final loss {steps[-1]['loss']:.4f}, {log['tokens_per_s']} tokens/s")
+
+
+if __name__ == "__main__":
+    main()