Skip to content

v0.8.2

Latest
Latest

Choose a tag to compare

View all tags
@JadenFiotto-Kaufman JadenFiotto-Kaufman released this 02 Nov 19:47
· 16 commits to main since this release
v0.8.2
b0f3b64
This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Verified
Learn about vigilant mode.

0.8.2 is a bug-fix release. It frees GPU memory that stopped traces, failed backward passes and edits kept allocated. It fixes several cases where vLLM returned wrong values without raising, under tensor parallelism and with concurrent requests. It supports tensor parallelism on transformers 5.21, and it pins vLLM to 0.27 and httpx below 1.0.

If you trace with vLLM under tensor parallelism, with LoRA, or with several requests at once, read the vLLM section below: some results from earlier versions may need re-running.

Highlights

Stopped traces, failed backward passes and edits no longer hold GPU memory

  • tracer.stop() (#790). The exception that ends a stopped trace stayed alive in a reference cycle. The cycle held every model frame the stop unwound, and with them those layers' activations and, during generation, the KV cache. GPU allocations don't trigger Python's cyclic garbage collector, so a loop of stopped traces kept growing: Llama-3.2-1B at 16×512 tokens, stopped at block 8, filled a 48 GB A6000 in about a dozen traces. A stopped trace now frees them when it ends. The vLLM runner keeps a stop for error reporting and is not changed by this fix.
  • A backward pass that fails partway (#793). If the backward raised while a with loss.backward(): block was still waiting on a .grad, for example on an out-of-memory error, the block stayed suspended and held the forward's tensors and autograd graph. When the trace ran on a thread that then exited, nothing could free them; on a server replica this held about 65 GiB for a week. The block is now unwound when the backward fails.
  • Edits (#783, @tivris). Values bound inside a model.edit() block stayed referenced until the next trace replaced them. With gpt2 and a 681-token prompt, a one-layer steering edit kept 545 MiB alive between traces. They are now freed when the trace ends.

vLLM: values that came back wrong without an error

If you used any of the setups below on 0.8.1 or earlier, re-run those results on 0.8.2.

Tensor parallelism (tensor_parallel_size > 1)

  • qkv_proj and gate_up_proj pack several projections into each rank's shard, and their gathered output came back in rank order ([q0 k0 v0 | q1 k1 v1]) instead of the single-GPU layout ([q | k | v]). Per-head slices and gate, up = x.chunk(2, -1) were wrong. The gather now returns the single-GPU layout, and edits are split back per rank (#817). On Qwen2.5-0.5B at tp=2, the relative error against tp=1 went from 0.78 to 0.0046 for qkv_proj.output and from 0.99 to 0.0076 for gate_up_proj.output.
  • With fewer KV heads than ranks, K and V came back once per rank. The gather now keeps one copy (#817).
  • With enable_lora=True, a block read only one rank's shard of each linear layer, and its edits went back unsplit. LoRA-wrapped layers are now gathered like the layers they wrap (#820).

Concurrent requests

  • Saves were tracked in one set shared by every request on the engine, so a request could lose a save or receive a name it never saved. In a sustained run, 20 of 60 requests came back with such a name. Each request now has its own set (#804).
  • A block reading logits_processor was given every request's rows, and an edit to one request's row could change another request's token. Reads and edits now see only the block's own row (#813).
  • A list, dict or set that an installed edit (model.edit()) captured from the surrounding code was one object shared by every request. Each request's output also carried every earlier request's entries. Each request's copy of the edit now starts with its own (#812).
  • After a KeyboardInterrupt during a synchronous trace or generate, the interrupted requests kept running and their outputs arrived with the next trace's, so its variables held the interrupted run's values. An interrupted run now aborts its requests (#808).

Tensor parallelism on transformers 5.21

transformers 5.21 shards embed_tokens on every model with a TP plan and allows uneven shards. On 0.8.1, reading embed_tokens.output raised an AssertionError on every sharded model, and a vocabulary that tp_size doesn't divide failed at the next collective. Both work now, and the two new sharding styles, embedding_colwise and replicate, are handled (#791).

Upgrade notes

  • vLLM is pinned to 0.27. The vllm extra now requires vllm>=0.27.1,<0.28 (#801). Without a pin, pip install "nnsight[vllm]" installed vLLM 0.32, where taps stop firing after the prefill without an error. Moving past 0.27 means porting to vLLM's V2 model runner, which is separate work.
  • httpx is capped below 1.0 (#792). The httpx 1.0 previews remove Timeout and AsyncClient, so every remote call failed with an AttributeError.
  • vLLM refuses three configurations when the model is built, each with a NotImplementedError that says what to use instead:
    • speculative_config (#795). A speculative step has several logits rows per request, and a block is scoped to one.
    • data_parallel_size > 1 (#814). On mode="async", requests on any engine but the first came back with no saves and no error. The synchronous engine already refused it.
    • Ray's compiled-graph executor, VLLM_USE_RAY_V2_EXECUTOR_BACKEND=0 (#815). It runs the forward on a different thread, so tracer.result failed and some saves came back unbound. vLLM's default Ray executor is unaffected.
  • Two vLLM behaviour changes (#795):
    • On a structured-output request, a block now reads and edits the logits after vLLM applies the grammar mask.
    • A request whose block raises under ignore_eos now runs to max_tokens with sampled tokens, instead of filling its output with eos tokens.
  • Edit replays (#783). If a model.edit() block reads a name that this trace's run of the block did not bind, it now raises NameError instead of reading the previous trace's value.
  • nnsight.ndif.get_remote_env: force_refresh is now keyword-only (#729).

Other fixes and improvements

vLLM

  • One request's failure no longer costs the other requests in the same collect their values. A block that raises after reading tracer.result (#803) and a saved value that can't be pickled (#810) are now errors on their own request only.
  • A trace's real error is no longer hidden when an installed edit ends with tracer.stop() (#805).
  • A block whose request failed outside the block, or an installed edit cleared mid-request, no longer stays suspended holding its scope (#806). Workers of finished requests are released at the end of each collect instead of staying held while the engine is idle (#807).
  • Closing an async trace's stream early now aborts the request instead of letting it run to max_tokens (#809).
  • nnsight-serve: a long-running server could run an earlier trace's code for a later multi-invoke trace from the same file, and its source caches grew with every client filename. Each build now runs inside nnsight.tracing.globals.block_scope, which restores the caches afterwards (#811).

Remote

  • The remote environment is cached per host, and nnsight.compare(host=...) compares against a specific server. Before, a trace with a per-call host (remote="http://host:port") read the configured server's environment. A trailing slash on a host no longer produces //request URLs (#729).

Batching and performance

  • A .skip() across several invokes now rebuilds the module's output as its own type. A namedtuple output no longer raises TypeError, and a dict subclass such as a transformers ModelOutput is no longer turned into a plain dict (#800, @abhishekcpp007).
  • Preparing each invoke now reads the generate-argument names once per pipeline class instead of on every invoke, which takes 1 µs instead of 83 µs. A 256-invoke trace on an A6000 went from 128.8 ms to 106.3 ms (#779).

Docs

Check your .output[0]. Under transformers 5, GPT-2 and Llama blocks return their hidden states as a plain tensor, not a tuple, so h[i].output[0] is the first batch row. A batched edit written as h[5].output[0][:] = 0 changes only the first prompt and raises no error; write h[5].output[:] = 0. Modules that still return tuples, such as vLLM's layers ((hidden, residual)), are unchanged. The docs and docstrings now use .output throughout (#796).

  • The 0.8 migration page covers two more breaking changes: .nns_output / .nns_input are now .E_output / .E_input, and automodel= no longer selects the model class; use TransformersModel(repo, task=...) (#798).
  • On vLLM, a block under tracer.iter or tracer.all() can see one step past a stop token or stop string, because async scheduling starts the next step before the token is known (#816).

What's Changed

New Contributors

Full Changelog: v0.8.1...v0.8.2

Contributors

JadenFiotto-Kaufman, abhishekcpp007, and tivris
Assets 2
Loading