diff --git a/mellea/backends/huggingface.py b/mellea/backends/huggingface.py index 627118084..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,16 +1785,75 @@ async def _generate_from_raw( result.generation.provider = self._provider result.raw.provider = self._provider - 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 + # 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. + # 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) + ): + 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") + 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." + ) - 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 + result.raw.response = GenerateDecoderOnlyOutput( + sequences=cast( + "torch.LongTensor", + outputs.sequences[i : i + 1, :].detach().clone(), + ), + 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." + ) + + # Reuse the cloned tensors. + if want_logits and mot_scores is not None: + result.generation.logits = mot_scores + + if want_raw_logits and mot_logits_raw is not None: + result.generation.raw_logits = mot_logits_raw action = actions[i] result.parsed_repr = ( @@ -1794,6 +1872,28 @@ async def _generate_from_raw( result._generate_log = generate_log results.append(result) + # 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 + + 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 reserved", + torch.cuda.memory_reserved(), + ) + 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 e96375061..f9706dd64 100644 --- a/test/backends/test_huggingface_unit.py +++ b/test/backends/test_huggingface_unit.py @@ -1309,3 +1309,195 @@ def _capture_grammar(schema, overrides=None): assert captured[0].get("whitespace_pattern") == r"[\x20\x0A\x0D\x09]{0,20}", ( f"Expected bounded whitespace_pattern to override False in {path_name}" ) + + +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})" + ) + # 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 clone, not a view" + 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_clones_when_logits_requested(): + """raw.response.scores is a tuple of clones when ModelOption.LOGITS is set. + + 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 + 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)" + ) + # 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 clone" + + +@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" + )