pichak

Fine-tuning hyperparameters derived from your machine and your task, with the measurement printed next to every number.

pip install pichak

from pichak import derive
from pichak.tune import TrainingStep, forbid_spill

forbid_spill("cuda:0")          # past the VRAM this process can have: OOM, never paging
opt = torch.optim.AdamW(trained_params, lr=1e-4)      # fresh; pichak sets the lr

plan = derive(
    data="train.jsonl",
    tokenizer=tok,
    # the loss of b rows padded to the longest sequence you train at
    step_fn=TrainingStep(lambda b: model(**rows(b, pad_to=seq_len)).loss, opt),
    grad_fn=lambda i: (gradients_of_fresh_micro_batch(i), scored_tokens(i)),
    named_weights=trained_weights,          # the adapter's factors, or the model's own
    full_finetune=False,                    # True when named_weights ARE the model's
)
print(plan.report())

Abridged, and assembled from the numbers the validation runs in validation/ measured (SmolLM2-135M full fine-tune, MetaMathQA, GTX 1060 6GB):

  micro_batch        2
                     a real forward+backward at each batch, doubling then
                     bisecting: 1=ok, 2=ok, 4=OOM, 3=ok. 4 refused, so 3 is the
                     largest that fit, peaking at 4.57GB of 4.57GB; kept 2, the
                     largest whose peak (3.63GB) leaves 12% of it for what one
                     step does not show: fragmentation over a long run,
                     evaluation, checkpoints

  noise_scale_tokens 4530
                     gradient noise scale tr(Sigma)/|G|^2 = 4,530 supervised
                     tokens, measured from 64 micro-batches (McCandlish et al. 2018)

  grad_accum         14
                     14 x micro-batch 2 x 166.0 scored tokens per row = 4,648
                     scored tokens per optimizer step, aimed at the gradient
                     noise scale

  learning_rate      0.000139
                     1/320 of the scale at which one Adam step would rewrite rather
                     than perturb ... Full fine-tune: 1/320 measured best on a 135M

Nothing under NOT MEASURED. Without grad_fn, tokens per step fall back to the 65,536 constant and the report says so.

---

Why not just use a good default

Because you cannot tell a good default from a bad one when it fails.

micro_batch=4 tells you nothing about whether 4 came from a measurement on your card, a heuristic in a blog post, or a number someone picked in 2023 for a different model. So when it OOMs at step 40 you bisect instead of reading.

Every number pichak returns carries the sentence that produced it. Anything it could not measure is collected under plan.constants() and printed together, so a constant can never quietly pass for a derivation.

plan.micro_batch          # 3
plan.why("micro_batch")   # the ramp, rung by rung, and which one failed
plan.constants()          # {} with grad_fn;  {'target_tokens': 65536} without

Tokens per step: measured, not 65,536

0.1.0 aimed every optimizer step at 65,536 tokens and said honestly that it was a guess at "the scale where gradient noise stops dominating". That scale has a name and a formula — the gradient noise scale B_noise = tr(Σ)/|G|² (McCandlish et al., 2018) — and it depends on the task. Measured on SmolLM2-135M, two independent draws of rows each:

task	loss scores	per row	B_noise draw 1	draw 2
MetaMathQA	the completion	166	4,378	4,683
Python code	every token	383	1,557	1,543
English web text	every token	383	2,125	—
Ukrainian Wikipedia	every token	383	865	814
True/False deduction	the completion	2	1	1

Draws agree within 1–7%. Tasks differ by ~4,700x. At 65,536 raw tokens a step, 0.1.0 spent 11x the noise scale per step on math and ~1,300x on the deduction task.

Two things make this universal rather than tuned:

• Scored tokens, not tokens read. A row of 900 prompt tokens and a one-token label is one scored token. profile_corpus counts what the loss scores at your seq_len, including what truncation cuts from the end of the completion.
• The corpus caps it. If B_noise would leave fewer than 300 updates in one pass, the step is cut to fit, and the learning rate follows it down.

The learning rate

An Adam update moves a parameter by roughly lr, whatever the gradient's scale. So there is a learning rate at which one step rewrites a weight matrix instead of nudging it: where lr * sqrt(numel(W)) reaches ||W||_F. Every sane learning rate is a fraction of that.

Adapters (LoRA/PiSSA factors), measured at rank 32:

rewrite/2    2.26e-3   destroyed the model twice (CE 1.74 -> 10.4, 1.11 -> 19.3)
rewrite/32   1.41e-4   the first scale that actually learned
rewrite/45             where a hand-tuned LoRA run at 2e-4 landed

A full fine-tune — the model's own weights — measured on SmolLM2-135M, MetaMathQA, 60 steps at the noise-scale batch, held-out CE from 0.793:

rewrite/32    1.39e-3   -> 0.733    the adapter anchor: 3.7x less learned
rewrite/320   1.39e-4   -> 0.568    best
rewrite/1483  3.00e-5   -> 0.605
rewrite/3200  1.39e-5   -> 0.639

The adapter anchor does not carry over: a LoRA step moves the factors, and W only through their product, which starts at zero. Train W itself and the same rate is a several-times bigger step. Pass full_finetune=True and pichak opens at rewrite/320.

The batch. The anchor is right at the noise scale. Below it the gradient is mostly noise and the step is cut by Adam's square-root rule (Malladi et al., 2022):

lr = rewrite / divisor * min(1, sqrt(B / B_noise))

It is never raised above the anchor: a bigger batch lowers the noise, not how far one step moves a weight.

It catches the failure that does not raise

On Windows, a batch that exceeds VRAM does not throw. The driver pages the excess to system memory and the step simply crawls. Measured on a GTX 1060 6GB with a 6-layer model at sequence 512:

	micro_batch	peak	seconds/step
watching only for OOM	14	18.24GB on a 6GB card	47.5
watching per-sample time too	3	5.05GB	2.0

Three guards, from strongest to weakest:

• forbid_spill() — a cap, not a check. It limits the process to the dedicated VRAM it can actually have, less a guard, so an over-allocation raises at the line that caused it. Measured: an uncapped probe on this card went to 5,999 of 6,144 MiB and put 1.1 GB in shared memory, silently. Capped, 23 minutes of training at up to 5 GB dedicated held shared memory at 58 MB — the 56 MB any idle CUDA process shows. See the limits below for what a cap on one process cannot do.
• Other programs' memory is not capacity. The certain check compares the peak with the memory this process could use — free plus its own — not the card's total. With browsers holding 2.9 GB of 6, a 5.5 GB peak used to pass.
• A slow step is a suspicion until it repeats. Real paging does not go away on a second timing; scheduler jitter does. A step costing microseconds used to stop the ramp at batch 1 about once in ninety runs.

The batch a training run fits, not a single step

A ramp is only as right as the step it runs. On a 135M full fine-tune with a 4.57 GB cap, three hand-written steps each cleared 3 rows that then OOMed at 3, for two reasons:

the step left out	ramp said	the run needed
rows at the run's longest sequence (ramped at 268 tokens, trained at 384)	3 fit	OOM
the gradient buffers an accumulating backward runs on top of	3 at 3.95 GB	OOM
— (TrainingStep: all of it)	3 at 4.57 of 4.57 GB, kept 2	60 steps, no OOM

TrainingStep(loss_of, optimizer) pads nothing for you — loss_of(b) must build b rows at your longest sequence — but it puts the optimizer's state and the gradient buffers in the peak, steps at learning rate 0 so the weights do not move, and hands the optimizer back fresh.

What it derives, and from what

	measurement
seq_len	tokenised p95 over 512 real rows of your corpus
loss_policy	whether the rows have a separable completion to mask against
micro_batch	real training steps, doubling then bisecting; 12% of the bytes kept back
seconds_per_step	the wall time of that rung
noise_scale_tokens	the gradient noise scale of your task, from grad_fn
grad_accum	rows per step to reach the noise scale in scored tokens, capped by the corpus
lora_rank	largest rung whose optimizer state fits the measured capacity
learning_rate	the trained weights' own norms, the adapter or full-fine-tune divisor, and the batch

Everything is optional. It will not invent a number it could not measure — a plan with three derived values and an honest gap is more useful than one with ten you cannot tell apart.

What it does not know yet

Said here so nobody finds it the hard way:

• The noise scale is measured once, at the start, and it grows fast. Same 135M full fine-tune on MetaMathQA, three runs, re-measured on held-out rows with 64 micro-batches of the run's own size:

  step	3 rows a micro-batch	2 rows	1 row (COARSE)
  0	4,530	4,530	4,530
  30	15,262 (3.4x)	24,454 (5.4x)	37,952 (8.4x)
  60	26,703 (5.9x)	22,784 (5.0x)	45,767 (10.1x)

  Every run says the same thing: 3-10x in 60 steps. So a batch planned at step 0 is a fifth of the noise scale or less by step 60, and the gradient drifts back into the noise-dominated regime the plan was meant to leave. The exact factor is as good as the probe: with 1-row micro-batches, 64 of them hold ~10,600 scored tokens against a noise scale of ~40,000, the estimate is coarse, and it read the highest. Until pichak tracks it for you, call measure_noise_scale again every few dozen steps, with enough micro-batches that k × b reaches the scale you expect, and grow grad_accum to follow it.

• B_noise at step 0 favours the easiest direction. On the True/False task it is 1 token: every row agrees on the answer format and the end-of-sequence token before any reasoning is learned. A batch sized from that is right only until the format is learned.

• The full-fine-tune divisor is one measurement — one 135M, one task, 60 steps. Treat /320 as the right order of magnitude, not a law.

• forbid_spill() caps this process, not the machine. It is set once, and it bounds only PyTorch's allocator. Memory outside it (CUDA context, cuBLAS workspaces: ~76 MB here) and other programs growing during the run (~130 MB over 20 minutes on this desktop) are what the 512 MB guard is for. With a 256 MB guard, one run in three had the Windows video memory manager demote 264 MB of the training process to system memory while PyTorch was still under its cap. A desktop app that grows by more than the guard can still do that.

Measurement tools, usable on their own

pichak gpu                    # virtualisation, launch latency, transfer bandwidth
pichak disk "models/*.safetensors"   # queue-depth sweep, page cache bypassed
pichak corpus train.jsonl mistralai/Mistral-Small-24B-Instruct-2501

from pichak.measure import disk_queue_depth, transfer_bandwidth, virtualisation

These exist because the numbers people quote are rarely the numbers their machine gives:

• A buffered disk benchmark reported 4079 MB/s on a drive rated 2100 — that was the page cache. Unbuffered, the same drive peaks at queue depth 3 and gets slower past it.
• A rented A6000 measured D2H 457 MB/s against H2D 1502, and pinned memory — the standard fix — bought nothing. Nothing about PCIe explains that; it was a fabric, and it decided where hidden states could live.
• One rented machine had 84.8us kernel launches against a normal 2-5. A workload issuing 400,000 launches per step would have spent nine hours on overhead and looked slower than a 2016 card.

Install

pip install pichak            # the plan, the disk sweep — no dependencies
pip install pichak[torch]     # the ramp, the noise scale, the learning rate, the GPU
pip install pichak[hf]        # + transformers, for the corpus CLI

Python 3.9+. The core has no dependencies at all; torch is only needed for the parts that touch a GPU.

Where the numbers come from

The 0.1.0 measurements were made while fine-tuning a 24B model on a 6GB GTX 1060 and on a rented A6000 held down to the same 6GB; the raw logs, including the runs that failed, are at flap-findings. The 0.2.0 measurements — noise scale across five tasks, the full-fine-tune divisor, the drift, and the spill cap — are reproducible from validation/ in this repository, with their JSON output beside each script.

Licence

MIT. By Oleksandr Pichak.