Senali D
Final submission: code, predictions, training
044f4e3 verified
|
Raw History Blame Contribute Delete
6.01 kB

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. 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/.

    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:

    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:

    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:

    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.