hahaws777/OLMo-2-0425-1B-CS336A5-LowVRAM

CS336A5 小显存适配 — OLMo-2-1B

用小显卡直接训练：在单张 RTX 5070 Ti 16GB 上完成的 OLMo-2 全参数强化学习模型。

This is the best recorded saved checkpoint from the project's local 100-rollout method suite: offpolicy_naive, rollout iteration 100. It was fine-tuned by hahaws777 from Ai2's OLMo-2-0425-1B base model, not trained from scratch and not derived from the official Instruct checkpoint.

Source code and reproduction notes: CS336A5 小显存适配. This is a student research implementation based on Stanford CS336 Assignment 5, not an official Stanford or Ai2 release.

Recorded evaluation

• —Strict answer reward: 497/1024 = 48.53515625%.
• —Format reward: 989/1024 = 96.58203125%.
• —Evaluation data: the first 1024 examples of the local GSM8K test file, in file order; this is not the complete 1319-example test split.
• —Prompt: r1_zero, supplied as r1_zero.prompt; one response per question, temperature 1.0, top-p 1.0, maximum 512 generated tokens, stop string </answer> included in the output.
• —Backend: colocated vLLM 0.19.1; evaluation sampling seed 100100 (training seed 0 + 100000 + iteration 100).
• —Grading: the assignment's r1_zero_reward_fn. It requires the exact substring </think> <answer> and a closing </answer> before grading the extracted answer. A mathematically correct but unparseable response can score zero.

These counts are derived from the recorded aggregate validation metrics, not from a new evaluation at publication time. The same test subset was inspected repeatedly during training and checkpoint selection; this is not an independent held-out assessment after selection. A single stochastic evaluation does not establish a general advantage over other methods. Results using other prompts, seeds, inference backends, or the complete test split are not directly comparable.

An older cloud run logged 518/1024 at iteration 160, but no corresponding saved weights were found; that number is not claimed for this release. Saved checkpoints without recorded evaluations have not all been reevaluated. See the published experiment evidence and the included evaluation.json for provenance.

Training and memory approach

• —Base revision: a1847dff35000b4271fa70afc5db10fd29fedbdf of allenai/OLMo-2-0425-1B.
• —Full-parameter fine-tuning, not LoRA and not 4-bit weight quantization. The checkpoint contains 1,484,916,736 parameters in BF16, approximately 2.97 GB of weight data on disk.
• —Training pool: first 6400 examples of the local GSM8K training file; r1_zero prompting, training seed 0.
• —100 rollout iterations, each sampling 32 questions with 8 responses each (256 responses). These are rollout iterations, not 100 dataset epochs or 100 optimizer updates.
• —Algorithm offpolicy_naive: group-mean reward baseline, group standard-deviation normalization, sequence-normalized loss, and no importance-ratio reweighting/clipping.
• —AdamW, learning rate 1e-5, betas (0.9, 0.95), weight decay 0, gradient norm clipping at 1.0, foreach=False.
• —Training batch size 8, gradient accumulation 8, microbatch size 1. The log records 1274 actual optimizer updates out of 3200 attempted update slots.
• —BF16 parameters, gradients, and native AdamW moments; FP32 probability, entropy, advantage, and loss calculations. No FP32 master-weight copy. This is a memory-saving numerical tradeoff, not a claim of equivalence to conventional FP32-state AdamW.
• —Decoder gradient checkpointing, output-head/log-probability chunking (32 tokens), SDPA attention, and alternating batched vLLM generation and policy training with sleep/offload and weight synchronization.

The 16 GB claim refers to this tested hardware and configuration, not every model size, context length, or batch size. Disk weight size is not training VRAM usage. Detailed memory methodology and limitations are in the reproduction guide.

Usage

The saved config was produced with Transformers 5.7.0. Use a compatible PyTorch installation for your GPU; do not replace an existing working training environment just to load this model. No custom remote model code is required.

from pathlib import Path
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "hahaws777/OLMo-2-0425-1B-CS336A5-LowVRAM"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, dtype=torch.bfloat16, attn_implementation="sdpa"
).to("cuda").eval()

template = Path(hf_hub_download(model_id, "r1_zero.prompt")).read_text(
    encoding="utf-8"
)
prompt = template.format(question="A box has 12 pencils. You add 6 more. How many pencils are in the box?")
inputs = tokenizer(prompt, return_tensors="pt", return_token_type_ids=False).to("cuda")
with torch.inference_mode():
    output = model.generate(
        **inputs,
        max_new_tokens=512,
        do_sample=False,
        use_cache=True,
        stop_strings=["</answer>"],
        tokenizer=tokenizer,
        pad_token_id=tokenizer.pad_token_id,
    )
print(tokenizer.decode(output[0, inputs.input_ids.shape[1]:], skip_special_tokens=True))

This is a greedy usage example, not reproduction of the stochastic vLLM evaluation above. No generated example output is asserted here. The prompt already ends with <think>; the continuation should complete the reasoning and answer tags. Do not assume a chat template is available. use_cache=True explicitly enables inference caching because the saved training config has it disabled.

Files and limitations

This repository contains full BF16 model weights, model/generation configs, tokenizer files, the exact prompt file, evaluation metadata, and a weight checksum manifest. It does not include optimizer, scheduler, or RNG state, and therefore is not an exact training-resume checkpoint. Source machine paths and raw training data are not included in the release metadata.

This is an English math research model, not a broadly instruction-tuned or safety-aligned assistant. It can make arithmetic and reasoning errors, ignore formatting, hallucinate, or produce inappropriate content. No general capability, safety, or deployment-readiness evaluation is claimed.

Attribution and license

The upstream OLMo-2 model is developed by the Allen Institute for AI (Ai2) and released under Apache-2.0. This fine-tuned model is distributed under the same license; see LICENSE and the upstream model card at the base revision. The weights were modified by reinforcement-learning fine-tuning; the original authors should not be credited with or considered endorsers of this particular fine-tuned release.

Model architecture and pretrained weights: Ai2 / Team OLMo. Assignment scaffolding and grading: Stanford CS336 and the upstream contributors. Local low-memory adaptation, integration, experiments, and this fine-tuned checkpoint: hahaws777's project. See the upstream sources for their respective notices and citations.