burtenshaw/tailsft-lfm350m-experiment

TailSFT with LFM2.5-350M and Trackio

Status, 2026-09-16: GPU preflight completed all three training arms and four evaluations. The strict final-answer grader and namespace/resume checks pass (50 tests). A full-length batch of 256 samples completed on the A100; the full comparison has been submitted. Full experiment results are not yet available.

Small, hackable reproduction of the TailSFT filtering method. The objective is higher pass@16 after supervised training. This experiment does not test downstream RL or claim to match the paper's reported scores.

• —Model: LiquidAI/LFM2.5-350M-Base, revision 9960764e30892e01f29a6dc23df2533fcd8bd5ae.
• —Training: 8,192 distinct math problems from OpenMathInstruct-2; 256 held-out validation problems; two epochs.
• —Arms: ordinary SFT, TailSFT dropping the most-improved 25%, random dropping 25%. Same initialization, data order, learning rate, steps, and training seed 42.
• —Evaluation: base and all three trained models, MATH-500, 16 samples/problem, batch 256, 3,072-token generation cap, temperature/top-p 1, no top-k. Math-Verify grading of the last explicit boxed answer in the completion; absent, unfinished, or unparseable answers count as incorrect.
• —Optimizer: FP32 parameters and AdamW state, BF16 autocast forward; batch 8, no accumulation/packing; LR 2e-5, weight decay 0.1, clipping 1, 3% warmup then constant.
• —Full GPU job.
• —Trackio dashboard.
• —Source, frozen data, and run artifacts.

This is a single training-seed comparison. Later confirmation seeds are separate experiments.

The part to hack

tail_loss.py contains the objective: calculate response-token cross-entropy, subtract each example's cached initial mean loss, drop the most negative margins, then divide retained token-loss sum by retained token count. drop_fraction=0 recovers ordinary SFT. Selection occurs within each physical batch.

Change the fraction or replace the ranking score to experiment. Keep the data, optimizer, and evaluation fixed for comparisons. Training uses ordinary PyTorch and Transformers; no custom kernels or distributed setup is required.

Reproduce locally on one CUDA GPU

The Hub repository stores these files under source/ and the prepared data under data/. The exact source revision for the current run is 86591e5c7c856c9c503b94b11dc1d22295cb3962; live status is recorded in the repository-root run-state.json. First install the HF CLI, authenticate with hf auth login, and download that immutable snapshot:

hf download burtenshaw/tailsft-lfm350m-experiment --repo-type dataset \
  --revision 86591e5c7c856c9c503b94b11dc1d22295cb3962 --include "source/*" "data/*" --local-dir tailsft-reproduction
cd tailsft-reproduction/source
uv venv --python 3.12 .venv
uv pip install --python .venv/bin/python -r requirements-linux.lock
source .venv/bin/activate
python -m pytest -q
# Edit configs/lfm350m.json for your own namespace before the following runs.
python train.py --config configs/lfm350m.json --data ../data --output ./outputs --mode preflight
python train.py --config configs/lfm350m.json --data ../data --output ./outputs-full --mode full

For your own run, edit trackio_space, trackio_project, artifact_repo, and experiment before starting. Authenticate with hf auth login for Trackio; add --push to persist checkpoints/results to your artifact repository. Use a new experiment ID whenever changing the model, data, or settings.

To regenerate the data:

.venv/bin/python prepare_data.py --config configs/lfm350m.json --output-dir ../data

Reproduce on HF Jobs

bootstrap_job.py embeds all 89 pinned Linux dependencies. It downloads one immutable source/data revision, runs the tests, and executes training. The namespace flags below save models, results, and Trackio logs to your account. Replace YOUR_USERNAME with your Hub username. --resume restores persisted checkpoints and raw evaluation samples from your artifact repository; keep the same experiment ID and settings when resuming.

hf jobs uv run --detach --flavor a100-large --timeout 10h --secrets HF_TOKEN \
  bootstrap_job.py --source-revision 86591e5c7c856c9c503b94b11dc1d22295cb3962 --mode full \
  --artifact-repo YOUR_USERNAME/tailsft-lfm350m-experiment \
  --trackio-space YOUR_USERNAME/tailsft-lfm350m-trackio \
  --trackio-project tailsft-lfm350m --experiment tailsft-lfm350m-my-run

The current A100 80GB hourly rate is $2.50; a ten-hour timeout bounds one full job's compute at roughly $25. A measured 256-sample, 3,072-token batch took 175.5 seconds, using 15.43 GB peak allocated GPU memory. This projects roughly 6.2 hours for evaluation alone; plan about 7–9 hours including training, with the ten-hour cutoff retained. This estimate is based on one base-model batch and is not a completion guarantee. Preflight runs are separate and are not experiment results. The original preflight used the permissive grader; its scores must not be interpreted as benchmark results.

Reproducibility and artifacts

data/manifest.json records immutable inputs, exact selection rules, row IDs, hashes, masks, EOS, and removal counts. Selection groups identical problems before splitting and removes normalized exact matches against MATH-500. This check does not exclude paraphrases or unknown pretraining exposure.

Prepared training data: mean length 529.58 tokens, maximum 1,822, and 3,538,839 target tokens. No target truncation. Data uses the explicit prompt in evaluate.build_prompt; it does not depend on a chat template.

Each run preserves configuration, exact installed dependencies, cached initial losses, per-step retained IDs/margins, full checkpoints with optimizer/RNG state, unfiltered validation loss, raw generations, graded outputs, and pass@1/2/4/8/16 by problem and difficulty. Checkpoints upload every 512 steps; evaluation uploads at roughly five-minute intervals. Trackio stores live curves and GPU/CPU metrics in a persistent HF bucket.

The GPU throughput check selects an evaluation batch size using runtime and memory only; its accuracy is not used to tune the recipe. Training and evaluation remain valid if they show no TailSFT gain. Test results will not be used to change the recipe. The original research recommendation is preserved in research-plan.md; this configuration supersedes its model choice.

Sources and attribution

• —TailSFT: Malladi et al., arXiv 2608.25756.
• —LFM weights: Liquid AI, under the LFM Open License v1.0. Fine-tuned weights are modified derivatives; checkpoint directories retain the license and a modification notice.
• —Training data: NVIDIA OpenMathInstruct-2, CC-BY-4.0. The snapshot is a selected, tokenized, answer-checked derivative with source IDs retained.
• —Evaluation: HuggingFaceH4/MATH-500; source provenance retained, without assigning it a new license.
• —Grading: Math-Verify, pinned 0.9.0.
• —Tracking: Trackio, pinned 0.38.0.