From f8831f0a1ee5a6e67a28aa570fd410e8f47f2e56 Mon Sep 17 00:00:00 2001 From: Vishal V Date: Mon, 10 Aug 2026 21:23:06 +0530 Subject: [PATCH 1/4] feat(backends): populate mot.raw.response on the generate_from_raw batch path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). Signed-off-by: Vishal V --- mellea/backends/huggingface.py | 60 ++++++ .../test_huggingface_raw_response_copy.py | 90 ++++++++ test/backends/test_huggingface_unit.py | 192 ++++++++++++++++++ 3 files changed, 342 insertions(+) create mode 100644 test/backends/test_huggingface_raw_response_copy.py diff --git a/mellea/backends/huggingface.py b/mellea/backends/huggingface.py index b44af1814..221448ae8 100644 --- a/mellea/backends/huggingface.py +++ b/mellea/backends/huggingface.py @@ -1763,6 +1763,37 @@ async def _generate_from_raw( result.generation.provider = self._provider result.raw.provider = self._provider + # Construct a per-MOT GenerateDecoderOnlyOutput slice using tensor views + # past_key_values, attentions, and hidden_states are omitted + mot_sequences = cast( + "torch.LongTensor", + outputs.sequences[i : i + 1, :] + if isinstance(outputs.sequences, torch.Tensor) + else None, + ) + mot_scores: tuple | None = None + if outputs.scores is not None: + mot_scores = tuple(s[i : i + 1, :] for s in outputs.scores) + mot_logits_raw: tuple | None = None + if outputs.logits is not None: + mot_logits_raw = tuple(s[i : i + 1, :] for s in outputs.logits) + + if "raw_batch_response_fields_omitted" not in self._warned_about: + self._warned_about.add("raw_batch_response_fields_omitted") + MelleaLogger.get_logger().debug( + "mot.raw.response.past_key_values, .attentions, and .hidden_states " + "are not available on the raw batch path and will always be None." + ) + + result.raw.response = GenerateDecoderOnlyOutput( + sequences=mot_sequences, + scores=mot_scores, + logits=mot_logits_raw, + attentions=None, + hidden_states=None, + past_key_values=None, + ) + if want_logits and outputs.scores is not None: # Clone each slice so this MOT does not hold a view into the shared batch allocation. result.generation.logits = tuple( @@ -1791,6 +1822,35 @@ async def _generate_from_raw( result._generate_log = generate_log results.append(result) + # Drop the shared batch output — per-MOT views in raw.response keep the + # underlying tensor storage alive + import gc + + if torch.cuda.is_available(): + MelleaLogger.get_logger().debug( + "GPU memory before raw batch cleanup: %d bytes allocated", + torch.cuda.memory_allocated(), + ) + if hasattr(outputs, "sequences") and outputs.sequences is not None: + del outputs.sequences + if hasattr(outputs, "scores") and outputs.scores is not None: + del outputs.scores + if hasattr(outputs, "logits") and outputs.logits is not None: + del outputs.logits + if hasattr(outputs, "attentions") and outputs.attentions is not None: + del outputs.attentions + if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: + del outputs.hidden_states + if hasattr(outputs, "past_key_values") and outputs.past_key_values is not None: + del outputs.past_key_values + gc.collect() + torch.cuda.empty_cache() + if torch.cuda.is_available(): + MelleaLogger.get_logger().debug( + "GPU memory after raw batch cleanup: %d bytes allocated", + torch.cuda.memory_allocated(), + ) + usage: dict[str, Any] | None = ( { "prompt_tokens": agg_prompt, diff --git a/test/backends/test_huggingface_raw_response_copy.py b/test/backends/test_huggingface_raw_response_copy.py new file mode 100644 index 000000000..7b90983f7 --- /dev/null +++ b/test/backends/test_huggingface_raw_response_copy.py @@ -0,0 +1,90 @@ +# Copyright IBM Corp. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for copy/deepcopy semantics of ModelOutputThunks that hold an HF +GenerateDecoderOnlyOutput in raw.response — the raw batch path specific case.""" + +import copy +from typing import Any + +import pytest + +torch = pytest.importorskip("torch", reason="torch not installed — install mellea[hf]") +pytest.importorskip( + "transformers", reason="transformers not installed — install mellea[hf]" +) +pytest.importorskip( + "llguidance", reason="llguidance not installed — install mellea[hf]" +) + +from transformers.generation.utils import GenerateDecoderOnlyOutput + +from mellea.core import ModelOutputThunk + + +def _make_mot_with_hf_raw_response() -> tuple[ModelOutputThunk, Any]: + """Build a MOT whose raw.response is a CPU-only GenerateDecoderOnlyOutput with a view.""" + full_batch = torch.arange(6, dtype=torch.long).reshape(2, 3) + # Simulate the view produced by the raw batch path for batch item 0. + row_view = full_batch[0:1, :] + hf_out = GenerateDecoderOnlyOutput( + sequences=row_view, + scores=None, + logits=None, + attentions=None, + hidden_states=None, + past_key_values=None, + ) + mot = ModelOutputThunk(value="hello") + mot.raw.response = hf_out + return mot, full_batch + + +def test_shallow_copy_raw_response_is_same_object(): + """copy.copy(mot).raw.response is the same object as mot.raw.response.""" + mot, _ = _make_mot_with_hf_raw_response() + copied = copy.copy(mot) + assert copied.raw.response is mot.raw.response, ( + "shallow copy must keep raw.response as the same object" + ) + + +def test_shallow_copy_raw_response_sequences_shares_storage(): + """Shallow-copied MOT preserves shared tensor storage for raw.response.sequences.""" + mot, full_batch = _make_mot_with_hf_raw_response() + copied = copy.copy(mot) + assert ( + copied.raw.response.sequences.untyped_storage().data_ptr() + == full_batch.untyped_storage().data_ptr() + ), ( + "shallow copy: raw.response.sequences must still share storage with the original batch" + ) + + +def test_deepcopy_raw_response_is_distinct_object(): + """copy.deepcopy(mot).raw.response is a distinct object from mot.raw.response.""" + mot, _ = _make_mot_with_hf_raw_response() + deep = copy.deepcopy(mot) + assert deep.raw.response is not mot.raw.response, ( + "deepcopy must produce a new raw.response object" + ) + + +def test_deepcopy_raw_response_sequences_does_not_share_storage(): + """Deepcopy breaks tensor storage sharing for raw.response.sequences.""" + mot, full_batch = _make_mot_with_hf_raw_response() + deep = copy.deepcopy(mot) + assert ( + deep.raw.response.sequences.untyped_storage().data_ptr() + != full_batch.untyped_storage().data_ptr() + ), "deepcopy: raw.response.sequences must NOT share storage with the original" + + +def test_deepcopy_raw_response_sequences_preserves_values(): + """Deepcopy preserves tensor values in raw.response.sequences despite storage isolation.""" + mot, _ = _make_mot_with_hf_raw_response() + original_values = mot.raw.response.sequences.clone() + deep = copy.deepcopy(mot) + assert torch.equal(deep.raw.response.sequences, original_values), ( + "deepcopy must preserve tensor values in raw.response.sequences" + ) diff --git a/test/backends/test_huggingface_unit.py b/test/backends/test_huggingface_unit.py index d76f5d375..770523f94 100644 --- a/test/backends/test_huggingface_unit.py +++ b/test/backends/test_huggingface_unit.py @@ -977,3 +977,195 @@ async def test_multimodal_blocks_in_intrinsic_ctx_raise_error( await LocalHFBackend._generate_from_intrinsic( backend, Intrinsic("answerability"), ctx, model_options={} ) + + +def _make_raw_fake_setup( + batch_size: int, vocab_size: int, n_tokens: int, prompt_len: int +): + """Return (backend, fake_encoding, fake_outputs, actions) for generate_from_raw tests.""" + backend = _make_backend() + fake_input_ids = torch.zeros(batch_size, prompt_len, dtype=torch.long) + fake_encoding = MagicMock() + fake_encoding.__getitem__ = lambda self, k: ( + fake_input_ids + if k == "input_ids" + else torch.ones(batch_size, prompt_len, dtype=torch.long) + ) + fake_encoding.to = MagicMock(return_value=fake_encoding) + backend._tokenizer = MagicMock(eos_token_id=0, vocab_size=vocab_size) + backend._tokenizer.__len__ = MagicMock(return_value=vocab_size) + backend._tokenizer.return_value = fake_encoding + decode_values = [f"result_{chr(ord('a') + i)}" for i in range(batch_size)] + backend._tokenizer.batch_decode = MagicMock(return_value=decode_values) + return backend, fake_encoding, fake_input_ids + + +@pytest.mark.asyncio +async def test_generate_from_raw_raw_response_set_per_mot(): + """Every MOT from generate_from_raw has raw.response set to a GenerateDecoderOnlyOutput. + + Asserts: + - raw.response is not None for each MOT. + - raw.response.sequences.shape == (1, full_seq_len). + - raw.response.sequences shares storage with the original batch sequences tensor (view, not clone). + - raw.response.past_key_values is None. + - raw.response.attentions is None. + - raw.response.hidden_states is None. + """ + batch_size = 2 + vocab_size = 32000 + n_tokens = 3 + prompt_len = 1 + full_seq_len = prompt_len + n_tokens + + backend, _fake_encoding, _fake_input_ids = _make_raw_fake_setup( + batch_size, vocab_size, n_tokens, prompt_len + ) + sequences = torch.zeros(batch_size, full_seq_len, dtype=torch.long) + fake_outputs = GenerateDecoderOnlyOutput( + sequences=sequences, + scores=None, + logits=None, + attentions=None, + hidden_states=None, + past_key_values=None, + ) + actions = [Message("user", "hello"), Message("user", "world")] + + with ( + patch( + "mellea.backends.huggingface.asyncio.to_thread", return_value=fake_outputs + ), + patch.object(backend, "do_generate_walks"), + patch.object(backend, "formatter") as mock_fmt, + ): + mock_fmt.print = MagicMock(return_value="prompt") + results = await backend.generate_from_raw( + actions, MagicMock(), model_options={} + ) + + assert len(results) == batch_size + for item_idx, result in enumerate(results): + assert result.raw.response is not None, ( + f"item {item_idx}: raw.response must be set" + ) + assert isinstance(result.raw.response, GenerateDecoderOnlyOutput), ( + f"item {item_idx}: raw.response must be GenerateDecoderOnlyOutput" + ) + assert result.raw.response.sequences.shape == (1, full_seq_len), ( + f"item {item_idx}: sequences shape must be (1, {full_seq_len})" + ) + # View — must share the same underlying storage as the original batch tensor. + assert ( + result.raw.response.sequences.untyped_storage().data_ptr() + == sequences.untyped_storage().data_ptr() + ), f"item {item_idx}: sequences must be a view, not a clone" + assert result.raw.response.past_key_values is None, ( + f"item {item_idx}: past_key_values must be None" + ) + assert result.raw.response.attentions is None, ( + f"item {item_idx}: attentions must be None" + ) + assert result.raw.response.hidden_states is None, ( + f"item {item_idx}: hidden_states must be None" + ) + + +@pytest.mark.asyncio +async def test_generate_from_raw_raw_response_scores_are_views_when_logits_requested(): + """raw.response.scores is a tuple of views when ModelOption.LOGITS is set. + + Each tensor in raw.response.scores must share storage with the corresponding + step tensor in the original batch scores (view, not clone), distinct from + generation.logits which holds clones. + """ + batch_size = 2 + vocab_size = 32000 + n_tokens = 3 + prompt_len = 1 + full_seq_len = prompt_len + n_tokens + + backend, _fake_encoding, _fake_input_ids = _make_raw_fake_setup( + batch_size, vocab_size, n_tokens, prompt_len + ) + sequences = torch.zeros(batch_size, full_seq_len, dtype=torch.long) + fake_scores = tuple(torch.randn(batch_size, vocab_size) for _ in range(n_tokens)) + fake_outputs = GenerateDecoderOnlyOutput( + sequences=sequences, + scores=fake_scores, + logits=None, + attentions=None, + hidden_states=None, + past_key_values=None, + ) + actions = [Message("user", "hello"), Message("user", "world")] + + with ( + patch( + "mellea.backends.huggingface.asyncio.to_thread", return_value=fake_outputs + ), + patch.object(backend, "do_generate_walks"), + patch.object(backend, "formatter") as mock_fmt, + ): + mock_fmt.print = MagicMock(return_value="prompt") + results = await backend.generate_from_raw( + actions, MagicMock(), model_options={ModelOption.LOGITS: True} + ) + + for item_idx, result in enumerate(results): + assert result.raw.response.scores is not None, ( + f"item {item_idx}: raw.response.scores must be set when LOGITS=True" + ) + assert len(result.raw.response.scores) == n_tokens, ( + f"item {item_idx}: one scores tensor per generation step" + ) + for tok_idx, t in enumerate(result.raw.response.scores): + assert t.shape == (1, vocab_size), ( + f"item {item_idx} token {tok_idx}: shape must be (1, vocab_size)" + ) + # View — must share storage with the original batch step tensor. + assert ( + t.untyped_storage().data_ptr() + == fake_scores[tok_idx].untyped_storage().data_ptr() + ), f"item {item_idx} token {tok_idx}: raw.response.scores must be a view" + + +@pytest.mark.asyncio +async def test_generate_from_raw_raw_response_scores_none_when_logits_not_requested(): + """raw.response.scores is None when ModelOption.LOGITS is not set.""" + batch_size = 1 + vocab_size = 32000 + n_tokens = 2 + prompt_len = 1 + full_seq_len = prompt_len + n_tokens + + backend, _fake_encoding, _fake_input_ids = _make_raw_fake_setup( + batch_size, vocab_size, n_tokens, prompt_len + ) + # When LOGITS is not set, model.generate() is called without output_scores=True, + # so outputs.scores will be None — simulate that here. + sequences = torch.zeros(batch_size, full_seq_len, dtype=torch.long) + fake_outputs = GenerateDecoderOnlyOutput( + sequences=sequences, + scores=None, + logits=None, + attentions=None, + hidden_states=None, + past_key_values=None, + ) + + with ( + patch( + "mellea.backends.huggingface.asyncio.to_thread", return_value=fake_outputs + ), + patch.object(backend, "do_generate_walks"), + patch.object(backend, "formatter") as mock_fmt, + ): + mock_fmt.print = MagicMock(return_value="prompt") + results = await backend.generate_from_raw( + [Message("user", "hi")], MagicMock(), model_options={} + ) + + assert results[0].raw.response.scores is None, ( + "raw.response.scores must be None when model.generate() returns no scores" + ) From f73166d7d834329a1d0ea6bab42e90f3ad65b7bd Mon Sep 17 00:00:00 2001 From: Vishal V <56761954+cptnm3@users.noreply.github.com> Date: Fri, 14 Aug 2026 01:50:50 +0530 Subject: [PATCH 2/4] fixup!: Apply suggestions from code review Co-authored-by: Nigel Jones Signed-off-by: Vishal V <56761954+cptnm3@users.noreply.github.com> --- mellea/backends/huggingface.py | 63 ++++++++++++++------------ test/backends/test_huggingface_unit.py | 12 ++--- 2 files changed, 40 insertions(+), 35 deletions(-) diff --git a/mellea/backends/huggingface.py b/mellea/backends/huggingface.py index 221448ae8..e24bd828d 100644 --- a/mellea/backends/huggingface.py +++ b/mellea/backends/huggingface.py @@ -1763,36 +1763,41 @@ async def _generate_from_raw( result.generation.provider = self._provider result.raw.provider = self._provider - # Construct a per-MOT GenerateDecoderOnlyOutput slice using tensor views - # past_key_values, attentions, and hidden_states are omitted - mot_sequences = cast( - "torch.LongTensor", - outputs.sequences[i : i + 1, :] - if isinstance(outputs.sequences, torch.Tensor) - else None, - ) - mot_scores: tuple | None = None - if outputs.scores is not None: - mot_scores = tuple(s[i : i + 1, :] for s in outputs.scores) - mot_logits_raw: tuple | None = None - if outputs.logits is not None: - mot_logits_raw = tuple(s[i : i + 1, :] for s in outputs.logits) - - if "raw_batch_response_fields_omitted" not in self._warned_about: - self._warned_about.add("raw_batch_response_fields_omitted") - MelleaLogger.get_logger().debug( - "mot.raw.response.past_key_values, .attentions, and .hidden_states " - "are not available on the raw batch path and will always be None." - ) + # Construct a per-MOT GenerateDecoderOnlyOutput slice holding clones of row i. + # past_key_values, attentions, and hidden_states are shared across the batch + # and cannot be sliced per MOT, so they are omitted. + if isinstance(outputs.sequences, torch.Tensor): + mot_scores: tuple[torch.Tensor, ...] | None = None + if outputs.scores is not None: + mot_scores = tuple( + s[i : i + 1, :].detach().clone() for s in outputs.scores + ) + mot_logits_raw: tuple[torch.Tensor, ...] | None = None + if outputs.logits is not None: + mot_logits_raw = tuple( + step_logits[i : i + 1, :].detach().clone() + for step_logits in outputs.logits + ) - result.raw.response = GenerateDecoderOnlyOutput( - sequences=mot_sequences, - scores=mot_scores, - logits=mot_logits_raw, - attentions=None, - hidden_states=None, - past_key_values=None, - ) + if "raw_batch_response_fields_omitted" not in self._warned_about: + self._warned_about.add("raw_batch_response_fields_omitted") + MelleaLogger.get_logger().debug( + "mot.raw.response.past_key_values, .attentions, and " + ".hidden_states are not available on the raw batch path " + "and will always be None." + ) + + result.raw.response = GenerateDecoderOnlyOutput( + sequences=cast( + "torch.LongTensor", + outputs.sequences[i : i + 1, :].detach().clone(), + ), + scores=mot_scores, + logits=mot_logits_raw, + attentions=None, + hidden_states=None, + past_key_values=None, + ) if want_logits and outputs.scores is not None: # Clone each slice so this MOT does not hold a view into the shared batch allocation. diff --git a/test/backends/test_huggingface_unit.py b/test/backends/test_huggingface_unit.py index 770523f94..d3f9367ef 100644 --- a/test/backends/test_huggingface_unit.py +++ b/test/backends/test_huggingface_unit.py @@ -1055,11 +1055,11 @@ async def test_generate_from_raw_raw_response_set_per_mot(): assert result.raw.response.sequences.shape == (1, full_seq_len), ( f"item {item_idx}: sequences shape must be (1, {full_seq_len})" ) - # View — must share the same underlying storage as the original batch tensor. + # Clone - must NOT share storage with the original batch tensor. assert ( result.raw.response.sequences.untyped_storage().data_ptr() - == sequences.untyped_storage().data_ptr() - ), f"item {item_idx}: sequences must be a view, not a clone" + != sequences.untyped_storage().data_ptr() + ), f"item {item_idx}: sequences must be a clone, not a view" assert result.raw.response.past_key_values is None, ( f"item {item_idx}: past_key_values must be None" ) @@ -1123,11 +1123,11 @@ async def test_generate_from_raw_raw_response_scores_are_views_when_logits_reque assert t.shape == (1, vocab_size), ( f"item {item_idx} token {tok_idx}: shape must be (1, vocab_size)" ) - # View — must share storage with the original batch step tensor. + # Clone - must NOT share storage with the original batch step tensor. assert ( t.untyped_storage().data_ptr() - == fake_scores[tok_idx].untyped_storage().data_ptr() - ), f"item {item_idx} token {tok_idx}: raw.response.scores must be a view" + != fake_scores[tok_idx].untyped_storage().data_ptr() + ), f"item {item_idx} token {tok_idx}: raw.response.scores must be a clone" @pytest.mark.asyncio From b2b484494ca34fd413ae800c08400492245731ef Mon Sep 17 00:00:00 2001 From: Vishal V Date: Sun, 16 Aug 2026 17:27:42 +0530 Subject: [PATCH 3/4] =?UTF-8?q?fixup!:=20populate=20mot.raw.response=20on?= =?UTF-8?q?=20the=20generate=5Ffrom=5Fraw=20batch=20path=20Changes:=20-=20?= =?UTF-8?q?Added=20isinstance(outputs,=20GenerateDecoderOnlyOutput)=20to?= =?UTF-8?q?=20the=20outer=20guard=20so=20beam-search=20output=20(GenerateB?= =?UTF-8?q?eamDecoderOnlyOutput)=20never=20gets=20silently=20mislabelled?= =?UTF-8?q?=20-=20Moved=20gc.collect()=20and=20torch.cuda.empty=5Fcache()?= =?UTF-8?q?=20inside=20torch.cuda.is=5Favailable().=20-=20Renamed=20test?= =?UTF-8?q?=5Fgenerate=5Ffrom=5Fraw=5Fraw=5Fresponse=5Fscores=5Fare=5Fview?= =?UTF-8?q?s=5Fwhen=5Flogits=5Frequested=20=E2=86=92=20test=5Fgenerate=5Ff?= =?UTF-8?q?rom=5Fraw=5Fraw=5Fresponse=5Fscores=5Fare=5Fclones=5Fwhen=5Flog?= =?UTF-8?q?its=5Frequested=20and=20updated=20its=20docstring=20to=20reflec?= =?UTF-8?q?t=20that=20raw.response.scores=20holds=20clones=20(not=20views)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Vishal V --- mellea/backends/huggingface.py | 35 +++++++++++++------------- test/backends/test_huggingface_unit.py | 10 ++++---- 2 files changed, 23 insertions(+), 22 deletions(-) diff --git a/mellea/backends/huggingface.py b/mellea/backends/huggingface.py index e24bd828d..3d8e0d2c1 100644 --- a/mellea/backends/huggingface.py +++ b/mellea/backends/huggingface.py @@ -1766,8 +1766,10 @@ async def _generate_from_raw( # Construct a per-MOT GenerateDecoderOnlyOutput slice holding clones of row i. # past_key_values, attentions, and hidden_states are shared across the batch # and cannot be sliced per MOT, so they are omitted. - if isinstance(outputs.sequences, torch.Tensor): - mot_scores: tuple[torch.Tensor, ...] | None = None + if isinstance(outputs, GenerateDecoderOnlyOutput) and isinstance( + outputs.sequences, torch.Tensor + ): + mot_scores: tuple | None = None if outputs.scores is not None: mot_scores = tuple( s[i : i + 1, :].detach().clone() for s in outputs.scores @@ -1792,8 +1794,8 @@ async def _generate_from_raw( "torch.LongTensor", outputs.sequences[i : i + 1, :].detach().clone(), ), - scores=mot_scores, - logits=mot_logits_raw, + scores=cast("tuple[torch.FloatTensor] | None", mot_scores), + logits=cast("tuple[torch.FloatTensor] | None", mot_logits_raw), attentions=None, hidden_states=None, past_key_values=None, @@ -1827,15 +1829,8 @@ async def _generate_from_raw( result._generate_log = generate_log results.append(result) - # Drop the shared batch output — per-MOT views in raw.response keep the - # underlying tensor storage alive - import gc - - if torch.cuda.is_available(): - MelleaLogger.get_logger().debug( - "GPU memory before raw batch cleanup: %d bytes allocated", - torch.cuda.memory_allocated(), - ) + # Drop the shared batch output — per-MOT clones in raw.response own their + # own storage, so releasing the batch frees its memory. if hasattr(outputs, "sequences") and outputs.sequences is not None: del outputs.sequences if hasattr(outputs, "scores") and outputs.scores is not None: @@ -1848,12 +1843,18 @@ async def _generate_from_raw( del outputs.hidden_states if hasattr(outputs, "past_key_values") and outputs.past_key_values is not None: del outputs.past_key_values - gc.collect() - torch.cuda.empty_cache() if torch.cuda.is_available(): + import gc + + MelleaLogger.get_logger().debug( + "GPU memory before raw batch cleanup: %d bytes reserved", + torch.cuda.memory_reserved(), + ) + gc.collect() + torch.cuda.empty_cache() MelleaLogger.get_logger().debug( - "GPU memory after raw batch cleanup: %d bytes allocated", - torch.cuda.memory_allocated(), + "GPU memory after raw batch cleanup: %d bytes reserved", + torch.cuda.memory_reserved(), ) usage: dict[str, Any] | None = ( diff --git a/test/backends/test_huggingface_unit.py b/test/backends/test_huggingface_unit.py index d3f9367ef..9b33343ab 100644 --- a/test/backends/test_huggingface_unit.py +++ b/test/backends/test_huggingface_unit.py @@ -1072,12 +1072,12 @@ async def test_generate_from_raw_raw_response_set_per_mot(): @pytest.mark.asyncio -async def test_generate_from_raw_raw_response_scores_are_views_when_logits_requested(): - """raw.response.scores is a tuple of views when ModelOption.LOGITS is set. +async def test_generate_from_raw_raw_response_scores_are_clones_when_logits_requested(): + """raw.response.scores is a tuple of clones when ModelOption.LOGITS is set. - Each tensor in raw.response.scores must share storage with the corresponding - step tensor in the original batch scores (view, not clone), distinct from - generation.logits which holds clones. + Each tensor in raw.response.scores must own compact per-row storage and must + not share storage with the corresponding batch step tensor — consistent with + generation.logits which also holds clones. """ batch_size = 2 vocab_size = 32000 From eb0ea8cd17dd16553cdbe24aa642b1a3932ce98f Mon Sep 17 00:00:00 2001 From: Vishal V Date: Wed, 19 Aug 2026 00:41:36 +0530 Subject: [PATCH 4/4] fixup!: add _use_cache guard before constructing per-mot object - Update the deletion code to drop the container properly specific to transformer"s implementation - document GenerateBeamDecoderOnlyOutput does not populate mot.raw.response Signed-off-by: Vishal V --- mellea/backends/huggingface.py | 138 ++++++++++++++++++++------------- 1 file changed, 86 insertions(+), 52 deletions(-) diff --git a/mellea/backends/huggingface.py b/mellea/backends/huggingface.py index 914bdb0f6..e4fab041f 100644 --- a/mellea/backends/huggingface.py +++ b/mellea/backends/huggingface.py @@ -1525,12 +1525,25 @@ class used during generation, if any. self.cache_put(cache_key, cache_info) # Clear KV cache and scores from HF output; retained via LRU cache above. - hf_output.past_key_values = None - hf_output.scores = None + # `ModelOutput` (`OrderedDict` subclass) does not sync `None` writes back + # to the mapping, so plain attribute assignment leaves the dict entry — and + # its tensor — alive. + for _field in ("past_key_values", "scores"): + if _field in hf_output: + dict.__delitem__(hf_output, _field) + try: + object.__delattr__(hf_output, _field) + except AttributeError: + pass # Clear the raw logits tensor (scores already cleared above if cached). if isinstance(hf_output, GenerateDecoderOnlyOutput): - hf_output.logits = None + if "logits" in hf_output: + dict.__delitem__(hf_output, "logits") + try: + object.__delattr__(hf_output, "logits") + except AttributeError: + pass # Only scan for tools if we are not doing structured output and tool calls were provided to the model. if _format is None and tool_calls: @@ -1608,20 +1621,26 @@ class used during generation, if any. if not self._use_caches and isinstance( mot.raw.response, GenerateDecoderOnlyOutput ): - import gc - hf_out = mot.raw.response - if hasattr(hf_out, "sequences") and hf_out.sequences is not None: - del hf_out.sequences - if hasattr(hf_out, "scores") and hf_out.scores is not None: - del hf_out.scores - if hasattr(hf_out, "logits") and hf_out.logits is not None: - del hf_out.logits + # GenerateDecoderOnlyOutput is a ModelOutput (OrderedDict subclass). + # ModelOutput.__setattr__ skips the dict write when value is None, and + # ModelOutput defines no __delattr__, so both `out.f = None` and + # `del out.f` leave the mapping entry — and its tensor — alive. + # Clear both the dict entry and the instance attribute to release them. + for field in ("sequences", "scores", "logits"): + if field in hf_out: + dict.__delitem__(hf_out, field) + try: + object.__delattr__(hf_out, field) + except AttributeError: + pass mot.raw.response = None - # Force Python GC and return CUDA memory to device - gc.collect() - torch.cuda.empty_cache() + if torch.cuda.is_available(): + import gc + + gc.collect() + torch.cuda.empty_cache() # Generate the log for this ModelOutputThunk. generate_log = GenerateLog() @@ -1766,23 +1785,40 @@ async def _generate_from_raw( result.generation.provider = self._provider result.raw.provider = self._provider + # Extract per-MOT tensors once. These are shared by generation + # and, where supported, raw.response. + mot_scores: tuple[torch.Tensor, ...] | None = None + if outputs.scores is not None: + mot_scores = tuple( + score[i].detach().clone() for score in outputs.scores + ) + + mot_logits_raw: tuple[torch.Tensor, ...] | None = None + if outputs.logits is not None: + mot_logits_raw = tuple( + logits[i].detach().clone() for logits in outputs.logits + ) + # Construct a per-MOT GenerateDecoderOnlyOutput slice holding clones of row i. # past_key_values, attentions, and hidden_states are shared across the batch # and cannot be sliced per MOT, so they are omitted. - if isinstance(outputs, GenerateDecoderOnlyOutput) and isinstance( - outputs.sequences, torch.Tensor + # Note: beam-search outputs (GenerateBeamDecoderOnlyOutput) are not supported here; + # mot.raw.response will be None. + if ( + self._use_caches + and isinstance(outputs, GenerateDecoderOnlyOutput) + and isinstance(outputs.sequences, torch.Tensor) ): - mot_scores: tuple | None = None - if outputs.scores is not None: - mot_scores = tuple( - s[i : i + 1, :].detach().clone() for s in outputs.scores - ) - mot_logits_raw: tuple[torch.Tensor, ...] | None = None - if outputs.logits is not None: - mot_logits_raw = tuple( - step_logits[i : i + 1, :].detach().clone() - for step_logits in outputs.logits - ) + response_scores = ( + tuple(score.unsqueeze(0) for score in mot_scores) + if mot_scores is not None + else None + ) + response_logits = ( + tuple(logits.unsqueeze(0) for logits in mot_logits_raw) + if mot_logits_raw is not None + else None + ) if "raw_batch_response_fields_omitted" not in self._warned_about: self._warned_about.add("raw_batch_response_fields_omitted") @@ -1797,23 +1833,27 @@ async def _generate_from_raw( "torch.LongTensor", outputs.sequences[i : i + 1, :].detach().clone(), ), - scores=cast("tuple[torch.FloatTensor] | None", mot_scores), - logits=cast("tuple[torch.FloatTensor] | None", mot_logits_raw), + scores=cast("tuple[torch.FloatTensor] | None", response_scores), + logits=cast("tuple[torch.FloatTensor] | None", response_logits), attentions=None, hidden_states=None, past_key_values=None, ) + elif self._use_caches: + # Beam search slicing is not feasible — raw.response is not populated. + warn_key = "raw_batch_beam_search_unsupported" + if warn_key not in self._warned_about: + self._warned_about.add(warn_key) + MelleaLogger.get_logger().debug( + "mot.raw.response is not available on the raw batch and will be None." + ) - if want_logits and outputs.scores is not None: - # Clone each slice so this MOT does not hold a view into the shared batch allocation. - result.generation.logits = tuple( - step_scores[i].detach().clone() for step_scores in outputs.scores - ) + # Reuse the cloned tensors. + if want_logits and mot_scores is not None: + result.generation.logits = mot_scores - if want_raw_logits and outputs.logits is not None: - result.generation.raw_logits = tuple( - step_logits[i].detach().clone() for step_logits in outputs.logits - ) + if want_raw_logits and mot_logits_raw is not None: + result.generation.raw_logits = mot_logits_raw action = actions[i] result.parsed_repr = ( @@ -1832,20 +1872,14 @@ async def _generate_from_raw( result._generate_log = generate_log results.append(result) - # Drop the shared batch output — per-MOT clones in raw.response own their - # own storage, so releasing the batch frees its memory. - if hasattr(outputs, "sequences") and outputs.sequences is not None: - del outputs.sequences - if hasattr(outputs, "scores") and outputs.scores is not None: - del outputs.scores - if hasattr(outputs, "logits") and outputs.logits is not None: - del outputs.logits - if hasattr(outputs, "attentions") and outputs.attentions is not None: - del outputs.attentions - if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: - del outputs.hidden_states - if hasattr(outputs, "past_key_values") and outputs.past_key_values is not None: - del outputs.past_key_values + # Drop all references that might pin the shared batch tensors. + # `sequences_to_decode` holds slice views of `outputs.sequences`, and + # `outputs` is a `GenerateDecoderOnlyOutput` — a `ModelOutput`, which + # subclasses `OrderedDict`. `del obj.attr` only removes the `__dict__` + # slot; the `OrderedDict` entry keeps a strong reference to the tensor. + # Setting `outputs = None` drops the whole container at once + del sequences_to_decode + outputs = None if torch.cuda.is_available(): import gc