nltower-100m-sft-tw1X-e2lr05

A 124M-parameter Llama-style model from the nl_tower study: the end-of-pretraining base model fine-tuned on solver-written reasoning traces for 2 epochs at learning rate 5e-5. It is the stage-1 starting point of the "staged" recipe (SFT, then pure RL without demonstrations), not a general-purpose assistant.

This repository is part of the nl_tower (100M) collection on the Hugging Face Hub (5 CPT models, 6 SFT models, the base model and 4 datasets).

nl_tower in brief

nl_tower is a synthetic reasoning task. Each puzzle places 6-12 people on distinct floors (1-40) of a tower. The prompt lists the occupied floors, each person's best friend, and clues that relate the floors of two terms ("greater than", "k floors above", sums, parity, "twice", highest/lowest). Many clues refer to people indirectly, for example "the best friend of the person on the second-highest floor"; such a clue becomes usable only after other clues have pinned part of the answer. Each puzzle has a unique solution that a step-by-step solver reaches without guessing.

The study asks whether reinforcement learning can teach reasoning compositions that the imitation data never shows:

  • Covered patterns (shown in the mid-training traces): friend ("Ava's best friend"), friend2 ("the best friend of Ava's best friend"), order ("the person directly above Ava"), rankn ("the person on the second-highest floor"), dfriend_occ ("the best friend of the person on floor 19").
  • Held-out patterns (appear only in RL prompts and tests): dfriend_rank ("the best friend of the person on the highest floor"), dfriend_above ("the best friend of the person directly above Ava"), and nested (two-level references).

Test rungs: easy (nltw_e, 6-7 people, covered only), covered mid (nltw_m0, 8-9 people, covered only), held-out mid (nltw_m, 8-9 people, 2-4 held-out steps), size-only hard (nltw_hb, 10-11 people, 2-4 held-out steps) and 12 people (nltw_hb12, test only).

Position in the pipeline

base pretraining (iteration 30,000, before the cooldown) --CPT-N (cooldown + N trace sequences per batch)--> CPT-N models --RL--> policies
base pretraining (iteration 38,147, end of the run) --SFT on the same traces--> SFT models --RL (no demos)--> policies

This model is one of the six staged SFT starting points (1 / 2 / 4 epochs x learning rate 2e-4 / 5e-5). Note that it starts from iteration 38,147 of the base run (the end of pretraining, after the full cooldown), not from sfanm/nltower-100m-base (iteration 30,000).

Training (SFT)

  • Start: end of the base pretraining run (iteration 38,147); that checkpoint is not part of this release.
  • Data: all 796,175 traces of sfanm/nltower-traces-tw1X (config traces), easy and covered-mid puzzles with covered patterns only.
  • Loss: token-level negative log-likelihood on the trace tokens plus EOS; prompt tokens get no loss. Sequences up to 2,048 tokens (the longest record has 1,218).
  • Optimiser: AdamW (beta1 0.9, beta2 0.95, weight decay 0.01), fp32 weights, peak learning rate 5e-5, cosine schedule with 5 % warmup.
  • Batch: 32 sequences per step on 4 GPUs (data parallel); 24,881 steps per epoch, 49,762 steps in total (2 epochs).

Architecture

property value
architecture LlamaForCausalLM (decoder-only)
parameters 123,587,328 (embeddings tied with the output layer)
layers / hidden size / heads 12 / 768 / 12 (head dim 64, no grouped-query attention)
MLP SwiGLU, intermediate size 2,048
normalisation / position RMSNorm (eps 1e-5) / RoPE (theta 10,000)
tokenizer GPT-2 BPE (50,257 tokens; embedding padded to 50,304)
special tokens <|endoftext|> (id 50256) is BOS, EOS and PAD
context length 2,048 tokens
weights model.safetensors, bfloat16

Prompt format and usage

The prompt is the puzzle text followed by "\n\nAnswer: ". The model writes its reasoning and ends with the full assignment in \boxed{Name=floor, ...}. A completion is correct when every Name=floor pair in the boxed answer matches the solution (order does not matter). The chat template only concatenates message contents (no role markers), so apply_chat_template and the raw prompt give the same tokens.

from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

repo = "sfanm/nltower-100m-sft-tw1X-e2lr05"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, torch_dtype="bfloat16")

test = load_dataset("sfanm/nltower-env", "tests", split="nltw_m0")
prompt = test[0]["prompt"]                      # question + "\n\nAnswer: "
ids = tok(prompt, return_tensors="pt").input_ids
out = model.generate(ids, max_new_tokens=1280, do_sample=False, pad_token_id=tok.eos_token_id)
print(tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True))   # ... \boxed{Name=floor, ...}
print("gold:", test[0]["answer"])

Evaluation

Zero-shot screen of this model before RL. Protocol: 1,024 test questions per rung; 64 samples per question at temperature 1.0 (top-p 1.0) and up to 1,280 new tokens with vLLM in bfloat16; pass@k uses the unbiased estimator 1 - C(n-c, k)/C(n, k) with n = 64; greedy is one temperature-0 sample. Test sets: sfanm/nltower-env, config tests.

test rung pass@1 pass@64 greedy
easy (nltw_e) 0.913 1.000 0.947
covered mid (nltw_m0) 0.788 1.000 0.853
held-out mid (nltw_m) 0.021 0.237 0.025

The covered rungs are nearly solved, while the held-out-pattern rungs stay low: imitation of covered patterns alone does not teach the held-out compositions.

RL runs that start from this model and have evaluations in sfanm/nltower-results (filter init == "staged-2ep-lr5e-5", the name of this model in the code), by setting (base: 8-9-person prompts, 6,000 updates; extended: 10-11-person prompts added, 12,000 updates) and group:

setting group runs
base staged 2
extended staged 2

Intended use and limitations

  • Intended for research on imitation followed by reinforcement learning, as the stage-1 model of the staged recipe.
  • Fine-tuning on traces largely overwrites general text ability (its loss on ordinary English prose is high); use it only on nl_tower prompts.
  • 124M parameters and a synthetic task: results may not transfer to larger models or natural tasks. Not instruction-tuned or safety-tuned; not for user-facing use.

Training data

Code

Code: https://github.com/yudasong/unified-training/tree/release (environment generator, solver, data builders, training recipes and evaluator).

License

No license has been chosen for these weights yet, so the metadata says license: other. The base model was pretrained on:

Use the weights only in ways that the terms of these datasets allow. The nl_tower data used after pretraining is the authors' own synthetic data (Apache-2.0, see the dataset repositories).

Authors and citation

Authors and citation: to be added.

This release (files and this card) was prepared with AI assistance (Claude, an Anthropic model) and is being reviewed by the authors.

Downloads last month
632
Safetensors
Model size
0.1B params
Tensor type
BF16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Datasets used to train sfanm/nltower-100m-sft-tw1X-e2lr05