Batching & remote

Two flags change how a document runs without changing what it is. --batch-size N bounds how many rows one model call holds. --engine ndif moves the weights off this machine. Neither is part of the experiment, and the page below runs one document all three ways and compares the numbers value by value.

The document is documents/real/weekdays_layer_sweep.json — Llama-3.2-1B, fp32, 32 swept points of 42 rows with two metrics each, so 2688 numbers per run. Layer sweep is that document end to end; here it is here because it is big enough to run three times and compare.

At a glance #

$ time causalab-mini run documents/real/weekdays_layer_sweep.json \
      --device-map cuda --out out/local
out/b-local/document.json
out/b-local/run.json
out/b-local/layers=0,pos=-4/document.json
out/b-local/layers=0,pos=-4/iia.json
…
out/b-local/layers=15,pos=-1/document.json
out/b-local/layers=15,pos=-1/iia.json
out/b-local/layers=15,pos=-1/logit_diff.json

real	0m16.210s

$ time causalab-mini run documents/real/weekdays_layer_sweep.json \
      --device-map cuda --batch-size 8 --out out/batch8
out/b-batch8/document.json
out/b-batch8/run.json
out/b-batch8/layers=0,pos=-4/document.json
out/b-batch8/layers=0,pos=-4/iia.json
…
out/b-batch8/layers=15,pos=-1/document.json
out/b-batch8/layers=15,pos=-1/iia.json
out/b-batch8/layers=15,pos=-1/logit_diff.json

real	0m16.172s

$ time NDIF_HOST=http://localhost:8001 causalab-mini run \
      documents/real/weekdays_layer_sweep.json --engine ndif --out out/ndif
✓ [5f3fa4433d264f09a122f27d83e1199e] COMPLETED    (4.0s) Your job has been completed.
out/b-ndif3/document.json
out/b-ndif3/run.json
out/b-ndif3/layers=0,pos=-4/document.json
…
out/b-ndif3/layers=15,pos=-1/document.json
out/b-ndif3/layers=15,pos=-1/iia.json
out/b-ndif3/layers=15,pos=-1/logit_diff.json

real	0m9.472s

The numbers #

Remote changed nothing. All 2688 values match to the bit, and every file the run wrote — 33 document.json and 64 metric tables — is byte for byte the same file. The serialization round trip, the block executing in another process, the plan being filled in from what came home: none of it perturbs the arithmetic.

Batching moved the last bits, and only of one metric. All 1344 iia values are identical — a match metric is an argmax, and the argmax does not move. 83 of 1344 logit_diff values are identical and the rest differ by at most 5.9e-05 on numbers around 2.4. Every point's mean is the same to four decimals, and the 65 byte-identical files are exactly the 33 document.json plus the 32 iia.json.

--batch-size: what it bounds #

A window of rows is a whole small pass: the same forwards in the same order over a slice of the rows, with the values published before it sliced the same way, so a write still meets the operand of its own row. What the windows read is concatenated back in row order, and everything downstream sees one pass.

The windowing is in the shared walk, not in an engine. An engine is handed a forward that is merely shorter and cannot tell the difference — every per-row thing a forward holds is a tuple with one entry per row, so a window is a slice of each. Positions are absolute indices into the padded width the client fixed once for all rows, so they survive the slice unchanged.

With no --batch-size there is one window, which is the pass exactly as compiled: every row in one model call. On this document that is 42 rows of 8 tokens, and what it costs against the weights is small:

171 MiB, on a document whose weights are about 4.6 GiB. Batching is not free — the two wall times above are the same to a tenth of a second and the arithmetic moves — so it is worth reaching for when the rows are what does not fit, and not before. The recorded run of weekdays_head_patching, 96 points over the same 42 rows, used --batch-size 16 (documents/real/README.md).

--engine ndif: what moves #

--engine ndif is not a different engine. It is the same NNterpEngine, built with dispatch=False so no weights are loaded on this machine, and executed with remote=True. The client holds a meta shell: enough of the model to resolve a component to a module path, to know a site's width, and to tokenize.

The whole plan runs inside one model.session(remote=True). A 32-point sweep is one job on the server, not 32 — the server log for the run above records one request, 4032 ms of execution inside the model actor, in a 9.5 s wall. The client's remaining time is building the meta shell, compiling 32 documents and tokenizing.

What comes home is plain: strings and tensors. The client already has the plan, so only what fills it in makes the trip, and the plan is filled in on arrival. That is why the output directory is identical rather than merely equivalent — the same writer writes it from the same plan either way.

Where the server is is nnsight's business, not mini's: NDIF_HOST overrides the configured base URL, and NDIF_API_KEY the key. There is no mini flag for either.

What run.json records #

Every output directory's root carries one. The three runs above, side by side:

{
 "engine": "NNterpEngine",
 "remote": false,
 "batch_size": null,
 "versions": {
  "torch": "2.14.0+cu126",
  "transformers": "5.17.0",
  "nnsight": "0.8.1.dev125+g524c33fc",
  "nnterp": "0.5.0.dev0+internalsaccessors",
  "pydantic": "2.13.5"
 },
 "causalab_mini": "76c8c9d0c0d1f3e0b8e2eaeb09564107d1c0a70cb9e8b7d6cecab09fdc92d8a1",
 "served_dtype": "torch.float32"
}

{
 "engine": "NNterpEngine",
 "remote": false,
 "batch_size": 8,
 "versions": {
  "torch": "2.14.0+cu126",
  "transformers": "5.17.0",
  "nnsight": "0.8.1.dev125+g524c33fc",
  "nnterp": "0.5.0.dev0+internalsaccessors",
  "pydantic": "2.13.5"
 },
 "causalab_mini": "76c8c9d0c0d1f3e0b8e2eaeb09564107d1c0a70cb9e8b7d6cecab09fdc92d8a1",
 "served_dtype": "torch.float32"
}

{
 "engine": "NNterpEngine",
 "remote": true,
 "batch_size": null,
 "versions": {
  "torch": "2.14.0+cu126",
  "transformers": "5.17.0",
  "nnsight": "0.8.1.dev125+g524c33fc",
  "nnterp": "0.5.0.dev0+internalsaccessors",
  "pydantic": "2.13.5"
 },
 "causalab_mini": "76c8c9d0c0d1f3e0b8e2eaeb09564107d1c0a70cb9e8b7d6cecab09fdc92d8a1",
 "served_dtype": "torch.float32"
}

Three fields differ, and the rest is the same because the same code ran:

Wall times #

The remote run is faster here because the replica is already hot and the client never loads 4.6 GiB of weights. That is the comparison NDIF is for, and it is the one the recorded numbers show too — documents/real/README.md records 8 / 12 / 33 / 61 s locally on an A100 for its four documents against 7 / 8 / 19 / 36 s through NDIF, each document one job.