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"
)
LocalHFBackend.post_processingclears three fields on theGenerateDecoderOnlyOutputheld inmot.raw.response, in each case to release GPU memory. These clears are done incorrectly don't appear to be clearing the memory.GenerateDecoderOnlyOutputis atransformers.utils.generic.ModelOutput, i.e. anOrderedDictsubclass that mirrors every field into the mapping.ModelOutput.__setattr__skips the mapping write when the value isNone, andModelOutputdefines no__delattr__. So bothout.f = Noneanddel out.fremove only the__dict__slot; the mapping entry keeps a strong reference. The attribute reads back asNonewhile the tensor stays reachable viaout["f"].Instead, we should route the deletes like:
This test should pass once the deletion changes are made: