Skip to content

bug(huggingface): nulling ModelOutput fields never frees the tensors (KV cache stays on the MOT) on chat path #1549

Description

@jakelorocco

LocalHFBackend.post_processing clears three fields on the GenerateDecoderOnlyOutput held in mot.raw.response, in each case to release GPU memory. These clears are done incorrectly don't appear to be clearing the memory.

GenerateDecoderOnlyOutput is a transformers.utils.generic.ModelOutput, i.e. an OrderedDict subclass that mirrors every field into the mapping. ModelOutput.__setattr__ skips the mapping write when the value is None, and ModelOutput defines no __delattr__. So both out.f = None and del out.f remove only the __dict__ slot; the mapping entry keeps a strong reference. The attribute reads back as None while the tensor stays reachable via out["f"].

Instead, we should route the deletes like:

hf_output["past_key_values"] = None
hf_output["scores"] = None
hf_output["logits"] = None
...

This test should pass once the deletion changes are made:

async def test_post_processing_clearing_raw_logits_actually_releases_them():
    """Clearing `hf_output.logits` must drop the tensors, not just the attribute.

    `GenerateDecoderOnlyOutput` is a `ModelOutput`, i.e. an `OrderedDict` subclass
    that mirrors every field into the mapping. `ModelOutput.__setattr__` skips the
    mapping write when the value is `None`, and `ModelOutput` defines no
    `__delattr__`, so `out.logits = None` and `del out.logits` both leave the
    mapping entry — and therefore the tensors — in place. Any code that nulls a
    field to free memory while keeping the container has to clear the mapping too.
    """
    backend = _make_backend(1)
    backend._use_caches = True  # keeps raw.response, so the container survives

    mot, refs = await _post_process_holding_only_weakrefs(backend, n_steps=2)
    gc.collect()
    gc.collect()

    assert mot.raw.response is not None, "test setup: raw.response should be retained"
    assert mot.raw.response.logits is None, "test setup: logits attribute was cleared"
    for step, ref in enumerate(refs["logits"]):
        assert ref() is None, (
            f"raw logits tensor for step {step} is still alive after hf_output.logits "
            "was set to None — the ModelOutput mapping entry still references it"
        )

Metadata

Metadata

Labels

area/backendsProvider-specific work: Ollama, HF, LiteLLM, OpenAI, Bedrock, vLLMbugSomething isn't workingp1High: important bugs (workaround exists) or high-value core features. Do soon, not on fire.

Type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions