Metadata-Version: 2.5
Name: pichak
Version: 0.2.0
Summary: Derive fine-tuning hyperparameters from measurements of your actual machine, and print the measurement behind every number.
Project-URL: Homepage, https://github.com/olesxg/pichak
Project-URL: Findings, https://github.com/olesxg/flap-findings
Author-email: Oleksandr Pichak <stasselust@gmail.com>
License: MIT
License-File: LICENSE
Keywords: batch-size,benchmark,fine-tuning,gpu,hyperparameters,learning-rate,llm,lora,vram
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Provides-Extra: dev
Requires-Dist: pytest>=7; extra == 'dev'
Provides-Extra: hf
Requires-Dist: safetensors>=0.4; extra == 'hf'
Requires-Dist: torch>=2.0; extra == 'hf'
Requires-Dist: transformers>=4.40; extra == 'hf'
Provides-Extra: torch
Requires-Dist: torch>=2.0; extra == 'torch'
Description-Content-Type: text/markdown

# pichak

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

```bash
pip install pichak
```

```python
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.

```python
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

```bash
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
```

```python
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

```bash
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](https://github.com/olesxg/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](https://github.com/olesxg).
