diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index c6c2e1fc37..022c796ac4 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -333,6 +333,20 @@ jobs: run: | python -m pytest -v --tb=short tests/test_callback_signature_drift.py + - name: batched left-padding generation guard (HARD GATE) + # Guards _fast_prepare_inputs_for_generation against the bug class of + # issues #1066 / #3699: position_ids taken from cache_position (which + # counts left-pad tokens) or the 2D attention mask truncated to its + # last column. Both shipped in cc4c5d77 and were fixed by #2216 and + # #4100; nothing tested this path, so each regression reached users. + # Layer 1 in the file is stdlib-ast-only (survives unsloth import + # breakage), layer 2 calls the real function on CPU via the + # tests/conftest.py CUDA spoof. Validated to fail on the pre-#2216 + # and pre-#4100 code states; staging proof on GPU-less runners: + # danielhanchen/unsloth-staging-2 PR 170 (green, gate passed in all combos) / PR 172 (red, gate failed in all combos). + run: | + python -m pytest -v --tb=short tests/utils/test_prepare_inputs_leftpad.py + - name: unsloth Bucket-A — CPU tests not in Repo tests (CPU) # CPU tests across 6 files under tests/saving/, tests/utils/, tests/python/ # that Repo tests (CPU) --ignores. AST/protobuf/regex plus tiny CPU model diff --git a/tests/utils/test_batched_leftpad_generation_gpu.py b/tests/utils/test_batched_leftpad_generation_gpu.py new file mode 100644 index 0000000000..c2610d4673 --- /dev/null +++ b/tests/utils/test_batched_leftpad_generation_gpu.py @@ -0,0 +1,107 @@ +"""End-to-end GPU guard for batched left-padded generation (issues #1066, #3699). + +For each prompt, greedy generation inside a left-padded batch must match +generating that prompt alone at batch size 1 for the first PREFIX_TOKENS +tokens, and the full output must not be gibberish. The bug class (#1066, +#3699) makes padded rows diverge immediately into garbage; in contrast, +benign batch-size-dependent kernel numerics can flip a greedy near-tie deep +into the sequence, so an exact full-length match would be flaky. Uses a small +instruct model (chat-templated prompts have high-margin argmaxes). + +Skipped automatically when CUDA is unavailable, so CPU CI is unaffected. +Run manually on any GPU box: + + python -m pytest tests/utils/test_batched_leftpad_generation_gpu.py -v +""" + +import pytest +import torch + +cuda_available = torch.cuda.is_available() + +pytestmark = pytest.mark.skipif(not cuda_available, reason = "requires a CUDA GPU") + +MODEL_NAME = "unsloth/Qwen2.5-0.5B-Instruct" +MAX_NEW_TOKENS = 32 +PREFIX_TOKENS = 16 + +PROMPTS = [ + "Give me a short introduction to large language model.", + "Here is an experiment log: " + + " ".join(f"run {i} completed with stable throughput and no anomalies;" for i in range(1, 41)) + + " In one sentence, what is the overall conclusion?", +] + + +@pytest.fixture(scope = "module") +def model_and_tokenizer(): + from unsloth import FastLanguageModel + + model, tokenizer = FastLanguageModel.from_pretrained( + model_name = MODEL_NAME, + max_seq_length = 2048, + load_in_4bit = True, + ) + FastLanguageModel.for_inference(model) + tokenizer.padding_side = "left" + if tokenizer.pad_token_id is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + return model, tokenizer + + +def _chat(tokenizer, prompt): + return tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize = False, + add_generation_prompt = True, + ) + + +def _generate(model, tokenizer, texts): + inputs = tokenizer(texts, return_tensors = "pt", padding = True, add_special_tokens = False).to( + "cuda" + ) + with torch.inference_mode(): + out = model.generate( + **inputs, + max_new_tokens = MAX_NEW_TOKENS, + do_sample = False, + temperature = None, + top_p = None, + top_k = None, + use_cache = True, + pad_token_id = tokenizer.pad_token_id, + ) + suffixes = out[:, inputs["input_ids"].shape[1] :] + return [row.tolist() for row in suffixes] + + +def _looks_gibberish(text): + if not text.strip(): + return True + exclam = text.count("!") / max(len(text), 1) + nonascii = sum(1 for c in text if ord(c) > 0x2FFF) / max(len(text), 1) + return exclam > 0.3 or nonascii > 0.5 + + +def test_batched_leftpad_matches_solo_generation(model_and_tokenizer): + model, tokenizer = model_and_tokenizer + + texts = [_chat(tokenizer, p) for p in PROMPTS] + solo = [_generate(model, tokenizer, [t])[0] for t in texts] + batched = _generate(model, tokenizer, texts) + + for i, prompt in enumerate(PROMPTS): + solo_text = tokenizer.decode(solo[i], skip_special_tokens = True) + batch_text = tokenizer.decode(batched[i], skip_special_tokens = True) + assert batched[i][:PREFIX_TOKENS] == solo[i][:PREFIX_TOKENS], ( + f"prompt {i} ({prompt[:30]!r}...) diverged from solo generation " + f"within the first {PREFIX_TOKENS} tokens inside a left-padded " + "batch; batched left-padded generation is broken again " + f"(issues #1066, #3699).\n" + f"solo : {solo_text!r}\nbatched: {batch_text!r}" + ) + assert not _looks_gibberish(batch_text), ( + f"prompt {i} produced gibberish in a left-padded batch " + f"(issues #1066, #3699): {batch_text!r}" + ) diff --git a/tests/utils/test_prepare_inputs_leftpad.py b/tests/utils/test_prepare_inputs_leftpad.py new file mode 100644 index 0000000000..b773b088b5 --- /dev/null +++ b/tests/utils/test_prepare_inputs_leftpad.py @@ -0,0 +1,409 @@ +"""Regression guard for batched left-padded generation (issues #1066, #3699). + +`_fast_prepare_inputs_for_generation` in unsloth/models/llama.py is shared by +every decoder family wired through `fix_prepare_inputs_for_generation` (llama, +qwen2/3, qwen3_moe, mistral, gemma/2, cohere, granite). Two historical bugs +lived in it: + + (a) the 2D attention mask was truncated to its last column during cached + decode, losing padding information (introduced cc4c5d77, fixed by #2216); + (b) position_ids were taken directly from cache_position, a global counter + that includes left-pad tokens, so padded rows generated garbage + (introduced cc4c5d77, reported in #1066/#3699, fixed by #4100). + +Two layers in this file, both CPU-only and deterministic: + + 1. AST structural checks (TestAstGuard section): parse llama.py with the + stdlib `ast` module only, no unsloth import, so they keep working even + when the package itself fails to import. + 2. Behavioral checks: call the real function with synthetic left-padded + masks and fake caches; unsloth is imported lazily inside each test. + +Both layers fail on the historical bug patterns (validated against the code +states immediately before #2216 and before #4100). Staging proof on GPU-less +hosted runners: danielhanchen/unsloth-staging-2 PRs 170 (green) and 171 (red). + +Companion manual GPU end-to-end check (skipped without CUDA): +tests/utils/test_batched_leftpad_generation_gpu.py +""" + +import ast +from pathlib import Path + +import pytest +import torch + +REPO_ROOT = Path(__file__).resolve().parents[2] +LLAMA_PY = REPO_ROOT / "unsloth" / "models" / "llama.py" + +FUNC_NAME = "_fast_prepare_inputs_for_generation" + + +# -------------------------------------------------------------------------- +# Layer 1: AST structural guard (stdlib only, no unsloth import) +# -------------------------------------------------------------------------- + +# Model files that call fix_prepare_inputs_for_generation(...) and therefore +# share the guarded function. glm4_moe and falcon_h1 are intentionally absent: +# GLM4 MoE does not patch the Llama-compatible generation path (MLA attention) +# and falcon_h1 ships its own _fast_prepare_inputs_for_generation variant. +WIRED_MODEL_FILES = [ + "mistral.py", + "gemma.py", + "gemma2.py", + "qwen2.py", + "qwen3.py", + "qwen3_moe.py", + "cohere.py", + "granite.py", +] + + +def _load_function(): + tree = ast.parse(LLAMA_PY.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == FUNC_NAME: + return node + raise AssertionError( + f"{FUNC_NAME} not found in {LLAMA_PY}; if it was renamed or moved, " + "update this guard so batched left-padded generation stays protected" + ) + + +def _names_in(node): + """All Name ids, attribute names and string constants in a subtree.""" + found = set() + for sub in ast.walk(node): + if isinstance(sub, ast.Name): + found.add(sub.id) + elif isinstance(sub, ast.Attribute): + found.add(sub.attr) + elif isinstance(sub, ast.Constant) and isinstance(sub.value, str): + found.add(sub.value) + return found + + +def _mentions_attention_mask(node): + return any("attention_mask" in name for name in _names_in(node)) + + +def _is_kwargs_position_ids_target(target): + return ( + isinstance(target, ast.Subscript) + and isinstance(target.value, ast.Name) + and target.value.id == "kwargs" + and isinstance(target.slice, ast.Constant) + and target.slice.value == "position_ids" + ) + + +def _walk_with_paths(node, path = ()): + yield node, path + for child in ast.iter_child_nodes(node): + yield from _walk_with_paths(child, path + (node,)) + + +def _find_mask_branch(func): + """The If whose test checks the 2D attention mask (dim() == 2).""" + for node in ast.walk(func): + if not isinstance(node, ast.If): + continue + test_names = _names_in(node.test) + if "dim" in test_names and any("attention_mask" in n for n in test_names): + return node + return None + + +def test_mask_derived_position_ids_branch_exists(): + func = _load_function() + branch = _find_mask_branch(func) + assert branch is not None, ( + f"{FUNC_NAME} no longer has a branch testing the 2D attention mask " + "(dim() == 2); position_ids must be derived per row from the mask for " + "left-padded batches (see PR #4100 / issues #1066, #3699)" + ) + + body_names = set() + for stmt in branch.body: + body_names |= _names_in(stmt) + assert "cumsum" in body_names and _mentions_attention_mask( + ast.Module(body = branch.body, type_ignores = []) + ), ( + "the attention-mask branch must compute position_ids via " + "attention_mask.cumsum(...); reintroducing cache_position-based " + "positions breaks left-padded batched generation (issue #3699)" + ) + assert ( + "masked_fill_" in body_names or "masked_fill" in body_names + ), "the attention-mask branch must mask pad positions (masked_fill on mask == 0)" + assigns_kwargs = any( + isinstance(stmt, ast.Assign) + and any(_is_kwargs_position_ids_target(t) for t in stmt.targets) + for stmt in ast.walk(ast.Module(body = branch.body, type_ignores = [])) + ) + assert ( + assigns_kwargs + ), 'the attention-mask branch must store the derived positions into kwargs["position_ids"]' + + +def test_cache_position_only_used_as_fallback_for_position_ids(): + func = _load_function() + branch = _find_mask_branch(func) + assert branch is not None + + orelse_nodes = set() + for stmt in branch.orelse: + for sub in ast.walk(stmt): + orelse_nodes.add(id(sub)) + + offenders = [] + for node, path in _walk_with_paths(func): + if not isinstance(node, ast.Assign): + continue + if not any(_is_kwargs_position_ids_target(t) for t in node.targets): + continue + value_names = _names_in(node.value) + # Direct use of cache_position, or the local alias `cp` the current + # implementation builds from it inside the fallback branch. + derives_from_cache_position = any("cache_position" in n for n in value_names) or bool( + value_names & {"cp"} + ) + if derives_from_cache_position and id(node) not in orelse_nodes: + offenders.append(ast.unparse(node)) + + assert not offenders, ( + 'kwargs["position_ids"] must never be assigned from cache_position ' + "outside the fallback (orelse) of the 2D attention-mask branch; " + "cache_position counts left-pad tokens, so padded rows generate " + f"garbage (issues #1066, #3699). Offending assignments: {offenders}" + ) + + +def test_attention_mask_never_truncated_to_last_column(): + func = _load_function() + offenders = [] + for node in ast.walk(func): + if not isinstance(node, ast.Assign): + continue + value = node.value + if not isinstance(value, ast.Subscript): + continue + if not _mentions_attention_mask(value.value): + continue + # Match a trailing [-1]-style column selection: [:, [-1]] or [:, -1:] + sl = value.slice + if isinstance(sl, ast.Tuple) and len(sl.elts) == 2: + col = sl.elts[1] + is_last_col_list = ( + isinstance(col, ast.List) + and len(col.elts) == 1 + and isinstance(col.elts[0], ast.UnaryOp) + ) + is_last_col_slice = ( + isinstance(col, ast.Slice) + and col.lower is not None + and isinstance(col.lower, ast.UnaryOp) + and getattr(getattr(col.lower, "operand", None), "value", None) == 1 + and col.upper is None + ) + if is_last_col_list or is_last_col_slice: + offenders.append(ast.unparse(node)) + assert not offenders, ( + "the 2D attention mask must not be truncated to its last column; this " + "was the pre-#2216 bug that drops padding information in cached decode " + f"(issue #1066). Offending assignments: {offenders}" + ) + + +def test_model_families_stay_wired_to_shared_prepare_inputs(): + missing = [] + for fname in WIRED_MODEL_FILES: + path = REPO_ROOT / "unsloth" / "models" / fname + if not path.exists(): + continue + if "fix_prepare_inputs_for_generation(" not in path.read_text(): + missing.append(fname) + assert not missing, ( + "these model files no longer call fix_prepare_inputs_for_generation, " + "so they lose the guarded left-padding-safe prepare_inputs path: " + f"{missing}" + ) + + +# -------------------------------------------------------------------------- +# Layer 2: behavioral guard (calls the real function, lazy unsloth import) +# -------------------------------------------------------------------------- + +PAST_LEN = 4 + +# Three rows with different amounts of left padding (0 = pad). +MASK = torch.tensor( + [ + [0, 0, 1, 1, 1], + [1, 1, 1, 1, 1], + [0, 1, 1, 1, 1], + ], + dtype = torch.long, +) +BS, SEQ = MASK.shape + +# Per-row positions: cumsum(-1) - 1 with pad slots filled with 1. +EXPECTED_PREFILL_POSITIONS = torch.tensor( + [ + [1, 1, 0, 1, 2], + [0, 1, 2, 3, 4], + [1, 0, 1, 2, 3], + ], + dtype = torch.long, +) + + +class FakeDynamicCache: + """Minimal stand-in for transformers DynamicCache with a non-empty cache.""" + + def __init__(self, seq_length): + self._seq_length = seq_length + + def __len__(self): + return 1 + + def get_seq_length(self): + return self._seq_length + + +class FakeModel: + """Bare-minimum `self` for _fast_prepare_inputs_for_generation.""" + + dtype = torch.float32 + config = None + + +class FakeModelWith4DMask(FakeModel): + """Variant exposing the HF 4D mask builder; records what it receives.""" + + def __init__(self): + self.mask_calls = [] + + def _prepare_4d_causal_attention_mask_with_cache_position( + self, + attention_mask, + sequence_length, + target_length, + dtype, + device, + cache_position, + batch_size, + config = None, + past_key_values = None, + ): + self.mask_calls.append( + { + "mask_shape": tuple(attention_mask.shape) if attention_mask is not None else None, + "sequence_length": sequence_length, + "target_length": target_length, + "batch_size": batch_size, + } + ) + return torch.zeros((batch_size, 1, sequence_length, target_length), dtype = dtype) + + +def _prepare(model, input_ids, attention_mask, **kwargs): + from unsloth.models import llama as llama_mod + return llama_mod._fast_prepare_inputs_for_generation( + model, input_ids, attention_mask = attention_mask, **kwargs + ) + + +def test_prefill_position_ids_derived_from_left_padded_mask(): + input_ids = torch.arange(BS * SEQ).reshape(BS, SEQ) + result = _prepare(FakeModel(), input_ids, MASK) + + position_ids = result.get("position_ids", None) + assert ( + position_ids is not None + ), "prefill with a left-padded 2D attention mask must populate position_ids" + assert torch.equal(position_ids.long().cpu(), EXPECTED_PREFILL_POSITIONS), ( + "prefill position_ids must be derived per row from the attention mask " + "(cumsum - 1, pads masked), so each row starts counting at its first " + f"real token; got {position_ids.tolist()}" + ) + assert result["input_ids"].shape == (BS, SEQ) + + +@pytest.mark.parametrize("pass_cache_position", [True, False]) +def test_cached_decode_position_ids_ignore_left_padding(pass_cache_position): + # Decode step: PAST_LEN tokens cached, current token is the mask's last + # column. Row 0 has 2 pads, so its current token sits at logical position 2, + # NOT at cache_position == PAST_LEN. This is exactly issue #3699. + input_ids = torch.arange(BS * SEQ).reshape(BS, SEQ) + kwargs = {"past_key_values": FakeDynamicCache(PAST_LEN)} + if pass_cache_position: + kwargs["cache_position"] = torch.arange(PAST_LEN, PAST_LEN + 1) + + result = _prepare(FakeModel(), input_ids, MASK, **kwargs) + + assert result["input_ids"].shape == ( + BS, + 1, + ), "cached decode must slice input_ids to the last token only" + position_ids = result.get("position_ids", None) + assert position_ids is not None + expected = torch.tensor([[2], [4], [3]], dtype = torch.long) + assert torch.equal(position_ids.long().cpu().reshape(BS, 1), expected), ( + "left-padded cached decode must derive per-row position_ids from the " + "attention mask, not from cache_position which counts pad tokens; got " + f"{position_ids.tolist()}, expected {expected.tolist()} " + "(row 0 has 2 pads: its position must be 2, not 4)" + ) + + +def test_cached_decode_does_not_truncate_2d_attention_mask(): + # Without a 4D mask builder the original 2D mask must survive untouched. + # The historical bug replaced it with attention_mask[:, [-1]]. + input_ids = torch.arange(BS * SEQ).reshape(BS, SEQ) + result = _prepare(FakeModel(), input_ids, MASK, past_key_values = FakeDynamicCache(PAST_LEN)) + mask_out = result["attention_mask"] + assert mask_out is not None + assert mask_out.dim() != 2 or mask_out.shape[-1] == SEQ, ( + "the 2D attention mask must not be truncated to its last column during " + f"cached decode (got shape {tuple(mask_out.shape)}); padding rows lose " + "their pad information otherwise" + ) + + +def test_cached_decode_4d_mask_builder_receives_full_target_length(): + model = FakeModelWith4DMask() + input_ids = torch.arange(BS * SEQ).reshape(BS, SEQ) + result = _prepare(model, input_ids, MASK, past_key_values = FakeDynamicCache(PAST_LEN)) + assert len(model.mask_calls) == 1 + call = model.mask_calls[0] + assert call["mask_shape"] == (BS, SEQ), ( + "the 4D mask builder must receive the full 2D padding mask, not a " + f"truncated one (got {call['mask_shape']})" + ) + assert call["sequence_length"] == 1 + assert call["target_length"] == SEQ, ( + "target_length must cover the whole mask so padded positions stay " + f"masked (got {call['target_length']})" + ) + assert result["attention_mask"].dim() == 4 + + +def test_caller_supplied_position_ids_are_passed_through(): + input_ids = torch.arange(BS * SEQ).reshape(BS, SEQ) + custom = torch.full((BS, SEQ), 7, dtype = torch.long) + result = _prepare(FakeModel(), input_ids, MASK, position_ids = custom) + assert torch.equal( + result["position_ids"], custom + ), "caller-supplied position_ids must not be overwritten" + + +def test_legacy_tuple_cache_still_takes_cached_decode_path(): + # Legacy cache format: tuple of (K, V) per layer; past length from K.shape[-2]. + k = torch.zeros((BS, 1, PAST_LEN, 8)) + legacy_cache = ((k, k.clone()),) + input_ids = torch.arange(BS * SEQ).reshape(BS, SEQ) + result = _prepare(FakeModel(), input_ids, MASK, past_key_values = legacy_cache) + assert result["input_ids"].shape == (BS, 1) + expected = torch.tensor([[2], [4], [3]], dtype = torch.long) + assert torch.equal(result["position_ids"].long().cpu().reshape(BS, 1), expected)