barrier¶
Line one block up with another, mid-run.
The blocks of a trace — one per with tracer.invoke(x): — run in an order the
model chooses, not the order they were written: each resumes when the model
reaches the location it asked for. That is what makes them a batch rather than
a sequence, and it is fine as long as they don't need anything from each other.
They need a meeting point as soon as one hands something to another. A block that
reads an activation and puts it somewhere the next block takes it from is only
correct if the read happens first, and neither block can see the other's progress
to know. A Barrier is that meeting point: every block that holds one
calls it, each waits, and the last to arrive releases them all — so everything
written above a barrier has happened before anything written below one.
Here two invokes share one value: the first reads a prompt's embeddings, the
second generates from them, and the barrier pins the read before the write so the
receiver isn't a NameError waiting to happen.
.. code-block:: python
from nnsight.modeling.transformers import TransformersModel
model = TransformersModel("openai-community/gpt2", dispatch=True)
with model.generate(max_new_tokens=3, do_sample=False) as tracer:
barrier = tracer.barrier(2)
with tracer.invoke("Madison Square Garden is in the city of"):
embeddings = model.transformer.wte.output
barrier()
tokens = tracer.result.save()
with tracer.invoke("_ _ _ _ _ _ _ _ _"):
barrier()
model.transformer.wte.output = embeddings
# The underscore prompt carries none of the meaning; the embeddings do, so
# both continuations end "New York City".
Waiting costs nothing but a greenlet switch — a block parks exactly as it does for a value it is waiting on the model to produce.
A meeting point for the blocks of one trace.
Built by barrier with
the number of blocks that will call it. Calling it parks the block until that
many have arrived; the last one through releases the rest and none of them
waits again.
n has to be the number of blocks that actually call it, and the two ways to
get it wrong fail differently. Too high and the round never releases: the run
still ends, and the blocks left waiting say so. Too low and the round releases
early, before the block holding the value has arrived, so the block it let
through raises a NameError on a value the barrier was there to sequence.
| ATTRIBUTE | DESCRIPTION |
|---|---|
n |
How many blocks have to arrive before any of them continues.
|