feat(backends): populate mot.raw.response on the HD generate_from_raw

[cptnm3]

Pull Request

Issue

Fixes #1331

Description

Each ModelOutputThunk returned by _generate_from_raw previously left mot.raw.response as None, making the HF raw path inconsistent with all other backends and with the HF chat path.

Changes:

• Inside the per-MOT loop, construct a GenerateDecoderOnlyOutput slice for row i using tensor views (no .clone()) so no additional GPU memory is allocated. sequences, scores, and logits are sliced; past_key_values, attentions, and hidden_states are set to None with a one-time debug log.
• After the loop, drop the shared outputs object — null out sequences, scores, logits, past_key_values, attentions, and hidden_states with hasattr guards, then call gc.collect() and torch.cuda.empty_cache(). Debug-log GPU memory before and after. Per-MOT views keep the underlying tensor storage alive via refcounting.

Tests added:

• test_generate_from_raw_raw_response_set_per_mot — asserts raw.response is set, sequences shape is (1, seq_len), storage is shared (view not clone), and omitted fields are None.
• test_generate_from_raw_raw_response_scores_are_views_when_logits_requested — asserts raw.response.scores is a tuple of views sharing storage with the original batch scores when LOGITS=True.
• test_generate_from_raw_raw_response_scores_none_when_logits_not_requested — asserts raw.response.scores is None when model.generate() returns no scores.
• New file test_huggingface_raw_response_copy.py with five tests covering shallow copy (shared raw.response identity and storage) and deepcopy (distinct object, broken storage sharing, preserved values).

Testing

• Tests added to the respective file if code was changed
• New code has 100% coverage if code was added
• Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

Attribution

• AI coding assistants used

---

Adding a new component, requirement, sampling strategy, or tool?

If your PR adds or modifies one of the types below, check the matching box. A checklist of type-specific review items will be posted as a comment.

• Component
• Requirement
• Sampling Strategy
• Tool

NOTE: Please ensure you have an issue that has been acknowledged by a core contributor and routed you to open a pull request against this repository. Otherwise, please open an issue before continuing with this pull request.

[jakelorocco]

Can you please double check the claims that we aren't storing unnecessary tensors / rows? I did a quick investigation (and then had Claude write some tests in https://github.com/jakelorocco/mellea/tree/test/backend-memory-regressions); I believe the whole tensor is saved when we have a single view into it.

Here's the corresponding analysis of issue/test:

Review finding	Test	On PR code
Views pin the whole batch (sequences)	test_raw_response_sequences_retain_only_their_own_row	FAIL — retains 256 B for a 32 B row
Views pin the whole batch (scores)	test_raw_response_scores_retain_only_their_own_row	FAIL — retains 2048 B for a 256 B row
Same, logits/RAW_LOGITS branch (untested by the PR)	test_raw_response_raw_logits_retain_only_their_own_row	FAIL — 1024 B for a 256 B row
Holding one MOT keeps the entire batch alive	test_batch_tensors_are_freed_once_only_mots_are_held	FAIL — batch tensor still alive via weakref
deepcopy duplicates the batch	test_deepcopy_of_result_does_not_duplicate_the_batch	FAIL — 256 B allocated for a 32 B row
Per-call gc.collect() + empty_cache()	test_generate_from_raw_does_not_force_gc_or_cuda_flush_without_cuda	FAIL — 1 full GC pass with no CUDA
Dead isinstance branch emits sequences=None	test_raw_response_is_never_emitted_with_null_sequences	FAIL
del out.f / out.f = None never frees (chat path)	test_post_processing_clearing_raw_logits_actually_releases_them	FAIL — tensors alive after clear
Mislabeled invariant in the PR's test	test_raw_response_scores_follow_generate_output_not_the_logits_option	PASS (documents the real rule)
One-time notice shouldn't spam per item/call	test_omitted_fields_notice_is_logged_once_per_backend	PASS (regression guard)

Maybe we are fine with one mot causing the full tensor to be saved, but that seems excessive to me (unless the fix is complicated / messy).

[planetf1]

Claude review.

Confirming @jakelorocco's finding, and correcting one row of the table.

The retention is real. A row view keeps the whole batch storage alive, not just its row, and nothing releases it: generate_from_raw (mellea/core/backend.py:232) never calls post_processing, and unlike the chat path at :1617 the raw path never sets raw.response = None. This PR's own test shows the size: at batch 2, vocab 32000, fp32, the view's storage is 256000 bytes where a per-row clone is 128000. The ratio is the batch size, so "generate N, keep the best one" retains N times what it needs.

On "maybe we are fine with one mot causing the full tensor to be saved": the file already decided it twice. :1798 reads # Clone each slice so this MOT does not hold a view into the shared batch allocation., and test_generate_from_raw_logits_sliced_per_item at :561 requires logits must be a clone, not a view for this same function. Nor is the fix messy: .detach().clone() on the three slices passes all 64 tests across both HF unit files, ruff format clean.

Correction: the isinstance branch is not dead. Removing it fails four test_multimodal_blocks_in_raw_ctx_not_checked cases, because test_huggingface_unit.py:904 mocks sequences as a plain list, so outputs.sequences[i : i + 1, :] raises TypeError: list indices must be integers or slices, not tuple. Keep the guard, clone inside it. It does read as dead code from the diff; only running it showed otherwise.

Two smaller additions on the cleanup block:

• memory_allocated() can't move here. It reports memory occupied by tensors, while empty_cache() releases unoccupied cached memory. memory_reserved() is the one that would move.
• It's CUDA-only twice over: torch.cuda.empty_cache() is a no-op off CUDA and both debug logs sit behind torch.cuda.is_available(), yet the backend selects cuda, then mps, then cpu at :399-405 and this function already special-cases mps at :1681. On mps or cpu it reduces to a gc.collect().

The two inline comments are on the new tests, which pin the view-sharing as the contract and would invert alongside the fix.

[planetf1]

Correcting my own comment above: "keep the guard, clone inside it" is not sufficient on its own.

[cptnm3]

Thanks @jakelorocco, @planetf1 for the comments. I'm working on incorporating requested changes.

[cptnm3]

Hi @jakelorocco @planetf1,
I have made the requested changes. @jakelorocco All the tests provided by you are green now, except test_post_processing_clearing_raw_logits_actually_releases_them. I have left out the changes needed for that test as my changes were purely on the generate_from_raw path and the required changes would affect the chat path. Please let me know if you would like me to add changes for the test test_post_processing_clearing_raw_logits_actually_releases_them to pass in this PR
Please take a look and let me know if there are any more changes required.
Thanks!

[jakelorocco]

This output deletion code apparently doesn't quite work due to transformer's implementation:

GenerateDecoderOnlyOutput subclasses transformers.utils.generic.ModelOutput, which subclasses OrderedDict. It overrides __setattr__/__setitem__ to keep the attribute and the dict entry in sync, but does not override __delattr__. So del outputs.sequences removes only the __dict__ slot; the OrderedDict entry keeps a strong reference. Verified locally:

del o.sequences
after del: has attr? False
after del: dict entry? True True
o["sequences"] is t: True        # tensor still alive
refcount: 4 -> 3                 # one slot dropped, dict ref retained

[jakelorocco]

I believe you need to do something like outputs["sequences"] = None.

[jakelorocco]

I think you also have to delete previous references like sequences_to_decode from above before the gc / cache clear will work.

So the final version would be something like:

del sequences_to_decode  # views into outputs.sequences
outputs = None           # drops the ModelOutput and its dict entries

[jakelorocco]

I think a test like this would work:

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"
        )

[jakelorocco]

I didn't realize that we are already duplicating logits here. Can we just offer a view into this or do some common extraction so that we don't reduplicate them above?

[jakelorocco]

We should probably also gate this saving / copying logic on _use_caches=True. Otherwise, everyone pays unconditionally for this. I believe this is how the chat completions api is treated.

[planetf1]

NIT: this guard excludes beam-search outputs (GenerateBeamDecoderOnlyOutput), so raw.response stays unset when num_beams is used. Could this support both decoder-only output types, or document the limitation?