From ad3aef7e253f2611f444b8356bee4e6f9856110f Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Mon, 8 Jun 2026 09:04:27 +0000 Subject: [PATCH 01/99] [feat] qwen2 generative-rec LM: HF-Qwen2 backbone + SID vocab - tzrec/models/generative_rec_lm.py, qwen2_rec_lm.py - tzrec/protos/models/generative_model.proto + model.proto oneof entry - tzrec/optim: lr_scheduler additions; optimizer.proto grad-accum/grad-clip - tzrec/tools/export_genreclm_to_hf.py (DCP -> HF export) Example scripts and design notes intentionally excluded (to be refactored). Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 439 +++++++++++++++++++++ tzrec/models/qwen2_rec_lm.py | 46 +++ tzrec/optim/lr_scheduler.py | 57 +++ tzrec/optim/lr_scheduler_test.py | 36 ++ tzrec/protos/model.proto | 8 + tzrec/protos/models/generative_model.proto | 72 ++++ tzrec/protos/optimizer.proto | 17 + tzrec/tools/export_genreclm_to_hf.py | 107 +++++ 8 files changed, 782 insertions(+) create mode 100644 tzrec/models/generative_rec_lm.py create mode 100644 tzrec/models/qwen2_rec_lm.py create mode 100644 tzrec/protos/models/generative_model.proto create mode 100644 tzrec/tools/export_genreclm_to_hf.py diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py new file mode 100644 index 000000000..83786d848 --- /dev/null +++ b/tzrec/models/generative_rec_lm.py @@ -0,0 +1,439 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 + +"""Generic generative-recommendation language-model base for TorchEasyRec. + +Implements the FINAL design (see FINAL_DESIGN_GENERATIVE_REC_LM.md): + + * Per-family subclasses (design §2 / G4): ``GenerativeRecLM`` is the + abstract base; each LLM family is a concrete subclass declaring a + ``CHAT_TEMPLATE`` class var (e.g. ``Qwen2RecLM`` in + ``tzrec/models/qwen2_rec_lm.py``). The pipeline config selects the + family via ``generative_rec_lm.class_name``; dispatch goes through the + BaseModel registry (subclasses auto-register by class name). + * Streaming sample format: each row carries two raw-int64 sequence features, + ``user_sequence`` (list[int]) and ``label`` (list[int]), both holding raw + SID indices in ``[1, sum(codebook)]``. + * The chat template is tokenised ONCE at ``__init__`` and cached as + ``nn.Module`` non-persistent buffers, so per-batch encoding is purely + integer arithmetic + tensor concatenation (no HF tokenizer in the hot + path). + * SID → token id by integer offset: ``token = sid + base_vocab - 1`` (the + SID atoms ``C0..C{sum-1}`` are added right after the original vocabulary, + no [SEP] in between; matches algr's ``add_tokens`` layout). + * Left padding with ``eos_token_id`` (L7 fix from §11 of the design doc) — + real content sits at the END of every row so the suffix slice captures + only ``[response + end_markers]`` and matches algr's pad-side exactly. + +(The old offline-tokenized ``tzrec/models/qwen2.py`` v1 wrapper and its +``Qwen2 qwen2 = 600`` proto entry have been REMOVED — proto field 600 is +reserved. This streaming pipeline is the only generative-rec path.) +""" + +from __future__ import annotations + +import os +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torchmetrics +from torch import nn +from transformers import AutoModelForCausalLM, AutoTokenizer + +from tzrec.datasets.utils import Batch +from tzrec.features.feature import BaseFeature +from tzrec.models.model import BaseModel +from tzrec.protos.model_pb2 import ModelConfig + + +def _encode_no_special(tokenizer, text: str) -> List[int]: + """Encode a fragment without prepending BOS / appending EOS specials. + + We're building the prompt manually from explicit ``<|im_start|>`` markers, + so we must NOT let the tokenizer's BOS/EOS handling double-emit them. + """ + return tokenizer.encode(text, add_special_tokens=False) + + +class GenerativeRecLM(BaseModel): + """Abstract base for HF-backed generative-recommendation LMs. + + Subclasses declare ``CHAT_TEMPLATE`` (design §5) and rarely override + ``predict()`` (e.g. a future ``MixtralRecLM`` must call the full HF + forward to capture ``aux_loss``). Everything else — model construction, + SID vocab extension, template caching, splice, algr-aligned forward, + loss/metrics — lives here. + + ``CHAT_TEMPLATE`` keys (all strings): + system_prefix / system_suffix — wrap the system instruction + user_prefix / user_suffix — wrap the user message + asst_prefix / asst_suffix — wrap the assistant answer + default_system_instruction — used when the proto doesn't + override ``system_instruction`` + """ + + CHAT_TEMPLATE: Optional[Dict[str, str]] = None + + def __new__(cls, model_config: ModelConfig, *args: Any, **kwargs: Any): + """Dispatch to the concrete family subclass. + + ``tzrec.main._create_model`` resolves the proto oneof message name + (``GenerativeRecLM``) to THIS class; the actual family is selected + by ``generative_rec_lm.class_name`` and looked up in the BaseModel + registry (every subclass auto-registers via the metaclass). + """ + if cls is GenerativeRecLM: + cfg = getattr(model_config, model_config.WhichOneof("model")) + class_name = cfg.class_name + # pyre-ignore [16] + sub_cls = BaseModel.create_class(class_name) + if not issubclass(sub_cls, GenerativeRecLM): + raise ValueError( + f"generative_rec_lm.class_name = {class_name!r} resolves " + f"to {sub_cls}, which is not a GenerativeRecLM subclass." + ) + return super().__new__(sub_cls) + return super().__new__(cls) + + def __init__( + self, + model_config: ModelConfig, + features: List[BaseFeature], + labels: List[str], + sample_weights: Optional[List[str]] = None, + **kwargs: Any, + ) -> None: + super().__init__(model_config, features, labels, sample_weights, **kwargs) + cfg = self._model_config # populated by BaseModel from WhichOneof + + tpl = type(self).CHAT_TEMPLATE + if tpl is None: + raise NotImplementedError( + f"{type(self).__name__} must declare a non-empty CHAT_TEMPLATE " + f"class var (design §2/G4); GenerativeRecLM itself is abstract " + f"— set generative_rec_lm.class_name to a concrete family " + f"(e.g. 'Qwen2RecLM')." + ) + + # --------- proto -> python knobs ------------------------------ + self._input_name: str = cfg.user_sequence_feature_name + self._label_name: str = cfg.label_feature_name + self._ignore_index: int = int(cfg.ignore_index) + codebook = list(cfg.codebook) + if len(codebook) == 0: + raise ValueError( + "GenerativeRecLM: codebook must be non-empty " + "(see design §3 — required field)" + ) + sid_atoms = sum(int(c) for c in codebook) + pad_mult = int(cfg.vocab_pad_to_multiple_of) or 128 + + # --------- backbone + tokenizer ------------------------------- + hf_model_id = cfg.hf_model_id + if not hf_model_id: + raise ValueError( + "GenerativeRecLM v1: hf_model_id is required " + "(architecture-spec path deferred to v1.x)" + ) + # torch_dtype="auto" preserves the safetensors-stored dtype (bf16 + # for Qwen2.5-0.5B). The default would silently upcast to fp32. + # On CPU the same flag is honoured; on GPU it avoids a 2× memory + # blow-up. + self.lm = AutoModelForCausalLM.from_pretrained( + hf_model_id, torch_dtype="auto" + ) + # ``use_fast=True`` is the modern default; explicit for clarity. + tokenizer = AutoTokenizer.from_pretrained(hf_model_id, use_fast=True) + + # --------- vocab extension (codebook required) ---------------- + # The SID-atom base is the tokenizer's next free id BEFORE adding + # ``C0..``. For Qwen2.5-0.5B that's 151665 = `model.config.vocab_size` + # (151936, includes ~300 reserved padding slots) minus the unused + # reserved span — so use ``len(tokenizer)`` directly, NOT + # ``model.config.vocab_size``. + base = len(tokenizer) + new_atoms = [f"C{i}" for i in range(sid_atoms)] + added = tokenizer.add_tokens(new_atoms) + if added != sid_atoms: + # The tokenizer already had some Cxxx tokens — we expect a fresh + # base, so this would silently break our offset arithmetic. + raise RuntimeError( + f"GenerativeRecLM: tokenizer was expected to grow by " + f"{sid_atoms} new atoms, only added {added}. " + f"Aborting to avoid silent SID-token mismatch." + ) + # Final vocab = base + sid_atoms, padded up to multiple of pad_mult. + # Matches algr's layout: SID atoms appended directly to the existing + # tokenizer vocab; offset arithmetic is `token = base + (sid - 1)`. + self.lm.resize_token_embeddings( + base + sid_atoms, pad_to_multiple_of=pad_mult + ) + + # L3 safety check: assert C0 lands at the recorded base. + c0_id = tokenizer.convert_tokens_to_ids("C0") + if c0_id != base: + raise RuntimeError( + f"GenerativeRecLM: SID atom layout mismatch — expected " + f"C0 at token id {base}, got {c0_id}. " + f"Splice arithmetic would produce wrong token ids." + ) + self._base_vocab = base # used in `_splice_input_ids` + + # L1 + L7 mitigation: pad with eos_token_id on the LEFT side. + pad_id = tokenizer.pad_token_id + if pad_id is None: + pad_id = tokenizer.eos_token_id + self._pad_token_id = int(pad_id) + + # --------- cache chat-template buffers ------------------------ + self._build_prompt_tokens(tokenizer, cfg) + + # Diagnostics for the first launch — useful when chasing splice bugs. + self._smoke_log_once = (os.environ.get("TZREC_GENRECLM_DEBUG", "0") == "1") + self._first_predict = True + + # ------------------------------------------------------------------ template + def _build_prompt_tokens(self, tokenizer, cfg) -> None: + """Tokenise the family chat template once; cache as buffers. + + Composes the proto's optional ``system_instruction`` / + ``user_prefix_text`` / ``user_suffix_text`` (algr's CN prompt + wrappers — L2/L4 mitigations) with the family's static fragments: + + tpl_system = system_prefix + system_instruction + system_suffix + tpl_user_prefix = user_prefix + user_prefix_text + tpl_user_suffix = user_suffix_text + user_suffix + tpl_asst_prefix / tpl_asst_suffix verbatim from the template + + Buffers are non-persistent — they live with the module (move with + ``model.to(...)``) but stay off the state_dict so HF safetensors + round-tripping isn't polluted by TER-only state. + """ + tpl = type(self).CHAT_TEMPLATE + sys_text = cfg.system_instruction or tpl["default_system_instruction"] + u_pre = cfg.user_prefix_text or "" + u_suf = cfg.user_suffix_text or "" + frags = { + "system": tpl["system_prefix"] + sys_text + tpl["system_suffix"], + "user_prefix": tpl["user_prefix"] + u_pre, + "user_suffix": u_suf + tpl["user_suffix"], + "asst_prefix": tpl["asst_prefix"], + "asst_suffix": tpl["asst_suffix"], + } + for slot_name, frag_str in frags.items(): + ids = torch.tensor( + _encode_no_special(tokenizer, frag_str), dtype=torch.long + ) + self.register_buffer(f"tpl_{slot_name}", ids, persistent=False) + # algr appends eos to BOTH input_ids and labels at train time + # (algr/models/qwen2_5/data.py:46-47) — i.e. the trailing eos is a + # SUPERVISED token. Cache it so the splice can mirror that exactly. + self.register_buffer( + "tpl_eos", + torch.tensor([int(tokenizer.eos_token_id)], dtype=torch.long), + persistent=False, + ) + + # ------------------------------------------------------------------ init_input + def init_input(self) -> None: + """No-op override. + + The HF backbone owns its own ``embed_tokens``; we don't use TER's + ``EmbeddingGroup`` at all. Token IDs flow through directly. + """ + self.embedding_group = None + + # ------------------------------------------------------------------ jagged -> rows + @staticmethod + def _jagged_to_row_list(jt) -> List[torch.Tensor]: + """Convert a TER JaggedTensor (values, lengths) to a list of 1-D + int64 row tensors. + + ``values`` may arrive as float (TER's ``sequence_raw_feature`` reads + ``list`` as float — see [[project-tzrec-qwen2-integration]] + gotcha §B). We cast to long here; SID values fit in float32 mantissa + for any realistic codebook size (< 2^24). + """ + values = jt.values() if callable(getattr(jt, "values", None)) else jt.values + lengths = jt.lengths() if callable(getattr(jt, "lengths", None)) else jt.lengths + if values.dim() == 2 and values.size(-1) == 1: + values = values.squeeze(-1) + values = values.long() + lengths = lengths.long() + out: List[torch.Tensor] = [] + start = 0 + for n in lengths.tolist(): + out.append(values[start : start + n]) + start += n + return out + + # ------------------------------------------------------------------ splice + def _splice_input_ids( + self, + user_seq_rows: List[torch.Tensor], + label_rows: List[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build ``(input_ids, labels, attention_mask)``, each ``(B, T_max)``. + + Left-padded with ``eos_token_id``. ``attention_mask`` is essential — + without it self-attention would let pad positions pollute real + positions' hidden states. CE is separately protected by ``-100`` + labels at pad slots, but the forward needs the mask too. + + SID → token: ``token = sid + base_vocab - 1`` (SID atoms ``C0..`` + start at position ``base_vocab``; SID indices are 1-indexed). + """ + assert len(user_seq_rows) == len(label_rows) + B = len(user_seq_rows) + dev = self.tpl_system.device + base = self._base_vocab + + rows_ids: List[torch.Tensor] = [] + rows_lab: List[torch.Tensor] = [] + for i in range(B): + # SID → token id with int math. Map ``sid`` to its corresponding + # ``C{sid-1}`` atom: ``token = base + (sid - 1)`` = ``sid + (base-1)``. + # Cast to long matches the buffer dtype. + u_tok = (user_seq_rows[i].to(dev) + (base - 1)) + a_tok = (label_rows[i].to(dev) + (base - 1)) + ids = torch.cat([ + self.tpl_system, self.tpl_user_prefix, u_tok, + self.tpl_user_suffix, self.tpl_asst_prefix, a_tok, + self.tpl_asst_suffix, self.tpl_eos, + ]) + ign = torch.full_like(ids, self._ignore_index) + start = ( + self.tpl_system.numel() + + self.tpl_user_prefix.numel() + + u_tok.numel() + + self.tpl_user_suffix.numel() + + self.tpl_asst_prefix.numel() + ) + ign[start : start + a_tok.numel()] = a_tok + # algr supervises the trailing eos (after the masked + # ``<|im_end|>\n`` markers) — data.py:46-47. Mirror it. + ign[-1] = self.tpl_eos[0] + rows_ids.append(ids) + rows_lab.append(ign) + + T = max(r.numel() for r in rows_ids) + input_ids = torch.full( + (B, T), self._pad_token_id, dtype=torch.long, device=dev + ) + labels = torch.full( + (B, T), self._ignore_index, dtype=torch.long, device=dev + ) + attention_mask = torch.zeros((B, T), dtype=torch.long, device=dev) + for i, (ids, ign) in enumerate(zip(rows_ids, rows_lab)): + n = ids.numel() + # LEFT padding: write rows to the END of each (T,) slot. + input_ids[i, -n:] = ids + labels[i, -n:] = ign + attention_mask[i, -n:] = 1 + return input_ids, labels, attention_mask + + @staticmethod + def _min_first_non_neg_index(labels: torch.Tensor) -> int: + """Verbatim port of algr's helper (al_sid/algr/models/qwen2_5/ + modeling_qwen.py:1267-1274). + + Returns the smallest position (across rows in the batch) where the + first non-(-100) label appears. Used to decide how many trailing + positions to feed into ``lm_head``. + """ + tmp = (labels >= 0).cumsum(dim=-1) + return int((tmp == 1).float().argmax(dim=-1).min().item()) + + # ------------------------------------------------------------------ predict + def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: + jt_u = batch.sequence_dense_features[self._input_name] + jt_l = batch.sequence_dense_features[self._label_name] + u_rows = self._jagged_to_row_list(jt_u) + l_rows = self._jagged_to_row_list(jt_l) + + input_ids, labels, attention_mask = self._splice_input_ids(u_rows, l_rows) + + if self._smoke_log_once and self._first_predict: + print( + f"[GENRECLM_DEBUG] first batch: B={input_ids.shape[0]} " + f"T={input_ids.shape[1]} pad_id={self._pad_token_id} " + f"ign={self._ignore_index} dev={input_ids.device} " + f"input_ids[0, -8:]={input_ids[0, -8:].tolist()} " + f"labels[0, -8:]={labels[0, -8:].tolist()}", + flush=True, + ) + self._first_predict = False + + outputs = self.lm.model( + input_ids=input_ids, attention_mask=attention_mask + ) + hidden = outputs.last_hidden_state # (B, T, D) + + # Suffix slice in BOTH train and eval. algr only slices when + # training (its eval goes through a separate beam-search predict), + # but for CE the slice is value-identical (positions outside the + # suffix all carry -100 labels) and it bounds the logits tensor to + # (B, T_suffix, V) — without it, eval at bsz=80 would materialise + # (B, T, 217k) logits plus HF loss_function's fp32 upcast and OOM. + if (labels >= 0).any(): + keep = labels.shape[1] - self._min_first_non_neg_index(labels) + 1 + sl = slice(-keep, None) + labels_sl = labels[:, sl] + else: + sl = slice(None) + labels_sl = labels + + logits = self.lm.lm_head(hidden[:, sl, :]) + + # ``loss_function`` is the HF ``ForCausalLMLoss`` callable hung off + # every ``…ForCausalLM`` class; does shift-by-one + CE with -100 + # ignore. Calling it here matches algr's training-step loss exactly. + loss = self.lm.loss_function( + logits=logits, + labels=labels_sl, + vocab_size=self.lm.config.vocab_size, + ) + return {"loss": loss, "logits": logits} + + # ------------------------------------------------------------------ loss + def init_loss(self) -> None: + return + + def loss( + self, + predictions: Dict[str, torch.Tensor], + batch: Batch, + ) -> Dict[str, torch.Tensor]: + return {"ce_loss": predictions["loss"]} + + # ------------------------------------------------------------------ metrics + # See [[project-tzrec-qwen2-integration]] gotcha §C: BaseModel only + # declares the eval-side metric methods; the train loop calls both + # families, so we must override both. + def init_metric(self) -> None: + # Mean CE over the eval set — gives `_evaluate` something to log + # (BaseModel.compute_metric iterates `_metric_modules` generically; + # torchmetrics handles the cross-rank sync at compute()). + self._metric_modules["ce_loss"] = torchmetrics.MeanMetric() + + def update_metric( + self, + predictions: Dict[str, torch.Tensor], + batch: Batch, + losses: Optional[Dict[str, torch.Tensor]] = None, + ) -> None: + self._metric_modules["ce_loss"].update(predictions["loss"].detach()) + + def init_train_metric(self) -> None: + return + + def update_train_metric( + self, + predictions: Dict[str, torch.Tensor], + batch: Batch, + losses: Optional[Dict[str, torch.Tensor]] = None, + ) -> None: + return diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py new file mode 100644 index 000000000..839637caa --- /dev/null +++ b/tzrec/models/qwen2_rec_lm.py @@ -0,0 +1,46 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 + +"""Qwen2/Qwen2.5 family subclass of ``GenerativeRecLM`` (design §5). + +Selected from the pipeline config via:: + + model_config { + generative_rec_lm { + class_name: "Qwen2RecLM" + ... + } + } + +The family contributes ONLY its chat-template fragments; the splice, +algr-aligned forward, vocab extension and checkpoint plumbing all live in +the ``GenerativeRecLM`` base. A new LLM family is one file like this one +(see ``FINAL_DESIGN_GENERATIVE_REC_LM.md`` §2/§5 — ``Qwen3RecLM`` should +subclass ``Qwen2RecLM`` and override only what differs). +""" + +from tzrec.models.generative_rec_lm import GenerativeRecLM + +# Verbatim Qwen2 ChatML fragments. ``default_system_instruction`` matches +# algr/models/qwen2_5/data.py:73 ("default_instruction") bit-for-bit — the +# L2 mitigation in design §11. +QWEN2_TEMPLATE = { + "system_prefix": "<|im_start|>system\n", + "system_suffix": "<|im_end|>\n", + "user_prefix": "<|im_start|>user\n", + "user_suffix": "<|im_end|>\n", + "asst_prefix": "<|im_start|>assistant\n", + "asst_suffix": "<|im_end|>\n", + "default_system_instruction": ( + "You are Qwen, created by Alibaba Cloud. You are a helpful assistant." + ), +} + + +class Qwen2RecLM(GenerativeRecLM): + """Qwen2 / Qwen2.5 generative-recommendation LM.""" + + CHAT_TEMPLATE = QWEN2_TEMPLATE diff --git a/tzrec/optim/lr_scheduler.py b/tzrec/optim/lr_scheduler.py index 38680f60c..fddbbea6b 100644 --- a/tzrec/optim/lr_scheduler.py +++ b/tzrec/optim/lr_scheduler.py @@ -159,6 +159,63 @@ def _get_lr(self) -> List[float]: return lr +class LinearDecayLR(BaseLR): + """Linear Decay LearningRate Scheduler. + + Decays the learning rate linearly from base_lr to min_learning_rate + over total_size steps or epochs, with optional linear warmup. Mirrors + HuggingFace Trainer's ``lr_scheduler_type: linear``. + + Args: + optimizer (Optimizer): an instance of Optimizer. + total_size (int): total number of steps or epochs for the decay. + min_learning_rate (float): minimum learning rate. + warmup_learning_rate (float): warmup start learning rate. + warmup_size (int): warmup steps or epochs. + by_epoch (bool): schedule by epoch or by step. + """ + + def __init__( + self, + optimizer: Optimizer, + total_size: int, + min_learning_rate: float = 0.0, + warmup_learning_rate: float = 0.0, + warmup_size: int = 0, + by_epoch: bool = False, + ) -> None: + if total_size <= 0: + raise ValueError(f"total_size must be positive, got {total_size}") + if warmup_size >= total_size: + raise ValueError( + f"warmup_size ({warmup_size}) must be smaller than " + f"total_size ({total_size})" + ) + self._total_size = total_size + self._min_learning_rate = min_learning_rate + self._warmup_learning_rate = warmup_learning_rate + self._warmup_size = warmup_size + super().__init__(optimizer, by_epoch=by_epoch) + + def _get_lr(self) -> List[float]: + """Calculates the learning rate.""" + step_count = max(self._step_count - 1, 0) + if step_count < self._warmup_size: + scale = step_count / self._warmup_size + return [ + (base_lr - self._warmup_learning_rate) * scale + + self._warmup_learning_rate + for base_lr in self.base_lrs + ] + t = min(step_count - self._warmup_size, self._total_size - self._warmup_size) + decay_scale = 1.0 - t / (self._total_size - self._warmup_size) + return [ + self._min_learning_rate + + (base_lr - self._min_learning_rate) * decay_scale + for base_lr in self.base_lrs + ] + + class CosineAnnealingLR(BaseLR): """Cosine Annealing LearningRate Scheduler. diff --git a/tzrec/optim/lr_scheduler_test.py b/tzrec/optim/lr_scheduler_test.py index 9c50cf653..e3551c2d1 100644 --- a/tzrec/optim/lr_scheduler_test.py +++ b/tzrec/optim/lr_scheduler_test.py @@ -83,6 +83,42 @@ def test_manual_step_lr_with_warmup(self) -> None: lr.step() self.assertAlmostEqual(opt.param_groups[0]["lr"], lr_gt) + def test_linear_decay_lr(self) -> None: + params = [torch.tensor([1.0, 2.0])] + opt = torch.optim.Adam(params, lr=0.01) + lr = lr_scheduler.LinearDecayLR(opt, total_size=4) + lr_gts = [0.0075, 0.005, 0.0025, 0.0, 0.0] + for lr_gt in lr_gts: + lr.step() + self.assertAlmostEqual(opt.param_groups[0]["lr"], lr_gt) + + def test_linear_decay_lr_with_min_lr(self) -> None: + params = [torch.tensor([1.0, 2.0])] + opt = torch.optim.Adam(params, lr=0.01) + lr = lr_scheduler.LinearDecayLR(opt, total_size=4, min_learning_rate=0.002) + lr_gts = [0.008, 0.006, 0.004, 0.002, 0.002] + for lr_gt in lr_gts: + lr.step() + self.assertAlmostEqual(opt.param_groups[0]["lr"], lr_gt) + + def test_linear_decay_lr_with_warmup(self) -> None: + params = [torch.tensor([1.0, 2.0])] + opt = torch.optim.Adam(params, lr=0.01) + lr = lr_scheduler.LinearDecayLR( + opt, total_size=6, warmup_size=2, warmup_learning_rate=0.002 + ) + self.assertFalse(lr.by_epoch) + # warmup step 0->1: scale=0.5, lr=0.002+(0.01-0.002)*0.5=0.006 + lr.step() + self.assertAlmostEqual(opt.param_groups[0]["lr"], 0.006) + # warmup step 1->2: scale=1.0, lr=0.01 + lr.step() + self.assertAlmostEqual(opt.param_groups[0]["lr"], 0.01) + # decay over remaining 4 steps: 0.0075, 0.005, 0.0025, 0.0 + for lr_gt in [0.0075, 0.005, 0.0025, 0.0, 0.0]: + lr.step() + self.assertAlmostEqual(opt.param_groups[0]["lr"], lr_gt) + def test_cosine_annealing_lr(self) -> None: params = [torch.tensor([1.0, 2.0])] opt = torch.optim.Adam(params, lr=0.01) diff --git a/tzrec/protos/model.proto b/tzrec/protos/model.proto index bef2062ea..05419f0e3 100644 --- a/tzrec/protos/model.proto +++ b/tzrec/protos/model.proto @@ -5,6 +5,7 @@ import "tzrec/protos/models/rank_model.proto"; import "tzrec/protos/models/multi_task_rank.proto"; import "tzrec/protos/models/match_model.proto"; import "tzrec/protos/models/general_rank_model.proto"; +import "tzrec/protos/models/generative_model.proto"; import "tzrec/protos/loss.proto"; import "tzrec/protos/metric.proto"; import "tzrec/protos/seq_encoder.proto"; @@ -76,8 +77,15 @@ message ModelConfig { TDM tdm = 400; RocketLaunching rocket_launching = 500; + + // Generative (causal-LM) models. + GenerativeRecLM generative_rec_lm = 601; } + // Field 600 was the removed v1 offline-tokenized `Qwen2` wrapper — + // do not reuse the number. + reserved 600; + optional uint32 num_class = 2 [default = 1]; repeated LossConfig losses = 3; diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto new file mode 100644 index 000000000..f57246ff9 --- /dev/null +++ b/tzrec/protos/models/generative_model.proto @@ -0,0 +1,72 @@ +syntax = "proto2"; +package tzrec.protos; + +// Generative (causal-LM) models. +// +// These differ from the rank/match/multi-task categories in that the model's +// output is a per-position vocabulary distribution rather than a single score +// or embedding. + +// GenerativeRecLM — the unified config for HF-backed generative recommendation +// LMs. Dispatch to a concrete subclass via `class_name`; see design doc +// §3 (FINAL_DESIGN_GENERATIVE_REC_LM.md). +// +// Sample contract (consumed by `predict()`): +// * user_sequence : list — raw SID indices in [1, sum(codebook)] +// * label : list — raw SID indices in [1, sum(codebook)] +// Both flow into TER as `sequence_raw_feature` parquet columns. The model +// converts SID → token id at batch time via integer offset arithmetic and +// splices them between the cached chat-template buffers built at __init__. +message GenerativeRecLM { + // Dispatches to the concrete Python subclass. Valid values: + // "Qwen2RecLM" — Qwen2 family (Qwen2.5-0.5B, etc.) + // Future: "Llama3RecLM", "MistralRecLM", "MixtralRecLM" (MoE override). + required string class_name = 1; + + // ----- Architecture source ----- + // Either supply an HF hub / local model id (loaded via + // `AutoModelForCausalLM.from_pretrained(..., torch_dtype="auto")`), or + // give an architecture spec to randomly initialise. Exactly one must be + // set; checked at runtime. + optional string hf_model_id = 2 [default = ""]; + // (architecture spec block is deferred to v1.x — `hf_model_id` is the + // only path validated in v1; matches algr's setup.) + + // ----- SID vocabulary ----- + // Required, non-empty. Each entry is the codebook size at one RQ layer + // (e.g. [512, 512] for a 2-layer RQ with 512 codes/layer). Total SID + // atoms = sum(codebook); these are added to the tokenizer / model + // vocabulary as `C0 .. C{sum-1}` directly after the base vocabulary + // (no [SEP]; matches algr's add_tokens layout). SID indices in the + // parquet are 1-indexed; the wrapper maps `sid → base_vocab + (sid-1)` + // where base_vocab = len(tokenizer) BEFORE the extension. + repeated uint32 codebook = 60; + + // Pad the post-extension vocabulary up to a multiple of this value. + // 128 matches algr's `model.resize_token_embeddings(..., pad_to_multiple_of=128)`. + optional uint32 vocab_pad_to_multiple_of = 61 [default = 128]; + + // ----- Chat template ----- + // Optional override for the system instruction string. When unset, the + // subclass's default (e.g. Qwen2RecLM's algr-matching default) is used. + optional string system_instruction = 70 [default = ""]; + + // Optional CN/EN text fragments wrapping the SID codes inside the user + // message. For algr's tiny dataset these are + // user_prefix_text = "当前用户的历史行为如下：" + // user_suffix_text = "，请预测用户在电商推荐场景后续行为的语义编码" + // Tokenised once at __init__ and concatenated INSIDE the cached + // ``tpl_user_prefix`` and ``tpl_user_suffix`` buffers (after the + // ``<|im_start|>user\n`` marker, before the ``<|im_end|>\n`` marker). + // Leaving them empty produces the algr DEFAULT_INSTRUCTION path. + optional string user_prefix_text = 71 [default = ""]; + optional string user_suffix_text = 72 [default = ""]; + + // ----- Sample feature names ----- + // Which `sequence_raw_feature` blocks in the parquet carry the SID lists. + required string user_sequence_feature_name = 51; + required string label_feature_name = 52; + + // Cross-entropy ignore index — matches PyTorch's F.cross_entropy default. + optional int32 ignore_index = 53 [default = -100]; +} diff --git a/tzrec/protos/optimizer.proto b/tzrec/protos/optimizer.proto index e020923b3..63fd6708f 100644 --- a/tzrec/protos/optimizer.proto +++ b/tzrec/protos/optimizer.proto @@ -20,6 +20,7 @@ message SparseOptimizer { ManualStepLR manual_step_learning_rate = 103; CosineAnnealingLR cosine_annealing_learning_rate = 104; CosineAnnealingWarmRestartsLR cosine_annealing_warm_restarts_learning_rate = 105; + LinearDecayLR linear_decay_learning_rate = 106; } } @@ -38,6 +39,7 @@ message DenseOptimizer { ManualStepLR manual_step_learning_rate = 103; CosineAnnealingLR cosine_annealing_learning_rate = 104; CosineAnnealingWarmRestartsLR cosine_annealing_warm_restarts_learning_rate = 105; + LinearDecayLR linear_decay_learning_rate = 106; } repeated PartOptimizer part_optimizers = 201; } @@ -58,6 +60,7 @@ message PartOptimizer { ManualStepLR manual_step_learning_rate = 103; CosineAnnealingLR cosine_annealing_learning_rate = 104; CosineAnnealingWarmRestartsLR cosine_annealing_warm_restarts_learning_rate = 105; + LinearDecayLR linear_decay_learning_rate = 106; } } @@ -233,6 +236,20 @@ message ManualStepLR { optional bool by_epoch = 4 [default = false]; } +message LinearDecayLR { + // total number of steps or epochs to decay from base_lr to + // min_learning_rate (mirrors HF Trainer's `lr_scheduler_type: linear`) + optional uint32 total_size = 1; + // minimum learning rate reached at total_size + optional float min_learning_rate = 2 [default = 0.0]; + // warmup start learning rate + optional float warmup_learning_rate = 3 [default = 0.0]; + // warmup steps or epochs + optional uint32 warmup_size = 4 [default = 0]; + // schedule by epoch or by step. + optional bool by_epoch = 5 [default = false]; +} + message CosineAnnealingLR { // total number of steps or epochs for cosine annealing optional uint32 T_max = 1; diff --git a/tzrec/tools/export_genreclm_to_hf.py b/tzrec/tools/export_genreclm_to_hf.py new file mode 100644 index 000000000..e73619492 --- /dev/null +++ b/tzrec/tools/export_genreclm_to_hf.py @@ -0,0 +1,107 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 + +"""Export a `GenerativeRecLM` TER DCP checkpoint to a HF-loadable directory. + +Design §6.4 (FINAL_DESIGN_GENERATIVE_REC_LM.md), simplified: instead of +hand-writing safetensors shards, we rebuild the model from the pipeline +config (which re-applies the SID vocab extension so shapes match), overlay +the DCP shards onto it, and let HF's ``save_pretrained`` deal with weight +tying, sharding and config serialisation. The extended tokenizer (C0.. at +``len(tokenizer)`` — TER layout, no [SEP]) is saved alongside so generation +consumers decode SID atoms with the SAME ids the model was trained on. + +DCP shard FQNs are ``model.lm.`` (TrainWrapper prefix ``model.`` + +wrapper attr ``lm.``); we restore through a TrainWrapper-shaped state dict +so no manual FQN surgery is needed. + +Usage (CPU-only; safe to run next to a live training):: + + PYTHONPATH=. python -m tzrec.tools.export_genreclm_to_hf \\ + --pipeline_config_path experiments//pipeline.config \\ + --checkpoint_path experiments//model.ckpt-40000 \\ + --export_dir experiments//export_hf_40000 +""" + +from __future__ import annotations + +import argparse +import os + +import torch +from google.protobuf import text_format +from torch.distributed.checkpoint import FileSystemReader, load + +from tzrec.models.generative_rec_lm import GenerativeRecLM # noqa: F401 +from tzrec.models.model import BaseModel +from tzrec.protos.pipeline_pb2 import EasyRecConfig +from transformers import AutoTokenizer + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--pipeline_config_path", required=True) + ap.add_argument("--checkpoint_path", required=True) + ap.add_argument("--export_dir", required=True) + args = ap.parse_args() + + pipeline_config = EasyRecConfig() + with open(args.pipeline_config_path) as f: + text_format.Merge(f.read(), pipeline_config) + model_config = pipeline_config.model_config + grl_cfg = getattr(model_config, model_config.WhichOneof("model")) + + # Rebuild the model exactly as training did (from_pretrained backbone + + # SID vocab extension), CPU-resident. + print(f"[export] building {grl_cfg.class_name} from {grl_cfg.hf_model_id}") + # pyre-ignore [16] + model_cls = BaseModel.create_class("GenerativeRecLM") + model = model_cls(model_config, features=[], labels=[]) + model.eval() + + # Overlay DCP shards. Shard keys are "model.lm." — present the + # state dict under the same prefix. + ckpt_model_dir = os.path.join(args.checkpoint_path, "model") + print(f"[export] overlaying DCP shards from {ckpt_model_dir}") + lm_sd = model.lm.state_dict() + prefixed = {f"model.lm.{k}": v for k, v in lm_sd.items()} + load(prefixed, storage_reader=FileSystemReader(ckpt_model_dir)) + model.lm.load_state_dict({k[len("model.lm."):]: v for k, v in prefixed.items()}) + + # Sanity: SID rows must differ from fresh init → confirm overlay landed. + with torch.no_grad(): + emb = model.lm.get_input_embeddings().weight + print( + f"[export] embed_tokens: shape={tuple(emb.shape)} " + f"dtype={emb.dtype} mean_abs={emb.abs().mean().item():.6f}" + ) + + os.makedirs(args.export_dir, exist_ok=True) + print(f"[export] save_pretrained -> {args.export_dir}") + model.lm.save_pretrained(args.export_dir) + + # Save the EXTENDED tokenizer (TER layout: C0 at len(base tokenizer), + # no [SEP]) so downstream generation maps SID atoms identically. + tokenizer = AutoTokenizer.from_pretrained(grl_cfg.hf_model_id, use_fast=True) + base = len(tokenizer) + tokenizer.add_tokens([f"C{i}" for i in range(sum(grl_cfg.codebook))]) + assert tokenizer.convert_tokens_to_ids("C0") == base + tokenizer.save_pretrained(args.export_dir) + with open(os.path.join(args.export_dir, "TER_EXPORT_INFO.txt"), "w") as f: + f.write( + f"source_checkpoint={args.checkpoint_path}\n" + f"pipeline_config={args.pipeline_config_path}\n" + f"sid_base_token_id={base}\n" + f"codebook={list(grl_cfg.codebook)}\n" + "note=C atoms appended directly after base vocab (NO [SEP]); " + "token_id = base + (sid - 1) for 1-indexed SIDs / base + k for C{k}.\n" + ) + print(f"[export] done; sid_base_token_id={base}") + return 0 + + +if __name__ == "__main__": + main() From a8eeb1247af7d302d0c13dfb3833315686b98780 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Tue, 9 Jun 2026 03:13:19 +0000 Subject: [PATCH 02/99] [refactor] generative-rec LM: base/subclass split + vectorized splice MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - GenerativeRecLM (base): architecture-agnostic plumbing — vocab extension, _tokenize_sids (SID->token-id offset map), _sid_token_rows (jagged read + tokenize-once + split, with data-boundary answer-width validation), device property, loss/metrics; _build_prompt_tokens / predict are abstract hooks. - Qwen2RecLM (subclass): ChatML template, causal-LM splice, decoder-only forward. _splice_input_ids builds input_ids/mask via pad_sequence and labels in one vectorized write (fixed answer width = len(codebook) levels). Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 341 +++++++++--------------------- tzrec/models/qwen2_rec_lm.py | 200 +++++++++++++++++- 2 files changed, 288 insertions(+), 253 deletions(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 83786d848..a6acc3736 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -9,8 +9,8 @@ Implements the FINAL design (see FINAL_DESIGN_GENERATIVE_REC_LM.md): * Per-family subclasses (design §2 / G4): ``GenerativeRecLM`` is the - abstract base; each LLM family is a concrete subclass declaring a - ``CHAT_TEMPLATE`` class var (e.g. ``Qwen2RecLM`` in + abstract base; each LLM family is a concrete subclass implementing the + ``_build_prompt_tokens`` and ``predict`` hooks (e.g. ``Qwen2RecLM`` in ``tzrec/models/qwen2_rec_lm.py``). The pipeline config selects the family via ``generative_rec_lm.class_name``; dispatch goes through the BaseModel registry (subclasses auto-register by class name). @@ -27,20 +27,15 @@ * Left padding with ``eos_token_id`` (L7 fix from §11 of the design doc) — real content sits at the END of every row so the suffix slice captures only ``[response + end_markers]`` and matches algr's pad-side exactly. - -(The old offline-tokenized ``tzrec/models/qwen2.py`` v1 wrapper and its -``Qwen2 qwen2 = 600`` proto entry have been REMOVED — proto field 600 is -reserved. This streaming pipeline is the only generative-rec path.) """ from __future__ import annotations import os -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional import torch import torchmetrics -from torch import nn from transformers import AutoModelForCausalLM, AutoTokenizer from tzrec.datasets.utils import Batch @@ -49,33 +44,24 @@ from tzrec.protos.model_pb2 import ModelConfig -def _encode_no_special(tokenizer, text: str) -> List[int]: - """Encode a fragment without prepending BOS / appending EOS specials. - - We're building the prompt manually from explicit ``<|im_start|>`` markers, - so we must NOT let the tokenizer's BOS/EOS handling double-emit them. - """ - return tokenizer.encode(text, add_special_tokens=False) - - class GenerativeRecLM(BaseModel): """Abstract base for HF-backed generative-recommendation LMs. - Subclasses declare ``CHAT_TEMPLATE`` (design §5) and rarely override - ``predict()`` (e.g. a future ``MixtralRecLM`` must call the full HF - forward to capture ``aux_loss``). Everything else — model construction, - SID vocab extension, template caching, splice, algr-aligned forward, - loss/metrics — lives here. - - ``CHAT_TEMPLATE`` keys (all strings): - system_prefix / system_suffix — wrap the system instruction - user_prefix / user_suffix — wrap the user message - asst_prefix / asst_suffix — wrap the assistant answer - default_system_instruction — used when the proto doesn't - override ``system_instruction`` - """ + The base owns the architecture-agnostic plumbing: model construction, + SID vocab extension, the shared sample data-prep (``_sid_token_rows`` / + ``_tokenize_sids`` — the streaming SID sample contract is the same for all + families, design §1), loss, and metrics. The two architecture-specific + pieces are abstract hooks that each family subclass implements (§15/§16): - CHAT_TEMPLATE: Optional[Dict[str, str]] = None + _build_prompt_tokens(tokenizer, cfg) — cache the prompt template + predict(batch) — build inputs + HF forward + + ``Qwen2RecLM`` (``tzrec/models/qwen2_rec_lm.py``) provides the decoder-only + chat implementation (ChatML splice + ``.model``/``.lm_head`` forward), + reusable by Llama/Mistral/Gemma/Phi-style families; GPT-NeoX/RWKV/Mamba/T5 + each need their own. The pipeline config selects the family via + ``generative_rec_lm.class_name`` (resolved through the BaseModel registry). + """ def __new__(cls, model_config: ModelConfig, *args: Any, **kwargs: Any): """Dispatch to the concrete family subclass. @@ -109,16 +95,7 @@ def __init__( super().__init__(model_config, features, labels, sample_weights, **kwargs) cfg = self._model_config # populated by BaseModel from WhichOneof - tpl = type(self).CHAT_TEMPLATE - if tpl is None: - raise NotImplementedError( - f"{type(self).__name__} must declare a non-empty CHAT_TEMPLATE " - f"class var (design §2/G4); GenerativeRecLM itself is abstract " - f"— set generative_rec_lm.class_name to a concrete family " - f"(e.g. 'Qwen2RecLM')." - ) - - # --------- proto -> python knobs ------------------------------ + # proto -> python knobs self._input_name: str = cfg.user_sequence_feature_name self._label_name: str = cfg.label_feature_name self._ignore_index: int = int(cfg.ignore_index) @@ -128,10 +105,14 @@ def __init__( "GenerativeRecLM: codebook must be non-empty " "(see design §3 — required field)" ) + # One entry per RQ level: ``len(codebook)`` = SID codes per item (the + # exact width of every answer), ``sum(codebook)`` = total SID atoms to + # append to the vocab. e.g. AL-GR is 3 levels x 8192 -> [8192,8192,8192]. + self._num_levels = len(codebook) sid_atoms = sum(int(c) for c in codebook) pad_mult = int(cfg.vocab_pad_to_multiple_of) or 128 - # --------- backbone + tokenizer ------------------------------- + # backbone + tokenizer hf_model_id = cfg.hf_model_id if not hf_model_id: raise ValueError( @@ -145,15 +126,10 @@ def __init__( self.lm = AutoModelForCausalLM.from_pretrained( hf_model_id, torch_dtype="auto" ) - # ``use_fast=True`` is the modern default; explicit for clarity. tokenizer = AutoTokenizer.from_pretrained(hf_model_id, use_fast=True) - # --------- vocab extension (codebook required) ---------------- - # The SID-atom base is the tokenizer's next free id BEFORE adding - # ``C0..``. For Qwen2.5-0.5B that's 151665 = `model.config.vocab_size` - # (151936, includes ~300 reserved padding slots) minus the unused - # reserved span — so use ``len(tokenizer)`` directly, NOT - # ``model.config.vocab_size``. + # vocab extension: base = tokenizer's next free id BEFORE adding C0.. + # (use len(tokenizer), NOT config.vocab_size which counts reserved slots). base = len(tokenizer) new_atoms = [f"C{i}" for i in range(sid_atoms)] added = tokenizer.add_tokens(new_atoms) @@ -165,14 +141,13 @@ def __init__( f"{sid_atoms} new atoms, only added {added}. " f"Aborting to avoid silent SID-token mismatch." ) - # Final vocab = base + sid_atoms, padded up to multiple of pad_mult. - # Matches algr's layout: SID atoms appended directly to the existing - # tokenizer vocab; offset arithmetic is `token = base + (sid - 1)`. + # SID atoms appended directly after the existing vocab (algr's layout); + # offset arithmetic is `token = base + (sid - 1)`. self.lm.resize_token_embeddings( base + sid_atoms, pad_to_multiple_of=pad_mult ) - # L3 safety check: assert C0 lands at the recorded base. + # assert C0 landed at the recorded base (the offset arithmetic relies on it) c0_id = tokenizer.convert_tokens_to_ids("C0") if c0_id != base: raise RuntimeError( @@ -180,64 +155,32 @@ def __init__( f"C0 at token id {base}, got {c0_id}. " f"Splice arithmetic would produce wrong token ids." ) - self._base_vocab = base # used in `_splice_input_ids` + self._base_vocab = base - # L1 + L7 mitigation: pad with eos_token_id on the LEFT side. + # pad token for the left-padded splice (fall back to eos) pad_id = tokenizer.pad_token_id if pad_id is None: pad_id = tokenizer.eos_token_id self._pad_token_id = int(pad_id) - # --------- cache chat-template buffers ------------------------ self._build_prompt_tokens(tokenizer, cfg) - # Diagnostics for the first launch — useful when chasing splice bugs. + # one-shot debug dump of the first spliced batch self._smoke_log_once = (os.environ.get("TZREC_GENRECLM_DEBUG", "0") == "1") self._first_predict = True - # ------------------------------------------------------------------ template def _build_prompt_tokens(self, tokenizer, cfg) -> None: - """Tokenise the family chat template once; cache as buffers. - - Composes the proto's optional ``system_instruction`` / - ``user_prefix_text`` / ``user_suffix_text`` (algr's CN prompt - wrappers — L2/L4 mitigations) with the family's static fragments: - - tpl_system = system_prefix + system_instruction + system_suffix - tpl_user_prefix = user_prefix + user_prefix_text - tpl_user_suffix = user_suffix_text + user_suffix - tpl_asst_prefix / tpl_asst_suffix verbatim from the template + """Family hook: cache the tokenised prompt template as buffers. - Buffers are non-persistent — they live with the module (move with - ``model.to(...)``) but stay off the state_dict so HF safetensors - round-tripping isn't polluted by TER-only state. + Called from ``__init__`` after vocab extension; the buffers it + registers are consumed by the family's ``predict``. Architecture- + specific — see design §15.1/§15.2. Subclasses MUST implement this. """ - tpl = type(self).CHAT_TEMPLATE - sys_text = cfg.system_instruction or tpl["default_system_instruction"] - u_pre = cfg.user_prefix_text or "" - u_suf = cfg.user_suffix_text or "" - frags = { - "system": tpl["system_prefix"] + sys_text + tpl["system_suffix"], - "user_prefix": tpl["user_prefix"] + u_pre, - "user_suffix": u_suf + tpl["user_suffix"], - "asst_prefix": tpl["asst_prefix"], - "asst_suffix": tpl["asst_suffix"], - } - for slot_name, frag_str in frags.items(): - ids = torch.tensor( - _encode_no_special(tokenizer, frag_str), dtype=torch.long - ) - self.register_buffer(f"tpl_{slot_name}", ids, persistent=False) - # algr appends eos to BOTH input_ids and labels at train time - # (algr/models/qwen2_5/data.py:46-47) — i.e. the trailing eos is a - # SUPERVISED token. Cache it so the splice can mirror that exactly. - self.register_buffer( - "tpl_eos", - torch.tensor([int(tokenizer.eos_token_id)], dtype=torch.long), - persistent=False, + raise NotImplementedError( + f"{type(self).__name__} must implement _build_prompt_tokens " + f"(GenerativeRecLM is abstract)." ) - # ------------------------------------------------------------------ init_input def init_input(self) -> None: """No-op override. @@ -246,160 +189,66 @@ def init_input(self) -> None: """ self.embedding_group = None - # ------------------------------------------------------------------ jagged -> rows - @staticmethod - def _jagged_to_row_list(jt) -> List[torch.Tensor]: - """Convert a TER JaggedTensor (values, lengths) to a list of 1-D - int64 row tensors. + @property + def device(self) -> torch.device: + """Device the HF backbone runs on — the single source for model I/O.""" + return self.lm.device + + def _tokenize_sids(self, sids: torch.Tensor) -> torch.Tensor: + """Map raw 1-indexed SID values to extended-vocab token ids. - ``values`` may arrive as float (TER's ``sequence_raw_feature`` reads - ``list`` as float — see [[project-tzrec-qwen2-integration]] - gotcha §B). We cast to long here; SID values fit in float32 mantissa - for any realistic codebook size (< 2^24). + Atom ``C{k}`` sits at ``base_vocab + k`` (atoms appended right after the + original vocab), so ``token_id = sid + base_vocab - 1``. The integer + counterpart of the HF tokenizer used for text; shape-agnostic. """ - values = jt.values() if callable(getattr(jt, "values", None)) else jt.values - lengths = jt.lengths() if callable(getattr(jt, "lengths", None)) else jt.lengths + return sids + (self._base_vocab - 1) + + def _sid_token_rows( + self, jt, expected_width: Optional[int] = None + ) -> List[torch.Tensor]: + """Read a SID jagged feature -> per-row token-id tensors. + + TER delivers the feature as a JaggedTensor (flat ``values`` + + ``lengths``); ``values`` may arrive as float / shape ``(N, 1)``. The + whole batch is tokenized once (``_tokenize_sids``) on the backbone + device, then split into rows. + + ``expected_width``, when set, enforces the sample contract here at the + data boundary: every row must have exactly that many codes (e.g. the + answer = ``num_levels``); a deviation is an anomalous sample. + """ + values = jt.values() + lengths = jt.lengths() if values.dim() == 2 and values.size(-1) == 1: values = values.squeeze(-1) - values = values.long() - lengths = lengths.long() - out: List[torch.Tensor] = [] - start = 0 - for n in lengths.tolist(): - out.append(values[start : start + n]) - start += n - return out - - # ------------------------------------------------------------------ splice - def _splice_input_ids( - self, - user_seq_rows: List[torch.Tensor], - label_rows: List[torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Build ``(input_ids, labels, attention_mask)``, each ``(B, T_max)``. - - Left-padded with ``eos_token_id``. ``attention_mask`` is essential — - without it self-attention would let pad positions pollute real - positions' hidden states. CE is separately protected by ``-100`` - labels at pad slots, but the forward needs the mask too. - - SID → token: ``token = sid + base_vocab - 1`` (SID atoms ``C0..`` - start at position ``base_vocab``; SID indices are 1-indexed). - """ - assert len(user_seq_rows) == len(label_rows) - B = len(user_seq_rows) - dev = self.tpl_system.device - base = self._base_vocab - - rows_ids: List[torch.Tensor] = [] - rows_lab: List[torch.Tensor] = [] - for i in range(B): - # SID → token id with int math. Map ``sid`` to its corresponding - # ``C{sid-1}`` atom: ``token = base + (sid - 1)`` = ``sid + (base-1)``. - # Cast to long matches the buffer dtype. - u_tok = (user_seq_rows[i].to(dev) + (base - 1)) - a_tok = (label_rows[i].to(dev) + (base - 1)) - ids = torch.cat([ - self.tpl_system, self.tpl_user_prefix, u_tok, - self.tpl_user_suffix, self.tpl_asst_prefix, a_tok, - self.tpl_asst_suffix, self.tpl_eos, - ]) - ign = torch.full_like(ids, self._ignore_index) - start = ( - self.tpl_system.numel() - + self.tpl_user_prefix.numel() - + u_tok.numel() - + self.tpl_user_suffix.numel() - + self.tpl_asst_prefix.numel() - ) - ign[start : start + a_tok.numel()] = a_tok - # algr supervises the trailing eos (after the masked - # ``<|im_end|>\n`` markers) — data.py:46-47. Mirror it. - ign[-1] = self.tpl_eos[0] - rows_ids.append(ids) - rows_lab.append(ign) - - T = max(r.numel() for r in rows_ids) - input_ids = torch.full( - (B, T), self._pad_token_id, dtype=torch.long, device=dev - ) - labels = torch.full( - (B, T), self._ignore_index, dtype=torch.long, device=dev - ) - attention_mask = torch.zeros((B, T), dtype=torch.long, device=dev) - for i, (ids, ign) in enumerate(zip(rows_ids, rows_lab)): - n = ids.numel() - # LEFT padding: write rows to the END of each (T,) slot. - input_ids[i, -n:] = ids - labels[i, -n:] = ign - attention_mask[i, -n:] = 1 - return input_ids, labels, attention_mask - - @staticmethod - def _min_first_non_neg_index(labels: torch.Tensor) -> int: - """Verbatim port of algr's helper (al_sid/algr/models/qwen2_5/ - modeling_qwen.py:1267-1274). - - Returns the smallest position (across rows in the batch) where the - first non-(-100) label appears. Used to decide how many trailing - positions to feed into ``lm_head``. - """ - tmp = (labels >= 0).cumsum(dim=-1) - return int((tmp == 1).float().argmax(dim=-1).min().item()) + # host-side split bounds, read before the H2D copy below + sizes = lengths.long().tolist() + if expected_width is not None: + bad = [i for i, n in enumerate(sizes) if n != expected_width] + if bad: + raise ValueError( + f"{type(self).__name__}: each SID item must be " + f"{expected_width} codes (len(codebook)); rows {bad} have " + f"{[sizes[i] for i in bad]} — anomalous sample(s)." + ) + # one vectorized SID->token map over the whole batch, on the backbone device + values = self._tokenize_sids(values.to(self.device).long()) + return list(torch.split(values, sizes)) - # ------------------------------------------------------------------ predict def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: - jt_u = batch.sequence_dense_features[self._input_name] - jt_l = batch.sequence_dense_features[self._label_name] - u_rows = self._jagged_to_row_list(jt_u) - l_rows = self._jagged_to_row_list(jt_l) - - input_ids, labels, attention_mask = self._splice_input_ids(u_rows, l_rows) - - if self._smoke_log_once and self._first_predict: - print( - f"[GENRECLM_DEBUG] first batch: B={input_ids.shape[0]} " - f"T={input_ids.shape[1]} pad_id={self._pad_token_id} " - f"ign={self._ignore_index} dev={input_ids.device} " - f"input_ids[0, -8:]={input_ids[0, -8:].tolist()} " - f"labels[0, -8:]={labels[0, -8:].tolist()}", - flush=True, - ) - self._first_predict = False + """Family hook: build inputs, run the HF forward, return ``{"loss": ...}``. - outputs = self.lm.model( - input_ids=input_ids, attention_mask=attention_mask - ) - hidden = outputs.last_hidden_state # (B, T, D) - - # Suffix slice in BOTH train and eval. algr only slices when - # training (its eval goes through a separate beam-search predict), - # but for CE the slice is value-identical (positions outside the - # suffix all carry -100 labels) and it bounds the logits tensor to - # (B, T_suffix, V) — without it, eval at bsz=80 would materialise - # (B, T, 217k) logits plus HF loss_function's fp32 upcast and OOM. - if (labels >= 0).any(): - keep = labels.shape[1] - self._min_first_non_neg_index(labels) + 1 - sl = slice(-keep, None) - labels_sl = labels[:, sl] - else: - sl = slice(None) - labels_sl = labels - - logits = self.lm.lm_head(hidden[:, sl, :]) - - # ``loss_function`` is the HF ``ForCausalLMLoss`` callable hung off - # every ``…ForCausalLM`` class; does shift-by-one + CE with -100 - # ignore. Calling it here matches algr's training-step loss exactly. - loss = self.lm.loss_function( - logits=logits, - labels=labels_sl, - vocab_size=self.lm.config.vocab_size, + Architecture-specific — the decoder-only implementation lives in + ``Qwen2RecLM``; other families (GPT-NeoX, Mamba, T5, …) need their own. + See design §15.4/§16. Subclasses MUST implement this. + """ + raise NotImplementedError( + f"{type(self).__name__} must implement predict " + f"(GenerativeRecLM is abstract)." ) - return {"loss": loss, "logits": logits} - # ------------------------------------------------------------------ loss def init_loss(self) -> None: + """No-op: the loss is computed inside ``predict`` (HF loss_function).""" return def loss( @@ -407,16 +256,13 @@ def loss( predictions: Dict[str, torch.Tensor], batch: Batch, ) -> Dict[str, torch.Tensor]: + """Surface the CE loss already computed in ``predict``.""" return {"ce_loss": predictions["loss"]} - # ------------------------------------------------------------------ metrics - # See [[project-tzrec-qwen2-integration]] gotcha §C: BaseModel only - # declares the eval-side metric methods; the train loop calls both - # families, so we must override both. + # BaseModel declares only the eval-side metric hooks, but the train loop + # calls both eval and train hooks, so both are overridden here. def init_metric(self) -> None: - # Mean CE over the eval set — gives `_evaluate` something to log - # (BaseModel.compute_metric iterates `_metric_modules` generically; - # torchmetrics handles the cross-rank sync at compute()). + """Register a mean-CE metric for the eval loop.""" self._metric_modules["ce_loss"] = torchmetrics.MeanMetric() def update_metric( @@ -425,9 +271,11 @@ def update_metric( batch: Batch, losses: Optional[Dict[str, torch.Tensor]] = None, ) -> None: + """Update the mean-CE metric with this batch's loss.""" self._metric_modules["ce_loss"].update(predictions["loss"].detach()) def init_train_metric(self) -> None: + """No-op: no train-time metric beyond the logged CE loss.""" return def update_train_metric( @@ -436,4 +284,5 @@ def update_train_metric( batch: Batch, losses: Optional[Dict[str, torch.Tensor]] = None, ) -> None: + """No-op: no train-time metric beyond the logged CE loss.""" return diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index 839637caa..dc4ecc229 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -15,18 +15,39 @@ } } -The family contributes ONLY its chat-template fragments; the splice, -algr-aligned forward, vocab extension and checkpoint plumbing all live in -the ``GenerativeRecLM`` base. A new LLM family is one file like this one -(see ``FINAL_DESIGN_GENERATIVE_REC_LM.md`` §2/§5 — ``Qwen3RecLM`` should -subclass ``Qwen2RecLM`` and override only what differs). +This subclass owns the decoder-only-chat implementation: the ChatML prompt +template, the causal-LM splice, and the ``.model``/``.lm_head`` forward +(design §15/§16). The ``GenerativeRecLM`` base owns the architecture-agnostic +plumbing (vocab extension, jagged→row, loss, metrics). + +The splice/forward here are generic to decoder-only families sharing Qwen2's +``.model``/``.lm_head`` layout (Llama/Mistral/Gemma/Phi — design §16), not +Qwen2-specific; only ``QWEN2_TEMPLATE`` is. When a second such family lands, +lift ``_splice_input_ids`` / ``_min_first_non_neg_index`` / ``predict`` (and +the ChatML ``_build_prompt_tokens``) into an intermediate +``DecoderOnlyChatRecLM`` base so each family is just its template. Until then +they live here. """ +from typing import Dict, List, Tuple + +import torch +from torch.nn.utils.rnn import pad_sequence + +from tzrec.datasets.utils import Batch from tzrec.models.generative_rec_lm import GenerativeRecLM + +def _encode_no_special(tokenizer, text: str) -> List[int]: + """Encode a fragment without prepending BOS / appending EOS specials. + + We're building the prompt manually from explicit ``<|im_start|>`` markers, + so we must NOT let the tokenizer's BOS/EOS handling double-emit them. + """ + return tokenizer.encode(text, add_special_tokens=False) + # Verbatim Qwen2 ChatML fragments. ``default_system_instruction`` matches -# algr/models/qwen2_5/data.py:73 ("default_instruction") bit-for-bit — the -# L2 mitigation in design §11. +# algr/models/qwen2_5/data.py:73 ("default_instruction") bit-for-bit. QWEN2_TEMPLATE = { "system_prefix": "<|im_start|>system\n", "system_suffix": "<|im_end|>\n", @@ -44,3 +65,168 @@ class Qwen2RecLM(GenerativeRecLM): """Qwen2 / Qwen2.5 generative-recommendation LM.""" CHAT_TEMPLATE = QWEN2_TEMPLATE + + def _build_prompt_tokens(self, tokenizer, cfg) -> None: + """Tokenise the family chat template once; cache as buffers. + + Composes the proto's optional ``system_instruction`` / + ``user_prefix_text`` / ``user_suffix_text`` (algr's CN prompt + wrappers) with the family's static fragments: + + tpl_system = system_prefix + system_instruction + system_suffix + tpl_user_prefix = user_prefix + user_prefix_text + tpl_user_suffix = user_suffix_text + user_suffix + tpl_asst_prefix / tpl_asst_suffix verbatim from the template + + Buffers are non-persistent — they live with the module (move with + ``model.to(...)``) but stay off the state_dict so HF safetensors + round-tripping isn't polluted by TER-only state. + """ + tpl = type(self).CHAT_TEMPLATE + sys_text = cfg.system_instruction or tpl["default_system_instruction"] + u_pre = cfg.user_prefix_text or "" + u_suf = cfg.user_suffix_text or "" + frags = { + "system": tpl["system_prefix"] + sys_text + tpl["system_suffix"], + "user_prefix": tpl["user_prefix"] + u_pre, + "user_suffix": u_suf + tpl["user_suffix"], + "asst_prefix": tpl["asst_prefix"], + "asst_suffix": tpl["asst_suffix"], + } + for slot_name, frag_str in frags.items(): + ids = torch.tensor( + _encode_no_special(tokenizer, frag_str), dtype=torch.long + ) + self.register_buffer(f"tpl_{slot_name}", ids, persistent=False) + # algr appends eos to BOTH input_ids and labels at train time + # (algr/models/qwen2_5/data.py:46-47) — i.e. the trailing eos is a + # SUPERVISED token. Cache it so the splice can mirror that exactly. + self.register_buffer( + "tpl_eos", + torch.tensor([int(tokenizer.eos_token_id)], dtype=torch.long), + persistent=False, + ) + + def _splice_input_ids( + self, + user_seq_rows: List[torch.Tensor], + label_rows: List[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build ``(input_ids, labels, attention_mask)``, each ``(B, T_max)``. + + Left-padded with ``eos_token_id``. ``attention_mask`` is essential — + without it self-attention would let pad positions pollute real + positions' hidden states. CE is separately protected by ``-100`` + labels at pad slots, but the forward needs the mask too. + + Every answer is exactly ``self._num_levels`` SID codes (one per codebook + level — validated at the data boundary in ``_sid_token_rows``), so the + supervised tail ``[answer | asst_suffix | eos]`` has a FIXED width and, + after left-padding, lands in the SAME columns for every row: ``labels`` + is built in one vectorized assignment (no per-row label loop). + ``input_ids`` still varies per row (the user history length differs). + + ``user_seq_rows`` / ``label_rows`` already hold extended-vocab token ids + on the model device (see ``_sid_token_rows``). + """ + assert len(user_seq_rows) == len(label_rows) + A = self._num_levels + + # input_ids: assembled per row (user history length varies), then + # left-padded into a (B, T) batch (real content right-aligned). + rows_ids = [ + torch.cat([ + self.tpl_system, self.tpl_user_prefix, user_seq_rows[i], + self.tpl_user_suffix, self.tpl_asst_prefix, label_rows[i], + self.tpl_asst_suffix, self.tpl_eos, + ]) + for i in range(len(user_seq_rows)) + ] + input_ids = pad_sequence( + rows_ids, batch_first=True, + padding_value=self._pad_token_id, padding_side="left", + ) + attention_mask = pad_sequence( + [torch.ones_like(r) for r in rows_ids], batch_first=True, + padding_value=0, padding_side="left", + ) + + # labels: the supervised tail is fixed-width, so left-padding aligns it + # to the same columns for every row -> one vectorized write. + # tail layout (from the end): [answer(A) | asst_suffix(s) | eos(1)]. + # ``tail <= T`` always holds: every row already contains those tokens. + B, T = input_ids.shape + s = self.tpl_asst_suffix.numel() + tail = A + s + 1 + labels = torch.full( + (B, T), self._ignore_index, dtype=torch.long, device=self.device + ) + labels[:, T - tail : T - tail + A] = torch.stack(label_rows) + labels[:, -1] = self.tpl_eos[0] # algr supervises the trailing eos + return input_ids, labels, attention_mask + + @staticmethod + def _min_first_non_neg_index(labels: torch.Tensor) -> int: + """Return the batch-min index of the first non-(-100) label. + + Verbatim port of algr's helper (al_sid/algr/models/qwen2_5/ + modeling_qwen.py:1267-1274) — the smallest position (across rows in + the batch) where the first non-(-100) label appears, used to decide + how many trailing positions to feed into ``lm_head``. + """ + tmp = (labels >= 0).cumsum(dim=-1) + return int((tmp == 1).float().argmax(dim=-1).min().item()) + + def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: + """Decoder-only forward: splice → ``.model`` → suffix-slice → CE loss.""" + # SID indices -> token ids once, at the data boundary (see + # _sid_token_rows); the splice then just assembles the prompt. + u_rows = self._sid_token_rows(batch.sequence_dense_features[self._input_name]) + l_rows = self._sid_token_rows( + batch.sequence_dense_features[self._label_name], + expected_width=self._num_levels, # answer = one item = num_levels codes + ) + + input_ids, labels, attention_mask = self._splice_input_ids(u_rows, l_rows) + + if self._smoke_log_once and self._first_predict: + print( + f"[GENRECLM_DEBUG] first batch: B={input_ids.shape[0]} " + f"T={input_ids.shape[1]} pad_id={self._pad_token_id} " + f"ign={self._ignore_index} dev={input_ids.device} " + f"input_ids[0, -8:]={input_ids[0, -8:].tolist()} " + f"labels[0, -8:]={labels[0, -8:].tolist()}", + flush=True, + ) + self._first_predict = False + + outputs = self.lm.model( + input_ids=input_ids, attention_mask=attention_mask + ) + hidden = outputs.last_hidden_state # (B, T, D) + + # Suffix slice in BOTH train and eval. algr only slices when + # training (its eval goes through a separate beam-search predict), + # but for CE the slice is value-identical (positions outside the + # suffix all carry -100 labels) and it bounds the logits tensor to + # (B, T_suffix, V) — without it, eval at bsz=80 would materialise + # (B, T, 217k) logits plus HF loss_function's fp32 upcast and OOM. + if (labels >= 0).any(): + keep = labels.shape[1] - self._min_first_non_neg_index(labels) + 1 + sl = slice(-keep, None) + labels_sl = labels[:, sl] + else: + sl = slice(None) + labels_sl = labels + + logits = self.lm.lm_head(hidden[:, sl, :]) + + # ``loss_function`` is the HF ``ForCausalLMLoss`` callable hung off + # every ``…ForCausalLM`` class; does shift-by-one + CE with -100 + # ignore. Calling it here matches algr's training-step loss exactly. + loss = self.lm.loss_function( + logits=logits, + labels=labels_sl, + vocab_size=self.lm.config.vocab_size, + ) + return {"loss": loss, "logits": logits} From 21c344a4e72dc1cec02f73aab5ebfb3d35701fac Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Tue, 9 Jun 2026 03:15:41 +0000 Subject: [PATCH 03/99] [test] generative-rec LM: base + Qwen2 splice/tokenize unit tests - generative_rec_lm_test: registry dispatch, abstract-hook errors, device property, _tokenize_sids offset map, _sid_token_rows split/cast/(N,1)-squeeze and answer-width validation (ok + violation). - qwen2_rec_lm_test: splice layout + label masking, left-padding/varied lengths, mask keeps trailing eos when pad==eos, _min_first_non_neg_index, _build_prompt_tokens buffer registration. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm_test.py | 101 +++++++++++++++++++++++ tzrec/models/qwen2_rec_lm_test.py | 110 +++++++++++++++++++++++++ 2 files changed, 211 insertions(+) create mode 100644 tzrec/models/generative_rec_lm_test.py create mode 100644 tzrec/models/qwen2_rec_lm_test.py diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py new file mode 100644 index 000000000..a97389082 --- /dev/null +++ b/tzrec/models/generative_rec_lm_test.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import types +import unittest + +import torch +from torch import nn + +from tzrec.models.generative_rec_lm import GenerativeRecLM +from tzrec.models.model import BaseModel +from tzrec.models.qwen2_rec_lm import Qwen2RecLM + + +class _FakeJT: + """Minimal stand-in for a TorchRec JaggedTensor (callable values/lengths).""" + + def __init__(self, values, lengths, dim2=False): + v = torch.tensor(values, dtype=torch.float) # TER delivers list as float + self._v = v.unsqueeze(-1) if dim2 else v + self._l = torch.tensor(lengths) + + def values(self): + return self._v + + def lengths(self): + return self._l + + +def _stub(num_levels=3, base_vocab=100, device="cpu"): + """A Qwen2RecLM with the base data-prep state wired up, but no HF backbone. + + Exercises the architecture-agnostic base methods (inherited by every family) + without downloading a model. + """ + m = object.__new__(Qwen2RecLM) + nn.Module.__init__(m) + m._base_vocab = base_vocab + m._num_levels = num_levels + m.lm = types.SimpleNamespace(device=torch.device(device)) + return m + + +class GenerativeRecLMTest(unittest.TestCase): + def test_registry_dispatch(self) -> None: + # importing qwen2_rec_lm auto-registers the family by class name + self.assertIs(BaseModel.create_class("Qwen2RecLM"), Qwen2RecLM) + self.assertTrue(issubclass(Qwen2RecLM, GenerativeRecLM)) + + def test_abstract_hooks_raise(self) -> None: + base = object.__new__(GenerativeRecLM) + with self.assertRaises(NotImplementedError): + base._build_prompt_tokens(None, None) + with self.assertRaises(NotImplementedError): + base.predict(None) + + def test_device_property(self) -> None: + self.assertEqual(_stub(device="cpu").device, torch.device("cpu")) + + def test_tokenize_sids(self) -> None: + m = _stub(base_vocab=100) # token = sid + base - 1 = sid + 99 + out = m._tokenize_sids(torch.tensor([1, 2, 3])) + self.assertEqual(out.tolist(), [100, 101, 102]) + self.assertEqual(out.dtype, torch.int64) + # shape-agnostic: a 2-D batch maps elementwise + out2 = m._tokenize_sids(torch.tensor([[1, 2], [3, 4]])) + self.assertEqual(out2.tolist(), [[100, 101], [102, 103]]) + + def test_sid_token_rows_split_and_cast(self) -> None: + m = _stub(base_vocab=100) + rows = m._sid_token_rows(_FakeJT([1, 2, 3, 4, 5], [3, 2])) + self.assertEqual([r.tolist() for r in rows], [[100, 101, 102], [103, 104]]) + self.assertTrue(all(r.dtype == torch.int64 for r in rows)) + + def test_sid_token_rows_squeezes_n1(self) -> None: + m = _stub(base_vocab=100) + rows = m._sid_token_rows(_FakeJT([1, 2, 3], [3], dim2=True)) # (N, 1) + self.assertEqual([r.tolist() for r in rows], [[100, 101, 102]]) + + def test_sid_token_rows_width_ok(self) -> None: + m = _stub(base_vocab=100, num_levels=3) + rows = m._sid_token_rows(_FakeJT([1, 2, 3, 4, 5, 6], [3, 3]), expected_width=3) + self.assertEqual([r.tolist() for r in rows], [[100, 101, 102], [103, 104, 105]]) + + def test_sid_token_rows_width_violation_raises(self) -> None: + m = _stub(base_vocab=100, num_levels=3) + with self.assertRaises(ValueError): + # second row has 2 codes, not 3 -> anomalous sample + m._sid_token_rows(_FakeJT([1, 2, 3, 4, 5], [3, 2]), expected_width=3) + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/qwen2_rec_lm_test.py new file mode 100644 index 000000000..020cb76a4 --- /dev/null +++ b/tzrec/models/qwen2_rec_lm_test.py @@ -0,0 +1,110 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import types +import unittest + +import torch +from torch import nn + +from tzrec.models.qwen2_rec_lm import Qwen2RecLM + + +def _stub(num_levels=3, base_vocab=100, pad_id=9, device="cpu"): + """A Qwen2RecLM with the splice-relevant state wired up, no HF backbone. + + Template buffers use tiny placeholder ids so the spliced layout is easy to + read; real buffers come from ``_build_prompt_tokens`` at init time. + """ + m = object.__new__(Qwen2RecLM) + nn.Module.__init__(m) + m._ignore_index = -100 + m._num_levels = num_levels + m._base_vocab = base_vocab + m._pad_token_id = pad_id + m.lm = types.SimpleNamespace(device=torch.device(device)) + for name, vals in { + "tpl_system": [10, 11], "tpl_user_prefix": [12], "tpl_user_suffix": [13], + "tpl_asst_prefix": [14], "tpl_asst_suffix": [15], "tpl_eos": [9], + }.items(): + m.register_buffer(name, torch.tensor(vals, dtype=torch.long), persistent=False) + return m + + +class Qwen2RecLMTest(unittest.TestCase): + def test_splice_layout_and_labels(self) -> None: + m = _stub() + u = [torch.tensor([100, 101, 102])] + a = [torch.tensor([200, 201, 202])] # 3 codes = num_levels + ids, labels, mask = m._splice_input_ids(u, a) + # [system | user_prefix | history | user_suffix | + # asst_prefix | answer | asst_suffix | eos] + self.assertEqual( + ids[0].tolist(), [10, 11, 12, 100, 101, 102, 13, 14, 200, 201, 202, 15, 9] + ) + # only the answer (cols 8-10) and the trailing eos (col 12) are supervised + self.assertEqual( + labels[0].tolist(), + [-100] * 8 + [200, 201, 202, -100, 9], + ) + self.assertEqual(mask[0].tolist(), [1] * 13) + + def test_left_padding_varied_lengths(self) -> None: + m = _stub() + u = [torch.tensor([100, 101, 102, 103]), torch.tensor([100])] + a = [torch.tensor([200, 201, 202]), torch.tensor([207, 208, 209])] + ids, labels, mask = m._splice_input_ids(u, a) + T = ids.shape[1] + n1 = 2 + 1 + 1 + 1 + 1 + 3 + 1 + 1 # shorter row's real length + # shorter row is left-padded: pad at the front, content right-aligned + self.assertEqual(ids[1, : T - n1].tolist(), [m._pad_token_id] * (T - n1)) + self.assertEqual(mask[1].tolist(), [0] * (T - n1) + [1] * n1) + self.assertEqual(labels[1, : T - n1].tolist(), [-100] * (T - n1)) + # every row's trailing eos is supervised and the answer ends just before + self.assertEqual(labels[:, -1].tolist(), [9, 9]) + + def test_mask_keeps_trailing_eos_when_pad_equals_eos(self) -> None: + # pad_id == eos value: the mask must NOT mask the real trailing eos + m = _stub(pad_id=9) # tpl_eos == 9 too + ids, _, mask = m._splice_input_ids( + [torch.tensor([100])], [torch.tensor([200, 201, 202])] + ) + self.assertEqual(int(ids[0, -1]), 9) + self.assertEqual(int(mask[0, -1]), 1) + self.assertEqual(mask[0].tolist(), [1] * ids.shape[1]) + + def test_min_first_non_neg_index(self) -> None: + labels = torch.tensor([[-100, -100, 5, 6], [-100, 7, 8, 9]]) + self.assertEqual(Qwen2RecLM._min_first_non_neg_index(labels), 1) + + def test_build_prompt_tokens_registers_buffers(self) -> None: + m = object.__new__(Qwen2RecLM) + nn.Module.__init__(m) + tok = types.SimpleNamespace( + eos_token_id=99, + encode=lambda text, add_special_tokens=False: [len(text)], + ) + cfg = types.SimpleNamespace( + system_instruction="", user_prefix_text="", user_suffix_text="" + ) + m._build_prompt_tokens(tok, cfg) + for name in [ + "tpl_system", "tpl_user_prefix", "tpl_user_suffix", + "tpl_asst_prefix", "tpl_asst_suffix", "tpl_eos", + ]: + buf = getattr(m, name) + self.assertIsInstance(buf, torch.Tensor) + self.assertEqual(buf.dtype, torch.int64) + self.assertEqual(m.tpl_eos.tolist(), [99]) # eos cached for supervision + + +if __name__ == "__main__": + unittest.main() From bfe657fb9d217291beff62b4b0e1e32d1eb04f79 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Tue, 9 Jun 2026 03:26:34 +0000 Subject: [PATCH 04/99] [feat] generative-rec LM: branch predict into train/eval loss vs inference beam search predict() dispatches on the TER inference flag (BaseModule.is_inference, set by main.py's set_is_inference before the predict/export wrappers): - Branch 1 (not is_inference -> train/eval): existing teacher-forced forward + suffix-slice + CE loss (moved to _predict_train). - Branch 2 (is_inference -> inference): _generate beam-searches the SID answer from an answer-less prompt (_splice_prompt_ids), emitting num_levels tokens/beam and mapping them back to raw SID indices -> {"generated_sids": (B, num_return, L)}. Beam params (_num_beams/_num_return) default to algr's 50/50 (optional proto fields). Tests cover prompt splice, is_inference routing, and token->SID map. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 7 ++ tzrec/models/qwen2_rec_lm.py | 87 +++++++++++++++++++--- tzrec/models/qwen2_rec_lm_test.py | 41 ++++++++++ tzrec/protos/models/generative_model.proto | 6 ++ 4 files changed, 132 insertions(+), 9 deletions(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index a6acc3736..705be9a3b 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -63,6 +63,10 @@ class GenerativeRecLM(BaseModel): ``generative_rec_lm.class_name`` (resolved through the BaseModel registry). """ + # predictions key the inference branch emits generated SIDs under, stable + # across families (PredictWrapper ``output_cols`` should reference it). + GENERATED_SIDS_KEY = "generated_sids" + def __new__(cls, model_config: ModelConfig, *args: Any, **kwargs: Any): """Dispatch to the concrete family subclass. @@ -111,6 +115,9 @@ def __init__( self._num_levels = len(codebook) sid_atoms = sum(int(c) for c in codebook) pad_mult = int(cfg.vocab_pad_to_multiple_of) or 128 + # beam-search params for the inference branch (proto defaults = 50/50) + self._num_beams = int(cfg.num_beams) + self._num_return = int(cfg.num_return_sequences) # backbone + tokenizer hf_model_id = cfg.hf_model_id diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index dc4ecc229..0b94ff87b 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -142,14 +142,7 @@ def _splice_input_ids( ]) for i in range(len(user_seq_rows)) ] - input_ids = pad_sequence( - rows_ids, batch_first=True, - padding_value=self._pad_token_id, padding_side="left", - ) - attention_mask = pad_sequence( - [torch.ones_like(r) for r in rows_ids], batch_first=True, - padding_value=0, padding_side="left", - ) + input_ids, attention_mask = self._left_pad(rows_ids) # labels: the supervised tail is fixed-width, so left-padding aligns it # to the same columns for every row -> one vectorized write. @@ -178,7 +171,19 @@ def _min_first_non_neg_index(labels: torch.Tensor) -> int: return int((tmp == 1).float().argmax(dim=-1).min().item()) def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: - """Decoder-only forward: splice → ``.model`` → suffix-slice → CE loss.""" + """Dispatch on the TER inference flag (``set_is_inference`` in main.py). + + Branch 1 (train / eval, ``not is_inference``) — teacher-forced forward + + CE loss (the metric path). + Branch 2 (inference, ``is_inference``) — beam-search the SID answer from + the prompt. + """ + if self.is_inference: + return self._generate(batch) + return self._predict_train(batch) + + def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: + """Branch 1: teacher-forced forward -> suffix-slice -> CE loss.""" # SID indices -> token ids once, at the data boundary (see # _sid_token_rows); the splice then just assembles the prompt. u_rows = self._sid_token_rows(batch.sequence_dense_features[self._input_name]) @@ -230,3 +235,67 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: vocab_size=self.lm.config.vocab_size, ) return {"loss": loss, "logits": logits} + + def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: + """Branch 2: beam-search the SID answer (no ground truth supplied). + + Builds the prompt (no answer), generates exactly ``num_levels`` new + tokens per beam, and maps them back to raw SID indices. Returns + ``generated_sids`` of shape ``(B, num_return, num_levels)``. + """ + u_rows = self._sid_token_rows(batch.sequence_dense_features[self._input_name]) + input_ids, attention_mask = self._splice_prompt_ids(u_rows) + out = self.lm.generate( + input_ids=input_ids, + attention_mask=attention_mask, + max_new_tokens=self._num_levels, + num_beams=self._num_beams, + num_return_sequences=self._num_return, + do_sample=False, + pad_token_id=self._pad_token_id, + ) + # keep only the generated tail; map token ids back to raw SID indices + # (inverse of _tokenize_sids: sid = token - base_vocab + 1). + new_tokens = out[:, input_ids.shape[1]:] + sids = new_tokens - (self._base_vocab - 1) + # generate() returns rows grouped batch-major: [b0_beam0, b0_beam1, ..., + # b1_beam0, ...], so this view groups beams under the right user. + sids = sids.view(input_ids.shape[0], self._num_return, self._num_levels) + return {self.GENERATED_SIDS_KEY: sids} + + def _splice_prompt_ids( + self, user_seq_rows: List[torch.Tensor] + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Assemble the answer-less prompt and left-pad into ``(B, T_max)``. + + Layout: ``[system | user_prefix | history | user_suffix | asst_prefix]`` + — everything up to (but not including) the answer, so generation + continues from the assistant turn. + """ + rows = [ + torch.cat([ + self.tpl_system, self.tpl_user_prefix, r, + self.tpl_user_suffix, self.tpl_asst_prefix, + ]) + for r in user_seq_rows + ] + return self._left_pad(rows) + + def _left_pad( + self, rows: List[torch.Tensor] + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Left-pad token rows into ``(input_ids, attention_mask)``, ``(B, T_max)``. + + Real content is right-aligned, pad at the front. ``attention_mask`` is + built from ``ones_like(row)`` (not ``!= pad``) so a real trailing eos is + never masked when ``pad_token_id == eos``. + """ + input_ids = pad_sequence( + rows, batch_first=True, + padding_value=self._pad_token_id, padding_side="left", + ) + attention_mask = pad_sequence( + [torch.ones_like(r) for r in rows], batch_first=True, + padding_value=0, padding_side="left", + ) + return input_ids, attention_mask diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/qwen2_rec_lm_test.py index 020cb76a4..f2c9c90f3 100644 --- a/tzrec/models/qwen2_rec_lm_test.py +++ b/tzrec/models/qwen2_rec_lm_test.py @@ -85,6 +85,47 @@ def test_min_first_non_neg_index(self) -> None: labels = torch.tensor([[-100, -100, 5, 6], [-100, 7, 8, 9]]) self.assertEqual(Qwen2RecLM._min_first_non_neg_index(labels), 1) + def test_splice_prompt_ids(self) -> None: + m = _stub() + ids, mask = m._splice_prompt_ids([torch.tensor([100, 101, 102])]) + # [system | user_prefix | history | user_suffix | asst_prefix], no answer + self.assertEqual(ids[0].tolist(), [10, 11, 12, 100, 101, 102, 13, 14]) + self.assertEqual(mask[0].tolist(), [1] * 8) + + def test_predict_routes_on_inference_flag(self) -> None: + m = _stub() + m._predict_train = lambda b: {"branch": "train"} + m._generate = lambda b: {"branch": "generate"} + m._is_inference = False # train / eval + self.assertEqual(Qwen2RecLM.predict(m, object())["branch"], "train") + m._is_inference = True # inference (set_is_inference in main.py) + self.assertEqual(Qwen2RecLM.predict(m, object())["branch"], "generate") + + def test_generate_maps_tokens_to_sids(self) -> None: + m = _stub(base_vocab=100) # sid = token - base + 1 = token - 99 + m._input_name = "user_sequence" + m._num_beams = m._num_return = 2 + + def fake_generate(input_ids, attention_mask, max_new_tokens, + num_beams, num_return_sequences, do_sample, pad_token_id): + prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) + new = torch.tensor([[200, 201, 202], [203, 204, 205]]) # 2 beams x 3 codes + return torch.cat([prompt, new], dim=1) + + m.lm.generate = fake_generate + + class _JT: + def values(self): + return torch.tensor([1, 2, 3], dtype=torch.float) + + def lengths(self): + return torch.tensor([3]) + + batch = types.SimpleNamespace(sequence_dense_features={"user_sequence": _JT()}) + sids = m._generate(batch)["generated_sids"] + self.assertEqual(tuple(sids.shape), (1, 2, 3)) # (B, num_return, num_levels) + self.assertEqual(sids[0].tolist(), [[101, 102, 103], [104, 105, 106]]) + def test_build_prompt_tokens_registers_buffers(self) -> None: m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index f57246ff9..af02224c2 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -69,4 +69,10 @@ message GenerativeRecLM { // Cross-entropy ignore index — matches PyTorch's F.cross_entropy default. optional int32 ignore_index = 53 [default = -100]; + + // ----- Inference (beam search) ----- + // Used only by the inference branch of predict() (tzrec.predict / export). + // Defaults match algr's predict config. + optional uint32 num_beams = 80 [default = 50]; + optional uint32 num_return_sequences = 81 [default = 50]; } From 9acfd86f613580a494da0aaf806dd3a74c847754 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Tue, 9 Jun 2026 07:07:23 +0000 Subject: [PATCH 05/99] [refactor] generative-rec LM: register families directly, drop class_name dispatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each LLM family is its own model_config oneof entry whose message-type name resolves straight to the same-named class (qwen2_rec_lm -> Qwen2RecLM), so GenerativeRecLM.__new__ and the class_name field are gone. Proto split: - GenerativeRecLMConfig = architecture-AGNOSTIC config (codebook, vocab pad, feature names, ignore_index, beam params), embedded as `common`. - The backbone `hf_model_id` is OWNED by the family message (Qwen2RecLM, default "Qwen/Qwen2.5-0.5B"), NOT `common` — the registered family IS the architecture commitment, so a backbone in the shared block could contradict it. (common field 1 reserved.) - Family-specific chat-template knobs also live on the family message. Base __init__ reads cfg.common.* (shared) and cfg.hf_model_id (family-owned). Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 84 +++++++-------- tzrec/models/generative_rec_lm_test.py | 14 +++ tzrec/models/qwen2_rec_lm.py | 9 +- tzrec/protos/model.proto | 4 +- tzrec/protos/models/generative_model.proto | 117 +++++++++------------ 5 files changed, 109 insertions(+), 119 deletions(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 705be9a3b..afd38fe55 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -11,9 +11,11 @@ * Per-family subclasses (design §2 / G4): ``GenerativeRecLM`` is the abstract base; each LLM family is a concrete subclass implementing the ``_build_prompt_tokens`` and ``predict`` hooks (e.g. ``Qwen2RecLM`` in - ``tzrec/models/qwen2_rec_lm.py``). The pipeline config selects the - family via ``generative_rec_lm.class_name``; dispatch goes through the - BaseModel registry (subclasses auto-register by class name). + ``tzrec/models/qwen2_rec_lm.py``). The pipeline config selects the family + by its own oneof entry (``qwen2_rec_lm``), whose message-type name resolves + directly to the same-named class via the BaseModel registry — no dispatch. + Shared config lives in ``GenerativeRecLMConfig`` (the family message's + ``common`` field); family-specific knobs sit on the family message. * Streaming sample format: each row carries two raw-int64 sequence features, ``user_sequence`` (list[int]) and ``label`` (list[int]), both holding raw SID indices in ``[1, sum(codebook)]``. @@ -56,38 +58,22 @@ class GenerativeRecLM(BaseModel): _build_prompt_tokens(tokenizer, cfg) — cache the prompt template predict(batch) — build inputs + HF forward + Family proto contract: every family message embeds + ``GenerativeRecLMConfig common = 1`` (shared config the base reads) and + supplies a backbone, by default via an ``hf_model_id`` field (overridable + through ``_backbone_id``). + ``Qwen2RecLM`` (``tzrec/models/qwen2_rec_lm.py``) provides the decoder-only chat implementation (ChatML splice + ``.model``/``.lm_head`` forward), reusable by Llama/Mistral/Gemma/Phi-style families; GPT-NeoX/RWKV/Mamba/T5 - each need their own. The pipeline config selects the family via - ``generative_rec_lm.class_name`` (resolved through the BaseModel registry). + each need their own. Each family registers directly (its oneof message-type + name == the class name); there is no ``class_name`` dispatch. """ # predictions key the inference branch emits generated SIDs under, stable # across families (PredictWrapper ``output_cols`` should reference it). GENERATED_SIDS_KEY = "generated_sids" - def __new__(cls, model_config: ModelConfig, *args: Any, **kwargs: Any): - """Dispatch to the concrete family subclass. - - ``tzrec.main._create_model`` resolves the proto oneof message name - (``GenerativeRecLM``) to THIS class; the actual family is selected - by ``generative_rec_lm.class_name`` and looked up in the BaseModel - registry (every subclass auto-registers via the metaclass). - """ - if cls is GenerativeRecLM: - cfg = getattr(model_config, model_config.WhichOneof("model")) - class_name = cfg.class_name - # pyre-ignore [16] - sub_cls = BaseModel.create_class(class_name) - if not issubclass(sub_cls, GenerativeRecLM): - raise ValueError( - f"generative_rec_lm.class_name = {class_name!r} resolves " - f"to {sub_cls}, which is not a GenerativeRecLM subclass." - ) - return super().__new__(sub_cls) - return super().__new__(cls) - def __init__( self, model_config: ModelConfig, @@ -97,39 +83,35 @@ def __init__( **kwargs: Any, ) -> None: super().__init__(model_config, features, labels, sample_weights, **kwargs) - cfg = self._model_config # populated by BaseModel from WhichOneof - - # proto -> python knobs - self._input_name: str = cfg.user_sequence_feature_name - self._label_name: str = cfg.label_feature_name - self._ignore_index: int = int(cfg.ignore_index) - codebook = list(cfg.codebook) + cfg = self._model_config # the family message (e.g. Qwen2RecLM) + common = cfg.common # GenerativeRecLMConfig — shared by all families + + # shared proto -> python knobs + self._input_name: str = common.user_sequence_feature_name + self._label_name: str = common.label_feature_name + self._ignore_index: int = int(common.ignore_index) + codebook = list(common.codebook) if len(codebook) == 0: raise ValueError( "GenerativeRecLM: codebook must be non-empty " "(see design §3 — required field)" ) - # One entry per RQ level: ``len(codebook)`` = SID codes per item (the - # exact width of every answer), ``sum(codebook)`` = total SID atoms to - # append to the vocab. e.g. AL-GR is 3 levels x 8192 -> [8192,8192,8192]. + # len(codebook) = SID codes per item (answer width); sum = vocab atoms. self._num_levels = len(codebook) sid_atoms = sum(int(c) for c in codebook) - pad_mult = int(cfg.vocab_pad_to_multiple_of) or 128 + pad_mult = int(common.vocab_pad_to_multiple_of) or 128 # beam-search params for the inference branch (proto defaults = 50/50) - self._num_beams = int(cfg.num_beams) - self._num_return = int(cfg.num_return_sequences) + self._num_beams = int(common.num_beams) + self._num_return = int(common.num_return_sequences) - # backbone + tokenizer - hf_model_id = cfg.hf_model_id + # backbone + tokenizer (the backbone is family-owned; see _backbone_id) + hf_model_id = self._backbone_id() if not hf_model_id: raise ValueError( - "GenerativeRecLM v1: hf_model_id is required " - "(architecture-spec path deferred to v1.x)" + f"{type(self).__name__}: empty backbone id (see _backbone_id)." ) - # torch_dtype="auto" preserves the safetensors-stored dtype (bf16 - # for Qwen2.5-0.5B). The default would silently upcast to fp32. - # On CPU the same flag is honoured; on GPU it avoids a 2× memory - # blow-up. + # torch_dtype="auto" keeps the stored dtype (bf16); the default upcasts + # to fp32 (2x memory on GPU). self.lm = AutoModelForCausalLM.from_pretrained( hf_model_id, torch_dtype="auto" ) @@ -176,6 +158,14 @@ def __init__( self._smoke_log_once = (os.environ.get("TZREC_GENRECLM_DEBUG", "0") == "1") self._first_predict = True + def _backbone_id(self) -> str: + """Family hook: the HF model id to load for ``self.lm``. + + Defaults to the family message's ``hf_model_id``; override if a family + sources its backbone differently. + """ + return self._model_config.hf_model_id + def _build_prompt_tokens(self, tokenizer, cfg) -> None: """Family hook: cache the tokenised prompt template as buffers. diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py index a97389082..630195004 100644 --- a/tzrec/models/generative_rec_lm_test.py +++ b/tzrec/models/generative_rec_lm_test.py @@ -55,6 +55,20 @@ def test_registry_dispatch(self) -> None: self.assertIs(BaseModel.create_class("Qwen2RecLM"), Qwen2RecLM) self.assertTrue(issubclass(Qwen2RecLM, GenerativeRecLM)) + def test_backbone_owned_by_family_proto(self) -> None: + # the backbone lives on the family message (its architecture), NOT in + # the shared common config; it defaults to the canonical Qwen2.5-0.5B. + from tzrec.protos.models.generative_model_pb2 import ( + GenerativeRecLMConfig, + ) + from tzrec.protos.models.generative_model_pb2 import ( + Qwen2RecLM as Qwen2RecLMProto, + ) + + self.assertEqual(Qwen2RecLMProto().hf_model_id, "Qwen/Qwen2.5-0.5B") + common_fields = [f.name for f in GenerativeRecLMConfig.DESCRIPTOR.fields] + self.assertNotIn("hf_model_id", common_fields) + def test_abstract_hooks_raise(self) -> None: base = object.__new__(GenerativeRecLM) with self.assertRaises(NotImplementedError): diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index 0b94ff87b..c3dad8e9a 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -6,12 +6,13 @@ """Qwen2/Qwen2.5 family subclass of ``GenerativeRecLM`` (design §5). -Selected from the pipeline config via:: +Selected from the pipeline config by its own oneof entry (the message-type name +resolves directly to this class — no ``class_name`` dispatch):: model_config { - generative_rec_lm { - class_name: "Qwen2RecLM" - ... + qwen2_rec_lm { + common { hf_model_id: "..." codebook: 8192 ... } + system_instruction: "..." } } diff --git a/tzrec/protos/model.proto b/tzrec/protos/model.proto index 05419f0e3..299722145 100644 --- a/tzrec/protos/model.proto +++ b/tzrec/protos/model.proto @@ -78,8 +78,8 @@ message ModelConfig { RocketLaunching rocket_launching = 500; - // Generative (causal-LM) models. - GenerativeRecLM generative_rec_lm = 601; + // Generative (causal-LM) models — one family message per LLM family. + Qwen2RecLM qwen2_rec_lm = 601; } // Field 600 was the removed v1 offline-tokenized `Qwen2` wrapper — diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index af02224c2..9e19579f5 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -3,76 +3,61 @@ package tzrec.protos; // Generative (causal-LM) models. // -// These differ from the rank/match/multi-task categories in that the model's -// output is a per-position vocabulary distribution rather than a single score -// or embedding. - -// GenerativeRecLM — the unified config for HF-backed generative recommendation -// LMs. Dispatch to a concrete subclass via `class_name`; see design doc -// §3 (FINAL_DESIGN_GENERATIVE_REC_LM.md). +// Each LLM family is registered DIRECTLY: the model_config oneof picks a family +// message (e.g. `qwen2_rec_lm`), whose message type name (`Qwen2RecLM`) resolves +// to the Python class of the same name. There is no `class_name` dispatch. // -// Sample contract (consumed by `predict()`): +// Shared, architecture-agnostic config lives in `GenerativeRecLMConfig` and is +// embedded as `common` in every family message; family-specific knobs (e.g. the +// chat template) live on the family message itself. Adding a family = a new +// message like `Qwen2RecLM` + a same-named Python subclass of GenerativeRecLM. + +// Architecture-agnostic config shared by ALL generative-rec families (the base +// reads this for everything except the backbone, which the family owns — see +// _backbone_id). Sample contract (consumed by `predict()`): // * user_sequence : list — raw SID indices in [1, sum(codebook)] // * label : list — raw SID indices in [1, sum(codebook)] -// Both flow into TER as `sequence_raw_feature` parquet columns. The model -// converts SID → token id at batch time via integer offset arithmetic and -// splices them between the cached chat-template buffers built at __init__. -message GenerativeRecLM { - // Dispatches to the concrete Python subclass. Valid values: - // "Qwen2RecLM" — Qwen2 family (Qwen2.5-0.5B, etc.) - // Future: "Llama3RecLM", "MistralRecLM", "MixtralRecLM" (MoE override). - required string class_name = 1; - - // ----- Architecture source ----- - // Either supply an HF hub / local model id (loaded via - // `AutoModelForCausalLM.from_pretrained(..., torch_dtype="auto")`), or - // give an architecture spec to randomly initialise. Exactly one must be - // set; checked at runtime. - optional string hf_model_id = 2 [default = ""]; - // (architecture spec block is deferred to v1.x — `hf_model_id` is the - // only path validated in v1; matches algr's setup.) - - // ----- SID vocabulary ----- - // Required, non-empty. Each entry is the codebook size at one RQ layer - // (e.g. [512, 512] for a 2-layer RQ with 512 codes/layer). Total SID - // atoms = sum(codebook); these are added to the tokenizer / model - // vocabulary as `C0 .. C{sum-1}` directly after the base vocabulary - // (no [SEP]; matches algr's add_tokens layout). SID indices in the - // parquet are 1-indexed; the wrapper maps `sid → base_vocab + (sid-1)` - // where base_vocab = len(tokenizer) BEFORE the extension. - repeated uint32 codebook = 60; - - // Pad the post-extension vocabulary up to a multiple of this value. - // 128 matches algr's `model.resize_token_embeddings(..., pad_to_multiple_of=128)`. - optional uint32 vocab_pad_to_multiple_of = 61 [default = 128]; - - // ----- Chat template ----- - // Optional override for the system instruction string. When unset, the - // subclass's default (e.g. Qwen2RecLM's algr-matching default) is used. - optional string system_instruction = 70 [default = ""]; - - // Optional CN/EN text fragments wrapping the SID codes inside the user - // message. For algr's tiny dataset these are - // user_prefix_text = "当前用户的历史行为如下：" - // user_suffix_text = "，请预测用户在电商推荐场景后续行为的语义编码" - // Tokenised once at __init__ and concatenated INSIDE the cached - // ``tpl_user_prefix`` and ``tpl_user_suffix`` buffers (after the - // ``<|im_start|>user\n`` marker, before the ``<|im_end|>\n`` marker). - // Leaving them empty produces the algr DEFAULT_INSTRUCTION path. - optional string user_prefix_text = 71 [default = ""]; - optional string user_suffix_text = 72 [default = ""]; - - // ----- Sample feature names ----- - // Which `sequence_raw_feature` blocks in the parquet carry the SID lists. - required string user_sequence_feature_name = 51; - required string label_feature_name = 52; +// Both flow into TER as `sequence_raw_feature` parquet columns; the model maps +// SID -> token id by integer offset at batch time. +message GenerativeRecLMConfig { + // Backbone (`hf_model_id`) is NOT here — it's the family's architecture + // commitment, so it lives on the family message. (Old class_name was 1.) + reserved 1; + + // SID vocabulary, one entry per RQ level: len = SID codes per item (answer + // width), sum = atoms appended as C0..C{sum-1} after the base vocab. + repeated uint32 codebook = 2; + // Pad the post-extension vocab up to a multiple of this value. + optional uint32 vocab_pad_to_multiple_of = 3 [default = 128]; + + // Which `sequence_raw_feature` parquet columns carry the SID lists. + required string user_sequence_feature_name = 4; + required string label_feature_name = 5; // Cross-entropy ignore index — matches PyTorch's F.cross_entropy default. - optional int32 ignore_index = 53 [default = -100]; + optional int32 ignore_index = 6 [default = -100]; + + // Inference (beam search) — used only by predict()'s inference branch. + optional uint32 num_beams = 7 [default = 50]; + optional uint32 num_return_sequences = 8 [default = 50]; +} - // ----- Inference (beam search) ----- - // Used only by the inference branch of predict() (tzrec.predict / export). - // Defaults match algr's predict config. - optional uint32 num_beams = 80 [default = 50]; - optional uint32 num_return_sequences = 81 [default = 50]; +// Qwen2 / Qwen2.5 family (Qwen2.5-0.5B, etc.). Registered directly via its +// message-type name; the Python class is `Qwen2RecLM`. +message Qwen2RecLM { + // Architecture-agnostic config shared by all generative-rec families. + optional GenerativeRecLMConfig common = 1; + + // Qwen2 backbone (HF hub id or local path; must be a Qwen2 model). Owned by + // this family message, not `common`. Default = canonical 0.5B. + optional string hf_model_id = 2 [default = "Qwen/Qwen2.5-0.5B"]; + + // ----- Chat template (family-specific) ----- + // Optional override for the system instruction. Empty -> the family default + // (Qwen2RecLM's algr-matching default). + optional string system_instruction = 10 [default = ""]; + // Optional CN/EN text wrapping the SID codes in the user message, e.g. + // "当前用户的历史行为如下：" / "，请预测用户在电商推荐场景后续行为的语义编码". + optional string user_prefix_text = 11 [default = ""]; + optional string user_suffix_text = 12 [default = ""]; } From 632fce44bcc071c34b150ffad21016243b45d418 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 10 Jun 2026 06:42:26 +0000 Subject: [PATCH 06/99] [feat] generative-rec LM: pre-allocate activation pool to fix per-rank VRAM imbalance One GPU reserved ~25GB more than the rest during 8-GPU training: the CUDA caching allocator stranded a whole segment generation on whichever rank drew its shortest batch first (variable seq-len + native allocator never shrinks reserved). Pre-size the pool up front so it never has to grow mid-run. - Qwen2RecLM warms the activation pool with a one-shot dummy fwd+bwd at the worst-case (batch_size, T_max) on the first training step (earliest the HF backbone is on-GPU); T_max = template frame + sequence_length + num_levels. - Pool length reuses the user_sequence feature's sequence_length (GenerativeRecLM._input_sequence_length); _sid_token_rows enforces it with a recency-preserving, item-aligned tail clip (keep newest items, drop oldest) so it is a guaranteed bound under FG_NONE, which does not truncate. - Thread data_config.batch_size into _create_model; move num_beams/num_return from base to subclass. Verified on 8xGPU: per-card spread 25GB -> 2.2GB, no OOM. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/main.py | 7 ++ tzrec/models/generative_rec_lm.py | 44 +++++++++++-- tzrec/models/qwen2_rec_lm.py | 102 ++++++++++++++++++++++++++++-- tzrec/models/qwen2_rec_lm_test.py | 90 ++++++++++++++++++++++++++ 4 files changed, 233 insertions(+), 10 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index 87f2984fb..b4e4fb68a 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -130,6 +130,7 @@ def _create_model( labels: List[str], sample_weights: Optional[List[str]] = None, sampler_type: Optional[str] = None, + batch_size: Optional[int] = None, ) -> BaseModel: """Build model. @@ -139,6 +140,9 @@ def _create_model( labels (list): list of label names. sample_weights (list): list of sample weight names. sampler_type (str): negative sampler type + batch_size (int): per-rank batch size (data_config.batch_size); most + models ignore it, generative LMs use it to pre-size their pool. + Return: model: a EasyRec Model. """ @@ -152,6 +156,7 @@ def _create_model( labels, sample_weights=sample_weights, sampler_type=sampler_type, + batch_size=batch_size, ) kernel = Kernel[KernelProto.Name(model_config.kernel)] @@ -643,6 +648,7 @@ def train_and_evaluate( list(data_config.label_fields), sample_weights=list(data_config.sample_weight_fields), sampler_type=sampler_type, + batch_size=data_config.batch_size, ) model = TrainWrapper( model, device=device, mixed_precision=train_config.mixed_precision @@ -822,6 +828,7 @@ def evaluate( list(data_config.label_fields), sample_weights=list(data_config.sample_weight_fields), sampler_type=sampler_type, + batch_size=data_config.batch_size, ) model = TrainWrapper( model, device=device, mixed_precision=train_config.mixed_precision diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index afd38fe55..1a357d0df 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -82,6 +82,10 @@ def __init__( sample_weights: Optional[List[str]] = None, **kwargs: Any, ) -> None: + # per-rank batch size, threaded from data_config.batch_size by + # _create_model (absorbed by BaseModule's **kwargs); used to pre-size + # the activation pool. 0 = unknown (e.g. export/predict construction). + self._batch_size: int = int(kwargs.get("batch_size") or 0) super().__init__(model_config, features, labels, sample_weights, **kwargs) cfg = self._model_config # the family message (e.g. Qwen2RecLM) common = cfg.common # GenerativeRecLMConfig — shared by all families @@ -90,6 +94,10 @@ def __init__( self._input_name: str = common.user_sequence_feature_name self._label_name: str = common.label_feature_name self._ignore_index: int = int(common.ignore_index) + # max history length (SID codes) for activation pre-allocation, taken + # from the user-sequence feature's truncation length (the data reader + # caps every row at it, so it's the guaranteed upper bound). 0 = off. + self._max_seq_length: int = self._input_sequence_length() codebook = list(common.codebook) if len(codebook) == 0: raise ValueError( @@ -100,9 +108,6 @@ def __init__( self._num_levels = len(codebook) sid_atoms = sum(int(c) for c in codebook) pad_mult = int(common.vocab_pad_to_multiple_of) or 128 - # beam-search params for the inference branch (proto defaults = 50/50) - self._num_beams = int(common.num_beams) - self._num_return = int(common.num_return_sequences) # backbone + tokenizer (the backbone is family-owned; see _backbone_id) hf_model_id = self._backbone_id() @@ -166,6 +171,19 @@ def _backbone_id(self) -> str: """ return self._model_config.hf_model_id + def _input_sequence_length(self) -> int: + """Truncation length (SID codes) of the user-sequence feature. + + The data reader caps every row's history at the feature's + ``sequence_length``, so it is the guaranteed upper bound used to + pre-size the activation pool (see ``Qwen2RecLM._warmup_alloc``). Returns + 0 if the feature has no length cap, which disables pre-allocation. + """ + for feature in self._features: + if feature.config.feature_name == self._input_name: + return int(getattr(feature, "sequence_length", 0) or 0) + return 0 + def _build_prompt_tokens(self, tokenizer, cfg) -> None: """Family hook: cache the tokenised prompt template as buffers. @@ -201,7 +219,10 @@ def _tokenize_sids(self, sids: torch.Tensor) -> torch.Tensor: return sids + (self._base_vocab - 1) def _sid_token_rows( - self, jt, expected_width: Optional[int] = None + self, + jt, + expected_width: Optional[int] = None, + max_codes: Optional[int] = None, ) -> List[torch.Tensor]: """Read a SID jagged feature -> per-row token-id tensors. @@ -213,6 +234,15 @@ def _sid_token_rows( ``expected_width``, when set, enforces the sample contract here at the data boundary: every row must have exactly that many codes (e.g. the answer = ``num_levels``); a deviation is an anomalous sample. + + ``max_codes``, when set, caps each row to its most-recent whole items — + the last ``floor(max_codes / num_levels) * num_levels`` codes, dropping + the oldest *head* (sequences are oldest->newest, so recent behaviour is + preserved). FG_NONE does not truncate, so this is what actually enforces + the feature's ``sequence_length`` — guaranteeing the pre-allocated pool + covers every batch. Done on host views before the H2D copy, and skipped + entirely unless some row overflows, so it's free in the common case and + shrinks downstream work (and forward ``T``) when it fires. """ values = jt.values() lengths = jt.lengths() @@ -228,6 +258,12 @@ def _sid_token_rows( f"{expected_width} codes (len(codebook)); rows {bad} have " f"{[sizes[i] for i in bad]} — anomalous sample(s)." ) + if max_codes: + keep = (max_codes // self._num_levels) * self._num_levels + if keep and any(n > keep for n in sizes): + rows = torch.split(values, sizes) # host views, no copy + values = torch.cat([r[-keep:] for r in rows]) # keep recent tail + sizes = [min(n, keep) for n in sizes] # one vectorized SID->token map over the whole batch, on the backbone device values = self._tokenize_sids(values.to(self.device).long()) return list(torch.split(values, sizes)) diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index c3dad8e9a..8e0c1e270 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -30,13 +30,15 @@ they live here. """ -from typing import Dict, List, Tuple +from typing import Any, Dict, List, Optional, Tuple import torch from torch.nn.utils.rnn import pad_sequence from tzrec.datasets.utils import Batch +from tzrec.features.feature import BaseFeature from tzrec.models.generative_rec_lm import GenerativeRecLM +from tzrec.protos.model_pb2 import ModelConfig def _encode_no_special(tokenizer, text: str) -> List[int]: @@ -67,6 +69,71 @@ class Qwen2RecLM(GenerativeRecLM): CHAT_TEMPLATE = QWEN2_TEMPLATE + def __init__( + self, + model_config: ModelConfig, + features: List[BaseFeature], + labels: List[str], + sample_weights: Optional[List[str]] = None, + **kwargs: Any, + ) -> None: + super().__init__(model_config, features, labels, sample_weights, **kwargs) + common = self._model_config.common + # generation params — read in the subclass because only this family's + # _generate (the inference branch) consumes them. + self._num_beams = int(common.num_beams) + self._num_return = int(common.num_return_sequences) + # worst-case spliced length (template-aware) and the one-shot warm-up + # latch; the base supplies max_seq_length, batch_size, and the tpl_* + # buffers (built in super().__init__ via _build_prompt_tokens). + self._max_total_len = self._compute_max_total_length() + self._pool_warmed = False + + def _compute_max_total_length(self) -> int: + """Full spliced length at the max history (0 if pre-allocation is off). + + Mirrors ``_splice_input_ids``: fixed ChatML frame + ``self._max_seq_length`` + history codes (the user-sequence feature's truncation length, supplied by + the base) + the ``num_levels``-code answer (the eos sits inside the + frame). This is the ``T`` the activation pool is pre-sized to. + """ + if self._max_seq_length <= 0: + return 0 + frame = ( + self.tpl_system.numel() + + self.tpl_user_prefix.numel() + + self.tpl_user_suffix.numel() + + self.tpl_asst_prefix.numel() + + self.tpl_asst_suffix.numel() + + self.tpl_eos.numel() + ) + return int(frame + self._max_seq_length + self._num_levels) + + def _warmup_alloc(self) -> None: + """One-shot: build the CUDA activation pool at the worst case (B, T_max). + + Runs a dummy forward+backward at the configured maximum length so the + caching allocator reserves its largest segments up front; every real + (shorter) batch is then served from that pool, so it never grows + mid-run — which is what stranded segments unevenly across ranks. Fires + from the first training step (earliest point the backbone is on-GPU); + the throwaway gradients are zeroed before the real step runs. + """ + batch_size = self._batch_size or 1 + device = self.device + tok = self._base_vocab # C0 — a valid extended-vocab id + u_rows = [ + torch.full((self._max_seq_length,), tok, dtype=torch.long, device=device) + for _ in range(batch_size) + ] + l_rows = [ + torch.full((self._num_levels,), tok, dtype=torch.long, device=device) + for _ in range(batch_size) + ] + input_ids, labels, attention_mask = self._splice_input_ids(u_rows, l_rows) + self._forward_loss(input_ids, labels, attention_mask)["loss"].backward() + self.lm.zero_grad(set_to_none=True) + def _build_prompt_tokens(self, tokenizer, cfg) -> None: """Tokenise the family chat template once; cache as buffers. @@ -185,9 +252,18 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: """Branch 1: teacher-forced forward -> suffix-slice -> CE loss.""" + # one-shot: pre-size the activation pool at the worst-case length on the + # first on-GPU training step (see _warmup_alloc); no-op once warmed or + # when the input feature has no truncation length (pre-allocation off). + if not self._pool_warmed and self._max_total_len > 0 and self.is_train: + self._warmup_alloc() + self._pool_warmed = True # SID indices -> token ids once, at the data boundary (see # _sid_token_rows); the splice then just assembles the prompt. - u_rows = self._sid_token_rows(batch.sequence_dense_features[self._input_name]) + u_rows = self._sid_token_rows( + batch.sequence_dense_features[self._input_name], + max_codes=self._max_seq_length, # cap to most-recent items (drop oldest) + ) l_rows = self._sid_token_rows( batch.sequence_dense_features[self._label_name], expected_width=self._num_levels, # answer = one item = num_levels codes @@ -206,9 +282,20 @@ def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: ) self._first_predict = False - outputs = self.lm.model( - input_ids=input_ids, attention_mask=attention_mask - ) + return self._forward_loss(input_ids, labels, attention_mask) + + def _forward_loss( + self, + input_ids: torch.Tensor, + labels: torch.Tensor, + attention_mask: torch.Tensor, + ) -> Dict[str, torch.Tensor]: + """Teacher-forced forward over spliced ids -> suffix-slice -> CE loss. + + Shared by ``_predict_train`` and the ``_warmup_alloc`` dummy step so the + pre-allocation reproduces the exact training allocation pattern. + """ + outputs = self.lm.model(input_ids=input_ids, attention_mask=attention_mask) hidden = outputs.last_hidden_state # (B, T, D) # Suffix slice in BOTH train and eval. algr only slices when @@ -244,7 +331,10 @@ def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: tokens per beam, and maps them back to raw SID indices. Returns ``generated_sids`` of shape ``(B, num_return, num_levels)``. """ - u_rows = self._sid_token_rows(batch.sequence_dense_features[self._input_name]) + u_rows = self._sid_token_rows( + batch.sequence_dense_features[self._input_name], + max_codes=self._max_seq_length, # cap to most-recent items (drop oldest) + ) input_ids, attention_mask = self._splice_prompt_ids(u_rows) out = self.lm.generate( input_ids=input_ids, diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/qwen2_rec_lm_test.py index f2c9c90f3..6f70f76f4 100644 --- a/tzrec/models/qwen2_rec_lm_test.py +++ b/tzrec/models/qwen2_rec_lm_test.py @@ -30,6 +30,7 @@ def _stub(num_levels=3, base_vocab=100, pad_id=9, device="cpu"): m._num_levels = num_levels m._base_vocab = base_vocab m._pad_token_id = pad_id + m._max_seq_length = 0 # no recency clip by default in unit stubs m.lm = types.SimpleNamespace(device=torch.device(device)) for name, vals in { "tpl_system": [10, 11], "tpl_user_prefix": [12], "tpl_user_suffix": [13], @@ -146,6 +147,95 @@ def test_build_prompt_tokens_registers_buffers(self) -> None: self.assertEqual(buf.dtype, torch.int64) self.assertEqual(m.tpl_eos.tolist(), [99]) # eos cached for supervision + def test_input_sequence_length_from_feature(self) -> None: + m = object.__new__(Qwen2RecLM) + m._input_name = "user_sequence" + f_user = types.SimpleNamespace( + config=types.SimpleNamespace(feature_name="user_sequence"), + sequence_length=300, + ) + f_label = types.SimpleNamespace( + config=types.SimpleNamespace(feature_name="label"), sequence_length=32 + ) + m._features = [f_label, f_user] + self.assertEqual(m._input_sequence_length(), 300) # truncation length + m._features = [f_label] # user-sequence feature absent + self.assertEqual(m._input_sequence_length(), 0) + f_user.sequence_length = None # no length cap -> pre-allocation disabled + m._features = [f_user] + self.assertEqual(m._input_sequence_length(), 0) + + def test_sid_token_rows_recency_clip(self) -> None: + m = _stub(num_levels=3, base_vocab=100) # token = sid + base - 1 = sid + 99 + + def _jt(n): # one row of n codes: values 1..n + return types.SimpleNamespace( + values=lambda: torch.arange(1, n + 1, dtype=torch.float), + lengths=lambda: torch.tensor([n]), + ) + + # 15 codes (5 items), cap 9 -> keep last 9 (items 3-5 = codes 7..15) + rows = m._sid_token_rows(_jt(15), max_codes=9) + self.assertEqual(rows[0].tolist(), [c + 99 for c in range(7, 16)]) + # item-aligned: cap 10 still keeps 9 (3 whole items), never cuts mid-item + rows = m._sid_token_rows(_jt(15), max_codes=10) + self.assertEqual(rows[0].tolist(), [c + 99 for c in range(7, 16)]) + # within cap -> untouched + rows = m._sid_token_rows(_jt(6), max_codes=9) + self.assertEqual(rows[0].tolist(), [c + 99 for c in range(1, 7)]) + # disabled (0/None) -> no clip + rows = m._sid_token_rows(_jt(15), max_codes=0) + self.assertEqual(rows[0].tolist(), [c + 99 for c in range(1, 16)]) + + def test_compute_max_total_length(self) -> None: + m = _stub(num_levels=3) + # frame = |system|2 + |user_prefix|1 + |user_suffix|1 + |asst_prefix|1 + # + |asst_suffix|1 + |eos|1 = 7; + max_history + answer(num_levels) + m._max_seq_length = 300 + self.assertEqual(m._compute_max_total_length(), 7 + 300 + 3) + m._max_seq_length = 0 # pre-allocation disabled + self.assertEqual(m._compute_max_total_length(), 0) + + def test_warmup_fires_once_in_training(self) -> None: + m = _stub() + m._is_inference = False # not inference + nn.Module.training=True -> is_train + m._smoke_log_once = False + m._input_name, m._label_name = "user_sequence", "label" + m._max_total_len = 50 + m._pool_warmed = False + calls = [] + m._warmup_alloc = lambda: calls.append(1) + m._sid_token_rows = lambda jt, expected_width=None, max_codes=None: [ + torch.tensor([100, 101, 102]) + ] + m._forward_loss = lambda i, lbl, a: {"loss": torch.tensor(0.0)} + batch = types.SimpleNamespace( + sequence_dense_features={"user_sequence": None, "label": None} + ) + m._predict_train(batch) + m._predict_train(batch) + self.assertEqual(len(calls), 1) # one-shot, latched by _pool_warmed + self.assertTrue(m._pool_warmed) + + def test_warmup_skipped_when_disabled(self) -> None: + m = _stub() + m._is_inference = False + m._smoke_log_once = False + m._input_name, m._label_name = "user_sequence", "label" + m._max_total_len = 0 # max_seq_length unset -> pre-allocation off + m._pool_warmed = False + calls = [] + m._warmup_alloc = lambda: calls.append(1) + m._sid_token_rows = lambda jt, expected_width=None, max_codes=None: [ + torch.tensor([100, 101, 102]) + ] + m._forward_loss = lambda i, lbl, a: {"loss": torch.tensor(0.0)} + batch = types.SimpleNamespace( + sequence_dense_features={"user_sequence": None, "label": None} + ) + m._predict_train(batch) + self.assertEqual(calls, []) # never warms up when disabled + if __name__ == "__main__": unittest.main() From b2f1890a803f697d1c9edf37705cdacb5c942af5 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 10 Jun 2026 06:56:25 +0000 Subject: [PATCH 07/99] [config] generative-rec LM example: set user_sequence sequence_length=300 Under FG_NONE the reader does not truncate, so the model enforces this length via the recency-preserving clip in _sid_token_rows and uses it to pre-size the activation pool. 300 = AL-GR-Tiny's realistic max history (100 items x 3 codes); the prior loose 1056 (algr max_length) would oversize the warm-up. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../generative_rec_lm_s1pretrained.config | 181 ++++++++++++++++++ 1 file changed, 181 insertions(+) create mode 100644 examples/generative_rec_lm_s1pretrained.config diff --git a/examples/generative_rec_lm_s1pretrained.config b/examples/generative_rec_lm_s1pretrained.config new file mode 100644 index 000000000..807cc4bc1 --- /dev/null +++ b/examples/generative_rec_lm_s1pretrained.config @@ -0,0 +1,181 @@ +# TorchEasyRec pipeline.config — `GenerativeRecLM` Qwen2.5-0.5B-Instruct +# translated from algr's +# /home/admin/workspace/al_sid/algr/config/qwen2.5_05b_3layer_s1_pretrained.json +# applied to the TINY parquet dataset +# /home/admin/workspace/aop_lab/data/AL-GR-Tiny/train_data_genreclm/*.parquet +# (produced by `examples/convert_s1tiny_to_parquet.py`; rows have +# `user_sequence: list` + `label: list`). +# +# algr → TER translation (knob by knob): +# per_device_train_batch_size : 80 → data_config.batch_size : 80 +# gradient_accumulation_steps : 4 → train_config.gradient_accumulation_steps : 4 *MUST MATCH* +# learning_rate : 5e-5 → dense_optimizer.adam_optimizer.lr : 5e-5 +# adam_beta1/2/epsilon : 0.9/0.999/1e-8 → adam_optimizer.{beta1,beta2,eps} +# lr_scheduler_type : linear → linear_decay_learning_rate (added to TER +# for this alignment; decays 5e-5 → 0). +# +# STEP-SEMANTICS NOTE (important): HF Trainer's `max_steps` counts +# OPTIMIZER steps, while TER's `num_steps` counts FORWARD steps and +# TZRecOptimizer skips `.step()` between accumulation boundaries +# (tzrec/optim/optimizer.py). TER's LR schedulers likewise tick once per +# forward step. With gradient_accumulation_steps=4, algr's max_steps +# 125000 therefore translates to: +# max_steps : 125000 → train_config.num_steps : 500000 (= 125000 × 4) +# save_steps : 10000 → train_config.save_checkpoints_steps : 40000 +# logging_steps : 1000 → train_config.log_step_count_steps : 4000 +# linear decay total : 125000 optimizer steps → total_size : 500000 forward steps +# +# Remaining benign divergence: HF Trainer scales the accumulated loss by +# 1/4 before backward (mean over micro-batches); TER sums the 4 +# micro-batch gradients. A constant gradient scale cancels in Adam's +# m̂/(√v̂+ε) update (up to ε), so trajectories match; logged ce_loss is +# per-micro-batch in both systems and stays directly comparable. +# bf16 : true → torch_dtype="auto" (the wrapper reads bf16 +# from the safetensors header). +# max_length / max_source_length / max_target_length +# : 1056/1024/32 → sequence_raw_feature.sequence_length : 1056 +# and capped by the data converter's row length +# (`--max-source-len` for the parquet driver). +# load_checkpoint_from : /home/admin/workspace/Qwen2.5-0.5B-Instruct +# → generative_rec_lm.hf_model_id : +# dataloader_num_workers : 4 → data_config.num_workers : 4 +# +# Launch (single H20:0, smoke): +# cd /home/admin/workspace/TorchEasyRec_qwen_smoke/TorchEasyRec +# export PYTHONPATH=$(pwd):$PYTHONPATH TRANSFORMERS_OFFLINE=1 HF_HUB_OFFLINE=1 +# torchrun --nnodes=1 --nproc-per-node=1 --master_port=32777 \ +# -m tzrec.train_eval \ +# --pipeline_config_path examples/generative_rec_lm_s1pretrained.config + +train_input_path: "/home/admin/workspace/aop_lab/data/AL-GR-Tiny/train_data_genreclm/*.parquet" +# held-out test split (s1_tiny_test.csv converted with the same script); +# evaluated (mean ce_loss) at every checkpoint save. +eval_input_path: "/home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data_genreclm/*.parquet" +model_dir: "experiments/generative_rec_lm_s1pretrained" + +train_config { + # GenerativeRecLM has no sparse params (the HF backbone owns its own + # `embed_tokens` and TER's EmbeddingGroup is unused). Define a no-op + # sparse_optimizer anyway since TER's TrainConfig requires it. + sparse_optimizer { + adagrad_optimizer { lr: 0.0 } + constant_learning_rate {} + } + dense_optimizer { + adam_optimizer { + lr: 5e-5 # algr `learning_rate: 5e-05` + beta1: 0.9 # algr `adam_beta1` + beta2: 0.999 # algr `adam_beta2` + eps: 1e-8 # algr `adam_epsilon` + } + # algr `lr_scheduler_type: linear` — decay 5e-5 → 0 over the whole + # run. TER schedulers tick per FORWARD step, so total_size is in + # forward steps (= 125000 optimizer steps × grad_accum 4). + linear_decay_learning_rate { + total_size: 500000 + } + } + # HF Trainer applies its DEFAULT `max_grad_norm: 1.0` global-norm + # clipping (algr's training_args don't override it; algr logs + # pre-clip grad_norm ~600 at step 1, so clipping is engaged). + # TER SUMS the grad_accum micro-grads where HF averages them, so the + # equivalent threshold is 1.0 × 4 = 4.0; the constant 4× gradient + # scale then cancels inside Adam's m̂/(√v̂+ε) update. + grad_clipping { + clipping_type: "norm" + max_gradient: 4.0 + norm_type: 2.0 + enable_global_grad_clip: true + } + # algr `max_steps: 125000` OPTIMIZER steps × grad_accum 4 = 500000 + # forward steps (see step-semantics note above). + num_steps: 500000 + # algr `gradient_accumulation_steps: 4` — MUST MATCH EXACTLY. + gradient_accumulation_steps: 4 + # algr `save_steps: 10000` optimizer steps × 4. + save_checkpoints_steps: 40000 + # algr `logging_steps: 1000` optimizer steps × 4. + log_step_count_steps: 4000 + # bf16 — TER reads from the safetensors header via torch_dtype="auto" + # at `from_pretrained` time. No flag here. +} + +eval_config { +} + +data_config { + # algr `per_device_train_batch_size: 80`. + batch_size: 80 + dataset_type: ParquetDataset + fg_mode: FG_NONE + # algr `dataloader_num_workers: 4`. + num_workers: 4 +} + +# The two sequence features map 1:1 to the parquet columns produced by +# `convert_s1tiny_to_parquet.py`. +feature_configs { + sequence_raw_feature { + feature_name: "user_sequence" + expression: "user:user_sequence" + # Under FG_NONE this length is NOT auto-truncated by the reader, so the + # model enforces it via a recency-preserving clip in `_sid_token_rows` + # (keep newest items, drop oldest) AND uses it to pre-size the CUDA + # activation pool (`_input_sequence_length`/`_warmup_alloc`). Set it to + # the data's realistic max history in codes (AL-GR-Tiny = 100 items × + # 3 = 300); a loose value (e.g. algr's 1056) would oversize the warm-up. + sequence_length: 300 + value_dim: 1 + } +} +feature_configs { + sequence_raw_feature { + feature_name: "label" + expression: "user:label" + # algr `max_target_length: 32`. Each item is 3 SIDs, so 32 covers + # up to 10 items; in practice labels are always 3 SIDs in this set. + sequence_length: 32 + value_dim: 1 + } +} + +model_config { + feature_groups { + group_name: "sids" + feature_names: "user_sequence" + feature_names: "label" + group_type: SEQUENCE + } + qwen2_rec_lm { + # Qwen2 backbone (owned by this family message, not `common`). + # algr `load_checkpoint_from: /home/admin/workspace/Qwen2.5-0.5B` + # (the BASE model — qwen2.5_05b_3layer_s1_pretrained.json does not use + # the Instruct variant). Omit to inherit the default "Qwen/Qwen2.5-0.5B". + hf_model_id: "/home/admin/workspace/Qwen2.5-0.5B" + + # architecture-agnostic config shared by all generative-rec families + common { + # SID codebook: ONE entry per RQ level. AL-GR is 3 levels x 8192 + # (verified from item_info: every codebook_lv* ∈ [0, 8191]), so + # len(codebook)=3 = SID codes per answer and sum=24 576 = atoms + # appended to the vocab (an exact fit — max SID atom is 24 575). + codebook: 8192 + codebook: 8192 + codebook: 8192 + vocab_pad_to_multiple_of: 128 + + # Sample feature names — match parquet columns. + user_sequence_feature_name: "user_sequence" + label_feature_name: "label" + ignore_index: -100 + } + + # algr's row.system / `default_instruction` — the CN recommender prompt. + # Matching this string bit-for-bit reproduces algr's input. + system_instruction: "你是一个推荐系统，根据用户的历史行为，预测用户在电商场景的下一步行为。我会给你一串连续行为的语义编码，按照用户点击的时间顺序排列，每个行为用三个词表示。" + + # algr's CN user-prompt wrappers around the SID list. + user_prefix_text: "当前用户的历史行为如下：" + user_suffix_text: "，请预测用户在电商推荐场景后续行为的语义编码" + } +} From e51651d97d00d8b81b49e101afd87d48dc8a498b Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Fri, 12 Jun 2026 02:28:27 +0000 Subject: [PATCH 08/99] [feat] generative-rec LM: save HF-format checkpoints alongside DCP + fix export tool schema Intermediate training checkpoints were only written in TER's DCP format, which isn't directly from_pretrained-loadable and reduced confidence when validating results. Now save an HF copy at each periodic checkpoint. - GenerativeRecLM.export_hf(dir): save self.lm + the rebuilt extended tokenizer (base + C0..C{sum-1}) straight from the live model. - main._train_and_evaluate: after each periodic ckpt_manager.save, rank-0 calls _model.export_hf(model_dir/hf_ckpt-{step}) when available (duck-typed; other models unaffected). - export_genreclm_to_hf: fix three new-schema bugs (class_name -> which_msg resolution, abstract GenerativeRecLM -> resolved family class + register Qwen2RecLM, grl_cfg.codebook -> grl_cfg.common.codebook). The codebook bug had exported weights without the tokenizer, breaking predict. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/main.py | 7 +++++++ tzrec/models/generative_rec_lm.py | 18 ++++++++++++++++++ tzrec/tools/export_genreclm_to_hf.py | 18 ++++++++++++------ 3 files changed, 37 insertions(+), 6 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index b4e4fb68a..e7b19bf03 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -466,6 +466,13 @@ def _train_and_evaluate( if i_step % save_checkpoints_steps == 0: last_ckpt_step = i_step ckpt_manager.save(i_step, model, optimizer, dataloader_state) + # also save an HF-loadable copy (GenerativeRecLM family only) + # so intermediate checkpoints are directly inspectable + # without the DCP->HF export round-trip. + if is_rank_zero and hasattr(_model, "export_hf"): + _model.export_hf( + os.path.join(model_dir, f"hf_ckpt-{i_step}") + ) if eval_dataloader is not None: _evaluate( model, diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 1a357d0df..5bb868ed7 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -184,6 +184,24 @@ def _input_sequence_length(self) -> int: return int(getattr(feature, "sequence_length", 0) or 0) return 0 + def export_hf(self, export_dir: str) -> None: + """Save the HF backbone + extended tokenizer as a HF-loadable dir. + + Mirrors ``tzrec.tools.export_genreclm_to_hf`` but straight from the + live in-memory model, so intermediate training checkpoints can be saved + in HF (``from_pretrained``-loadable) format alongside the DCP + checkpoints (see the save hook in ``main._train_and_evaluate``). Call on + rank 0 only; the dense backbone is replicated, so rank 0 holds the full + weights. The extended tokenizer is rebuilt (base + C0..C{sum-1}) so SID + atoms decode with the same ids the model trained on. + """ + os.makedirs(export_dir, exist_ok=True) + self.lm.save_pretrained(export_dir) + tokenizer = AutoTokenizer.from_pretrained(self._backbone_id(), use_fast=True) + sid_atoms = sum(int(c) for c in self._model_config.common.codebook) + tokenizer.add_tokens([f"C{i}" for i in range(sid_atoms)]) + tokenizer.save_pretrained(export_dir) + def _build_prompt_tokens(self, tokenizer, cfg) -> None: """Family hook: cache the tokenised prompt template as buffers. diff --git a/tzrec/tools/export_genreclm_to_hf.py b/tzrec/tools/export_genreclm_to_hf.py index e73619492..19bf8daf0 100644 --- a/tzrec/tools/export_genreclm_to_hf.py +++ b/tzrec/tools/export_genreclm_to_hf.py @@ -4,7 +4,7 @@ # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 -"""Export a `GenerativeRecLM` TER DCP checkpoint to a HF-loadable directory. +r"""Export a `GenerativeRecLM` TER DCP checkpoint to a HF-loadable directory. Design §6.4 (FINAL_DESIGN_GENERATIVE_REC_LM.md), simplified: instead of hand-writing safetensors shards, we rebuild the model from the pipeline @@ -34,14 +34,17 @@ import torch from google.protobuf import text_format from torch.distributed.checkpoint import FileSystemReader, load +from transformers import AutoTokenizer from tzrec.models.generative_rec_lm import GenerativeRecLM # noqa: F401 from tzrec.models.model import BaseModel +from tzrec.models.qwen2_rec_lm import Qwen2RecLM # noqa: F401 (register family) from tzrec.protos.pipeline_pb2 import EasyRecConfig -from transformers import AutoTokenizer +from tzrec.utils import config_util def main() -> int: + """Export a GenerativeRecLM DCP checkpoint to a HF-loadable dir.""" ap = argparse.ArgumentParser() ap.add_argument("--pipeline_config_path", required=True) ap.add_argument("--checkpoint_path", required=True) @@ -53,12 +56,15 @@ def main() -> int: text_format.Merge(f.read(), pipeline_config) model_config = pipeline_config.model_config grl_cfg = getattr(model_config, model_config.WhichOneof("model")) + # Resolve the family class from the oneof message-type name (e.g. + # "Qwen2RecLM"), the same way main._create_model does — no class_name field. + model_cls_name = config_util.which_msg(model_config, "model") # Rebuild the model exactly as training did (from_pretrained backbone + # SID vocab extension), CPU-resident. - print(f"[export] building {grl_cfg.class_name} from {grl_cfg.hf_model_id}") + print(f"[export] building {model_cls_name} from {grl_cfg.hf_model_id}") # pyre-ignore [16] - model_cls = BaseModel.create_class("GenerativeRecLM") + model_cls = BaseModel.create_class(model_cls_name) model = model_cls(model_config, features=[], labels=[]) model.eval() @@ -87,7 +93,7 @@ def main() -> int: # no [SEP]) so downstream generation maps SID atoms identically. tokenizer = AutoTokenizer.from_pretrained(grl_cfg.hf_model_id, use_fast=True) base = len(tokenizer) - tokenizer.add_tokens([f"C{i}" for i in range(sum(grl_cfg.codebook))]) + tokenizer.add_tokens([f"C{i}" for i in range(sum(grl_cfg.common.codebook))]) assert tokenizer.convert_tokens_to_ids("C0") == base tokenizer.save_pretrained(args.export_dir) with open(os.path.join(args.export_dir, "TER_EXPORT_INFO.txt"), "w") as f: @@ -95,7 +101,7 @@ def main() -> int: f"source_checkpoint={args.checkpoint_path}\n" f"pipeline_config={args.pipeline_config_path}\n" f"sid_base_token_id={base}\n" - f"codebook={list(grl_cfg.codebook)}\n" + f"codebook={list(grl_cfg.common.codebook)}\n" "note=C atoms appended directly after base vocab (NO [SEP]); " "token_id = base + (sid - 1) for 1-indexed SIDs / base + k for C{k}.\n" ) From 59b3c5b0406813ae158fff1f99bd4945b71c1a00 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Fri, 12 Jun 2026 02:32:17 +0000 Subject: [PATCH 09/99] [refactor] generative-rec LM export: reuse model.export_hf in the offline tool The standalone export tool duplicated the backbone+tokenizer save that GenerativeRecLM.export_hf now owns. Make the tool do only the DCP overlay then call model.export_hf, so offline and in-training HF exports are one code path (byte-identical) and there is a single source of truth. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/tools/export_genreclm_to_hf.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/tzrec/tools/export_genreclm_to_hf.py b/tzrec/tools/export_genreclm_to_hf.py index 19bf8daf0..ffe285f48 100644 --- a/tzrec/tools/export_genreclm_to_hf.py +++ b/tzrec/tools/export_genreclm_to_hf.py @@ -34,7 +34,6 @@ import torch from google.protobuf import text_format from torch.distributed.checkpoint import FileSystemReader, load -from transformers import AutoTokenizer from tzrec.models.generative_rec_lm import GenerativeRecLM # noqa: F401 from tzrec.models.model import BaseModel @@ -85,17 +84,12 @@ def main() -> int: f"dtype={emb.dtype} mean_abs={emb.abs().mean().item():.6f}" ) - os.makedirs(args.export_dir, exist_ok=True) - print(f"[export] save_pretrained -> {args.export_dir}") - model.lm.save_pretrained(args.export_dir) - - # Save the EXTENDED tokenizer (TER layout: C0 at len(base tokenizer), - # no [SEP]) so downstream generation maps SID atoms identically. - tokenizer = AutoTokenizer.from_pretrained(grl_cfg.hf_model_id, use_fast=True) - base = len(tokenizer) - tokenizer.add_tokens([f"C{i}" for i in range(sum(grl_cfg.common.codebook))]) - assert tokenizer.convert_tokens_to_ids("C0") == base - tokenizer.save_pretrained(args.export_dir) + # Reuse the model's HF-save (backbone + extended tokenizer) — the SAME code + # path as the in-training checkpoint hook, so offline and online exports are + # byte-identical. + print(f"[export] save_pretrained + tokenizer -> {args.export_dir}") + model.export_hf(args.export_dir) + base = model._base_vocab with open(os.path.join(args.export_dir, "TER_EXPORT_INFO.txt"), "w") as f: f.write( f"source_checkpoint={args.checkpoint_path}\n" From 679aa177a99addc2033080435cd5549dd946ec9c Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Fri, 12 Jun 2026 02:40:16 +0000 Subject: [PATCH 10/99] [refactor] generative-rec LM: fold HF export into main.export, delete standalone tool The standalone tzrec/tools/export_genreclm_to_hf.py re-implemented the model build + checkpoint restore that TER's export() already does. Delete it and add a GenerativeRecLM branch directly in main.export: after the existing checkpoint resolution, overlay the DCP shards and call model.export_hf (the same single save path used by the in-training checkpoint hook). No separate tool, no duplicated save; `python -m tzrec.export` now produces the HF dir for GenerativeRecLM models. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/main.py | 20 +++++ tzrec/models/generative_rec_lm.py | 15 ++-- tzrec/tools/export_genreclm_to_hf.py | 107 --------------------------- 3 files changed, 28 insertions(+), 114 deletions(-) delete mode 100644 tzrec/tools/export_genreclm_to_hf.py diff --git a/tzrec/main.py b/tzrec/main.py index e7b19bf03..e5177c326 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -24,6 +24,8 @@ from torch import nn, optim from torch.amp import GradScaler from torch.distributed._shard.sharded_tensor import ShardedTensor +from torch.distributed.checkpoint import FileSystemReader +from torch.distributed.checkpoint import load as dcp_load from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter from torchrec.optim.apply_optimizer_in_backward import ( @@ -961,6 +963,24 @@ def export( else: checkpoint_path, _ = ckpt_manager.latest_checkpoint() + # GenerativeRecLM family: export to a HF (from_pretrained-loadable) dir + # instead of TorchScript. Overlay the DCP checkpoint (training FQNs are + # ``model.lm.``) and reuse the model's ``export_hf`` save path — the + # same code path as the in-training checkpoint hook, so there is no separate + # export tool and no duplicated save. + grl_model = model.model + if hasattr(grl_model, "export_hf"): + ckpt_model_dir = os.path.join(checkpoint_path, "model") + lm_sd = grl_model.lm.state_dict() + prefixed = {f"model.lm.{k}": v for k, v in lm_sd.items()} + dcp_load(prefixed, storage_reader=FileSystemReader(ckpt_model_dir)) + grl_model.lm.load_state_dict( + {k[len("model.lm.") :]: v for k, v in prefixed.items()} + ) + if is_rank_zero: + grl_model.export_hf(export_dir) + return + if isinstance(model.model, MatchModel): for name, module in model.model.named_children(): if isinstance(module, MatchTower) or isinstance(module, MatchTowerWoEG): diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 5bb868ed7..16e87b83e 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -187,13 +187,14 @@ def _input_sequence_length(self) -> int: def export_hf(self, export_dir: str) -> None: """Save the HF backbone + extended tokenizer as a HF-loadable dir. - Mirrors ``tzrec.tools.export_genreclm_to_hf`` but straight from the - live in-memory model, so intermediate training checkpoints can be saved - in HF (``from_pretrained``-loadable) format alongside the DCP - checkpoints (see the save hook in ``main._train_and_evaluate``). Call on - rank 0 only; the dense backbone is replicated, so rank 0 holds the full - weights. The extended tokenizer is rebuilt (base + C0..C{sum-1}) so SID - atoms decode with the same ids the model trained on. + Single source of truth for HF (``from_pretrained``-loadable) export, + reused by BOTH ``main.export`` (the standard TER export path, after it + overlays the DCP checkpoint) and the in-training checkpoint hook in + ``main._train_and_evaluate`` — so offline and online exports are + identical and there is no separate export tool. Call on rank 0 only; + the dense backbone is replicated, so rank 0 holds the full weights. The + extended tokenizer is rebuilt (base + C0..C{sum-1}) so SID atoms decode + with the same ids the model trained on. """ os.makedirs(export_dir, exist_ok=True) self.lm.save_pretrained(export_dir) diff --git a/tzrec/tools/export_genreclm_to_hf.py b/tzrec/tools/export_genreclm_to_hf.py deleted file mode 100644 index ffe285f48..000000000 --- a/tzrec/tools/export_genreclm_to_hf.py +++ /dev/null @@ -1,107 +0,0 @@ -# Copyright (c) 2026, Alibaba Group; -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# http://www.apache.org/licenses/LICENSE-2.0 - -r"""Export a `GenerativeRecLM` TER DCP checkpoint to a HF-loadable directory. - -Design §6.4 (FINAL_DESIGN_GENERATIVE_REC_LM.md), simplified: instead of -hand-writing safetensors shards, we rebuild the model from the pipeline -config (which re-applies the SID vocab extension so shapes match), overlay -the DCP shards onto it, and let HF's ``save_pretrained`` deal with weight -tying, sharding and config serialisation. The extended tokenizer (C0.. at -``len(tokenizer)`` — TER layout, no [SEP]) is saved alongside so generation -consumers decode SID atoms with the SAME ids the model was trained on. - -DCP shard FQNs are ``model.lm.`` (TrainWrapper prefix ``model.`` + -wrapper attr ``lm.``); we restore through a TrainWrapper-shaped state dict -so no manual FQN surgery is needed. - -Usage (CPU-only; safe to run next to a live training):: - - PYTHONPATH=. python -m tzrec.tools.export_genreclm_to_hf \\ - --pipeline_config_path experiments//pipeline.config \\ - --checkpoint_path experiments//model.ckpt-40000 \\ - --export_dir experiments//export_hf_40000 -""" - -from __future__ import annotations - -import argparse -import os - -import torch -from google.protobuf import text_format -from torch.distributed.checkpoint import FileSystemReader, load - -from tzrec.models.generative_rec_lm import GenerativeRecLM # noqa: F401 -from tzrec.models.model import BaseModel -from tzrec.models.qwen2_rec_lm import Qwen2RecLM # noqa: F401 (register family) -from tzrec.protos.pipeline_pb2 import EasyRecConfig -from tzrec.utils import config_util - - -def main() -> int: - """Export a GenerativeRecLM DCP checkpoint to a HF-loadable dir.""" - ap = argparse.ArgumentParser() - ap.add_argument("--pipeline_config_path", required=True) - ap.add_argument("--checkpoint_path", required=True) - ap.add_argument("--export_dir", required=True) - args = ap.parse_args() - - pipeline_config = EasyRecConfig() - with open(args.pipeline_config_path) as f: - text_format.Merge(f.read(), pipeline_config) - model_config = pipeline_config.model_config - grl_cfg = getattr(model_config, model_config.WhichOneof("model")) - # Resolve the family class from the oneof message-type name (e.g. - # "Qwen2RecLM"), the same way main._create_model does — no class_name field. - model_cls_name = config_util.which_msg(model_config, "model") - - # Rebuild the model exactly as training did (from_pretrained backbone + - # SID vocab extension), CPU-resident. - print(f"[export] building {model_cls_name} from {grl_cfg.hf_model_id}") - # pyre-ignore [16] - model_cls = BaseModel.create_class(model_cls_name) - model = model_cls(model_config, features=[], labels=[]) - model.eval() - - # Overlay DCP shards. Shard keys are "model.lm." — present the - # state dict under the same prefix. - ckpt_model_dir = os.path.join(args.checkpoint_path, "model") - print(f"[export] overlaying DCP shards from {ckpt_model_dir}") - lm_sd = model.lm.state_dict() - prefixed = {f"model.lm.{k}": v for k, v in lm_sd.items()} - load(prefixed, storage_reader=FileSystemReader(ckpt_model_dir)) - model.lm.load_state_dict({k[len("model.lm."):]: v for k, v in prefixed.items()}) - - # Sanity: SID rows must differ from fresh init → confirm overlay landed. - with torch.no_grad(): - emb = model.lm.get_input_embeddings().weight - print( - f"[export] embed_tokens: shape={tuple(emb.shape)} " - f"dtype={emb.dtype} mean_abs={emb.abs().mean().item():.6f}" - ) - - # Reuse the model's HF-save (backbone + extended tokenizer) — the SAME code - # path as the in-training checkpoint hook, so offline and online exports are - # byte-identical. - print(f"[export] save_pretrained + tokenizer -> {args.export_dir}") - model.export_hf(args.export_dir) - base = model._base_vocab - with open(os.path.join(args.export_dir, "TER_EXPORT_INFO.txt"), "w") as f: - f.write( - f"source_checkpoint={args.checkpoint_path}\n" - f"pipeline_config={args.pipeline_config_path}\n" - f"sid_base_token_id={base}\n" - f"codebook={list(grl_cfg.common.codebook)}\n" - "note=C atoms appended directly after base vocab (NO [SEP]); " - "token_id = base + (sid - 1) for 1-indexed SIDs / base + k for C{k}.\n" - ) - print(f"[export] done; sid_base_token_id={base}") - return 0 - - -if __name__ == "__main__": - main() From 8fff84679c85996ffa4d3196f6bf4e6a532a1f3e Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Tue, 16 Jun 2026 09:56:57 +0000 Subject: [PATCH 11/99] [refactor] generative-rec LM: pipeline-driven HF import/export + first-step pad pre-sizing Import: __init__ builds the empty extended arch (AutoConfig + from_config, no weight download); new init_from_pretrained() is the sole from_pretrained, invoked by the pipeline on cold start via a no-op BaseModel hook (no hasattr duck-typing). Restore/eval/export load weights from DCP, skipping the ~1GB download. Export: training writes DCP only; each model.ckpt-N/ co-locates HF config+tokenizer (no weights), gated by export_config.export_format == HF and owned by CheckpointManager. tzrec.export converts via a standalone dcp_to_hf (recorded backbone-prefix recorded as data + suffix-match self-heal + strict 1:1 validation, never a silent partial load). Adds ExportFormat to export.proto; TORCHSCRIPT path untouched. Training: replace _warmup_alloc -- a separate unscaled forward+backward that corrupted the first optimizer step (the lr5e-5 HR 2.6-2.9->1.14-flat regression) -- with first-step left-padding to the worst-case length in _predict_train / _splice_input_ids / _left_pad. The pad rides the real step (positions masked + labelled -100, so loss/grad are identical) while pre-sizing the activation pool to (B, T_max), keeping per-rank reservations uniform. Tests: rewrite the warmup unit tests to assert the first-step padding contract. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/main.py | 101 ++++++++++------- tzrec/models/generative_rec_lm.py | 111 ++++++++++++------ tzrec/models/model.py | 10 ++ tzrec/models/qwen2_rec_lm.py | 77 +++++++------ tzrec/models/qwen2_rec_lm_test.py | 37 ++++-- tzrec/protos/export.proto | 11 ++ tzrec/utils/checkpoint_util.py | 20 +++- tzrec/utils/export_util.py | 182 ++++++++++++++++++++++++++++++ 8 files changed, 419 insertions(+), 130 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index e5177c326..8ff9dc4f0 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -24,8 +24,6 @@ from torch import nn, optim from torch.amp import GradScaler from torch.distributed._shard.sharded_tensor import ShardedTensor -from torch.distributed.checkpoint import FileSystemReader -from torch.distributed.checkpoint import load as dcp_load from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter from torchrec.optim.apply_optimizer_in_backward import ( @@ -73,6 +71,7 @@ from tzrec.optim import optimizer_builder from tzrec.optim.lr_scheduler import BaseLR from tzrec.optim.optimizer import TZRecOptimizer +from tzrec.protos import export_pb2 from tzrec.protos.data_pb2 import DataConfig, DatasetType from tzrec.protos.eval_pb2 import EvalConfig from tzrec.protos.feature_pb2 import FeatureConfig @@ -467,14 +466,15 @@ def _train_and_evaluate( if save_checkpoints_steps > 0 and i_step > 0: if i_step % save_checkpoints_steps == 0: last_ckpt_step = i_step - ckpt_manager.save(i_step, model, optimizer, dataloader_state) - # also save an HF-loadable copy (GenerativeRecLM family only) - # so intermediate checkpoints are directly inspectable - # without the DCP->HF export round-trip. - if is_rank_zero and hasattr(_model, "export_hf"): - _model.export_hf( - os.path.join(model_dir, f"hf_ckpt-{i_step}") - ) + ckpt_manager.save( + i_step, + model, + optimizer, + dataloader_state, + ) + # Training writes DCP (+ co-located HF config/tokenizer when + # export_format == HF); HF weights are produced on demand by + # `tzrec.export` (dcp_to_hf) from the self-contained ckpt dir. if eval_dataloader is not None: _evaluate( model, @@ -493,7 +493,12 @@ def _train_and_evaluate( if save_checkpoints_epochs > 0 and i_step > 0: if (i_epoch + 1) % save_checkpoints_epochs == 0: last_ckpt_step = i_step - ckpt_manager.save(i_step, model, optimizer, dataloader_state) + ckpt_manager.save( + i_step, + model, + optimizer, + dataloader_state, + ) if eval_dataloader is not None: _evaluate( model, @@ -528,7 +533,12 @@ def _train_and_evaluate( if train_config.is_profiling: prof.stop() if last_ckpt_step != i_step: - ckpt_manager.save(i_step, model, optimizer, dataloader_state) + ckpt_manager.save( + i_step, + model, + optimizer, + dataloader_state, + ) if eval_dataloader is not None: _evaluate( model, @@ -659,6 +669,16 @@ def train_and_evaluate( sampler_type=sampler_type, batch_size=data_config.batch_size, ) + # Cold-start gate (training-only). `_create_model` builds the EMPTY extended + # architecture (GenerativeRecLM: from_config, no weight download). On a fresh + # run (no checkpoint to resume/fine-tune — same `ckpt_path is None` signal the + # DCP-restore branch at L420-433 keys off) we load the pretrained HF weights + # ONCE here, before TrainWrapper/DMP wrapping and before any DCP restore. On + # resume/fine-tune (`ckpt_path` set) we skip it: DCP `load_state_dict` fills + # the weights. eval/export never reach this path — they always DCP-restore an + # empty model. (See design §1.) + if ckpt_path is None: + model.init_from_pretrained() # no-op unless the model has a pretrained source model = TrainWrapper( model, device=device, mixed_precision=train_config.mixed_precision ) @@ -932,6 +952,33 @@ def export( if asset_files: assets = asset_files.split(",") + ckpt_manager = checkpoint_util.CheckpointManager( + pipeline_config.model_dir, export_config=pipeline_config.export_config + ) + if not checkpoint_path: + if ( + pipeline_config.HasField("export_config") + and pipeline_config.export_config.exporter_type == "best" + ): + checkpoint_path, _ = ckpt_manager.best_checkpoint() + else: + checkpoint_path, _ = ckpt_manager.latest_checkpoint() + + # Explicit export-format branch (driven by export_config.export_format). + # HF: a standalone DCP->HF conversion — NO model build, NO DCP restore, NO + # from_pretrained. `dcp_to_hf` reads everything from the self-contained + # checkpoint dir (DCP weights + co-located config/tokenizer) and writes a + # `from_pretrained`-loadable dir. Done before _create_model so we never + # instantiate the model for HF export (design §2). + if pipeline_config.export_config.export_format == export_pb2.ExportFormat.HF: + if checkpoint_path is None: + raise ValueError("HF export: no checkpoint found to convert.") + if is_rank_zero: + from tzrec.utils.export_util import dcp_to_hf + + dcp_to_hf(checkpoint_path, export_dir) + return + data_config = pipeline_config.data_config # Build feature @@ -951,36 +998,6 @@ def export( model.set_is_inference(True) model = InferWrapper(model) - if not checkpoint_path: - ckpt_manager = checkpoint_util.CheckpointManager( - pipeline_config.model_dir, export_config=pipeline_config.export_config - ) - if ( - pipeline_config.HasField("export_config") - and pipeline_config.export_config.exporter_type == "best" - ): - checkpoint_path, _ = ckpt_manager.best_checkpoint() - else: - checkpoint_path, _ = ckpt_manager.latest_checkpoint() - - # GenerativeRecLM family: export to a HF (from_pretrained-loadable) dir - # instead of TorchScript. Overlay the DCP checkpoint (training FQNs are - # ``model.lm.``) and reuse the model's ``export_hf`` save path — the - # same code path as the in-training checkpoint hook, so there is no separate - # export tool and no duplicated save. - grl_model = model.model - if hasattr(grl_model, "export_hf"): - ckpt_model_dir = os.path.join(checkpoint_path, "model") - lm_sd = grl_model.lm.state_dict() - prefixed = {f"model.lm.{k}": v for k, v in lm_sd.items()} - dcp_load(prefixed, storage_reader=FileSystemReader(ckpt_model_dir)) - grl_model.lm.load_state_dict( - {k[len("model.lm.") :]: v for k, v in prefixed.items()} - ) - if is_rank_zero: - grl_model.export_hf(export_dir) - return - if isinstance(model.model, MatchModel): for name, module in model.model.named_children(): if isinstance(module, MatchTower) or isinstance(module, MatchTowerWoEG): diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 16e87b83e..027a94715 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -38,7 +38,7 @@ import torch import torchmetrics -from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer from tzrec.datasets.utils import Batch from tzrec.features.feature import BaseFeature @@ -82,10 +82,6 @@ def __init__( sample_weights: Optional[List[str]] = None, **kwargs: Any, ) -> None: - # per-rank batch size, threaded from data_config.batch_size by - # _create_model (absorbed by BaseModule's **kwargs); used to pre-size - # the activation pool. 0 = unknown (e.g. export/predict construction). - self._batch_size: int = int(kwargs.get("batch_size") or 0) super().__init__(model_config, features, labels, sample_weights, **kwargs) cfg = self._model_config # the family message (e.g. Qwen2RecLM) common = cfg.common # GenerativeRecLMConfig — shared by all families @@ -115,11 +111,25 @@ def __init__( raise ValueError( f"{type(self).__name__}: empty backbone id (see _backbone_id)." ) - # torch_dtype="auto" keeps the stored dtype (bf16); the default upcasts + # Build the EMPTY extended architecture only — NO weight download. The + # config/tokenizer reads are lightweight (they define the module shapes + # and vocab the DCP checkpoint expects); the GB-scale pretrained weights + # load exactly once, at cold start, via `init_from_pretrained` (the + # PIPELINE is the caller). Every restore/eval/export then fills weights + # from DCP and skips the download. See design §1 (import). + # NOTE: named `hf_cfg` (not `cfg`) — `cfg` above is the proto family + # message consumed by `_build_prompt_tokens` below; do not shadow it. + hf_cfg = AutoConfig.from_pretrained(hf_model_id) + # from_config respects torch_dtype; keep the stored dtype (bf16) so the + # DCP checkpoint's tensors match on load_state_dict. The default upcasts # to fp32 (2x memory on GPU). - self.lm = AutoModelForCausalLM.from_pretrained( - hf_model_id, torch_dtype="auto" + self.lm = AutoModelForCausalLM.from_config( + hf_cfg, torch_dtype=hf_cfg.torch_dtype or torch.bfloat16 ) + if next(self.lm.parameters()).dtype != torch.bfloat16: + # Some configs/builders ignore torch_dtype; force bf16 so the empty + # arch matches the checkpoint's dtype for DCP load_state_dict. + self.lm = self.lm.to(torch.bfloat16) tokenizer = AutoTokenizer.from_pretrained(hf_model_id, use_fast=True) # vocab extension: base = tokenizer's next free id BEFORE adding C0.. @@ -136,11 +146,18 @@ def __init__( f"Aborting to avoid silent SID-token mismatch." ) # SID atoms appended directly after the existing vocab (algr's layout); - # offset arithmetic is `token = base + (sid - 1)`. + # offset arithmetic is `token = base + (sid - 1)`. Stash the resize + # target + pad so `init_from_pretrained` re-extends to the SAME shape. + self._target_vocab = base + sid_atoms + self._vocab_pad_mult = pad_mult self.lm.resize_token_embeddings( - base + sid_atoms, pad_to_multiple_of=pad_mult + self._target_vocab, pad_to_multiple_of=self._vocab_pad_mult ) + # keep the extended tokenizer (base vocab + C0..C{sum-1}) so export can + # save it alongside the backbone without rebuilding the SID atoms. + self._hf_tokenizer = tokenizer + # assert C0 landed at the recorded base (the offset arithmetic relies on it) c0_id = tokenizer.convert_tokens_to_ids("C0") if c0_id != base: @@ -171,37 +188,47 @@ def _backbone_id(self) -> str: """ return self._model_config.hf_model_id + def init_from_pretrained(self) -> None: + """Load the pretrained HF backbone weights into ``self.lm``. + + The SINGLE place ``AutoModelForCausalLM.from_pretrained`` runs. The + PIPELINE (``tzrec/main.py``) calls this exactly once, only at COLD START + (no checkpoint to resume). On resume/eval/export the empty arch built in + ``__init__`` is weight-filled by native DCP ``load_state_dict`` instead, + so the GB-scale download is skipped. Re-extends the vocab to the SAME + target/pad as ``__init__`` so the module shapes stay identical. + """ + # torch_dtype="auto" keeps the stored dtype (bf16); the default upcasts + # to fp32 (2x memory on GPU). + lm = AutoModelForCausalLM.from_pretrained( + self._backbone_id(), torch_dtype="auto" + ) + lm.resize_token_embeddings( + self._target_vocab, pad_to_multiple_of=self._vocab_pad_mult + ) + self.lm = lm + def _input_sequence_length(self) -> int: """Truncation length (SID codes) of the user-sequence feature. The data reader caps every row's history at the feature's ``sequence_length``, so it is the guaranteed upper bound used to - pre-size the activation pool (see ``Qwen2RecLM._warmup_alloc``). Returns - 0 if the feature has no length cap, which disables pre-allocation. + pre-size the activation pool (see ``Qwen2RecLM._predict_train``'s + first-step padding). Returns 0 if the feature has no length cap, which + disables pre-allocation. """ for feature in self._features: if feature.config.feature_name == self._input_name: return int(getattr(feature, "sequence_length", 0) or 0) return 0 - def export_hf(self, export_dir: str) -> None: - """Save the HF backbone + extended tokenizer as a HF-loadable dir. - - Single source of truth for HF (``from_pretrained``-loadable) export, - reused by BOTH ``main.export`` (the standard TER export path, after it - overlays the DCP checkpoint) and the in-training checkpoint hook in - ``main._train_and_evaluate`` — so offline and online exports are - identical and there is no separate export tool. Call on rank 0 only; - the dense backbone is replicated, so rank 0 holds the full weights. The - extended tokenizer is rebuilt (base + C0..C{sum-1}) so SID atoms decode - with the same ids the model trained on. - """ - os.makedirs(export_dir, exist_ok=True) - self.lm.save_pretrained(export_dir) - tokenizer = AutoTokenizer.from_pretrained(self._backbone_id(), use_fast=True) - sid_atoms = sum(int(c) for c in self._model_config.common.codebook) - tokenizer.add_tokens([f"C{i}" for i in range(sid_atoms)]) - tokenizer.save_pretrained(export_dir) + def hf_backbone(self): + """The HF backbone module (``export_util.write_hf_assets``/``dcp_to_hf``).""" + return self.lm + + def hf_tokenizer(self): + """The extended tokenizer (base vocab + C0..C{sum-1}) to serialize.""" + return self._hf_tokenizer def _build_prompt_tokens(self, tokenizer, cfg) -> None: """Family hook: cache the tokenised prompt template as buffers. @@ -216,12 +243,26 @@ def _build_prompt_tokens(self, tokenizer, cfg) -> None: ) def init_input(self) -> None: - """No-op override. - - The HF backbone owns its own ``embed_tokens``; we don't use TER's - ``EmbeddingGroup`` at all. Token IDs flow through directly. + """Build the native sparse EmbeddingGroup only if declared. + + The HF backbone owns its own ``embed_tokens`` and SID token ids flow + through it directly, so the dense-only GenerativeRecLM case has no + sparse feature groups and keeps ``embedding_group = None`` (the dense + forward never touches it). When a future model declares sparse + ``feature_groups``, this builds the native ``EmbeddingGroup`` (same as + ``RankModel.init_input``) so those params flow through the native + planner/DMP/DCP path alongside the replicated backbone. """ - self.embedding_group = None + if self._feature_groups: + # NOTE: enables the native sparse path for future dense+sparse models. + from tzrec.modules.embedding import EmbeddingGroup + + self.embedding_group = EmbeddingGroup( + self._features, + self._feature_groups, + ) + else: + self.embedding_group = None @property def device(self) -> torch.device: diff --git a/tzrec/models/model.py b/tzrec/models/model.py index 40da5335a..6b676735b 100644 --- a/tzrec/models/model.py +++ b/tzrec/models/model.py @@ -94,6 +94,16 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: """ raise NotImplementedError + def init_from_pretrained(self) -> None: + """Load pretrained weights at cold start (no checkpoint to restore). + + Lifecycle hook the training pipeline calls only on a fresh run + (``ckpt_path is None``), before distributed wrapping. The default is a + no-op; models backed by an external pretrained source (e.g. an HF + backbone) override it. Resume/eval/export never reach it -- they restore + weights from the checkpoint. + """ + def init_loss(self) -> None: """Initialize loss modules.""" raise NotImplementedError diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index 8e0c1e270..1b5d6322c 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -109,31 +109,6 @@ def _compute_max_total_length(self) -> int: ) return int(frame + self._max_seq_length + self._num_levels) - def _warmup_alloc(self) -> None: - """One-shot: build the CUDA activation pool at the worst case (B, T_max). - - Runs a dummy forward+backward at the configured maximum length so the - caching allocator reserves its largest segments up front; every real - (shorter) batch is then served from that pool, so it never grows - mid-run — which is what stranded segments unevenly across ranks. Fires - from the first training step (earliest point the backbone is on-GPU); - the throwaway gradients are zeroed before the real step runs. - """ - batch_size = self._batch_size or 1 - device = self.device - tok = self._base_vocab # C0 — a valid extended-vocab id - u_rows = [ - torch.full((self._max_seq_length,), tok, dtype=torch.long, device=device) - for _ in range(batch_size) - ] - l_rows = [ - torch.full((self._num_levels,), tok, dtype=torch.long, device=device) - for _ in range(batch_size) - ] - input_ids, labels, attention_mask = self._splice_input_ids(u_rows, l_rows) - self._forward_loss(input_ids, labels, attention_mask)["loss"].backward() - self.lm.zero_grad(set_to_none=True) - def _build_prompt_tokens(self, tokenizer, cfg) -> None: """Tokenise the family chat template once; cache as buffers. @@ -179,6 +154,7 @@ def _splice_input_ids( self, user_seq_rows: List[torch.Tensor], label_rows: List[torch.Tensor], + pad_to: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Build ``(input_ids, labels, attention_mask)``, each ``(B, T_max)``. @@ -196,6 +172,10 @@ def _splice_input_ids( ``user_seq_rows`` / ``label_rows`` already hold extended-vocab token ids on the model device (see ``_sid_token_rows``). + + ``pad_to`` left-extends every row to at least that length (first-step + activation-pool pre-sizing). The supervised tail stays end-aligned, so + labels and the suffix-slice are unchanged — only more masked left-pad. """ assert len(user_seq_rows) == len(label_rows) A = self._num_levels @@ -210,7 +190,7 @@ def _splice_input_ids( ]) for i in range(len(user_seq_rows)) ] - input_ids, attention_mask = self._left_pad(rows_ids) + input_ids, attention_mask = self._left_pad(rows_ids, pad_to=pad_to) # labels: the supervised tail is fixed-width, so left-padding aligns it # to the same columns for every row -> one vectorized write. @@ -252,12 +232,6 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: """Branch 1: teacher-forced forward -> suffix-slice -> CE loss.""" - # one-shot: pre-size the activation pool at the worst-case length on the - # first on-GPU training step (see _warmup_alloc); no-op once warmed or - # when the input feature has no truncation length (pre-allocation off). - if not self._pool_warmed and self._max_total_len > 0 and self.is_train: - self._warmup_alloc() - self._pool_warmed = True # SID indices -> token ids once, at the data boundary (see # _sid_token_rows); the splice then just assembles the prompt. u_rows = self._sid_token_rows( @@ -269,7 +243,23 @@ def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: expected_width=self._num_levels, # answer = one item = num_levels codes ) - input_ids, labels, attention_mask = self._splice_input_ids(u_rows, l_rows) + # One-shot pool pre-sizing: on the FIRST training step, pad the splice to + # the worst-case length so the caching allocator reserves its largest + # (B, T_max) activation segments up front — every shorter batch is then + # served from that pool, keeping per-rank reservations uniform (no + # mid-run growth, the source of the cross-rank imbalance). This rides the + # REAL step: the extra positions are attention-masked and labelled -100, + # so loss and gradient are identical to the unpadded batch. Replaces the + # old _warmup_alloc, whose separate unscaled forward+backward corrupted + # the first optimizer step (the 1.14-flat HR regression). + pad_to = 0 + if not self._pool_warmed and self._max_total_len > 0 and self.is_train: + pad_to = self._max_total_len + self._pool_warmed = True + + input_ids, labels, attention_mask = self._splice_input_ids( + u_rows, l_rows, pad_to=pad_to + ) if self._smoke_log_once and self._first_predict: print( @@ -290,11 +280,7 @@ def _forward_loss( labels: torch.Tensor, attention_mask: torch.Tensor, ) -> Dict[str, torch.Tensor]: - """Teacher-forced forward over spliced ids -> suffix-slice -> CE loss. - - Shared by ``_predict_train`` and the ``_warmup_alloc`` dummy step so the - pre-allocation reproduces the exact training allocation pattern. - """ + """Teacher-forced forward over spliced ids -> suffix-slice -> CE loss.""" outputs = self.lm.model(input_ids=input_ids, attention_mask=attention_mask) hidden = outputs.last_hidden_state # (B, T, D) @@ -373,13 +359,17 @@ def _splice_prompt_ids( return self._left_pad(rows) def _left_pad( - self, rows: List[torch.Tensor] + self, rows: List[torch.Tensor], pad_to: int = 0 ) -> Tuple[torch.Tensor, torch.Tensor]: """Left-pad token rows into ``(input_ids, attention_mask)``, ``(B, T_max)``. Real content is right-aligned, pad at the front. ``attention_mask`` is built from ``ones_like(row)`` (not ``!= pad``) so a real trailing eos is never masked when ``pad_token_id == eos``. + + ``pad_to`` left-extends the batch to at least that many columns (the + first-step activation-pool pre-sizing). Extending on the LEFT keeps the + end-aligned supervised tail in place, so labels/suffix-slice are intact. """ input_ids = pad_sequence( rows, batch_first=True, @@ -389,4 +379,13 @@ def _left_pad( [torch.ones_like(r) for r in rows], batch_first=True, padding_value=0, padding_side="left", ) + if pad_to > input_ids.shape[1]: + B, extra = input_ids.shape[0], pad_to - input_ids.shape[1] + input_ids = torch.cat( + [input_ids.new_full((B, extra), self._pad_token_id), input_ids], + dim=1, + ) + attention_mask = torch.cat( + [attention_mask.new_zeros((B, extra)), attention_mask], dim=1 + ) return input_ids, attention_mask diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/qwen2_rec_lm_test.py index 6f70f76f4..7f77fa765 100644 --- a/tzrec/models/qwen2_rec_lm_test.py +++ b/tzrec/models/qwen2_rec_lm_test.py @@ -196,45 +196,58 @@ def test_compute_max_total_length(self) -> None: m._max_seq_length = 0 # pre-allocation disabled self.assertEqual(m._compute_max_total_length(), 0) - def test_warmup_fires_once_in_training(self) -> None: + def test_first_step_pads_to_max_then_actual_length(self) -> None: m = _stub() m._is_inference = False # not inference + nn.Module.training=True -> is_train m._smoke_log_once = False m._input_name, m._label_name = "user_sequence", "label" m._max_total_len = 50 m._pool_warmed = False - calls = [] - m._warmup_alloc = lambda: calls.append(1) + seen_lens = [] + + def fwd(i, lbl, a): + seen_lens.append(i.shape[1]) + return {"loss": torch.tensor(0.0)} + m._sid_token_rows = lambda jt, expected_width=None, max_codes=None: [ torch.tensor([100, 101, 102]) ] - m._forward_loss = lambda i, lbl, a: {"loss": torch.tensor(0.0)} + m._forward_loss = fwd batch = types.SimpleNamespace( sequence_dense_features={"user_sequence": None, "label": None} ) - m._predict_train(batch) - m._predict_train(batch) - self.assertEqual(len(calls), 1) # one-shot, latched by _pool_warmed + m._predict_train(batch) # first step: pre-size to worst case + m._predict_train(batch) # subsequent step: natural length + # one-shot: first step left-pads to _max_total_len, latched by + # _pool_warmed; later steps use the actual (shorter) length. + self.assertEqual(seen_lens[0], 50) + self.assertLess(seen_lens[1], 50) self.assertTrue(m._pool_warmed) - def test_warmup_skipped_when_disabled(self) -> None: + def test_no_forced_padding_when_disabled(self) -> None: m = _stub() m._is_inference = False m._smoke_log_once = False m._input_name, m._label_name = "user_sequence", "label" m._max_total_len = 0 # max_seq_length unset -> pre-allocation off m._pool_warmed = False - calls = [] - m._warmup_alloc = lambda: calls.append(1) + seen_lens = [] + + def fwd(i, lbl, a): + seen_lens.append(i.shape[1]) + return {"loss": torch.tensor(0.0)} + m._sid_token_rows = lambda jt, expected_width=None, max_codes=None: [ torch.tensor([100, 101, 102]) ] - m._forward_loss = lambda i, lbl, a: {"loss": torch.tensor(0.0)} + m._forward_loss = fwd batch = types.SimpleNamespace( sequence_dense_features={"user_sequence": None, "label": None} ) m._predict_train(batch) - self.assertEqual(calls, []) # never warms up when disabled + # disabled: natural length, never forced to max; flag stays unlatched. + self.assertLess(seen_lens[0], 50) + self.assertFalse(m._pool_warmed) if __name__ == "__main__": diff --git a/tzrec/protos/export.proto b/tzrec/protos/export.proto index 1139ea4b5..178417744 100644 --- a/tzrec/protos/export.proto +++ b/tzrec/protos/export.proto @@ -1,6 +1,14 @@ syntax = "proto2"; package tzrec.protos; +// serialization format produced by `tzrec.export`. +// TORCHSCRIPT: native scripted_model.pt (+ TRT/AOTI), the default. +// HF: a HuggingFace `from_pretrained`-loadable dir (GenerativeRecLM family). +enum ExportFormat { + TORCHSCRIPT = 0; + HF = 1; +} + message ExportConfig { // type of exporter [latest | best] when train_and_evaluation // latest: regularly exports the serving graph and checkpoints @@ -21,4 +29,7 @@ message ExportConfig { optional bool cudnn_allow_tf32 = 5 [default = true]; // whether to use torch.backends.cuda.matmul.allow_tf32 optional bool cuda_matmul_allow_tf32 = 6 [default = false]; + // serialization format produced by `tzrec.export`; selects the export + // branch in main.export (TORCHSCRIPT native graph vs HF backbone dir). + optional ExportFormat export_format = 7 [default = TORCHSCRIPT]; } diff --git a/tzrec/utils/checkpoint_util.py b/tzrec/utils/checkpoint_util.py index 3bcbeded9..8d4edf490 100644 --- a/tzrec/utils/checkpoint_util.py +++ b/tzrec/utils/checkpoint_util.py @@ -306,9 +306,23 @@ def save( optimizer: Optional[optim.Optimizer] = None, dataloader_state: Optional[Dict[str, int]] = None, ) -> str: - """Save a checkpoint at the given step, then request an async prune.""" + """Save a checkpoint at the given step, then request an async prune. + + When ``export_format == HF``, co-locates the HF config + tokenizer (no + weights) in this checkpoint dir so each ``model.ckpt-N/`` is + self-contained and convertible to HF (design §2); ``write_hf_assets`` + no-ops for non-HF models, so gating on the export format is enough. + """ ckpt_dir = os.path.join(self._model_dir, f"model.ckpt-{step}") save_model(ckpt_dir, model, optimizer) + if ( + self._export_config is not None + and self._export_config.export_format == export_pb2.ExportFormat.HF + ): + # Local import avoids a circular import (export_util imports us). + from tzrec.utils.export_util import write_hf_assets + + write_hf_assets(model, ckpt_dir) if dataloader_state is not None: save_dataloader_state(ckpt_dir, dataloader_state) self.prune() @@ -723,7 +737,9 @@ def restore_model( def save_model( - checkpoint_dir: str, model: nn.Module, optimizer: Optional[optim.Optimizer] = None + checkpoint_dir: str, + model: nn.Module, + optimizer: Optional[optim.Optimizer] = None, ) -> None: """Save model state. diff --git a/tzrec/utils/export_util.py b/tzrec/utils/export_util.py index aee51e590..513383bb5 100644 --- a/tzrec/utils/export_util.py +++ b/tzrec/utils/export_util.py @@ -77,6 +77,188 @@ from tzrec.utils.state_dict_util import fix_mch_state, init_parameters +# HF config/tokenizer asset files co-located in each checkpoint dir (no +# weights) and copied into the converted HF export dir. Missing files are +# skipped — different tokenizers emit different subsets (e.g. BPE merges vs. +# sentencepiece vocab). +_HF_ASSET_FILES = ( + "config.json", + "generation_config.json", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + "added_tokens.json", + "special_tokens_map.json", +) + +_HF_EXPORT_META_FILENAME = "hf_export_meta.json" + +# Max wrapper layers to unwrap to reach the HF-backed model: DMP(.module) -> +# TrainWrapper(.model) -> model is 3 hops; 4 is a small cycle guard. +_MAX_WRAPPER_DEPTH = 4 + + +def _unwrap_hf_model(wrapped_model: nn.Module) -> Optional[nn.Module]: + """Walk DMP/TrainWrapper layers down to the model exposing ``hf_backbone``. + + Handles DMP (``.module`` -> TrainWrapper), TrainWrapper (``.model`` -> + GenerativeRecLM), and the GenerativeRecLM itself. Returns ``None`` if no + backbone is found in the wrapper chain (i.e. not an HF-backed model, so + callers no-op). + """ + m = wrapped_model + for _ in range(_MAX_WRAPPER_DEPTH): + if hasattr(m, "hf_backbone"): + return m + if hasattr(m, "module"): # DMP / DDP-style wrapper + m = m.module + elif hasattr(m, "model"): # Train/Predict/Script wrapper + m = m.model + else: + break + return None + + +def write_hf_assets(wrapped_model: nn.Module, save_dir: str) -> None: + """Co-locate the HF config + tokenizer (NO weights) in a checkpoint dir. + + Writes the architecture (``config.json`` + ``generation_config.json``) and + the extended tokenizer alongside the DCP ``model/`` dir so each + ``model.ckpt-N/`` is self-describing and convertible to HF independent of + the current code (design §2). Also records the backbone's true FQN prefix + off the LIVE module graph into ``hf_export_meta.json`` so ``dcp_to_hf`` can + strip it without hard-coding a wrapper/attribute convention. Rank 0 only — + the dense backbone is data-parallel-replicated, so rank 0 holds it. + + ``wrapped_model`` may be the DMP (``.module.model``), the TrainWrapper + (``.model``), or the GenerativeRecLM itself; the inner model is found by + walking the wrappers down to the one exposing ``hf_backbone``. The recorded + prefix is read off ``wrapped_model.named_modules()`` so it matches the FQNs + that ``save_model`` writes from ``wrapped_model.state_dict()``. + """ + if int(os.environ.get("RANK", 0)) != 0: + return + inner = _unwrap_hf_model(wrapped_model) + if inner is None: # not an HF-backed model -> nothing to co-locate + return + os.makedirs(save_dir, exist_ok=True) + + backbone = inner.hf_backbone() + backbone.config.save_pretrained(save_dir) + gen_cfg = getattr(backbone, "generation_config", None) + if gen_cfg is not None: + gen_cfg.save_pretrained(save_dir) + inner.hf_tokenizer().save_pretrained(save_dir) + + # Record the backbone's FQN prefix as it appears in the SAVED state_dict + # (data, not a magic string): robust to wrapper/parallelism and to families + # that don't name the backbone `self.lm`. `named_modules()` FQNs carry the + # DMP/DDP wrapper prefix (e.g. `_dmp_wrapped_module.module.`) that + # `state_dict()` strips, so map it through the same helper `save_model`'s + # DCP keys went through, or the recorded prefix would not match them. + raw_prefix = next( + (n for n, m in wrapped_model.named_modules() if m is backbone), "" + ) + prefix = checkpoint_util._strip_dmp_prefix(raw_prefix) + meta = {"backbone_state_dict_prefix": prefix + ("." if prefix else "")} + with open(os.path.join(save_dir, _HF_EXPORT_META_FILENAME), "w") as f: + json.dump(meta, f, indent=2) + + +def dcp_to_hf(ckpt_dir: str, out_dir: str) -> None: + """Convert a self-contained checkpoint dir to a ``from_pretrained`` HF dir. + + Standalone — builds no live model and runs no ``from_pretrained`` weight + download (design §2). Reads EVERYTHING from the one checkpoint dir: + materialize the DCP state dict, map the wrapper-prefixed keys onto the bare + HF keys, strict-validate against the architecture built from the co-located + ``config.json``, then write safetensors + copy the config/tokenizer. The + numbered steps below walk through it. + """ + from transformers import AutoConfig, AutoModelForCausalLM + from torch.distributed.checkpoint.state_dict_loader import ( + _load_state_dict_from_keys, + ) + + model_ckpt_path = os.path.join(ckpt_dir, "model") + if not os.path.exists(model_ckpt_path): + raise RuntimeError(f"dcp_to_hf: model DCP dir [{model_ckpt_path}] not exists.") + + # 1. materialize the full (replicated) DCP state dict into a plain dict. + # No keys => load everything; non-distributed => full tensors locally. + raw_state: Dict[str, torch.Tensor] = _load_state_dict_from_keys( + checkpoint_id=model_ckpt_path + ) + + # 2. recorded backbone prefix (data, not a magic string); None => derive. + meta_path = os.path.join(ckpt_dir, _HF_EXPORT_META_FILENAME) + prefix: Optional[str] = None + if os.path.exists(meta_path): + with open(meta_path, "r") as f: + prefix = json.load(f).get("backbone_state_dict_prefix") + + # 3. empty backbone from the checkpoint's OWN config -> the target key set. + cfg = AutoConfig.from_pretrained(ckpt_dir) + empty = AutoModelForCausalLM.from_config(cfg) + target_keys: Set[str] = set(empty.state_dict().keys()) + + def _strip_recorded_prefix() -> Optional[Dict[str, torch.Tensor]]: + """Strip the recorded prefix; None unless it yields an EXACT match.""" + if not prefix: + return None + out = {k[len(prefix) :]: v for k, v in raw_state.items() if k.startswith(prefix)} + return out if set(out) == target_keys else None + + def _derive_by_suffix() -> Optional[Dict[str, torch.Tensor]]: + """Each target key is a unique suffix of exactly one DCP key; None if not.""" + out: Dict[str, torch.Tensor] = {} + for tk in target_keys: + matches = [k for k in raw_state if k == tk or k.endswith("." + tk)] + if len(matches) != 1: + return None + out[tk] = raw_state[matches[0]] + return out + + # The recorded prefix is a hint, not gospel: if it does not map exactly onto + # the architecture (stale metadata, a wrapper/parallelism change since the + # checkpoint was written), self-heal by suffix-matching the actual DCP keys + # to the architecture's keys -- which is independent of any prefix convention. + mapped = _strip_recorded_prefix() + if mapped is None: + if prefix: + logger.warning( + f"dcp_to_hf: recorded prefix [{prefix}] did not map exactly onto " + "the architecture; deriving the backbone prefix by suffix-matching." + ) + mapped = _derive_by_suffix() + + # STRICT validation: exact 1:1 with the architecture, fail loudly. Reached + # only when neither path matched -- a genuine architecture/checkpoint + # mismatch, never a silent partial load. + if mapped is None or set(mapped.keys()) != target_keys: + got = set(mapped.keys()) if mapped is not None else set() + missing = sorted(target_keys - got) + extra = sorted(got - target_keys) + raise RuntimeError( + "dcp_to_hf: cannot map the DCP state dict onto the backbone " + f"architecture (recorded prefix={prefix!r}). missing={missing[:10]} " + f"extra={extra[:10]}. Refusing to write a partially-loaded HF model." + ) + + # 4. write weights as safetensors. Clone so save_file gets contiguous, + # storage-owning tensors. + os.makedirs(out_dir, exist_ok=True) + mapped = {k: v.contiguous().clone() for k, v in mapped.items()} + save_file(mapped, os.path.join(out_dir, "model.safetensors")) + + # 5. copy the co-located HF config + tokenizer assets. + for fname in _HF_ASSET_FILES: + src = os.path.join(ckpt_dir, fname) + if os.path.exists(src): + shutil.copy(src, os.path.join(out_dir, fname)) + + def export_model( pipeline_config: EasyRecConfig, model: BaseModule, From 9b43143dc710b47c60098788e34304c7ada50783 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Tue, 16 Jun 2026 11:57:03 +0000 Subject: [PATCH 12/99] [refactor] generative-rec LM: encapsulate __init__ into helpers, trim comments __init__ was ~105 lines of mixed concerns + heavy narration. Extract three behavior-preserving private helpers and condense comments to the load-bearing rationale only: - _read_common_config(common) -> sid_atoms: proto knobs + codebook guard. - _build_backbone() -> module: empty bf16 arch (no weight download) + empty-id guard. - _build_extended_tokenizer(sid_atoms) -> (tokenizer, base): add C0.. atoms, resize self.lm, the added==sid_atoms and C0-at-base guards. __init__ now reads as super().__init__ -> read config -> build backbone -> build tokenizer -> pad/prompt/debug tail. All instance attributes, guard messages, and effect order are unchanged; no behavior change. 22 model unit tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 110 +++++++++++++++--------------- 1 file changed, 55 insertions(+), 55 deletions(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 027a94715..da5f6dcd0 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -83,16 +83,35 @@ def __init__( **kwargs: Any, ) -> None: super().__init__(model_config, features, labels, sample_weights, **kwargs) - cfg = self._model_config # the family message (e.g. Qwen2RecLM) + cfg = self._model_config # family message (e.g. Qwen2RecLM) common = cfg.common # GenerativeRecLMConfig — shared by all families - # shared proto -> python knobs + sid_atoms = self._read_common_config(common) + + self.lm = self._build_backbone() + tokenizer, base = self._build_extended_tokenizer(sid_atoms) + self._hf_tokenizer = tokenizer + self._base_vocab = base + + # pad token for the left-padded splice (fall back to eos) + pad_id = tokenizer.pad_token_id + if pad_id is None: + pad_id = tokenizer.eos_token_id + self._pad_token_id = int(pad_id) + + self._build_prompt_tokens(tokenizer, cfg) + + # one-shot debug dump of the first spliced batch + self._smoke_log_once = os.environ.get("TZREC_GENRECLM_DEBUG", "0") == "1" + self._first_predict = True + + def _read_common_config(self, common: Any) -> int: + """Parse shared proto knobs into attributes; return the SID atom count.""" self._input_name: str = common.user_sequence_feature_name self._label_name: str = common.label_feature_name self._ignore_index: int = int(common.ignore_index) - # max history length (SID codes) for activation pre-allocation, taken - # from the user-sequence feature's truncation length (the data reader - # caps every row at it, so it's the guaranteed upper bound). 0 = off. + # max history (SID codes) for activation pre-sizing = the user-sequence + # feature's truncation length (the reader caps every row at it). 0 = off. self._max_seq_length: int = self._input_sequence_length() codebook = list(common.codebook) if len(codebook) == 0: @@ -102,63 +121,56 @@ def __init__( ) # len(codebook) = SID codes per item (answer width); sum = vocab atoms. self._num_levels = len(codebook) - sid_atoms = sum(int(c) for c in codebook) - pad_mult = int(common.vocab_pad_to_multiple_of) or 128 + self._vocab_pad_mult = int(common.vocab_pad_to_multiple_of) or 128 + return sum(int(c) for c in codebook) - # backbone + tokenizer (the backbone is family-owned; see _backbone_id) + def _build_backbone(self) -> Any: + """Build the EMPTY extended architecture in bf16 — no weight download. + + Config reads only define the module shapes the DCP checkpoint expects; + the GB-scale pretrained weights load once at cold start via + ``init_from_pretrained``, then DCP fills them on every restore/eval. + """ hf_model_id = self._backbone_id() if not hf_model_id: raise ValueError( f"{type(self).__name__}: empty backbone id (see _backbone_id)." ) - # Build the EMPTY extended architecture only — NO weight download. The - # config/tokenizer reads are lightweight (they define the module shapes - # and vocab the DCP checkpoint expects); the GB-scale pretrained weights - # load exactly once, at cold start, via `init_from_pretrained` (the - # PIPELINE is the caller). Every restore/eval/export then fills weights - # from DCP and skips the download. See design §1 (import). - # NOTE: named `hf_cfg` (not `cfg`) — `cfg` above is the proto family - # message consumed by `_build_prompt_tokens` below; do not shadow it. hf_cfg = AutoConfig.from_pretrained(hf_model_id) - # from_config respects torch_dtype; keep the stored dtype (bf16) so the - # DCP checkpoint's tensors match on load_state_dict. The default upcasts - # to fp32 (2x memory on GPU). - self.lm = AutoModelForCausalLM.from_config( + # keep the stored dtype (bf16) so the DCP checkpoint's tensors match on + # load_state_dict; the default upcasts to fp32 (2x memory on GPU). + lm = AutoModelForCausalLM.from_config( hf_cfg, torch_dtype=hf_cfg.torch_dtype or torch.bfloat16 ) - if next(self.lm.parameters()).dtype != torch.bfloat16: - # Some configs/builders ignore torch_dtype; force bf16 so the empty - # arch matches the checkpoint's dtype for DCP load_state_dict. - self.lm = self.lm.to(torch.bfloat16) - tokenizer = AutoTokenizer.from_pretrained(hf_model_id, use_fast=True) - - # vocab extension: base = tokenizer's next free id BEFORE adding C0.. - # (use len(tokenizer), NOT config.vocab_size which counts reserved slots). + if next(lm.parameters()).dtype != torch.bfloat16: + # Some configs/builders ignore torch_dtype; force bf16 to match DCP. + lm = lm.to(torch.bfloat16) + return lm + + def _build_extended_tokenizer(self, sid_atoms: int) -> tuple[Any, int]: + """Add the SID atoms ``C0..C{sid_atoms-1}`` and resize ``self.lm``. + + Returns ``(tokenizer, base)`` where ``base`` is the tokenizer's next free + id BEFORE adding the atoms — use ``len(tokenizer)``, NOT + ``config.vocab_size`` (which counts reserved slots). The atoms append + directly after the existing vocab (algr's layout), so the splice offset + is ``token = base + (sid - 1)``. + """ + tokenizer = AutoTokenizer.from_pretrained(self._backbone_id(), use_fast=True) base = len(tokenizer) - new_atoms = [f"C{i}" for i in range(sid_atoms)] - added = tokenizer.add_tokens(new_atoms) + added = tokenizer.add_tokens([f"C{i}" for i in range(sid_atoms)]) if added != sid_atoms: - # The tokenizer already had some Cxxx tokens — we expect a fresh - # base, so this would silently break our offset arithmetic. + # pre-existing Cxxx tokens would silently break the offset arithmetic. raise RuntimeError( f"GenerativeRecLM: tokenizer was expected to grow by " f"{sid_atoms} new atoms, only added {added}. " f"Aborting to avoid silent SID-token mismatch." ) - # SID atoms appended directly after the existing vocab (algr's layout); - # offset arithmetic is `token = base + (sid - 1)`. Stash the resize - # target + pad so `init_from_pretrained` re-extends to the SAME shape. + # stash the resize target so init_from_pretrained re-extends identically. self._target_vocab = base + sid_atoms - self._vocab_pad_mult = pad_mult self.lm.resize_token_embeddings( self._target_vocab, pad_to_multiple_of=self._vocab_pad_mult ) - - # keep the extended tokenizer (base vocab + C0..C{sum-1}) so export can - # save it alongside the backbone without rebuilding the SID atoms. - self._hf_tokenizer = tokenizer - - # assert C0 landed at the recorded base (the offset arithmetic relies on it) c0_id = tokenizer.convert_tokens_to_ids("C0") if c0_id != base: raise RuntimeError( @@ -166,19 +178,7 @@ def __init__( f"C0 at token id {base}, got {c0_id}. " f"Splice arithmetic would produce wrong token ids." ) - self._base_vocab = base - - # pad token for the left-padded splice (fall back to eos) - pad_id = tokenizer.pad_token_id - if pad_id is None: - pad_id = tokenizer.eos_token_id - self._pad_token_id = int(pad_id) - - self._build_prompt_tokens(tokenizer, cfg) - - # one-shot debug dump of the first spliced batch - self._smoke_log_once = (os.environ.get("TZREC_GENRECLM_DEBUG", "0") == "1") - self._first_predict = True + return tokenizer, base def _backbone_id(self) -> str: """Family hook: the HF model id to load for ``self.lm``. From 5526dc9a4b42f4bad0558fe0e5ece6fab8ae5af1 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 17 Jun 2026 02:36:21 +0000 Subject: [PATCH 13/99] =?UTF-8?q?[fix]=20generative-rec=20LM:=20build=20LM?= =?UTF-8?q?=20in=20fp32=20(master=20weights)=20=E2=80=94=20fixes=20lr1e-5?= =?UTF-8?q?=20collapse?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause (overnight investigation): the LM was built in bf16 params with mixed_precision unset, so the optimizer updated bf16 weights directly with no fp32 master. Adam's small updates at lr=1e-5 (~1e-5) fall below the bf16 ULP of the weights and round to zero -> weights freeze -> training collapses (eval ce hard-plateaus ~5.76, HR ~0). lr=5e-5 masked it (5x larger updates clear the ULP). ALGR's HF-Trainer bf16:true keeps an fp32 master, so it trains fine at lr1e-5. Fix: build the LM in fp32 in BOTH paths (`_build_backbone` from_config and `init_from_pretrained` from_pretrained) so the optimizer keeps fp32 master weights. Set `mixed_precision:"BF16"` in the run config for bf16 *compute* speed (autocast) on the fp32 master — the standard AMP pattern, mirroring ALGR. Proven by a single-variable A/B (only precision changed): bf16-params: lr1e-5 = 0.02 flat (collapse) | lr5e-5 = 2.71 rising fp32-master: lr1e-5 = 1.17 rising (1ep) | lr5e-5 = 2.68 rising The fix recovers lr1e-5 ~60x without regressing lr5e-5 (2.68 ~= 2.71). Tradeoff: fp32 master -> fp32 DCP checkpoint (~2x) + ~64GB/GPU (vs 54). Old bf16 checkpoints still restore (upcast to fp32). Details in ai_report/MASTER_EXPERIMENT_REPORT.md. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 26 ++++++++++++++------------ 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index da5f6dcd0..d83d68b01 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -125,7 +125,7 @@ def _read_common_config(self, common: Any) -> int: return sum(int(c) for c in codebook) def _build_backbone(self) -> Any: - """Build the EMPTY extended architecture in bf16 — no weight download. + """Build the EMPTY extended architecture in fp32 (master) — no download. Config reads only define the module shapes the DCP checkpoint expects; the GB-scale pretrained weights load once at cold start via @@ -137,14 +137,14 @@ def _build_backbone(self) -> Any: f"{type(self).__name__}: empty backbone id (see _backbone_id)." ) hf_cfg = AutoConfig.from_pretrained(hf_model_id) - # keep the stored dtype (bf16) so the DCP checkpoint's tensors match on - # load_state_dict; the default upcasts to fp32 (2x memory on GPU). - lm = AutoModelForCausalLM.from_config( - hf_cfg, torch_dtype=hf_cfg.torch_dtype or torch.bfloat16 - ) - if next(lm.parameters()).dtype != torch.bfloat16: - # Some configs/builders ignore torch_dtype; force bf16 to match DCP. - lm = lm.to(torch.bfloat16) + # Build in fp32 so the optimizer keeps fp32 MASTER weights. With bf16 + # params, Adam's small updates (e.g. at lr=1e-5) fall below the bf16 ULP + # and round to zero -> weights freeze -> training collapses (the lr1e-5 + # bug). Use mixed_precision:"BF16" (autocast) for bf16 *compute* speed on + # the fp32 master; the DCP checkpoint is then fp32 (consistent on restore). + lm = AutoModelForCausalLM.from_config(hf_cfg, torch_dtype=torch.float32) + if next(lm.parameters()).dtype != torch.float32: + lm = lm.to(torch.float32) return lm def _build_extended_tokenizer(self, sid_atoms: int) -> tuple[Any, int]: @@ -198,10 +198,12 @@ def init_from_pretrained(self) -> None: so the GB-scale download is skipped. Re-extends the vocab to the SAME target/pad as ``__init__`` so the module shapes stay identical. """ - # torch_dtype="auto" keeps the stored dtype (bf16); the default upcasts - # to fp32 (2x memory on GPU). + # Load fp32 (master weights), NOT torch_dtype="auto" (which keeps the + # backbone's stored bf16): bf16 params underflow Adam's lr=1e-5 updates -> + # collapse; fp32 master fixes it. Must match _build_backbone's fp32 so the + # cold-start and restore arches agree. mixed_precision:"BF16" -> bf16 compute. lm = AutoModelForCausalLM.from_pretrained( - self._backbone_id(), torch_dtype="auto" + self._backbone_id(), torch_dtype=torch.float32 ) lm.resize_token_embeddings( self._target_vocab, pad_to_multiple_of=self._vocab_pad_mult From d72436a282e941435984ab9011204444b27be696 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 17 Jun 2026 04:12:13 +0000 Subject: [PATCH 14/99] [fix] generative-rec LM: SID validity gate + code-review fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Layer A — _generate now validates each beam against the per-level SID bands; malformed candidates (early EOS / non-SID / wrong-level atom) collapse to a -1 sentinel that can never match a real item, and the fixed-width canvas removes the reshape crash when beams stop early. The gate + token<->SID inversion live in the base GenerativeRecLM (_validate_sid_candidates, _sid_level_bands) as one source of truth, reused by every family. Code-review findings: - #2 (perf): cache self._suffix_keep; _forward_loss slices a constant suffix instead of recomputing it per step -> drops 2 GPU->CPU syncs/step. - #5: isolate the rank-0 write_hf_assets call (try/except + log) so an asset write error can't abort before the next collective and hang other ranks. - #7: make the dense-only contract explicit (embedding_group = None); remove the never-called init_input (calling it would wrongly build an unused sharded table for the SEQUENCE feature, which flows as raw token ids). - #8: drop the dead batch_size plumbing from _create_model + call sites. - #9: _resolve_pad_token_id asserts pad/eos present (clear error, not int(None)). - dtype: single _PARAM_DTYPE source of truth (both builders fp32-master). Also: streamline comments, pin transformers==4.51.2 (<5.0). Tests: 25 genrec unit tests (Layer-A valid/malformed/narrow-tail, _suffix_keep equivalence, pad resolution). Co-Authored-By: Claude Opus 4.8 (1M context) --- requirements/runtime.txt | 1 + tzrec/main.py | 6 - tzrec/models/generative_rec_lm.py | 229 ++++++++++++++----------- tzrec/models/generative_rec_lm_test.py | 20 +++ tzrec/models/qwen2_rec_lm.py | 154 ++++++----------- tzrec/models/qwen2_rec_lm_test.py | 102 +++++++++-- tzrec/utils/checkpoint_util.py | 12 +- 7 files changed, 305 insertions(+), 219 deletions(-) diff --git a/requirements/runtime.txt b/requirements/runtime.txt index 0be1b001b..3c605b062 100644 --- a/requirements/runtime.txt +++ b/requirements/runtime.txt @@ -22,3 +22,4 @@ tensorboard torch==2.11.0 torchmetrics==1.0.3 torchrec==1.6.0 +transformers==4.51.2 # generative-rec LM backbone (HF); pin <5.0 — 5.x flips from_pretrained dtype default to "auto"(bf16) diff --git a/tzrec/main.py b/tzrec/main.py index a1654e665..3e3158803 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -131,7 +131,6 @@ def _create_model( labels: List[str], sample_weights: Optional[List[str]] = None, sampler_type: Optional[str] = None, - batch_size: Optional[int] = None, ) -> BaseModel: """Build model. @@ -141,8 +140,6 @@ def _create_model( labels (list): list of label names. sample_weights (list): list of sample weight names. sampler_type (str): negative sampler type - batch_size (int): per-rank batch size (data_config.batch_size); most - models ignore it, generative LMs use it to pre-size their pool. Return: model: a EasyRec Model. @@ -157,7 +154,6 @@ def _create_model( labels, sample_weights=sample_weights, sampler_type=sampler_type, - batch_size=batch_size, ) kernel = Kernel[KernelProto.Name(model_config.kernel)] @@ -665,7 +661,6 @@ def train_and_evaluate( list(data_config.label_fields), sample_weights=list(data_config.sample_weight_fields), sampler_type=sampler_type, - batch_size=data_config.batch_size, ) # Cold-start gate (training-only). `_create_model` builds the EMPTY extended # architecture (GenerativeRecLM: from_config, no weight download). On a fresh @@ -855,7 +850,6 @@ def evaluate( list(data_config.label_fields), sample_weights=list(data_config.sample_weight_fields), sampler_type=sampler_type, - batch_size=data_config.batch_size, ) model = TrainWrapper( model, device=device, mixed_precision=train_config.mixed_precision diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index d83d68b01..9a263779b 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -6,29 +6,22 @@ """Generic generative-recommendation language-model base for TorchEasyRec. -Implements the FINAL design (see FINAL_DESIGN_GENERATIVE_REC_LM.md): - - * Per-family subclasses (design §2 / G4): ``GenerativeRecLM`` is the - abstract base; each LLM family is a concrete subclass implementing the - ``_build_prompt_tokens`` and ``predict`` hooks (e.g. ``Qwen2RecLM`` in - ``tzrec/models/qwen2_rec_lm.py``). The pipeline config selects the family - by its own oneof entry (``qwen2_rec_lm``), whose message-type name resolves - directly to the same-named class via the BaseModel registry — no dispatch. - Shared config lives in ``GenerativeRecLMConfig`` (the family message's - ``common`` field); family-specific knobs sit on the family message. - * Streaming sample format: each row carries two raw-int64 sequence features, - ``user_sequence`` (list[int]) and ``label`` (list[int]), both holding raw - SID indices in ``[1, sum(codebook)]``. - * The chat template is tokenised ONCE at ``__init__`` and cached as - ``nn.Module`` non-persistent buffers, so per-batch encoding is purely - integer arithmetic + tensor concatenation (no HF tokenizer in the hot - path). - * SID → token id by integer offset: ``token = sid + base_vocab - 1`` (the - SID atoms ``C0..C{sum-1}`` are added right after the original vocabulary, - no [SEP] in between; matches algr's ``add_tokens`` layout). - * Left padding with ``eos_token_id`` (L7 fix from §11 of the design doc) — - real content sits at the END of every row so the suffix slice captures - only ``[response + end_markers]`` and matches algr's pad-side exactly. +* Per-family subclasses: ``GenerativeRecLM`` is the abstract base; each LLM +family is a concrete subclass implementing the ``_build_prompt_tokens`` and +``predict`` hooks (e.g. ``Qwen2RecLM``). The pipeline config selects the +family by its own oneof entry, whose message-type name resolves directly to +the same-named class via the BaseModel registry. Shared config lives in +``GenerativeRecLMConfig`` (the family message's ``common`` field). +* Streaming sample format: each row carries two raw-int64 sequence features, +``user_sequence`` and ``label``, both holding raw SID indices in +``[1, sum(codebook)]``. +* The chat template is tokenised ONCE at ``__init__`` and cached as +non-persistent buffers, so per-batch encoding is integer arithmetic only +(no HF tokenizer in the hot path). +* SID -> token id by integer offset: ``token = sid + base_vocab - 1`` (the SID +atoms ``C0..C{sum-1}`` are added right after the original vocabulary). +* Left padding with ``eos_token_id``: real content sits at the END of every +row so the suffix slice captures only ``[response + end_markers]``. """ from __future__ import annotations @@ -49,31 +42,33 @@ class GenerativeRecLM(BaseModel): """Abstract base for HF-backed generative-recommendation LMs. - The base owns the architecture-agnostic plumbing: model construction, - SID vocab extension, the shared sample data-prep (``_sid_token_rows`` / - ``_tokenize_sids`` — the streaming SID sample contract is the same for all - families, design §1), loss, and metrics. The two architecture-specific - pieces are abstract hooks that each family subclass implements (§15/§16): + The base owns the architecture-agnostic plumbing: model construction, SID + vocab extension, the shared sample data-prep (``_sid_token_rows`` / + ``_tokenize_sids``), loss, and metrics. Each family subclass implements two + architecture-specific hooks: _build_prompt_tokens(tokenizer, cfg) — cache the prompt template predict(batch) — build inputs + HF forward - Family proto contract: every family message embeds - ``GenerativeRecLMConfig common = 1`` (shared config the base reads) and - supplies a backbone, by default via an ``hf_model_id`` field (overridable - through ``_backbone_id``). + Family proto contract: every family message embeds ``GenerativeRecLMConfig + common = 1`` and supplies a backbone (by default an ``hf_model_id`` field, + overridable via ``_backbone_id``). - ``Qwen2RecLM`` (``tzrec/models/qwen2_rec_lm.py``) provides the decoder-only - chat implementation (ChatML splice + ``.model``/``.lm_head`` forward), - reusable by Llama/Mistral/Gemma/Phi-style families; GPT-NeoX/RWKV/Mamba/T5 - each need their own. Each family registers directly (its oneof message-type - name == the class name); there is no ``class_name`` dispatch. + ``Qwen2RecLM`` provides the decoder-only chat implementation, reusable by + Llama/Mistral/Gemma/Phi-style families; GPT-NeoX/RWKV/Mamba/T5 each need + their own. Each family registers directly (oneof message-type name == class + name). """ # predictions key the inference branch emits generated SIDs under, stable # across families (PredictWrapper ``output_cols`` should reference it). GENERATED_SIDS_KEY = "generated_sids" + # Single source of truth for the backbone PARAM dtype: fp32 MASTER weights. + # The optimizer needs fp32 to avoid bf16-ULP underflow at small lr; bf16 + # *compute* comes from mixed_precision:"BF16" autocast, not the param dtype. + _PARAM_DTYPE = torch.float32 + def __init__( self, model_config: ModelConfig, @@ -93,18 +88,37 @@ def __init__( self._hf_tokenizer = tokenizer self._base_vocab = base - # pad token for the left-padded splice (fall back to eos) - pad_id = tokenizer.pad_token_id - if pad_id is None: - pad_id = tokenizer.eos_token_id - self._pad_token_id = int(pad_id) + self._pad_token_id = self._resolve_pad_token_id(tokenizer) self._build_prompt_tokens(tokenizer, cfg) + # Dense-only: the HF backbone owns its embeddings and SID token ids flow + # through it directly (the SEQUENCE feature is consumed as raw token ids in + # _sid_token_rows, NOT via an EmbeddingGroup), so there is no native sparse + # path. Set explicitly (vs RankModel.init_input) to keep the contract clear. + self.embedding_group = None + # one-shot debug dump of the first spliced batch self._smoke_log_once = os.environ.get("TZREC_GENRECLM_DEBUG", "0") == "1" self._first_predict = True + @staticmethod + def _resolve_pad_token_id(tokenizer: Any) -> int: + """Pad id for the left-padded splice, falling back to eos. + + Qwen2.5 has both, but the base is reused by Llama/Mistral/Gemma/Phi — fail + loudly if a tokenizer has neither rather than an opaque ``int(None)``. + """ + pad_id = tokenizer.pad_token_id + if pad_id is None: + pad_id = tokenizer.eos_token_id + if pad_id is None: + raise ValueError( + "GenerativeRecLM: tokenizer has neither pad_token_id nor " + "eos_token_id; cannot choose a pad id for the left-padded splice." + ) + return int(pad_id) + def _read_common_config(self, common: Any) -> int: """Parse shared proto knobs into attributes; return the SID atom count.""" self._input_name: str = common.user_sequence_feature_name @@ -115,12 +129,15 @@ def _read_common_config(self, common: Any) -> int: self._max_seq_length: int = self._input_sequence_length() codebook = list(common.codebook) if len(codebook) == 0: - raise ValueError( - "GenerativeRecLM: codebook must be non-empty " - "(see design §3 — required field)" - ) + raise ValueError("GenerativeRecLM: codebook must be non-empty.") # len(codebook) = SID codes per item (answer width); sum = vocab atoms. self._num_levels = len(codebook) + # Per-level SID validity bands, cached as (num_levels,) buffers so the + # inference gate (_validate_sid_candidates) can reject malformed beams with + # a vectorized band check and no host sync. Non-persistent: derived config. + lo, hi = self._sid_level_bands(codebook) + self.register_buffer("_sid_lvl_lo", lo, persistent=False) + self.register_buffer("_sid_lvl_hi", hi, persistent=False) self._vocab_pad_mult = int(common.vocab_pad_to_multiple_of) or 128 return sum(int(c) for c in codebook) @@ -137,14 +154,13 @@ def _build_backbone(self) -> Any: f"{type(self).__name__}: empty backbone id (see _backbone_id)." ) hf_cfg = AutoConfig.from_pretrained(hf_model_id) - # Build in fp32 so the optimizer keeps fp32 MASTER weights. With bf16 - # params, Adam's small updates (e.g. at lr=1e-5) fall below the bf16 ULP - # and round to zero -> weights freeze -> training collapses (the lr1e-5 - # bug). Use mixed_precision:"BF16" (autocast) for bf16 *compute* speed on - # the fp32 master; the DCP checkpoint is then fp32 (consistent on restore). - lm = AutoModelForCausalLM.from_config(hf_cfg, torch_dtype=torch.float32) - if next(lm.parameters()).dtype != torch.float32: - lm = lm.to(torch.float32) + # Build in fp32 so the optimizer keeps fp32 MASTER weights: with bf16 + # params, Adam's small updates (e.g. lr=1e-5) fall below the bf16 ULP and + # round to zero, freezing the weights. Use mixed_precision:"BF16" for bf16 + # compute speed; the DCP checkpoint stays fp32 (consistent on restore). + lm = AutoModelForCausalLM.from_config(hf_cfg, torch_dtype=self._PARAM_DTYPE) + if next(lm.parameters()).dtype != self._PARAM_DTYPE: + lm = lm.to(self._PARAM_DTYPE) return lm def _build_extended_tokenizer(self, sid_atoms: int) -> tuple[Any, int]: @@ -153,8 +169,8 @@ def _build_extended_tokenizer(self, sid_atoms: int) -> tuple[Any, int]: Returns ``(tokenizer, base)`` where ``base`` is the tokenizer's next free id BEFORE adding the atoms — use ``len(tokenizer)``, NOT ``config.vocab_size`` (which counts reserved slots). The atoms append - directly after the existing vocab (algr's layout), so the splice offset - is ``token = base + (sid - 1)``. + directly after the existing vocab, so the splice offset is + ``token = base + (sid - 1)``. """ tokenizer = AutoTokenizer.from_pretrained(self._backbone_id(), use_fast=True) base = len(tokenizer) @@ -198,12 +214,11 @@ def init_from_pretrained(self) -> None: so the GB-scale download is skipped. Re-extends the vocab to the SAME target/pad as ``__init__`` so the module shapes stay identical. """ - # Load fp32 (master weights), NOT torch_dtype="auto" (which keeps the - # backbone's stored bf16): bf16 params underflow Adam's lr=1e-5 updates -> - # collapse; fp32 master fixes it. Must match _build_backbone's fp32 so the - # cold-start and restore arches agree. mixed_precision:"BF16" -> bf16 compute. + # Load fp32 master weights, NOT torch_dtype="auto" (which keeps the stored + # bf16): bf16 params underflow Adam's lr=1e-5 updates. Must match + # _build_backbone's fp32 so the cold-start and restore arches agree. lm = AutoModelForCausalLM.from_pretrained( - self._backbone_id(), torch_dtype=torch.float32 + self._backbone_id(), torch_dtype=self._PARAM_DTYPE ) lm.resize_token_embeddings( self._target_vocab, pad_to_multiple_of=self._vocab_pad_mult @@ -244,28 +259,6 @@ def _build_prompt_tokens(self, tokenizer, cfg) -> None: f"(GenerativeRecLM is abstract)." ) - def init_input(self) -> None: - """Build the native sparse EmbeddingGroup only if declared. - - The HF backbone owns its own ``embed_tokens`` and SID token ids flow - through it directly, so the dense-only GenerativeRecLM case has no - sparse feature groups and keeps ``embedding_group = None`` (the dense - forward never touches it). When a future model declares sparse - ``feature_groups``, this builds the native ``EmbeddingGroup`` (same as - ``RankModel.init_input``) so those params flow through the native - planner/DMP/DCP path alongside the replicated backbone. - """ - if self._feature_groups: - # NOTE: enables the native sparse path for future dense+sparse models. - from tzrec.modules.embedding import EmbeddingGroup - - self.embedding_group = EmbeddingGroup( - self._features, - self._feature_groups, - ) - else: - self.embedding_group = None - @property def device(self) -> torch.device: """Device the HF backbone runs on — the single source for model I/O.""" @@ -280,6 +273,48 @@ def _tokenize_sids(self, sids: torch.Tensor) -> torch.Tensor: """ return sids + (self._base_vocab - 1) + @staticmethod + def _sid_level_bands(codebook: Any) -> tuple[torch.Tensor, torch.Tensor]: + """Per-level closed SID bands ``(lo, hi)`` as ``(num_levels,)`` long tensors. + + Level ``j`` occupies a DISJOINT band ``[offset_j + 1, offset_j + + codebook[j]]`` where ``offset_j = sum(codebook[:j])``. Single source of + truth: both ``_read_common_config`` and the tests build bands from this. + """ + lo, hi, acc = [], [], 0 + for c in codebook: + lo.append(acc + 1) + hi.append(acc + int(c)) + acc += int(c) + return ( + torch.tensor(lo, dtype=torch.long), + torch.tensor(hi, dtype=torch.long), + ) + + def _validate_sid_candidates( + self, new_tokens: torch.Tensor, batch_size: int + ) -> torch.Tensor: + """Map a generated token tail back to SIDs and reject malformed candidates. + + Inference-side counterpart of ``_tokenize_sids``. ``new_tokens`` is the + per-beam generated tail ``(B*num_return, w)`` (``w`` may be < ``num_levels`` + when beams stop early). Returns ``(batch_size, num_return, num_levels)`` raw + SIDs with every MALFORMED candidate set to the ``-1`` sentinel — early EOS, + a non-SID token, or a wrong-level atom (outside THAT level's band). ``-1`` + can never match a real item, and the fixed-width canvas keeps the reshape + rectangular even when every beam stopped early. + """ + sids = new_tokens - (self._base_vocab - 1) + canvas = sids.new_full((sids.shape[0], self._num_levels), -1) + w = min(sids.shape[1], self._num_levels) + canvas[:, :w] = sids[:, :w] + in_band = (canvas >= self._sid_lvl_lo) & (canvas <= self._sid_lvl_hi) + canvas[~in_band.all(dim=1)] = -1 + # view groups beams under the right user (generate() returns rows + # batch-major: [b0_beam0, b0_beam1, ..., b1_beam0, ...]); the in-place mask + # above preserves that order (no filter/scatter that would scramble it). + return canvas.view(batch_size, -1, self._num_levels) + def _sid_token_rows( self, jt, @@ -288,23 +323,17 @@ def _sid_token_rows( ) -> List[torch.Tensor]: """Read a SID jagged feature -> per-row token-id tensors. - TER delivers the feature as a JaggedTensor (flat ``values`` + - ``lengths``); ``values`` may arrive as float / shape ``(N, 1)``. The - whole batch is tokenized once (``_tokenize_sids``) on the backbone - device, then split into rows. - - ``expected_width``, when set, enforces the sample contract here at the - data boundary: every row must have exactly that many codes (e.g. the - answer = ``num_levels``); a deviation is an anomalous sample. - - ``max_codes``, when set, caps each row to its most-recent whole items — - the last ``floor(max_codes / num_levels) * num_levels`` codes, dropping - the oldest *head* (sequences are oldest->newest, so recent behaviour is - preserved). FG_NONE does not truncate, so this is what actually enforces - the feature's ``sequence_length`` — guaranteeing the pre-allocated pool - covers every batch. Done on host views before the H2D copy, and skipped - entirely unless some row overflows, so it's free in the common case and - shrinks downstream work (and forward ``T``) when it fires. + TER delivers the feature as a JaggedTensor (flat ``values`` + ``lengths``); + ``values`` may arrive as float / shape ``(N, 1)``. The whole batch is + tokenized once on the backbone device, then split into rows. + + ``expected_width``, when set, enforces the sample contract: every row must + have exactly that many codes (the answer = ``num_levels``). + + ``max_codes``, when set, caps each row to its most-recent whole items (the + last ``floor(max_codes / num_levels) * num_levels`` codes, dropping the + oldest head) so the pre-allocated pool covers every batch. Done on host + views before the H2D copy, skipped unless a row overflows. """ values = jt.values() lengths = jt.lengths() diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py index 630195004..934df831b 100644 --- a/tzrec/models/generative_rec_lm_test.py +++ b/tzrec/models/generative_rec_lm_test.py @@ -55,6 +55,26 @@ def test_registry_dispatch(self) -> None: self.assertIs(BaseModel.create_class("Qwen2RecLM"), Qwen2RecLM) self.assertTrue(issubclass(Qwen2RecLM, GenerativeRecLM)) + def test_resolve_pad_token_id(self) -> None: + tok = types.SimpleNamespace + # pad present -> pad + self.assertEqual( + GenerativeRecLM._resolve_pad_token_id(tok(pad_token_id=5, eos_token_id=9)), + 5, + ) + # pad absent -> eos fallback + self.assertEqual( + GenerativeRecLM._resolve_pad_token_id( + tok(pad_token_id=None, eos_token_id=9) + ), + 9, + ) + # neither -> a clear error, not an opaque int(None) TypeError + with self.assertRaisesRegex(ValueError, "neither pad_token_id nor"): + GenerativeRecLM._resolve_pad_token_id( + tok(pad_token_id=None, eos_token_id=None) + ) + def test_backbone_owned_by_family_proto(self) -> None: # the backbone lives on the family message (its architecture), NOT in # the shared common config; it defaults to the canonical Qwen2.5-0.5B. diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index 1b5d6322c..eb3e512c0 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -4,10 +4,10 @@ # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 -"""Qwen2/Qwen2.5 family subclass of ``GenerativeRecLM`` (design §5). +"""Qwen2/Qwen2.5 family subclass of ``GenerativeRecLM``. Selected from the pipeline config by its own oneof entry (the message-type name -resolves directly to this class — no ``class_name`` dispatch):: +resolves directly to this class):: model_config { qwen2_rec_lm { @@ -17,17 +17,13 @@ } This subclass owns the decoder-only-chat implementation: the ChatML prompt -template, the causal-LM splice, and the ``.model``/``.lm_head`` forward -(design §15/§16). The ``GenerativeRecLM`` base owns the architecture-agnostic -plumbing (vocab extension, jagged→row, loss, metrics). - -The splice/forward here are generic to decoder-only families sharing Qwen2's -``.model``/``.lm_head`` layout (Llama/Mistral/Gemma/Phi — design §16), not -Qwen2-specific; only ``QWEN2_TEMPLATE`` is. When a second such family lands, -lift ``_splice_input_ids`` / ``_min_first_non_neg_index`` / ``predict`` (and -the ChatML ``_build_prompt_tokens``) into an intermediate -``DecoderOnlyChatRecLM`` base so each family is just its template. Until then -they live here. +template, the causal-LM splice, and the ``.model``/``.lm_head`` forward. The +``GenerativeRecLM`` base owns the architecture-agnostic plumbing (vocab +extension, jagged->row, loss, metrics). + +The splice/forward are generic to decoder-only families sharing Qwen2's +``.model``/``.lm_head`` layout (Llama/Mistral/Gemma/Phi); only ``QWEN2_TEMPLATE`` +is Qwen2-specific. """ from typing import Any, Dict, List, Optional, Tuple @@ -49,8 +45,7 @@ def _encode_no_special(tokenizer, text: str) -> List[int]: """ return tokenizer.encode(text, add_special_tokens=False) -# Verbatim Qwen2 ChatML fragments. ``default_system_instruction`` matches -# algr/models/qwen2_5/data.py:73 ("default_instruction") bit-for-bit. +# Verbatim Qwen2 ChatML fragments. QWEN2_TEMPLATE = { "system_prefix": "<|im_start|>system\n", "system_suffix": "<|im_end|>\n", @@ -79,23 +74,23 @@ def __init__( ) -> None: super().__init__(model_config, features, labels, sample_weights, **kwargs) common = self._model_config.common - # generation params — read in the subclass because only this family's - # _generate (the inference branch) consumes them. + # generation params, consumed only by this family's _generate. self._num_beams = int(common.num_beams) self._num_return = int(common.num_return_sequences) - # worst-case spliced length (template-aware) and the one-shot warm-up - # latch; the base supplies max_seq_length, batch_size, and the tpl_* - # buffers (built in super().__init__ via _build_prompt_tokens). + # worst-case spliced length for the first-step activation-pool pre-sizing. self._max_total_len = self._compute_max_total_length() self._pool_warmed = False + # CE suffix width. The supervised tail is fixed — [answer(num_levels) | + # asst_suffix | eos] — so _forward_loss slices a CONSTANT number of trailing + # positions (the tail + 1 for the shift-by-one CE alignment) instead of + # recomputing it per step, which cost two GPU->CPU syncs every step. + self._suffix_keep = self._num_levels + self.tpl_asst_suffix.numel() + 2 def _compute_max_total_length(self) -> int: """Full spliced length at the max history (0 if pre-allocation is off). - Mirrors ``_splice_input_ids``: fixed ChatML frame + ``self._max_seq_length`` - history codes (the user-sequence feature's truncation length, supplied by - the base) + the ``num_levels``-code answer (the eos sits inside the - frame). This is the ``T`` the activation pool is pre-sized to. + Fixed ChatML frame + ``self._max_seq_length`` history codes + the + ``num_levels``-code answer: the ``T`` the activation pool is pre-sized to. """ if self._max_seq_length <= 0: return 0 @@ -113,17 +108,16 @@ def _build_prompt_tokens(self, tokenizer, cfg) -> None: """Tokenise the family chat template once; cache as buffers. Composes the proto's optional ``system_instruction`` / - ``user_prefix_text`` / ``user_suffix_text`` (algr's CN prompt - wrappers) with the family's static fragments: + ``user_prefix_text`` / ``user_suffix_text`` with the family's static + fragments:: tpl_system = system_prefix + system_instruction + system_suffix tpl_user_prefix = user_prefix + user_prefix_text tpl_user_suffix = user_suffix_text + user_suffix tpl_asst_prefix / tpl_asst_suffix verbatim from the template - Buffers are non-persistent — they live with the module (move with - ``model.to(...)``) but stay off the state_dict so HF safetensors - round-tripping isn't polluted by TER-only state. + Buffers are non-persistent: they move with ``model.to(...)`` but stay off + the state_dict so HF safetensors round-tripping isn't polluted. """ tpl = type(self).CHAT_TEMPLATE sys_text = cfg.system_instruction or tpl["default_system_instruction"] @@ -141,9 +135,7 @@ def _build_prompt_tokens(self, tokenizer, cfg) -> None: _encode_no_special(tokenizer, frag_str), dtype=torch.long ) self.register_buffer(f"tpl_{slot_name}", ids, persistent=False) - # algr appends eos to BOTH input_ids and labels at train time - # (algr/models/qwen2_5/data.py:46-47) — i.e. the trailing eos is a - # SUPERVISED token. Cache it so the splice can mirror that exactly. + # the trailing eos is a SUPERVISED token; cache it for the splice. self.register_buffer( "tpl_eos", torch.tensor([int(tokenizer.eos_token_id)], dtype=torch.long), @@ -159,23 +151,17 @@ def _splice_input_ids( """Build ``(input_ids, labels, attention_mask)``, each ``(B, T_max)``. Left-padded with ``eos_token_id``. ``attention_mask`` is essential — - without it self-attention would let pad positions pollute real - positions' hidden states. CE is separately protected by ``-100`` - labels at pad slots, but the forward needs the mask too. - - Every answer is exactly ``self._num_levels`` SID codes (one per codebook - level — validated at the data boundary in ``_sid_token_rows``), so the - supervised tail ``[answer | asst_suffix | eos]`` has a FIXED width and, - after left-padding, lands in the SAME columns for every row: ``labels`` - is built in one vectorized assignment (no per-row label loop). - ``input_ids`` still varies per row (the user history length differs). - - ``user_seq_rows`` / ``label_rows`` already hold extended-vocab token ids - on the model device (see ``_sid_token_rows``). - - ``pad_to`` left-extends every row to at least that length (first-step - activation-pool pre-sizing). The supervised tail stays end-aligned, so - labels and the suffix-slice are unchanged — only more masked left-pad. + without it self-attention lets pad positions pollute real positions; CE + is separately protected by ``-100`` labels at pad slots. + + Every answer is exactly ``self._num_levels`` SID codes, so the supervised + tail ``[answer | asst_suffix | eos]`` has a FIXED width and lands in the + same columns for every row after left-padding -> ``labels`` is one + vectorized write. ``input_ids`` still varies per row (history length). + + ``user_seq_rows`` / ``label_rows`` already hold token ids on the model + device (see ``_sid_token_rows``). ``pad_to`` left-extends every row for + first-step pool pre-sizing; the supervised tail stays end-aligned. """ assert len(user_seq_rows) == len(label_rows) A = self._num_levels @@ -203,21 +189,9 @@ def _splice_input_ids( (B, T), self._ignore_index, dtype=torch.long, device=self.device ) labels[:, T - tail : T - tail + A] = torch.stack(label_rows) - labels[:, -1] = self.tpl_eos[0] # algr supervises the trailing eos + labels[:, -1] = self.tpl_eos[0] # supervise the trailing eos return input_ids, labels, attention_mask - @staticmethod - def _min_first_non_neg_index(labels: torch.Tensor) -> int: - """Return the batch-min index of the first non-(-100) label. - - Verbatim port of algr's helper (al_sid/algr/models/qwen2_5/ - modeling_qwen.py:1267-1274) — the smallest position (across rows in - the batch) where the first non-(-100) label appears, used to decide - how many trailing positions to feed into ``lm_head``. - """ - tmp = (labels >= 0).cumsum(dim=-1) - return int((tmp == 1).float().argmax(dim=-1).min().item()) - def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: """Dispatch on the TER inference flag (``set_is_inference`` in main.py). @@ -232,8 +206,7 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: """Branch 1: teacher-forced forward -> suffix-slice -> CE loss.""" - # SID indices -> token ids once, at the data boundary (see - # _sid_token_rows); the splice then just assembles the prompt. + # SID indices -> token ids once at the data boundary (_sid_token_rows). u_rows = self._sid_token_rows( batch.sequence_dense_features[self._input_name], max_codes=self._max_seq_length, # cap to most-recent items (drop oldest) @@ -245,13 +218,10 @@ def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: # One-shot pool pre-sizing: on the FIRST training step, pad the splice to # the worst-case length so the caching allocator reserves its largest - # (B, T_max) activation segments up front — every shorter batch is then - # served from that pool, keeping per-rank reservations uniform (no - # mid-run growth, the source of the cross-rank imbalance). This rides the - # REAL step: the extra positions are attention-masked and labelled -100, - # so loss and gradient are identical to the unpadded batch. Replaces the - # old _warmup_alloc, whose separate unscaled forward+backward corrupted - # the first optimizer step (the 1.14-flat HR regression). + # (B, T_max) activation segments up front, keeping per-rank reservations + # uniform (no mid-run growth). The extra positions are attention-masked + # and labelled -100, so loss and gradient are identical to the unpadded + # batch. pad_to = 0 if not self._pool_warmed and self._max_total_len > 0 and self.is_train: pad_to = self._max_total_len @@ -284,25 +254,17 @@ def _forward_loss( outputs = self.lm.model(input_ids=input_ids, attention_mask=attention_mask) hidden = outputs.last_hidden_state # (B, T, D) - # Suffix slice in BOTH train and eval. algr only slices when - # training (its eval goes through a separate beam-search predict), - # but for CE the slice is value-identical (positions outside the - # suffix all carry -100 labels) and it bounds the logits tensor to - # (B, T_suffix, V) — without it, eval at bsz=80 would materialise - # (B, T, 217k) logits plus HF loss_function's fp32 upcast and OOM. - if (labels >= 0).any(): - keep = labels.shape[1] - self._min_first_non_neg_index(labels) + 1 - sl = slice(-keep, None) - labels_sl = labels[:, sl] - else: - sl = slice(None) - labels_sl = labels - + # Suffix slice in BOTH train and eval. The supervised tail is fixed-width + # (see _splice_input_ids), so slice a CONSTANT number of trailing positions + # — a per-step recompute would cost two GPU->CPU syncs. Outside the suffix + # every label is -100 (CE value-identical), and this bounds the logits to + # (B, suffix, V) instead of (B, T, vocab) + the fp32 upcast (eval at bsz=80 + # would otherwise OOM). + sl = slice(-self._suffix_keep, None) + labels_sl = labels[:, sl] logits = self.lm.lm_head(hidden[:, sl, :]) - # ``loss_function`` is the HF ``ForCausalLMLoss`` callable hung off - # every ``…ForCausalLM`` class; does shift-by-one + CE with -100 - # ignore. Calling it here matches algr's training-step loss exactly. + # HF ForCausalLMLoss: shift-by-one + CE with -100 ignore. loss = self.lm.loss_function( logits=logits, labels=labels_sl, @@ -313,9 +275,10 @@ def _forward_loss( def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: """Branch 2: beam-search the SID answer (no ground truth supplied). - Builds the prompt (no answer), generates exactly ``num_levels`` new - tokens per beam, and maps them back to raw SID indices. Returns - ``generated_sids`` of shape ``(B, num_return, num_levels)``. + Builds the prompt (no answer), generates up to ``num_levels`` new tokens + per beam, and hands the generated tail to the base + ``_validate_sid_candidates`` (token->SID, malformed beams -> ``-1``). + Returns ``generated_sids`` of shape ``(B, num_return, num_levels)``. """ u_rows = self._sid_token_rows( batch.sequence_dense_features[self._input_name], @@ -331,13 +294,8 @@ def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: do_sample=False, pad_token_id=self._pad_token_id, ) - # keep only the generated tail; map token ids back to raw SID indices - # (inverse of _tokenize_sids: sid = token - base_vocab + 1). - new_tokens = out[:, input_ids.shape[1]:] - sids = new_tokens - (self._base_vocab - 1) - # generate() returns rows grouped batch-major: [b0_beam0, b0_beam1, ..., - # b1_beam0, ...], so this view groups beams under the right user. - sids = sids.view(input_ids.shape[0], self._num_return, self._num_levels) + new_tokens = out[:, input_ids.shape[1]:] # the generated tail + sids = self._validate_sid_candidates(new_tokens, input_ids.shape[0]) return {self.GENERATED_SIDS_KEY: sids} def _splice_prompt_ids( diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/qwen2_rec_lm_test.py index 7f77fa765..146c7d03f 100644 --- a/tzrec/models/qwen2_rec_lm_test.py +++ b/tzrec/models/qwen2_rec_lm_test.py @@ -15,14 +15,20 @@ import torch from torch import nn +from tzrec.models.generative_rec_lm_test import _FakeJT from tzrec.models.qwen2_rec_lm import Qwen2RecLM -def _stub(num_levels=3, base_vocab=100, pad_id=9, device="cpu"): +def _stub(num_levels=3, base_vocab=100, pad_id=9, device="cpu", per_level=4): """A Qwen2RecLM with the splice-relevant state wired up, no HF backbone. Template buffers use tiny placeholder ids so the spliced layout is easy to read; real buffers come from ``_build_prompt_tokens`` at init time. + + ``per_level`` sets the (uniform) codebook size used to derive the per-level + SID validity bands the inference gate checks — real bands come from + ``_read_common_config``. With ``per_level=4`` / ``num_levels=3`` the bands + are lo=[1,5,9], hi=[4,8,12] (level j -> sid in [j*4+1, (j+1)*4]). """ m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) @@ -37,9 +43,25 @@ def _stub(num_levels=3, base_vocab=100, pad_id=9, device="cpu"): "tpl_asst_prefix": [14], "tpl_asst_suffix": [15], "tpl_eos": [9], }.items(): m.register_buffer(name, torch.tensor(vals, dtype=torch.long), persistent=False) + lo, hi = Qwen2RecLM._sid_level_bands([per_level] * num_levels) + m.register_buffer("_sid_lvl_lo", lo, persistent=False) + m.register_buffer("_sid_lvl_hi", hi, persistent=False) return m +def _gen_batch(): + """A one-row inference batch (history SIDs [1, 2, 3]) for ``_generate`` tests.""" + return types.SimpleNamespace( + sequence_dense_features={"user_sequence": _FakeJT([1, 2, 3], [3])} + ) + + +def _first_non_neg_index(labels): + """The per-step suffix bound _forward_loss's cached _suffix_keep replaces.""" + tmp = (labels >= 0).cumsum(dim=-1) + return int((tmp == 1).float().argmax(dim=-1).min().item()) + + class Qwen2RecLMTest(unittest.TestCase): def test_splice_layout_and_labels(self) -> None: m = _stub() @@ -82,9 +104,21 @@ def test_mask_keeps_trailing_eos_when_pad_equals_eos(self) -> None: self.assertEqual(int(mask[0, -1]), 1) self.assertEqual(mask[0].tolist(), [1] * ids.shape[1]) - def test_min_first_non_neg_index(self) -> None: - labels = torch.tensor([[-100, -100, 5, 6], [-100, 7, 8, 9]]) - self.assertEqual(Qwen2RecLM._min_first_non_neg_index(labels), 1) + def test_suffix_keep_matches_dynamic_slice(self) -> None: + # _forward_loss caches self._suffix_keep instead of recomputing the suffix + # bound per step (two GPU->CPU syncs); prove the constant equals the old + # dynamic computation and that the slice drops nothing supervised. + m = _stub() # num_levels=3, asst_suffix=[15] (numel 1) -> suffix_keep=6 + suffix_keep = m._num_levels + m.tpl_asst_suffix.numel() + 2 + _, labels, _ = m._splice_input_ids( + [torch.tensor([100, 101, 102])], [torch.tensor([200, 201, 202])] + ) + # the constant matches the per-step bound it replaces + self.assertEqual( + suffix_keep, labels.shape[1] - _first_non_neg_index(labels) + 1 + ) + # everything before the kept suffix is unsupervised (-100): nothing dropped + self.assertTrue(bool((labels[:, :-suffix_keep] < 0).all())) def test_splice_prompt_ids(self) -> None: m = _stub() @@ -110,22 +144,62 @@ def test_generate_maps_tokens_to_sids(self) -> None: def fake_generate(input_ids, attention_mask, max_new_tokens, num_beams, num_return_sequences, do_sample, pad_token_id): prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) - new = torch.tensor([[200, 201, 202], [203, 204, 205]]) # 2 beams x 3 codes + # 2 beams x 3 codes, every atom INSIDE its level's band + # (bands lo=[1,5,9] hi=[4,8,12]; token = sid + 99): + # pos0 token in [100,103], pos1 in [104,107], pos2 in [108,111]. + new = torch.tensor([[100, 104, 108], [103, 107, 111]]) + return torch.cat([prompt, new], dim=1) + + m.lm.generate = fake_generate + sids = m._generate(_gen_batch())["generated_sids"] + self.assertEqual(tuple(sids.shape), (1, 2, 3)) # (B, num_return, num_levels) + self.assertEqual(sids[0].tolist(), [[1, 5, 9], [4, 8, 12]]) + + def test_generate_rejects_malformed_candidates(self) -> None: + # Layer-A gate: every malformed candidate -> the -1 sentinel, in place. + # bands lo=[1,5,9] hi=[4,8,12]; token = sid + 99. + m = _stub(base_vocab=100) + m._input_name = "user_sequence" + m._num_beams = m._num_return = 4 + + def fake_generate(input_ids, attention_mask, max_new_tokens, + num_beams, num_return_sequences, do_sample, pad_token_id): + prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) + new = torch.tensor([ + [100, 104, 108], # all in-band -> valid -> [1, 5, 9] + [100, 104, 9], # pos2 = eos/pad token (sid -90) -> invalid + [108, 104, 100], # wrong-level scramble (pos0 = lvl-2 code) -> invalid + [100, 104, 112], # pos2 sid 13 > band hi 12 -> invalid + ]) return torch.cat([prompt, new], dim=1) m.lm.generate = fake_generate + sids = m._generate(_gen_batch())["generated_sids"] + self.assertEqual(tuple(sids.shape), (1, 4, 3)) + # valid candidate kept at its rank; every malformed one -> all -1 (in place) + self.assertEqual( + sids[0].tolist(), + [[1, 5, 9], [-1, -1, -1], [-1, -1, -1], [-1, -1, -1]], + ) - class _JT: - def values(self): - return torch.tensor([1, 2, 3], dtype=torch.float) + def test_generate_narrow_tail_no_crash(self) -> None: + # every beam emits EOS before num_levels -> generate() returns a tail + # narrower than num_levels; the canvas keeps the reshape rectangular. + m = _stub(base_vocab=100) + m._input_name = "user_sequence" + m._num_beams = m._num_return = 2 - def lengths(self): - return torch.tensor([3]) + def fake_generate(input_ids, attention_mask, max_new_tokens, + num_beams, num_return_sequences, do_sample, pad_token_id): + prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) + new = torch.tensor([[100, 104], [103, 107]]) # width 2 < num_levels 3 + return torch.cat([prompt, new], dim=1) - batch = types.SimpleNamespace(sequence_dense_features={"user_sequence": _JT()}) - sids = m._generate(batch)["generated_sids"] - self.assertEqual(tuple(sids.shape), (1, 2, 3)) # (B, num_return, num_levels) - self.assertEqual(sids[0].tolist(), [[101, 102, 103], [104, 105, 106]]) + m.lm.generate = fake_generate + sids = m._generate(_gen_batch())["generated_sids"] + self.assertEqual(tuple(sids.shape), (1, 2, 3)) # rectangular, no crash + # the missing 3rd atom stays -1 -> out of band -> whole candidate -1 + self.assertEqual(sids[0].tolist(), [[-1, -1, -1], [-1, -1, -1]]) def test_build_prompt_tokens_registers_buffers(self) -> None: m = object.__new__(Qwen2RecLM) diff --git a/tzrec/utils/checkpoint_util.py b/tzrec/utils/checkpoint_util.py index a11f3b842..ba074da0f 100644 --- a/tzrec/utils/checkpoint_util.py +++ b/tzrec/utils/checkpoint_util.py @@ -335,7 +335,17 @@ def save( # Local import avoids a circular import (export_util imports us). from tzrec.utils.export_util import write_hf_assets - write_hf_assets(model, ckpt_dir) + # HF assets are convenience metadata; the DCP weights (save_model + # above) are already durable. Isolate a failure so it can't abort the + # save before the next collective (save_dataloader_state's all_gather) + # and one-sidedly hang the other ranks. + try: + write_hf_assets(model, ckpt_dir) + except Exception as e: # noqa: BLE001 + logger.warning( + f"write_hf_assets failed for {ckpt_dir}: {e} — checkpoint " + f"weights are saved; skipping HF assets." + ) if dataloader_state is not None: save_dataloader_state(ckpt_dir, dataloader_state) self.prune() From 97390b2104a842c01ea257725f0167723c29f942 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 17 Jun 2026 06:48:01 +0000 Subject: [PATCH 15/99] [refactor] _validate_sid_candidates: F.pad + masked_fill MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the explicit canvas alloc + min/slice-copy + boolean-index assignment with F.pad (fixed-width -1 padding) and masked_fill (row invalidation): 6 logic lines -> 4, no in-place indexing, clearer intent. Behavior is unchanged — verified value-identical to the previous form on 4 edge cases (valid, narrow tails w=1/2, all-invalid) + 2000 random fuzz batches. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 9a263779b..e2f5e46c8 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -30,6 +30,7 @@ from typing import Any, Dict, List, Optional import torch +import torch.nn.functional as F import torchmetrics from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer @@ -301,19 +302,21 @@ def _validate_sid_candidates( when beams stop early). Returns ``(batch_size, num_return, num_levels)`` raw SIDs with every MALFORMED candidate set to the ``-1`` sentinel — early EOS, a non-SID token, or a wrong-level atom (outside THAT level's band). ``-1`` - can never match a real item, and the fixed-width canvas keeps the reshape + can never match a real item, and the fixed-width padding keeps the reshape rectangular even when every beam stopped early. """ sids = new_tokens - (self._base_vocab - 1) - canvas = sids.new_full((sids.shape[0], self._num_levels), -1) - w = min(sids.shape[1], self._num_levels) - canvas[:, :w] = sids[:, :w] - in_band = (canvas >= self._sid_lvl_lo) & (canvas <= self._sid_lvl_hi) - canvas[~in_band.all(dim=1)] = -1 + # pad each candidate to exactly num_levels with the -1 sentinel (an + # early-EOS beam returns fewer tokens) so the reshape stays rectangular. + sids = F.pad(sids, (0, self._num_levels - sids.shape[1]), value=-1) + # valid only if every position-j atom is in level j's band; any violation + # invalidates the WHOLE candidate -> -1. + invalid = ((sids < self._sid_lvl_lo) | (sids > self._sid_lvl_hi)).any(dim=1) + sids = sids.masked_fill(invalid.unsqueeze(1), -1) # view groups beams under the right user (generate() returns rows - # batch-major: [b0_beam0, b0_beam1, ..., b1_beam0, ...]); the in-place mask - # above preserves that order (no filter/scatter that would scramble it). - return canvas.view(batch_size, -1, self._num_levels) + # batch-major: [b0_beam0, b0_beam1, ..., b1_beam0, ...]); the row-wise mask + # preserves that order (no filter/scatter that would scramble it). + return sids.view(batch_size, -1, self._num_levels) def _sid_token_rows( self, From 17d4350216441876ac113e53de4476957fce8f23 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 17 Jun 2026 06:53:05 +0000 Subject: [PATCH 16/99] [refactor] generative-rec LM: trim over-detailed comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Condense the multi-line inline rationale comments (CE-suffix width, one-shot pool pre-sizing, fp32-master build, SID-band gate, beam-order reshape) and the two longest docstrings (init_from_pretrained, _validate_sid_candidates) to terse 1-3 liners, keeping the load-bearing why (fp32-master underflow, -1 can't match a real item, suffix-slice OOM, pad==eos mask). Comments only — no behavior change; ruff clean, 25 genrec tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 56 ++++++++++++------------------- tzrec/models/qwen2_rec_lm.py | 30 ++++++----------- 2 files changed, 31 insertions(+), 55 deletions(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index e2f5e46c8..18002408f 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -93,10 +93,8 @@ def __init__( self._build_prompt_tokens(tokenizer, cfg) - # Dense-only: the HF backbone owns its embeddings and SID token ids flow - # through it directly (the SEQUENCE feature is consumed as raw token ids in - # _sid_token_rows, NOT via an EmbeddingGroup), so there is no native sparse - # path. Set explicitly (vs RankModel.init_input) to keep the contract clear. + # Dense-only: the HF backbone owns its embeddings and SID ids flow through + # it directly (no EmbeddingGroup). Set explicit (vs RankModel.init_input). self.embedding_group = None # one-shot debug dump of the first spliced batch @@ -107,8 +105,7 @@ def __init__( def _resolve_pad_token_id(tokenizer: Any) -> int: """Pad id for the left-padded splice, falling back to eos. - Qwen2.5 has both, but the base is reused by Llama/Mistral/Gemma/Phi — fail - loudly if a tokenizer has neither rather than an opaque ``int(None)``. + Fail loudly if a tokenizer has neither (vs an opaque ``int(None)``). """ pad_id = tokenizer.pad_token_id if pad_id is None: @@ -133,9 +130,8 @@ def _read_common_config(self, common: Any) -> int: raise ValueError("GenerativeRecLM: codebook must be non-empty.") # len(codebook) = SID codes per item (answer width); sum = vocab atoms. self._num_levels = len(codebook) - # Per-level SID validity bands, cached as (num_levels,) buffers so the - # inference gate (_validate_sid_candidates) can reject malformed beams with - # a vectorized band check and no host sync. Non-persistent: derived config. + # Per-level SID validity bands (buffers) for the inference gate + # (_validate_sid_candidates). Non-persistent: derived config. lo, hi = self._sid_level_bands(codebook) self.register_buffer("_sid_lvl_lo", lo, persistent=False) self.register_buffer("_sid_lvl_hi", hi, persistent=False) @@ -155,10 +151,8 @@ def _build_backbone(self) -> Any: f"{type(self).__name__}: empty backbone id (see _backbone_id)." ) hf_cfg = AutoConfig.from_pretrained(hf_model_id) - # Build in fp32 so the optimizer keeps fp32 MASTER weights: with bf16 - # params, Adam's small updates (e.g. lr=1e-5) fall below the bf16 ULP and - # round to zero, freezing the weights. Use mixed_precision:"BF16" for bf16 - # compute speed; the DCP checkpoint stays fp32 (consistent on restore). + # fp32 MASTER weights: bf16 params underflow Adam's small (lr=1e-5) updates + # and freeze. bf16 compute comes from mixed_precision:"BF16"; ckpt stays fp32. lm = AutoModelForCausalLM.from_config(hf_cfg, torch_dtype=self._PARAM_DTYPE) if next(lm.parameters()).dtype != self._PARAM_DTYPE: lm = lm.to(self._PARAM_DTYPE) @@ -208,16 +202,13 @@ def _backbone_id(self) -> str: def init_from_pretrained(self) -> None: """Load the pretrained HF backbone weights into ``self.lm``. - The SINGLE place ``AutoModelForCausalLM.from_pretrained`` runs. The - PIPELINE (``tzrec/main.py``) calls this exactly once, only at COLD START - (no checkpoint to resume). On resume/eval/export the empty arch built in - ``__init__`` is weight-filled by native DCP ``load_state_dict`` instead, - so the GB-scale download is skipped. Re-extends the vocab to the SAME - target/pad as ``__init__`` so the module shapes stay identical. + The single ``from_pretrained`` call, run once at COLD START (no checkpoint + to resume); on resume/eval/export the empty ``__init__`` arch is filled by + DCP instead, skipping the download. Re-extends the vocab to ``__init__``'s + target so the shapes match. """ - # Load fp32 master weights, NOT torch_dtype="auto" (which keeps the stored - # bf16): bf16 params underflow Adam's lr=1e-5 updates. Must match - # _build_backbone's fp32 so the cold-start and restore arches agree. + # fp32 master (not "auto", which keeps the stored bf16); must match + # _build_backbone so cold-start and restore arches agree. lm = AutoModelForCausalLM.from_pretrained( self._backbone_id(), torch_dtype=self._PARAM_DTYPE ) @@ -251,9 +242,8 @@ def hf_tokenizer(self): def _build_prompt_tokens(self, tokenizer, cfg) -> None: """Family hook: cache the tokenised prompt template as buffers. - Called from ``__init__`` after vocab extension; the buffers it - registers are consumed by the family's ``predict``. Architecture- - specific — see design §15.1/§15.2. Subclasses MUST implement this. + Called from ``__init__`` after vocab extension; consumed by the family's + ``predict``. Subclasses MUST implement this. """ raise NotImplementedError( f"{type(self).__name__} must implement _build_prompt_tokens " @@ -297,13 +287,10 @@ def _validate_sid_candidates( ) -> torch.Tensor: """Map a generated token tail back to SIDs and reject malformed candidates. - Inference-side counterpart of ``_tokenize_sids``. ``new_tokens`` is the - per-beam generated tail ``(B*num_return, w)`` (``w`` may be < ``num_levels`` - when beams stop early). Returns ``(batch_size, num_return, num_levels)`` raw - SIDs with every MALFORMED candidate set to the ``-1`` sentinel — early EOS, - a non-SID token, or a wrong-level atom (outside THAT level's band). ``-1`` - can never match a real item, and the fixed-width padding keeps the reshape - rectangular even when every beam stopped early. + ``new_tokens`` is the per-beam tail ``(B*num_return, w)`` (``w`` may be < + ``num_levels`` when beams stop early). Returns ``(batch_size, num_return, + num_levels)`` SIDs with every malformed candidate (early EOS / non-SID / + wrong-level atom) set to ``-1`` — which can never match a real item. """ sids = new_tokens - (self._base_vocab - 1) # pad each candidate to exactly num_levels with the -1 sentinel (an @@ -313,9 +300,8 @@ def _validate_sid_candidates( # invalidates the WHOLE candidate -> -1. invalid = ((sids < self._sid_lvl_lo) | (sids > self._sid_lvl_hi)).any(dim=1) sids = sids.masked_fill(invalid.unsqueeze(1), -1) - # view groups beams under the right user (generate() returns rows - # batch-major: [b0_beam0, b0_beam1, ..., b1_beam0, ...]); the row-wise mask - # preserves that order (no filter/scatter that would scramble it). + # generate() returns rows batch-major ([b0_beam0, b0_beam1, ...]) so this + # groups beams per user; the row-wise mask above preserved that order. return sids.view(batch_size, -1, self._num_levels) def _sid_token_rows( diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index eb3e512c0..b06b22b14 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -80,10 +80,8 @@ def __init__( # worst-case spliced length for the first-step activation-pool pre-sizing. self._max_total_len = self._compute_max_total_length() self._pool_warmed = False - # CE suffix width. The supervised tail is fixed — [answer(num_levels) | - # asst_suffix | eos] — so _forward_loss slices a CONSTANT number of trailing - # positions (the tail + 1 for the shift-by-one CE alignment) instead of - # recomputing it per step, which cost two GPU->CPU syncs every step. + # CE suffix width: the supervised tail [answer | asst_suffix | eos] is + # fixed, so _forward_loss slices a constant suffix (no per-step sync). self._suffix_keep = self._num_levels + self.tpl_asst_suffix.numel() + 2 def _compute_max_total_length(self) -> int: @@ -178,10 +176,8 @@ def _splice_input_ids( ] input_ids, attention_mask = self._left_pad(rows_ids, pad_to=pad_to) - # labels: the supervised tail is fixed-width, so left-padding aligns it - # to the same columns for every row -> one vectorized write. - # tail layout (from the end): [answer(A) | asst_suffix(s) | eos(1)]. - # ``tail <= T`` always holds: every row already contains those tokens. + # supervised tail is fixed-width -> same columns every row -> one write. + # tail from the end: [answer(A) | asst_suffix(s) | eos(1)]. B, T = input_ids.shape s = self.tpl_asst_suffix.numel() tail = A + s + 1 @@ -216,12 +212,9 @@ def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: expected_width=self._num_levels, # answer = one item = num_levels codes ) - # One-shot pool pre-sizing: on the FIRST training step, pad the splice to - # the worst-case length so the caching allocator reserves its largest - # (B, T_max) activation segments up front, keeping per-rank reservations - # uniform (no mid-run growth). The extra positions are attention-masked - # and labelled -100, so loss and gradient are identical to the unpadded - # batch. + # One-shot pool pre-sizing: pad the FIRST train step to the worst-case + # length so the allocator reserves its largest segments up front (no + # mid-run growth). Extra positions are masked + -100 -> loss/grad unchanged. pad_to = 0 if not self._pool_warmed and self._max_total_len > 0 and self.is_train: pad_to = self._max_total_len @@ -254,12 +247,9 @@ def _forward_loss( outputs = self.lm.model(input_ids=input_ids, attention_mask=attention_mask) hidden = outputs.last_hidden_state # (B, T, D) - # Suffix slice in BOTH train and eval. The supervised tail is fixed-width - # (see _splice_input_ids), so slice a CONSTANT number of trailing positions - # — a per-step recompute would cost two GPU->CPU syncs. Outside the suffix - # every label is -100 (CE value-identical), and this bounds the logits to - # (B, suffix, V) instead of (B, T, vocab) + the fp32 upcast (eval at bsz=80 - # would otherwise OOM). + # Slice the fixed-width supervised suffix (constant -> no per-step sync). + # Outside it every label is -100 (CE unchanged); it also bounds the logits + # to (B, suffix, V) — the full (B, T, vocab) + fp32 upcast would OOM. sl = slice(-self._suffix_keep, None) labels_sl = labels[:, sl] logits = self.lm.lm_head(hidden[:, sl, :]) From 187122751c842aab5c3a1760f3a8aa0a3cfdf769 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 17 Jun 2026 06:54:14 +0000 Subject: [PATCH 17/99] [refactor] generative-rec LM: extract _detokenize_sids The SID->token offset now has both directions as named helpers next to each other: _tokenize_sids (sid -> token) and its inverse _detokenize_sids (token -> sid), used by _validate_sid_candidates instead of the inline `new_tokens - (base_vocab - 1)`. One owner for the offset constant. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 18002408f..77194df77 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -264,6 +264,10 @@ def _tokenize_sids(self, sids: torch.Tensor) -> torch.Tensor: """ return sids + (self._base_vocab - 1) + def _detokenize_sids(self, tokens: torch.Tensor) -> torch.Tensor: + """Inverse of ``_tokenize_sids``: token id -> raw 1-indexed SID.""" + return tokens - (self._base_vocab - 1) + @staticmethod def _sid_level_bands(codebook: Any) -> tuple[torch.Tensor, torch.Tensor]: """Per-level closed SID bands ``(lo, hi)`` as ``(num_levels,)`` long tensors. @@ -292,7 +296,7 @@ def _validate_sid_candidates( num_levels)`` SIDs with every malformed candidate (early EOS / non-SID / wrong-level atom) set to ``-1`` — which can never match a real item. """ - sids = new_tokens - (self._base_vocab - 1) + sids = self._detokenize_sids(new_tokens) # pad each candidate to exactly num_levels with the -1 sentinel (an # early-EOS beam returns fewer tokens) so the reshape stays rectangular. sids = F.pad(sids, (0, self._num_levels - sids.shape[1]), value=-1) From 59312bf7253ef05ca0d9fcd1fe447b68a20e7418 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 17 Jun 2026 07:16:53 +0000 Subject: [PATCH 18/99] [docs] generative-rec LM: fix misleading "reader caps" comments The data reader/DataParser does NOT truncate sequences under FG_NONE (only pyfg does, in FG_NORMAL/FG_DAG; the native EmbeddingGroup does via to_padded_dense at forward time). GenerativeRecLM uses FG_NONE and no EmbeddingGroup, so the history cap is enforced model-side in _sid_token_rows (item-aligned to whole num_levels items). Correct _read_common_config + _input_sequence_length to say so instead of "the reader caps every row". Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 77194df77..306afb01f 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -123,7 +123,8 @@ def _read_common_config(self, common: Any) -> int: self._label_name: str = common.label_feature_name self._ignore_index: int = int(common.ignore_index) # max history (SID codes) for activation pre-sizing = the user-sequence - # feature's truncation length (the reader caps every row at it). 0 = off. + # feature's sequence_length. FG_NONE doesn't truncate, so _sid_token_rows + # enforces this cap (item-aligned) model-side. 0 = off. self._max_seq_length: int = self._input_sequence_length() codebook = list(common.codebook) if len(codebook) == 0: @@ -218,13 +219,12 @@ def init_from_pretrained(self) -> None: self.lm = lm def _input_sequence_length(self) -> int: - """Truncation length (SID codes) of the user-sequence feature. + """The user-sequence feature's ``sequence_length`` (SID codes), or 0. - The data reader caps every row's history at the feature's - ``sequence_length``, so it is the guaranteed upper bound used to - pre-size the activation pool (see ``Qwen2RecLM._predict_train``'s - first-step padding). Returns 0 if the feature has no length cap, which - disables pre-allocation. + FG_NONE does NOT truncate, so this is the cap ``_sid_token_rows`` enforces + model-side (item-aligned), and the upper bound the activation pool is + pre-sized to (see ``Qwen2RecLM._predict_train``). 0 if the feature has no + length cap, which disables pre-allocation. """ for feature in self._features: if feature.config.feature_name == self._input_name: From 0487481a85e316b1bdfcc1f6d27e687e17e23103 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 18 Jun 2026 02:52:52 +0000 Subject: [PATCH 19/99] [feat] genrec LM: ALGR-style escalating-beam decode (dynamic_beam) Port ALGR's dynamic_beams schedule (beam width doubles per SID level 50->100->200->400, returns num_beams*2**num_levels candidates) as a torch-only KV-cached kernel; faithful to ALGR's escalating beam and exploits the fixed-length, EOS-free SID answer. - tzrec/models/escalating_beam.py: escalating_beam_search kernel - qwen2_rec_lm.py: _dynamic_beam_search delegates; _generate dispatch on the dynamic_beam flag - generative_model.proto: dynamic_beam flag - examples/generative_rec_lm_predict.py: --dynamic_beam / --codebook - tests: exhaustive==brute-force top-k + validity/left-pad (19 pass) Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/generative_rec_lm_predict.py | 241 +++++++++++++++++++++ tzrec/models/escalating_beam.py | 121 +++++++++++ tzrec/models/qwen2_rec_lm.py | 87 ++++++-- tzrec/models/qwen2_rec_lm_test.py | 163 ++++++++++++-- tzrec/protos/models/generative_model.proto | 6 + 5 files changed, 579 insertions(+), 39 deletions(-) create mode 100644 examples/generative_rec_lm_predict.py create mode 100644 tzrec/models/escalating_beam.py diff --git a/examples/generative_rec_lm_predict.py b/examples/generative_rec_lm_predict.py new file mode 100644 index 000000000..e694257ac --- /dev/null +++ b/examples/generative_rec_lm_predict.py @@ -0,0 +1,241 @@ +"""Beam-search SID generation for an exported `GenerativeRecLM` checkpoint. + +Mirrors algr's predict flow (config/qwen_predict_tiny_*.json: beam=50, +num_return=50, max_new_tokens=3, greedy) and emits `output_sids.jsonl` +lines in the exact shape `calc_hr_fast.py` consumes:: + + {"_generated_text_": ["C{a}C{b}C{c}", ... x num_return], "answer": "item1;item2"} + +The generated C-codes carry the dataset's layer offsets verbatim +(lv2 ∈ [8192, 16383], lv3 ∈ [16384, ...]) because the model is trained on +offset codes — decoding is just ``token_id - sid_base``. + +Prompts are rebuilt with the SAME splice as training, minus the answer: +``system + user_prefix + SID atoms + user_suffix + asst_prefix``, +left-padded with eos. The SID base is read from the exported tokenizer +(``convert_tokens_to_ids("C0")``), so tokenizer-layout drift is impossible. + +Run sharded across GPUs (one process per GPU, simple row interleave):: + + PYTHONPATH=. python -m examples.generative_rec_lm_predict \\ + --export_dir experiments//export_hf_40000 \\ + --test_parquet_dir /home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data_genreclm \\ + --test_csv /home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data/s1_tiny_test.csv \\ + --out logs/ter_predict_40000/output_sids.rank0.jsonl \\ + --rank 0 --world_size 8 +""" + +from __future__ import annotations + +import argparse +import csv +import glob +import json +import os +import sys +import time +from typing import List + +import pyarrow.parquet as pq +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from tzrec.models.escalating_beam import escalating_beam_search + +csv.field_size_limit(10 * 1024 * 1024) + + +def _enc(tok, text: str) -> List[int]: + return tok.encode(text, add_special_tokens=False) + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--export_dir", required=True) + ap.add_argument("--test_parquet_dir", required=True) + ap.add_argument( + "--test_csv", + required=True, + help="raw test CSV; supplies the ground-truth `answer` item-id strings", + ) + ap.add_argument("--out", required=True) + ap.add_argument("--rank", type=int, default=0) + ap.add_argument("--world_size", type=int, default=1) + ap.add_argument( + "--bsz", + type=int, + default=4, + help="algr per_device_eval_batch_size=4; 8 OOMs at beam 50", + ) + ap.add_argument("--num_beams", type=int, default=50) + ap.add_argument("--num_return", type=int, default=50) + ap.add_argument("--max_new_tokens", type=int, default=3) + ap.add_argument( + "--dynamic_beam", + action="store_true", + help="ALGR-style escalating beam (width doubles per level); " + "returns num_beams * 2**max_new_tokens candidates, " + "ignoring --num_return.", + ) + ap.add_argument( + "--codebook", + type=int, + default=8192, + help="per-level SID codebook size (for --dynamic_beam band masking)", + ) + ap.add_argument( + "--max_rows", type=int, default=0, help="cap rows (0 = all); for smoke tests" + ) + # algr CN prompt fragments (must match the training config) + ap.add_argument( + "--system-instruction", + default=( + "你是一个推荐系统，根据用户的历史行为，预测用户在电商场景的下一步行为。" + "我会给你一串连续行为的语义编码，按照用户点击的时间顺序排列，每个行为用三个词表示。" + ), + ) + ap.add_argument("--user-prefix-text", default="当前用户的历史行为如下：") + ap.add_argument( + "--user-suffix-text", default="，请预测用户在电商推荐场景后续行为的语义编码" + ) + args = ap.parse_args() + + device = "cuda" if torch.cuda.is_available() else "cpu" + tok = AutoTokenizer.from_pretrained(args.export_dir, use_fast=True) + base = tok.convert_tokens_to_ids("C0") + assert base is not None and base > 0, "exported tokenizer lacks C atoms" + eos = tok.eos_token_id + + model = ( + AutoModelForCausalLM.from_pretrained(args.export_dir, torch_dtype="auto") + .to(device) + .eval() + ) + print( + f"[predict] model loaded vocab={model.config.vocab_size} sid_base={base} dev={device}", + flush=True, + ) + + # Per-level token-space SID bands (level j atom in [base+j*cb, base+(j+1)*cb-1]). + # Only needed by --dynamic_beam; the dataset offset codes are already disjoint. + L, cb = args.max_new_tokens, args.codebook + lo_tok = torch.tensor([base + j * cb for j in range(L)], device=device) + hi_tok = torch.tensor([base + (j + 1) * cb - 1 for j in range(L)], device=device) + if args.dynamic_beam: + print( + f"[predict] dynamic beam: widths " + f"{[args.num_beams * 2 ** (j + 1) for j in range(L)]} " + f"-> {args.num_beams * 2**L} candidates/row", + flush=True, + ) + + # Cached template fragments — identical composition to training splice. + tpl_system = _enc(tok, f"<|im_start|>system\n{args.system_instruction}<|im_end|>\n") + tpl_user_prefix = _enc(tok, f"<|im_start|>user\n{args.user_prefix_text}") + tpl_user_suffix = _enc(tok, f"{args.user_suffix_text}<|im_end|>\n") + tpl_asst_prefix = _enc(tok, "<|im_start|>assistant\n") + + # Rows: user_sequence SIDs from parquet, answer strings from the CSV + # (the converter wrote every CSV row, so ordering is 1:1). + paths = sorted(glob.glob(os.path.join(args.test_parquet_dir, "*.parquet"))) + user_rows: List[List[int]] = [] + for p in paths: + user_rows.extend( + pq.read_table(p, columns=["user_sequence"]) + .column("user_sequence") + .to_pylist() + ) + answers: List[str] = [] + with open(args.test_csv, encoding="utf-8", newline="") as f: + for row in csv.DictReader(f): + answers.append((row.get("answer") or "").strip()) + assert len(user_rows) == len(answers), (len(user_rows), len(answers)) + + idxs = list(range(len(user_rows)))[args.rank :: args.world_size] + if args.max_rows: + idxs = idxs[: args.max_rows] + print(f"[predict] rank={args.rank}/{args.world_size} rows={len(idxs)}", flush=True) + + os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True) + t0 = time.time() + n_done = 0 + with open(args.out, "w", encoding="utf-8") as out_f: + for bstart in range(0, len(idxs), args.bsz): + bidx = idxs[bstart : bstart + args.bsz] + prompts = [] + for i in bidx: + u_tok = [s + base - 1 for s in user_rows[i]] + prompts.append( + tpl_system + + tpl_user_prefix + + u_tok + + tpl_user_suffix + + tpl_asst_prefix + ) + T = max(len(p) for p in prompts) + input_ids = torch.full((len(prompts), T), eos, dtype=torch.long) + attn = torch.zeros((len(prompts), T), dtype=torch.long) + for r, p in enumerate(prompts): # left pad + input_ids[r, -len(p) :] = torch.tensor(p, dtype=torch.long) + attn[r, -len(p) :] = 1 + input_ids, attn = input_ids.to(device), attn.to(device) + + with torch.no_grad(): + if args.dynamic_beam: + new_tokens = escalating_beam_search( + model, + input_ids, + attn, + num_beams=args.num_beams, + lo_tok=lo_tok, + hi_tok=hi_tok, + ) # (B * num_beams*2**L, L) + else: + gen = model.generate( + input_ids=input_ids, + attention_mask=attn, + max_new_tokens=args.max_new_tokens, + num_beams=args.num_beams, + num_return_sequences=args.num_return, + do_sample=False, + pad_token_id=eos, + ) + new_tokens = gen[:, T:] # (B * num_return, max_new_tokens) + nret = new_tokens.shape[0] // len(bidx) + new_tokens = new_tokens.view(len(bidx), nret, -1).tolist() + + for row_i, seqs in zip(bidx, new_tokens): + texts = [] + for seq in seqs: + parts = [] + for t in seq: + if t >= base: + parts.append(f"C{t - base}") + else: + parts.append(tok.decode([t], skip_special_tokens=True)) + texts.append("".join(parts)) + out_f.write( + json.dumps( + {"_generated_text_": texts, "answer": answers[row_i]}, + ensure_ascii=False, + ) + + "\n" + ) + n_done += len(bidx) + if (bstart // args.bsz) % 25 == 0: + rate = n_done / max(time.time() - t0, 1e-6) + print( + f"[predict] {n_done}/{len(idxs)} rows ({rate:.1f} rows/s)", + flush=True, + ) + + print( + f"[predict] DONE rank={args.rank} rows={n_done} wall={time.time() - t0:.0f}s " + f"out={args.out}", + flush=True, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tzrec/models/escalating_beam.py b/tzrec/models/escalating_beam.py new file mode 100644 index 000000000..45d816938 --- /dev/null +++ b/tzrec/models/escalating_beam.py @@ -0,0 +1,121 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""ALGR-style escalating-beam SID decode (torch-only, no tzrec deps). + +A faithful port of ALGR's ``dynamic_beams`` schedule: the beam width doubles at +every SID level (``num_beams`` -> ``2*num_beams`` -> ...), keeping +``num_beams * 2**(j+1)`` candidates after level ``j`` and returning +``num_beams * 2**num_levels`` per row. The aggressive early pruning (only the +top ``2*num_beams`` level-0 prefixes survive) is what distinguishes it from a +fixed-width beam that keeps every level-0 code. + +Lives in its own torch-only module so both the production path +(``Qwen2RecLM._dynamic_beam_search``) and the offline predict harness share one +tested implementation. +""" + +from __future__ import annotations + +import torch + + +@torch.no_grad() +def escalating_beam_search( + model, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + *, + num_beams: int, + lo_tok: torch.Tensor, + hi_tok: torch.Tensor, +) -> torch.Tensor: + """Decode SID answers with the escalating beam. + + Args: + model: an HF causal LM exposing ``.model`` / ``.lm_head`` (Qwen2 layout). + input_ids: left-padded prompt ids ``(B, P)``. + attention_mask: prompt mask ``(B, P)``. + num_beams: base beam width; doubles per level. + lo_tok: inclusive lower per-level token-space band edge, ``(num_levels,)`` + (``num_levels`` is inferred from its length). + hi_tok: inclusive upper per-level token-space band edge, ``(num_levels,)``. + + Returns: + The generated SID token tail ``(B * num_beams * 2**num_levels, + num_levels)``, score-ordered best-first per row. The SID answer is + exactly ``num_levels`` codes with no in-answer EOS, so every beam emits + exactly ``num_levels`` tokens — no finished-beam bookkeeping is needed, + and band masking makes every candidate well-formed by construction. + """ + device = input_ids.device + bsz = input_ids.shape[0] + num_levels = lo_tok.shape[0] + # candidates kept after level j (doubling), capped to what the band + the + # surviving prefixes can actually supply (guards tiny codebooks). + widths, prev = [], 1 + for j in range(num_levels): + avail = prev * int(hi_tok[j] - lo_tok[j] + 1) + widths.append(min(num_beams * (2 ** (j + 1)), avail)) + prev = widths[-1] + + def _band_logp(logits: torch.Tensor, j: int) -> torch.Tensor: + ids = torch.arange(logits.shape[-1], device=device) + keep = (ids >= lo_tok[j]) & (ids <= hi_tok[j]) + logp = torch.log_softmax(logits.float(), dim=-1) + return logp.masked_fill(~keep, float("-inf")) + + # 1. prompt forward (bsz beams) -> level-0 logits. + pos = (attention_mask.long().cumsum(-1) - 1).clamp(min=0) + h = model.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=pos, + use_cache=True, + ) + past = h.past_key_values + vocab = model.config.vocab_size + scores = _band_logp(model.lm_head(h.last_hidden_state[:, -1, :]), 0) + beam_scores, tok = scores.topk(widths[0], dim=-1) # (B, W0) + seq = tok.reshape(-1, 1) + beam_scores = beam_scores.reshape(-1) + parent = torch.arange(bsz, device=device).repeat_interleave(widths[0]) + past.reorder_cache(parent) + am = attention_mask.repeat_interleave(widths[0], dim=0) + cur_w = widths[0] + + # 2. levels 1..n-1: forward the last chosen atom with cache, escalate width. + for j in range(1, num_levels): + am = torch.cat([am, am.new_ones(bsz * cur_w, 1)], dim=1) + step_pos = (am.long().cumsum(-1) - 1)[:, -1:].clamp(min=0) + cache_pos = torch.tensor([past.get_seq_length()], device=device) + h = model.model( + input_ids=seq[:, -1:], + attention_mask=am, + position_ids=step_pos, + past_key_values=past, + use_cache=True, + cache_position=cache_pos, + ) + scores = _band_logp(model.lm_head(h.last_hidden_state[:, -1, :]), j) + scores = scores + beam_scores[:, None] # (B*cur_w, V) cumulative + # global top-widths[j] per row over the (cur_w * V) continuations. + beam_scores, idx = scores.view(bsz, cur_w * vocab).topk(widths[j], dim=-1) + parent_local = torch.div(idx, vocab, rounding_mode="floor") # in [0,cur_w) + tok = idx % vocab + row_base = torch.arange(bsz, device=device)[:, None] * cur_w + parent = (parent_local + row_base).reshape(-1) + past.reorder_cache(parent) + am = am[parent] + seq = torch.cat([seq[parent], tok.reshape(-1, 1)], dim=1) + beam_scores = beam_scores.reshape(-1) + cur_w = widths[j] + return seq # (B*cur_w, num_levels) diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index b06b22b14..37f908526 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -33,6 +33,7 @@ from tzrec.datasets.utils import Batch from tzrec.features.feature import BaseFeature +from tzrec.models.escalating_beam import escalating_beam_search from tzrec.models.generative_rec_lm import GenerativeRecLM from tzrec.protos.model_pb2 import ModelConfig @@ -45,6 +46,7 @@ def _encode_no_special(tokenizer, text: str) -> List[int]: """ return tokenizer.encode(text, add_special_tokens=False) + # Verbatim Qwen2 ChatML fragments. QWEN2_TEMPLATE = { "system_prefix": "<|im_start|>system\n", @@ -77,6 +79,8 @@ def __init__( # generation params, consumed only by this family's _generate. self._num_beams = int(common.num_beams) self._num_return = int(common.num_return_sequences) + # opt-in ALGR-style escalating beam (width doubles per SID level). + self._dynamic_beam = bool(common.dynamic_beam) # worst-case spliced length for the first-step activation-pool pre-sizing. self._max_total_len = self._compute_max_total_length() self._pool_warmed = False @@ -167,11 +171,18 @@ def _splice_input_ids( # input_ids: assembled per row (user history length varies), then # left-padded into a (B, T) batch (real content right-aligned). rows_ids = [ - torch.cat([ - self.tpl_system, self.tpl_user_prefix, user_seq_rows[i], - self.tpl_user_suffix, self.tpl_asst_prefix, label_rows[i], - self.tpl_asst_suffix, self.tpl_eos, - ]) + torch.cat( + [ + self.tpl_system, + self.tpl_user_prefix, + user_seq_rows[i], + self.tpl_user_suffix, + self.tpl_asst_prefix, + label_rows[i], + self.tpl_asst_suffix, + self.tpl_eos, + ] + ) for i in range(len(user_seq_rows)) ] input_ids, attention_mask = self._left_pad(rows_ids, pad_to=pad_to) @@ -275,19 +286,40 @@ def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: max_codes=self._max_seq_length, # cap to most-recent items (drop oldest) ) input_ids, attention_mask = self._splice_prompt_ids(u_rows) - out = self.lm.generate( - input_ids=input_ids, - attention_mask=attention_mask, - max_new_tokens=self._num_levels, - num_beams=self._num_beams, - num_return_sequences=self._num_return, - do_sample=False, - pad_token_id=self._pad_token_id, - ) - new_tokens = out[:, input_ids.shape[1]:] # the generated tail + if self._dynamic_beam: + new_tokens = self._dynamic_beam_search(input_ids, attention_mask) + else: + out = self.lm.generate( + input_ids=input_ids, + attention_mask=attention_mask, + max_new_tokens=self._num_levels, + num_beams=self._num_beams, + num_return_sequences=self._num_return, + do_sample=False, + pad_token_id=self._pad_token_id, + ) + new_tokens = out[:, input_ids.shape[1] :] # the generated tail sids = self._validate_sid_candidates(new_tokens, input_ids.shape[0]) return {self.GENERATED_SIDS_KEY: sids} + def _dynamic_beam_search( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor + ) -> torch.Tensor: + """ALGR-style escalating-beam decode (delegates to the shared kernel). + + Returns the generated tail ``(B * num_beams * 2**num_levels, num_levels)`` + score-ordered best-first per row, ready for ``_validate_sid_candidates``. + See ``escalating_beam_search`` for the schedule. + """ + return escalating_beam_search( + self.lm, + input_ids, + attention_mask, + num_beams=self._num_beams, + lo_tok=self._tokenize_sids(self._sid_lvl_lo), + hi_tok=self._tokenize_sids(self._sid_lvl_hi), + ) + def _splice_prompt_ids( self, user_seq_rows: List[torch.Tensor] ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -298,10 +330,15 @@ def _splice_prompt_ids( continues from the assistant turn. """ rows = [ - torch.cat([ - self.tpl_system, self.tpl_user_prefix, r, - self.tpl_user_suffix, self.tpl_asst_prefix, - ]) + torch.cat( + [ + self.tpl_system, + self.tpl_user_prefix, + r, + self.tpl_user_suffix, + self.tpl_asst_prefix, + ] + ) for r in user_seq_rows ] return self._left_pad(rows) @@ -320,12 +357,16 @@ def _left_pad( end-aligned supervised tail in place, so labels/suffix-slice are intact. """ input_ids = pad_sequence( - rows, batch_first=True, - padding_value=self._pad_token_id, padding_side="left", + rows, + batch_first=True, + padding_value=self._pad_token_id, + padding_side="left", ) attention_mask = pad_sequence( - [torch.ones_like(r) for r in rows], batch_first=True, - padding_value=0, padding_side="left", + [torch.ones_like(r) for r in rows], + batch_first=True, + padding_value=0, + padding_side="left", ) if pad_to > input_ids.shape[1]: B, extra = input_ids.shape[0], pad_to - input_ids.shape[1] diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/qwen2_rec_lm_test.py index 146c7d03f..8017fe342 100644 --- a/tzrec/models/qwen2_rec_lm_test.py +++ b/tzrec/models/qwen2_rec_lm_test.py @@ -36,11 +36,16 @@ def _stub(num_levels=3, base_vocab=100, pad_id=9, device="cpu", per_level=4): m._num_levels = num_levels m._base_vocab = base_vocab m._pad_token_id = pad_id + m._dynamic_beam = False # default = HF fixed-width beam path m._max_seq_length = 0 # no recency clip by default in unit stubs m.lm = types.SimpleNamespace(device=torch.device(device)) for name, vals in { - "tpl_system": [10, 11], "tpl_user_prefix": [12], "tpl_user_suffix": [13], - "tpl_asst_prefix": [14], "tpl_asst_suffix": [15], "tpl_eos": [9], + "tpl_system": [10, 11], + "tpl_user_prefix": [12], + "tpl_user_suffix": [13], + "tpl_asst_prefix": [14], + "tpl_asst_suffix": [15], + "tpl_eos": [9], }.items(): m.register_buffer(name, torch.tensor(vals, dtype=torch.long), persistent=False) lo, hi = Qwen2RecLM._sid_level_bands([per_level] * num_levels) @@ -141,8 +146,15 @@ def test_generate_maps_tokens_to_sids(self) -> None: m._input_name = "user_sequence" m._num_beams = m._num_return = 2 - def fake_generate(input_ids, attention_mask, max_new_tokens, - num_beams, num_return_sequences, do_sample, pad_token_id): + def fake_generate( + input_ids, + attention_mask, + max_new_tokens, + num_beams, + num_return_sequences, + do_sample, + pad_token_id, + ): prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) # 2 beams x 3 codes, every atom INSIDE its level's band # (bands lo=[1,5,9] hi=[4,8,12]; token = sid + 99): @@ -162,15 +174,28 @@ def test_generate_rejects_malformed_candidates(self) -> None: m._input_name = "user_sequence" m._num_beams = m._num_return = 4 - def fake_generate(input_ids, attention_mask, max_new_tokens, - num_beams, num_return_sequences, do_sample, pad_token_id): + def fake_generate( + input_ids, + attention_mask, + max_new_tokens, + num_beams, + num_return_sequences, + do_sample, + pad_token_id, + ): prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) - new = torch.tensor([ - [100, 104, 108], # all in-band -> valid -> [1, 5, 9] - [100, 104, 9], # pos2 = eos/pad token (sid -90) -> invalid - [108, 104, 100], # wrong-level scramble (pos0 = lvl-2 code) -> invalid - [100, 104, 112], # pos2 sid 13 > band hi 12 -> invalid - ]) + new = torch.tensor( + [ + [100, 104, 108], # all in-band -> valid -> [1, 5, 9] + [100, 104, 9], # pos2 = eos/pad token (sid -90) -> invalid + [ + 108, + 104, + 100, + ], # wrong-level scramble (pos0 = lvl-2 code) -> invalid + [100, 104, 112], # pos2 sid 13 > band hi 12 -> invalid + ] + ) return torch.cat([prompt, new], dim=1) m.lm.generate = fake_generate @@ -189,8 +214,15 @@ def test_generate_narrow_tail_no_crash(self) -> None: m._input_name = "user_sequence" m._num_beams = m._num_return = 2 - def fake_generate(input_ids, attention_mask, max_new_tokens, - num_beams, num_return_sequences, do_sample, pad_token_id): + def fake_generate( + input_ids, + attention_mask, + max_new_tokens, + num_beams, + num_return_sequences, + do_sample, + pad_token_id, + ): prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) new = torch.tensor([[100, 104], [103, 107]]) # width 2 < num_levels 3 return torch.cat([prompt, new], dim=1) @@ -213,8 +245,12 @@ def test_build_prompt_tokens_registers_buffers(self) -> None: ) m._build_prompt_tokens(tok, cfg) for name in [ - "tpl_system", "tpl_user_prefix", "tpl_user_suffix", - "tpl_asst_prefix", "tpl_asst_suffix", "tpl_eos", + "tpl_system", + "tpl_user_prefix", + "tpl_user_suffix", + "tpl_asst_prefix", + "tpl_asst_suffix", + "tpl_eos", ]: buf = getattr(m, name) self.assertIsInstance(buf, torch.Tensor) @@ -324,5 +360,100 @@ def fwd(i, lbl, a): self.assertFalse(m._pool_warmed) +def _real_lm_stub(num_levels=3, base_vocab=20, per_level=4, num_beams=2): + """A Qwen2RecLM carrying a real (tiny, random) Qwen2 backbone. + + Needed by the dynamic-beam tests, which exercise the actual KV-cached + forward / cache-reorder path (the other tests mock ``lm.generate``). The SID + atoms occupy the last ``num_levels * per_level`` token ids, matching the + base-vocab + appended-codebook layout. + """ + from transformers import Qwen2Config, Qwen2ForCausalLM + + m = object.__new__(Qwen2RecLM) + nn.Module.__init__(m) + m._num_levels = num_levels + m._base_vocab = base_vocab + m._num_beams = num_beams + cfg = Qwen2Config( + vocab_size=base_vocab + per_level * num_levels, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=64, + ) + torch.manual_seed(0) + m.lm = Qwen2ForCausalLM(cfg).eval() + lo, hi = Qwen2RecLM._sid_level_bands([per_level] * num_levels) + m.register_buffer("_sid_lvl_lo", lo, persistent=False) + m.register_buffer("_sid_lvl_hi", hi, persistent=False) + return m + + +class Qwen2DynamicBeamTest(unittest.TestCase): + def test_width_schedule_and_final_count(self) -> None: + # widths double per level, returning num_beams * 2**num_levels candidates. + m = _real_lm_stub(num_levels=3, base_vocab=20, per_level=8, num_beams=2) + ids = torch.tensor([[1, 2, 3, 4]]) + new = m._dynamic_beam_search(ids, torch.ones_like(ids)) + # base 2: widths [4, 8, 16] (none capped: per_level 8 is roomy) -> 16 final + self.assertEqual(tuple(new.shape), (2 * 2**3, 3)) + + def test_every_candidate_is_in_band(self) -> None: + # band masking guarantees well-formed SIDs: validate -> no -1 sentinels. + m = _real_lm_stub(num_levels=3, base_vocab=20, per_level=4, num_beams=2) + ids = torch.tensor([[5, 6, 7]]) + new = m._dynamic_beam_search(ids, torch.ones_like(ids)) + sids = m._validate_sid_candidates(new, batch_size=1) + self.assertEqual(tuple(sids.shape), (1, new.shape[0], 3)) + self.assertFalse(bool((sids < 0).any())) # every candidate well-formed + + def test_left_padding_two_rows(self) -> None: + # ragged batch (row 1 left-padded): both rows yield valid, full beam sets. + m = _real_lm_stub(num_levels=2, base_vocab=20, per_level=4, num_beams=2) + ids = torch.tensor([[5, 6, 7, 8], [0, 0, 9, 10]]) + am = torch.tensor([[1, 1, 1, 1], [0, 0, 1, 1]]) + new = m._dynamic_beam_search(ids, am) + self.assertEqual(tuple(new.shape), (2 * 2 * 2**2, 2)) # B=2, 2*2^2=8 each + sids = m._validate_sid_candidates(new, batch_size=2) + self.assertEqual(tuple(sids.shape), (2, 8, 2)) + self.assertFalse(bool((sids < 0).any())) + + def test_exhaustive_matches_bruteforce_topk(self) -> None: + # When the schedule covers the whole tree (no pruning) the escalating + # beam is EXACT: its candidate set must equal all SID combos and be + # ordered by true (full-recompute) cumulative log-prob. This validates + # cache-stepping, band masking, scoring, and ordering end-to-end. + per, base = 3, 20 + m = _real_lm_stub(num_levels=2, base_vocab=base, per_level=per, num_beams=3) + # widths [min(6,3)=3, min(12,9)=9] -> exhaustive over all 3*3=9 SIDs + ids = torch.tensor([[5, 6, 7, 8]]) + am = torch.ones_like(ids) + got = [tuple(r) for r in m._dynamic_beam_search(ids, am).tolist()] + self.assertEqual(len(got), per * per) + lm = m.lm + lo0, lo1 = base, base + per # level token bands [20,22] and [23,25] + ref = {} + with torch.no_grad(): + logp0 = torch.log_softmax( + lm(ids, attention_mask=am).logits[0, -1].float(), -1 + ) + for t0 in range(lo0, lo0 + per): + s2 = torch.cat([ids, torch.tensor([[t0]])], 1) + logp1 = torch.log_softmax( + lm(s2, attention_mask=torch.ones_like(s2)).logits[0, -1].float(), -1 + ) + for t1 in range(lo1, lo1 + per): + ref[(t0, t1)] = (logp0[t0] + logp1[t1]).item() + # 1. exhaustive: the returned set is exactly every SID combination + self.assertEqual(set(got), set(ref)) + # 2. ordered best-first by the true score (tolerant of float-noise ties) + s = [ref[c] for c in got] + self.assertTrue(all(s[i] >= s[i + 1] - 1e-4 for i in range(len(s) - 1))) + self.assertEqual(got[0], max(ref, key=ref.get)) # top-1 is the global best + + if __name__ == "__main__": unittest.main() diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index 9e19579f5..268b6b832 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -40,6 +40,12 @@ message GenerativeRecLMConfig { // Inference (beam search) — used only by predict()'s inference branch. optional uint32 num_beams = 7 [default = 50]; optional uint32 num_return_sequences = 8 [default = 50]; + // When set, decode with the ALGR-style escalating beam instead of HF's + // fixed-width beam search: the beam width doubles at every SID level + // (num_beams -> 2*num_beams -> ... ), returning num_beams * 2**num_levels + // candidates (num_return_sequences is ignored). Faithful parity with ALGR's + // dynamic_beams schedule; exploits the fixed-length, EOS-free SID answer. + optional bool dynamic_beam = 9 [default = false]; } // Qwen2 / Qwen2.5 family (Qwen2.5-0.5B, etc.). Registered directly via its From af8788880ceb97efd34c5f2edb549588ea417cb0 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 18 Jun 2026 02:52:52 +0000 Subject: [PATCH 20/99] [examples] genrec LM: data converters + configs + smoke/train scaffolding Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/convert_s1full_to_parquet.py | 121 +++++++++ examples/convert_s1tiny_test_to_parquet.py | 149 +++++++++++ examples/convert_s1tiny_to_parquet.py | 173 ++++++++++++ examples/convert_s2test_to_parquet.py | 90 +++++++ .../generative_rec_lm_s2pretrained.config | 84 ++++++ examples/generative_rec_lm_smoke.py | 149 +++++++++++ examples/generative_rec_lm_train_loop.py | 230 ++++++++++++++++ .../generative_rec_lm_train_loop_parquet.py | 247 ++++++++++++++++++ 8 files changed, 1243 insertions(+) create mode 100644 examples/convert_s1full_to_parquet.py create mode 100644 examples/convert_s1tiny_test_to_parquet.py create mode 100644 examples/convert_s1tiny_to_parquet.py create mode 100644 examples/convert_s2test_to_parquet.py create mode 100644 examples/generative_rec_lm_s2pretrained.config create mode 100644 examples/generative_rec_lm_smoke.py create mode 100644 examples/generative_rec_lm_train_loop.py create mode 100644 examples/generative_rec_lm_train_loop_parquet.py diff --git a/examples/convert_s1full_to_parquet.py b/examples/convert_s1full_to_parquet.py new file mode 100644 index 000000000..1133ad220 --- /dev/null +++ b/examples/convert_s1full_to_parquet.py @@ -0,0 +1,121 @@ +"""Convert full-s1 (``s1_splits/part_*.csv``, 41 shards, ~246 GB) into +TorchEasyRec genrec parquet — the multi-CSV streaming sibling of +``convert_s1tiny_to_parquet.py``. + +Identical row logic (extract every ``C\\d+`` -> 1-indexed SID; columns +``user_sequence`` / ``label`` as ``list``), but iterates a GLOB of CSV +parts with a CONTINUOUS shard index so the 41 files land in one parquet dir. +Streaming (``csv.DictReader`` + shard-buffered flush) so memory stays bounded +regardless of the 246 GB input; the dropped CN-prompt text means the parquet is +a small fraction of the CSV size. + +Usage on remote:: + + /opt/conda/bin/python -m examples.convert_s1full_to_parquet \\ + --csv_glob '/home/admin/workspace/aop_lab/data/AL-GR-v1/s1_splits/*.csv' \\ + --out_dir /home/admin/workspace/aop_lab/data/AL-GR-v1/train_data_genreclm_s1full \\ + --shard_size 200000 --log_every 1000000 +""" + +from __future__ import annotations + +import argparse +import csv +import glob +import os +import re +import sys +import time +from typing import Iterator, List, Tuple + +import pyarrow as pa +import pyarrow.parquet as pq + +csv.field_size_limit(10 * 1024 * 1024) +_SID_RE = re.compile(r"C(\d+)") + + +def _extract_sids(text: str) -> List[int]: + return [int(m) + 1 for m in _SID_RE.findall(text)] + + +def _iter_rows( + csv_paths: List[str], max_rows: int +) -> Iterator[Tuple[List[int], List[int]]]: + n = 0 + for p in csv_paths: + with open(p, "r", encoding="utf-8", newline="") as f: + for row in csv.DictReader(f): + if max_rows and n >= max_rows: + return + u = _extract_sids(row.get("user") or "") + lab = _extract_sids(row.get("answer") or "") + if not u or not lab: + continue + n += 1 + yield u, lab + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--csv_glob", required=True) + ap.add_argument("--out_dir", required=True) + ap.add_argument("--shard_size", type=int, default=200_000) + ap.add_argument("--max_rows", type=int, default=0) + ap.add_argument("--log_every", type=int, default=1_000_000) + args = ap.parse_args() + + paths = sorted(glob.glob(args.csv_glob)) + print(f"{len(paths)} csv parts matched", flush=True) + os.makedirs(args.out_dir, exist_ok=True) + schema = pa.schema( + [ + pa.field("user_sequence", pa.list_(pa.int64()), nullable=False), + pa.field("label", pa.list_(pa.int64()), nullable=False), + ] + ) + + def flush(rows: List[Tuple[List[int], List[int]]], idx: int) -> None: + if not rows: + return + tbl = pa.table( + {"user_sequence": [r[0] for r in rows], "label": [r[1] for r in rows]}, + schema=schema, + ) + pq.write_table( + tbl, os.path.join(args.out_dir, f"shard-{idx:05d}.parquet"), + compression="zstd", + ) + + t0 = time.time() + rows: List[Tuple[List[int], List[int]]] = [] + idx = 0 + total = 0 + max_sid = 0 + for u, lab in _iter_rows(paths, args.max_rows): + rows.append((u, lab)) + total += 1 + max_sid = max(max_sid, max(u), max(lab)) + if len(rows) >= args.shard_size: + flush(rows, idx) + rows = [] + idx += 1 + if args.log_every and total % args.log_every == 0: + print( + f"[progress] rows={total} shards={idx} max_sid={max_sid} " + f"wall={time.time() - t0:.0f}s", + flush=True, + ) + if rows: + flush(rows, idx) + idx += 1 + print( + f"[done] rows={total} shards={idx} max_sid={max_sid} " + f"wall={time.time() - t0:.0f}s out_dir={args.out_dir}", + flush=True, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/convert_s1tiny_test_to_parquet.py b/examples/convert_s1tiny_test_to_parquet.py new file mode 100644 index 000000000..b557caaf6 --- /dev/null +++ b/examples/convert_s1tiny_test_to_parquet.py @@ -0,0 +1,149 @@ +"""Convert algr's ``s1_tiny_test.csv`` into TorchEasyRec-shaped parquet. + +The TEST split differs from the train split: the ``answer`` column holds +ground-truth ITEM IDs (e.g. ``AIB2f`` or ``9nXWC;I6kld;...`` — +semicolon-separated when multiple), not SID ``C``-codes. algr evaluates +this split via beam-search generation + ``calc_hr_fast.py`` item mapping. + +For TER's periodic CE-loss eval we need ``(user_sequence, label)`` SID rows, +so this script maps the FIRST ground-truth item of each row to its SID +triple via ``item_info/tiny_item_sid_final.csv``: + + sid_triple(item) = [lv1 + 1, lv2 + 8192 + 1, lv3 + 16384 + 1] + +(the per-layer offsets follow the dataset's C-token scheme +``C{lv1}C{lv2+8192}C{lv3+16384}``; the +1 converts the 0-indexed C-code to +the 1-indexed SID contract of ``GenerativeRecLM`` — identical to +``convert_s1tiny_to_parquet.py``'s ``Ck → SID = k + 1``.) + +Rows whose first answer item is missing from the item map are skipped +(counted and reported). + +Usage on remote:: + + /opt/conda/bin/python -m examples.convert_s1tiny_test_to_parquet \\ + --csv /home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data/s1_tiny_test.csv \\ + --item_sid_csv /home/admin/workspace/aop_lab/data/AL-GR-Tiny/item_info/tiny_item_sid_final.csv \\ + --out_dir /home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data_genreclm +""" + +from __future__ import annotations + +import argparse +import csv +import os +import re +import sys +import time +from typing import Dict, List, Tuple + +import pyarrow as pa +import pyarrow.parquet as pq + +csv.field_size_limit(10 * 1024 * 1024) + +_SID_RE = re.compile(r"C(\d+)") + +_LV2_OFFSET = 8192 +_LV3_OFFSET = 16384 + + +def _extract_sids(text: str) -> List[int]: + """``Ck → SID = k + 1`` (1-indexed; matches the train converter).""" + return [int(m) + 1 for m in _SID_RE.findall(text)] + + +def _load_item_map(path: str) -> Dict[str, Tuple[int, int, int]]: + t0 = time.time() + item_map: Dict[str, Tuple[int, int, int]] = {} + bad = 0 + with open(path, "r", encoding="utf-8", newline="") as f: + r = csv.DictReader(f) + for row in r: + try: + item_map[row["item_id"]] = ( + int(row["codebook_lv1"]), + int(row["codebook_lv2"]), + int(row["codebook_lv3"]), + ) + except (ValueError, TypeError, KeyError): + # the 25M-row map contains a few malformed/empty rows + bad += 1 + print( + f"[item_map] {len(item_map)} items loaded " + f"({bad} malformed rows skipped) in {time.time()-t0:.0f}s", + flush=True, + ) + return item_map + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--csv", required=True) + ap.add_argument("--item_sid_csv", required=True) + ap.add_argument("--out_dir", required=True) + ap.add_argument("--shard_size", type=int, default=200_000) + args = ap.parse_args() + + os.makedirs(args.out_dir, exist_ok=True) + schema = pa.schema( + [ + pa.field("user_sequence", pa.list_(pa.int64()), nullable=False), + pa.field("label", pa.list_(pa.int64()), nullable=False), + ] + ) + + item_map = _load_item_map(args.item_sid_csv) + + rows: List[Tuple[List[int], List[int]]] = [] + total = read = miss_item = miss_user = 0 + max_sid = 0 + with open(args.csv, "r", encoding="utf-8", newline="") as f: + for row in csv.DictReader(f): + read += 1 + user_sids = _extract_sids(row.get("user") or "") + if not user_sids: + miss_user += 1 + continue + first_item = (row.get("answer") or "").split(";")[0].strip() + lv = item_map.get(first_item) + if lv is None: + miss_item += 1 + continue + label_sids = [ + lv[0] + 1, + lv[1] + _LV2_OFFSET + 1, + lv[2] + _LV3_OFFSET + 1, + ] + rows.append((user_sids, label_sids)) + total += 1 + max_sid = max(max_sid, max(user_sids), max(label_sids)) + + shard_idx = 0 + for i in range(0, len(rows), args.shard_size): + chunk = rows[i : i + args.shard_size] + tbl = pa.table( + { + "user_sequence": [r[0] for r in chunk], + "label": [r[1] for r in chunk], + }, + schema=schema, + ) + pq.write_table( + tbl, + os.path.join(args.out_dir, f"shard-{shard_idx:05d}.parquet"), + compression="zstd", + ) + shard_idx += 1 + + print( + f"[done] read={read} rows_written={total} shards={shard_idx} " + f"skipped_missing_item={miss_item} skipped_no_user_sids={miss_user} " + f"max_sid={max_sid} out_dir={args.out_dir}", + flush=True, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/convert_s1tiny_to_parquet.py b/examples/convert_s1tiny_to_parquet.py new file mode 100644 index 000000000..99a54fd5b --- /dev/null +++ b/examples/convert_s1tiny_to_parquet.py @@ -0,0 +1,173 @@ +"""Convert algr's ``s1_tiny.csv`` into TorchEasyRec-shaped parquet. + +algr's row format:: + + system,user,answer + ,...C4805C8364C16402C4487C12277...,C1517C12109C16399 + +Each item is encoded as 3 contiguous ``C{i}`` codes (one per RQ-VAE layer). +The ``user`` column holds a CN sentence prefix + the user-history SID codes ++ a CN sentence suffix; the ``answer`` column holds just the target SID +codes for the next item. + +This script extracts all ``C\\d+`` matches from each field, converts to +1-indexed SID integers (``Ck → SID = k + 1``), and writes a parquet shard +per N rows with two columns matching ``GenerativeRecLM``'s contract: + + user_sequence : list + label : list + +The chat-template prompt strings are NOT carried over — ``GenerativeRecLM`` +re-builds them from cached buffers at splice time, using its own +``system_instruction``. To preserve algr's bit-exact prompt for parity +testing, set ``GenerativeRecLM.system_instruction`` to algr's CN prefix in +the pipeline.config (proto field already exists; see §3 of the design). + +Usage on remote:: + + /opt/conda/bin/python -m examples.convert_s1tiny_to_parquet \\ + --csv /home/admin/workspace/aop_lab/data/AL-GR-Tiny/train_data/s1_tiny.csv \\ + --out_dir /home/admin/workspace/aop_lab/data/AL-GR-Tiny/train_data_genreclm \\ + --shard_size 200000 --max_rows 0 \\ + --log_every 200000 +""" + +from __future__ import annotations + +import argparse +import csv +import os +import re +import sys +import time +from typing import Iterator, List, Tuple + +import pyarrow as pa +import pyarrow.parquet as pq + +# Allow rows up to ~6 MB of CSV text (the long Chinese prompt + thousand SIDs). +csv.field_size_limit(10 * 1024 * 1024) + +# Match every Cxxx in a string. SIDs are non-negative integers; we use a +# bounded ``+`` quantifier so it's anchored on whole tokens. +_SID_RE = re.compile(r"C(\d+)") + + +def _extract_sids(text: str) -> List[int]: + """Pull every ``C\\d+`` out of `text` and return as 1-indexed SIDs. + + ``Ck → SID = k + 1`` so that SID range is [1, sum(codebook)] which is what + ``GenerativeRecLM._splice_input_ids`` expects (SID=1 maps to atom C0 via + ``token = sid + base - 1`` = ``base + (sid - 1)``). + """ + return [int(m) + 1 for m in _SID_RE.findall(text)] + + +def _iter_rows(csv_path: str, max_rows: int) -> Iterator[Tuple[List[int], List[int]]]: + with open(csv_path, "r", encoding="utf-8", newline="") as f: + r = csv.DictReader(f) + for i, row in enumerate(r): + if max_rows and i >= max_rows: + return + user_sids = _extract_sids(row.get("user") or "") + label_sids = _extract_sids(row.get("answer") or "") + if not user_sids or not label_sids: + continue + yield user_sids, label_sids + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--csv", required=True) + ap.add_argument("--out_dir", required=True) + ap.add_argument( + "--shard_size", type=int, default=200_000, + help="rows per parquet shard", + ) + ap.add_argument( + "--max_rows", type=int, default=0, + help="cap total rows; 0 = whole CSV", + ) + ap.add_argument("--log_every", type=int, default=100_000) + args = ap.parse_args() + + os.makedirs(args.out_dir, exist_ok=True) + # Two int64 lists, no nulls. ``list_(int64())`` lets pyarrow store these + # as ListArray(int64) inside the parquet, which TER's + # ``sequence_raw_feature`` reads directly into a JaggedTensor. + schema = pa.schema( + [ + pa.field("user_sequence", pa.list_(pa.int64()), nullable=False), + pa.field("label", pa.list_(pa.int64()), nullable=False), + ] + ) + + t0 = time.time() + rows_in_shard: List[Tuple[List[int], List[int]]] = [] + shard_idx = 0 + total = 0 + max_sid_seen = 0 + max_user_len = 0 + max_label_len = 0 + + def flush(rows, idx): + if not rows: + return + user_col = [r[0] for r in rows] + label_col = [r[1] for r in rows] + tbl = pa.table( + {"user_sequence": user_col, "label": label_col}, schema=schema, + ) + path = os.path.join(args.out_dir, f"shard-{idx:05d}.parquet") + pq.write_table(tbl, path, compression="zstd") + + for u, lab in _iter_rows(args.csv, args.max_rows): + rows_in_shard.append((u, lab)) + total += 1 + max_sid_seen = max(max_sid_seen, max(u), max(lab)) + max_user_len = max(max_user_len, len(u)) + max_label_len = max(max_label_len, len(lab)) + + if len(rows_in_shard) >= args.shard_size: + flush(rows_in_shard, shard_idx) + rows_in_shard = [] + shard_idx += 1 + print( + f"[shard {shard_idx-1}] flushed; total={total} " + f"max_sid={max_sid_seen} max_user_len={max_user_len} " + f"max_label_len={max_label_len} " + f"wall={time.time()-t0:.0f}s", + flush=True, + ) + + if args.log_every and total % args.log_every == 0: + print( + f"[progress] rows={total} max_sid={max_sid_seen} " + f"max_user_len={max_user_len} max_label_len={max_label_len} " + f"wall={time.time()-t0:.0f}s", + flush=True, + ) + + if rows_in_shard: + flush(rows_in_shard, shard_idx) + shard_idx += 1 + + print( + f"[done] rows_written={total} shards={shard_idx} " + f"max_sid={max_sid_seen} max_user_len={max_user_len} " + f"max_label_len={max_label_len} " + f"wall={time.time()-t0:.1f}s out_dir={args.out_dir}", + flush=True, + ) + if max_sid_seen > 65536: + print( + f"[WARN] max_sid={max_sid_seen} > 65536 — algr's tokenizer adds " + f"only 65 536 atoms. Vocab extension in TER must be at least " + f"{max_sid_seen}.", + flush=True, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/convert_s2test_to_parquet.py b/examples/convert_s2test_to_parquet.py new file mode 100644 index 000000000..c3c5c59b0 --- /dev/null +++ b/examples/convert_s2test_to_parquet.py @@ -0,0 +1,90 @@ +"""Convert s2 test CSV to TER parquet (user_sequence only; label = [0] placeholder). + +s2 test CSV answers are item IDs (not SID codes), so label cannot be derived. +The predict script only reads `user_sequence` from parquet, so label is a dummy. + +Usage on remote:: + + python3 /home/admin/workspace/TorchEasyRec_qwen_smoke/TorchEasyRec/examples/convert_s2test_to_parquet.py \ + --csv /home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data/s2_tiny_test.csv \ + --out_dir /home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data_genreclm_s2 +""" +from __future__ import annotations + +import argparse +import csv +import os +import re +import sys +import time +from typing import Iterator, List, Tuple + +import pyarrow as pa +import pyarrow.parquet as pq + +csv.field_size_limit(10 * 1024 * 1024) + +_SID_RE = re.compile(r"C(\d+)") + + +def _extract_user_sids(text: str) -> List[int]: + return [int(m) + 1 for m in _SID_RE.findall(text)] + + +def _iter_rows(csv_path: str, max_rows: int) -> Iterator[List[int]]: + with open(csv_path, newline="", encoding="utf-8") as f: + reader = csv.DictReader(f) + for i, row in enumerate(reader): + if max_rows and i >= max_rows: + break + user_sids = _extract_user_sids(row["user"]) + if not user_sids: + continue + yield user_sids + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--csv", required=True) + ap.add_argument("--out_dir", required=True) + ap.add_argument("--shard_size", type=int, default=200_000) + ap.add_argument("--max_rows", type=int, default=0) + args = ap.parse_args() + + os.makedirs(args.out_dir, exist_ok=True) + schema = pa.schema([ + pa.field("user_sequence", pa.list_(pa.int64()), nullable=False), + pa.field("label", pa.list_(pa.int64()), nullable=False), + ]) + + t0 = time.time() + rows_in_shard: List[Tuple[List[int], List[int]]] = [] + shard_idx = 0 + total = 0 + + def flush(rows, idx): + if not rows: + return + user_col = [r[0] for r in rows] + label_col = [r[1] for r in rows] + tbl = pa.table({"user_sequence": user_col, "label": label_col}, schema=schema) + pq.write_table(tbl, os.path.join(args.out_dir, f"shard-{idx:05d}.parquet"), compression="zstd") + + for u in _iter_rows(args.csv, args.max_rows): + rows_in_shard.append((u, [0])) # dummy label + total += 1 + if len(rows_in_shard) >= args.shard_size: + flush(rows_in_shard, shard_idx) + rows_in_shard = [] + shard_idx += 1 + + if rows_in_shard: + flush(rows_in_shard, shard_idx) + shard_idx += 1 + + print(f"[done] rows_written={total} shards={shard_idx} wall={time.time()-t0:.1f}s out_dir={args.out_dir}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/generative_rec_lm_s2pretrained.config b/examples/generative_rec_lm_s2pretrained.config new file mode 100644 index 000000000..05309e013 --- /dev/null +++ b/examples/generative_rec_lm_s2pretrained.config @@ -0,0 +1,84 @@ +train_input_path: "/home/admin/workspace/aop_lab/data/AL-GR-Tiny/train_data_genreclm_s2/*.parquet" +eval_input_path: "/home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data_genreclm_s2/*.parquet" +model_dir: "experiments/generative_rec_lm_s2pretrained" +train_config { + sparse_optimizer { + adagrad_optimizer { + lr: 0.0 + } + constant_learning_rate { + } + } + dense_optimizer { + adam_optimizer { + lr: 5e-05 + beta1: 0.9 + beta2: 0.999 + eps: 1e-08 + } + linear_decay_learning_rate { + total_size: 500000 + } + } + num_steps: 24000 + save_checkpoints_steps: 8000 + log_step_count_steps: 1000 + gradient_accumulation_steps: 4 + grad_clipping { + clipping_type: "norm" + max_gradient: 4.0 + norm_type: 2.0 + enable_global_grad_clip: true + } +} +eval_config { +} +data_config { + batch_size: 80 + dataset_type: ParquetDataset + num_workers: 4 + fg_mode: FG_NONE +} +feature_configs { + sequence_raw_feature { + feature_name: "user_sequence" + expression: "user:user_sequence" + value_dim: 1 + sequence_length: 1056 + } +} +feature_configs { + sequence_raw_feature { + feature_name: "label" + expression: "user:label" + value_dim: 1 + sequence_length: 32 + } +} +model_config { + feature_groups { + group_name: "sids" + feature_names: "user_sequence" + feature_names: "label" + group_type: SEQUENCE + } + qwen2_rec_lm { + # Qwen2 backbone (owned by this family message, not `common`). + hf_model_id: "/home/admin/workspace/Qwen2.5-0.5B" + common { + user_sequence_feature_name: "user_sequence" + label_feature_name: "label" + ignore_index: -100 + # SID codebook: one entry per RQ level (AL-GR = 3 levels x 8192; verified + # from item_info codebook_lv* ranges). len = SID codes/answer, sum = vocab + # atoms. + codebook: 8192 + codebook: 8192 + codebook: 8192 + vocab_pad_to_multiple_of: 128 + } + system_instruction: "你是一个推荐系统，根据用户的历史行为，预测用户在电商场景的下一步行为。我会给你一串连续行为的语义编码，按照用户点击的时间顺序排列，每个行为用三个词表示。" + user_prefix_text: "当前用户的历史行为如下：" + user_suffix_text: "，请预测用户在电商推荐场景后续行为的语义编码" + } +} diff --git a/examples/generative_rec_lm_smoke.py b/examples/generative_rec_lm_smoke.py new file mode 100644 index 000000000..87954990f --- /dev/null +++ b/examples/generative_rec_lm_smoke.py @@ -0,0 +1,149 @@ +"""Smoke test for ``GenerativeRecLM`` — exercises the full forward path +(splice + base forward + suffix slice + HF loss_function) without going +through ``tzrec.train_eval``. + +Validates: + 1. Proto + dispatch wire up. + 2. Vocab extension preserves SID-token offset arithmetic. + 3. Left-pad with eos_token_id (L7 fix) produces a finite, non-trivial CE. + 4. SID → token mapping (``token = sid + base_vocab - 1``) hits the new atoms. + +Usage: + cd /workspace/fangtinglin/codework/feat/support_qwen/TorchEasyRec + /opt/conda/bin/python -m examples.generative_rec_lm_smoke +""" + +from __future__ import annotations + +import os +import sys +import time + +import torch +from google.protobuf import text_format + +from tzrec.datasets.utils import Batch +from tzrec.models import generative_rec_lm # noqa: F401 registers class +from tzrec.models.model import BaseModel +from tzrec.protos.model_pb2 import ModelConfig + + +# Tiny config — small codebook so vocab extension stays cheap on CPU. +HF_MODEL_ID = "/workspace/fangtinglin/_hf_stage/Qwen2.5-0.5B" +CODEBOOK = [64, 64] # 128 SID atoms total → small vocab grow on CPU +USER_FEATURE = "user_sequence" +LABEL_FEATURE = "label" + + +def _make_proto(): + cfg = ModelConfig() + grl = cfg.generative_rec_lm + grl.class_name = "Qwen2RecLM" + grl.hf_model_id = HF_MODEL_ID + for c in CODEBOOK: + grl.codebook.append(c) + grl.user_sequence_feature_name = USER_FEATURE + grl.label_feature_name = LABEL_FEATURE + grl.ignore_index = -100 + return cfg + + +class _DummyJagged: + """Minimal stand-in for torchrec's JaggedTensor — just enough surface + for ``GenerativeRecLM._jagged_to_row_list`` to consume.""" + def __init__(self, values: torch.Tensor, lengths: torch.Tensor): + self._values = values + self._lengths = lengths + + def values(self) -> torch.Tensor: + return self._values + + def lengths(self) -> torch.Tensor: + return self._lengths + + +def _make_batch(B: int, sum_codebook: int, user_len: int, label_len: int) -> Batch: + # SID indices are 1-indexed in [1, sum(codebook)]. + torch.manual_seed(0) + user_vals = torch.randint(1, sum_codebook + 1, (B * user_len,), dtype=torch.long) + label_vals = torch.randint(1, sum_codebook + 1, (B * label_len,), dtype=torch.long) + user_lens = torch.full((B,), user_len, dtype=torch.long) + label_lens = torch.full((B,), label_len, dtype=torch.long) + + seq_dense = { + USER_FEATURE: _DummyJagged(user_vals, user_lens), + LABEL_FEATURE: _DummyJagged(label_vals, label_lens), + } + # Batch is a dataclass-like NamedTuple — construct with just the field we need. + # tzrec.datasets.utils.Batch is a NamedTuple of many fields, all default to {}. + return Batch(sequence_dense_features=seq_dense) + + +def main() -> int: + os.environ["TZREC_GENRECLM_DEBUG"] = "1" + print(f"[smoke] torch={torch.__version__} cuda_avail={torch.cuda.is_available()}", flush=True) + + cfg = _make_proto() + print(f"[smoke] proto: class_name={cfg.generative_rec_lm.class_name} " + f"codebook={list(cfg.generative_rec_lm.codebook)}", flush=True) + + # --- dispatch --- + model_cls = BaseModel.create_class("GenerativeRecLM") + print(f"[smoke] dispatched -> {model_cls.__name__}", flush=True) + + # --- construct --- + t0 = time.time() + model = model_cls( + cfg, + features=[], + labels=[], + sample_weights=None, + ) + model.train() + print(f"[smoke] construct ok ({time.time()-t0:.1f}s); " + f"base_vocab={model._base_vocab} " + f"final_vocab={model.lm.config.vocab_size} " + f"pad_id={model._pad_token_id}", flush=True) + print(f"[smoke] tpl_system numel={model.tpl_system.numel()} " + f"tpl_user_prefix numel={model.tpl_user_prefix.numel()} " + f"tpl_asst_prefix numel={model.tpl_asst_prefix.numel()}", flush=True) + + # --- batch --- + sum_codebook = sum(CODEBOOK) + batch = _make_batch(B=2, sum_codebook=sum_codebook, user_len=5, label_len=4) + + # --- forward --- + t0 = time.time() + with torch.no_grad(): + pred = model.predict(batch) + elapsed = time.time() - t0 + loss = pred["loss"] + logits = pred["logits"] + print(f"[smoke] forward ok ({elapsed:.2f}s); " + f"loss={float(loss):.4f} logits.shape={tuple(logits.shape)}", flush=True) + + # Acceptance criteria + ok = True + if not torch.isfinite(loss): + print("[smoke] FAIL: loss is non-finite") + ok = False + if float(loss) < 0.1 or float(loss) > 100.0: + print(f"[smoke] FAIL: loss out of plausible range ({float(loss)})") + ok = False + if logits.shape[0] != 2: + print(f"[smoke] FAIL: logits batch dim wrong ({logits.shape})") + ok = False + + # --- one backward step to make sure grads flow through the chat-template + # buffers + base model + lm_head without error --- + t0 = time.time() + pred = model.predict(batch) + pred["loss"].backward() + print(f"[smoke] backward ok ({time.time()-t0:.2f}s)", flush=True) + + print(f"[smoke] {'PASS' if ok else 'FAIL'}", flush=True) + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/generative_rec_lm_train_loop.py b/examples/generative_rec_lm_train_loop.py new file mode 100644 index 000000000..a019923f0 --- /dev/null +++ b/examples/generative_rec_lm_train_loop.py @@ -0,0 +1,230 @@ +"""Multi-step CPU training loop for ``GenerativeRecLM``. + +Exercises ~hundreds of forward+backward+optimizer steps on synthetic batches +to verify the model learns. The data is a tiny next-SID-prediction task: +each row maps a sampled ``user_sequence`` of SIDs to a deterministic +``label`` SID list derived from it via a fixed permutation, so a sufficiently +expressive LM can drive the CE loss down toward 0. + +Emits one ``[step …]`` line every ``--log-every`` steps so a Monitor wrapper +can stream loss progress. Failure modes (Traceback, NaN/Inf loss) also go to +stdout so the monitor sees them. + +Usage: + cd /workspace/fangtinglin/codework/feat/support_qwen/TorchEasyRec + /opt/conda/bin/python -m examples.generative_rec_lm_train_loop \\ + --steps 500 --bsz 2 --user-len 6 --label-len 3 \\ + --log-every 5 +""" + +from __future__ import annotations + +import argparse +import math +import os +import sys +import time +from typing import List + +import torch + +from tzrec.datasets.utils import Batch +from tzrec.models import generative_rec_lm # noqa: F401 registers class +from tzrec.models.model import BaseModel +from tzrec.protos.model_pb2 import ModelConfig + + +HF_MODEL_ID_DEFAULT = "/workspace/fangtinglin/_hf_stage/Qwen2.5-0.5B" + + +class _DummyJagged: + def __init__(self, values: torch.Tensor, lengths: torch.Tensor): + self._values = values + self._lengths = lengths + + def values(self) -> torch.Tensor: + return self._values + + def lengths(self) -> torch.Tensor: + return self._lengths + + +def _make_proto(codebook: List[int], hf_model_id: str) -> ModelConfig: + cfg = ModelConfig() + grl = cfg.generative_rec_lm + grl.class_name = "Qwen2RecLM" + grl.hf_model_id = hf_model_id + for c in codebook: + grl.codebook.append(c) + grl.user_sequence_feature_name = "user_sequence" + grl.label_feature_name = "label" + grl.ignore_index = -100 + return cfg + + +def _sample_batch( + rng: torch.Generator, + sum_codebook: int, + bsz: int, + user_len: int, + label_len: int, +) -> Batch: + """Synthesise a deterministic-mapping batch: ``label = (user mod K) + 1`` + for the first ``label_len`` positions. Easy to memorise; gives the LM + something to drive CE down.""" + user_vals = torch.randint( + 1, sum_codebook + 1, (bsz, user_len), generator=rng, dtype=torch.long, + ) + # Deterministic label: pick the first label_len user positions, modded + # back into [1, sum_codebook]. Small enough to fit in CE in ~hundreds of + # steps even on a frozen-most-of-the-net base. + label_vals = ((user_vals[:, :label_len] % sum_codebook) + 1).contiguous() + user_flat = user_vals.reshape(-1) + label_flat = label_vals.reshape(-1) + user_lens = torch.full((bsz,), user_len, dtype=torch.long) + label_lens = torch.full((bsz,), label_len, dtype=torch.long) + return Batch(sequence_dense_features={ + "user_sequence": _DummyJagged(user_flat, user_lens), + "label": _DummyJagged(label_flat, label_lens), + }) + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--steps", type=int, default=500) + ap.add_argument("--bsz", type=int, default=2) + ap.add_argument("--user-len", type=int, default=6) + ap.add_argument("--label-len", type=int, default=3) + ap.add_argument("--log-every", type=int, default=5) + # 1e-5 matches algr's `qwen2.5_05b_3layer_s1tiny.json` (see + # [[project-tzrec-qwen2-integration]]). + ap.add_argument("--lr", type=float, default=1e-5) + ap.add_argument("--hf-model-id", default=HF_MODEL_ID_DEFAULT) + ap.add_argument( + "--codebook", default="64,64", + help="comma-separated codebook sizes; 64,64 fast / 32768,32768 algr-like", + ) + ap.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"]) + args = ap.parse_args() + + if args.device == "auto": + device = "cuda" if torch.cuda.is_available() else "cpu" + else: + device = args.device + print( + f"[init] torch={torch.__version__} device={device} " + f"steps={args.steps} bsz={args.bsz} user_len={args.user_len} " + f"label_len={args.label_len} lr={args.lr} hf={args.hf_model_id}", + flush=True, + ) + + codebook = [int(x) for x in args.codebook.split(",") if x] + sum_codebook = sum(codebook) + cfg = _make_proto(codebook, args.hf_model_id) + + model_cls = BaseModel.create_class("GenerativeRecLM") + t0 = time.time() + model = model_cls(cfg, features=[], labels=[], sample_weights=None) + model.train() + if device == "cuda": + model.to("cuda") + torch.cuda.reset_peak_memory_stats() + print( + f"[init] construct ok ({time.time()-t0:.1f}s) " + f"base_vocab={model._base_vocab} " + f"final_vocab={model.lm.config.vocab_size} " + f"sum_codebook={sum_codebook}", + flush=True, + ) + + optim = torch.optim.AdamW( + model.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-8, + ) + + rng = torch.Generator(device="cpu").manual_seed(0) + # Warmup: discount the first 3 steps from throughput accounting (CUDA + # kernel autotune, cudnn benchmark, allocator warm-up). + warmup_steps = 3 if device == "cuda" else 0 + start = None + losses: List[float] = [] + step_walls: List[float] = [] + for step in range(1, args.steps + 1): + batch = _sample_batch( + rng, sum_codebook, args.bsz, args.user_len, args.label_len, + ) + if device == "cuda": + # Move tensors to GPU — small batches, ok to do per-step. + for k, v in batch.sequence_dense_features.items(): + v._values = v._values.cuda(non_blocking=True) + v._lengths = v._lengths.cuda(non_blocking=True) + torch.cuda.synchronize() + t_step = time.perf_counter() + pred = model.predict(batch) + loss = pred["loss"] + + loss_val = loss.detach().item() + if not math.isfinite(loss_val): + print(f"[step {step}] FAIL: non-finite loss {loss_val}", flush=True) + return 1 + + optim.zero_grad(set_to_none=True) + loss.backward() + optim.step() + if device == "cuda": + torch.cuda.synchronize() + dt = time.perf_counter() - t_step + losses.append(loss.detach().item()) + + if step == warmup_steps: + start = time.time() + step_walls = [] + if device == "cuda": + torch.cuda.reset_peak_memory_stats() + elif step > warmup_steps: + step_walls.append(dt) + + if step % args.log_every == 0 or step == 1: + mem_str = "" + if device == "cuda": + peak_gb = torch.cuda.max_memory_allocated() / 1024**3 + resv_gb = torch.cuda.max_memory_reserved() / 1024**3 + mem_str = f" peak_alloc={peak_gb:.2f}GB peak_reserved={resv_gb:.2f}GB" + if start is not None and step_walls: + avg_dt = sum(step_walls) / len(step_walls) + it_s = 1.0 / avg_dt + wall_str = f" avg_step={avg_dt*1000:.1f}ms it/s={it_s:.2f}" + else: + wall_str = " (warmup)" + print( + f"[step {step}/{args.steps}] ce_loss={loss_val:.4f} " + f"avg_last_{min(20, len(losses))}={sum(losses[-20:])/min(20, len(losses)):.4f}" + f"{wall_str}{mem_str}", + flush=True, + ) + + # ---- final report ---- + start_avg = sum(losses[:5]) / min(5, len(losses)) + end_avg = sum(losses[-5:]) / min(5, len(losses)) + if step_walls: + median_dt = sorted(step_walls)[len(step_walls) // 2] + avg_dt = sum(step_walls) / len(step_walls) + peak_str = "" + if device == "cuda": + peak_gb = torch.cuda.max_memory_allocated() / 1024**3 + resv_gb = torch.cuda.max_memory_reserved() / 1024**3 + peak_str = ( + f" peak_alloc={peak_gb:.2f}GB peak_reserved={resv_gb:.2f}GB" + ) + print( + f"[done] steps={args.steps} bsz={args.bsz} " + f"start_avg5={start_avg:.4f} end_avg5={end_avg:.4f} " + f"delta={start_avg-end_avg:+.4f} " + f"avg_step={avg_dt*1000:.1f}ms median_step={median_dt*1000:.1f}ms " + f"it/s={1.0/avg_dt:.2f}{peak_str}", + flush=True, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/generative_rec_lm_train_loop_parquet.py b/examples/generative_rec_lm_train_loop_parquet.py new file mode 100644 index 000000000..434f27600 --- /dev/null +++ b/examples/generative_rec_lm_train_loop_parquet.py @@ -0,0 +1,247 @@ +"""Parquet-fed multi-step training loop for ``GenerativeRecLM``. + +Same instrumentation as ``generative_rec_lm_train_loop.py`` (GPU memory + +throughput) but consumes real ``user_sequence + label`` rows from a parquet +directory instead of synthesising them. Used for apples-to-apples +comparison with algr on the same data. + +Usage on remote:: + + /opt/conda/bin/python -m examples.generative_rec_lm_train_loop_parquet \\ + --parquet_dir /home/admin/workspace/aop_lab/data/AL-GR-Tiny/train_data_genreclm_100k \\ + --steps 25 --log-every 1 --bsz 10 --device cuda \\ + --codebook 32768,32768 \\ + --hf-model-id /home/admin/workspace/Qwen2.5-0.5B-Instruct +""" + +from __future__ import annotations + +import argparse +import glob +import math +import os +import sys +import time +from typing import List + +import pyarrow.parquet as pq +import torch + +from tzrec.datasets.utils import Batch +from tzrec.models import generative_rec_lm # noqa: F401 registers class +from tzrec.models.model import BaseModel +from tzrec.protos.model_pb2 import ModelConfig + + +class _DummyJagged: + def __init__(self, values: torch.Tensor, lengths: torch.Tensor): + self._values = values + self._lengths = lengths + + def values(self) -> torch.Tensor: + return self._values + + def lengths(self) -> torch.Tensor: + return self._lengths + + +class _ParquetIter: + """Round-robin iterator over rows in a parquet directory. + + Loads each shard into memory as Arrow Tables (small for our 100K split). + Yields ``(user_sids: list[int], label_sids: list[int])`` rows. Loops to + the start when exhausted so the training loop never runs out of data. + """ + def __init__(self, parquet_dir: str, max_source_len: int): + paths = sorted(glob.glob(os.path.join(parquet_dir, "*.parquet"))) + assert paths, f"no parquet files under {parquet_dir!r}" + self._rows: List = [] + for p in paths: + t = pq.read_table(p) + u_col = t["user_sequence"].to_pylist() + l_col = t["label"].to_pylist() + for u, lab in zip(u_col, l_col): + # algr truncates user prompts to ``max_source_length`` tokens + # via ``prompt_ids[:max_source_length]``. We mirror that at + # the SID-list level (each SID is one token after splice). + if max_source_len and len(u) > max_source_len: + u = u[-max_source_len:] # keep most-recent items + self._rows.append((u, lab)) + self._n = len(self._rows) + self._i = 0 + + def __len__(self) -> int: + return self._n + + def next_batch(self, bsz: int) -> Batch: + rows = [] + for _ in range(bsz): + rows.append(self._rows[self._i]) + self._i = (self._i + 1) % self._n + user_lens = torch.tensor([len(r[0]) for r in rows], dtype=torch.long) + label_lens = torch.tensor([len(r[1]) for r in rows], dtype=torch.long) + # Flatten to JaggedTensor (values, lengths) shape + user_vals: List[int] = [] + for r in rows: + user_vals.extend(r[0]) + label_vals: List[int] = [] + for r in rows: + label_vals.extend(r[1]) + return Batch(sequence_dense_features={ + "user_sequence": _DummyJagged( + torch.tensor(user_vals, dtype=torch.long), user_lens, + ), + "label": _DummyJagged( + torch.tensor(label_vals, dtype=torch.long), label_lens, + ), + }) + + +def _make_proto( + codebook: List[int], + hf_model_id: str, + system_instruction: str = "", + user_prefix_text: str = "", + user_suffix_text: str = "", +) -> ModelConfig: + cfg = ModelConfig() + grl = cfg.generative_rec_lm + grl.class_name = "Qwen2RecLM" + grl.hf_model_id = hf_model_id + for c in codebook: + grl.codebook.append(c) + grl.user_sequence_feature_name = "user_sequence" + grl.label_feature_name = "label" + grl.ignore_index = -100 + if system_instruction: + grl.system_instruction = system_instruction + if user_prefix_text: + grl.user_prefix_text = user_prefix_text + if user_suffix_text: + grl.user_suffix_text = user_suffix_text + return cfg + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--parquet_dir", required=True) + ap.add_argument("--steps", type=int, default=25) + ap.add_argument("--bsz", type=int, default=10) + ap.add_argument("--log-every", type=int, default=1) + ap.add_argument("--lr", type=float, default=1e-5) + ap.add_argument("--max-source-len", type=int, default=1020, + help="cap user_sequence SID count (mirrors algr max_source_length)") + ap.add_argument("--codebook", default="32768,32768") + ap.add_argument("--hf-model-id", required=True) + ap.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"]) + ap.add_argument("--system-instruction", default="", + help="CN/EN system prompt; matches algr's `default_instruction` / row.system") + ap.add_argument("--user-prefix-text", default="", + help="text wrapped before SID list in user message (e.g. algr CN prefix)") + ap.add_argument("--user-suffix-text", default="", + help="text wrapped after SID list in user message (e.g. algr CN suffix)") + args = ap.parse_args() + + device = "cuda" if args.device == "auto" and torch.cuda.is_available() else args.device + if device == "auto": + device = "cpu" + print( + f"[init] torch={torch.__version__} device={device} steps={args.steps} " + f"bsz={args.bsz} lr={args.lr} parquet={args.parquet_dir}", + flush=True, + ) + + t_load = time.time() + data = _ParquetIter(args.parquet_dir, args.max_source_len) + print(f"[init] parquet loaded ({time.time()-t_load:.1f}s) rows={len(data)}", flush=True) + + codebook = [int(x) for x in args.codebook.split(",") if x] + cfg = _make_proto( + codebook, args.hf_model_id, + system_instruction=args.system_instruction, + user_prefix_text=args.user_prefix_text, + user_suffix_text=args.user_suffix_text, + ) + model_cls = BaseModel.create_class("GenerativeRecLM") + t0 = time.time() + model = model_cls(cfg, features=[], labels=[], sample_weights=None) + model.train() + if device == "cuda": + model.to("cuda") + torch.cuda.reset_peak_memory_stats() + print( + f"[init] construct ok ({time.time()-t0:.1f}s) base_vocab={model._base_vocab} " + f"final_vocab={model.lm.config.vocab_size}", + flush=True, + ) + + optim = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-8) + + warmup = 3 if device == "cuda" else 0 + start = None + step_walls: List[float] = [] + losses: List[float] = [] + for step in range(1, args.steps + 1): + batch = data.next_batch(args.bsz) + if device == "cuda": + for v in batch.sequence_dense_features.values(): + v._values = v._values.cuda(non_blocking=True) + v._lengths = v._lengths.cuda(non_blocking=True) + torch.cuda.synchronize() + t = time.perf_counter() + pred = model.predict(batch) + loss = pred["loss"] + loss_val = loss.detach().item() + if not math.isfinite(loss_val): + print(f"[step {step}] FAIL: non-finite loss {loss_val}", flush=True) + return 1 + optim.zero_grad(set_to_none=True) + loss.backward() + optim.step() + if device == "cuda": + torch.cuda.synchronize() + dt = time.perf_counter() - t + losses.append(loss_val) + if step == warmup: + start = time.time() + step_walls = [] + if device == "cuda": + torch.cuda.reset_peak_memory_stats() + elif step > warmup: + step_walls.append(dt) + if step % args.log_every == 0 or step == 1: + mem_str = "" + if device == "cuda": + pa = torch.cuda.max_memory_allocated() / 1024**3 + pr = torch.cuda.max_memory_reserved() / 1024**3 + mem_str = f" peak_alloc={pa:.2f}GB peak_reserved={pr:.2f}GB" + wall_str = ( + f" avg_step={sum(step_walls)/len(step_walls)*1000:.1f}ms " + f"it/s={1.0/(sum(step_walls)/len(step_walls)):.2f}" + if step_walls else " (warmup)" + ) + ph_t = batch.sequence_dense_features["user_sequence"]._values.shape[0] + print( + f"[step {step}/{args.steps}] ce_loss={loss_val:.4f} " + f"avg_last_{min(20, len(losses))}={sum(losses[-20:])/min(20, len(losses)):.4f}" + f" sum_lens={ph_t}{wall_str}{mem_str}", + flush=True, + ) + if step_walls: + pa = (torch.cuda.max_memory_allocated() / 1024**3) if device == "cuda" else 0 + pr = (torch.cuda.max_memory_reserved() / 1024**3) if device == "cuda" else 0 + print( + f"[done] steps={args.steps} bsz={args.bsz} " + f"start_avg5={sum(losses[:5])/min(5, len(losses)):.4f} " + f"end_avg5={sum(losses[-5:])/min(5, len(losses)):.4f} " + f"avg_step={sum(step_walls)/len(step_walls)*1000:.1f}ms " + f"median_step={sorted(step_walls)[len(step_walls)//2]*1000:.1f}ms " + f"it/s={1.0/(sum(step_walls)/len(step_walls)):.2f}" + f" peak_alloc={pa:.2f}GB peak_reserved={pr:.2f}GB", + flush=True, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 446d857753b72f7d5c25ad379275526cad73cea5 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 24 Jun 2026 09:01:35 +0000 Subject: [PATCH 21/99] [feat] generative-rec LM: feature retrieval via init_input/build_input (JAGGED_SEQUENCE) Route the genrec-LM's SID retrieval through the framework's standard init_input/build_input + EmbeddingGroup path (the HSTU idiom) instead of reaching into batch.sequence_dense_features directly. - proto: add GenerativeRecLMConfig.history_group_name / label_group_name (group-name knobs, defaults "user_seq"/"answer"), keyed by GROUP name like HSTU, decoupled from feature names. - generative_rec_lm: init_input builds the param-free raw-passthrough EmbeddingGroup; build_input reads "{group}.sequence"/".sequence_length" and tokenizes; _sid_token_rows now takes (values, lengths); group knobs read in _read_common_config. - qwen2_rec_lm: _predict_train / _generate consume build_input. - example config: one SEQUENCE group -> two single-feature JAGGED_SEQUENCE groups. - tests updated (+ build_input coverage). Integration-verified that a top-level sequence_raw_feature passes raw through a JAGGED_SEQUENCE EmbeddingGroup. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../generative_rec_lm_s1pretrained.config | 14 +++- tzrec/models/generative_rec_lm.py | 64 ++++++++++++++++--- tzrec/models/generative_rec_lm_test.py | 50 +++++++++++++-- tzrec/models/qwen2_rec_lm.py | 20 ++---- tzrec/models/qwen2_rec_lm_test.py | 46 +++++++------ tzrec/protos/models/generative_model.proto | 8 +++ 6 files changed, 152 insertions(+), 50 deletions(-) diff --git a/examples/generative_rec_lm_s1pretrained.config b/examples/generative_rec_lm_s1pretrained.config index 807cc4bc1..d0f8d8658 100644 --- a/examples/generative_rec_lm_s1pretrained.config +++ b/examples/generative_rec_lm_s1pretrained.config @@ -140,11 +140,21 @@ feature_configs { } model_config { + # Two single-feature JAGGED_SEQUENCE groups: the model retrieves the raw SID + # streams via the EmbeddingGroup (build_input) and keys them by GROUP name + # (HSTU idiom). Group names match the common.history_group_name / + # label_group_name defaults ("user_seq"/"answer") and are SEPARATE from the + # feature names. Two groups (not one) because jagged-group members share nnz + # and user_sequence (≤300) ≠ label (3). feature_groups { - group_name: "sids" + group_name: "user_seq" feature_names: "user_sequence" + group_type: JAGGED_SEQUENCE + } + feature_groups { + group_name: "answer" feature_names: "label" - group_type: SEQUENCE + group_type: JAGGED_SEQUENCE } qwen2_rec_lm { # Qwen2 backbone (owned by this family message, not `common`). diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 306afb01f..6549939bb 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -37,6 +37,7 @@ from tzrec.datasets.utils import Batch from tzrec.features.feature import BaseFeature from tzrec.models.model import BaseModel +from tzrec.modules.embedding import EmbeddingGroup from tzrec.protos.model_pb2 import ModelConfig @@ -93,9 +94,9 @@ def __init__( self._build_prompt_tokens(tokenizer, cfg) - # Dense-only: the HF backbone owns its embeddings and SID ids flow through - # it directly (no EmbeddingGroup). Set explicit (vs RankModel.init_input). - self.embedding_group = None + # Build the (param-free, raw-passthrough) EmbeddingGroup for the SID + # JAGGED_SEQUENCE groups — the HSTU retrieval idiom (see init_input). + self.init_input() # one-shot debug dump of the first spliced batch self._smoke_log_once = os.environ.get("TZREC_GENRECLM_DEBUG", "0") == "1" @@ -121,6 +122,11 @@ def _read_common_config(self, common: Any) -> int: """Parse shared proto knobs into attributes; return the SID atom count.""" self._input_name: str = common.user_sequence_feature_name self._label_name: str = common.label_feature_name + # Which JAGGED_SEQUENCE feature_group carries each SID stream; build_input + # keys the EmbeddingGroup output by these GROUP names (HSTU idiom), + # decoupled from the feature names above. + self._history_group: str = common.history_group_name + self._label_group: str = common.label_group_name self._ignore_index: int = int(common.ignore_index) # max history (SID codes) for activation pre-sizing = the user-sequence # feature's sequence_length. FG_NONE doesn't truncate, so _sid_token_rows @@ -308,17 +314,57 @@ def _validate_sid_candidates( # groups beams per user; the row-wise mask above preserved that order. return sids.view(batch_size, -1, self._num_levels) + def init_input(self) -> None: + """Build the EmbeddingGroup for the raw SID JAGGED_SEQUENCE groups. + + Raw (passthrough) features carry no embedding tables, so this + EmbeddingGroup holds no params (DMP-neutral); it exists purely to + retrieve the raw SID sequences as flat ``(values, lengths)`` — the same + path HSTU uses. The HF backbone still owns the token embeddings; SID ids + flow through it directly (no embedding lookup here). + """ + self.embedding_group = EmbeddingGroup(self._features, self._feature_groups) + + def build_input(self, batch: Batch) -> Dict[str, List[torch.Tensor]]: + """Retrieve per-row SID token sequences via the EmbeddingGroup. + + ``embedding_group(batch)`` returns, per JAGGED_SEQUENCE group, the flat + raw values ``"{group}.sequence"`` + ``"{group}.sequence_length"`` (the + HSTU idiom). We map SID indices -> extended-vocab token ids and split to + rows. The EmbeddingGroup output is keyed by GROUP name + (``_history_group`` / ``_label_group``); the returned dict is keyed by + FEATURE name (what the family ``predict`` consumes). The label is omitted + in inference, where no ground truth is supplied. + """ + g = self.embedding_group(batch) + rows: Dict[str, List[torch.Tensor]] = { + self._input_name: self._sid_token_rows( + g[f"{self._history_group}.sequence"], + g[f"{self._history_group}.sequence_length"], + max_codes=self._max_seq_length, + ), + } + if not self.is_inference: + rows[self._label_name] = self._sid_token_rows( + g[f"{self._label_group}.sequence"], + g[f"{self._label_group}.sequence_length"], + expected_width=self._num_levels, + ) + return rows + def _sid_token_rows( self, - jt, + values: torch.Tensor, + lengths: torch.Tensor, expected_width: Optional[int] = None, max_codes: Optional[int] = None, ) -> List[torch.Tensor]: - """Read a SID jagged feature -> per-row token-id tensors. + """Map flat SID ``(values, lengths)`` -> per-row token-id tensors. - TER delivers the feature as a JaggedTensor (flat ``values`` + ``lengths``); - ``values`` may arrive as float / shape ``(N, 1)``. The whole batch is - tokenized once on the backbone device, then split into rows. + ``build_input`` supplies a JAGGED_SEQUENCE group's flat + ``"{group}.sequence"`` values + ``"{group}.sequence_length"``; ``values`` + may arrive as float / shape ``(N, 1)``. The whole batch is tokenized once + on the backbone device, then split into rows. ``expected_width``, when set, enforces the sample contract: every row must have exactly that many codes (the answer = ``num_levels``). @@ -328,8 +374,6 @@ def _sid_token_rows( oldest head) so the pre-allocated pool covers every batch. Done on host views before the H2D copy, skipped unless a row overflows. """ - values = jt.values() - lengths = jt.lengths() if values.dim() == 2 and values.size(-1) == 1: values = values.squeeze(-1) # host-side split bounds, read before the H2D copy below diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py index 934df831b..890f2cc3e 100644 --- a/tzrec/models/generative_rec_lm_test.py +++ b/tzrec/models/generative_rec_lm_test.py @@ -110,25 +110,67 @@ def test_tokenize_sids(self) -> None: def test_sid_token_rows_split_and_cast(self) -> None: m = _stub(base_vocab=100) - rows = m._sid_token_rows(_FakeJT([1, 2, 3, 4, 5], [3, 2])) + jt = _FakeJT([1, 2, 3, 4, 5], [3, 2]) + rows = m._sid_token_rows(jt.values(), jt.lengths()) self.assertEqual([r.tolist() for r in rows], [[100, 101, 102], [103, 104]]) self.assertTrue(all(r.dtype == torch.int64 for r in rows)) def test_sid_token_rows_squeezes_n1(self) -> None: m = _stub(base_vocab=100) - rows = m._sid_token_rows(_FakeJT([1, 2, 3], [3], dim2=True)) # (N, 1) + jt = _FakeJT([1, 2, 3], [3], dim2=True) # (N, 1) + rows = m._sid_token_rows(jt.values(), jt.lengths()) self.assertEqual([r.tolist() for r in rows], [[100, 101, 102]]) def test_sid_token_rows_width_ok(self) -> None: m = _stub(base_vocab=100, num_levels=3) - rows = m._sid_token_rows(_FakeJT([1, 2, 3, 4, 5, 6], [3, 3]), expected_width=3) + jt = _FakeJT([1, 2, 3, 4, 5, 6], [3, 3]) + rows = m._sid_token_rows(jt.values(), jt.lengths(), expected_width=3) self.assertEqual([r.tolist() for r in rows], [[100, 101, 102], [103, 104, 105]]) def test_sid_token_rows_width_violation_raises(self) -> None: m = _stub(base_vocab=100, num_levels=3) with self.assertRaises(ValueError): # second row has 2 codes, not 3 -> anomalous sample - m._sid_token_rows(_FakeJT([1, 2, 3, 4, 5], [3, 2]), expected_width=3) + jt = _FakeJT([1, 2, 3, 4, 5], [3, 2]) + m._sid_token_rows(jt.values(), jt.lengths(), expected_width=3) + + def test_build_input_keys_by_group_returns_by_feature(self) -> None: + # build_input reads the EmbeddingGroup output by GROUP name + # ("{group}.sequence" / ".sequence_length") and returns rows keyed by + # FEATURE name, tokenizing SID -> token id (sid + base - 1). + m = _stub(base_vocab=100, num_levels=3) + m._input_name, m._label_name = "user_sequence", "label" + m._history_group, m._label_group = "user_seq", "answer" + m._max_seq_length = 0 + m._is_inference = False # train: the answer is retrieved too + out = { + "user_seq.sequence": torch.tensor([1.0, 2.0, 3.0, 4.0]), + "user_seq.sequence_length": torch.tensor([2, 2]), + "answer.sequence": torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]), + "answer.sequence_length": torch.tensor([3, 3]), + } + m.embedding_group = lambda b: out + rows = m.build_input(object()) + self.assertEqual( + [r.tolist() for r in rows["user_sequence"]], [[100, 101], [102, 103]] + ) + self.assertEqual( + [r.tolist() for r in rows["label"]], [[100, 101, 102], [103, 104, 105]] + ) + + def test_build_input_skips_label_in_inference(self) -> None: + m = _stub(base_vocab=100, num_levels=3) + m._input_name, m._label_name = "user_sequence", "label" + m._history_group, m._label_group = "user_seq", "answer" + m._max_seq_length = 0 + m._is_inference = True # inference: history only, no ground-truth label + m.embedding_group = lambda b: { + "user_seq.sequence": torch.tensor([1.0, 2.0, 3.0]), + "user_seq.sequence_length": torch.tensor([3]), + } + rows = m.build_input(object()) + self.assertEqual([r.tolist() for r in rows["user_sequence"]], [[100, 101, 102]]) + self.assertNotIn("label", rows) if __name__ == "__main__": diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index 37f908526..588523ef9 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -213,15 +213,11 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: """Branch 1: teacher-forced forward -> suffix-slice -> CE loss.""" - # SID indices -> token ids once at the data boundary (_sid_token_rows). - u_rows = self._sid_token_rows( - batch.sequence_dense_features[self._input_name], - max_codes=self._max_seq_length, # cap to most-recent items (drop oldest) - ) - l_rows = self._sid_token_rows( - batch.sequence_dense_features[self._label_name], - expected_width=self._num_levels, # answer = one item = num_levels codes - ) + # Retrieve SID token rows via the EmbeddingGroup (build_input), keyed by + # feature name. Train needs both the history and the teacher-forced answer. + rows = self.build_input(batch) + u_rows = rows[self._input_name] + l_rows = rows[self._label_name] # One-shot pool pre-sizing: pad the FIRST train step to the worst-case # length so the allocator reserves its largest segments up front (no @@ -281,10 +277,8 @@ def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: ``_validate_sid_candidates`` (token->SID, malformed beams -> ``-1``). Returns ``generated_sids`` of shape ``(B, num_return, num_levels)``. """ - u_rows = self._sid_token_rows( - batch.sequence_dense_features[self._input_name], - max_codes=self._max_seq_length, # cap to most-recent items (drop oldest) - ) + # history rows via the EmbeddingGroup (build_input); inference skips label. + u_rows = self.build_input(batch)[self._input_name] input_ids, attention_mask = self._splice_prompt_ids(u_rows) if self._dynamic_beam: new_tokens = self._dynamic_beam_search(input_ids, attention_mask) diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/qwen2_rec_lm_test.py index 8017fe342..b8c4ac489 100644 --- a/tzrec/models/qwen2_rec_lm_test.py +++ b/tzrec/models/qwen2_rec_lm_test.py @@ -163,6 +163,9 @@ def fake_generate( return torch.cat([prompt, new], dim=1) m.lm.generate = fake_generate + # build_input (mocked) supplies the tokenized history rows; the batch is + # opaque to it. SIDs [1,2,3] tokenize to [100,101,102] at base_vocab=100. + m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 101, 102])]} sids = m._generate(_gen_batch())["generated_sids"] self.assertEqual(tuple(sids.shape), (1, 2, 3)) # (B, num_return, num_levels) self.assertEqual(sids[0].tolist(), [[1, 5, 9], [4, 8, 12]]) @@ -199,6 +202,9 @@ def fake_generate( return torch.cat([prompt, new], dim=1) m.lm.generate = fake_generate + # build_input (mocked) supplies the tokenized history rows; the batch is + # opaque to it. SIDs [1,2,3] tokenize to [100,101,102] at base_vocab=100. + m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 101, 102])]} sids = m._generate(_gen_batch())["generated_sids"] self.assertEqual(tuple(sids.shape), (1, 4, 3)) # valid candidate kept at its rank; every malformed one -> all -1 (in place) @@ -228,6 +234,9 @@ def fake_generate( return torch.cat([prompt, new], dim=1) m.lm.generate = fake_generate + # build_input (mocked) supplies the tokenized history rows; the batch is + # opaque to it. SIDs [1,2,3] tokenize to [100,101,102] at base_vocab=100. + m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 101, 102])]} sids = m._generate(_gen_batch())["generated_sids"] self.assertEqual(tuple(sids.shape), (1, 2, 3)) # rectangular, no crash # the missing 3rd atom stays -1 -> out of band -> whole candidate -1 @@ -278,23 +287,20 @@ def test_input_sequence_length_from_feature(self) -> None: def test_sid_token_rows_recency_clip(self) -> None: m = _stub(num_levels=3, base_vocab=100) # token = sid + base - 1 = sid + 99 - def _jt(n): # one row of n codes: values 1..n - return types.SimpleNamespace( - values=lambda: torch.arange(1, n + 1, dtype=torch.float), - lengths=lambda: torch.tensor([n]), - ) + def _vl(n): # one row of n codes (values 1..n) as flat (values, lengths) + return torch.arange(1, n + 1, dtype=torch.float), torch.tensor([n]) # 15 codes (5 items), cap 9 -> keep last 9 (items 3-5 = codes 7..15) - rows = m._sid_token_rows(_jt(15), max_codes=9) + rows = m._sid_token_rows(*_vl(15), max_codes=9) self.assertEqual(rows[0].tolist(), [c + 99 for c in range(7, 16)]) # item-aligned: cap 10 still keeps 9 (3 whole items), never cuts mid-item - rows = m._sid_token_rows(_jt(15), max_codes=10) + rows = m._sid_token_rows(*_vl(15), max_codes=10) self.assertEqual(rows[0].tolist(), [c + 99 for c in range(7, 16)]) # within cap -> untouched - rows = m._sid_token_rows(_jt(6), max_codes=9) + rows = m._sid_token_rows(*_vl(6), max_codes=9) self.assertEqual(rows[0].tolist(), [c + 99 for c in range(1, 7)]) # disabled (0/None) -> no clip - rows = m._sid_token_rows(_jt(15), max_codes=0) + rows = m._sid_token_rows(*_vl(15), max_codes=0) self.assertEqual(rows[0].tolist(), [c + 99 for c in range(1, 16)]) def test_compute_max_total_length(self) -> None: @@ -319,13 +325,12 @@ def fwd(i, lbl, a): seen_lens.append(i.shape[1]) return {"loss": torch.tensor(0.0)} - m._sid_token_rows = lambda jt, expected_width=None, max_codes=None: [ - torch.tensor([100, 101, 102]) - ] + m.build_input = lambda b: { + m._input_name: [torch.tensor([100, 101, 102])], + m._label_name: [torch.tensor([200, 201, 202])], + } m._forward_loss = fwd - batch = types.SimpleNamespace( - sequence_dense_features={"user_sequence": None, "label": None} - ) + batch = object() m._predict_train(batch) # first step: pre-size to worst case m._predict_train(batch) # subsequent step: natural length # one-shot: first step left-pads to _max_total_len, latched by @@ -347,13 +352,12 @@ def fwd(i, lbl, a): seen_lens.append(i.shape[1]) return {"loss": torch.tensor(0.0)} - m._sid_token_rows = lambda jt, expected_width=None, max_codes=None: [ - torch.tensor([100, 101, 102]) - ] + m.build_input = lambda b: { + m._input_name: [torch.tensor([100, 101, 102])], + m._label_name: [torch.tensor([200, 201, 202])], + } m._forward_loss = fwd - batch = types.SimpleNamespace( - sequence_dense_features={"user_sequence": None, "label": None} - ) + batch = object() m._predict_train(batch) # disabled: natural length, never forced to max; flag stays unlatched. self.assertLess(seen_lens[0], 50) diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index 268b6b832..9aaa3fcb5 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -46,6 +46,14 @@ message GenerativeRecLMConfig { // candidates (num_return_sequences is ignored). Faithful parity with ALGR's // dynamic_beams schedule; exploits the fixed-length, EOS-free SID answer. optional bool dynamic_beam = 9 [default = false]; + + // Which feature_group (group_type: JAGGED_SEQUENCE) carries the history / + // answer SID stream. The model keys the EmbeddingGroup output by GROUP name + // (like HSTU's grouped_features["candidate.sequence"]), decoupled from the + // feature names — a group may bundle >1 features. Defaults match the + // canonical example config's group names. + optional string history_group_name = 10 [default = "user_seq"]; + optional string label_group_name = 11 [default = "answer"]; } // Qwen2 / Qwen2.5 family (Qwen2.5-0.5B, etc.). Registered directly via its From 7c888a07a75962772b4e45481d8afbac2a67c8db Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 24 Jun 2026 09:06:05 +0000 Subject: [PATCH 22/99] [feat] generative-rec LM: make generated_sids_key + param_dtype proto-configurable Promote the two GenerativeRecLM class constants to GenerativeRecLMConfig knobs, keeping their previous values as defaults: - generated_sids_key (default "generated_sids") -> self._generated_sids_key, used by _generate's output dict. - param_dtype (default "float32") -> self._param_dtype via _DTYPE_BY_NAME {float32, bfloat16, float16}; used by _build_backbone / init_from_pretrained. An unknown value raises a clear ValueError. Both read in _read_common_config; tests cover defaults + dtype mapping + validation. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 35 +++++++++++++++------- tzrec/models/generative_rec_lm_test.py | 34 +++++++++++++++++++++ tzrec/models/qwen2_rec_lm.py | 2 +- tzrec/models/qwen2_rec_lm_test.py | 1 + tzrec/protos/models/generative_model.proto | 9 ++++++ 5 files changed, 70 insertions(+), 11 deletions(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 6549939bb..66972b232 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -62,14 +62,20 @@ class GenerativeRecLM(BaseModel): name). """ - # predictions key the inference branch emits generated SIDs under, stable - # across families (PredictWrapper ``output_cols`` should reference it). + # Default predictions key the inference branch emits generated SIDs under, + # stable across families (PredictWrapper ``output_cols`` should reference it). + # Overridable via ``common.generated_sids_key`` -> ``self._generated_sids_key``. GENERATED_SIDS_KEY = "generated_sids" - # Single source of truth for the backbone PARAM dtype: fp32 MASTER weights. - # The optimizer needs fp32 to avoid bf16-ULP underflow at small lr; bf16 - # *compute* comes from mixed_precision:"BF16" autocast, not the param dtype. - _PARAM_DTYPE = torch.float32 + # Backbone PARAM dtype options (the fp32 MASTER weights). fp32 avoids + # bf16-ULP underflow of Adam's small (lr=1e-5) updates; bf16 *compute* comes + # from mixed_precision:"BF16" autocast, NOT the param dtype. Selected by + # ``common.param_dtype`` (default "float32") -> ``self._param_dtype``. + _DTYPE_BY_NAME = { + "float32": torch.float32, + "bfloat16": torch.bfloat16, + "float16": torch.float16, + } def __init__( self, @@ -128,6 +134,15 @@ def _read_common_config(self, common: Any) -> int: self._history_group: str = common.history_group_name self._label_group: str = common.label_group_name self._ignore_index: int = int(common.ignore_index) + # Inference output key + backbone param dtype (configurable; default + # "generated_sids" / "float32" = the fp32-master weights). + self._generated_sids_key: str = common.generated_sids_key + if common.param_dtype not in self._DTYPE_BY_NAME: + raise ValueError( + f"{type(self).__name__}: param_dtype must be one of " + f"{list(self._DTYPE_BY_NAME)}, got {common.param_dtype!r}." + ) + self._param_dtype: torch.dtype = self._DTYPE_BY_NAME[common.param_dtype] # max history (SID codes) for activation pre-sizing = the user-sequence # feature's sequence_length. FG_NONE doesn't truncate, so _sid_token_rows # enforces this cap (item-aligned) model-side. 0 = off. @@ -160,9 +175,9 @@ def _build_backbone(self) -> Any: hf_cfg = AutoConfig.from_pretrained(hf_model_id) # fp32 MASTER weights: bf16 params underflow Adam's small (lr=1e-5) updates # and freeze. bf16 compute comes from mixed_precision:"BF16"; ckpt stays fp32. - lm = AutoModelForCausalLM.from_config(hf_cfg, torch_dtype=self._PARAM_DTYPE) - if next(lm.parameters()).dtype != self._PARAM_DTYPE: - lm = lm.to(self._PARAM_DTYPE) + lm = AutoModelForCausalLM.from_config(hf_cfg, torch_dtype=self._param_dtype) + if next(lm.parameters()).dtype != self._param_dtype: + lm = lm.to(self._param_dtype) return lm def _build_extended_tokenizer(self, sid_atoms: int) -> tuple[Any, int]: @@ -217,7 +232,7 @@ def init_from_pretrained(self) -> None: # fp32 master (not "auto", which keeps the stored bf16); must match # _build_backbone so cold-start and restore arches agree. lm = AutoModelForCausalLM.from_pretrained( - self._backbone_id(), torch_dtype=self._PARAM_DTYPE + self._backbone_id(), torch_dtype=self._param_dtype ) lm.resize_token_embeddings( self._target_vocab, pad_to_multiple_of=self._vocab_pad_mult diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py index 890f2cc3e..2f333b464 100644 --- a/tzrec/models/generative_rec_lm_test.py +++ b/tzrec/models/generative_rec_lm_test.py @@ -89,6 +89,40 @@ def test_backbone_owned_by_family_proto(self) -> None: common_fields = [f.name for f in GenerativeRecLMConfig.DESCRIPTOR.fields] self.assertNotIn("hf_model_id", common_fields) + def test_configurable_knob_defaults(self) -> None: + # generated_sids_key / param_dtype are proto knobs whose defaults are the + # previous class-constant values. + from tzrec.protos.models.generative_model_pb2 import GenerativeRecLMConfig + + c = GenerativeRecLMConfig() + self.assertEqual(c.generated_sids_key, "generated_sids") + self.assertEqual(c.param_dtype, "float32") + self.assertIs(Qwen2RecLM._DTYPE_BY_NAME["float32"], torch.float32) + self.assertIs(Qwen2RecLM._DTYPE_BY_NAME["bfloat16"], torch.bfloat16) + + def test_read_common_config_reads_knobs(self) -> None: + m = object.__new__(Qwen2RecLM) + nn.Module.__init__(m) + m._features = [] # _input_sequence_length -> 0 + common = types.SimpleNamespace( + user_sequence_feature_name="user_sequence", + label_feature_name="label", + history_group_name="user_seq", + label_group_name="answer", + ignore_index=-100, + generated_sids_key="my_sids", + param_dtype="bfloat16", + codebook=[4, 4, 4], + vocab_pad_to_multiple_of=128, + ) + m._read_common_config(common) + self.assertEqual(m._generated_sids_key, "my_sids") # configurable + self.assertIs(m._param_dtype, torch.bfloat16) # name -> torch dtype + # unknown dtype -> a clear error, not a KeyError + common.param_dtype = "float64" + with self.assertRaisesRegex(ValueError, "param_dtype must be one of"): + m._read_common_config(common) + def test_abstract_hooks_raise(self) -> None: base = object.__new__(GenerativeRecLM) with self.assertRaises(NotImplementedError): diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index 588523ef9..deda355f5 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -294,7 +294,7 @@ def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: ) new_tokens = out[:, input_ids.shape[1] :] # the generated tail sids = self._validate_sid_candidates(new_tokens, input_ids.shape[0]) - return {self.GENERATED_SIDS_KEY: sids} + return {self._generated_sids_key: sids} def _dynamic_beam_search( self, input_ids: torch.Tensor, attention_mask: torch.Tensor diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/qwen2_rec_lm_test.py index b8c4ac489..ece8a3247 100644 --- a/tzrec/models/qwen2_rec_lm_test.py +++ b/tzrec/models/qwen2_rec_lm_test.py @@ -38,6 +38,7 @@ def _stub(num_levels=3, base_vocab=100, pad_id=9, device="cpu", per_level=4): m._pad_token_id = pad_id m._dynamic_beam = False # default = HF fixed-width beam path m._max_seq_length = 0 # no recency clip by default in unit stubs + m._generated_sids_key = "generated_sids" # configurable; default key m.lm = types.SimpleNamespace(device=torch.device(device)) for name, vals in { "tpl_system": [10, 11], diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index 9aaa3fcb5..c92469954 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -54,6 +54,15 @@ message GenerativeRecLMConfig { // canonical example config's group names. optional string history_group_name = 10 [default = "user_seq"]; optional string label_group_name = 11 [default = "answer"]; + + // Predictions key the inference branch emits generated SIDs under (stable + // across families; PredictWrapper output_cols should reference it). + optional string generated_sids_key = 12 [default = "generated_sids"]; + // Backbone PARAM dtype = the fp32 MASTER weights. "float32" avoids bf16-ULP + // underflow of Adam's small (lr=1e-5) updates; bf16 COMPUTE comes from + // mixed_precision:"BF16" autocast, NOT the param dtype. One of: + // float32 | bfloat16 | float16. + optional string param_dtype = 13 [default = "float32"]; } // Qwen2 / Qwen2.5 family (Qwen2.5-0.5B, etc.). Registered directly via its From bc7c3f9d01f8efafb0d8b411337c97bdc44c353b Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 24 Jun 2026 09:07:31 +0000 Subject: [PATCH 23/99] [refactor] generative-rec LM: drop now-redundant GENERATED_SIDS_KEY constant The value now lives in the GenerativeRecLMConfig.generated_sids_key proto default ("generated_sids"); _generate already emits self._generated_sids_key. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 66972b232..0c6dbcaa3 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -62,11 +62,6 @@ class GenerativeRecLM(BaseModel): name). """ - # Default predictions key the inference branch emits generated SIDs under, - # stable across families (PredictWrapper ``output_cols`` should reference it). - # Overridable via ``common.generated_sids_key`` -> ``self._generated_sids_key``. - GENERATED_SIDS_KEY = "generated_sids" - # Backbone PARAM dtype options (the fp32 MASTER weights). fp32 avoids # bf16-ULP underflow of Adam's small (lr=1e-5) updates; bf16 *compute* comes # from mixed_precision:"BF16" autocast, NOT the param dtype. Selected by From 214d44d61758674b090a69f16876ed00f6157a27 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 24 Jun 2026 11:24:36 +0000 Subject: [PATCH 24/99] [refactor] generative-rec LM: collapse param_dtype double-lookup to one .get() /simplify cleanup: membership-check + dict-index -> single .get() + None-guard (same ValueError). No behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 0c6dbcaa3..c6cf06197 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -132,12 +132,13 @@ def _read_common_config(self, common: Any) -> int: # Inference output key + backbone param dtype (configurable; default # "generated_sids" / "float32" = the fp32-master weights). self._generated_sids_key: str = common.generated_sids_key - if common.param_dtype not in self._DTYPE_BY_NAME: + param_dtype = self._DTYPE_BY_NAME.get(common.param_dtype) + if param_dtype is None: raise ValueError( f"{type(self).__name__}: param_dtype must be one of " f"{list(self._DTYPE_BY_NAME)}, got {common.param_dtype!r}." ) - self._param_dtype: torch.dtype = self._DTYPE_BY_NAME[common.param_dtype] + self._param_dtype: torch.dtype = param_dtype # max history (SID codes) for activation pre-sizing = the user-sequence # feature's sequence_length. FG_NONE doesn't truncate, so _sid_token_rows # enforces this cap (item-aligned) model-side. 0 = off. From b71efce4f393491696e78a5d26e986b446b52d8c Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 24 Jun 2026 12:11:44 +0000 Subject: [PATCH 25/99] [feat] generative-rec LM: max_sequence_length as a model-config knob (HSTU-style) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Promote the model's history budget to GenerativeRecLMConfig.max_sequence_length (field 14), mirroring HSTU's DlrmHSTU.max_seq_len — a model knob distinct from the user_sequence feature's sequence_length. _max_seq_length now reads common.max_sequence_length, falling back to the feature's sequence_length when 0 (backward-compatible). It drives the recency-preserving truncation cap (_sid_token_rows) + the activation-pool pre-size. Example config sets it in common; tests cover the model-knob and the fallback. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../generative_rec_lm_s1pretrained.config | 6 ++++ tzrec/models/generative_rec_lm.py | 26 +++++++++----- tzrec/models/generative_rec_lm_test.py | 36 +++++++++++++++++++ tzrec/protos/models/generative_model.proto | 8 +++++ 4 files changed, 67 insertions(+), 9 deletions(-) diff --git a/examples/generative_rec_lm_s1pretrained.config b/examples/generative_rec_lm_s1pretrained.config index d0f8d8658..a7e005094 100644 --- a/examples/generative_rec_lm_s1pretrained.config +++ b/examples/generative_rec_lm_s1pretrained.config @@ -178,6 +178,12 @@ model_config { user_sequence_feature_name: "user_sequence" label_feature_name: "label" ignore_index: -100 + + # Model history budget (HSTU-style model knob): truncation cap + + # activation-pool pre-size. Distinct from the user_sequence feature's + # sequence_length (kept as data/export metadata); they coincide at 300 + # here. 0 would fall back to that feature's sequence_length. + max_sequence_length: 300 } # algr's row.system / `default_instruction` — the CN recommender prompt. diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index c6cf06197..4c56644ce 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -81,7 +81,7 @@ def __init__( **kwargs: Any, ) -> None: super().__init__(model_config, features, labels, sample_weights, **kwargs) - cfg = self._model_config # family message (e.g. Qwen2RecLM) + cfg = self._model_config common = cfg.common # GenerativeRecLMConfig — shared by all families sid_atoms = self._read_common_config(common) @@ -139,10 +139,14 @@ def _read_common_config(self, common: Any) -> int: f"{list(self._DTYPE_BY_NAME)}, got {common.param_dtype!r}." ) self._param_dtype: torch.dtype = param_dtype - # max history (SID codes) for activation pre-sizing = the user-sequence - # feature's sequence_length. FG_NONE doesn't truncate, so _sid_token_rows - # enforces this cap (item-aligned) model-side. 0 = off. - self._max_seq_length: int = self._input_sequence_length() + # Model's history budget (SID codes): the truncation cap (_sid_token_rows, + # item-aligned) AND the activation-pool pre-size. HSTU-style model knob + # (common.max_sequence_length); 0 -> fall back to the user_sequence + # feature's sequence_length (then off if that's also unset). FG_NONE does + # not truncate, so this cap is enforced model-side. + self._max_seq_length: int = ( + int(common.max_sequence_length) or self._input_sequence_length() + ) codebook = list(common.codebook) if len(codebook) == 0: raise ValueError("GenerativeRecLM: codebook must be non-empty.") @@ -238,10 +242,9 @@ def init_from_pretrained(self) -> None: def _input_sequence_length(self) -> int: """The user-sequence feature's ``sequence_length`` (SID codes), or 0. - FG_NONE does NOT truncate, so this is the cap ``_sid_token_rows`` enforces - model-side (item-aligned), and the upper bound the activation pool is - pre-sized to (see ``Qwen2RecLM._predict_train``). 0 if the feature has no - length cap, which disables pre-allocation. + The FALLBACK for ``_max_seq_length`` when ``common.max_sequence_length`` + is 0; the model knob takes precedence (see ``_read_common_config``). 0 if + the feature has no length cap, which disables the cap + pre-allocation. """ for feature in self._features: if feature.config.feature_name == self._input_name: @@ -397,6 +400,11 @@ def _sid_token_rows( f"{expected_width} codes (len(codebook)); rows {bad} have " f"{[sizes[i] for i in bad]} — anomalous sample(s)." ) + + # TODO: The truncation logic should not be placed here, but should be + # handled in FG. Since FG currently cannot control the truncation + # direction (it keeps the HEAD), this may result in truncating the most + # recent SIDs. Check it. if max_codes: keep = (max_codes // self._num_levels) * self._num_levels if keep and any(n > keep for n in sizes): diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py index 2f333b464..3465c72f0 100644 --- a/tzrec/models/generative_rec_lm_test.py +++ b/tzrec/models/generative_rec_lm_test.py @@ -114,15 +114,51 @@ def test_read_common_config_reads_knobs(self) -> None: param_dtype="bfloat16", codebook=[4, 4, 4], vocab_pad_to_multiple_of=128, + max_sequence_length=288, ) m._read_common_config(common) self.assertEqual(m._generated_sids_key, "my_sids") # configurable self.assertIs(m._param_dtype, torch.bfloat16) # name -> torch dtype + self.assertEqual(m._max_seq_length, 288) # model knob used # unknown dtype -> a clear error, not a KeyError common.param_dtype = "float64" with self.assertRaisesRegex(ValueError, "param_dtype must be one of"): m._read_common_config(common) + def test_max_sequence_length_model_knob_and_fallback(self) -> None: + # _max_seq_length comes from the model knob (HSTU-style); 0 falls back to + # the user_sequence feature's sequence_length. + def _common(max_seq): + return types.SimpleNamespace( + user_sequence_feature_name="user_sequence", + label_feature_name="label", + history_group_name="user_seq", + label_group_name="answer", + ignore_index=-100, + generated_sids_key="generated_sids", + param_dtype="float32", + codebook=[4, 4, 4], + vocab_pad_to_multiple_of=128, + max_sequence_length=max_seq, + ) + + f_user = types.SimpleNamespace( + config=types.SimpleNamespace(feature_name="user_sequence"), + sequence_length=256, + ) + # knob 0 -> fall back to the feature's sequence_length + m = object.__new__(Qwen2RecLM) + nn.Module.__init__(m) + m._features = [f_user] + m._read_common_config(_common(0)) + self.assertEqual(m._max_seq_length, 256) + # knob set -> overrides the feature + m2 = object.__new__(Qwen2RecLM) + nn.Module.__init__(m2) + m2._features = [f_user] + m2._read_common_config(_common(128)) + self.assertEqual(m2._max_seq_length, 128) + def test_abstract_hooks_raise(self) -> None: base = object.__new__(GenerativeRecLM) with self.assertRaises(NotImplementedError): diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index c92469954..e933154b7 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -63,6 +63,14 @@ message GenerativeRecLMConfig { // mixed_precision:"BF16" autocast, NOT the param dtype. One of: // float32 | bfloat16 | float16. optional string param_dtype = 13 [default = "float32"]; + + // Model's history budget (SID codes): the truncation cap enforced model-side + // (_sid_token_rows, item-aligned, recency-preserving) AND the activation-pool + // pre-size (_compute_max_total_length). A model knob like HSTU's + // DlrmHSTU.max_seq_len, distinct from the user_sequence feature's + // sequence_length. 0 = fall back to that feature's sequence_length (then off + // if it too is unset). + optional uint32 max_sequence_length = 14 [default = 0]; } // Qwen2 / Qwen2.5 family (Qwen2.5-0.5B, etc.). Registered directly via its From b8f20cb231bcd88dbf498c19d9beb0deee966b71 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 24 Jun 2026 12:17:53 +0000 Subject: [PATCH 26/99] [refactor] generative-rec LM: drop _input_sequence_length; max_sequence_length is the sole source common.max_sequence_length is now the single source for the model's history budget; remove the feature-derived fallback (_input_sequence_length) and its test. 0 = off (no cap / no activation-pool pre-allocation). Proto + example-config comments updated accordingly. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../generative_rec_lm_s1pretrained.config | 3 +-- tzrec/models/generative_rec_lm.py | 21 +++------------- tzrec/models/generative_rec_lm_test.py | 25 +++++++------------ tzrec/models/qwen2_rec_lm_test.py | 18 ------------- tzrec/protos/models/generative_model.proto | 3 +-- 5 files changed, 14 insertions(+), 56 deletions(-) diff --git a/examples/generative_rec_lm_s1pretrained.config b/examples/generative_rec_lm_s1pretrained.config index a7e005094..cdeeb0fcf 100644 --- a/examples/generative_rec_lm_s1pretrained.config +++ b/examples/generative_rec_lm_s1pretrained.config @@ -181,8 +181,7 @@ model_config { # Model history budget (HSTU-style model knob): truncation cap + # activation-pool pre-size. Distinct from the user_sequence feature's - # sequence_length (kept as data/export metadata); they coincide at 300 - # here. 0 would fall back to that feature's sequence_length. + # sequence_length (kept as data/export metadata). 0 = off (no cap). max_sequence_length: 300 } diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 4c56644ce..aec036e46 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -140,13 +140,10 @@ def _read_common_config(self, common: Any) -> int: ) self._param_dtype: torch.dtype = param_dtype # Model's history budget (SID codes): the truncation cap (_sid_token_rows, - # item-aligned) AND the activation-pool pre-size. HSTU-style model knob - # (common.max_sequence_length); 0 -> fall back to the user_sequence - # feature's sequence_length (then off if that's also unset). FG_NONE does + # item-aligned, recency-preserving) AND the activation-pool pre-size. + # HSTU-style model knob; 0 = off (no cap, no pre-allocation). FG_NONE does # not truncate, so this cap is enforced model-side. - self._max_seq_length: int = ( - int(common.max_sequence_length) or self._input_sequence_length() - ) + self._max_seq_length: int = int(common.max_sequence_length) codebook = list(common.codebook) if len(codebook) == 0: raise ValueError("GenerativeRecLM: codebook must be non-empty.") @@ -239,18 +236,6 @@ def init_from_pretrained(self) -> None: ) self.lm = lm - def _input_sequence_length(self) -> int: - """The user-sequence feature's ``sequence_length`` (SID codes), or 0. - - The FALLBACK for ``_max_seq_length`` when ``common.max_sequence_length`` - is 0; the model knob takes precedence (see ``_read_common_config``). 0 if - the feature has no length cap, which disables the cap + pre-allocation. - """ - for feature in self._features: - if feature.config.feature_name == self._input_name: - return int(getattr(feature, "sequence_length", 0) or 0) - return 0 - def hf_backbone(self): """The HF backbone module (``export_util.write_hf_assets``/``dcp_to_hf``).""" return self.lm diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py index 3465c72f0..363f3f2ce 100644 --- a/tzrec/models/generative_rec_lm_test.py +++ b/tzrec/models/generative_rec_lm_test.py @@ -103,7 +103,7 @@ def test_configurable_knob_defaults(self) -> None: def test_read_common_config_reads_knobs(self) -> None: m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) - m._features = [] # _input_sequence_length -> 0 + m._features = [] common = types.SimpleNamespace( user_sequence_feature_name="user_sequence", label_feature_name="label", @@ -125,9 +125,8 @@ def test_read_common_config_reads_knobs(self) -> None: with self.assertRaisesRegex(ValueError, "param_dtype must be one of"): m._read_common_config(common) - def test_max_sequence_length_model_knob_and_fallback(self) -> None: - # _max_seq_length comes from the model knob (HSTU-style); 0 falls back to - # the user_sequence feature's sequence_length. + def test_max_sequence_length_model_knob(self) -> None: + # _max_seq_length is the model knob; 0 = off (no feature fallback). def _common(max_seq): return types.SimpleNamespace( user_sequence_feature_name="user_sequence", @@ -142,22 +141,16 @@ def _common(max_seq): max_sequence_length=max_seq, ) - f_user = types.SimpleNamespace( - config=types.SimpleNamespace(feature_name="user_sequence"), - sequence_length=256, - ) - # knob 0 -> fall back to the feature's sequence_length m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) - m._features = [f_user] - m._read_common_config(_common(0)) - self.assertEqual(m._max_seq_length, 256) - # knob set -> overrides the feature + m._features = [] + m._read_common_config(_common(128)) + self.assertEqual(m._max_seq_length, 128) # model knob used m2 = object.__new__(Qwen2RecLM) nn.Module.__init__(m2) - m2._features = [f_user] - m2._read_common_config(_common(128)) - self.assertEqual(m2._max_seq_length, 128) + m2._features = [] + m2._read_common_config(_common(0)) + self.assertEqual(m2._max_seq_length, 0) # 0 = off, no fallback def test_abstract_hooks_raise(self) -> None: base = object.__new__(GenerativeRecLM) diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/qwen2_rec_lm_test.py index ece8a3247..efad14a42 100644 --- a/tzrec/models/qwen2_rec_lm_test.py +++ b/tzrec/models/qwen2_rec_lm_test.py @@ -267,24 +267,6 @@ def test_build_prompt_tokens_registers_buffers(self) -> None: self.assertEqual(buf.dtype, torch.int64) self.assertEqual(m.tpl_eos.tolist(), [99]) # eos cached for supervision - def test_input_sequence_length_from_feature(self) -> None: - m = object.__new__(Qwen2RecLM) - m._input_name = "user_sequence" - f_user = types.SimpleNamespace( - config=types.SimpleNamespace(feature_name="user_sequence"), - sequence_length=300, - ) - f_label = types.SimpleNamespace( - config=types.SimpleNamespace(feature_name="label"), sequence_length=32 - ) - m._features = [f_label, f_user] - self.assertEqual(m._input_sequence_length(), 300) # truncation length - m._features = [f_label] # user-sequence feature absent - self.assertEqual(m._input_sequence_length(), 0) - f_user.sequence_length = None # no length cap -> pre-allocation disabled - m._features = [f_user] - self.assertEqual(m._input_sequence_length(), 0) - def test_sid_token_rows_recency_clip(self) -> None: m = _stub(num_levels=3, base_vocab=100) # token = sid + base - 1 = sid + 99 diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index e933154b7..84b49d302 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -68,8 +68,7 @@ message GenerativeRecLMConfig { // (_sid_token_rows, item-aligned, recency-preserving) AND the activation-pool // pre-size (_compute_max_total_length). A model knob like HSTU's // DlrmHSTU.max_seq_len, distinct from the user_sequence feature's - // sequence_length. 0 = fall back to that feature's sequence_length (then off - // if it too is unset). + // sequence_length. 0 = off (no cap, no pre-allocation). optional uint32 max_sequence_length = 14 [default = 0]; } From 15bc2c02b8a07e39c758c789a3075d432cb2a400 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Wed, 24 Jun 2026 12:25:31 +0000 Subject: [PATCH 27/99] [feat] generative-rec LM: make max_sequence_length required; migrate s2 example config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit max_sequence_length is now a REQUIRED GenerativeRecLMConfig field (like HSTU's DlrmHSTU.max_seq_len) — every genrec config must set it (0 = explicitly off). Migrate examples/generative_rec_lm_s2pretrained.config to the new format to keep it valid + consistent: SEQUENCE group -> two JAGGED_SEQUENCE groups (user_seq / answer) + max_sequence_length: 1056. Both example configs verified to parse. Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/generative_rec_lm_s2pretrained.config | 14 ++++++++++++-- tzrec/protos/models/generative_model.proto | 5 +++-- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/examples/generative_rec_lm_s2pretrained.config b/examples/generative_rec_lm_s2pretrained.config index 05309e013..ec1cfd7fa 100644 --- a/examples/generative_rec_lm_s2pretrained.config +++ b/examples/generative_rec_lm_s2pretrained.config @@ -56,11 +56,18 @@ feature_configs { } } model_config { + # Two single-feature JAGGED_SEQUENCE groups (raw SID passthrough; see + # s1pretrained config / build_input). Names match history_group_name / + # label_group_name defaults. feature_groups { - group_name: "sids" + group_name: "user_seq" feature_names: "user_sequence" + group_type: JAGGED_SEQUENCE + } + feature_groups { + group_name: "answer" feature_names: "label" - group_type: SEQUENCE + group_type: JAGGED_SEQUENCE } qwen2_rec_lm { # Qwen2 backbone (owned by this family message, not `common`). @@ -76,6 +83,9 @@ model_config { codebook: 8192 codebook: 8192 vocab_pad_to_multiple_of: 128 + # Model history budget (required): truncation cap + activation-pool pre-size + # (matches the user_sequence feature's sequence_length here). + max_sequence_length: 1056 } system_instruction: "你是一个推荐系统，根据用户的历史行为，预测用户在电商场景的下一步行为。我会给你一串连续行为的语义编码，按照用户点击的时间顺序排列，每个行为用三个词表示。" user_prefix_text: "当前用户的历史行为如下：" diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index 84b49d302..901ea9508 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -68,8 +68,9 @@ message GenerativeRecLMConfig { // (_sid_token_rows, item-aligned, recency-preserving) AND the activation-pool // pre-size (_compute_max_total_length). A model knob like HSTU's // DlrmHSTU.max_seq_len, distinct from the user_sequence feature's - // sequence_length. 0 = off (no cap, no pre-allocation). - optional uint32 max_sequence_length = 14 [default = 0]; + // sequence_length. REQUIRED (like HSTU's DlrmHSTU.max_seq_len); set it to 0 + // to explicitly disable the cap + activation-pool pre-allocation. + required uint32 max_sequence_length = 14; } // Qwen2 / Qwen2.5 family (Qwen2.5-0.5B, etc.). Registered directly via its From b373e9e52fc901f549408e807a550050816cc97d Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 25 Jun 2026 06:59:20 +0000 Subject: [PATCH 28/99] [refactor] generative-rec LM: answer is a data_config.label_field, not a feature Move the answer/target SID stream out of feature_configs into data_config.label_fields (a list column -> batch.jagged_labels[label]). build_input now reads HISTORY from the user_seq JAGGED_SEQUENCE group (EmbeddingGroup) and the ANSWER from batch.jagged_labels[self._label_name]. Drop the now-unused label_group_name proto knob (reserved 11). This is semantically correct (the answer is the target, not an input feature) and lets the label be absent at inference without the EmbeddingGroup requiring it. Example configs (s1/s2) migrated; tests updated. (Also wrapped the hf_backbone/hf_tokenizer export-only docstrings.) Co-Authored-By: Claude Opus 4.8 (1M context) --- .../generative_rec_lm_s1pretrained.config | 40 ++++++------------- .../generative_rec_lm_s2pretrained.config | 22 +++------- tzrec/models/generative_rec_lm.py | 40 +++++++++++-------- tzrec/models/generative_rec_lm_test.py | 29 +++++++------- tzrec/protos/models/generative_model.proto | 16 ++++---- 5 files changed, 65 insertions(+), 82 deletions(-) diff --git a/examples/generative_rec_lm_s1pretrained.config b/examples/generative_rec_lm_s1pretrained.config index cdeeb0fcf..1794d6ffc 100644 --- a/examples/generative_rec_lm_s1pretrained.config +++ b/examples/generative_rec_lm_s1pretrained.config @@ -110,52 +110,38 @@ data_config { fg_mode: FG_NONE # algr `dataloader_num_workers: 4`. num_workers: 4 + # The answer SID column is the TARGET, so it's a label_field (a list + # column -> batch.jagged_labels["label"]), NOT a feature. build_input reads it + # via common.label_feature_name. Being a label_field, it can be absent at + # inference (no ground truth) without the EmbeddingGroup needing it. + label_fields: ["label"] } -# The two sequence features map 1:1 to the parquet columns produced by -# `convert_s1tiny_to_parquet.py`. +# user_sequence (the history INPUT) maps 1:1 to the parquet column produced by +# `convert_s1tiny_to_parquet.py`. (The "label" column is a label_field, above.) feature_configs { sequence_raw_feature { feature_name: "user_sequence" expression: "user:user_sequence" # Under FG_NONE this length is NOT auto-truncated by the reader, so the # model enforces it via a recency-preserving clip in `_sid_token_rows` - # (keep newest items, drop oldest) AND uses it to pre-size the CUDA - # activation pool (`_input_sequence_length`/`_warmup_alloc`). Set it to - # the data's realistic max history in codes (AL-GR-Tiny = 100 items × - # 3 = 300); a loose value (e.g. algr's 1056) would oversize the warm-up. + # (keep newest items, drop oldest). The model's cap is common. + # max_sequence_length (below), not this field. sequence_length: 300 value_dim: 1 } } -feature_configs { - sequence_raw_feature { - feature_name: "label" - expression: "user:label" - # algr `max_target_length: 32`. Each item is 3 SIDs, so 32 covers - # up to 10 items; in practice labels are always 3 SIDs in this set. - sequence_length: 32 - value_dim: 1 - } -} model_config { - # Two single-feature JAGGED_SEQUENCE groups: the model retrieves the raw SID - # streams via the EmbeddingGroup (build_input) and keys them by GROUP name - # (HSTU idiom). Group names match the common.history_group_name / - # label_group_name defaults ("user_seq"/"answer") and are SEPARATE from the - # feature names. Two groups (not one) because jagged-group members share nnz - # and user_sequence (≤300) ≠ label (3). + # One JAGGED_SEQUENCE group for the history INPUT: build_input retrieves the + # raw SID stream via the EmbeddingGroup and keys it by GROUP name (HSTU idiom); + # the name matches common.history_group_name default ("user_seq"). The answer + # is NOT a group — it's the label_field above (batch.jagged_labels). feature_groups { group_name: "user_seq" feature_names: "user_sequence" group_type: JAGGED_SEQUENCE } - feature_groups { - group_name: "answer" - feature_names: "label" - group_type: JAGGED_SEQUENCE - } qwen2_rec_lm { # Qwen2 backbone (owned by this family message, not `common`). # algr `load_checkpoint_from: /home/admin/workspace/Qwen2.5-0.5B` diff --git a/examples/generative_rec_lm_s2pretrained.config b/examples/generative_rec_lm_s2pretrained.config index ec1cfd7fa..a58e9af97 100644 --- a/examples/generative_rec_lm_s2pretrained.config +++ b/examples/generative_rec_lm_s2pretrained.config @@ -38,6 +38,9 @@ data_config { dataset_type: ParquetDataset num_workers: 4 fg_mode: FG_NONE + # answer = TARGET -> label_field (list -> batch.jagged_labels["label"]), + # not a feature; read by build_input via common.label_feature_name. + label_fields: ["label"] } feature_configs { sequence_raw_feature { @@ -47,28 +50,15 @@ feature_configs { sequence_length: 1056 } } -feature_configs { - sequence_raw_feature { - feature_name: "label" - expression: "user:label" - value_dim: 1 - sequence_length: 32 - } -} model_config { - # Two single-feature JAGGED_SEQUENCE groups (raw SID passthrough; see - # s1pretrained config / build_input). Names match history_group_name / - # label_group_name defaults. + # One JAGGED_SEQUENCE group for the history INPUT (raw SID passthrough; see + # s1pretrained / build_input). Name matches history_group_name default. The + # answer is the label_field above (batch.jagged_labels), not a group. feature_groups { group_name: "user_seq" feature_names: "user_sequence" group_type: JAGGED_SEQUENCE } - feature_groups { - group_name: "answer" - feature_names: "label" - group_type: JAGGED_SEQUENCE - } qwen2_rec_lm { # Qwen2 backbone (owned by this family message, not `common`). hf_model_id: "/home/admin/workspace/Qwen2.5-0.5B" diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index aec036e46..ebfc9097a 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -123,11 +123,11 @@ def _read_common_config(self, common: Any) -> int: """Parse shared proto knobs into attributes; return the SID atom count.""" self._input_name: str = common.user_sequence_feature_name self._label_name: str = common.label_feature_name - # Which JAGGED_SEQUENCE feature_group carries each SID stream; build_input - # keys the EmbeddingGroup output by these GROUP names (HSTU idiom), - # decoupled from the feature names above. + # The history is a JAGGED_SEQUENCE feature_group; build_input keys the + # EmbeddingGroup output by this GROUP name (HSTU idiom). The answer is NOT + # a feature — it's a data_config.label_field read from batch.jagged_labels + # by self._label_name (see build_input). self._history_group: str = common.history_group_name - self._label_group: str = common.label_group_name self._ignore_index: int = int(common.ignore_index) # Inference output key + backbone param dtype (configurable; default # "generated_sids" / "float32" = the fp32-master weights). @@ -237,11 +237,14 @@ def init_from_pretrained(self) -> None: self.lm = lm def hf_backbone(self): - """The HF backbone module (``export_util.write_hf_assets``/``dcp_to_hf``).""" + """The HF backbone module. Only use when export.""" return self.lm def hf_tokenizer(self): - """The extended tokenizer (base vocab + C0..C{sum-1}) to serialize.""" + """The extended tokenizer (base vocab + C0..C{sum-1}) to serialize. + + Only used at export (export_util.write_hf_assets). + """ return self._hf_tokenizer def _build_prompt_tokens(self, tokenizer, cfg) -> None: @@ -325,15 +328,17 @@ def init_input(self) -> None: self.embedding_group = EmbeddingGroup(self._features, self._feature_groups) def build_input(self, batch: Batch) -> Dict[str, List[torch.Tensor]]: - """Retrieve per-row SID token sequences via the EmbeddingGroup. - - ``embedding_group(batch)`` returns, per JAGGED_SEQUENCE group, the flat - raw values ``"{group}.sequence"`` + ``"{group}.sequence_length"`` (the - HSTU idiom). We map SID indices -> extended-vocab token ids and split to - rows. The EmbeddingGroup output is keyed by GROUP name - (``_history_group`` / ``_label_group``); the returned dict is keyed by - FEATURE name (what the family ``predict`` consumes). The label is omitted - in inference, where no ground truth is supplied. + """Retrieve per-row SID token sequences. + + HISTORY is a JAGGED_SEQUENCE feature_group: ``embedding_group(batch)`` + returns its flat raw values ``"{group}.sequence"`` + + ``"{group}.sequence_length"`` (the HSTU idiom). The ANSWER is a + data_config.label_field: its JaggedTensor comes from + ``batch.jagged_labels[self._label_name]`` (so it can be absent at + inference, where no ground truth is supplied — unlike a feature_group, + which the EmbeddingGroup would require every forward). Both are mapped + SID -> extended-vocab token ids and split to rows; the returned dict is + keyed by FEATURE name (what the family ``predict`` consumes). """ g = self.embedding_group(batch) rows: Dict[str, List[torch.Tensor]] = { @@ -344,9 +349,10 @@ def build_input(self, batch: Batch) -> Dict[str, List[torch.Tensor]]: ), } if not self.is_inference: + jt = batch.jagged_labels[self._label_name] rows[self._label_name] = self._sid_token_rows( - g[f"{self._label_group}.sequence"], - g[f"{self._label_group}.sequence_length"], + jt.values(), + jt.lengths(), expected_width=self._num_levels, ) return rows diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py index 363f3f2ce..ad4c2ec8b 100644 --- a/tzrec/models/generative_rec_lm_test.py +++ b/tzrec/models/generative_rec_lm_test.py @@ -108,7 +108,6 @@ def test_read_common_config_reads_knobs(self) -> None: user_sequence_feature_name="user_sequence", label_feature_name="label", history_group_name="user_seq", - label_group_name="answer", ignore_index=-100, generated_sids_key="my_sids", param_dtype="bfloat16", @@ -132,7 +131,6 @@ def _common(max_seq): user_sequence_feature_name="user_sequence", label_feature_name="label", history_group_name="user_seq", - label_group_name="answer", ignore_index=-100, generated_sids_key="generated_sids", param_dtype="float32", @@ -197,23 +195,23 @@ def test_sid_token_rows_width_violation_raises(self) -> None: jt = _FakeJT([1, 2, 3, 4, 5], [3, 2]) m._sid_token_rows(jt.values(), jt.lengths(), expected_width=3) - def test_build_input_keys_by_group_returns_by_feature(self) -> None: - # build_input reads the EmbeddingGroup output by GROUP name - # ("{group}.sequence" / ".sequence_length") and returns rows keyed by - # FEATURE name, tokenizing SID -> token id (sid + base - 1). + def test_build_input_history_group_label_field(self) -> None: + # history: EmbeddingGroup output keyed by GROUP name ("{group}.sequence" + # / ".sequence_length"). answer: batch.jagged_labels[label_name]. Both + # tokenized (sid -> sid + base - 1); returned dict keyed by FEATURE name. m = _stub(base_vocab=100, num_levels=3) m._input_name, m._label_name = "user_sequence", "label" - m._history_group, m._label_group = "user_seq", "answer" + m._history_group = "user_seq" m._max_seq_length = 0 - m._is_inference = False # train: the answer is retrieved too - out = { + m._is_inference = False # train: the answer label_field is read too + m.embedding_group = lambda b: { "user_seq.sequence": torch.tensor([1.0, 2.0, 3.0, 4.0]), "user_seq.sequence_length": torch.tensor([2, 2]), - "answer.sequence": torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]), - "answer.sequence_length": torch.tensor([3, 3]), } - m.embedding_group = lambda b: out - rows = m.build_input(object()) + batch = types.SimpleNamespace( + jagged_labels={"label": _FakeJT([1, 2, 3, 4, 5, 6], [3, 3])} + ) + rows = m.build_input(batch) self.assertEqual( [r.tolist() for r in rows["user_sequence"]], [[100, 101], [102, 103]] ) @@ -224,14 +222,15 @@ def test_build_input_keys_by_group_returns_by_feature(self) -> None: def test_build_input_skips_label_in_inference(self) -> None: m = _stub(base_vocab=100, num_levels=3) m._input_name, m._label_name = "user_sequence", "label" - m._history_group, m._label_group = "user_seq", "answer" + m._history_group = "user_seq" m._max_seq_length = 0 m._is_inference = True # inference: history only, no ground-truth label m.embedding_group = lambda b: { "user_seq.sequence": torch.tensor([1.0, 2.0, 3.0]), "user_seq.sequence_length": torch.tensor([3]), } - rows = m.build_input(object()) + # jagged_labels intentionally empty — the label is absent at inference + rows = m.build_input(types.SimpleNamespace(jagged_labels={})) self.assertEqual([r.tolist() for r in rows["user_sequence"]], [[100, 101, 102]]) self.assertNotIn("label", rows) diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index 901ea9508..17fb62886 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -30,7 +30,9 @@ message GenerativeRecLMConfig { // Pad the post-extension vocab up to a multiple of this value. optional uint32 vocab_pad_to_multiple_of = 3 [default = 128]; - // Which `sequence_raw_feature` parquet columns carry the SID lists. + // Names of the SID-list columns. user_sequence is a `sequence_raw_feature` + // (an input, in feature_configs / a JAGGED_SEQUENCE group); label is a + // `data_config.label_field` (the target) read from batch.jagged_labels. required string user_sequence_feature_name = 4; required string label_feature_name = 5; @@ -47,13 +49,13 @@ message GenerativeRecLMConfig { // dynamic_beams schedule; exploits the fixed-length, EOS-free SID answer. optional bool dynamic_beam = 9 [default = false]; - // Which feature_group (group_type: JAGGED_SEQUENCE) carries the history / - // answer SID stream. The model keys the EmbeddingGroup output by GROUP name - // (like HSTU's grouped_features["candidate.sequence"]), decoupled from the - // feature names — a group may bundle >1 features. Defaults match the - // canonical example config's group names. + // Which feature_group (group_type: JAGGED_SEQUENCE) carries the HISTORY SID + // stream. The model keys the EmbeddingGroup output by GROUP name (like HSTU's + // grouped_features["candidate.sequence"]), decoupled from the feature names. + // (The answer is NOT a feature_group: it's a data_config.label_field read from + // batch.jagged_labels — see label_feature_name.) optional string history_group_name = 10 [default = "user_seq"]; - optional string label_group_name = 11 [default = "answer"]; + reserved 11; // was label_group_name (answer moved to data_config.label_fields) // Predictions key the inference branch emits generated SIDs under (stable // across families; PredictWrapper output_cols should reference it). From 0730645625bc08713fec0228b789a0038c51ad84 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 25 Jun 2026 07:41:53 +0000 Subject: [PATCH 29/99] [refactor] genreclm: derive SID-column names instead of restating them common.user_sequence_feature_name / label_feature_name were redundant with config that already names these columns: * history feature -> the single member of the history_group (feature_groups[history_group].feature_names[0]) * answer label -> the first data_config.label_field (self._labels[0], the RankModel idiom) Drop both proto knobs (reserved 4, 5); _read_common_config now derives them via a new _history_feature_name() helper (fails loudly on a misconfigured history_group). Configs only declare the history feature_group + label_fields. Tests cover the derivation + the unknown-history-group error path. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../generative_rec_lm_s1pretrained.config | 13 +++-- .../generative_rec_lm_s2pretrained.config | 6 +- tzrec/models/generative_rec_lm.py | 32 +++++++++-- tzrec/models/generative_rec_lm_test.py | 55 +++++++++++++++---- tzrec/protos/models/generative_model.proto | 16 +++--- 5 files changed, 92 insertions(+), 30 deletions(-) diff --git a/examples/generative_rec_lm_s1pretrained.config b/examples/generative_rec_lm_s1pretrained.config index 1794d6ffc..0874ef50f 100644 --- a/examples/generative_rec_lm_s1pretrained.config +++ b/examples/generative_rec_lm_s1pretrained.config @@ -111,9 +111,9 @@ data_config { # algr `dataloader_num_workers: 4`. num_workers: 4 # The answer SID column is the TARGET, so it's a label_field (a list - # column -> batch.jagged_labels["label"]), NOT a feature. build_input reads it - # via common.label_feature_name. Being a label_field, it can be absent at - # inference (no ground truth) without the EmbeddingGroup needing it. + # column -> batch.jagged_labels["label"]), NOT a feature. build_input reads the + # FIRST label_field (like RankModel's labels[0]). Being a label_field, it can be + # absent at inference (no ground truth) without the EmbeddingGroup needing it. label_fields: ["label"] } @@ -160,9 +160,10 @@ model_config { codebook: 8192 vocab_pad_to_multiple_of: 128 - # Sample feature names — match parquet columns. - user_sequence_feature_name: "user_sequence" - label_feature_name: "label" + # SID-column names are derived, not configured: the history feature is + # the single member of the history_group below; the answer is the first + # data_config.label_field. Only the history GROUP needs naming + # (history_group_name default "user_seq" matches the group below). ignore_index: -100 # Model history budget (HSTU-style model knob): truncation cap + diff --git a/examples/generative_rec_lm_s2pretrained.config b/examples/generative_rec_lm_s2pretrained.config index a58e9af97..cf8f813a9 100644 --- a/examples/generative_rec_lm_s2pretrained.config +++ b/examples/generative_rec_lm_s2pretrained.config @@ -39,7 +39,7 @@ data_config { num_workers: 4 fg_mode: FG_NONE # answer = TARGET -> label_field (list -> batch.jagged_labels["label"]), - # not a feature; read by build_input via common.label_feature_name. + # not a feature; build_input reads the first label_field (like RankModel labels[0]). label_fields: ["label"] } feature_configs { @@ -63,8 +63,8 @@ model_config { # Qwen2 backbone (owned by this family message, not `common`). hf_model_id: "/home/admin/workspace/Qwen2.5-0.5B" common { - user_sequence_feature_name: "user_sequence" - label_feature_name: "label" + # SID-column names are derived (history = the history_group's single feature; + # answer = the first data_config.label_field), not configured here. ignore_index: -100 # SID codebook: one entry per RQ level (AL-GR = 3 levels x 8192; verified # from item_info codebook_lv* ranges). len = SID codes/answer, sum = vocab diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index ebfc9097a..7d29e9f22 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -121,13 +121,16 @@ def _resolve_pad_token_id(tokenizer: Any) -> int: def _read_common_config(self, common: Any) -> int: """Parse shared proto knobs into attributes; return the SID atom count.""" - self._input_name: str = common.user_sequence_feature_name - self._label_name: str = common.label_feature_name # The history is a JAGGED_SEQUENCE feature_group; build_input keys the # EmbeddingGroup output by this GROUP name (HSTU idiom). The answer is NOT - # a feature — it's a data_config.label_field read from batch.jagged_labels - # by self._label_name (see build_input). + # a feature — it's a data_config.label_field read from batch.jagged_labels. self._history_group: str = common.history_group_name + # SID-column names are derived, not restated: the history feature is the + # single member of the history group; the answer is the first + # data_config.label_field (like RankModel's labels[0]). _input_name is only + # the rows-dict key; _label_name also indexes batch.jagged_labels. + self._input_name: str = self._history_feature_name() + self._label_name: str = self._labels[0] if self._labels else "" self._ignore_index: int = int(common.ignore_index) # Inference output key + backbone param dtype (configurable; default # "generated_sids" / "float32" = the fp32-master weights). @@ -157,6 +160,27 @@ def _read_common_config(self, common: Any) -> int: self._vocab_pad_mult = int(common.vocab_pad_to_multiple_of) or 128 return sum(int(c) for c in codebook) + def _history_feature_name(self) -> str: + """The history feature's name = the single member of the history group. + + The history JAGGED_SEQUENCE group carries one SID stream, so its feature + name is ``feature_names[0]``; build_input uses it only as the rows-dict + key. Fails loudly on a misconfigured group (vs a silent KeyError later). + """ + for g in self._feature_groups: + if g.group_name == self._history_group: + if not g.feature_names: + raise ValueError( + f"{type(self).__name__}: history feature_group " + f"{self._history_group!r} has no feature_names." + ) + return g.feature_names[0] + raise ValueError( + f"{type(self).__name__}: history_group_name {self._history_group!r} " + f"matches no feature_group; declare a JAGGED_SEQUENCE group with that " + f"group_name." + ) + def _build_backbone(self) -> Any: """Build the EMPTY extended architecture in fp32 (master) — no download. diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py index ad4c2ec8b..2a9c62bdd 100644 --- a/tzrec/models/generative_rec_lm_test.py +++ b/tzrec/models/generative_rec_lm_test.py @@ -104,9 +104,15 @@ def test_read_common_config_reads_knobs(self) -> None: m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) m._features = [] + # SID-column names are DERIVED, not configured: history = the history + # group's single feature; answer = the first data_config.label_field. + m._labels = ["label"] + m._feature_groups = [ + types.SimpleNamespace( + group_name="user_seq", feature_names=["user_sequence"] + ) + ] common = types.SimpleNamespace( - user_sequence_feature_name="user_sequence", - label_feature_name="label", history_group_name="user_seq", ignore_index=-100, generated_sids_key="my_sids", @@ -116,6 +122,8 @@ def test_read_common_config_reads_knobs(self) -> None: max_sequence_length=288, ) m._read_common_config(common) + self.assertEqual(m._input_name, "user_sequence") # from history group + self.assertEqual(m._label_name, "label") # from label_fields[0] self.assertEqual(m._generated_sids_key, "my_sids") # configurable self.assertIs(m._param_dtype, torch.bfloat16) # name -> torch dtype self.assertEqual(m._max_seq_length, 288) # model knob used @@ -124,12 +132,31 @@ def test_read_common_config_reads_knobs(self) -> None: with self.assertRaisesRegex(ValueError, "param_dtype must be one of"): m._read_common_config(common) + def test_read_common_config_unknown_history_group_raises(self) -> None: + # history_group_name must match a declared feature_group, else fail loudly. + m = object.__new__(Qwen2RecLM) + nn.Module.__init__(m) + m._features = [] + m._labels = ["label"] + m._feature_groups = [ + types.SimpleNamespace(group_name="other", feature_names=["x"]) + ] + common = types.SimpleNamespace( + history_group_name="user_seq", + ignore_index=-100, + generated_sids_key="generated_sids", + param_dtype="float32", + codebook=[4, 4, 4], + vocab_pad_to_multiple_of=128, + max_sequence_length=0, + ) + with self.assertRaisesRegex(ValueError, "matches no feature_group"): + m._read_common_config(common) + def test_max_sequence_length_model_knob(self) -> None: # _max_seq_length is the model knob; 0 = off (no feature fallback). def _common(max_seq): return types.SimpleNamespace( - user_sequence_feature_name="user_sequence", - label_feature_name="label", history_group_name="user_seq", ignore_index=-100, generated_sids_key="generated_sids", @@ -139,14 +166,22 @@ def _common(max_seq): max_sequence_length=max_seq, ) - m = object.__new__(Qwen2RecLM) - nn.Module.__init__(m) - m._features = [] + def _wired(): + m = object.__new__(Qwen2RecLM) + nn.Module.__init__(m) + m._features = [] + m._labels = ["label"] + m._feature_groups = [ + types.SimpleNamespace( + group_name="user_seq", feature_names=["user_sequence"] + ) + ] + return m + + m = _wired() m._read_common_config(_common(128)) self.assertEqual(m._max_seq_length, 128) # model knob used - m2 = object.__new__(Qwen2RecLM) - nn.Module.__init__(m2) - m2._features = [] + m2 = _wired() m2._read_common_config(_common(0)) self.assertEqual(m2._max_seq_length, 0) # 0 = off, no fallback diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index 17fb62886..5174baa36 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -30,11 +30,13 @@ message GenerativeRecLMConfig { // Pad the post-extension vocab up to a multiple of this value. optional uint32 vocab_pad_to_multiple_of = 3 [default = 128]; - // Names of the SID-list columns. user_sequence is a `sequence_raw_feature` - // (an input, in feature_configs / a JAGGED_SEQUENCE group); label is a - // `data_config.label_field` (the target) read from batch.jagged_labels. - required string user_sequence_feature_name = 4; - required string label_feature_name = 5; + // SID-column names are NOT restated here — they're derived from the config + // that already declares them: the history feature is the single member of the + // history_group (its feature_names[0]); the answer is the first + // data_config.label_field (read from batch.jagged_labels, like RankModel's + // labels[0]). + reserved 4; // was user_sequence_feature_name (-> history_group's feature) + reserved 5; // was label_feature_name (-> data_config.label_fields[0]) // Cross-entropy ignore index — matches PyTorch's F.cross_entropy default. optional int32 ignore_index = 6 [default = -100]; @@ -52,8 +54,8 @@ message GenerativeRecLMConfig { // Which feature_group (group_type: JAGGED_SEQUENCE) carries the HISTORY SID // stream. The model keys the EmbeddingGroup output by GROUP name (like HSTU's // grouped_features["candidate.sequence"]), decoupled from the feature names. - // (The answer is NOT a feature_group: it's a data_config.label_field read from - // batch.jagged_labels — see label_feature_name.) + // (The answer is NOT a feature_group: it's the first data_config.label_field, + // read from batch.jagged_labels.) optional string history_group_name = 10 [default = "user_seq"]; reserved 11; // was label_group_name (answer moved to data_config.label_fields) From d68f594fddb5f910588b5417f2984110a0e25d4b Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 25 Jun 2026 07:58:51 +0000 Subject: [PATCH 30/99] [refactor] genreclm proto: drop reserved markers for removed fields MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per preference, removed fields are deleted outright rather than held as `reserved`. Strips reserved 1 / 4 / 5 / 11 from GenerativeRecLMConfig; field numbers of live fields are unchanged (gaps are legal in proto2). No behavior change — pb2 regenerates, configs reference field names. Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/protos/models/generative_model.proto | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index 5174baa36..255a46ab2 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -21,8 +21,7 @@ package tzrec.protos; // SID -> token id by integer offset at batch time. message GenerativeRecLMConfig { // Backbone (`hf_model_id`) is NOT here — it's the family's architecture - // commitment, so it lives on the family message. (Old class_name was 1.) - reserved 1; + // commitment, so it lives on the family message. // SID vocabulary, one entry per RQ level: len = SID codes per item (answer // width), sum = atoms appended as C0..C{sum-1} after the base vocab. @@ -35,8 +34,6 @@ message GenerativeRecLMConfig { // history_group (its feature_names[0]); the answer is the first // data_config.label_field (read from batch.jagged_labels, like RankModel's // labels[0]). - reserved 4; // was user_sequence_feature_name (-> history_group's feature) - reserved 5; // was label_feature_name (-> data_config.label_fields[0]) // Cross-entropy ignore index — matches PyTorch's F.cross_entropy default. optional int32 ignore_index = 6 [default = -100]; @@ -57,7 +54,6 @@ message GenerativeRecLMConfig { // (The answer is NOT a feature_group: it's the first data_config.label_field, // read from batch.jagged_labels.) optional string history_group_name = 10 [default = "user_seq"]; - reserved 11; // was label_group_name (answer moved to data_config.label_fields) // Predictions key the inference branch emits generated SIDs under (stable // across families; PredictWrapper output_cols should reference it). From af2d7bb4d7e12e2e209d08c70002878decd3bb11 Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 25 Jun 2026 08:08:52 +0000 Subject: [PATCH 31/99] [refactor] genreclm: drop history_group_name; use the single feature_group genrec declares exactly one feature_group now (the JAGGED_SEQUENCE history; the answer is a data_config.label_field, not a group), so naming it via common.history_group_name was redundant. _read_common_config now takes the single declared feature_group as the history (group_name + its one member), via a new _history_feature_group() that fails loudly on a missing/empty group. Removes the proto field (no reserved, per preference); configs no longer mention the knob (group_name is now free); tests cover the derivation + the no-feature_group error path. Also fixes the GenerativeRecLMConfig doc, which still claimed the answer was a sequence_raw_feature column. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../generative_rec_lm_s1pretrained.config | 17 +++--- .../generative_rec_lm_s2pretrained.config | 12 ++-- tzrec/models/generative_rec_lm.py | 56 ++++++++++--------- tzrec/models/generative_rec_lm_test.py | 21 +++---- tzrec/protos/models/generative_model.proto | 27 ++++----- 5 files changed, 66 insertions(+), 67 deletions(-) diff --git a/examples/generative_rec_lm_s1pretrained.config b/examples/generative_rec_lm_s1pretrained.config index 0874ef50f..ff31c9747 100644 --- a/examples/generative_rec_lm_s1pretrained.config +++ b/examples/generative_rec_lm_s1pretrained.config @@ -133,10 +133,11 @@ feature_configs { } model_config { - # One JAGGED_SEQUENCE group for the history INPUT: build_input retrieves the - # raw SID stream via the EmbeddingGroup and keys it by GROUP name (HSTU idiom); - # the name matches common.history_group_name default ("user_seq"). The answer - # is NOT a group — it's the label_field above (batch.jagged_labels). + # The SINGLE feature_group = the history INPUT: build_input retrieves the raw + # SID stream via the EmbeddingGroup and keys it by GROUP name (HSTU idiom). The + # model takes the only declared group as the history (group_name is free — it + # need not be "user_seq"). The answer is NOT a group — it's the label_field + # above (batch.jagged_labels). feature_groups { group_name: "user_seq" feature_names: "user_sequence" @@ -160,10 +161,10 @@ model_config { codebook: 8192 vocab_pad_to_multiple_of: 128 - # SID-column names are derived, not configured: the history feature is - # the single member of the history_group below; the answer is the first - # data_config.label_field. Only the history GROUP needs naming - # (history_group_name default "user_seq" matches the group below). + # The history feature_group and the SID-column names are all derived, + # not configured: the history is the single feature_group above and its + # one member is the history feature; the answer is the first + # data_config.label_field. ignore_index: -100 # Model history budget (HSTU-style model knob): truncation cap + diff --git a/examples/generative_rec_lm_s2pretrained.config b/examples/generative_rec_lm_s2pretrained.config index cf8f813a9..ee3dce1d2 100644 --- a/examples/generative_rec_lm_s2pretrained.config +++ b/examples/generative_rec_lm_s2pretrained.config @@ -51,9 +51,10 @@ feature_configs { } } model_config { - # One JAGGED_SEQUENCE group for the history INPUT (raw SID passthrough; see - # s1pretrained / build_input). Name matches history_group_name default. The - # answer is the label_field above (batch.jagged_labels), not a group. + # The SINGLE feature_group = the history INPUT (raw SID passthrough; see + # s1pretrained / build_input). The model takes the only declared group as the + # history (group_name is free). The answer is the label_field above + # (batch.jagged_labels), not a group. feature_groups { group_name: "user_seq" feature_names: "user_sequence" @@ -63,8 +64,9 @@ model_config { # Qwen2 backbone (owned by this family message, not `common`). hf_model_id: "/home/admin/workspace/Qwen2.5-0.5B" common { - # SID-column names are derived (history = the history_group's single feature; - # answer = the first data_config.label_field), not configured here. + # The history feature_group and SID-column names are derived (history = the + # single feature_group above + its one member; answer = the first + # data_config.label_field), not configured here. ignore_index: -100 # SID codebook: one entry per RQ level (AL-GR = 3 levels x 8192; verified # from item_info codebook_lv* ranges). len = SID codes/answer, sum = vocab diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 7d29e9f22..84c89aa67 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -121,15 +121,16 @@ def _resolve_pad_token_id(tokenizer: Any) -> int: def _read_common_config(self, common: Any) -> int: """Parse shared proto knobs into attributes; return the SID atom count.""" - # The history is a JAGGED_SEQUENCE feature_group; build_input keys the - # EmbeddingGroup output by this GROUP name (HSTU idiom). The answer is NOT - # a feature — it's a data_config.label_field read from batch.jagged_labels. - self._history_group: str = common.history_group_name - # SID-column names are derived, not restated: the history feature is the - # single member of the history group; the answer is the first - # data_config.label_field (like RankModel's labels[0]). _input_name is only - # the rows-dict key; _label_name also indexes batch.jagged_labels. - self._input_name: str = self._history_feature_name() + # The history is the single declared feature_group (a JAGGED_SEQUENCE + # group); build_input keys the EmbeddingGroup output by its GROUP name + # (HSTU idiom). Nothing about it is restated in `common`: the group name + # and its one member (the history feature) come straight from the group; + # the answer is the first data_config.label_field (like RankModel's + # labels[0]). _input_name is only the rows-dict key; _label_name also + # indexes batch.jagged_labels. The answer is NOT a feature_group. + hist = self._history_feature_group() + self._history_group: str = hist.group_name + self._input_name: str = hist.feature_names[0] self._label_name: str = self._labels[0] if self._labels else "" self._ignore_index: int = int(common.ignore_index) # Inference output key + backbone param dtype (configurable; default @@ -160,26 +161,27 @@ def _read_common_config(self, common: Any) -> int: self._vocab_pad_mult = int(common.vocab_pad_to_multiple_of) or 128 return sum(int(c) for c in codebook) - def _history_feature_name(self) -> str: - """The history feature's name = the single member of the history group. + def _history_feature_group(self) -> Any: + """The history feature_group = the single declared feature_group. - The history JAGGED_SEQUENCE group carries one SID stream, so its feature - name is ``feature_names[0]``; build_input uses it only as the rows-dict - key. Fails loudly on a misconfigured group (vs a silent KeyError later). + genrec declares exactly one feature_group — the JAGGED_SEQUENCE group + carrying the history SID stream (the answer is a data_config.label_field, + not a group). Its ``group_name`` keys the EmbeddingGroup output and its one + member is the history feature. Fails loudly on a missing/empty group (vs a + silent IndexError/KeyError later). """ - for g in self._feature_groups: - if g.group_name == self._history_group: - if not g.feature_names: - raise ValueError( - f"{type(self).__name__}: history feature_group " - f"{self._history_group!r} has no feature_names." - ) - return g.feature_names[0] - raise ValueError( - f"{type(self).__name__}: history_group_name {self._history_group!r} " - f"matches no feature_group; declare a JAGGED_SEQUENCE group with that " - f"group_name." - ) + if not self._feature_groups: + raise ValueError( + f"{type(self).__name__}: no feature_group declared; genrec needs " + f"one JAGGED_SEQUENCE group carrying the history SID stream." + ) + g = self._feature_groups[0] + if not g.feature_names: + raise ValueError( + f"{type(self).__name__}: history feature_group {g.group_name!r} " + f"has no feature_names." + ) + return g def _build_backbone(self) -> Any: """Build the EMPTY extended architecture in fp32 (master) — no download. diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py index 2a9c62bdd..e6266fe1e 100644 --- a/tzrec/models/generative_rec_lm_test.py +++ b/tzrec/models/generative_rec_lm_test.py @@ -104,8 +104,9 @@ def test_read_common_config_reads_knobs(self) -> None: m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) m._features = [] - # SID-column names are DERIVED, not configured: history = the history - # group's single feature; answer = the first data_config.label_field. + # The history group + SID-column names are DERIVED, not configured: the + # history is the single feature_group + its one member; the answer is the + # first data_config.label_field. m._labels = ["label"] m._feature_groups = [ types.SimpleNamespace( @@ -113,7 +114,6 @@ def test_read_common_config_reads_knobs(self) -> None: ) ] common = types.SimpleNamespace( - history_group_name="user_seq", ignore_index=-100, generated_sids_key="my_sids", param_dtype="bfloat16", @@ -122,7 +122,8 @@ def test_read_common_config_reads_knobs(self) -> None: max_sequence_length=288, ) m._read_common_config(common) - self.assertEqual(m._input_name, "user_sequence") # from history group + self.assertEqual(m._history_group, "user_seq") # the single group + self.assertEqual(m._input_name, "user_sequence") # its one member self.assertEqual(m._label_name, "label") # from label_fields[0] self.assertEqual(m._generated_sids_key, "my_sids") # configurable self.assertIs(m._param_dtype, torch.bfloat16) # name -> torch dtype @@ -132,17 +133,14 @@ def test_read_common_config_reads_knobs(self) -> None: with self.assertRaisesRegex(ValueError, "param_dtype must be one of"): m._read_common_config(common) - def test_read_common_config_unknown_history_group_raises(self) -> None: - # history_group_name must match a declared feature_group, else fail loudly. + def test_read_common_config_no_feature_group_raises(self) -> None: + # the history is the single declared feature_group; none -> fail loudly. m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) m._features = [] m._labels = ["label"] - m._feature_groups = [ - types.SimpleNamespace(group_name="other", feature_names=["x"]) - ] + m._feature_groups = [] # no group declared common = types.SimpleNamespace( - history_group_name="user_seq", ignore_index=-100, generated_sids_key="generated_sids", param_dtype="float32", @@ -150,14 +148,13 @@ def test_read_common_config_unknown_history_group_raises(self) -> None: vocab_pad_to_multiple_of=128, max_sequence_length=0, ) - with self.assertRaisesRegex(ValueError, "matches no feature_group"): + with self.assertRaisesRegex(ValueError, "no feature_group declared"): m._read_common_config(common) def test_max_sequence_length_model_knob(self) -> None: # _max_seq_length is the model knob; 0 = off (no feature fallback). def _common(max_seq): return types.SimpleNamespace( - history_group_name="user_seq", ignore_index=-100, generated_sids_key="generated_sids", param_dtype="float32", diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index 255a46ab2..3146bafcd 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -15,10 +15,12 @@ package tzrec.protos; // Architecture-agnostic config shared by ALL generative-rec families (the base // reads this for everything except the backbone, which the family owns — see // _backbone_id). Sample contract (consumed by `predict()`): -// * user_sequence : list — raw SID indices in [1, sum(codebook)] -// * label : list — raw SID indices in [1, sum(codebook)] -// Both flow into TER as `sequence_raw_feature` parquet columns; the model maps -// SID -> token id by integer offset at batch time. +// * history : list — raw SID indices in [1, sum(codebook)]; the single +// JAGGED_SEQUENCE feature_group (a `sequence_raw_feature` column). +// * answer : list — raw SID indices; a `data_config.label_field` +// (read from batch.jagged_labels), NOT a feature. +// The model maps SID -> token id by integer offset at batch time. Column names +// are NOT configured here (see below). message GenerativeRecLMConfig { // Backbone (`hf_model_id`) is NOT here — it's the family's architecture // commitment, so it lives on the family message. @@ -29,11 +31,13 @@ message GenerativeRecLMConfig { // Pad the post-extension vocab up to a multiple of this value. optional uint32 vocab_pad_to_multiple_of = 3 [default = 128]; - // SID-column names are NOT restated here — they're derived from the config - // that already declares them: the history feature is the single member of the - // history_group (its feature_names[0]); the answer is the first + // Neither the history feature_group nor the SID-column names are configured + // here — they're derived from the config that already declares them: the + // history is the single declared feature_group (a JAGGED_SEQUENCE group; + // build_input keys the EmbeddingGroup output by its GROUP name, the HSTU + // idiom), and its one member is the history feature; the answer is the first // data_config.label_field (read from batch.jagged_labels, like RankModel's - // labels[0]). + // labels[0]). The answer is NOT a feature_group. // Cross-entropy ignore index — matches PyTorch's F.cross_entropy default. optional int32 ignore_index = 6 [default = -100]; @@ -48,13 +52,6 @@ message GenerativeRecLMConfig { // dynamic_beams schedule; exploits the fixed-length, EOS-free SID answer. optional bool dynamic_beam = 9 [default = false]; - // Which feature_group (group_type: JAGGED_SEQUENCE) carries the HISTORY SID - // stream. The model keys the EmbeddingGroup output by GROUP name (like HSTU's - // grouped_features["candidate.sequence"]), decoupled from the feature names. - // (The answer is NOT a feature_group: it's the first data_config.label_field, - // read from batch.jagged_labels.) - optional string history_group_name = 10 [default = "user_seq"]; - // Predictions key the inference branch emits generated SIDs under (stable // across families; PredictWrapper output_cols should reference it). optional string generated_sids_key = 12 [default = "generated_sids"]; From c168cc31fafb30360fde894da5ff0cd8eb9f27ca Mon Sep 17 00:00:00 2001 From: shuqi <597191244@qq.com> Date: Thu, 25 Jun 2026 08:22:04 +0000 Subject: [PATCH 32/99] [refactor] genreclm: trim redundant inline comments Remove inline comments that merely restate the adjacent code (pool-build, debug-dump, dtype/knob assignments, the cited max_sequence_length block, etc.), keeping only the non-obvious "why" notes: fp32-master rationale, first-step pool warming, the OOM-avoiding suffix slice, the batch-major generate() order, and the whole-candidate SID reject. Also fix two now-stale docstrings (the module sample-format + _sid_token_rows describing the answer as a feature_group rather than a data_config.label_field) and the init_input "groups" plural. Verified by a 3-lens review pass (over-removal / redundancy / accuracy). Co-Authored-By: Claude Opus 4.8 (1M context) --- tzrec/models/generative_rec_lm.py | 58 +++++++++++-------------------- tzrec/models/qwen2_rec_lm.py | 14 ++------ 2 files changed, 23 insertions(+), 49 deletions(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 84c89aa67..aa250b8bc 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -12,9 +12,9 @@ family by its own oneof entry, whose message-type name resolves directly to the same-named class via the BaseModel registry. Shared config lives in ``GenerativeRecLMConfig`` (the family message's ``common`` field). -* Streaming sample format: each row carries two raw-int64 sequence features, -``user_sequence`` and ``label``, both holding raw SID indices in -``[1, sum(codebook)]``. +* Streaming sample format: each row carries the history ``user_sequence`` (a +raw-int64 sequence feature) and the ``label`` answer (a ``data_config`` +label_field), both holding raw SID indices in ``[1, sum(codebook)]``. * The chat template is tokenised ONCE at ``__init__`` and cached as non-persistent buffers, so per-batch encoding is integer arithmetic only (no HF tokenizer in the hot path). @@ -95,11 +95,8 @@ def __init__( self._build_prompt_tokens(tokenizer, cfg) - # Build the (param-free, raw-passthrough) EmbeddingGroup for the SID - # JAGGED_SEQUENCE groups — the HSTU retrieval idiom (see init_input). self.init_input() - # one-shot debug dump of the first spliced batch self._smoke_log_once = os.environ.get("TZREC_GENRECLM_DEBUG", "0") == "1" self._first_predict = True @@ -121,20 +118,13 @@ def _resolve_pad_token_id(tokenizer: Any) -> int: def _read_common_config(self, common: Any) -> int: """Parse shared proto knobs into attributes; return the SID atom count.""" - # The history is the single declared feature_group (a JAGGED_SEQUENCE - # group); build_input keys the EmbeddingGroup output by its GROUP name - # (HSTU idiom). Nothing about it is restated in `common`: the group name - # and its one member (the history feature) come straight from the group; - # the answer is the first data_config.label_field (like RankModel's - # labels[0]). _input_name is only the rows-dict key; _label_name also - # indexes batch.jagged_labels. The answer is NOT a feature_group. + # History = the single feature_group (its group_name keys the + # EmbeddingGroup output); answer = the first data_config.label_field. hist = self._history_feature_group() self._history_group: str = hist.group_name self._input_name: str = hist.feature_names[0] self._label_name: str = self._labels[0] if self._labels else "" self._ignore_index: int = int(common.ignore_index) - # Inference output key + backbone param dtype (configurable; default - # "generated_sids" / "float32" = the fp32-master weights). self._generated_sids_key: str = common.generated_sids_key param_dtype = self._DTYPE_BY_NAME.get(common.param_dtype) if param_dtype is None: @@ -143,18 +133,13 @@ def _read_common_config(self, common: Any) -> int: f"{list(self._DTYPE_BY_NAME)}, got {common.param_dtype!r}." ) self._param_dtype: torch.dtype = param_dtype - # Model's history budget (SID codes): the truncation cap (_sid_token_rows, - # item-aligned, recency-preserving) AND the activation-pool pre-size. - # HSTU-style model knob; 0 = off (no cap, no pre-allocation). FG_NONE does - # not truncate, so this cap is enforced model-side. self._max_seq_length: int = int(common.max_sequence_length) codebook = list(common.codebook) if len(codebook) == 0: raise ValueError("GenerativeRecLM: codebook must be non-empty.") - # len(codebook) = SID codes per item (answer width); sum = vocab atoms. self._num_levels = len(codebook) - # Per-level SID validity bands (buffers) for the inference gate - # (_validate_sid_candidates). Non-persistent: derived config. + # inference-gate validity bands; non-persistent — derived from codebook, + # kept off the state_dict (HF safetensors round-trip). lo, hi = self._sid_level_bands(codebook) self.register_buffer("_sid_lvl_lo", lo, persistent=False) self.register_buffer("_sid_lvl_hi", hi, persistent=False) @@ -331,19 +316,18 @@ def _validate_sid_candidates( wrong-level atom) set to ``-1`` — which can never match a real item. """ sids = self._detokenize_sids(new_tokens) - # pad each candidate to exactly num_levels with the -1 sentinel (an - # early-EOS beam returns fewer tokens) so the reshape stays rectangular. + # early-EOS beams return < num_levels tokens; pad to num_levels with -1. sids = F.pad(sids, (0, self._num_levels - sids.shape[1]), value=-1) - # valid only if every position-j atom is in level j's band; any violation - # invalidates the WHOLE candidate -> -1. + # any single out-of-band atom invalidates the WHOLE candidate (.any over + # the levels -> masked_fill blanks the entire row to -1, matching no item). invalid = ((sids < self._sid_lvl_lo) | (sids > self._sid_lvl_hi)).any(dim=1) sids = sids.masked_fill(invalid.unsqueeze(1), -1) - # generate() returns rows batch-major ([b0_beam0, b0_beam1, ...]) so this - # groups beams per user; the row-wise mask above preserved that order. + # generate() returns rows batch-major ([b0_beam0, b0_beam1, ...]); group + # the beams per user. return sids.view(batch_size, -1, self._num_levels) def init_input(self) -> None: - """Build the EmbeddingGroup for the raw SID JAGGED_SEQUENCE groups. + """Build the EmbeddingGroup for the single raw SID JAGGED_SEQUENCE group. Raw (passthrough) features carry no embedding tables, so this EmbeddingGroup holds no params (DMP-neutral); it exists purely to @@ -392,10 +376,12 @@ def _sid_token_rows( ) -> List[torch.Tensor]: """Map flat SID ``(values, lengths)`` -> per-row token-id tensors. - ``build_input`` supplies a JAGGED_SEQUENCE group's flat - ``"{group}.sequence"`` values + ``"{group}.sequence_length"``; ``values`` - may arrive as float / shape ``(N, 1)``. The whole batch is tokenized once - on the backbone device, then split into rows. + ``build_input`` supplies flat SID ``(values, lengths)``: the history from + the JAGGED_SEQUENCE group's ``"{group}.sequence"`` / + ``"{group}.sequence_length"``, and the answer from + ``batch.jagged_labels[label].values()/.lengths()``. ``values`` may arrive + as float / shape ``(N, 1)``. The whole batch is tokenized once on the + backbone device, then split into rows. ``expected_width``, when set, enforces the sample contract: every row must have exactly that many codes (the answer = ``num_levels``). @@ -407,7 +393,6 @@ def _sid_token_rows( """ if values.dim() == 2 and values.size(-1) == 1: values = values.squeeze(-1) - # host-side split bounds, read before the H2D copy below sizes = lengths.long().tolist() if expected_width is not None: bad = [i for i, n in enumerate(sizes) if n != expected_width] @@ -425,10 +410,9 @@ def _sid_token_rows( if max_codes: keep = (max_codes // self._num_levels) * self._num_levels if keep and any(n > keep for n in sizes): - rows = torch.split(values, sizes) # host views, no copy - values = torch.cat([r[-keep:] for r in rows]) # keep recent tail + rows = torch.split(values, sizes) + values = torch.cat([r[-keep:] for r in rows]) sizes = [min(n, keep) for n in sizes] - # one vectorized SID->token map over the whole batch, on the backbone device values = self._tokenize_sids(values.to(self.device).long()) return list(torch.split(values, sizes)) diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index deda355f5..ffb732d82 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -76,12 +76,9 @@ def __init__( ) -> None: super().__init__(model_config, features, labels, sample_weights, **kwargs) common = self._model_config.common - # generation params, consumed only by this family's _generate. self._num_beams = int(common.num_beams) self._num_return = int(common.num_return_sequences) - # opt-in ALGR-style escalating beam (width doubles per SID level). self._dynamic_beam = bool(common.dynamic_beam) - # worst-case spliced length for the first-step activation-pool pre-sizing. self._max_total_len = self._compute_max_total_length() self._pool_warmed = False # CE suffix width: the supervised tail [answer | asst_suffix | eos] is @@ -168,8 +165,6 @@ def _splice_input_ids( assert len(user_seq_rows) == len(label_rows) A = self._num_levels - # input_ids: assembled per row (user history length varies), then - # left-padded into a (B, T) batch (real content right-aligned). rows_ids = [ torch.cat( [ @@ -187,11 +182,9 @@ def _splice_input_ids( ] input_ids, attention_mask = self._left_pad(rows_ids, pad_to=pad_to) - # supervised tail is fixed-width -> same columns every row -> one write. - # tail from the end: [answer(A) | asst_suffix(s) | eos(1)]. B, T = input_ids.shape s = self.tpl_asst_suffix.numel() - tail = A + s + 1 + tail = A + s + 1 # [answer(A) | asst_suffix(s) | eos(1)], end-aligned labels = torch.full( (B, T), self._ignore_index, dtype=torch.long, device=self.device ) @@ -213,8 +206,6 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: """Branch 1: teacher-forced forward -> suffix-slice -> CE loss.""" - # Retrieve SID token rows via the EmbeddingGroup (build_input), keyed by - # feature name. Train needs both the history and the teacher-forced answer. rows = self.build_input(batch) u_rows = rows[self._input_name] l_rows = rows[self._label_name] @@ -277,7 +268,6 @@ def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: ``_validate_sid_candidates`` (token->SID, malformed beams -> ``-1``). Returns ``generated_sids`` of shape ``(B, num_return, num_levels)``. """ - # history rows via the EmbeddingGroup (build_input); inference skips label. u_rows = self.build_input(batch)[self._input_name] input_ids, attention_mask = self._splice_prompt_ids(u_rows) if self._dynamic_beam: @@ -292,7 +282,7 @@ def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: do_sample=False, pad_token_id=self._pad_token_id, ) - new_tokens = out[:, input_ids.shape[1] :] # the generated tail + new_tokens = out[:, input_ids.shape[1] :] sids = self._validate_sid_candidates(new_tokens, input_ids.shape[0]) return {self._generated_sids_key: sids} From 2d97b5c33acb4844615dfb867b276add05258df0 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Mon, 27 Jul 2026 09:21:20 +0000 Subject: [PATCH 33/99] [fix] GenerativeRecLM: apply per-level SID offsets Use local 1-based codes at the data boundary and derive level offsets inside the model. Remove branch-only generative recommendation examples that are not intended for submission. --- examples/convert_s1full_to_parquet.py | 121 --------- examples/convert_s1tiny_test_to_parquet.py | 149 ----------- examples/convert_s1tiny_to_parquet.py | 173 ------------ examples/convert_s2test_to_parquet.py | 90 ------- examples/generative_rec_lm_predict.py | 241 ----------------- .../generative_rec_lm_s1pretrained.config | 184 ------------- .../generative_rec_lm_s2pretrained.config | 86 ------ examples/generative_rec_lm_smoke.py | 149 ----------- examples/generative_rec_lm_train_loop.py | 230 ---------------- .../generative_rec_lm_train_loop_parquet.py | 247 ------------------ tzrec/models/generative_rec_lm.py | 97 ++++--- tzrec/models/generative_rec_lm_test.py | 114 +++++--- tzrec/models/qwen2_rec_lm.py | 5 +- tzrec/models/qwen2_rec_lm_test.py | 197 +++++++++----- tzrec/protos/models/generative_model.proto | 21 +- 15 files changed, 292 insertions(+), 1812 deletions(-) delete mode 100644 examples/convert_s1full_to_parquet.py delete mode 100644 examples/convert_s1tiny_test_to_parquet.py delete mode 100644 examples/convert_s1tiny_to_parquet.py delete mode 100644 examples/convert_s2test_to_parquet.py delete mode 100644 examples/generative_rec_lm_predict.py delete mode 100644 examples/generative_rec_lm_s1pretrained.config delete mode 100644 examples/generative_rec_lm_s2pretrained.config delete mode 100644 examples/generative_rec_lm_smoke.py delete mode 100644 examples/generative_rec_lm_train_loop.py delete mode 100644 examples/generative_rec_lm_train_loop_parquet.py diff --git a/examples/convert_s1full_to_parquet.py b/examples/convert_s1full_to_parquet.py deleted file mode 100644 index 1133ad220..000000000 --- a/examples/convert_s1full_to_parquet.py +++ /dev/null @@ -1,121 +0,0 @@ -"""Convert full-s1 (``s1_splits/part_*.csv``, 41 shards, ~246 GB) into -TorchEasyRec genrec parquet — the multi-CSV streaming sibling of -``convert_s1tiny_to_parquet.py``. - -Identical row logic (extract every ``C\\d+`` -> 1-indexed SID; columns -``user_sequence`` / ``label`` as ``list``), but iterates a GLOB of CSV -parts with a CONTINUOUS shard index so the 41 files land in one parquet dir. -Streaming (``csv.DictReader`` + shard-buffered flush) so memory stays bounded -regardless of the 246 GB input; the dropped CN-prompt text means the parquet is -a small fraction of the CSV size. - -Usage on remote:: - - /opt/conda/bin/python -m examples.convert_s1full_to_parquet \\ - --csv_glob '/home/admin/workspace/aop_lab/data/AL-GR-v1/s1_splits/*.csv' \\ - --out_dir /home/admin/workspace/aop_lab/data/AL-GR-v1/train_data_genreclm_s1full \\ - --shard_size 200000 --log_every 1000000 -""" - -from __future__ import annotations - -import argparse -import csv -import glob -import os -import re -import sys -import time -from typing import Iterator, List, Tuple - -import pyarrow as pa -import pyarrow.parquet as pq - -csv.field_size_limit(10 * 1024 * 1024) -_SID_RE = re.compile(r"C(\d+)") - - -def _extract_sids(text: str) -> List[int]: - return [int(m) + 1 for m in _SID_RE.findall(text)] - - -def _iter_rows( - csv_paths: List[str], max_rows: int -) -> Iterator[Tuple[List[int], List[int]]]: - n = 0 - for p in csv_paths: - with open(p, "r", encoding="utf-8", newline="") as f: - for row in csv.DictReader(f): - if max_rows and n >= max_rows: - return - u = _extract_sids(row.get("user") or "") - lab = _extract_sids(row.get("answer") or "") - if not u or not lab: - continue - n += 1 - yield u, lab - - -def main() -> int: - ap = argparse.ArgumentParser() - ap.add_argument("--csv_glob", required=True) - ap.add_argument("--out_dir", required=True) - ap.add_argument("--shard_size", type=int, default=200_000) - ap.add_argument("--max_rows", type=int, default=0) - ap.add_argument("--log_every", type=int, default=1_000_000) - args = ap.parse_args() - - paths = sorted(glob.glob(args.csv_glob)) - print(f"{len(paths)} csv parts matched", flush=True) - os.makedirs(args.out_dir, exist_ok=True) - schema = pa.schema( - [ - pa.field("user_sequence", pa.list_(pa.int64()), nullable=False), - pa.field("label", pa.list_(pa.int64()), nullable=False), - ] - ) - - def flush(rows: List[Tuple[List[int], List[int]]], idx: int) -> None: - if not rows: - return - tbl = pa.table( - {"user_sequence": [r[0] for r in rows], "label": [r[1] for r in rows]}, - schema=schema, - ) - pq.write_table( - tbl, os.path.join(args.out_dir, f"shard-{idx:05d}.parquet"), - compression="zstd", - ) - - t0 = time.time() - rows: List[Tuple[List[int], List[int]]] = [] - idx = 0 - total = 0 - max_sid = 0 - for u, lab in _iter_rows(paths, args.max_rows): - rows.append((u, lab)) - total += 1 - max_sid = max(max_sid, max(u), max(lab)) - if len(rows) >= args.shard_size: - flush(rows, idx) - rows = [] - idx += 1 - if args.log_every and total % args.log_every == 0: - print( - f"[progress] rows={total} shards={idx} max_sid={max_sid} " - f"wall={time.time() - t0:.0f}s", - flush=True, - ) - if rows: - flush(rows, idx) - idx += 1 - print( - f"[done] rows={total} shards={idx} max_sid={max_sid} " - f"wall={time.time() - t0:.0f}s out_dir={args.out_dir}", - flush=True, - ) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/examples/convert_s1tiny_test_to_parquet.py b/examples/convert_s1tiny_test_to_parquet.py deleted file mode 100644 index b557caaf6..000000000 --- a/examples/convert_s1tiny_test_to_parquet.py +++ /dev/null @@ -1,149 +0,0 @@ -"""Convert algr's ``s1_tiny_test.csv`` into TorchEasyRec-shaped parquet. - -The TEST split differs from the train split: the ``answer`` column holds -ground-truth ITEM IDs (e.g. ``AIB2f`` or ``9nXWC;I6kld;...`` — -semicolon-separated when multiple), not SID ``C``-codes. algr evaluates -this split via beam-search generation + ``calc_hr_fast.py`` item mapping. - -For TER's periodic CE-loss eval we need ``(user_sequence, label)`` SID rows, -so this script maps the FIRST ground-truth item of each row to its SID -triple via ``item_info/tiny_item_sid_final.csv``: - - sid_triple(item) = [lv1 + 1, lv2 + 8192 + 1, lv3 + 16384 + 1] - -(the per-layer offsets follow the dataset's C-token scheme -``C{lv1}C{lv2+8192}C{lv3+16384}``; the +1 converts the 0-indexed C-code to -the 1-indexed SID contract of ``GenerativeRecLM`` — identical to -``convert_s1tiny_to_parquet.py``'s ``Ck → SID = k + 1``.) - -Rows whose first answer item is missing from the item map are skipped -(counted and reported). - -Usage on remote:: - - /opt/conda/bin/python -m examples.convert_s1tiny_test_to_parquet \\ - --csv /home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data/s1_tiny_test.csv \\ - --item_sid_csv /home/admin/workspace/aop_lab/data/AL-GR-Tiny/item_info/tiny_item_sid_final.csv \\ - --out_dir /home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data_genreclm -""" - -from __future__ import annotations - -import argparse -import csv -import os -import re -import sys -import time -from typing import Dict, List, Tuple - -import pyarrow as pa -import pyarrow.parquet as pq - -csv.field_size_limit(10 * 1024 * 1024) - -_SID_RE = re.compile(r"C(\d+)") - -_LV2_OFFSET = 8192 -_LV3_OFFSET = 16384 - - -def _extract_sids(text: str) -> List[int]: - """``Ck → SID = k + 1`` (1-indexed; matches the train converter).""" - return [int(m) + 1 for m in _SID_RE.findall(text)] - - -def _load_item_map(path: str) -> Dict[str, Tuple[int, int, int]]: - t0 = time.time() - item_map: Dict[str, Tuple[int, int, int]] = {} - bad = 0 - with open(path, "r", encoding="utf-8", newline="") as f: - r = csv.DictReader(f) - for row in r: - try: - item_map[row["item_id"]] = ( - int(row["codebook_lv1"]), - int(row["codebook_lv2"]), - int(row["codebook_lv3"]), - ) - except (ValueError, TypeError, KeyError): - # the 25M-row map contains a few malformed/empty rows - bad += 1 - print( - f"[item_map] {len(item_map)} items loaded " - f"({bad} malformed rows skipped) in {time.time()-t0:.0f}s", - flush=True, - ) - return item_map - - -def main() -> int: - ap = argparse.ArgumentParser() - ap.add_argument("--csv", required=True) - ap.add_argument("--item_sid_csv", required=True) - ap.add_argument("--out_dir", required=True) - ap.add_argument("--shard_size", type=int, default=200_000) - args = ap.parse_args() - - os.makedirs(args.out_dir, exist_ok=True) - schema = pa.schema( - [ - pa.field("user_sequence", pa.list_(pa.int64()), nullable=False), - pa.field("label", pa.list_(pa.int64()), nullable=False), - ] - ) - - item_map = _load_item_map(args.item_sid_csv) - - rows: List[Tuple[List[int], List[int]]] = [] - total = read = miss_item = miss_user = 0 - max_sid = 0 - with open(args.csv, "r", encoding="utf-8", newline="") as f: - for row in csv.DictReader(f): - read += 1 - user_sids = _extract_sids(row.get("user") or "") - if not user_sids: - miss_user += 1 - continue - first_item = (row.get("answer") or "").split(";")[0].strip() - lv = item_map.get(first_item) - if lv is None: - miss_item += 1 - continue - label_sids = [ - lv[0] + 1, - lv[1] + _LV2_OFFSET + 1, - lv[2] + _LV3_OFFSET + 1, - ] - rows.append((user_sids, label_sids)) - total += 1 - max_sid = max(max_sid, max(user_sids), max(label_sids)) - - shard_idx = 0 - for i in range(0, len(rows), args.shard_size): - chunk = rows[i : i + args.shard_size] - tbl = pa.table( - { - "user_sequence": [r[0] for r in chunk], - "label": [r[1] for r in chunk], - }, - schema=schema, - ) - pq.write_table( - tbl, - os.path.join(args.out_dir, f"shard-{shard_idx:05d}.parquet"), - compression="zstd", - ) - shard_idx += 1 - - print( - f"[done] read={read} rows_written={total} shards={shard_idx} " - f"skipped_missing_item={miss_item} skipped_no_user_sids={miss_user} " - f"max_sid={max_sid} out_dir={args.out_dir}", - flush=True, - ) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/examples/convert_s1tiny_to_parquet.py b/examples/convert_s1tiny_to_parquet.py deleted file mode 100644 index 99a54fd5b..000000000 --- a/examples/convert_s1tiny_to_parquet.py +++ /dev/null @@ -1,173 +0,0 @@ -"""Convert algr's ``s1_tiny.csv`` into TorchEasyRec-shaped parquet. - -algr's row format:: - - system,user,answer - ,...C4805C8364C16402C4487C12277...,C1517C12109C16399 - -Each item is encoded as 3 contiguous ``C{i}`` codes (one per RQ-VAE layer). -The ``user`` column holds a CN sentence prefix + the user-history SID codes -+ a CN sentence suffix; the ``answer`` column holds just the target SID -codes for the next item. - -This script extracts all ``C\\d+`` matches from each field, converts to -1-indexed SID integers (``Ck → SID = k + 1``), and writes a parquet shard -per N rows with two columns matching ``GenerativeRecLM``'s contract: - - user_sequence : list - label : list - -The chat-template prompt strings are NOT carried over — ``GenerativeRecLM`` -re-builds them from cached buffers at splice time, using its own -``system_instruction``. To preserve algr's bit-exact prompt for parity -testing, set ``GenerativeRecLM.system_instruction`` to algr's CN prefix in -the pipeline.config (proto field already exists; see §3 of the design). - -Usage on remote:: - - /opt/conda/bin/python -m examples.convert_s1tiny_to_parquet \\ - --csv /home/admin/workspace/aop_lab/data/AL-GR-Tiny/train_data/s1_tiny.csv \\ - --out_dir /home/admin/workspace/aop_lab/data/AL-GR-Tiny/train_data_genreclm \\ - --shard_size 200000 --max_rows 0 \\ - --log_every 200000 -""" - -from __future__ import annotations - -import argparse -import csv -import os -import re -import sys -import time -from typing import Iterator, List, Tuple - -import pyarrow as pa -import pyarrow.parquet as pq - -# Allow rows up to ~6 MB of CSV text (the long Chinese prompt + thousand SIDs). -csv.field_size_limit(10 * 1024 * 1024) - -# Match every Cxxx in a string. SIDs are non-negative integers; we use a -# bounded ``+`` quantifier so it's anchored on whole tokens. -_SID_RE = re.compile(r"C(\d+)") - - -def _extract_sids(text: str) -> List[int]: - """Pull every ``C\\d+`` out of `text` and return as 1-indexed SIDs. - - ``Ck → SID = k + 1`` so that SID range is [1, sum(codebook)] which is what - ``GenerativeRecLM._splice_input_ids`` expects (SID=1 maps to atom C0 via - ``token = sid + base - 1`` = ``base + (sid - 1)``). - """ - return [int(m) + 1 for m in _SID_RE.findall(text)] - - -def _iter_rows(csv_path: str, max_rows: int) -> Iterator[Tuple[List[int], List[int]]]: - with open(csv_path, "r", encoding="utf-8", newline="") as f: - r = csv.DictReader(f) - for i, row in enumerate(r): - if max_rows and i >= max_rows: - return - user_sids = _extract_sids(row.get("user") or "") - label_sids = _extract_sids(row.get("answer") or "") - if not user_sids or not label_sids: - continue - yield user_sids, label_sids - - -def main() -> int: - ap = argparse.ArgumentParser() - ap.add_argument("--csv", required=True) - ap.add_argument("--out_dir", required=True) - ap.add_argument( - "--shard_size", type=int, default=200_000, - help="rows per parquet shard", - ) - ap.add_argument( - "--max_rows", type=int, default=0, - help="cap total rows; 0 = whole CSV", - ) - ap.add_argument("--log_every", type=int, default=100_000) - args = ap.parse_args() - - os.makedirs(args.out_dir, exist_ok=True) - # Two int64 lists, no nulls. ``list_(int64())`` lets pyarrow store these - # as ListArray(int64) inside the parquet, which TER's - # ``sequence_raw_feature`` reads directly into a JaggedTensor. - schema = pa.schema( - [ - pa.field("user_sequence", pa.list_(pa.int64()), nullable=False), - pa.field("label", pa.list_(pa.int64()), nullable=False), - ] - ) - - t0 = time.time() - rows_in_shard: List[Tuple[List[int], List[int]]] = [] - shard_idx = 0 - total = 0 - max_sid_seen = 0 - max_user_len = 0 - max_label_len = 0 - - def flush(rows, idx): - if not rows: - return - user_col = [r[0] for r in rows] - label_col = [r[1] for r in rows] - tbl = pa.table( - {"user_sequence": user_col, "label": label_col}, schema=schema, - ) - path = os.path.join(args.out_dir, f"shard-{idx:05d}.parquet") - pq.write_table(tbl, path, compression="zstd") - - for u, lab in _iter_rows(args.csv, args.max_rows): - rows_in_shard.append((u, lab)) - total += 1 - max_sid_seen = max(max_sid_seen, max(u), max(lab)) - max_user_len = max(max_user_len, len(u)) - max_label_len = max(max_label_len, len(lab)) - - if len(rows_in_shard) >= args.shard_size: - flush(rows_in_shard, shard_idx) - rows_in_shard = [] - shard_idx += 1 - print( - f"[shard {shard_idx-1}] flushed; total={total} " - f"max_sid={max_sid_seen} max_user_len={max_user_len} " - f"max_label_len={max_label_len} " - f"wall={time.time()-t0:.0f}s", - flush=True, - ) - - if args.log_every and total % args.log_every == 0: - print( - f"[progress] rows={total} max_sid={max_sid_seen} " - f"max_user_len={max_user_len} max_label_len={max_label_len} " - f"wall={time.time()-t0:.0f}s", - flush=True, - ) - - if rows_in_shard: - flush(rows_in_shard, shard_idx) - shard_idx += 1 - - print( - f"[done] rows_written={total} shards={shard_idx} " - f"max_sid={max_sid_seen} max_user_len={max_user_len} " - f"max_label_len={max_label_len} " - f"wall={time.time()-t0:.1f}s out_dir={args.out_dir}", - flush=True, - ) - if max_sid_seen > 65536: - print( - f"[WARN] max_sid={max_sid_seen} > 65536 — algr's tokenizer adds " - f"only 65 536 atoms. Vocab extension in TER must be at least " - f"{max_sid_seen}.", - flush=True, - ) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/examples/convert_s2test_to_parquet.py b/examples/convert_s2test_to_parquet.py deleted file mode 100644 index c3c5c59b0..000000000 --- a/examples/convert_s2test_to_parquet.py +++ /dev/null @@ -1,90 +0,0 @@ -"""Convert s2 test CSV to TER parquet (user_sequence only; label = [0] placeholder). - -s2 test CSV answers are item IDs (not SID codes), so label cannot be derived. -The predict script only reads `user_sequence` from parquet, so label is a dummy. - -Usage on remote:: - - python3 /home/admin/workspace/TorchEasyRec_qwen_smoke/TorchEasyRec/examples/convert_s2test_to_parquet.py \ - --csv /home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data/s2_tiny_test.csv \ - --out_dir /home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data_genreclm_s2 -""" -from __future__ import annotations - -import argparse -import csv -import os -import re -import sys -import time -from typing import Iterator, List, Tuple - -import pyarrow as pa -import pyarrow.parquet as pq - -csv.field_size_limit(10 * 1024 * 1024) - -_SID_RE = re.compile(r"C(\d+)") - - -def _extract_user_sids(text: str) -> List[int]: - return [int(m) + 1 for m in _SID_RE.findall(text)] - - -def _iter_rows(csv_path: str, max_rows: int) -> Iterator[List[int]]: - with open(csv_path, newline="", encoding="utf-8") as f: - reader = csv.DictReader(f) - for i, row in enumerate(reader): - if max_rows and i >= max_rows: - break - user_sids = _extract_user_sids(row["user"]) - if not user_sids: - continue - yield user_sids - - -def main() -> int: - ap = argparse.ArgumentParser() - ap.add_argument("--csv", required=True) - ap.add_argument("--out_dir", required=True) - ap.add_argument("--shard_size", type=int, default=200_000) - ap.add_argument("--max_rows", type=int, default=0) - args = ap.parse_args() - - os.makedirs(args.out_dir, exist_ok=True) - schema = pa.schema([ - pa.field("user_sequence", pa.list_(pa.int64()), nullable=False), - pa.field("label", pa.list_(pa.int64()), nullable=False), - ]) - - t0 = time.time() - rows_in_shard: List[Tuple[List[int], List[int]]] = [] - shard_idx = 0 - total = 0 - - def flush(rows, idx): - if not rows: - return - user_col = [r[0] for r in rows] - label_col = [r[1] for r in rows] - tbl = pa.table({"user_sequence": user_col, "label": label_col}, schema=schema) - pq.write_table(tbl, os.path.join(args.out_dir, f"shard-{idx:05d}.parquet"), compression="zstd") - - for u in _iter_rows(args.csv, args.max_rows): - rows_in_shard.append((u, [0])) # dummy label - total += 1 - if len(rows_in_shard) >= args.shard_size: - flush(rows_in_shard, shard_idx) - rows_in_shard = [] - shard_idx += 1 - - if rows_in_shard: - flush(rows_in_shard, shard_idx) - shard_idx += 1 - - print(f"[done] rows_written={total} shards={shard_idx} wall={time.time()-t0:.1f}s out_dir={args.out_dir}") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/examples/generative_rec_lm_predict.py b/examples/generative_rec_lm_predict.py deleted file mode 100644 index e694257ac..000000000 --- a/examples/generative_rec_lm_predict.py +++ /dev/null @@ -1,241 +0,0 @@ -"""Beam-search SID generation for an exported `GenerativeRecLM` checkpoint. - -Mirrors algr's predict flow (config/qwen_predict_tiny_*.json: beam=50, -num_return=50, max_new_tokens=3, greedy) and emits `output_sids.jsonl` -lines in the exact shape `calc_hr_fast.py` consumes:: - - {"_generated_text_": ["C{a}C{b}C{c}", ... x num_return], "answer": "item1;item2"} - -The generated C-codes carry the dataset's layer offsets verbatim -(lv2 ∈ [8192, 16383], lv3 ∈ [16384, ...]) because the model is trained on -offset codes — decoding is just ``token_id - sid_base``. - -Prompts are rebuilt with the SAME splice as training, minus the answer: -``system + user_prefix + SID atoms + user_suffix + asst_prefix``, -left-padded with eos. The SID base is read from the exported tokenizer -(``convert_tokens_to_ids("C0")``), so tokenizer-layout drift is impossible. - -Run sharded across GPUs (one process per GPU, simple row interleave):: - - PYTHONPATH=. python -m examples.generative_rec_lm_predict \\ - --export_dir experiments//export_hf_40000 \\ - --test_parquet_dir /home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data_genreclm \\ - --test_csv /home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data/s1_tiny_test.csv \\ - --out logs/ter_predict_40000/output_sids.rank0.jsonl \\ - --rank 0 --world_size 8 -""" - -from __future__ import annotations - -import argparse -import csv -import glob -import json -import os -import sys -import time -from typing import List - -import pyarrow.parquet as pq -import torch -from transformers import AutoModelForCausalLM, AutoTokenizer - -from tzrec.models.escalating_beam import escalating_beam_search - -csv.field_size_limit(10 * 1024 * 1024) - - -def _enc(tok, text: str) -> List[int]: - return tok.encode(text, add_special_tokens=False) - - -def main() -> int: - ap = argparse.ArgumentParser() - ap.add_argument("--export_dir", required=True) - ap.add_argument("--test_parquet_dir", required=True) - ap.add_argument( - "--test_csv", - required=True, - help="raw test CSV; supplies the ground-truth `answer` item-id strings", - ) - ap.add_argument("--out", required=True) - ap.add_argument("--rank", type=int, default=0) - ap.add_argument("--world_size", type=int, default=1) - ap.add_argument( - "--bsz", - type=int, - default=4, - help="algr per_device_eval_batch_size=4; 8 OOMs at beam 50", - ) - ap.add_argument("--num_beams", type=int, default=50) - ap.add_argument("--num_return", type=int, default=50) - ap.add_argument("--max_new_tokens", type=int, default=3) - ap.add_argument( - "--dynamic_beam", - action="store_true", - help="ALGR-style escalating beam (width doubles per level); " - "returns num_beams * 2**max_new_tokens candidates, " - "ignoring --num_return.", - ) - ap.add_argument( - "--codebook", - type=int, - default=8192, - help="per-level SID codebook size (for --dynamic_beam band masking)", - ) - ap.add_argument( - "--max_rows", type=int, default=0, help="cap rows (0 = all); for smoke tests" - ) - # algr CN prompt fragments (must match the training config) - ap.add_argument( - "--system-instruction", - default=( - "你是一个推荐系统，根据用户的历史行为，预测用户在电商场景的下一步行为。" - "我会给你一串连续行为的语义编码，按照用户点击的时间顺序排列，每个行为用三个词表示。" - ), - ) - ap.add_argument("--user-prefix-text", default="当前用户的历史行为如下：") - ap.add_argument( - "--user-suffix-text", default="，请预测用户在电商推荐场景后续行为的语义编码" - ) - args = ap.parse_args() - - device = "cuda" if torch.cuda.is_available() else "cpu" - tok = AutoTokenizer.from_pretrained(args.export_dir, use_fast=True) - base = tok.convert_tokens_to_ids("C0") - assert base is not None and base > 0, "exported tokenizer lacks C atoms" - eos = tok.eos_token_id - - model = ( - AutoModelForCausalLM.from_pretrained(args.export_dir, torch_dtype="auto") - .to(device) - .eval() - ) - print( - f"[predict] model loaded vocab={model.config.vocab_size} sid_base={base} dev={device}", - flush=True, - ) - - # Per-level token-space SID bands (level j atom in [base+j*cb, base+(j+1)*cb-1]). - # Only needed by --dynamic_beam; the dataset offset codes are already disjoint. - L, cb = args.max_new_tokens, args.codebook - lo_tok = torch.tensor([base + j * cb for j in range(L)], device=device) - hi_tok = torch.tensor([base + (j + 1) * cb - 1 for j in range(L)], device=device) - if args.dynamic_beam: - print( - f"[predict] dynamic beam: widths " - f"{[args.num_beams * 2 ** (j + 1) for j in range(L)]} " - f"-> {args.num_beams * 2**L} candidates/row", - flush=True, - ) - - # Cached template fragments — identical composition to training splice. - tpl_system = _enc(tok, f"<|im_start|>system\n{args.system_instruction}<|im_end|>\n") - tpl_user_prefix = _enc(tok, f"<|im_start|>user\n{args.user_prefix_text}") - tpl_user_suffix = _enc(tok, f"{args.user_suffix_text}<|im_end|>\n") - tpl_asst_prefix = _enc(tok, "<|im_start|>assistant\n") - - # Rows: user_sequence SIDs from parquet, answer strings from the CSV - # (the converter wrote every CSV row, so ordering is 1:1). - paths = sorted(glob.glob(os.path.join(args.test_parquet_dir, "*.parquet"))) - user_rows: List[List[int]] = [] - for p in paths: - user_rows.extend( - pq.read_table(p, columns=["user_sequence"]) - .column("user_sequence") - .to_pylist() - ) - answers: List[str] = [] - with open(args.test_csv, encoding="utf-8", newline="") as f: - for row in csv.DictReader(f): - answers.append((row.get("answer") or "").strip()) - assert len(user_rows) == len(answers), (len(user_rows), len(answers)) - - idxs = list(range(len(user_rows)))[args.rank :: args.world_size] - if args.max_rows: - idxs = idxs[: args.max_rows] - print(f"[predict] rank={args.rank}/{args.world_size} rows={len(idxs)}", flush=True) - - os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True) - t0 = time.time() - n_done = 0 - with open(args.out, "w", encoding="utf-8") as out_f: - for bstart in range(0, len(idxs), args.bsz): - bidx = idxs[bstart : bstart + args.bsz] - prompts = [] - for i in bidx: - u_tok = [s + base - 1 for s in user_rows[i]] - prompts.append( - tpl_system - + tpl_user_prefix - + u_tok - + tpl_user_suffix - + tpl_asst_prefix - ) - T = max(len(p) for p in prompts) - input_ids = torch.full((len(prompts), T), eos, dtype=torch.long) - attn = torch.zeros((len(prompts), T), dtype=torch.long) - for r, p in enumerate(prompts): # left pad - input_ids[r, -len(p) :] = torch.tensor(p, dtype=torch.long) - attn[r, -len(p) :] = 1 - input_ids, attn = input_ids.to(device), attn.to(device) - - with torch.no_grad(): - if args.dynamic_beam: - new_tokens = escalating_beam_search( - model, - input_ids, - attn, - num_beams=args.num_beams, - lo_tok=lo_tok, - hi_tok=hi_tok, - ) # (B * num_beams*2**L, L) - else: - gen = model.generate( - input_ids=input_ids, - attention_mask=attn, - max_new_tokens=args.max_new_tokens, - num_beams=args.num_beams, - num_return_sequences=args.num_return, - do_sample=False, - pad_token_id=eos, - ) - new_tokens = gen[:, T:] # (B * num_return, max_new_tokens) - nret = new_tokens.shape[0] // len(bidx) - new_tokens = new_tokens.view(len(bidx), nret, -1).tolist() - - for row_i, seqs in zip(bidx, new_tokens): - texts = [] - for seq in seqs: - parts = [] - for t in seq: - if t >= base: - parts.append(f"C{t - base}") - else: - parts.append(tok.decode([t], skip_special_tokens=True)) - texts.append("".join(parts)) - out_f.write( - json.dumps( - {"_generated_text_": texts, "answer": answers[row_i]}, - ensure_ascii=False, - ) - + "\n" - ) - n_done += len(bidx) - if (bstart // args.bsz) % 25 == 0: - rate = n_done / max(time.time() - t0, 1e-6) - print( - f"[predict] {n_done}/{len(idxs)} rows ({rate:.1f} rows/s)", - flush=True, - ) - - print( - f"[predict] DONE rank={args.rank} rows={n_done} wall={time.time() - t0:.0f}s " - f"out={args.out}", - flush=True, - ) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/examples/generative_rec_lm_s1pretrained.config b/examples/generative_rec_lm_s1pretrained.config deleted file mode 100644 index ff31c9747..000000000 --- a/examples/generative_rec_lm_s1pretrained.config +++ /dev/null @@ -1,184 +0,0 @@ -# TorchEasyRec pipeline.config — `GenerativeRecLM` Qwen2.5-0.5B-Instruct -# translated from algr's -# /home/admin/workspace/al_sid/algr/config/qwen2.5_05b_3layer_s1_pretrained.json -# applied to the TINY parquet dataset -# /home/admin/workspace/aop_lab/data/AL-GR-Tiny/train_data_genreclm/*.parquet -# (produced by `examples/convert_s1tiny_to_parquet.py`; rows have -# `user_sequence: list` + `label: list`). -# -# algr → TER translation (knob by knob): -# per_device_train_batch_size : 80 → data_config.batch_size : 80 -# gradient_accumulation_steps : 4 → train_config.gradient_accumulation_steps : 4 *MUST MATCH* -# learning_rate : 5e-5 → dense_optimizer.adam_optimizer.lr : 5e-5 -# adam_beta1/2/epsilon : 0.9/0.999/1e-8 → adam_optimizer.{beta1,beta2,eps} -# lr_scheduler_type : linear → linear_decay_learning_rate (added to TER -# for this alignment; decays 5e-5 → 0). -# -# STEP-SEMANTICS NOTE (important): HF Trainer's `max_steps` counts -# OPTIMIZER steps, while TER's `num_steps` counts FORWARD steps and -# TZRecOptimizer skips `.step()` between accumulation boundaries -# (tzrec/optim/optimizer.py). TER's LR schedulers likewise tick once per -# forward step. With gradient_accumulation_steps=4, algr's max_steps -# 125000 therefore translates to: -# max_steps : 125000 → train_config.num_steps : 500000 (= 125000 × 4) -# save_steps : 10000 → train_config.save_checkpoints_steps : 40000 -# logging_steps : 1000 → train_config.log_step_count_steps : 4000 -# linear decay total : 125000 optimizer steps → total_size : 500000 forward steps -# -# Remaining benign divergence: HF Trainer scales the accumulated loss by -# 1/4 before backward (mean over micro-batches); TER sums the 4 -# micro-batch gradients. A constant gradient scale cancels in Adam's -# m̂/(√v̂+ε) update (up to ε), so trajectories match; logged ce_loss is -# per-micro-batch in both systems and stays directly comparable. -# bf16 : true → torch_dtype="auto" (the wrapper reads bf16 -# from the safetensors header). -# max_length / max_source_length / max_target_length -# : 1056/1024/32 → sequence_raw_feature.sequence_length : 1056 -# and capped by the data converter's row length -# (`--max-source-len` for the parquet driver). -# load_checkpoint_from : /home/admin/workspace/Qwen2.5-0.5B-Instruct -# → generative_rec_lm.hf_model_id : -# dataloader_num_workers : 4 → data_config.num_workers : 4 -# -# Launch (single H20:0, smoke): -# cd /home/admin/workspace/TorchEasyRec_qwen_smoke/TorchEasyRec -# export PYTHONPATH=$(pwd):$PYTHONPATH TRANSFORMERS_OFFLINE=1 HF_HUB_OFFLINE=1 -# torchrun --nnodes=1 --nproc-per-node=1 --master_port=32777 \ -# -m tzrec.train_eval \ -# --pipeline_config_path examples/generative_rec_lm_s1pretrained.config - -train_input_path: "/home/admin/workspace/aop_lab/data/AL-GR-Tiny/train_data_genreclm/*.parquet" -# held-out test split (s1_tiny_test.csv converted with the same script); -# evaluated (mean ce_loss) at every checkpoint save. -eval_input_path: "/home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data_genreclm/*.parquet" -model_dir: "experiments/generative_rec_lm_s1pretrained" - -train_config { - # GenerativeRecLM has no sparse params (the HF backbone owns its own - # `embed_tokens` and TER's EmbeddingGroup is unused). Define a no-op - # sparse_optimizer anyway since TER's TrainConfig requires it. - sparse_optimizer { - adagrad_optimizer { lr: 0.0 } - constant_learning_rate {} - } - dense_optimizer { - adam_optimizer { - lr: 5e-5 # algr `learning_rate: 5e-05` - beta1: 0.9 # algr `adam_beta1` - beta2: 0.999 # algr `adam_beta2` - eps: 1e-8 # algr `adam_epsilon` - } - # algr `lr_scheduler_type: linear` — decay 5e-5 → 0 over the whole - # run. TER schedulers tick per FORWARD step, so total_size is in - # forward steps (= 125000 optimizer steps × grad_accum 4). - linear_decay_learning_rate { - total_size: 500000 - } - } - # HF Trainer applies its DEFAULT `max_grad_norm: 1.0` global-norm - # clipping (algr's training_args don't override it; algr logs - # pre-clip grad_norm ~600 at step 1, so clipping is engaged). - # TER SUMS the grad_accum micro-grads where HF averages them, so the - # equivalent threshold is 1.0 × 4 = 4.0; the constant 4× gradient - # scale then cancels inside Adam's m̂/(√v̂+ε) update. - grad_clipping { - clipping_type: "norm" - max_gradient: 4.0 - norm_type: 2.0 - enable_global_grad_clip: true - } - # algr `max_steps: 125000` OPTIMIZER steps × grad_accum 4 = 500000 - # forward steps (see step-semantics note above). - num_steps: 500000 - # algr `gradient_accumulation_steps: 4` — MUST MATCH EXACTLY. - gradient_accumulation_steps: 4 - # algr `save_steps: 10000` optimizer steps × 4. - save_checkpoints_steps: 40000 - # algr `logging_steps: 1000` optimizer steps × 4. - log_step_count_steps: 4000 - # bf16 — TER reads from the safetensors header via torch_dtype="auto" - # at `from_pretrained` time. No flag here. -} - -eval_config { -} - -data_config { - # algr `per_device_train_batch_size: 80`. - batch_size: 80 - dataset_type: ParquetDataset - fg_mode: FG_NONE - # algr `dataloader_num_workers: 4`. - num_workers: 4 - # The answer SID column is the TARGET, so it's a label_field (a list - # column -> batch.jagged_labels["label"]), NOT a feature. build_input reads the - # FIRST label_field (like RankModel's labels[0]). Being a label_field, it can be - # absent at inference (no ground truth) without the EmbeddingGroup needing it. - label_fields: ["label"] -} - -# user_sequence (the history INPUT) maps 1:1 to the parquet column produced by -# `convert_s1tiny_to_parquet.py`. (The "label" column is a label_field, above.) -feature_configs { - sequence_raw_feature { - feature_name: "user_sequence" - expression: "user:user_sequence" - # Under FG_NONE this length is NOT auto-truncated by the reader, so the - # model enforces it via a recency-preserving clip in `_sid_token_rows` - # (keep newest items, drop oldest). The model's cap is common. - # max_sequence_length (below), not this field. - sequence_length: 300 - value_dim: 1 - } -} - -model_config { - # The SINGLE feature_group = the history INPUT: build_input retrieves the raw - # SID stream via the EmbeddingGroup and keys it by GROUP name (HSTU idiom). The - # model takes the only declared group as the history (group_name is free — it - # need not be "user_seq"). The answer is NOT a group — it's the label_field - # above (batch.jagged_labels). - feature_groups { - group_name: "user_seq" - feature_names: "user_sequence" - group_type: JAGGED_SEQUENCE - } - qwen2_rec_lm { - # Qwen2 backbone (owned by this family message, not `common`). - # algr `load_checkpoint_from: /home/admin/workspace/Qwen2.5-0.5B` - # (the BASE model — qwen2.5_05b_3layer_s1_pretrained.json does not use - # the Instruct variant). Omit to inherit the default "Qwen/Qwen2.5-0.5B". - hf_model_id: "/home/admin/workspace/Qwen2.5-0.5B" - - # architecture-agnostic config shared by all generative-rec families - common { - # SID codebook: ONE entry per RQ level. AL-GR is 3 levels x 8192 - # (verified from item_info: every codebook_lv* ∈ [0, 8191]), so - # len(codebook)=3 = SID codes per answer and sum=24 576 = atoms - # appended to the vocab (an exact fit — max SID atom is 24 575). - codebook: 8192 - codebook: 8192 - codebook: 8192 - vocab_pad_to_multiple_of: 128 - - # The history feature_group and the SID-column names are all derived, - # not configured: the history is the single feature_group above and its - # one member is the history feature; the answer is the first - # data_config.label_field. - ignore_index: -100 - - # Model history budget (HSTU-style model knob): truncation cap + - # activation-pool pre-size. Distinct from the user_sequence feature's - # sequence_length (kept as data/export metadata). 0 = off (no cap). - max_sequence_length: 300 - } - - # algr's row.system / `default_instruction` — the CN recommender prompt. - # Matching this string bit-for-bit reproduces algr's input. - system_instruction: "你是一个推荐系统，根据用户的历史行为，预测用户在电商场景的下一步行为。我会给你一串连续行为的语义编码，按照用户点击的时间顺序排列，每个行为用三个词表示。" - - # algr's CN user-prompt wrappers around the SID list. - user_prefix_text: "当前用户的历史行为如下：" - user_suffix_text: "，请预测用户在电商推荐场景后续行为的语义编码" - } -} diff --git a/examples/generative_rec_lm_s2pretrained.config b/examples/generative_rec_lm_s2pretrained.config deleted file mode 100644 index ee3dce1d2..000000000 --- a/examples/generative_rec_lm_s2pretrained.config +++ /dev/null @@ -1,86 +0,0 @@ -train_input_path: "/home/admin/workspace/aop_lab/data/AL-GR-Tiny/train_data_genreclm_s2/*.parquet" -eval_input_path: "/home/admin/workspace/aop_lab/data/AL-GR-Tiny/test_data_genreclm_s2/*.parquet" -model_dir: "experiments/generative_rec_lm_s2pretrained" -train_config { - sparse_optimizer { - adagrad_optimizer { - lr: 0.0 - } - constant_learning_rate { - } - } - dense_optimizer { - adam_optimizer { - lr: 5e-05 - beta1: 0.9 - beta2: 0.999 - eps: 1e-08 - } - linear_decay_learning_rate { - total_size: 500000 - } - } - num_steps: 24000 - save_checkpoints_steps: 8000 - log_step_count_steps: 1000 - gradient_accumulation_steps: 4 - grad_clipping { - clipping_type: "norm" - max_gradient: 4.0 - norm_type: 2.0 - enable_global_grad_clip: true - } -} -eval_config { -} -data_config { - batch_size: 80 - dataset_type: ParquetDataset - num_workers: 4 - fg_mode: FG_NONE - # answer = TARGET -> label_field (list -> batch.jagged_labels["label"]), - # not a feature; build_input reads the first label_field (like RankModel labels[0]). - label_fields: ["label"] -} -feature_configs { - sequence_raw_feature { - feature_name: "user_sequence" - expression: "user:user_sequence" - value_dim: 1 - sequence_length: 1056 - } -} -model_config { - # The SINGLE feature_group = the history INPUT (raw SID passthrough; see - # s1pretrained / build_input). The model takes the only declared group as the - # history (group_name is free). The answer is the label_field above - # (batch.jagged_labels), not a group. - feature_groups { - group_name: "user_seq" - feature_names: "user_sequence" - group_type: JAGGED_SEQUENCE - } - qwen2_rec_lm { - # Qwen2 backbone (owned by this family message, not `common`). - hf_model_id: "/home/admin/workspace/Qwen2.5-0.5B" - common { - # The history feature_group and SID-column names are derived (history = the - # single feature_group above + its one member; answer = the first - # data_config.label_field), not configured here. - ignore_index: -100 - # SID codebook: one entry per RQ level (AL-GR = 3 levels x 8192; verified - # from item_info codebook_lv* ranges). len = SID codes/answer, sum = vocab - # atoms. - codebook: 8192 - codebook: 8192 - codebook: 8192 - vocab_pad_to_multiple_of: 128 - # Model history budget (required): truncation cap + activation-pool pre-size - # (matches the user_sequence feature's sequence_length here). - max_sequence_length: 1056 - } - system_instruction: "你是一个推荐系统，根据用户的历史行为，预测用户在电商场景的下一步行为。我会给你一串连续行为的语义编码，按照用户点击的时间顺序排列，每个行为用三个词表示。" - user_prefix_text: "当前用户的历史行为如下：" - user_suffix_text: "，请预测用户在电商推荐场景后续行为的语义编码" - } -} diff --git a/examples/generative_rec_lm_smoke.py b/examples/generative_rec_lm_smoke.py deleted file mode 100644 index 87954990f..000000000 --- a/examples/generative_rec_lm_smoke.py +++ /dev/null @@ -1,149 +0,0 @@ -"""Smoke test for ``GenerativeRecLM`` — exercises the full forward path -(splice + base forward + suffix slice + HF loss_function) without going -through ``tzrec.train_eval``. - -Validates: - 1. Proto + dispatch wire up. - 2. Vocab extension preserves SID-token offset arithmetic. - 3. Left-pad with eos_token_id (L7 fix) produces a finite, non-trivial CE. - 4. SID → token mapping (``token = sid + base_vocab - 1``) hits the new atoms. - -Usage: - cd /workspace/fangtinglin/codework/feat/support_qwen/TorchEasyRec - /opt/conda/bin/python -m examples.generative_rec_lm_smoke -""" - -from __future__ import annotations - -import os -import sys -import time - -import torch -from google.protobuf import text_format - -from tzrec.datasets.utils import Batch -from tzrec.models import generative_rec_lm # noqa: F401 registers class -from tzrec.models.model import BaseModel -from tzrec.protos.model_pb2 import ModelConfig - - -# Tiny config — small codebook so vocab extension stays cheap on CPU. -HF_MODEL_ID = "/workspace/fangtinglin/_hf_stage/Qwen2.5-0.5B" -CODEBOOK = [64, 64] # 128 SID atoms total → small vocab grow on CPU -USER_FEATURE = "user_sequence" -LABEL_FEATURE = "label" - - -def _make_proto(): - cfg = ModelConfig() - grl = cfg.generative_rec_lm - grl.class_name = "Qwen2RecLM" - grl.hf_model_id = HF_MODEL_ID - for c in CODEBOOK: - grl.codebook.append(c) - grl.user_sequence_feature_name = USER_FEATURE - grl.label_feature_name = LABEL_FEATURE - grl.ignore_index = -100 - return cfg - - -class _DummyJagged: - """Minimal stand-in for torchrec's JaggedTensor — just enough surface - for ``GenerativeRecLM._jagged_to_row_list`` to consume.""" - def __init__(self, values: torch.Tensor, lengths: torch.Tensor): - self._values = values - self._lengths = lengths - - def values(self) -> torch.Tensor: - return self._values - - def lengths(self) -> torch.Tensor: - return self._lengths - - -def _make_batch(B: int, sum_codebook: int, user_len: int, label_len: int) -> Batch: - # SID indices are 1-indexed in [1, sum(codebook)]. - torch.manual_seed(0) - user_vals = torch.randint(1, sum_codebook + 1, (B * user_len,), dtype=torch.long) - label_vals = torch.randint(1, sum_codebook + 1, (B * label_len,), dtype=torch.long) - user_lens = torch.full((B,), user_len, dtype=torch.long) - label_lens = torch.full((B,), label_len, dtype=torch.long) - - seq_dense = { - USER_FEATURE: _DummyJagged(user_vals, user_lens), - LABEL_FEATURE: _DummyJagged(label_vals, label_lens), - } - # Batch is a dataclass-like NamedTuple — construct with just the field we need. - # tzrec.datasets.utils.Batch is a NamedTuple of many fields, all default to {}. - return Batch(sequence_dense_features=seq_dense) - - -def main() -> int: - os.environ["TZREC_GENRECLM_DEBUG"] = "1" - print(f"[smoke] torch={torch.__version__} cuda_avail={torch.cuda.is_available()}", flush=True) - - cfg = _make_proto() - print(f"[smoke] proto: class_name={cfg.generative_rec_lm.class_name} " - f"codebook={list(cfg.generative_rec_lm.codebook)}", flush=True) - - # --- dispatch --- - model_cls = BaseModel.create_class("GenerativeRecLM") - print(f"[smoke] dispatched -> {model_cls.__name__}", flush=True) - - # --- construct --- - t0 = time.time() - model = model_cls( - cfg, - features=[], - labels=[], - sample_weights=None, - ) - model.train() - print(f"[smoke] construct ok ({time.time()-t0:.1f}s); " - f"base_vocab={model._base_vocab} " - f"final_vocab={model.lm.config.vocab_size} " - f"pad_id={model._pad_token_id}", flush=True) - print(f"[smoke] tpl_system numel={model.tpl_system.numel()} " - f"tpl_user_prefix numel={model.tpl_user_prefix.numel()} " - f"tpl_asst_prefix numel={model.tpl_asst_prefix.numel()}", flush=True) - - # --- batch --- - sum_codebook = sum(CODEBOOK) - batch = _make_batch(B=2, sum_codebook=sum_codebook, user_len=5, label_len=4) - - # --- forward --- - t0 = time.time() - with torch.no_grad(): - pred = model.predict(batch) - elapsed = time.time() - t0 - loss = pred["loss"] - logits = pred["logits"] - print(f"[smoke] forward ok ({elapsed:.2f}s); " - f"loss={float(loss):.4f} logits.shape={tuple(logits.shape)}", flush=True) - - # Acceptance criteria - ok = True - if not torch.isfinite(loss): - print("[smoke] FAIL: loss is non-finite") - ok = False - if float(loss) < 0.1 or float(loss) > 100.0: - print(f"[smoke] FAIL: loss out of plausible range ({float(loss)})") - ok = False - if logits.shape[0] != 2: - print(f"[smoke] FAIL: logits batch dim wrong ({logits.shape})") - ok = False - - # --- one backward step to make sure grads flow through the chat-template - # buffers + base model + lm_head without error --- - t0 = time.time() - pred = model.predict(batch) - pred["loss"].backward() - print(f"[smoke] backward ok ({time.time()-t0:.2f}s)", flush=True) - - print(f"[smoke] {'PASS' if ok else 'FAIL'}", flush=True) - return 0 if ok else 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/examples/generative_rec_lm_train_loop.py b/examples/generative_rec_lm_train_loop.py deleted file mode 100644 index a019923f0..000000000 --- a/examples/generative_rec_lm_train_loop.py +++ /dev/null @@ -1,230 +0,0 @@ -"""Multi-step CPU training loop for ``GenerativeRecLM``. - -Exercises ~hundreds of forward+backward+optimizer steps on synthetic batches -to verify the model learns. The data is a tiny next-SID-prediction task: -each row maps a sampled ``user_sequence`` of SIDs to a deterministic -``label`` SID list derived from it via a fixed permutation, so a sufficiently -expressive LM can drive the CE loss down toward 0. - -Emits one ``[step …]`` line every ``--log-every`` steps so a Monitor wrapper -can stream loss progress. Failure modes (Traceback, NaN/Inf loss) also go to -stdout so the monitor sees them. - -Usage: - cd /workspace/fangtinglin/codework/feat/support_qwen/TorchEasyRec - /opt/conda/bin/python -m examples.generative_rec_lm_train_loop \\ - --steps 500 --bsz 2 --user-len 6 --label-len 3 \\ - --log-every 5 -""" - -from __future__ import annotations - -import argparse -import math -import os -import sys -import time -from typing import List - -import torch - -from tzrec.datasets.utils import Batch -from tzrec.models import generative_rec_lm # noqa: F401 registers class -from tzrec.models.model import BaseModel -from tzrec.protos.model_pb2 import ModelConfig - - -HF_MODEL_ID_DEFAULT = "/workspace/fangtinglin/_hf_stage/Qwen2.5-0.5B" - - -class _DummyJagged: - def __init__(self, values: torch.Tensor, lengths: torch.Tensor): - self._values = values - self._lengths = lengths - - def values(self) -> torch.Tensor: - return self._values - - def lengths(self) -> torch.Tensor: - return self._lengths - - -def _make_proto(codebook: List[int], hf_model_id: str) -> ModelConfig: - cfg = ModelConfig() - grl = cfg.generative_rec_lm - grl.class_name = "Qwen2RecLM" - grl.hf_model_id = hf_model_id - for c in codebook: - grl.codebook.append(c) - grl.user_sequence_feature_name = "user_sequence" - grl.label_feature_name = "label" - grl.ignore_index = -100 - return cfg - - -def _sample_batch( - rng: torch.Generator, - sum_codebook: int, - bsz: int, - user_len: int, - label_len: int, -) -> Batch: - """Synthesise a deterministic-mapping batch: ``label = (user mod K) + 1`` - for the first ``label_len`` positions. Easy to memorise; gives the LM - something to drive CE down.""" - user_vals = torch.randint( - 1, sum_codebook + 1, (bsz, user_len), generator=rng, dtype=torch.long, - ) - # Deterministic label: pick the first label_len user positions, modded - # back into [1, sum_codebook]. Small enough to fit in CE in ~hundreds of - # steps even on a frozen-most-of-the-net base. - label_vals = ((user_vals[:, :label_len] % sum_codebook) + 1).contiguous() - user_flat = user_vals.reshape(-1) - label_flat = label_vals.reshape(-1) - user_lens = torch.full((bsz,), user_len, dtype=torch.long) - label_lens = torch.full((bsz,), label_len, dtype=torch.long) - return Batch(sequence_dense_features={ - "user_sequence": _DummyJagged(user_flat, user_lens), - "label": _DummyJagged(label_flat, label_lens), - }) - - -def main() -> int: - ap = argparse.ArgumentParser() - ap.add_argument("--steps", type=int, default=500) - ap.add_argument("--bsz", type=int, default=2) - ap.add_argument("--user-len", type=int, default=6) - ap.add_argument("--label-len", type=int, default=3) - ap.add_argument("--log-every", type=int, default=5) - # 1e-5 matches algr's `qwen2.5_05b_3layer_s1tiny.json` (see - # [[project-tzrec-qwen2-integration]]). - ap.add_argument("--lr", type=float, default=1e-5) - ap.add_argument("--hf-model-id", default=HF_MODEL_ID_DEFAULT) - ap.add_argument( - "--codebook", default="64,64", - help="comma-separated codebook sizes; 64,64 fast / 32768,32768 algr-like", - ) - ap.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"]) - args = ap.parse_args() - - if args.device == "auto": - device = "cuda" if torch.cuda.is_available() else "cpu" - else: - device = args.device - print( - f"[init] torch={torch.__version__} device={device} " - f"steps={args.steps} bsz={args.bsz} user_len={args.user_len} " - f"label_len={args.label_len} lr={args.lr} hf={args.hf_model_id}", - flush=True, - ) - - codebook = [int(x) for x in args.codebook.split(",") if x] - sum_codebook = sum(codebook) - cfg = _make_proto(codebook, args.hf_model_id) - - model_cls = BaseModel.create_class("GenerativeRecLM") - t0 = time.time() - model = model_cls(cfg, features=[], labels=[], sample_weights=None) - model.train() - if device == "cuda": - model.to("cuda") - torch.cuda.reset_peak_memory_stats() - print( - f"[init] construct ok ({time.time()-t0:.1f}s) " - f"base_vocab={model._base_vocab} " - f"final_vocab={model.lm.config.vocab_size} " - f"sum_codebook={sum_codebook}", - flush=True, - ) - - optim = torch.optim.AdamW( - model.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-8, - ) - - rng = torch.Generator(device="cpu").manual_seed(0) - # Warmup: discount the first 3 steps from throughput accounting (CUDA - # kernel autotune, cudnn benchmark, allocator warm-up). - warmup_steps = 3 if device == "cuda" else 0 - start = None - losses: List[float] = [] - step_walls: List[float] = [] - for step in range(1, args.steps + 1): - batch = _sample_batch( - rng, sum_codebook, args.bsz, args.user_len, args.label_len, - ) - if device == "cuda": - # Move tensors to GPU — small batches, ok to do per-step. - for k, v in batch.sequence_dense_features.items(): - v._values = v._values.cuda(non_blocking=True) - v._lengths = v._lengths.cuda(non_blocking=True) - torch.cuda.synchronize() - t_step = time.perf_counter() - pred = model.predict(batch) - loss = pred["loss"] - - loss_val = loss.detach().item() - if not math.isfinite(loss_val): - print(f"[step {step}] FAIL: non-finite loss {loss_val}", flush=True) - return 1 - - optim.zero_grad(set_to_none=True) - loss.backward() - optim.step() - if device == "cuda": - torch.cuda.synchronize() - dt = time.perf_counter() - t_step - losses.append(loss.detach().item()) - - if step == warmup_steps: - start = time.time() - step_walls = [] - if device == "cuda": - torch.cuda.reset_peak_memory_stats() - elif step > warmup_steps: - step_walls.append(dt) - - if step % args.log_every == 0 or step == 1: - mem_str = "" - if device == "cuda": - peak_gb = torch.cuda.max_memory_allocated() / 1024**3 - resv_gb = torch.cuda.max_memory_reserved() / 1024**3 - mem_str = f" peak_alloc={peak_gb:.2f}GB peak_reserved={resv_gb:.2f}GB" - if start is not None and step_walls: - avg_dt = sum(step_walls) / len(step_walls) - it_s = 1.0 / avg_dt - wall_str = f" avg_step={avg_dt*1000:.1f}ms it/s={it_s:.2f}" - else: - wall_str = " (warmup)" - print( - f"[step {step}/{args.steps}] ce_loss={loss_val:.4f} " - f"avg_last_{min(20, len(losses))}={sum(losses[-20:])/min(20, len(losses)):.4f}" - f"{wall_str}{mem_str}", - flush=True, - ) - - # ---- final report ---- - start_avg = sum(losses[:5]) / min(5, len(losses)) - end_avg = sum(losses[-5:]) / min(5, len(losses)) - if step_walls: - median_dt = sorted(step_walls)[len(step_walls) // 2] - avg_dt = sum(step_walls) / len(step_walls) - peak_str = "" - if device == "cuda": - peak_gb = torch.cuda.max_memory_allocated() / 1024**3 - resv_gb = torch.cuda.max_memory_reserved() / 1024**3 - peak_str = ( - f" peak_alloc={peak_gb:.2f}GB peak_reserved={resv_gb:.2f}GB" - ) - print( - f"[done] steps={args.steps} bsz={args.bsz} " - f"start_avg5={start_avg:.4f} end_avg5={end_avg:.4f} " - f"delta={start_avg-end_avg:+.4f} " - f"avg_step={avg_dt*1000:.1f}ms median_step={median_dt*1000:.1f}ms " - f"it/s={1.0/avg_dt:.2f}{peak_str}", - flush=True, - ) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/examples/generative_rec_lm_train_loop_parquet.py b/examples/generative_rec_lm_train_loop_parquet.py deleted file mode 100644 index 434f27600..000000000 --- a/examples/generative_rec_lm_train_loop_parquet.py +++ /dev/null @@ -1,247 +0,0 @@ -"""Parquet-fed multi-step training loop for ``GenerativeRecLM``. - -Same instrumentation as ``generative_rec_lm_train_loop.py`` (GPU memory + -throughput) but consumes real ``user_sequence + label`` rows from a parquet -directory instead of synthesising them. Used for apples-to-apples -comparison with algr on the same data. - -Usage on remote:: - - /opt/conda/bin/python -m examples.generative_rec_lm_train_loop_parquet \\ - --parquet_dir /home/admin/workspace/aop_lab/data/AL-GR-Tiny/train_data_genreclm_100k \\ - --steps 25 --log-every 1 --bsz 10 --device cuda \\ - --codebook 32768,32768 \\ - --hf-model-id /home/admin/workspace/Qwen2.5-0.5B-Instruct -""" - -from __future__ import annotations - -import argparse -import glob -import math -import os -import sys -import time -from typing import List - -import pyarrow.parquet as pq -import torch - -from tzrec.datasets.utils import Batch -from tzrec.models import generative_rec_lm # noqa: F401 registers class -from tzrec.models.model import BaseModel -from tzrec.protos.model_pb2 import ModelConfig - - -class _DummyJagged: - def __init__(self, values: torch.Tensor, lengths: torch.Tensor): - self._values = values - self._lengths = lengths - - def values(self) -> torch.Tensor: - return self._values - - def lengths(self) -> torch.Tensor: - return self._lengths - - -class _ParquetIter: - """Round-robin iterator over rows in a parquet directory. - - Loads each shard into memory as Arrow Tables (small for our 100K split). - Yields ``(user_sids: list[int], label_sids: list[int])`` rows. Loops to - the start when exhausted so the training loop never runs out of data. - """ - def __init__(self, parquet_dir: str, max_source_len: int): - paths = sorted(glob.glob(os.path.join(parquet_dir, "*.parquet"))) - assert paths, f"no parquet files under {parquet_dir!r}" - self._rows: List = [] - for p in paths: - t = pq.read_table(p) - u_col = t["user_sequence"].to_pylist() - l_col = t["label"].to_pylist() - for u, lab in zip(u_col, l_col): - # algr truncates user prompts to ``max_source_length`` tokens - # via ``prompt_ids[:max_source_length]``. We mirror that at - # the SID-list level (each SID is one token after splice). - if max_source_len and len(u) > max_source_len: - u = u[-max_source_len:] # keep most-recent items - self._rows.append((u, lab)) - self._n = len(self._rows) - self._i = 0 - - def __len__(self) -> int: - return self._n - - def next_batch(self, bsz: int) -> Batch: - rows = [] - for _ in range(bsz): - rows.append(self._rows[self._i]) - self._i = (self._i + 1) % self._n - user_lens = torch.tensor([len(r[0]) for r in rows], dtype=torch.long) - label_lens = torch.tensor([len(r[1]) for r in rows], dtype=torch.long) - # Flatten to JaggedTensor (values, lengths) shape - user_vals: List[int] = [] - for r in rows: - user_vals.extend(r[0]) - label_vals: List[int] = [] - for r in rows: - label_vals.extend(r[1]) - return Batch(sequence_dense_features={ - "user_sequence": _DummyJagged( - torch.tensor(user_vals, dtype=torch.long), user_lens, - ), - "label": _DummyJagged( - torch.tensor(label_vals, dtype=torch.long), label_lens, - ), - }) - - -def _make_proto( - codebook: List[int], - hf_model_id: str, - system_instruction: str = "", - user_prefix_text: str = "", - user_suffix_text: str = "", -) -> ModelConfig: - cfg = ModelConfig() - grl = cfg.generative_rec_lm - grl.class_name = "Qwen2RecLM" - grl.hf_model_id = hf_model_id - for c in codebook: - grl.codebook.append(c) - grl.user_sequence_feature_name = "user_sequence" - grl.label_feature_name = "label" - grl.ignore_index = -100 - if system_instruction: - grl.system_instruction = system_instruction - if user_prefix_text: - grl.user_prefix_text = user_prefix_text - if user_suffix_text: - grl.user_suffix_text = user_suffix_text - return cfg - - -def main() -> int: - ap = argparse.ArgumentParser() - ap.add_argument("--parquet_dir", required=True) - ap.add_argument("--steps", type=int, default=25) - ap.add_argument("--bsz", type=int, default=10) - ap.add_argument("--log-every", type=int, default=1) - ap.add_argument("--lr", type=float, default=1e-5) - ap.add_argument("--max-source-len", type=int, default=1020, - help="cap user_sequence SID count (mirrors algr max_source_length)") - ap.add_argument("--codebook", default="32768,32768") - ap.add_argument("--hf-model-id", required=True) - ap.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"]) - ap.add_argument("--system-instruction", default="", - help="CN/EN system prompt; matches algr's `default_instruction` / row.system") - ap.add_argument("--user-prefix-text", default="", - help="text wrapped before SID list in user message (e.g. algr CN prefix)") - ap.add_argument("--user-suffix-text", default="", - help="text wrapped after SID list in user message (e.g. algr CN suffix)") - args = ap.parse_args() - - device = "cuda" if args.device == "auto" and torch.cuda.is_available() else args.device - if device == "auto": - device = "cpu" - print( - f"[init] torch={torch.__version__} device={device} steps={args.steps} " - f"bsz={args.bsz} lr={args.lr} parquet={args.parquet_dir}", - flush=True, - ) - - t_load = time.time() - data = _ParquetIter(args.parquet_dir, args.max_source_len) - print(f"[init] parquet loaded ({time.time()-t_load:.1f}s) rows={len(data)}", flush=True) - - codebook = [int(x) for x in args.codebook.split(",") if x] - cfg = _make_proto( - codebook, args.hf_model_id, - system_instruction=args.system_instruction, - user_prefix_text=args.user_prefix_text, - user_suffix_text=args.user_suffix_text, - ) - model_cls = BaseModel.create_class("GenerativeRecLM") - t0 = time.time() - model = model_cls(cfg, features=[], labels=[], sample_weights=None) - model.train() - if device == "cuda": - model.to("cuda") - torch.cuda.reset_peak_memory_stats() - print( - f"[init] construct ok ({time.time()-t0:.1f}s) base_vocab={model._base_vocab} " - f"final_vocab={model.lm.config.vocab_size}", - flush=True, - ) - - optim = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-8) - - warmup = 3 if device == "cuda" else 0 - start = None - step_walls: List[float] = [] - losses: List[float] = [] - for step in range(1, args.steps + 1): - batch = data.next_batch(args.bsz) - if device == "cuda": - for v in batch.sequence_dense_features.values(): - v._values = v._values.cuda(non_blocking=True) - v._lengths = v._lengths.cuda(non_blocking=True) - torch.cuda.synchronize() - t = time.perf_counter() - pred = model.predict(batch) - loss = pred["loss"] - loss_val = loss.detach().item() - if not math.isfinite(loss_val): - print(f"[step {step}] FAIL: non-finite loss {loss_val}", flush=True) - return 1 - optim.zero_grad(set_to_none=True) - loss.backward() - optim.step() - if device == "cuda": - torch.cuda.synchronize() - dt = time.perf_counter() - t - losses.append(loss_val) - if step == warmup: - start = time.time() - step_walls = [] - if device == "cuda": - torch.cuda.reset_peak_memory_stats() - elif step > warmup: - step_walls.append(dt) - if step % args.log_every == 0 or step == 1: - mem_str = "" - if device == "cuda": - pa = torch.cuda.max_memory_allocated() / 1024**3 - pr = torch.cuda.max_memory_reserved() / 1024**3 - mem_str = f" peak_alloc={pa:.2f}GB peak_reserved={pr:.2f}GB" - wall_str = ( - f" avg_step={sum(step_walls)/len(step_walls)*1000:.1f}ms " - f"it/s={1.0/(sum(step_walls)/len(step_walls)):.2f}" - if step_walls else " (warmup)" - ) - ph_t = batch.sequence_dense_features["user_sequence"]._values.shape[0] - print( - f"[step {step}/{args.steps}] ce_loss={loss_val:.4f} " - f"avg_last_{min(20, len(losses))}={sum(losses[-20:])/min(20, len(losses)):.4f}" - f" sum_lens={ph_t}{wall_str}{mem_str}", - flush=True, - ) - if step_walls: - pa = (torch.cuda.max_memory_allocated() / 1024**3) if device == "cuda" else 0 - pr = (torch.cuda.max_memory_reserved() / 1024**3) if device == "cuda" else 0 - print( - f"[done] steps={args.steps} bsz={args.bsz} " - f"start_avg5={sum(losses[:5])/min(5, len(losses)):.4f} " - f"end_avg5={sum(losses[-5:])/min(5, len(losses)):.4f} " - f"avg_step={sum(step_walls)/len(step_walls)*1000:.1f}ms " - f"median_step={sorted(step_walls)[len(step_walls)//2]*1000:.1f}ms " - f"it/s={1.0/(sum(step_walls)/len(step_walls)):.2f}" - f" peak_alloc={pa:.2f}GB peak_reserved={pr:.2f}GB", - flush=True, - ) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index aa250b8bc..d8cbb1930 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -14,12 +14,14 @@ ``GenerativeRecLMConfig`` (the family message's ``common`` field). * Streaming sample format: each row carries the history ``user_sequence`` (a raw-int64 sequence feature) and the ``label`` answer (a ``data_config`` -label_field), both holding raw SID indices in ``[1, sum(codebook)]``. +label_field), both holding local, 1-based per-level codes in +``[1, codebook[level]]``. Rows contain whole items in level order. * The chat template is tokenised ONCE at ``__init__`` and cached as non-persistent buffers, so per-batch encoding is integer arithmetic only (no HF tokenizer in the hot path). -* SID -> token id by integer offset: ``token = sid + base_vocab - 1`` (the SID -atoms ``C0..C{sum-1}`` are added right after the original vocabulary). +* Raw code -> token id by integer offsets: +``token = base_vocab + level_offset[level] + code - 1``. The model derives +``level_offsets`` from ``codebook``; sample writers must not pre-apply them. * Left padding with ``eos_token_id``: real content sits at the END of every row so the suffix slice captures only ``[response + end_markers]``. """ @@ -138,11 +140,10 @@ def _read_common_config(self, common: Any) -> int: if len(codebook) == 0: raise ValueError("GenerativeRecLM: codebook must be non-empty.") self._num_levels = len(codebook) - # inference-gate validity bands; non-persistent — derived from codebook, - # kept off the state_dict (HF safetensors round-trip). + # Level layout is derived from codebook and kept out of state_dict. lo, hi = self._sid_level_bands(codebook) - self.register_buffer("_sid_lvl_lo", lo, persistent=False) - self.register_buffer("_sid_lvl_hi", hi, persistent=False) + self.register_buffer("_level_offsets", lo - 1, persistent=False) + self.register_buffer("_codebook_sizes", hi - lo + 1, persistent=False) self._vocab_pad_mult = int(common.vocab_pad_to_multiple_of) or 128 return sum(int(c) for c in codebook) @@ -274,26 +275,31 @@ def device(self) -> torch.device: """Device the HF backbone runs on — the single source for model I/O.""" return self.lm.device - def _tokenize_sids(self, sids: torch.Tensor) -> torch.Tensor: - """Map raw 1-indexed SID values to extended-vocab token ids. + def _tokenize_sids( + self, codes: torch.Tensor, level_ids: torch.Tensor + ) -> torch.Tensor: + """Map local 1-based per-level codes to extended-vocab token ids. + + ``level_ids`` must broadcast against ``codes`` and identify each code's + RQ level. Atom ``C{k}`` sits at ``base_vocab + k``, therefore: - Atom ``C{k}`` sits at ``base_vocab + k`` (atoms appended right after the - original vocab), so ``token_id = sid + base_vocab - 1``. The integer - counterpart of the HF tokenizer used for text; shape-agnostic. + ``token = code + level_offsets[level] + base_vocab - 1``. """ - return sids + (self._base_vocab - 1) + return codes + self._level_offsets[level_ids] + (self._base_vocab - 1) - def _detokenize_sids(self, tokens: torch.Tensor) -> torch.Tensor: - """Inverse of ``_tokenize_sids``: token id -> raw 1-indexed SID.""" - return tokens - (self._base_vocab - 1) + def _detokenize_sids( + self, tokens: torch.Tensor, level_ids: torch.Tensor + ) -> torch.Tensor: + """Inverse of ``_tokenize_sids``: token ids to local 1-based codes.""" + return tokens - (self._base_vocab - 1) - self._level_offsets[level_ids] @staticmethod def _sid_level_bands(codebook: Any) -> tuple[torch.Tensor, torch.Tensor]: - """Per-level closed SID bands ``(lo, hi)`` as ``(num_levels,)`` long tensors. + """Per-level flattened 1-based SID bands as two long tensors. Level ``j`` occupies a DISJOINT band ``[offset_j + 1, offset_j + codebook[j]]`` where ``offset_j = sum(codebook[:j])``. Single source of - truth: both ``_read_common_config`` and the tests build bands from this. + truth for the derived ``_level_offsets`` and ``_codebook_sizes`` buffers. """ lo, hi, acc = [], [], 0 for c in codebook: @@ -305,26 +311,40 @@ def _sid_level_bands(codebook: Any) -> tuple[torch.Tensor, torch.Tensor]: torch.tensor(hi, dtype=torch.long), ) + def _sid_token_bands(self) -> tuple[torch.Tensor, torch.Tensor]: + """Return the inclusive token-id band for every SID level.""" + level_ids = torch.arange(self._num_levels, device=self.device) + return ( + self._tokenize_sids(torch.ones_like(level_ids), level_ids), + self._tokenize_sids(self._codebook_sizes, level_ids), + ) + def _validate_sid_candidates( self, new_tokens: torch.Tensor, batch_size: int ) -> torch.Tensor: - """Map a generated token tail back to SIDs and reject malformed candidates. + """Decode generated tokens to local 1-based codes and reject bad beams. ``new_tokens`` is the per-beam tail ``(B*num_return, w)`` (``w`` may be < ``num_levels`` when beams stop early). Returns ``(batch_size, num_return, - num_levels)`` SIDs with every malformed candidate (early EOS / non-SID / - wrong-level atom) set to ``-1`` — which can never match a real item. + num_levels)`` local codes. Every malformed candidate (early EOS / + non-SID / wrong-level atom) is set to ``-1``, which cannot match a real + 1-based code. """ - sids = self._detokenize_sids(new_tokens) + level_ids = torch.arange(new_tokens.shape[1], device=new_tokens.device) + codes = self._detokenize_sids(new_tokens, level_ids) # early-EOS beams return < num_levels tokens; pad to num_levels with -1. - sids = F.pad(sids, (0, self._num_levels - sids.shape[1]), value=-1) + codes = F.pad( + codes, + (0, self._num_levels - codes.shape[1]), + value=-1, + ) # any single out-of-band atom invalidates the WHOLE candidate (.any over # the levels -> masked_fill blanks the entire row to -1, matching no item). - invalid = ((sids < self._sid_lvl_lo) | (sids > self._sid_lvl_hi)).any(dim=1) - sids = sids.masked_fill(invalid.unsqueeze(1), -1) + invalid = ((codes < 1) | (codes > self._codebook_sizes)).any(dim=1) + codes = codes.masked_fill(invalid.unsqueeze(1), -1) # generate() returns rows batch-major ([b0_beam0, b0_beam1, ...]); group # the beams per user. - return sids.view(batch_size, -1, self._num_levels) + return codes.view(batch_size, -1, self._num_levels) def init_input(self) -> None: """Build the EmbeddingGroup for the single raw SID JAGGED_SEQUENCE group. @@ -383,8 +403,10 @@ def _sid_token_rows( as float / shape ``(N, 1)``. The whole batch is tokenized once on the backbone device, then split into rows. - ``expected_width``, when set, enforces the sample contract: every row must - have exactly that many codes (the answer = ``num_levels``). + All values are local 1-based codes. Rows must contain whole items so the + per-level offsets restart at level 0 for each row. ``expected_width``, + when set, additionally requires exactly that many codes (the answer = + ``num_levels``). ``max_codes``, when set, caps each row to its most-recent whole items (the last ``floor(max_codes / num_levels) * num_levels`` codes, dropping the @@ -394,6 +416,13 @@ def _sid_token_rows( if values.dim() == 2 and values.size(-1) == 1: values = values.squeeze(-1) sizes = lengths.long().tolist() + misaligned = [i for i, n in enumerate(sizes) if n % self._num_levels != 0] + if misaligned: + raise ValueError( + f"{type(self).__name__}: SID rows must contain whole " + f"{self._num_levels}-level items; rows {misaligned} have lengths " + f"{[sizes[i] for i in misaligned]}." + ) if expected_width is not None: bad = [i for i, n in enumerate(sizes) if n != expected_width] if bad: @@ -413,8 +442,16 @@ def _sid_token_rows( rows = torch.split(values, sizes) values = torch.cat([r[-keep:] for r in rows]) sizes = [min(n, keep) for n in sizes] - values = self._tokenize_sids(values.to(self.device).long()) - return list(torch.split(values, sizes)) + codes = values.to(self.device).long() + level_ids = torch.arange(codes.numel(), device=self.device) % self._num_levels + invalid = (codes < 1) | (codes > self._codebook_sizes[level_ids]) + if invalid.any().item(): + raise ValueError( + f"{type(self).__name__}: SID codes must be local 1-based values " + f"in [1, codebook[level]]." + ) + tokens = self._tokenize_sids(codes, level_ids) + return list(torch.split(tokens, sizes)) def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: """Family hook: build inputs, run the HF forward, return ``{"loss": ...}``. diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py index e6266fe1e..8e7ef9370 100644 --- a/tzrec/models/generative_rec_lm_test.py +++ b/tzrec/models/generative_rec_lm_test.py @@ -35,17 +35,21 @@ def lengths(self): return self._l -def _stub(num_levels=3, base_vocab=100, device="cpu"): +def _stub(codebook=None, base_vocab=100, device="cpu"): """A Qwen2RecLM with the base data-prep state wired up, but no HF backbone. Exercises the architecture-agnostic base methods (inherited by every family) without downloading a model. """ + codebook = codebook or [2, 3, 4] m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) m._base_vocab = base_vocab - m._num_levels = num_levels + m._num_levels = len(codebook) m.lm = types.SimpleNamespace(device=torch.device(device)) + lo, hi = Qwen2RecLM._sid_level_bands(codebook) + m.register_buffer("_level_offsets", lo - 1, persistent=False) + m.register_buffer("_codebook_sizes", hi - lo + 1, persistent=False) return m @@ -117,17 +121,22 @@ def test_read_common_config_reads_knobs(self) -> None: ignore_index=-100, generated_sids_key="my_sids", param_dtype="bfloat16", - codebook=[4, 4, 4], + codebook=[2, 3, 4], vocab_pad_to_multiple_of=128, max_sequence_length=288, ) - m._read_common_config(common) + sid_atoms = m._read_common_config(common) self.assertEqual(m._history_group, "user_seq") # the single group self.assertEqual(m._input_name, "user_sequence") # its one member self.assertEqual(m._label_name, "label") # from label_fields[0] self.assertEqual(m._generated_sids_key, "my_sids") # configurable self.assertIs(m._param_dtype, torch.bfloat16) # name -> torch dtype self.assertEqual(m._max_seq_length, 288) # model knob used + self.assertEqual(sid_atoms, 9) + self.assertEqual(m._level_offsets.tolist(), [0, 2, 5]) + self.assertEqual(m._codebook_sizes.tolist(), [2, 3, 4]) + self.assertNotIn("_level_offsets", m.state_dict()) + self.assertNotIn("_codebook_sizes", m.state_dict()) # unknown dtype -> a clear error, not a KeyError common.param_dtype = "float64" with self.assertRaisesRegex(ValueError, "param_dtype must be one of"): @@ -193,77 +202,122 @@ def test_device_property(self) -> None: self.assertEqual(_stub(device="cpu").device, torch.device("cpu")) def test_tokenize_sids(self) -> None: - m = _stub(base_vocab=100) # token = sid + base - 1 = sid + 99 - out = m._tokenize_sids(torch.tensor([1, 2, 3])) - self.assertEqual(out.tolist(), [100, 101, 102]) + m = _stub(base_vocab=100) + codes = torch.tensor([[1, 1, 1], [2, 3, 4]]) + level_ids = torch.arange(3) + out = m._tokenize_sids(codes, level_ids) + self.assertEqual(out.tolist(), [[100, 102, 105], [101, 104, 108]]) self.assertEqual(out.dtype, torch.int64) - # shape-agnostic: a 2-D batch maps elementwise - out2 = m._tokenize_sids(torch.tensor([[1, 2], [3, 4]])) - self.assertEqual(out2.tolist(), [[100, 101], [102, 103]]) + self.assertEqual( + m._detokenize_sids(out, level_ids).tolist(), + codes.tolist(), + ) + + def test_sid_token_bands_use_same_level_aware_mapping(self) -> None: + m = _stub(base_vocab=100) + lo, hi = m._sid_token_bands() + self.assertEqual(lo.tolist(), [100, 102, 105]) + self.assertEqual(hi.tolist(), [101, 104, 108]) def test_sid_token_rows_split_and_cast(self) -> None: m = _stub(base_vocab=100) - jt = _FakeJT([1, 2, 3, 4, 5], [3, 2]) + jt = _FakeJT([1, 1, 1, 2, 3, 4, 2, 1, 3], [6, 3]) rows = m._sid_token_rows(jt.values(), jt.lengths()) - self.assertEqual([r.tolist() for r in rows], [[100, 101, 102], [103, 104]]) + self.assertEqual( + [r.tolist() for r in rows], + [[100, 102, 105, 101, 104, 108], [101, 102, 107]], + ) self.assertTrue(all(r.dtype == torch.int64 for r in rows)) def test_sid_token_rows_squeezes_n1(self) -> None: m = _stub(base_vocab=100) - jt = _FakeJT([1, 2, 3], [3], dim2=True) # (N, 1) + jt = _FakeJT([1, 1, 1], [3], dim2=True) # (N, 1) rows = m._sid_token_rows(jt.values(), jt.lengths()) - self.assertEqual([r.tolist() for r in rows], [[100, 101, 102]]) + self.assertEqual([r.tolist() for r in rows], [[100, 102, 105]]) def test_sid_token_rows_width_ok(self) -> None: - m = _stub(base_vocab=100, num_levels=3) - jt = _FakeJT([1, 2, 3, 4, 5, 6], [3, 3]) + m = _stub(base_vocab=100) + jt = _FakeJT([1, 1, 1, 2, 3, 4], [3, 3]) rows = m._sid_token_rows(jt.values(), jt.lengths(), expected_width=3) - self.assertEqual([r.tolist() for r in rows], [[100, 101, 102], [103, 104, 105]]) + self.assertEqual( + [r.tolist() for r in rows], + [[100, 102, 105], [101, 104, 108]], + ) def test_sid_token_rows_width_violation_raises(self) -> None: - m = _stub(base_vocab=100, num_levels=3) + m = _stub(base_vocab=100) with self.assertRaises(ValueError): - # second row has 2 codes, not 3 -> anomalous sample - jt = _FakeJT([1, 2, 3, 4, 5], [3, 2]) + # second row has 6 codes, not 3 -> anomalous answer sample + jt = _FakeJT([1, 1, 1, 1, 1, 1, 1, 1, 1], [3, 6]) m._sid_token_rows(jt.values(), jt.lengths(), expected_width=3) + def test_sid_token_rows_rejects_partial_items(self) -> None: + m = _stub(base_vocab=100) + with self.assertRaisesRegex(ValueError, "whole 3-level items"): + # The total is divisible by 3, but neither row starts a whole item. + jt = _FakeJT([1, 1, 1, 1, 1, 1], [2, 4]) + m._sid_token_rows(jt.values(), jt.lengths()) + + def test_sid_token_rows_rejects_out_of_range_codes(self) -> None: + m = _stub(base_vocab=100) + for values in ( + [0, 1, 1], + [1, 0, 1], + [1, 1, 0], + [-1, 1, 1], + [1, -1, 1], + [1, 1, -1], + [3, 1, 1], + [1, 4, 1], + [1, 1, 5], + [1, 3, 6], # legacy global 1-based representation + ): + with self.subTest(values=values): + with self.assertRaisesRegex(ValueError, "local 1-based"): + jt = _FakeJT(values, [3]) + m._sid_token_rows(jt.values(), jt.lengths()) + def test_build_input_history_group_label_field(self) -> None: # history: EmbeddingGroup output keyed by GROUP name ("{group}.sequence" # / ".sequence_length"). answer: batch.jagged_labels[label_name]. Both - # tokenized (sid -> sid + base - 1); returned dict keyed by FEATURE name. - m = _stub(base_vocab=100, num_levels=3) + # level-offset tokenized; returned dict keyed by FEATURE name. + m = _stub(base_vocab=100) m._input_name, m._label_name = "user_sequence", "label" m._history_group = "user_seq" m._max_seq_length = 0 m._is_inference = False # train: the answer label_field is read too m.embedding_group = lambda b: { - "user_seq.sequence": torch.tensor([1.0, 2.0, 3.0, 4.0]), - "user_seq.sequence_length": torch.tensor([2, 2]), + "user_seq.sequence": torch.tensor( + [1.0, 1.0, 1.0, 2.0, 3.0, 4.0, 2.0, 1.0, 3.0] + ), + "user_seq.sequence_length": torch.tensor([6, 3]), } batch = types.SimpleNamespace( - jagged_labels={"label": _FakeJT([1, 2, 3, 4, 5, 6], [3, 3])} + jagged_labels={"label": _FakeJT([2, 3, 4, 1, 2, 3], [3, 3])} ) rows = m.build_input(batch) self.assertEqual( - [r.tolist() for r in rows["user_sequence"]], [[100, 101], [102, 103]] + [r.tolist() for r in rows["user_sequence"]], + [[100, 102, 105, 101, 104, 108], [101, 102, 107]], ) self.assertEqual( - [r.tolist() for r in rows["label"]], [[100, 101, 102], [103, 104, 105]] + [r.tolist() for r in rows["label"]], + [[101, 104, 108], [100, 103, 107]], ) def test_build_input_skips_label_in_inference(self) -> None: - m = _stub(base_vocab=100, num_levels=3) + m = _stub(base_vocab=100) m._input_name, m._label_name = "user_sequence", "label" m._history_group = "user_seq" m._max_seq_length = 0 m._is_inference = True # inference: history only, no ground-truth label m.embedding_group = lambda b: { - "user_seq.sequence": torch.tensor([1.0, 2.0, 3.0]), + "user_seq.sequence": torch.tensor([1.0, 1.0, 1.0]), "user_seq.sequence_length": torch.tensor([3]), } # jagged_labels intentionally empty — the label is absent at inference rows = m.build_input(types.SimpleNamespace(jagged_labels={})) - self.assertEqual([r.tolist() for r in rows["user_sequence"]], [[100, 101, 102]]) + self.assertEqual([r.tolist() for r in rows["user_sequence"]], [[100, 102, 105]]) self.assertNotIn("label", rows) diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index ffb732d82..4f3e34ebe 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -295,13 +295,14 @@ def _dynamic_beam_search( score-ordered best-first per row, ready for ``_validate_sid_candidates``. See ``escalating_beam_search`` for the schedule. """ + lo_tok, hi_tok = self._sid_token_bands() return escalating_beam_search( self.lm, input_ids, attention_mask, num_beams=self._num_beams, - lo_tok=self._tokenize_sids(self._sid_lvl_lo), - hi_tok=self._tokenize_sids(self._sid_lvl_hi), + lo_tok=lo_tok, + hi_tok=hi_tok, ) def _splice_prompt_ids( diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/qwen2_rec_lm_test.py index efad14a42..ff8cdca5c 100644 --- a/tzrec/models/qwen2_rec_lm_test.py +++ b/tzrec/models/qwen2_rec_lm_test.py @@ -19,21 +19,20 @@ from tzrec.models.qwen2_rec_lm import Qwen2RecLM -def _stub(num_levels=3, base_vocab=100, pad_id=9, device="cpu", per_level=4): +def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): """A Qwen2RecLM with the splice-relevant state wired up, no HF backbone. Template buffers use tiny placeholder ids so the spliced layout is easy to read; real buffers come from ``_build_prompt_tokens`` at init time. - ``per_level`` sets the (uniform) codebook size used to derive the per-level - SID validity bands the inference gate checks — real bands come from - ``_read_common_config``. With ``per_level=4`` / ``num_levels=3`` the bands - are lo=[1,5,9], hi=[4,8,12] (level j -> sid in [j*4+1, (j+1)*4]). + The non-uniform default makes incorrect ``level * uniform_size`` offset + arithmetic visible: sizes=[2,3,4], offsets=[0,2,5]. """ + codebook = codebook or [2, 3, 4] m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) m._ignore_index = -100 - m._num_levels = num_levels + m._num_levels = len(codebook) m._base_vocab = base_vocab m._pad_token_id = pad_id m._dynamic_beam = False # default = HF fixed-width beam path @@ -49,9 +48,9 @@ def _stub(num_levels=3, base_vocab=100, pad_id=9, device="cpu", per_level=4): "tpl_eos": [9], }.items(): m.register_buffer(name, torch.tensor(vals, dtype=torch.long), persistent=False) - lo, hi = Qwen2RecLM._sid_level_bands([per_level] * num_levels) - m.register_buffer("_sid_lvl_lo", lo, persistent=False) - m.register_buffer("_sid_lvl_hi", hi, persistent=False) + lo, hi = Qwen2RecLM._sid_level_bands(codebook) + m.register_buffer("_level_offsets", lo - 1, persistent=False) + m.register_buffer("_codebook_sizes", hi - lo + 1, persistent=False) return m @@ -143,7 +142,7 @@ def test_predict_routes_on_inference_flag(self) -> None: self.assertEqual(Qwen2RecLM.predict(m, object())["branch"], "generate") def test_generate_maps_tokens_to_sids(self) -> None: - m = _stub(base_vocab=100) # sid = token - base + 1 = token - 99 + m = _stub(base_vocab=100) m._input_name = "user_sequence" m._num_beams = m._num_return = 2 @@ -157,26 +156,23 @@ def fake_generate( pad_token_id, ): prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) - # 2 beams x 3 codes, every atom INSIDE its level's band - # (bands lo=[1,5,9] hi=[4,8,12]; token = sid + 99): - # pos0 token in [100,103], pos1 in [104,107], pos2 in [108,111]. - new = torch.tensor([[100, 104, 108], [103, 107, 111]]) + # codebook=[2,3,4], offsets=[0,2,5]: + # local [1,1,1] / [2,3,4] -> the min/max token of each level. + new = torch.tensor([[100, 102, 105], [101, 104, 108]]) return torch.cat([prompt, new], dim=1) m.lm.generate = fake_generate - # build_input (mocked) supplies the tokenized history rows; the batch is - # opaque to it. SIDs [1,2,3] tokenize to [100,101,102] at base_vocab=100. - m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 101, 102])]} + # build_input is mocked, so the batch is opaque to this generation test. + m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} sids = m._generate(_gen_batch())["generated_sids"] self.assertEqual(tuple(sids.shape), (1, 2, 3)) # (B, num_return, num_levels) - self.assertEqual(sids[0].tolist(), [[1, 5, 9], [4, 8, 12]]) + self.assertEqual(sids[0].tolist(), [[1, 1, 1], [2, 3, 4]]) def test_generate_rejects_malformed_candidates(self) -> None: # Layer-A gate: every malformed candidate -> the -1 sentinel, in place. - # bands lo=[1,5,9] hi=[4,8,12]; token = sid + 99. m = _stub(base_vocab=100) m._input_name = "user_sequence" - m._num_beams = m._num_return = 4 + m._num_beams = m._num_return = 6 def fake_generate( input_ids, @@ -190,28 +186,31 @@ def fake_generate( prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) new = torch.tensor( [ - [100, 104, 108], # all in-band -> valid -> [1, 5, 9] + [100, 102, 105], # valid -> local [1, 1, 1] + [101, 104, 108], # valid -> local [2, 3, 4] + [102, 102, 105], # pos0 above level-0 band + [100, 101, 105], # pos1 below level-1 band + [100, 102, 109], # pos2 above level-2 band [100, 104, 9], # pos2 = eos/pad token (sid -90) -> invalid - [ - 108, - 104, - 100, - ], # wrong-level scramble (pos0 = lvl-2 code) -> invalid - [100, 104, 112], # pos2 sid 13 > band hi 12 -> invalid ] ) return torch.cat([prompt, new], dim=1) m.lm.generate = fake_generate - # build_input (mocked) supplies the tokenized history rows; the batch is - # opaque to it. SIDs [1,2,3] tokenize to [100,101,102] at base_vocab=100. - m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 101, 102])]} + m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} sids = m._generate(_gen_batch())["generated_sids"] - self.assertEqual(tuple(sids.shape), (1, 4, 3)) + self.assertEqual(tuple(sids.shape), (1, 6, 3)) # valid candidate kept at its rank; every malformed one -> all -1 (in place) self.assertEqual( sids[0].tolist(), - [[1, 5, 9], [-1, -1, -1], [-1, -1, -1], [-1, -1, -1]], + [ + [1, 1, 1], + [2, 3, 4], + [-1, -1, -1], + [-1, -1, -1], + [-1, -1, -1], + [-1, -1, -1], + ], ) def test_generate_narrow_tail_no_crash(self) -> None: @@ -231,13 +230,11 @@ def fake_generate( pad_token_id, ): prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) - new = torch.tensor([[100, 104], [103, 107]]) # width 2 < num_levels 3 + new = torch.tensor([[100, 102], [101, 104]]) # width 2 < num_levels 3 return torch.cat([prompt, new], dim=1) m.lm.generate = fake_generate - # build_input (mocked) supplies the tokenized history rows; the batch is - # opaque to it. SIDs [1,2,3] tokenize to [100,101,102] at base_vocab=100. - m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 101, 102])]} + m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} sids = m._generate(_gen_batch())["generated_sids"] self.assertEqual(tuple(sids.shape), (1, 2, 3)) # rectangular, no crash # the missing 3rd atom stays -1 -> out of band -> whole candidate -1 @@ -268,26 +265,66 @@ def test_build_prompt_tokens_registers_buffers(self) -> None: self.assertEqual(m.tpl_eos.tolist(), [99]) # eos cached for supervision def test_sid_token_rows_recency_clip(self) -> None: - m = _stub(num_levels=3, base_vocab=100) # token = sid + base - 1 = sid + 99 + m = _stub(base_vocab=100) + values = torch.tensor( + [ + 1, + 1, + 1, + 2, + 3, + 4, + 1, + 2, + 3, + 2, + 1, + 4, + 1, + 3, + 2, + ], + dtype=torch.float, + ) - def _vl(n): # one row of n codes (values 1..n) as flat (values, lengths) - return torch.arange(1, n + 1, dtype=torch.float), torch.tensor([n]) + def _vl(): + return values, torch.tensor([values.numel()]) - # 15 codes (5 items), cap 9 -> keep last 9 (items 3-5 = codes 7..15) - rows = m._sid_token_rows(*_vl(15), max_codes=9) - self.assertEqual(rows[0].tolist(), [c + 99 for c in range(7, 16)]) + # 15 codes (5 items), cap 9 -> keep the most recent three whole items. + expected_tail = [100, 103, 107, 101, 102, 108, 100, 104, 106] + rows = m._sid_token_rows(*_vl(), max_codes=9) + self.assertEqual(rows[0].tolist(), expected_tail) # item-aligned: cap 10 still keeps 9 (3 whole items), never cuts mid-item - rows = m._sid_token_rows(*_vl(15), max_codes=10) - self.assertEqual(rows[0].tolist(), [c + 99 for c in range(7, 16)]) + rows = m._sid_token_rows(*_vl(), max_codes=10) + self.assertEqual(rows[0].tolist(), expected_tail) # within cap -> untouched - rows = m._sid_token_rows(*_vl(6), max_codes=9) - self.assertEqual(rows[0].tolist(), [c + 99 for c in range(1, 7)]) + rows = m._sid_token_rows(values[:6], torch.tensor([6]), max_codes=9) + self.assertEqual(rows[0].tolist(), [100, 102, 105, 101, 104, 108]) # disabled (0/None) -> no clip - rows = m._sid_token_rows(*_vl(15), max_codes=0) - self.assertEqual(rows[0].tolist(), [c + 99 for c in range(1, 16)]) + rows = m._sid_token_rows(*_vl(), max_codes=0) + self.assertEqual( + rows[0].tolist(), + [ + 100, + 102, + 105, + 101, + 104, + 108, + 100, + 103, + 107, + 101, + 102, + 108, + 100, + 104, + 106, + ], + ) def test_compute_max_total_length(self) -> None: - m = _stub(num_levels=3) + m = _stub() # frame = |system|2 + |user_prefix|1 + |user_suffix|1 + |asst_prefix|1 # + |asst_suffix|1 + |eos|1 = 7; + max_history + answer(num_levels) m._max_seq_length = 300 @@ -347,23 +384,24 @@ def fwd(i, lbl, a): self.assertFalse(m._pool_warmed) -def _real_lm_stub(num_levels=3, base_vocab=20, per_level=4, num_beams=2): +def _real_lm_stub(codebook=None, base_vocab=20, num_beams=2): """A Qwen2RecLM carrying a real (tiny, random) Qwen2 backbone. Needed by the dynamic-beam tests, which exercise the actual KV-cached forward / cache-reorder path (the other tests mock ``lm.generate``). The SID - atoms occupy the last ``num_levels * per_level`` token ids, matching the - base-vocab + appended-codebook layout. + atoms occupy the last ``sum(codebook)`` token ids, matching the base-vocab + plus appended-codebook layout. """ from transformers import Qwen2Config, Qwen2ForCausalLM + codebook = codebook or [2, 3, 4] m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) - m._num_levels = num_levels + m._num_levels = len(codebook) m._base_vocab = base_vocab m._num_beams = num_beams cfg = Qwen2Config( - vocab_size=base_vocab + per_level * num_levels, + vocab_size=base_vocab + sum(codebook), hidden_size=32, intermediate_size=64, num_hidden_layers=2, @@ -373,66 +411,76 @@ def _real_lm_stub(num_levels=3, base_vocab=20, per_level=4, num_beams=2): ) torch.manual_seed(0) m.lm = Qwen2ForCausalLM(cfg).eval() - lo, hi = Qwen2RecLM._sid_level_bands([per_level] * num_levels) - m.register_buffer("_sid_lvl_lo", lo, persistent=False) - m.register_buffer("_sid_lvl_hi", hi, persistent=False) + lo, hi = Qwen2RecLM._sid_level_bands(codebook) + m.register_buffer("_level_offsets", lo - 1, persistent=False) + m.register_buffer("_codebook_sizes", hi - lo + 1, persistent=False) return m class Qwen2DynamicBeamTest(unittest.TestCase): def test_width_schedule_and_final_count(self) -> None: # widths double per level, returning num_beams * 2**num_levels candidates. - m = _real_lm_stub(num_levels=3, base_vocab=20, per_level=8, num_beams=2) + m = _real_lm_stub(codebook=[8, 7, 6], base_vocab=20, num_beams=2) ids = torch.tensor([[1, 2, 3, 4]]) new = m._dynamic_beam_search(ids, torch.ones_like(ids)) - # base 2: widths [4, 8, 16] (none capped: per_level 8 is roomy) -> 16 final + # base 2: widths [4, 8, 16], with enough combinations at each level. self.assertEqual(tuple(new.shape), (2 * 2**3, 3)) def test_every_candidate_is_in_band(self) -> None: # band masking guarantees well-formed SIDs: validate -> no -1 sentinels. - m = _real_lm_stub(num_levels=3, base_vocab=20, per_level=4, num_beams=2) + codebook = [2, 3, 4] + m = _real_lm_stub(codebook=codebook, base_vocab=20, num_beams=2) + lo, hi = m._sid_token_bands() + self.assertEqual(lo.tolist(), [20, 22, 25]) + self.assertEqual(hi.tolist(), [21, 24, 28]) ids = torch.tensor([[5, 6, 7]]) new = m._dynamic_beam_search(ids, torch.ones_like(ids)) sids = m._validate_sid_candidates(new, batch_size=1) self.assertEqual(tuple(sids.shape), (1, new.shape[0], 3)) - self.assertFalse(bool((sids < 0).any())) # every candidate well-formed + for level, size in enumerate(codebook): + self.assertTrue(bool((sids[..., level] >= 1).all())) + self.assertTrue(bool((sids[..., level] <= size).all())) def test_left_padding_two_rows(self) -> None: # ragged batch (row 1 left-padded): both rows yield valid, full beam sets. - m = _real_lm_stub(num_levels=2, base_vocab=20, per_level=4, num_beams=2) + codebook = [2, 3] + m = _real_lm_stub(codebook=codebook, base_vocab=20, num_beams=2) ids = torch.tensor([[5, 6, 7, 8], [0, 0, 9, 10]]) am = torch.tensor([[1, 1, 1, 1], [0, 0, 1, 1]]) new = m._dynamic_beam_search(ids, am) - self.assertEqual(tuple(new.shape), (2 * 2 * 2**2, 2)) # B=2, 2*2^2=8 each + # The requested width is 8, but only 2*3=6 distinct pairs exist. + self.assertEqual(tuple(new.shape), (2 * 6, 2)) sids = m._validate_sid_candidates(new, batch_size=2) - self.assertEqual(tuple(sids.shape), (2, 8, 2)) - self.assertFalse(bool((sids < 0).any())) + self.assertEqual(tuple(sids.shape), (2, 6, 2)) + for level, size in enumerate(codebook): + self.assertTrue(bool((sids[..., level] >= 1).all())) + self.assertTrue(bool((sids[..., level] <= size).all())) def test_exhaustive_matches_bruteforce_topk(self) -> None: # When the schedule covers the whole tree (no pruning) the escalating # beam is EXACT: its candidate set must equal all SID combos and be # ordered by true (full-recompute) cumulative log-prob. This validates # cache-stepping, band masking, scoring, and ordering end-to-end. - per, base = 3, 20 - m = _real_lm_stub(num_levels=2, base_vocab=base, per_level=per, num_beams=3) - # widths [min(6,3)=3, min(12,9)=9] -> exhaustive over all 3*3=9 SIDs + codebook, base = [2, 3], 20 + m = _real_lm_stub(codebook=codebook, base_vocab=base, num_beams=3) + # widths [min(6,2)=2, min(12,6)=6] -> exhaustive over all 2*3 SIDs ids = torch.tensor([[5, 6, 7, 8]]) am = torch.ones_like(ids) got = [tuple(r) for r in m._dynamic_beam_search(ids, am).tolist()] - self.assertEqual(len(got), per * per) + self.assertEqual(len(got), codebook[0] * codebook[1]) lm = m.lm - lo0, lo1 = base, base + per # level token bands [20,22] and [23,25] + lo0, lo1 = base, base + codebook[0] ref = {} with torch.no_grad(): logp0 = torch.log_softmax( lm(ids, attention_mask=am).logits[0, -1].float(), -1 ) - for t0 in range(lo0, lo0 + per): + for t0 in range(lo0, lo0 + codebook[0]): s2 = torch.cat([ids, torch.tensor([[t0]])], 1) logp1 = torch.log_softmax( lm(s2, attention_mask=torch.ones_like(s2)).logits[0, -1].float(), -1 ) - for t1 in range(lo1, lo1 + per): + for t1 in range(lo1, lo1 + codebook[1]): ref[(t0, t1)] = (logp0[t0] + logp1[t1]).item() # 1. exhaustive: the returned set is exactly every SID combination self.assertEqual(set(got), set(ref)) @@ -440,6 +488,11 @@ def test_exhaustive_matches_bruteforce_topk(self) -> None: s = [ref[c] for c in got] self.assertTrue(all(s[i] >= s[i + 1] - 1e-4 for i in range(len(s) - 1))) self.assertEqual(got[0], max(ref, key=ref.get)) # top-1 is the global best + decoded = m._validate_sid_candidates(torch.tensor(got), batch_size=1)[0] + self.assertEqual( + {tuple(row) for row in decoded.tolist()}, + {(a, b) for a in range(1, 3) for b in range(1, 4)}, + ) if __name__ == "__main__": diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index 3146bafcd..99919442d 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -15,18 +15,23 @@ package tzrec.protos; // Architecture-agnostic config shared by ALL generative-rec families (the base // reads this for everything except the backbone, which the family owns — see // _backbone_id). Sample contract (consumed by `predict()`): -// * history : list — raw SID indices in [1, sum(codebook)]; the single -// JAGGED_SEQUENCE feature_group (a `sequence_raw_feature` column). -// * answer : list — raw SID indices; a `data_config.label_field` -// (read from batch.jagged_labels), NOT a feature. -// The model maps SID -> token id by integer offset at batch time. Column names -// are NOT configured here (see below). +// * history : list — local 1-based per-level codes in +// [1, codebook[level]], laid out as whole items in level order; +// the single JAGGED_SEQUENCE feature_group. +// * answer : list — one item's local 1-based codes; a +// `data_config.label_field` (read from batch.jagged_labels), +// NOT a feature. +// The model derives level_offsets from codebook and maps each code at batch +// time: token_id = base_vocab + level_offsets[level] + code - 1. Sample +// writers must not pre-apply level_offsets. Column names are not configured +// here (see below). message GenerativeRecLMConfig { // Backbone (`hf_model_id`) is NOT here — it's the family's architecture // commitment, so it lives on the family message. - // SID vocabulary, one entry per RQ level: len = SID codes per item (answer - // width), sum = atoms appended as C0..C{sum-1} after the base vocab. + // SID vocabulary, one entry per RQ level: each public code is in + // [1, codebook[level]]. len = codes per item (answer width); sum = atoms + // appended as C0..C{sum-1} after the base vocab. repeated uint32 codebook = 2; // Pad the post-extension vocab up to a multiple of this value. optional uint32 vocab_pad_to_multiple_of = 3 [default = 128]; From 230dd06a972b65690191358bb09ac8c3b1889b84 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Mon, 27 Jul 2026 11:31:34 +0000 Subject: [PATCH 34/99] [bugfix] genrec LM: fix beam OOM, dropped ignore_index, and export bloat The escalating beam scored candidates with a full-vocab log_softmax followed by a masked_fill, so every level materialized three (beams, vocab) fp32 tensors even though a level's SID codes occupy one contiguous vocab slice. At the default num_beams that is tens of GiB of transients per level. It now normalizes with logsumexp and slices the band, which is numerically identical and shrinks the topk domain by the vocab-to-codebook ratio. The same loop reordered the KV cache after the final level, building the decode's largest cache only to discard it; that reorder is now skipped on the last iteration. _forward_loss filled labels with the configured ignore_index but never passed it to HF's loss_function, which defaults to -100, so any other value either raised on a negative target or silently supervised every padded position. It is passed through now. dcp_to_hf materialized and randomly initialized a throwaway backbone on CPU purely to read its key set, and wrote the tied lm_head as a second full copy of the embedding matrix. The architecture is now built on the meta device and tied keys are dropped after validation, since from_pretrained re-ties them. transformers was imported at module scope under tzrec/models/, and auto_import re-raises any failure there, so a missing or incompatible install broke import tzrec for every model and entry point. The import is deferred to first use, matching how faiss is handled. Also moves the HF export helpers to tzrec/utils/hf_export_util.py and the beam kernel to tzrec/modules/, rejects a history feature_group that is not JAGGED_SEQUENCE, and covers the export path and the CE objective, neither of which had a test. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/main.py | 30 +- tzrec/models/generative_rec_lm.py | 359 ++++++++----------- tzrec/models/generative_rec_lm_test.py | 123 +++++-- tzrec/models/model.py | 2 +- tzrec/models/qwen2_rec_lm.py | 240 +++++-------- tzrec/models/qwen2_rec_lm_test.py | 153 +++++--- tzrec/{models => modules}/escalating_beam.py | 80 ++--- tzrec/optim/lr_scheduler.py | 3 +- tzrec/protos/model.proto | 3 +- tzrec/protos/models/generative_model.proto | 86 ++--- tzrec/protos/optimizer.proto | 4 +- tzrec/utils/checkpoint_util.py | 42 +-- tzrec/utils/export_util.py | 182 ---------- tzrec/utils/hf_export_util.py | 180 ++++++++++ tzrec/utils/hf_export_util_test.py | 175 +++++++++ 15 files changed, 885 insertions(+), 777 deletions(-) rename tzrec/{models => modules}/escalating_beam.py (56%) create mode 100644 tzrec/utils/hf_export_util.py create mode 100644 tzrec/utils/hf_export_util_test.py diff --git a/tzrec/main.py b/tzrec/main.py index 908df7930..70714e81a 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -146,7 +146,6 @@ def _create_model( labels (list): list of label names. sample_weights (list): list of sample weight names. sampler_type (str): negative sampler type - Return: model: a EasyRec Model. """ @@ -739,16 +738,9 @@ def train_and_evaluate( sample_weights=list(data_config.sample_weight_fields), sampler_type=sampler_type, ) - # Cold-start gate (training-only). `_create_model` builds the EMPTY extended - # architecture (GenerativeRecLM: from_config, no weight download). On a fresh - # run (no checkpoint to resume/fine-tune — same `ckpt_path is None` signal the - # DCP-restore branch at L420-433 keys off) we load the pretrained HF weights - # ONCE here, before TrainWrapper/DMP wrapping and before any DCP restore. On - # resume/fine-tune (`ckpt_path` set) we skip it: DCP `load_state_dict` fills - # the weights. eval/export never reach this path — they always DCP-restore an - # empty model. (See design §1.) + # Cold start only; a resumed or fine-tuned run gets its weights from DCP. if ckpt_path is None: - model.init_from_pretrained() # no-op unless the model has a pretrained source + model.init_from_pretrained() model = TrainWrapper( model, device=device, mixed_precision=train_config.mixed_precision ) @@ -1046,17 +1038,19 @@ def export( else: checkpoint_path, _ = ckpt_manager.latest_checkpoint() - # Explicit export-format branch (driven by export_config.export_format). - # HF: a standalone DCP->HF conversion — NO model build, NO DCP restore, NO - # from_pretrained. `dcp_to_hf` reads everything from the self-contained - # checkpoint dir (DCP weights + co-located config/tokenizer) and writes a - # `from_pretrained`-loadable dir. Done before _create_model so we never - # instantiate the model for HF export (design §2). + # HF export converts the checkpoint dir directly -- no model build, no DCP restore. if pipeline_config.export_config.export_format == export_pb2.ExportFormat.HF: - if checkpoint_path is None: + if not checkpoint_path: raise ValueError("HF export: no checkpoint found to convert.") + if not os.path.exists(os.path.join(checkpoint_path, "config.json")): + raise ValueError( + f"HF export: {checkpoint_path} has no co-located HF assets; it " + f"was not written by an HF-backed model." + ) + if assets: + logger.warning(f"HF export ignores asset_files: {assets}.") if is_rank_zero: - from tzrec.utils.export_util import dcp_to_hf + from tzrec.utils.hf_export_util import dcp_to_hf dcp_to_hf(checkpoint_path, export_dir) return diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index d8cbb1930..4ae488bb9 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -3,71 +3,69 @@ # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """Generic generative-recommendation language-model base for TorchEasyRec. -* Per-family subclasses: ``GenerativeRecLM`` is the abstract base; each LLM -family is a concrete subclass implementing the ``_build_prompt_tokens`` and -``predict`` hooks (e.g. ``Qwen2RecLM``). The pipeline config selects the -family by its own oneof entry, whose message-type name resolves directly to -the same-named class via the BaseModel registry. Shared config lives in -``GenerativeRecLMConfig`` (the family message's ``common`` field). -* Streaming sample format: each row carries the history ``user_sequence`` (a -raw-int64 sequence feature) and the ``label`` answer (a ``data_config`` -label_field), both holding local, 1-based per-level codes in -``[1, codebook[level]]``. Rows contain whole items in level order. -* The chat template is tokenised ONCE at ``__init__`` and cached as -non-persistent buffers, so per-batch encoding is integer arithmetic only -(no HF tokenizer in the hot path). -* Raw code -> token id by integer offsets: -``token = base_vocab + level_offset[level] + code - 1``. The model derives -``level_offsets`` from ``codebook``; sample writers must not pre-apply them. -* Left padding with ``eos_token_id``: real content sits at the END of every -row so the suffix slice captures only ``[response + end_markers]``. +``GenerativeRecLM`` owns the architecture-agnostic plumbing; a concrete family +subclass (e.g. ``Qwen2RecLM``) implements ``_build_prompt_tokens`` and +``predict``. Shared config and the sample contract live in +``GenerativeRecLMConfig`` (see ``protos/models/generative_model.proto``). """ -from __future__ import annotations - -import os -from typing import Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import torch import torch.nn.functional as F import torchmetrics -from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer from tzrec.datasets.utils import Batch from tzrec.features.feature import BaseFeature from tzrec.models.model import BaseModel from tzrec.modules.embedding import EmbeddingGroup +from tzrec.protos import model_pb2 from tzrec.protos.model_pb2 import ModelConfig +from tzrec.protos.models import generative_model_pb2 +if TYPE_CHECKING: + from transformers import PreTrainedModel, PreTrainedTokenizerBase -class GenerativeRecLM(BaseModel): - """Abstract base for HF-backed generative-recommendation LMs. - The base owns the architecture-agnostic plumbing: model construction, SID - vocab extension, the shared sample data-prep (``_sid_token_rows`` / - ``_tokenize_sids``), loss, and metrics. Each family subclass implements two - architecture-specific hooks: +def _hf_auto_classes() -> Tuple[Any, Any, Any]: + """Import the HF Auto classes lazily. - _build_prompt_tokens(tokenizer, cfg) — cache the prompt template - predict(batch) — build inputs + HF forward + ``transformers`` is an optional, multi-hundred-MB dependency needed only by + this model family. Importing it at module scope would make it mandatory for + the whole package, because ``load_class.auto_import`` re-raises any import + failure under ``tzrec/models/`` out of ``import tzrec``. - Family proto contract: every family message embeds ``GenerativeRecLMConfig - common = 1`` and supplies a backbone (by default an ``hf_model_id`` field, - overridable via ``_backbone_id``). + Raises: + ImportError: if ``transformers`` is not installed. + """ + try: + from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + except ImportError as e: + raise ImportError( + "transformers is required for generative-recommendation LMs. " + "Install via `pip install transformers==4.51.2`." + ) from e + return AutoConfig, AutoModelForCausalLM, AutoTokenizer - ``Qwen2RecLM`` provides the decoder-only chat implementation, reusable by - Llama/Mistral/Gemma/Phi-style families; GPT-NeoX/RWKV/Mamba/T5 each need - their own. Each family registers directly (oneof message-type name == class - name). + +class GenerativeRecLM(BaseModel): + """Abstract base for HF-backed generative-recommendation LMs. + + Owns model construction, SID vocab extension, sample data-prep, loss and + metrics; subclasses implement ``_build_prompt_tokens`` and ``predict``. The + family's proto message must supply ``common`` (``GenerativeRecLMConfig``) + and an ``hf_model_id`` field, both read directly by this base. """ - # Backbone PARAM dtype options (the fp32 MASTER weights). fp32 avoids - # bf16-ULP underflow of Adam's small (lr=1e-5) updates; bf16 *compute* comes - # from mixed_precision:"BF16" autocast, NOT the param dtype. Selected by - # ``common.param_dtype`` (default "float32") -> ``self._param_dtype``. + # See `common.param_dtype` in the proto for why fp32 is the default. _DTYPE_BY_NAME = { "float32": torch.float32, "bfloat16": torch.bfloat16, @@ -84,30 +82,20 @@ def __init__( ) -> None: super().__init__(model_config, features, labels, sample_weights, **kwargs) cfg = self._model_config - common = cfg.common # GenerativeRecLMConfig — shared by all families - - sid_atoms = self._read_common_config(common) + sid_atoms = self._read_common_config(cfg.common) self.lm = self._build_backbone() tokenizer, base = self._build_extended_tokenizer(sid_atoms) self._hf_tokenizer = tokenizer self._base_vocab = base - self._pad_token_id = self._resolve_pad_token_id(tokenizer) self._build_prompt_tokens(tokenizer, cfg) - self.init_input() - self._smoke_log_once = os.environ.get("TZREC_GENRECLM_DEBUG", "0") == "1" - self._first_predict = True - @staticmethod - def _resolve_pad_token_id(tokenizer: Any) -> int: - """Pad id for the left-padded splice, falling back to eos. - - Fail loudly if a tokenizer has neither (vs an opaque ``int(None)``). - """ + def _resolve_pad_token_id(tokenizer: "PreTrainedTokenizerBase") -> int: + """Pad id for the left-padded splice, falling back to eos.""" pad_id = tokenizer.pad_token_id if pad_id is None: pad_id = tokenizer.eos_token_id @@ -118,10 +106,10 @@ def _resolve_pad_token_id(tokenizer: Any) -> int: ) return int(pad_id) - def _read_common_config(self, common: Any) -> int: + def _read_common_config( + self, common: generative_model_pb2.GenerativeRecLMConfig + ) -> int: """Parse shared proto knobs into attributes; return the SID atom count.""" - # History = the single feature_group (its group_name keys the - # EmbeddingGroup output); answer = the first data_config.label_field. hist = self._history_feature_group() self._history_group: str = hist.group_name self._input_name: str = hist.feature_names[0] @@ -136,25 +124,26 @@ def _read_common_config(self, common: Any) -> int: ) self._param_dtype: torch.dtype = param_dtype self._max_seq_length: int = int(common.max_sequence_length) - codebook = list(common.codebook) - if len(codebook) == 0: + codebook = [int(c) for c in common.codebook] + if not codebook: raise ValueError("GenerativeRecLM: codebook must be non-empty.") + if any(c <= 0 for c in codebook): + raise ValueError( + f"GenerativeRecLM: every codebook size must be positive, got " + f"{codebook}." + ) self._num_levels = len(codebook) - # Level layout is derived from codebook and kept out of state_dict. - lo, hi = self._sid_level_bands(codebook) - self.register_buffer("_level_offsets", lo - 1, persistent=False) - self.register_buffer("_codebook_sizes", hi - lo + 1, persistent=False) - self._vocab_pad_mult = int(common.vocab_pad_to_multiple_of) or 128 - return sum(int(c) for c in codebook) - - def _history_feature_group(self) -> Any: + offsets, sizes = self._sid_level_layout(codebook) + self.register_buffer("_level_offsets", offsets, persistent=False) + self.register_buffer("_codebook_sizes", sizes, persistent=False) + self._vocab_pad_mult = int(common.vocab_pad_to_multiple_of) + return sum(codebook) + + def _history_feature_group(self) -> model_pb2.FeatureGroupConfig: """The history feature_group = the single declared feature_group. - genrec declares exactly one feature_group — the JAGGED_SEQUENCE group - carrying the history SID stream (the answer is a data_config.label_field, - not a group). Its ``group_name`` keys the EmbeddingGroup output and its one - member is the history feature. Fails loudly on a missing/empty group (vs a - silent IndexError/KeyError later). + Must be JAGGED_SEQUENCE: a SEQUENCE group would silently deliver + padded-dense values instead of the flat raw SID stream. """ if not self._feature_groups: raise ValueError( @@ -167,42 +156,47 @@ def _history_feature_group(self) -> Any: f"{type(self).__name__}: history feature_group {g.group_name!r} " f"has no feature_names." ) + if g.group_type != model_pb2.JAGGED_SEQUENCE: + raise ValueError( + f"{type(self).__name__}: history feature_group {g.group_name!r} " + f"must be JAGGED_SEQUENCE, got " + f"{model_pb2.FeatureGroupType.Name(g.group_type)}." + ) return g - def _build_backbone(self) -> Any: - """Build the EMPTY extended architecture in fp32 (master) — no download. + def _build_backbone(self) -> "PreTrainedModel": + """Build the EMPTY extended architecture -- no weight download. - Config reads only define the module shapes the DCP checkpoint expects; - the GB-scale pretrained weights load once at cold start via - ``init_from_pretrained``, then DCP fills them on every restore/eval. + Only the module shapes matter here; the weights arrive from + ``init_from_pretrained`` (cold start) or DCP (restore/eval). """ - hf_model_id = self._backbone_id() + auto_config, auto_causal_lm, _ = _hf_auto_classes() + hf_model_id = self._model_config.hf_model_id if not hf_model_id: - raise ValueError( - f"{type(self).__name__}: empty backbone id (see _backbone_id)." - ) - hf_cfg = AutoConfig.from_pretrained(hf_model_id) - # fp32 MASTER weights: bf16 params underflow Adam's small (lr=1e-5) updates - # and freeze. bf16 compute comes from mixed_precision:"BF16"; ckpt stays fp32. - lm = AutoModelForCausalLM.from_config(hf_cfg, torch_dtype=self._param_dtype) + raise ValueError(f"{type(self).__name__}: empty hf_model_id.") + hf_cfg = auto_config.from_pretrained(hf_model_id) + lm = auto_causal_lm.from_config(hf_cfg, torch_dtype=self._param_dtype) if next(lm.parameters()).dtype != self._param_dtype: lm = lm.to(self._param_dtype) return lm - def _build_extended_tokenizer(self, sid_atoms: int) -> tuple[Any, int]: + def _build_extended_tokenizer( + self, sid_atoms: int + ) -> Tuple["PreTrainedTokenizerBase", int]: """Add the SID atoms ``C0..C{sid_atoms-1}`` and resize ``self.lm``. Returns ``(tokenizer, base)`` where ``base`` is the tokenizer's next free - id BEFORE adding the atoms — use ``len(tokenizer)``, NOT - ``config.vocab_size`` (which counts reserved slots). The atoms append - directly after the existing vocab, so the splice offset is - ``token = base + (sid - 1)``. + id BEFORE adding the atoms -- use ``len(tokenizer)``, NOT + ``config.vocab_size`` (which counts reserved slots). """ - tokenizer = AutoTokenizer.from_pretrained(self._backbone_id(), use_fast=True) + _, _, auto_tokenizer = _hf_auto_classes() + tokenizer = auto_tokenizer.from_pretrained( + self._model_config.hf_model_id, use_fast=True + ) base = len(tokenizer) added = tokenizer.add_tokens([f"C{i}" for i in range(sid_atoms)]) if added != sid_atoms: - # pre-existing Cxxx tokens would silently break the offset arithmetic. + # a pre-existing Cxxx token would shift the atoms off `base`. raise RuntimeError( f"GenerativeRecLM: tokenizer was expected to grow by " f"{sid_atoms} new atoms, only added {added}. " @@ -211,59 +205,54 @@ def _build_extended_tokenizer(self, sid_atoms: int) -> tuple[Any, int]: # stash the resize target so init_from_pretrained re-extends identically. self._target_vocab = base + sid_atoms self.lm.resize_token_embeddings( - self._target_vocab, pad_to_multiple_of=self._vocab_pad_mult + self._target_vocab, pad_to_multiple_of=self._vocab_pad_mult or None ) c0_id = tokenizer.convert_tokens_to_ids("C0") if c0_id != base: raise RuntimeError( - f"GenerativeRecLM: SID atom layout mismatch — expected " + f"GenerativeRecLM: SID atom layout mismatch -- expected " f"C0 at token id {base}, got {c0_id}. " f"Splice arithmetic would produce wrong token ids." ) return tokenizer, base - def _backbone_id(self) -> str: - """Family hook: the HF model id to load for ``self.lm``. - - Defaults to the family message's ``hf_model_id``; override if a family - sources its backbone differently. - """ - return self._model_config.hf_model_id - def init_from_pretrained(self) -> None: """Load the pretrained HF backbone weights into ``self.lm``. - The single ``from_pretrained`` call, run once at COLD START (no checkpoint - to resume); on resume/eval/export the empty ``__init__`` arch is filled by - DCP instead, skipping the download. Re-extends the vocab to ``__init__``'s - target so the shapes match. + Re-extends the vocab to ``__init__``'s target so the shapes match. + + Every rank runs this, so every rank must apply the identical resize or + DDP's parameter-shape check fails. The newly-created SID embedding rows + are drawn from the global RNG and therefore differ per rank; DDP's + ``_sync_module_states`` broadcast from rank 0 is what makes them agree. """ - # fp32 master (not "auto", which keeps the stored bf16); must match - # _build_backbone so cold-start and restore arches agree. - lm = AutoModelForCausalLM.from_pretrained( - self._backbone_id(), torch_dtype=self._param_dtype + _, auto_causal_lm, _ = _hf_auto_classes() + # drop the empty arch first: holding both peaks at 2x model host RAM. + self.lm = None + lm = auto_causal_lm.from_pretrained( + self._model_config.hf_model_id, + torch_dtype=self._param_dtype, + low_cpu_mem_usage=True, ) lm.resize_token_embeddings( - self._target_vocab, pad_to_multiple_of=self._vocab_pad_mult + self._target_vocab, pad_to_multiple_of=self._vocab_pad_mult or None ) self.lm = lm - def hf_backbone(self): - """The HF backbone module. Only use when export.""" + def hf_backbone(self) -> "PreTrainedModel": + """The HF backbone module, for checkpoint/export asset writing.""" return self.lm - def hf_tokenizer(self): - """The extended tokenizer (base vocab + C0..C{sum-1}) to serialize. - - Only used at export (export_util.write_hf_assets). - """ + def hf_tokenizer(self) -> "PreTrainedTokenizerBase": + """The extended tokenizer (base vocab + C0..C{sum-1}) to serialize.""" return self._hf_tokenizer - def _build_prompt_tokens(self, tokenizer, cfg) -> None: + def _build_prompt_tokens( + self, tokenizer: "PreTrainedTokenizerBase", cfg: Any + ) -> None: """Family hook: cache the tokenised prompt template as buffers. - Called from ``__init__`` after vocab extension; consumed by the family's - ``predict``. Subclasses MUST implement this. + Called from ``__init__`` after vocab extension; consumed by ``predict``. """ raise NotImplementedError( f"{type(self).__name__} must implement _build_prompt_tokens " @@ -272,7 +261,7 @@ def _build_prompt_tokens(self, tokenizer, cfg) -> None: @property def device(self) -> torch.device: - """Device the HF backbone runs on — the single source for model I/O.""" + """Device the HF backbone runs on.""" return self.lm.device def _tokenize_sids( @@ -281,9 +270,7 @@ def _tokenize_sids( """Map local 1-based per-level codes to extended-vocab token ids. ``level_ids`` must broadcast against ``codes`` and identify each code's - RQ level. Atom ``C{k}`` sits at ``base_vocab + k``, therefore: - - ``token = code + level_offsets[level] + base_vocab - 1``. + RQ level. """ return codes + self._level_offsets[level_ids] + (self._base_vocab - 1) @@ -294,24 +281,12 @@ def _detokenize_sids( return tokens - (self._base_vocab - 1) - self._level_offsets[level_ids] @staticmethod - def _sid_level_bands(codebook: Any) -> tuple[torch.Tensor, torch.Tensor]: - """Per-level flattened 1-based SID bands as two long tensors. + def _sid_level_layout(codebook: List[int]) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-level ``(flat_offset, size)``; the levels occupy disjoint bands.""" + sizes = torch.tensor(codebook, dtype=torch.long) + return torch.cumsum(sizes, 0) - sizes, sizes - Level ``j`` occupies a DISJOINT band ``[offset_j + 1, offset_j + - codebook[j]]`` where ``offset_j = sum(codebook[:j])``. Single source of - truth for the derived ``_level_offsets`` and ``_codebook_sizes`` buffers. - """ - lo, hi, acc = [], [], 0 - for c in codebook: - lo.append(acc + 1) - hi.append(acc + int(c)) - acc += int(c) - return ( - torch.tensor(lo, dtype=torch.long), - torch.tensor(hi, dtype=torch.long), - ) - - def _sid_token_bands(self) -> tuple[torch.Tensor, torch.Tensor]: + def _sid_token_bands(self) -> Tuple[torch.Tensor, torch.Tensor]: """Return the inclusive token-id band for every SID level.""" level_ids = torch.arange(self._num_levels, device=self.device) return ( @@ -324,51 +299,34 @@ def _validate_sid_candidates( ) -> torch.Tensor: """Decode generated tokens to local 1-based codes and reject bad beams. - ``new_tokens`` is the per-beam tail ``(B*num_return, w)`` (``w`` may be < - ``num_levels`` when beams stop early). Returns ``(batch_size, num_return, + ``new_tokens`` is the per-beam tail ``(B*C, w)`` (``w`` may be < + ``num_levels`` when beams stop early). Returns ``(batch_size, C, num_levels)`` local codes. Every malformed candidate (early EOS / non-SID / wrong-level atom) is set to ``-1``, which cannot match a real 1-based code. """ level_ids = torch.arange(new_tokens.shape[1], device=new_tokens.device) codes = self._detokenize_sids(new_tokens, level_ids) - # early-EOS beams return < num_levels tokens; pad to num_levels with -1. - codes = F.pad( - codes, - (0, self._num_levels - codes.shape[1]), - value=-1, - ) - # any single out-of-band atom invalidates the WHOLE candidate (.any over - # the levels -> masked_fill blanks the entire row to -1, matching no item). + codes = F.pad(codes, (0, self._num_levels - codes.shape[1]), value=-1) + # one out-of-band atom invalidates the whole candidate row. invalid = ((codes < 1) | (codes > self._codebook_sizes)).any(dim=1) codes = codes.masked_fill(invalid.unsqueeze(1), -1) - # generate() returns rows batch-major ([b0_beam0, b0_beam1, ...]); group - # the beams per user. + # decoders return rows batch-major ([b0_c0, b0_c1, ...]); group per user. return codes.view(batch_size, -1, self._num_levels) def init_input(self) -> None: """Build the EmbeddingGroup for the single raw SID JAGGED_SEQUENCE group. - Raw (passthrough) features carry no embedding tables, so this - EmbeddingGroup holds no params (DMP-neutral); it exists purely to - retrieve the raw SID sequences as flat ``(values, lengths)`` — the same - path HSTU uses. The HF backbone still owns the token embeddings; SID ids - flow through it directly (no embedding lookup here). + Raw passthrough features carry no embedding tables, so this group holds + no params (DMP-neutral); it only retrieves the flat ``(values, lengths)``. """ self.embedding_group = EmbeddingGroup(self._features, self._feature_groups) def build_input(self, batch: Batch) -> Dict[str, List[torch.Tensor]]: - """Retrieve per-row SID token sequences. - - HISTORY is a JAGGED_SEQUENCE feature_group: ``embedding_group(batch)`` - returns its flat raw values ``"{group}.sequence"`` + - ``"{group}.sequence_length"`` (the HSTU idiom). The ANSWER is a - data_config.label_field: its JaggedTensor comes from - ``batch.jagged_labels[self._label_name]`` (so it can be absent at - inference, where no ground truth is supplied — unlike a feature_group, - which the EmbeddingGroup would require every forward). Both are mapped - SID -> extended-vocab token ids and split to rows; the returned dict is - keyed by FEATURE name (what the family ``predict`` consumes). + """Retrieve per-row SID token sequences, keyed by feature name. + + The answer is a ``data_config.label_field`` rather than a feature_group + so it can be absent at inference, where no ground truth is supplied. """ g = self.embedding_group(batch) rows: Dict[str, List[torch.Tensor]] = { @@ -379,11 +337,14 @@ def build_input(self, batch: Batch) -> Dict[str, List[torch.Tensor]]: ), } if not self.is_inference: + if not self._label_name: + raise ValueError( + f"{type(self).__name__}: training needs the answer SIDs; " + f"declare it as the first data_config.label_field." + ) jt = batch.jagged_labels[self._label_name] rows[self._label_name] = self._sid_token_rows( - jt.values(), - jt.lengths(), - expected_width=self._num_levels, + jt.values(), jt.lengths(), expected_width=self._num_levels ) return rows @@ -396,22 +357,13 @@ def _sid_token_rows( ) -> List[torch.Tensor]: """Map flat SID ``(values, lengths)`` -> per-row token-id tensors. - ``build_input`` supplies flat SID ``(values, lengths)``: the history from - the JAGGED_SEQUENCE group's ``"{group}.sequence"`` / - ``"{group}.sequence_length"``, and the answer from - ``batch.jagged_labels[label].values()/.lengths()``. ``values`` may arrive - as float / shape ``(N, 1)``. The whole batch is tokenized once on the - backbone device, then split into rows. - - All values are local 1-based codes. Rows must contain whole items so the - per-level offsets restart at level 0 for each row. ``expected_width``, - when set, additionally requires exactly that many codes (the answer = - ``num_levels``). - - ``max_codes``, when set, caps each row to its most-recent whole items (the - last ``floor(max_codes / num_levels) * num_levels`` codes, dropping the - oldest head) so the pre-allocated pool covers every batch. Done on host - views before the H2D copy, skipped unless a row overflows. + ``values`` may arrive as float / shape ``(N, 1)``. Rows must contain + whole items so the per-level offsets restart at level 0 for each row; + ``expected_width``, when set, requires exactly that many codes per row. + + ``max_codes``, when set, caps each row to its most-recent whole items + (the last ``floor(max_codes / num_levels) * num_levels`` codes) so the + pre-allocated pool covers every batch. """ if values.dim() == 2 and values.size(-1) == 1: values = values.squeeze(-1) @@ -429,13 +381,10 @@ def _sid_token_rows( raise ValueError( f"{type(self).__name__}: each SID item must be " f"{expected_width} codes (len(codebook)); rows {bad} have " - f"{[sizes[i] for i in bad]} — anomalous sample(s)." + f"{[sizes[i] for i in bad]} -- anomalous sample(s)." ) - # TODO: The truncation logic should not be placed here, but should be - # handled in FG. Since FG currently cannot control the truncation - # direction (it keeps the HEAD), this may result in truncating the most - # recent SIDs. Check it. + # TODO(shuqi): move truncation into FG once FG can keep the TAIL, not the HEAD. if max_codes: keep = (max_codes // self._num_levels) * self._num_levels if keep and any(n > keep for n in sizes): @@ -453,18 +402,6 @@ def _sid_token_rows( tokens = self._tokenize_sids(codes, level_ids) return list(torch.split(tokens, sizes)) - def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: - """Family hook: build inputs, run the HF forward, return ``{"loss": ...}``. - - Architecture-specific — the decoder-only implementation lives in - ``Qwen2RecLM``; other families (GPT-NeoX, Mamba, T5, …) need their own. - See design §15.4/§16. Subclasses MUST implement this. - """ - raise NotImplementedError( - f"{type(self).__name__} must implement predict " - f"(GenerativeRecLM is abstract)." - ) - def init_loss(self) -> None: """No-op: the loss is computed inside ``predict`` (HF loss_function).""" return @@ -477,8 +414,6 @@ def loss( """Surface the CE loss already computed in ``predict``.""" return {"ce_loss": predictions["loss"]} - # BaseModel declares only the eval-side metric hooks, but the train loop - # calls both eval and train hooks, so both are overridden here. def init_metric(self) -> None: """Register a mean-CE metric for the eval loop.""" self._metric_modules["ce_loss"] = torchmetrics.MeanMetric() @@ -492,15 +427,11 @@ def update_metric( """Update the mean-CE metric with this batch's loss.""" self._metric_modules["ce_loss"].update(predictions["loss"].detach()) - def init_train_metric(self) -> None: - """No-op: no train-time metric beyond the logged CE loss.""" - return - + # NOTE: BaseModel declares no such hook, but the train loop calls it. def update_train_metric( self, predictions: Dict[str, torch.Tensor], batch: Batch, - losses: Optional[Dict[str, torch.Tensor]] = None, ) -> None: """No-op: no train-time metric beyond the logged CE loss.""" return diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py index 8e7ef9370..eb0f1ddc1 100644 --- a/tzrec/models/generative_rec_lm_test.py +++ b/tzrec/models/generative_rec_lm_test.py @@ -18,10 +18,11 @@ from tzrec.models.generative_rec_lm import GenerativeRecLM from tzrec.models.model import BaseModel from tzrec.models.qwen2_rec_lm import Qwen2RecLM +from tzrec.protos import model_pb2 class _FakeJT: - """Minimal stand-in for a TorchRec JaggedTensor (callable values/lengths).""" + """Minimal stand-in for a TorchRec JaggedTensor.""" def __init__(self, values, lengths, dim2=False): v = torch.tensor(values, dtype=torch.float) # TER delivers list as float @@ -36,37 +37,42 @@ def lengths(self): def _stub(codebook=None, base_vocab=100, device="cpu"): - """A Qwen2RecLM with the base data-prep state wired up, but no HF backbone. - - Exercises the architecture-agnostic base methods (inherited by every family) - without downloading a model. - """ + """A Qwen2RecLM with the base data-prep state wired up, but no HF backbone.""" codebook = codebook or [2, 3, 4] m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) m._base_vocab = base_vocab m._num_levels = len(codebook) m.lm = types.SimpleNamespace(device=torch.device(device)) - lo, hi = Qwen2RecLM._sid_level_bands(codebook) - m.register_buffer("_level_offsets", lo - 1, persistent=False) - m.register_buffer("_codebook_sizes", hi - lo + 1, persistent=False) + offsets, sizes = Qwen2RecLM._sid_level_layout(codebook) + m.register_buffer("_level_offsets", offsets, persistent=False) + m.register_buffer("_codebook_sizes", sizes, persistent=False) return m class GenerativeRecLMTest(unittest.TestCase): def test_registry_dispatch(self) -> None: - # importing qwen2_rec_lm auto-registers the family by class name self.assertIs(BaseModel.create_class("Qwen2RecLM"), Qwen2RecLM) self.assertTrue(issubclass(Qwen2RecLM, GenerativeRecLM)) + def test_model_config_oneof_resolves_to_the_class(self) -> None: + # the path _create_model takes: oneof -> message type name -> class. + from tzrec.utils import config_util + + cfg = model_pb2.ModelConfig() + cfg.qwen2_rec_lm.common.codebook.extend([2, 3]) + cfg.qwen2_rec_lm.common.max_sequence_length = 8 # required field + self.assertEqual(config_util.which_msg(cfg, "model"), "Qwen2RecLM") + self.assertIs( + BaseModel.create_class(config_util.which_msg(cfg, "model")), Qwen2RecLM + ) + def test_resolve_pad_token_id(self) -> None: tok = types.SimpleNamespace - # pad present -> pad self.assertEqual( GenerativeRecLM._resolve_pad_token_id(tok(pad_token_id=5, eos_token_id=9)), 5, ) - # pad absent -> eos fallback self.assertEqual( GenerativeRecLM._resolve_pad_token_id( tok(pad_token_id=None, eos_token_id=9) @@ -80,8 +86,6 @@ def test_resolve_pad_token_id(self) -> None: ) def test_backbone_owned_by_family_proto(self) -> None: - # the backbone lives on the family message (its architecture), NOT in - # the shared common config; it defaults to the canonical Qwen2.5-0.5B. from tzrec.protos.models.generative_model_pb2 import ( GenerativeRecLMConfig, ) @@ -94,8 +98,6 @@ def test_backbone_owned_by_family_proto(self) -> None: self.assertNotIn("hf_model_id", common_fields) def test_configurable_knob_defaults(self) -> None: - # generated_sids_key / param_dtype are proto knobs whose defaults are the - # previous class-constant values. from tzrec.protos.models.generative_model_pb2 import GenerativeRecLMConfig c = GenerativeRecLMConfig() @@ -108,13 +110,12 @@ def test_read_common_config_reads_knobs(self) -> None: m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) m._features = [] - # The history group + SID-column names are DERIVED, not configured: the - # history is the single feature_group + its one member; the answer is the - # first data_config.label_field. m._labels = ["label"] m._feature_groups = [ types.SimpleNamespace( - group_name="user_seq", feature_names=["user_sequence"] + group_name="user_seq", + feature_names=["user_sequence"], + group_type=model_pb2.JAGGED_SEQUENCE, ) ] common = types.SimpleNamespace( @@ -129,9 +130,9 @@ def test_read_common_config_reads_knobs(self) -> None: self.assertEqual(m._history_group, "user_seq") # the single group self.assertEqual(m._input_name, "user_sequence") # its one member self.assertEqual(m._label_name, "label") # from label_fields[0] - self.assertEqual(m._generated_sids_key, "my_sids") # configurable - self.assertIs(m._param_dtype, torch.bfloat16) # name -> torch dtype - self.assertEqual(m._max_seq_length, 288) # model knob used + self.assertEqual(m._generated_sids_key, "my_sids") + self.assertIs(m._param_dtype, torch.bfloat16) + self.assertEqual(m._max_seq_length, 288) self.assertEqual(sid_atoms, 9) self.assertEqual(m._level_offsets.tolist(), [0, 2, 5]) self.assertEqual(m._codebook_sizes.tolist(), [2, 3, 4]) @@ -143,12 +144,11 @@ def test_read_common_config_reads_knobs(self) -> None: m._read_common_config(common) def test_read_common_config_no_feature_group_raises(self) -> None: - # the history is the single declared feature_group; none -> fail loudly. m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) m._features = [] m._labels = ["label"] - m._feature_groups = [] # no group declared + m._feature_groups = [] common = types.SimpleNamespace( ignore_index=-100, generated_sids_key="generated_sids", @@ -160,8 +160,66 @@ def test_read_common_config_no_feature_group_raises(self) -> None: with self.assertRaisesRegex(ValueError, "no feature_group declared"): m._read_common_config(common) + def test_read_common_config_rejects_bad_group_and_codebook(self) -> None: + def _wired(group_type, feature_names=("user_sequence",)): + m = object.__new__(Qwen2RecLM) + nn.Module.__init__(m) + m._features = [] + m._labels = ["label"] + m._feature_groups = [ + types.SimpleNamespace( + group_name="user_seq", + feature_names=list(feature_names), + group_type=group_type, + ) + ] + return m + + def _common(codebook): + return types.SimpleNamespace( + ignore_index=-100, + generated_sids_key="generated_sids", + param_dtype="float32", + codebook=codebook, + vocab_pad_to_multiple_of=128, + max_sequence_length=0, + ) + + # a SEQUENCE group emits the same key with padded-dense semantics. + with self.assertRaisesRegex(ValueError, "must be JAGGED_SEQUENCE"): + _wired(model_pb2.SEQUENCE)._read_common_config(_common([2, 3])) + with self.assertRaisesRegex(ValueError, "has no feature_names"): + _wired(model_pb2.JAGGED_SEQUENCE, ())._read_common_config(_common([2, 3])) + with self.assertRaisesRegex(ValueError, "codebook must be non-empty"): + _wired(model_pb2.JAGGED_SEQUENCE)._read_common_config(_common([])) + with self.assertRaisesRegex(ValueError, "codebook size must be positive"): + _wired(model_pb2.JAGGED_SEQUENCE)._read_common_config(_common([2, 0])) + + def test_vocab_pad_zero_disables_padding(self) -> None: + m = object.__new__(Qwen2RecLM) + nn.Module.__init__(m) + m._features = [] + m._labels = ["label"] + m._feature_groups = [ + types.SimpleNamespace( + group_name="user_seq", + feature_names=["user_sequence"], + group_type=model_pb2.JAGGED_SEQUENCE, + ) + ] + m._read_common_config( + types.SimpleNamespace( + ignore_index=-100, + generated_sids_key="generated_sids", + param_dtype="float32", + codebook=[2, 3], + vocab_pad_to_multiple_of=0, + max_sequence_length=0, + ) + ) + self.assertEqual(m._vocab_pad_mult, 0) # not silently rewritten to 128 + def test_max_sequence_length_model_knob(self) -> None: - # _max_seq_length is the model knob; 0 = off (no feature fallback). def _common(max_seq): return types.SimpleNamespace( ignore_index=-100, @@ -179,14 +237,16 @@ def _wired(): m._labels = ["label"] m._feature_groups = [ types.SimpleNamespace( - group_name="user_seq", feature_names=["user_sequence"] + group_name="user_seq", + feature_names=["user_sequence"], + group_type=model_pb2.JAGGED_SEQUENCE, ) ] return m m = _wired() m._read_common_config(_common(128)) - self.assertEqual(m._max_seq_length, 128) # model knob used + self.assertEqual(m._max_seq_length, 128) m2 = _wired() m2._read_common_config(_common(0)) self.assertEqual(m2._max_seq_length, 0) # 0 = off, no fallback @@ -247,7 +307,6 @@ def test_sid_token_rows_width_ok(self) -> None: def test_sid_token_rows_width_violation_raises(self) -> None: m = _stub(base_vocab=100) with self.assertRaises(ValueError): - # second row has 6 codes, not 3 -> anomalous answer sample jt = _FakeJT([1, 1, 1, 1, 1, 1, 1, 1, 1], [3, 6]) m._sid_token_rows(jt.values(), jt.lengths(), expected_width=3) @@ -270,7 +329,7 @@ def test_sid_token_rows_rejects_out_of_range_codes(self) -> None: [3, 1, 1], [1, 4, 1], [1, 1, 5], - [1, 3, 6], # legacy global 1-based representation + [1, 3, 6], # global cross-level codes, not local per-level ): with self.subTest(values=values): with self.assertRaisesRegex(ValueError, "local 1-based"): @@ -278,9 +337,6 @@ def test_sid_token_rows_rejects_out_of_range_codes(self) -> None: m._sid_token_rows(jt.values(), jt.lengths()) def test_build_input_history_group_label_field(self) -> None: - # history: EmbeddingGroup output keyed by GROUP name ("{group}.sequence" - # / ".sequence_length"). answer: batch.jagged_labels[label_name]. Both - # level-offset tokenized; returned dict keyed by FEATURE name. m = _stub(base_vocab=100) m._input_name, m._label_name = "user_sequence", "label" m._history_group = "user_seq" @@ -315,7 +371,6 @@ def test_build_input_skips_label_in_inference(self) -> None: "user_seq.sequence": torch.tensor([1.0, 1.0, 1.0]), "user_seq.sequence_length": torch.tensor([3]), } - # jagged_labels intentionally empty — the label is absent at inference rows = m.build_input(types.SimpleNamespace(jagged_labels={})) self.assertEqual([r.tolist() for r in rows["user_sequence"]], [[100, 102, 105]]) self.assertNotIn("label", rows) diff --git a/tzrec/models/model.py b/tzrec/models/model.py index 8a7ae1d35..a18c47a41 100644 --- a/tzrec/models/model.py +++ b/tzrec/models/model.py @@ -420,7 +420,7 @@ class CudaAutocastWrapper(nn.Module): proper dtype casts. CUTLASS HSTU attention requires bf16/fp16 inputs. When ``device`` is set, it is passed as a second positional argument - to ``inner.forward(x, device)`` — this binds the device for models + to ``inner.forward(x, device)`` -- this binds the device for models like ``ScriptWrapper`` whose forward takes ``(data, device)``. ``_mixed_dtype_id: Final[int]`` encodes the dtype so that diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index 4f3e34ebe..dc8ab1fdd 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -3,51 +3,34 @@ # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """Qwen2/Qwen2.5 family subclass of ``GenerativeRecLM``. -Selected from the pipeline config by its own oneof entry (the message-type name -resolves directly to this class):: - - model_config { - qwen2_rec_lm { - common { hf_model_id: "..." codebook: 8192 ... } - system_instruction: "..." - } - } - -This subclass owns the decoder-only-chat implementation: the ChatML prompt -template, the causal-LM splice, and the ``.model``/``.lm_head`` forward. The -``GenerativeRecLM`` base owns the architecture-agnostic plumbing (vocab -extension, jagged->row, loss, metrics). - -The splice/forward are generic to decoder-only families sharing Qwen2's -``.model``/``.lm_head`` layout (Llama/Mistral/Gemma/Phi); only ``QWEN2_TEMPLATE`` -is Qwen2-specific. +Owns the decoder-only-chat implementation: the ChatML prompt template, the +causal-LM splice, and the ``.model``/``.lm_head`` forward. """ -from typing import Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import torch from torch.nn.utils.rnn import pad_sequence from tzrec.datasets.utils import Batch from tzrec.features.feature import BaseFeature -from tzrec.models.escalating_beam import escalating_beam_search from tzrec.models.generative_rec_lm import GenerativeRecLM +from tzrec.modules.escalating_beam import escalating_beam_search from tzrec.protos.model_pb2 import ModelConfig +from tzrec.protos.models import generative_model_pb2 +if TYPE_CHECKING: + from transformers import PreTrainedTokenizerBase -def _encode_no_special(tokenizer, text: str) -> List[int]: - """Encode a fragment without prepending BOS / appending EOS specials. - We're building the prompt manually from explicit ``<|im_start|>`` markers, - so we must NOT let the tokenizer's BOS/EOS handling double-emit them. - """ - return tokenizer.encode(text, add_special_tokens=False) - - -# Verbatim Qwen2 ChatML fragments. QWEN2_TEMPLATE = { "system_prefix": "<|im_start|>system\n", "system_suffix": "<|im_end|>\n", @@ -79,17 +62,21 @@ def __init__( self._num_beams = int(common.num_beams) self._num_return = int(common.num_return_sequences) self._dynamic_beam = bool(common.dynamic_beam) + if not self._dynamic_beam and self._num_return > self._num_beams: + raise ValueError( + f"{type(self).__name__}: num_return_sequences " + f"({self._num_return}) must not exceed num_beams " + f"({self._num_beams})." + ) self._max_total_len = self._compute_max_total_length() self._pool_warmed = False - # CE suffix width: the supervised tail [answer | asst_suffix | eos] is - # fixed, so _forward_loss slices a constant suffix (no per-step sync). + # +2 = trailing eos + HF's shift-by-one; constant width avoids a per-step sync. self._suffix_keep = self._num_levels + self.tpl_asst_suffix.numel() + 2 def _compute_max_total_length(self) -> int: """Full spliced length at the max history (0 if pre-allocation is off). - Fixed ChatML frame + ``self._max_seq_length`` history codes + the - ``num_levels``-code answer: the ``T`` the activation pool is pre-sized to. + The ``T`` the activation pool is pre-sized to. """ if self._max_seq_length <= 0: return 0 @@ -103,35 +90,32 @@ def _compute_max_total_length(self) -> int: ) return int(frame + self._max_seq_length + self._num_levels) - def _build_prompt_tokens(self, tokenizer, cfg) -> None: + def _build_prompt_tokens( + self, + tokenizer: "PreTrainedTokenizerBase", + cfg: generative_model_pb2.Qwen2RecLM, + ) -> None: """Tokenise the family chat template once; cache as buffers. Composes the proto's optional ``system_instruction`` / ``user_prefix_text`` / ``user_suffix_text`` with the family's static - fragments:: - - tpl_system = system_prefix + system_instruction + system_suffix - tpl_user_prefix = user_prefix + user_prefix_text - tpl_user_suffix = user_suffix_text + user_suffix - tpl_asst_prefix / tpl_asst_suffix verbatim from the template - - Buffers are non-persistent: they move with ``model.to(...)`` but stay off - the state_dict so HF safetensors round-tripping isn't polluted. + fragments. Buffers are non-persistent: they move with ``model.to(...)`` + but stay off the state_dict so HF safetensors round-tripping isn't + polluted. """ tpl = type(self).CHAT_TEMPLATE sys_text = cfg.system_instruction or tpl["default_system_instruction"] - u_pre = cfg.user_prefix_text or "" - u_suf = cfg.user_suffix_text or "" frags = { "system": tpl["system_prefix"] + sys_text + tpl["system_suffix"], - "user_prefix": tpl["user_prefix"] + u_pre, - "user_suffix": u_suf + tpl["user_suffix"], + "user_prefix": tpl["user_prefix"] + (cfg.user_prefix_text or ""), + "user_suffix": (cfg.user_suffix_text or "") + tpl["user_suffix"], "asst_prefix": tpl["asst_prefix"], "asst_suffix": tpl["asst_suffix"], } for slot_name, frag_str in frags.items(): + # explicit <|im_start|> markers frame the prompt; no auto BOS/EOS. ids = torch.tensor( - _encode_no_special(tokenizer, frag_str), dtype=torch.long + tokenizer.encode(frag_str, add_special_tokens=False), dtype=torch.long ) self.register_buffer(f"tpl_{slot_name}", ids, persistent=False) # the trailing eos is a SUPERVISED token; cache it for the splice. @@ -141,6 +125,24 @@ def _build_prompt_tokens(self, tokenizer, cfg) -> None: persistent=False, ) + def _prompt_rows(self, user_seq_rows: List[torch.Tensor]) -> List[torch.Tensor]: + """Per-row ``[system | user_prefix | history | user_suffix | asst_prefix]``. + + Shared by the teacher-forced splice and the answer-less inference prompt. + """ + return [ + torch.cat( + [ + self.tpl_system, + self.tpl_user_prefix, + row, + self.tpl_user_suffix, + self.tpl_asst_prefix, + ] + ) + for row in user_seq_rows + ] + def _splice_input_ids( self, user_seq_rows: List[torch.Tensor], @@ -149,90 +151,56 @@ def _splice_input_ids( ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Build ``(input_ids, labels, attention_mask)``, each ``(B, T_max)``. - Left-padded with ``eos_token_id``. ``attention_mask`` is essential — - without it self-attention lets pad positions pollute real positions; CE - is separately protected by ``-100`` labels at pad slots. - - Every answer is exactly ``self._num_levels`` SID codes, so the supervised - tail ``[answer | asst_suffix | eos]`` has a FIXED width and lands in the - same columns for every row after left-padding -> ``labels`` is one - vectorized write. ``input_ids`` still varies per row (history length). + Rows already hold on-device token ids (see ``_sid_token_rows``). - ``user_seq_rows`` / ``label_rows`` already hold token ids on the model - device (see ``_sid_token_rows``). ``pad_to`` left-extends every row for - first-step pool pre-sizing; the supervised tail stays end-aligned. + Every answer is exactly ``self._num_levels`` SID codes, so the tail + ``[answer | asst_suffix | eos]`` has a FIXED width and lands in the same + columns for every row after left-padding -> ``labels`` is one vectorized + write. Only the answer and the trailing eos are supervised: a decode + emits exactly ``num_levels`` tokens and never has to produce + ``asst_suffix``. ``pad_to`` left-extends every row for pool pre-sizing, + keeping the supervised tail end-aligned. """ - assert len(user_seq_rows) == len(label_rows) - A = self._num_levels - - rows_ids = [ - torch.cat( - [ - self.tpl_system, - self.tpl_user_prefix, - user_seq_rows[i], - self.tpl_user_suffix, - self.tpl_asst_prefix, - label_rows[i], - self.tpl_asst_suffix, - self.tpl_eos, - ] + if len(user_seq_rows) != len(label_rows): + raise ValueError( + f"{type(self).__name__}: history/answer row count mismatch " + f"({len(user_seq_rows)} vs {len(label_rows)})." ) - for i in range(len(user_seq_rows)) + rows_ids = [ + torch.cat([prompt, label_rows[i], self.tpl_asst_suffix, self.tpl_eos]) + for i, prompt in enumerate(self._prompt_rows(user_seq_rows)) ] input_ids, attention_mask = self._left_pad(rows_ids, pad_to=pad_to) B, T = input_ids.shape - s = self.tpl_asst_suffix.numel() - tail = A + s + 1 # [answer(A) | asst_suffix(s) | eos(1)], end-aligned + answer_width = self._num_levels + tail = answer_width + self.tpl_asst_suffix.numel() + 1 labels = torch.full( (B, T), self._ignore_index, dtype=torch.long, device=self.device ) - labels[:, T - tail : T - tail + A] = torch.stack(label_rows) - labels[:, -1] = self.tpl_eos[0] # supervise the trailing eos + labels[:, T - tail : T - tail + answer_width] = torch.stack(label_rows) + labels[:, -1] = self.tpl_eos[0] return input_ids, labels, attention_mask def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: - """Dispatch on the TER inference flag (``set_is_inference`` in main.py). - - Branch 1 (train / eval, ``not is_inference``) — teacher-forced forward + - CE loss (the metric path). - Branch 2 (inference, ``is_inference``) — beam-search the SID answer from - the prompt. - """ + """Dispatch on the TER inference flag (``set_is_inference`` in main.py).""" if self.is_inference: return self._generate(batch) return self._predict_train(batch) def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: - """Branch 1: teacher-forced forward -> suffix-slice -> CE loss.""" + """Build the teacher-forced splice for a batch and return the CE loss.""" rows = self.build_input(batch) - u_rows = rows[self._input_name] - l_rows = rows[self._label_name] - # One-shot pool pre-sizing: pad the FIRST train step to the worst-case - # length so the allocator reserves its largest segments up front (no - # mid-run growth). Extra positions are masked + -100 -> loss/grad unchanged. + # Pre-size the caching allocator on step 1; the extra columns are masked. pad_to = 0 if not self._pool_warmed and self._max_total_len > 0 and self.is_train: pad_to = self._max_total_len self._pool_warmed = True input_ids, labels, attention_mask = self._splice_input_ids( - u_rows, l_rows, pad_to=pad_to + rows[self._input_name], rows[self._label_name], pad_to=pad_to ) - - if self._smoke_log_once and self._first_predict: - print( - f"[GENRECLM_DEBUG] first batch: B={input_ids.shape[0]} " - f"T={input_ids.shape[1]} pad_id={self._pad_token_id} " - f"ign={self._ignore_index} dev={input_ids.device} " - f"input_ids[0, -8:]={input_ids[0, -8:].tolist()} " - f"labels[0, -8:]={labels[0, -8:].tolist()}", - flush=True, - ) - self._first_predict = False - return self._forward_loss(input_ids, labels, attention_mask) def _forward_loss( @@ -245,28 +213,24 @@ def _forward_loss( outputs = self.lm.model(input_ids=input_ids, attention_mask=attention_mask) hidden = outputs.last_hidden_state # (B, T, D) - # Slice the fixed-width supervised suffix (constant -> no per-step sync). - # Outside it every label is -100 (CE unchanged); it also bounds the logits - # to (B, suffix, V) — the full (B, T, vocab) + fp32 upcast would OOM. - sl = slice(-self._suffix_keep, None) - labels_sl = labels[:, sl] - logits = self.lm.lm_head(hidden[:, sl, :]) + # Bound the logits to the supervised suffix; a full (B, T, V) upcast OOMs. + suffix = slice(-self._suffix_keep, None) + logits = self.lm.lm_head(hidden[:, suffix, :]) - # HF ForCausalLMLoss: shift-by-one + CE with -100 ignore. loss = self.lm.loss_function( logits=logits, - labels=labels_sl, + labels=labels[:, suffix], vocab_size=self.lm.config.vocab_size, + ignore_index=self._ignore_index, ) - return {"loss": loss, "logits": logits} + return {"loss": loss} def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: - """Branch 2: beam-search the SID answer (no ground truth supplied). + """Beam-search the SID answer (no ground truth supplied). - Builds the prompt (no answer), generates up to ``num_levels`` new tokens - per beam, and hands the generated tail to the base - ``_validate_sid_candidates`` (token->SID, malformed beams -> ``-1``). - Returns ``generated_sids`` of shape ``(B, num_return, num_levels)``. + Returns ``generated_sids`` of shape ``(B, C, num_levels)``, where ``C`` + is ``num_return_sequences`` on the HF path and the escalating beam's + final width when ``dynamic_beam`` is set. """ u_rows = self.build_input(batch)[self._input_name] input_ids, attention_mask = self._splice_prompt_ids(u_rows) @@ -289,12 +253,7 @@ def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: def _dynamic_beam_search( self, input_ids: torch.Tensor, attention_mask: torch.Tensor ) -> torch.Tensor: - """ALGR-style escalating-beam decode (delegates to the shared kernel). - - Returns the generated tail ``(B * num_beams * 2**num_levels, num_levels)`` - score-ordered best-first per row, ready for ``_validate_sid_candidates``. - See ``escalating_beam_search`` for the schedule. - """ + """Escalating-beam decode; see ``escalating_beam_search`` for the schedule.""" lo_tok, hi_tok = self._sid_token_bands() return escalating_beam_search( self.lm, @@ -308,38 +267,17 @@ def _dynamic_beam_search( def _splice_prompt_ids( self, user_seq_rows: List[torch.Tensor] ) -> Tuple[torch.Tensor, torch.Tensor]: - """Assemble the answer-less prompt and left-pad into ``(B, T_max)``. - - Layout: ``[system | user_prefix | history | user_suffix | asst_prefix]`` - — everything up to (but not including) the answer, so generation - continues from the assistant turn. - """ - rows = [ - torch.cat( - [ - self.tpl_system, - self.tpl_user_prefix, - r, - self.tpl_user_suffix, - self.tpl_asst_prefix, - ] - ) - for r in user_seq_rows - ] - return self._left_pad(rows) + """Assemble the answer-less prompt and left-pad into ``(B, T_max)``.""" + return self._left_pad(self._prompt_rows(user_seq_rows)) def _left_pad( self, rows: List[torch.Tensor], pad_to: int = 0 ) -> Tuple[torch.Tensor, torch.Tensor]: """Left-pad token rows into ``(input_ids, attention_mask)``, ``(B, T_max)``. - Real content is right-aligned, pad at the front. ``attention_mask`` is - built from ``ones_like(row)`` (not ``!= pad``) so a real trailing eos is - never masked when ``pad_token_id == eos``. - - ``pad_to`` left-extends the batch to at least that many columns (the - first-step activation-pool pre-sizing). Extending on the LEFT keeps the - end-aligned supervised tail in place, so labels/suffix-slice are intact. + ``attention_mask`` is built from ``ones_like(row)`` (not ``!= pad``) so a + real trailing eos is never masked when ``pad_token_id == eos``. ``pad_to`` + extends on the LEFT, keeping the end-aligned supervised tail in place. """ input_ids = pad_sequence( rows, diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/qwen2_rec_lm_test.py index ff8cdca5c..d4778fdd7 100644 --- a/tzrec/models/qwen2_rec_lm_test.py +++ b/tzrec/models/qwen2_rec_lm_test.py @@ -22,11 +22,8 @@ def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): """A Qwen2RecLM with the splice-relevant state wired up, no HF backbone. - Template buffers use tiny placeholder ids so the spliced layout is easy to - read; real buffers come from ``_build_prompt_tokens`` at init time. - - The non-uniform default makes incorrect ``level * uniform_size`` offset - arithmetic visible: sizes=[2,3,4], offsets=[0,2,5]. + The non-uniform default codebook makes incorrect ``level * uniform_size`` + offset arithmetic visible: sizes=[2,3,4], offsets=[0,2,5]. """ codebook = codebook or [2, 3, 4] m = object.__new__(Qwen2RecLM) @@ -35,9 +32,9 @@ def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): m._num_levels = len(codebook) m._base_vocab = base_vocab m._pad_token_id = pad_id - m._dynamic_beam = False # default = HF fixed-width beam path - m._max_seq_length = 0 # no recency clip by default in unit stubs - m._generated_sids_key = "generated_sids" # configurable; default key + m._dynamic_beam = False + m._max_seq_length = 0 + m._generated_sids_key = "generated_sids" m.lm = types.SimpleNamespace(device=torch.device(device)) for name, vals in { "tpl_system": [10, 11], @@ -48,21 +45,21 @@ def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): "tpl_eos": [9], }.items(): m.register_buffer(name, torch.tensor(vals, dtype=torch.long), persistent=False) - lo, hi = Qwen2RecLM._sid_level_bands(codebook) - m.register_buffer("_level_offsets", lo - 1, persistent=False) - m.register_buffer("_codebook_sizes", hi - lo + 1, persistent=False) + offsets, sizes = Qwen2RecLM._sid_level_layout(codebook) + m.register_buffer("_level_offsets", offsets, persistent=False) + m.register_buffer("_codebook_sizes", sizes, persistent=False) return m def _gen_batch(): - """A one-row inference batch (history SIDs [1, 2, 3]) for ``_generate`` tests.""" + """An opaque one-row batch; the ``_generate`` tests mock ``build_input``.""" return types.SimpleNamespace( sequence_dense_features={"user_sequence": _FakeJT([1, 2, 3], [3])} ) def _first_non_neg_index(labels): - """The per-step suffix bound _forward_loss's cached _suffix_keep replaces.""" + """First supervised label column, minimized over rows.""" tmp = (labels >= 0).cumsum(dim=-1) return int((tmp == 1).float().argmax(dim=-1).min().item()) @@ -73,8 +70,7 @@ def test_splice_layout_and_labels(self) -> None: u = [torch.tensor([100, 101, 102])] a = [torch.tensor([200, 201, 202])] # 3 codes = num_levels ids, labels, mask = m._splice_input_ids(u, a) - # [system | user_prefix | history | user_suffix | - # asst_prefix | answer | asst_suffix | eos] + # sys|user_prefix|history|user_suffix|asst_prefix|answer|asst_suffix|eos self.assertEqual( ids[0].tolist(), [10, 11, 12, 100, 101, 102, 13, 14, 200, 201, 202, 15, 9] ) @@ -92,15 +88,13 @@ def test_left_padding_varied_lengths(self) -> None: ids, labels, mask = m._splice_input_ids(u, a) T = ids.shape[1] n1 = 2 + 1 + 1 + 1 + 1 + 3 + 1 + 1 # shorter row's real length - # shorter row is left-padded: pad at the front, content right-aligned self.assertEqual(ids[1, : T - n1].tolist(), [m._pad_token_id] * (T - n1)) self.assertEqual(mask[1].tolist(), [0] * (T - n1) + [1] * n1) self.assertEqual(labels[1, : T - n1].tolist(), [-100] * (T - n1)) - # every row's trailing eos is supervised and the answer ends just before + # the trailing eos is supervised in every row self.assertEqual(labels[:, -1].tolist(), [9, 9]) def test_mask_keeps_trailing_eos_when_pad_equals_eos(self) -> None: - # pad_id == eos value: the mask must NOT mask the real trailing eos m = _stub(pad_id=9) # tpl_eos == 9 too ids, _, mask = m._splice_input_ids( [torch.tensor([100])], [torch.tensor([200, 201, 202])] @@ -110,19 +104,15 @@ def test_mask_keeps_trailing_eos_when_pad_equals_eos(self) -> None: self.assertEqual(mask[0].tolist(), [1] * ids.shape[1]) def test_suffix_keep_matches_dynamic_slice(self) -> None: - # _forward_loss caches self._suffix_keep instead of recomputing the suffix - # bound per step (two GPU->CPU syncs); prove the constant equals the old - # dynamic computation and that the slice drops nothing supervised. + # the constant suffix width must cover every supervised column. m = _stub() # num_levels=3, asst_suffix=[15] (numel 1) -> suffix_keep=6 suffix_keep = m._num_levels + m.tpl_asst_suffix.numel() + 2 _, labels, _ = m._splice_input_ids( [torch.tensor([100, 101, 102])], [torch.tensor([200, 201, 202])] ) - # the constant matches the per-step bound it replaces self.assertEqual( suffix_keep, labels.shape[1] - _first_non_neg_index(labels) + 1 ) - # everything before the kept suffix is unsupervised (-100): nothing dropped self.assertTrue(bool((labels[:, :-suffix_keep] < 0).all())) def test_splice_prompt_ids(self) -> None: @@ -138,7 +128,7 @@ def test_predict_routes_on_inference_flag(self) -> None: m._generate = lambda b: {"branch": "generate"} m._is_inference = False # train / eval self.assertEqual(Qwen2RecLM.predict(m, object())["branch"], "train") - m._is_inference = True # inference (set_is_inference in main.py) + m._is_inference = True # inference self.assertEqual(Qwen2RecLM.predict(m, object())["branch"], "generate") def test_generate_maps_tokens_to_sids(self) -> None: @@ -156,8 +146,7 @@ def fake_generate( pad_token_id, ): prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) - # codebook=[2,3,4], offsets=[0,2,5]: - # local [1,1,1] / [2,3,4] -> the min/max token of each level. + # offsets [0,2,5]: local [1,1,1]/[2,3,4] -> each level's min/max token new = torch.tensor([[100, 102, 105], [101, 104, 108]]) return torch.cat([prompt, new], dim=1) @@ -169,7 +158,7 @@ def fake_generate( self.assertEqual(sids[0].tolist(), [[1, 1, 1], [2, 3, 4]]) def test_generate_rejects_malformed_candidates(self) -> None: - # Layer-A gate: every malformed candidate -> the -1 sentinel, in place. + # every malformed candidate collapses to the -1 sentinel, in place. m = _stub(base_vocab=100) m._input_name = "user_sequence" m._num_beams = m._num_return = 6 @@ -200,7 +189,6 @@ def fake_generate( m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} sids = m._generate(_gen_batch())["generated_sids"] self.assertEqual(tuple(sids.shape), (1, 6, 3)) - # valid candidate kept at its rank; every malformed one -> all -1 (in place) self.assertEqual( sids[0].tolist(), [ @@ -214,8 +202,7 @@ def fake_generate( ) def test_generate_narrow_tail_no_crash(self) -> None: - # every beam emits EOS before num_levels -> generate() returns a tail - # narrower than num_levels; the canvas keeps the reshape rectangular. + # early EOS -> a tail narrower than num_levels must still reshape cleanly m = _stub(base_vocab=100) m._input_name = "user_sequence" m._num_beams = m._num_return = 2 @@ -236,7 +223,7 @@ def fake_generate( m.lm.generate = fake_generate m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} sids = m._generate(_gen_batch())["generated_sids"] - self.assertEqual(tuple(sids.shape), (1, 2, 3)) # rectangular, no crash + self.assertEqual(tuple(sids.shape), (1, 2, 3)) # the missing 3rd atom stays -1 -> out of band -> whole candidate -1 self.assertEqual(sids[0].tolist(), [[-1, -1, -1], [-1, -1, -1]]) @@ -325,8 +312,7 @@ def _vl(): def test_compute_max_total_length(self) -> None: m = _stub() - # frame = |system|2 + |user_prefix|1 + |user_suffix|1 + |asst_prefix|1 - # + |asst_suffix|1 + |eos|1 = 7; + max_history + answer(num_levels) + # frame = 2 system + 1 each user_pfx/sfx, asst_pfx/sfx, eos = 7 m._max_seq_length = 300 self.assertEqual(m._compute_max_total_length(), 7 + 300 + 3) m._max_seq_length = 0 # pre-allocation disabled @@ -335,7 +321,6 @@ def test_compute_max_total_length(self) -> None: def test_first_step_pads_to_max_then_actual_length(self) -> None: m = _stub() m._is_inference = False # not inference + nn.Module.training=True -> is_train - m._smoke_log_once = False m._input_name, m._label_name = "user_sequence", "label" m._max_total_len = 50 m._pool_warmed = False @@ -353,8 +338,6 @@ def fwd(i, lbl, a): batch = object() m._predict_train(batch) # first step: pre-size to worst case m._predict_train(batch) # subsequent step: natural length - # one-shot: first step left-pads to _max_total_len, latched by - # _pool_warmed; later steps use the actual (shorter) length. self.assertEqual(seen_lens[0], 50) self.assertLess(seen_lens[1], 50) self.assertTrue(m._pool_warmed) @@ -362,7 +345,6 @@ def fwd(i, lbl, a): def test_no_forced_padding_when_disabled(self) -> None: m = _stub() m._is_inference = False - m._smoke_log_once = False m._input_name, m._label_name = "user_sequence", "label" m._max_total_len = 0 # max_seq_length unset -> pre-allocation off m._pool_warmed = False @@ -379,18 +361,89 @@ def fwd(i, lbl, a): m._forward_loss = fwd batch = object() m._predict_train(batch) - # disabled: natural length, never forced to max; flag stays unlatched. self.assertLess(seen_lens[0], 50) self.assertFalse(m._pool_warmed) + def test_splice_row_count_mismatch_raises(self) -> None: + m = _stub() + with self.assertRaisesRegex(ValueError, "row count mismatch"): + m._splice_input_ids([torch.tensor([100])], []) + + +class Qwen2ForwardLossTest(unittest.TestCase): + """The training objective, run for real against a tiny Qwen2 backbone.""" + + def _model(self, ignore_index=-100): + m = _real_lm_stub(codebook=[2, 3, 4], base_vocab=20, num_beams=2) + m._ignore_index = ignore_index + m._pad_token_id = 0 + for name, vals in { + "tpl_system": [1, 2], + "tpl_user_prefix": [3], + "tpl_user_suffix": [4], + "tpl_asst_prefix": [5], + "tpl_asst_suffix": [6], + "tpl_eos": [7], + }.items(): + m.register_buffer( + name, torch.tensor(vals, dtype=torch.long), persistent=False + ) + m._suffix_keep = m._num_levels + m.tpl_asst_suffix.numel() + 2 + return m + + def _rows(self): + # ragged histories so left padding is exercised + u = [torch.tensor([20, 22, 25]), torch.tensor([21, 23, 26, 20, 24, 27])] + a = [torch.tensor([20, 22, 25]), torch.tensor([21, 24, 28])] + return u, a + + def test_suffix_slice_matches_full_sequence_loss(self) -> None: + # the fixed-width suffix slice must give the same CE as full-T logits + m = self._model() + ids, labels, mask = m._splice_input_ids(*self._rows()) + with torch.no_grad(): + got = m._forward_loss(ids, labels, mask)["loss"] + full = m.lm(input_ids=ids, attention_mask=mask).logits + ref = m.lm.loss_function( + logits=full, labels=labels, vocab_size=m.lm.config.vocab_size + ) + self.assertTrue(torch.allclose(got, ref, atol=1e-6)) + + def test_loss_is_invariant_to_extra_left_padding(self) -> None: + # the pool-warmup pad_to must not perturb the objective + m = self._model() + u, a = self._rows() + with torch.no_grad(): + base = m._forward_loss(*m._splice_input_ids(u, a))["loss"] + padded = m._forward_loss(*m._splice_input_ids(u, a, pad_to=40))["loss"] + self.assertTrue(torch.allclose(base, padded, atol=1e-6)) + + def test_forward_loss_honours_configured_ignore_index(self) -> None: + # ignore_index must reach loss_function or every pad slot is supervised + m = self._model(ignore_index=-100) + u, a = self._rows() + with torch.no_grad(): + default = m._forward_loss(*m._splice_input_ids(u, a))["loss"] + m._ignore_index = -7 + with torch.no_grad(): + ids, labels, mask = m._splice_input_ids(u, a) + custom = m._forward_loss(ids, labels, mask)["loss"] + self.assertEqual(int((labels == -7).sum() > 0), 1) + self.assertTrue(torch.allclose(default, custom, atol=1e-6)) + + def test_forward_loss_returns_only_the_loss(self) -> None: + # returned logits would stay alive across the next step's fwd/bwd + m = self._model() + with torch.no_grad(): + out = m._forward_loss(*m._splice_input_ids(*self._rows())) + self.assertEqual(list(out), ["loss"]) + def _real_lm_stub(codebook=None, base_vocab=20, num_beams=2): """A Qwen2RecLM carrying a real (tiny, random) Qwen2 backbone. - Needed by the dynamic-beam tests, which exercise the actual KV-cached - forward / cache-reorder path (the other tests mock ``lm.generate``). The SID - atoms occupy the last ``sum(codebook)`` token ids, matching the base-vocab - plus appended-codebook layout. + Needed by the dynamic-beam tests, which exercise the real KV-cached forward + and cache-reorder path; the other tests mock ``lm.generate``. """ from transformers import Qwen2Config, Qwen2ForCausalLM @@ -411,15 +464,14 @@ def _real_lm_stub(codebook=None, base_vocab=20, num_beams=2): ) torch.manual_seed(0) m.lm = Qwen2ForCausalLM(cfg).eval() - lo, hi = Qwen2RecLM._sid_level_bands(codebook) - m.register_buffer("_level_offsets", lo - 1, persistent=False) - m.register_buffer("_codebook_sizes", hi - lo + 1, persistent=False) + offsets, sizes = Qwen2RecLM._sid_level_layout(codebook) + m.register_buffer("_level_offsets", offsets, persistent=False) + m.register_buffer("_codebook_sizes", sizes, persistent=False) return m class Qwen2DynamicBeamTest(unittest.TestCase): def test_width_schedule_and_final_count(self) -> None: - # widths double per level, returning num_beams * 2**num_levels candidates. m = _real_lm_stub(codebook=[8, 7, 6], base_vocab=20, num_beams=2) ids = torch.tensor([[1, 2, 3, 4]]) new = m._dynamic_beam_search(ids, torch.ones_like(ids)) @@ -457,10 +509,8 @@ def test_left_padding_two_rows(self) -> None: self.assertTrue(bool((sids[..., level] <= size).all())) def test_exhaustive_matches_bruteforce_topk(self) -> None: - # When the schedule covers the whole tree (no pruning) the escalating - # beam is EXACT: its candidate set must equal all SID combos and be - # ordered by true (full-recompute) cumulative log-prob. This validates - # cache-stepping, band masking, scoring, and ordering end-to-end. + # with no pruning the beam is EXACT: every SID combo, ordered by the + # true (full-recompute) cumulative log-prob. codebook, base = [2, 3], 20 m = _real_lm_stub(codebook=codebook, base_vocab=base, num_beams=3) # widths [min(6,2)=2, min(12,6)=6] -> exhaustive over all 2*3 SIDs @@ -482,9 +532,8 @@ def test_exhaustive_matches_bruteforce_topk(self) -> None: ) for t1 in range(lo1, lo1 + codebook[1]): ref[(t0, t1)] = (logp0[t0] + logp1[t1]).item() - # 1. exhaustive: the returned set is exactly every SID combination self.assertEqual(set(got), set(ref)) - # 2. ordered best-first by the true score (tolerant of float-noise ties) + # ordered best-first by the true score (tolerant of float-noise ties) s = [ref[c] for c in got] self.assertTrue(all(s[i] >= s[i + 1] - 1e-4 for i in range(len(s) - 1))) self.assertEqual(got[0], max(ref, key=ref.get)) # top-1 is the global best diff --git a/tzrec/models/escalating_beam.py b/tzrec/modules/escalating_beam.py similarity index 56% rename from tzrec/models/escalating_beam.py rename to tzrec/modules/escalating_beam.py index 45d816938..35b992e34 100644 --- a/tzrec/models/escalating_beam.py +++ b/tzrec/modules/escalating_beam.py @@ -9,28 +9,22 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""ALGR-style escalating-beam SID decode (torch-only, no tzrec deps). +"""Escalating-beam SID decode (torch-only, no tzrec deps). -A faithful port of ALGR's ``dynamic_beams`` schedule: the beam width doubles at -every SID level (``num_beams`` -> ``2*num_beams`` -> ...), keeping -``num_beams * 2**(j+1)`` candidates after level ``j`` and returning -``num_beams * 2**num_levels`` per row. The aggressive early pruning (only the -top ``2*num_beams`` level-0 prefixes survive) is what distinguishes it from a -fixed-width beam that keeps every level-0 code. - -Lives in its own torch-only module so both the production path -(``Qwen2RecLM._dynamic_beam_search``) and the offline predict harness share one -tested implementation. +The beam width doubles at every SID level, so early levels are pruned hard. """ -from __future__ import annotations +from typing import TYPE_CHECKING, List, Tuple import torch +if TYPE_CHECKING: + from transformers import PreTrainedModel + @torch.no_grad() def escalating_beam_search( - model, + model: "PreTrainedModel", input_ids: torch.Tensor, attention_mask: torch.Tensor, *, @@ -50,30 +44,35 @@ def escalating_beam_search( hi_tok: inclusive upper per-level token-space band edge, ``(num_levels,)``. Returns: - The generated SID token tail ``(B * num_beams * 2**num_levels, - num_levels)``, score-ordered best-first per row. The SID answer is - exactly ``num_levels`` codes with no in-answer EOS, so every beam emits - exactly ``num_levels`` tokens — no finished-beam bookkeeping is needed, - and band masking makes every candidate well-formed by construction. + The generated SID token tail ``(B * W, num_levels)`` score-ordered + best-first per row, where ``W`` is ``num_beams * 2**num_levels`` capped + to the number of distinct SIDs the codebook can supply. The answer is + fixed-length and EOS-free, so no finished-beam bookkeeping is needed. """ device = input_ids.device bsz = input_ids.shape[0] num_levels = lo_tok.shape[0] - # candidates kept after level j (doubling), capped to what the band + the - # surviving prefixes can actually supply (guards tiny codebooks). - widths, prev = [], 1 - for j in range(num_levels): - avail = prev * int(hi_tok[j] - lo_tok[j] + 1) - widths.append(min(num_beams * (2 ** (j + 1)), avail)) + # Hoist the band edges to host once to keep the level loop sync-free. + bands: List[Tuple[int, int]] = [ + (int(lo_tok[j]), int(hi_tok[j])) for j in range(num_levels) + ] + # Cap the doubling at what band x surviving prefixes supply (tiny codebooks). + widths: List[int] = [] + prev = 1 + for j, (lo, hi) in enumerate(bands): + widths.append(min(num_beams * (2 ** (j + 1)), prev * (hi - lo + 1))) prev = widths[-1] def _band_logp(logits: torch.Tensor, j: int) -> torch.Tensor: - ids = torch.arange(logits.shape[-1], device=device) - keep = (ids >= lo_tok[j]) & (ids <= hi_tok[j]) - logp = torch.log_softmax(logits.float(), dim=-1) - return logp.masked_fill(~keep, float("-inf")) + """Full-vocab log-probs, narrowed to level ``j``'s band ``(R, band)``. + + Slicing after normalizing keeps the exact cross-beam ranking of a + full-vocab ``log_softmax`` without materializing one per level. + """ + lo, hi = bands[j] + log_z = torch.logsumexp(logits.float(), dim=-1, keepdim=True) + return logits[:, lo : hi + 1].float() - log_z - # 1. prompt forward (bsz beams) -> level-0 logits. pos = (attention_mask.long().cumsum(-1) - 1).clamp(min=0) h = model.model( input_ids=input_ids, @@ -82,17 +81,15 @@ def _band_logp(logits: torch.Tensor, j: int) -> torch.Tensor: use_cache=True, ) past = h.past_key_values - vocab = model.config.vocab_size scores = _band_logp(model.lm_head(h.last_hidden_state[:, -1, :]), 0) - beam_scores, tok = scores.topk(widths[0], dim=-1) # (B, W0) - seq = tok.reshape(-1, 1) + beam_scores, local = scores.topk(widths[0], dim=-1) # (B, W0) + seq = (local + bands[0][0]).reshape(-1, 1) beam_scores = beam_scores.reshape(-1) parent = torch.arange(bsz, device=device).repeat_interleave(widths[0]) past.reorder_cache(parent) am = attention_mask.repeat_interleave(widths[0], dim=0) cur_w = widths[0] - # 2. levels 1..n-1: forward the last chosen atom with cache, escalate width. for j in range(1, num_levels): am = torch.cat([am, am.new_ones(bsz * cur_w, 1)], dim=1) step_pos = (am.long().cumsum(-1) - 1)[:, -1:].clamp(min=0) @@ -105,17 +102,20 @@ def _band_logp(logits: torch.Tensor, j: int) -> torch.Tensor: use_cache=True, cache_position=cache_pos, ) + lo_j, hi_j = bands[j] + band = hi_j - lo_j + 1 scores = _band_logp(model.lm_head(h.last_hidden_state[:, -1, :]), j) - scores = scores + beam_scores[:, None] # (B*cur_w, V) cumulative - # global top-widths[j] per row over the (cur_w * V) continuations. - beam_scores, idx = scores.view(bsz, cur_w * vocab).topk(widths[j], dim=-1) - parent_local = torch.div(idx, vocab, rounding_mode="floor") # in [0,cur_w) - tok = idx % vocab + scores = scores + beam_scores[:, None] # (B*cur_w, band) cumulative + beam_scores, idx = scores.view(bsz, cur_w * band).topk(widths[j], dim=-1) + parent_local = torch.div(idx, band, rounding_mode="floor") + tok = lo_j + idx % band row_base = torch.arange(bsz, device=device)[:, None] * cur_w parent = (parent_local + row_base).reshape(-1) - past.reorder_cache(parent) - am = am[parent] seq = torch.cat([seq[parent], tok.reshape(-1, 1)], dim=1) beam_scores = beam_scores.reshape(-1) cur_w = widths[j] + if j + 1 < num_levels: + # the last level never reads the cache; skip the largest reorder copy. + past.reorder_cache(parent) + am = am[parent] return seq # (B*cur_w, num_levels) diff --git a/tzrec/optim/lr_scheduler.py b/tzrec/optim/lr_scheduler.py index fddbbea6b..20c40bdeb 100644 --- a/tzrec/optim/lr_scheduler.py +++ b/tzrec/optim/lr_scheduler.py @@ -210,8 +210,7 @@ def _get_lr(self) -> List[float]: t = min(step_count - self._warmup_size, self._total_size - self._warmup_size) decay_scale = 1.0 - t / (self._total_size - self._warmup_size) return [ - self._min_learning_rate - + (base_lr - self._min_learning_rate) * decay_scale + self._min_learning_rate + (base_lr - self._min_learning_rate) * decay_scale for base_lr in self.base_lrs ] diff --git a/tzrec/protos/model.proto b/tzrec/protos/model.proto index f25964953..123669f47 100644 --- a/tzrec/protos/model.proto +++ b/tzrec/protos/model.proto @@ -83,8 +83,7 @@ message ModelConfig { SidRqvae sid_rqvae = 600; SidRqkmeans sid_rqkmeans = 601; - // Generative (causal-LM) models — one family message per LLM family. - // 700-block keeps clear of the SID family's 600-block. + // Generative (causal-LM) models; the 700-block keeps clear of the SID 600s. Qwen2RecLM qwen2_rec_lm = 700; } diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index 99919442d..4867bcb6c 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -1,64 +1,46 @@ syntax = "proto2"; package tzrec.protos; -// Generative (causal-LM) models. +// Generative (causal-LM) recommendation models. // -// Each LLM family is registered DIRECTLY: the model_config oneof picks a family -// message (e.g. `qwen2_rec_lm`), whose message type name (`Qwen2RecLM`) resolves -// to the Python class of the same name. There is no `class_name` dispatch. -// -// Shared, architecture-agnostic config lives in `GenerativeRecLMConfig` and is -// embedded as `common` in every family message; family-specific knobs (e.g. the -// chat template) live on the family message itself. Adding a family = a new -// message like `Qwen2RecLM` + a same-named Python subclass of GenerativeRecLM. +// Shared, architecture-agnostic config lives in `GenerativeRecLMConfig`, embedded +// as `common` in every family message; family-specific knobs (the backbone, the +// chat template) live on the family message itself. -// Architecture-agnostic config shared by ALL generative-rec families (the base -// reads this for everything except the backbone, which the family owns — see -// _backbone_id). Sample contract (consumed by `predict()`): -// * history : list — local 1-based per-level codes in +// Architecture-agnostic config shared by all generative-rec families. +// Sample contract: +// * history : list -- local 1-based per-level codes in // [1, codebook[level]], laid out as whole items in level order; -// the single JAGGED_SEQUENCE feature_group. -// * answer : list — one item's local 1-based codes; a -// `data_config.label_field` (read from batch.jagged_labels), -// NOT a feature. +// the single JAGGED_SEQUENCE feature_group (its one member). +// * answer : list -- one item's local 1-based codes; the FIRST +// `data_config.label_field`, NOT a feature. // The model derives level_offsets from codebook and maps each code at batch // time: token_id = base_vocab + level_offsets[level] + code - 1. Sample -// writers must not pre-apply level_offsets. Column names are not configured -// here (see below). +// writers must not pre-apply level_offsets. message GenerativeRecLMConfig { - // Backbone (`hf_model_id`) is NOT here — it's the family's architecture - // commitment, so it lives on the family message. - // SID vocabulary, one entry per RQ level: each public code is in // [1, codebook[level]]. len = codes per item (answer width); sum = atoms // appended as C0..C{sum-1} after the base vocab. repeated uint32 codebook = 2; - // Pad the post-extension vocab up to a multiple of this value. + // Pad the post-extension vocab up to a multiple of this value; 0 disables + // padding. optional uint32 vocab_pad_to_multiple_of = 3 [default = 128]; - // Neither the history feature_group nor the SID-column names are configured - // here — they're derived from the config that already declares them: the - // history is the single declared feature_group (a JAGGED_SEQUENCE group; - // build_input keys the EmbeddingGroup output by its GROUP name, the HSTU - // idiom), and its one member is the history feature; the answer is the first - // data_config.label_field (read from batch.jagged_labels, like RankModel's - // labels[0]). The answer is NOT a feature_group. - // Cross-entropy ignore index — matches PyTorch's F.cross_entropy default. + // Cross-entropy ignore index -- matches PyTorch's F.cross_entropy default. optional int32 ignore_index = 6 [default = -100]; - // Inference (beam search) — used only by predict()'s inference branch. + // Beam search (inference only). num_return_sequences must not exceed + // num_beams. optional uint32 num_beams = 7 [default = 50]; optional uint32 num_return_sequences = 8 [default = 50]; - // When set, decode with the ALGR-style escalating beam instead of HF's - // fixed-width beam search: the beam width doubles at every SID level - // (num_beams -> 2*num_beams -> ... ), returning num_beams * 2**num_levels - // candidates (num_return_sequences is ignored). Faithful parity with ALGR's - // dynamic_beams schedule; exploits the fixed-length, EOS-free SID answer. + // When set, decode with the escalating beam: the width doubles at every SID + // level, returning up to num_beams * 2**num_levels candidates (capped by the + // codebook). num_return_sequences is ignored. optional bool dynamic_beam = 9 [default = false]; - // Predictions key the inference branch emits generated SIDs under (stable - // across families; PredictWrapper output_cols should reference it). + // Prediction key the generated SIDs are emitted under; reference it from + // PredictWrapper output_cols. optional string generated_sids_key = 12 [default = "generated_sids"]; // Backbone PARAM dtype = the fp32 MASTER weights. "float32" avoids bf16-ULP // underflow of Adam's small (lr=1e-5) updates; bf16 COMPUTE comes from @@ -66,31 +48,25 @@ message GenerativeRecLMConfig { // float32 | bfloat16 | float16. optional string param_dtype = 13 [default = "float32"]; - // Model's history budget (SID codes): the truncation cap enforced model-side - // (_sid_token_rows, item-aligned, recency-preserving) AND the activation-pool - // pre-size (_compute_max_total_length). A model knob like HSTU's - // DlrmHSTU.max_seq_len, distinct from the user_sequence feature's - // sequence_length. REQUIRED (like HSTU's DlrmHSTU.max_seq_len); set it to 0 - // to explicitly disable the cap + activation-pool pre-allocation. + // Model's history budget in SID codes: the item-aligned, recency-preserving + // truncation cap AND the activation-pool pre-size. Distinct from the history + // feature's own sequence_length. Set to 0 to disable both. required uint32 max_sequence_length = 14; } -// Qwen2 / Qwen2.5 family (Qwen2.5-0.5B, etc.). Registered directly via its -// message-type name; the Python class is `Qwen2RecLM`. +// Qwen2 / Qwen2.5 family (Qwen2.5-0.5B, etc.). message Qwen2RecLM { - // Architecture-agnostic config shared by all generative-rec families. optional GenerativeRecLMConfig common = 1; - // Qwen2 backbone (HF hub id or local path; must be a Qwen2 model). Owned by - // this family message, not `common`. Default = canonical 0.5B. + // Qwen2 backbone: HF hub id or local path; must be a Qwen2 model. optional string hf_model_id = 2 [default = "Qwen/Qwen2.5-0.5B"]; - // ----- Chat template (family-specific) ----- - // Optional override for the system instruction. Empty -> the family default - // (Qwen2RecLM's algr-matching default). + // ----- Chat template ----- + // Optional override for the system instruction. Empty -> the Qwen2 default. optional string system_instruction = 10 [default = ""]; - // Optional CN/EN text wrapping the SID codes in the user message, e.g. - // "当前用户的历史行为如下：" / "，请预测用户在电商推荐场景后续行为的语义编码". + // Optional text wrapping the SID codes in the user message, e.g. + // "The user's history is as follows: " / ", predict the semantic id of the + // next item the user will interact with". optional string user_prefix_text = 11 [default = ""]; optional string user_suffix_text = 12 [default = ""]; } diff --git a/tzrec/protos/optimizer.proto b/tzrec/protos/optimizer.proto index 63fd6708f..7ee350bd9 100644 --- a/tzrec/protos/optimizer.proto +++ b/tzrec/protos/optimizer.proto @@ -238,7 +238,9 @@ message ManualStepLR { message LinearDecayLR { // total number of steps or epochs to decay from base_lr to - // min_learning_rate (mirrors HF Trainer's `lr_scheduler_type: linear`) + // min_learning_rate (mirrors HF Trainer's `lr_scheduler_type: linear`). + // NOTE: unlike decay_size/T_max in the schedulers above, this horizon is + // measured from step 0 and INCLUDES warmup_size. Required (must be > 0). optional uint32 total_size = 1; // minimum learning rate reached at total_size optional float min_learning_rate = 2 [default = 0.0]; diff --git a/tzrec/utils/checkpoint_util.py b/tzrec/utils/checkpoint_util.py index 0d1b0699d..1853b7e90 100644 --- a/tzrec/utils/checkpoint_util.py +++ b/tzrec/utils/checkpoint_util.py @@ -386,31 +386,25 @@ def save( ) -> str: """Save a checkpoint at the given step, then request an async prune. - When ``export_format == HF``, co-locates the HF config + tokenizer (no - weights) in this checkpoint dir so each ``model.ckpt-N/`` is - self-contained and convertible to HF (design §2); ``write_hf_assets`` - no-ops for non-HF models, so gating on the export format is enough. + For HF-backed models, co-locates the HF config + tokenizer (no weights) + in this checkpoint dir so each ``model.ckpt-N/`` is self-contained and + convertible to HF. Deliberately not gated on ``export_format``: that is + an export-time knob, and gating it would make a run trained with the + default format permanently unexportable to HF. """ ckpt_dir = os.path.join(self._model_dir, f"model.ckpt-{step}") save_model(ckpt_dir, model, optimizer) - if ( - self._export_config is not None - and self._export_config.export_format == export_pb2.ExportFormat.HF - ): - # Local import avoids a circular import (export_util imports us). - from tzrec.utils.export_util import write_hf_assets - - # HF assets are convenience metadata; the DCP weights (save_model - # above) are already durable. Isolate a failure so it can't abort the - # save before the next collective (save_dataloader_state's all_gather) - # and one-sidedly hang the other ranks. - try: - write_hf_assets(model, ckpt_dir) - except Exception as e: # noqa: BLE001 - logger.warning( - f"write_hf_assets failed for {ckpt_dir}: {e} — checkpoint " - f"weights are saved; skipping HF assets." - ) + # Local import avoids a circular import (hf_export_util imports us). + from tzrec.utils.hf_export_util import write_hf_assets + + # a raise here skips save_dataloader_state's all_gather and hangs other ranks. + try: + write_hf_assets(model, ckpt_dir) + except Exception as e: # noqa: BLE001 + logger.warning( + f"write_hf_assets failed for {ckpt_dir}: {e} -- checkpoint " + f"weights are saved; skipping HF assets." + ) if dataloader_state is not None: save_dataloader_state(ckpt_dir, dataloader_state) self._last_ckpt_dir = ckpt_dir @@ -1071,9 +1065,7 @@ def restore_model( def save_model( - checkpoint_dir: str, - model: nn.Module, - optimizer: Optional[optim.Optimizer] = None, + checkpoint_dir: str, model: nn.Module, optimizer: Optional[optim.Optimizer] = None ) -> None: """Save model state. diff --git a/tzrec/utils/export_util.py b/tzrec/utils/export_util.py index 34f14c839..e81c49c0e 100644 --- a/tzrec/utils/export_util.py +++ b/tzrec/utils/export_util.py @@ -84,188 +84,6 @@ from tzrec.utils.state_dict_util import fix_mch_state, init_parameters -# HF config/tokenizer asset files co-located in each checkpoint dir (no -# weights) and copied into the converted HF export dir. Missing files are -# skipped — different tokenizers emit different subsets (e.g. BPE merges vs. -# sentencepiece vocab). -_HF_ASSET_FILES = ( - "config.json", - "generation_config.json", - "tokenizer.json", - "tokenizer_config.json", - "vocab.json", - "merges.txt", - "added_tokens.json", - "special_tokens_map.json", -) - -_HF_EXPORT_META_FILENAME = "hf_export_meta.json" - -# Max wrapper layers to unwrap to reach the HF-backed model: DMP(.module) -> -# TrainWrapper(.model) -> model is 3 hops; 4 is a small cycle guard. -_MAX_WRAPPER_DEPTH = 4 - - -def _unwrap_hf_model(wrapped_model: nn.Module) -> Optional[nn.Module]: - """Walk DMP/TrainWrapper layers down to the model exposing ``hf_backbone``. - - Handles DMP (``.module`` -> TrainWrapper), TrainWrapper (``.model`` -> - GenerativeRecLM), and the GenerativeRecLM itself. Returns ``None`` if no - backbone is found in the wrapper chain (i.e. not an HF-backed model, so - callers no-op). - """ - m = wrapped_model - for _ in range(_MAX_WRAPPER_DEPTH): - if hasattr(m, "hf_backbone"): - return m - if hasattr(m, "module"): # DMP / DDP-style wrapper - m = m.module - elif hasattr(m, "model"): # Train/Predict/Script wrapper - m = m.model - else: - break - return None - - -def write_hf_assets(wrapped_model: nn.Module, save_dir: str) -> None: - """Co-locate the HF config + tokenizer (NO weights) in a checkpoint dir. - - Writes the architecture (``config.json`` + ``generation_config.json``) and - the extended tokenizer alongside the DCP ``model/`` dir so each - ``model.ckpt-N/`` is self-describing and convertible to HF independent of - the current code (design §2). Also records the backbone's true FQN prefix - off the LIVE module graph into ``hf_export_meta.json`` so ``dcp_to_hf`` can - strip it without hard-coding a wrapper/attribute convention. Rank 0 only — - the dense backbone is data-parallel-replicated, so rank 0 holds it. - - ``wrapped_model`` may be the DMP (``.module.model``), the TrainWrapper - (``.model``), or the GenerativeRecLM itself; the inner model is found by - walking the wrappers down to the one exposing ``hf_backbone``. The recorded - prefix is read off ``wrapped_model.named_modules()`` so it matches the FQNs - that ``save_model`` writes from ``wrapped_model.state_dict()``. - """ - if int(os.environ.get("RANK", 0)) != 0: - return - inner = _unwrap_hf_model(wrapped_model) - if inner is None: # not an HF-backed model -> nothing to co-locate - return - os.makedirs(save_dir, exist_ok=True) - - backbone = inner.hf_backbone() - backbone.config.save_pretrained(save_dir) - gen_cfg = getattr(backbone, "generation_config", None) - if gen_cfg is not None: - gen_cfg.save_pretrained(save_dir) - inner.hf_tokenizer().save_pretrained(save_dir) - - # Record the backbone's FQN prefix as it appears in the SAVED state_dict - # (data, not a magic string): robust to wrapper/parallelism and to families - # that don't name the backbone `self.lm`. `named_modules()` FQNs carry the - # DMP/DDP wrapper prefix (e.g. `_dmp_wrapped_module.module.`) that - # `state_dict()` strips, so map it through the same helper `save_model`'s - # DCP keys went through, or the recorded prefix would not match them. - raw_prefix = next( - (n for n, m in wrapped_model.named_modules() if m is backbone), "" - ) - prefix = checkpoint_util._strip_dmp_prefix(raw_prefix) - meta = {"backbone_state_dict_prefix": prefix + ("." if prefix else "")} - with open(os.path.join(save_dir, _HF_EXPORT_META_FILENAME), "w") as f: - json.dump(meta, f, indent=2) - - -def dcp_to_hf(ckpt_dir: str, out_dir: str) -> None: - """Convert a self-contained checkpoint dir to a ``from_pretrained`` HF dir. - - Standalone — builds no live model and runs no ``from_pretrained`` weight - download (design §2). Reads EVERYTHING from the one checkpoint dir: - materialize the DCP state dict, map the wrapper-prefixed keys onto the bare - HF keys, strict-validate against the architecture built from the co-located - ``config.json``, then write safetensors + copy the config/tokenizer. The - numbered steps below walk through it. - """ - from transformers import AutoConfig, AutoModelForCausalLM - from torch.distributed.checkpoint.state_dict_loader import ( - _load_state_dict_from_keys, - ) - - model_ckpt_path = os.path.join(ckpt_dir, "model") - if not os.path.exists(model_ckpt_path): - raise RuntimeError(f"dcp_to_hf: model DCP dir [{model_ckpt_path}] not exists.") - - # 1. materialize the full (replicated) DCP state dict into a plain dict. - # No keys => load everything; non-distributed => full tensors locally. - raw_state: Dict[str, torch.Tensor] = _load_state_dict_from_keys( - checkpoint_id=model_ckpt_path - ) - - # 2. recorded backbone prefix (data, not a magic string); None => derive. - meta_path = os.path.join(ckpt_dir, _HF_EXPORT_META_FILENAME) - prefix: Optional[str] = None - if os.path.exists(meta_path): - with open(meta_path, "r") as f: - prefix = json.load(f).get("backbone_state_dict_prefix") - - # 3. empty backbone from the checkpoint's OWN config -> the target key set. - cfg = AutoConfig.from_pretrained(ckpt_dir) - empty = AutoModelForCausalLM.from_config(cfg) - target_keys: Set[str] = set(empty.state_dict().keys()) - - def _strip_recorded_prefix() -> Optional[Dict[str, torch.Tensor]]: - """Strip the recorded prefix; None unless it yields an EXACT match.""" - if not prefix: - return None - out = {k[len(prefix) :]: v for k, v in raw_state.items() if k.startswith(prefix)} - return out if set(out) == target_keys else None - - def _derive_by_suffix() -> Optional[Dict[str, torch.Tensor]]: - """Each target key is a unique suffix of exactly one DCP key; None if not.""" - out: Dict[str, torch.Tensor] = {} - for tk in target_keys: - matches = [k for k in raw_state if k == tk or k.endswith("." + tk)] - if len(matches) != 1: - return None - out[tk] = raw_state[matches[0]] - return out - - # The recorded prefix is a hint, not gospel: if it does not map exactly onto - # the architecture (stale metadata, a wrapper/parallelism change since the - # checkpoint was written), self-heal by suffix-matching the actual DCP keys - # to the architecture's keys -- which is independent of any prefix convention. - mapped = _strip_recorded_prefix() - if mapped is None: - if prefix: - logger.warning( - f"dcp_to_hf: recorded prefix [{prefix}] did not map exactly onto " - "the architecture; deriving the backbone prefix by suffix-matching." - ) - mapped = _derive_by_suffix() - - # STRICT validation: exact 1:1 with the architecture, fail loudly. Reached - # only when neither path matched -- a genuine architecture/checkpoint - # mismatch, never a silent partial load. - if mapped is None or set(mapped.keys()) != target_keys: - got = set(mapped.keys()) if mapped is not None else set() - missing = sorted(target_keys - got) - extra = sorted(got - target_keys) - raise RuntimeError( - "dcp_to_hf: cannot map the DCP state dict onto the backbone " - f"architecture (recorded prefix={prefix!r}). missing={missing[:10]} " - f"extra={extra[:10]}. Refusing to write a partially-loaded HF model." - ) - - # 4. write weights as safetensors. Clone so save_file gets contiguous, - # storage-owning tensors. - os.makedirs(out_dir, exist_ok=True) - mapped = {k: v.contiguous().clone() for k, v in mapped.items()} - save_file(mapped, os.path.join(out_dir, "model.safetensors")) - - # 5. copy the co-located HF config + tokenizer assets. - for fname in _HF_ASSET_FILES: - src = os.path.join(ckpt_dir, fname) - if os.path.exists(src): - shutil.copy(src, os.path.join(out_dir, fname)) - - def ensure_input_tile_for_distributed_embedding() -> None: """Ensure distributed embedding export uses INPUT_TILE=3.""" if not acc_utils.use_distributed_embedding(): diff --git a/tzrec/utils/hf_export_util.py b/tzrec/utils/hf_export_util.py new file mode 100644 index 000000000..c02a2e32d --- /dev/null +++ b/tzrec/utils/hf_export_util.py @@ -0,0 +1,180 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""HuggingFace export for HF-backed models (``GenerativeRecLM`` family). + +Kept out of ``export_util`` (TorchScript/TRT/AOTI) so ``checkpoint_util`` can +call ``write_hf_assets`` without a circular import. +""" + +import json +import os +import shutil +from typing import Dict, Optional, Set + +import torch +from safetensors.torch import save_file +from torch import nn + +from tzrec.utils import checkpoint_util +from tzrec.utils.logging_util import logger + +# Missing files are skipped -- tokenizers emit different subsets. +_HF_ASSET_FILES = ( + "config.json", + "generation_config.json", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + "added_tokens.json", + "special_tokens_map.json", +) + +_HF_EXPORT_META_FILENAME = "hf_export_meta.json" + + +def _unwrap_hf_model(wrapped_model: nn.Module) -> Optional[nn.Module]: + """Walk DMP/TrainWrapper layers down to the model exposing ``hf_backbone``. + + Returns ``None`` if no backbone is found in the chain -- not an HF-backed + model, so callers no-op. + """ + m = wrapped_model + while not hasattr(m, "hf_backbone"): + if hasattr(m, "module"): # DMP / DDP-style wrapper + m = m.module + elif hasattr(m, "model"): # Train/Predict/Script wrapper + m = m.model + else: + return None + return m + + +def write_hf_assets(wrapped_model: nn.Module, save_dir: str) -> None: + """Co-locate the HF config + tokenizer (NO weights) in a checkpoint dir. + + Records the backbone's FQN prefix, read off ``wrapped_model``'s live module + graph, into ``hf_export_meta.json`` so ``dcp_to_hf`` can strip it without + hard-coding a wrapper convention. Rank 0 only -- the dense backbone is + data-parallel-replicated. + """ + if int(os.environ.get("RANK", 0)) != 0: + return + inner = _unwrap_hf_model(wrapped_model) + if inner is None: + return + os.makedirs(save_dir, exist_ok=True) + + backbone = inner.hf_backbone() + backbone.config.save_pretrained(save_dir) + gen_cfg = getattr(backbone, "generation_config", None) + if gen_cfg is not None: + gen_cfg.save_pretrained(save_dir) + inner.hf_tokenizer().save_pretrained(save_dir) + + # named_modules() FQNs carry the DMP prefix that state_dict() strips. + raw_prefix = next( + (n for n, m in wrapped_model.named_modules() if m is backbone), "" + ) + prefix = checkpoint_util._strip_dmp_prefix(raw_prefix) + meta = {"backbone_state_dict_prefix": prefix + ("." if prefix else "")} + with open(os.path.join(save_dir, _HF_EXPORT_META_FILENAME), "w") as f: + json.dump(meta, f, indent=2) + + +def dcp_to_hf(ckpt_dir: str, out_dir: str) -> None: + """Convert a self-contained checkpoint dir to a ``from_pretrained`` HF dir. + + Reads everything from ``ckpt_dir`` -- no live model is built and no weights + are downloaded. Keys that do not map 1:1 onto the architecture in the + co-located ``config.json`` raise rather than write a partial model. + """ + from torch.distributed.checkpoint.state_dict_loader import ( + _load_state_dict_from_keys, + ) + from transformers import AutoConfig, AutoModelForCausalLM + + model_ckpt_path = os.path.join(ckpt_dir, "model") + if not os.path.exists(model_ckpt_path): + raise RuntimeError(f"dcp_to_hf: model DCP dir [{model_ckpt_path}] not exists.") + + # No keys => load every key; non-distributed => full tensors locally. + raw_state: Dict[str, torch.Tensor] = _load_state_dict_from_keys( + checkpoint_id=model_ckpt_path + ) + + meta_path = os.path.join(ckpt_dir, _HF_EXPORT_META_FILENAME) + prefix: Optional[str] = None + if os.path.exists(meta_path): + with open(meta_path, "r") as f: + prefix = json.load(f).get("backbone_state_dict_prefix") + + cfg = AutoConfig.from_pretrained(ckpt_dir) + with torch.device("meta"): + empty = AutoModelForCausalLM.from_config(cfg) + target_keys: Set[str] = set(empty.state_dict().keys()) + tied_keys: Set[str] = set(getattr(empty, "_tied_weights_keys", None) or []) + del empty + + def _strip_recorded_prefix( + state: Dict[str, torch.Tensor], + ) -> Optional[Dict[str, torch.Tensor]]: + """Strip the recorded prefix; None unless it yields an EXACT match.""" + if not prefix: + return None + out = {k[len(prefix) :]: v for k, v in state.items() if k.startswith(prefix)} + return out if set(out) == target_keys else None + + def _derive_by_suffix( + state: Dict[str, torch.Tensor], + ) -> Optional[Dict[str, torch.Tensor]]: + """Each target key is a unique suffix of exactly one DCP key; None if not.""" + out: Dict[str, torch.Tensor] = {} + for tk in target_keys: + matches = [k for k in state if k == tk or k.endswith("." + tk)] + if len(matches) != 1: + return None + out[tk] = state[matches[0]] + return out + + # A stale recorded prefix falls back to prefix-free suffix matching. + mapped = _strip_recorded_prefix(raw_state) + if mapped is None: + if prefix: + logger.warning( + f"dcp_to_hf: recorded prefix [{prefix}] did not map exactly onto " + "the architecture; deriving the backbone prefix by suffix-matching." + ) + mapped = _derive_by_suffix(raw_state) + + if mapped is None or set(mapped.keys()) != target_keys: + got = set(mapped.keys()) if mapped is not None else set() + missing = sorted(target_keys - got) + extra = sorted(got - target_keys) + raise RuntimeError( + "dcp_to_hf: cannot map the DCP state dict onto the backbone " + f"architecture (recorded prefix={prefix!r}). missing={missing[:10]} " + f"extra={extra[:10]}. Refusing to write a partially-loaded HF model." + ) + + # Tied heads are dropped after validation; from_pretrained re-ties them. + if getattr(cfg, "tie_word_embeddings", False): + mapped = {k: v for k, v in mapped.items() if k not in tied_keys} + mapped = {k: v.contiguous() for k, v in mapped.items()} # save_file rejects views + del raw_state + os.makedirs(out_dir, exist_ok=True) + save_file(mapped, os.path.join(out_dir, "model.safetensors")) + + for fname in _HF_ASSET_FILES: + src = os.path.join(ckpt_dir, fname) + if os.path.exists(src): + shutil.copy(src, os.path.join(out_dir, fname)) diff --git a/tzrec/utils/hf_export_util_test.py b/tzrec/utils/hf_export_util_test.py new file mode 100644 index 000000000..76e150fec --- /dev/null +++ b/tzrec/utils/hf_export_util_test.py @@ -0,0 +1,175 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +import shutil +import unittest + +import torch +from safetensors.torch import load_file +from torch import nn + +from tzrec.utils.checkpoint_util import save_model +from tzrec.utils.hf_export_util import ( + _HF_EXPORT_META_FILENAME, + _unwrap_hf_model, + dcp_to_hf, + write_hf_assets, +) +from tzrec.utils.test_util import make_test_dir + + +class _FakeTokenizer: + """Writes the two tokenizer asset files `write_hf_assets` copies.""" + + def save_pretrained(self, save_dir): + for name in ("tokenizer.json", "tokenizer_config.json"): + with open(os.path.join(save_dir, name), "w") as f: + f.write("{}") + + +class _GenRec(nn.Module): + """Stand-in for GenerativeRecLM: an HF backbone plus unrelated params.""" + + def __init__(self, lm): + super().__init__() + self.lm = lm + self.other = nn.Linear(4, 4) + + def hf_backbone(self): + return self.lm + + def hf_tokenizer(self): + return _FakeTokenizer() + + +class _TrainWrapper(nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + + +class _DmpLike(nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + +def _tiny_lm(tie=True): + from transformers import AutoModelForCausalLM, Qwen2Config + + cfg = Qwen2Config( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + tie_word_embeddings=tie, + max_position_embeddings=128, + ) + return AutoModelForCausalLM.from_config(cfg, torch_dtype=torch.float32) + + +class HfExportUtilTest(unittest.TestCase): + def setUp(self) -> None: + self.test_dir = make_test_dir() + os.environ.setdefault("RANK", "0") + + def tearDown(self) -> None: + shutil.rmtree(self.test_dir, ignore_errors=True) + + def test_unwrap_walks_dmp_and_train_wrapper(self) -> None: + inner = _GenRec(_tiny_lm()) + self.assertIs(_unwrap_hf_model(inner), inner) + self.assertIs(_unwrap_hf_model(_TrainWrapper(inner)), inner) + self.assertIs(_unwrap_hf_model(_DmpLike(_TrainWrapper(inner))), inner) + + def test_unwrap_returns_none_for_non_hf_model(self) -> None: + self.assertIsNone(_unwrap_hf_model(_TrainWrapper(nn.Linear(4, 4)))) + + def test_write_hf_assets_noop_for_non_hf_model(self) -> None: + save_dir = os.path.join(self.test_dir, "plain") + write_hf_assets(_TrainWrapper(nn.Linear(4, 4)), save_dir) + self.assertFalse(os.path.exists(save_dir)) + + def _save_ckpt(self, wrapped, name="model.ckpt-1"): + ckpt_dir = os.path.join(self.test_dir, name) + save_model(ckpt_dir, wrapped) + write_hf_assets(wrapped, ckpt_dir) + return ckpt_dir + + def test_write_hf_assets_records_state_dict_prefix(self) -> None: + lm = _tiny_lm() + wrapped = _TrainWrapper(_GenRec(lm)) + ckpt_dir = self._save_ckpt(wrapped) + for name in ("config.json", "tokenizer.json", _HF_EXPORT_META_FILENAME): + self.assertTrue(os.path.exists(os.path.join(ckpt_dir, name)), name) + with open(os.path.join(ckpt_dir, _HF_EXPORT_META_FILENAME)) as f: + prefix = json.load(f)["backbone_state_dict_prefix"] + self.assertEqual(prefix, "model.lm.") + # the prefix must reconstruct the exact FQNs save_model wrote + saved = set(wrapped.state_dict()) + self.assertTrue(all(prefix + k in saved for k in lm.state_dict())) + + def test_dcp_to_hf_round_trip_drops_tied_head(self) -> None: + from transformers import AutoModelForCausalLM + + lm = _tiny_lm(tie=True) + ckpt_dir = self._save_ckpt(_DmpLike(_TrainWrapper(_GenRec(lm)))) + out_dir = os.path.join(self.test_dir, "hf_out") + dcp_to_hf(ckpt_dir, out_dir) + + st = load_file(os.path.join(out_dir, "model.safetensors")) + self.assertNotIn("lm_head.weight", st) + self.assertIn("model.embed_tokens.weight", st) + back = AutoModelForCausalLM.from_pretrained(out_dir) + self.assertEqual( + back.lm_head.weight.data_ptr(), back.model.embed_tokens.weight.data_ptr() + ) + for k, v in lm.state_dict().items(): + self.assertTrue(torch.equal(back.state_dict()[k], v), k) + + def test_dcp_to_hf_self_heals_a_stale_prefix(self) -> None: + lm = _tiny_lm() + ckpt_dir = self._save_ckpt(_TrainWrapper(_GenRec(lm))) + meta_path = os.path.join(ckpt_dir, _HF_EXPORT_META_FILENAME) + with open(meta_path, "w") as f: + json.dump({"backbone_state_dict_prefix": "bogus.wrapper."}, f) + out_dir = os.path.join(self.test_dir, "hf_out_stale") + dcp_to_hf(ckpt_dir, out_dir) # falls back to suffix matching + st = load_file(os.path.join(out_dir, "model.safetensors")) + self.assertTrue( + torch.equal(st["model.embed_tokens.weight"], lm.model.embed_tokens.weight) + ) + + def test_dcp_to_hf_refuses_a_mismatched_architecture(self) -> None: + ckpt_dir = self._save_ckpt(_TrainWrapper(_GenRec(_tiny_lm()))) + # widen the recorded architecture so the checkpoint can no longer fill it + cfg_path = os.path.join(ckpt_dir, "config.json") + with open(cfg_path) as f: + cfg = json.load(f) + cfg["num_hidden_layers"] = 4 + with open(cfg_path, "w") as f: + json.dump(cfg, f) + with self.assertRaisesRegex(RuntimeError, "Refusing to write"): + dcp_to_hf(ckpt_dir, os.path.join(self.test_dir, "hf_out_bad")) + + def test_dcp_to_hf_missing_dcp_dir(self) -> None: + empty = os.path.join(self.test_dir, "no_dcp") + os.makedirs(empty, exist_ok=True) + with self.assertRaisesRegex(RuntimeError, "not exists"): + dcp_to_hf(empty, os.path.join(self.test_dir, "hf_out_missing")) + + +if __name__ == "__main__": + unittest.main() From 848ed880a4de2afddc49b1dfc7e94d26c7019374 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Mon, 27 Jul 2026 12:22:20 +0000 Subject: [PATCH 35/99] [refactor] genrec LM: test the beam kernel directly and drop test duplication escalating_beam_search is a reusable, torch-only kernel but had no test beside it: its only coverage reached it through Qwen2RecLM, so the schedule, the KV-cache reorder and the left-padding arithmetic were only ever exercised incidentally. Mutating the production code showed the cost -- ten of thirteen mutations, including a misordered cache reorder, a flat width schedule, cross-row beam mixing and a non-deterministic decode, passed the old suite untouched. It now has its own test file that imports nothing from tzrec.models. A duck-typed backbone records the row count of each forward, which makes the per-level beam width observable without exposing it from the kernel, and the exhaustive comparison against a full-recompute ranking runs at three levels so the cache reorder is on a verified path. The Qwen2RecLM tests keep only what belongs to the wrapper: that it forwards the right SID bands, and that band restriction leaves no rejected candidate behind. The three generate tests collapse into one parameterized table that also asserts the arguments handed to lm.generate, the recency-clip test moves next to the base-class method it covers, and the repo's only cross-test-module import goes away with the unused batch fixture that motivated it. None of the production code changes; the same thirteen mutations now all fail. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/generative_rec_lm_test.py | 37 ++ tzrec/models/qwen2_rec_lm_test.py | 464 ++++++++++--------------- tzrec/modules/escalating_beam_test.py | 204 +++++++++++ tzrec/utils/hf_export_util_test.py | 6 +- 4 files changed, 437 insertions(+), 274 deletions(-) create mode 100644 tzrec/modules/escalating_beam_test.py diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py index eb0f1ddc1..80154a760 100644 --- a/tzrec/models/generative_rec_lm_test.py +++ b/tzrec/models/generative_rec_lm_test.py @@ -336,6 +336,43 @@ def test_sid_token_rows_rejects_out_of_range_codes(self) -> None: jt = _FakeJT(values, [3]) m._sid_token_rows(jt.values(), jt.lengths()) + def test_sid_token_rows_recency_clip(self) -> None: + m = _stub(base_vocab=100) + items = [(1, 1, 1), (2, 3, 4), (1, 2, 3), (2, 1, 4), (1, 3, 2)] + toks = [ + [100, 102, 105], + [101, 104, 108], + [100, 103, 107], + [101, 102, 108], + [100, 104, 106], + ] + values = torch.tensor(items, dtype=torch.float).flatten() + lengths = torch.tensor([values.numel()]) + + # 15 codes (5 items), cap 9 -> keep the most recent three whole items. + rows = m._sid_token_rows(values, lengths, max_codes=9) + self.assertEqual(rows[0].tolist(), sum(toks[2:], [])) + # item-aligned: cap 10 still keeps 9 (3 whole items), never cuts mid-item + rows = m._sid_token_rows(values, lengths, max_codes=10) + self.assertEqual(rows[0].tolist(), sum(toks[2:], [])) + # within cap -> untouched + rows = m._sid_token_rows(values[:6], torch.tensor([6]), max_codes=9) + self.assertEqual(rows[0].tolist(), sum(toks[:2], [])) + # disabled (0/None) -> no clip + rows = m._sid_token_rows(values, lengths, max_codes=0) + self.assertEqual(rows[0].tolist(), sum(toks, [])) + + def test_validate_sid_candidates_groups_batch_major(self) -> None: + m = _stub(base_vocab=100) + # decoders emit rows batch-major: [b0_c0, b0_c1, b1_c0, b1_c1] + tokens = torch.tensor( + [[100, 102, 105], [101, 104, 108], [101, 103, 106], [100, 104, 107]] + ) + sids = m._validate_sid_candidates(tokens, batch_size=2) + self.assertEqual(tuple(sids.shape), (2, 2, 3)) + self.assertEqual(sids[0].tolist(), [[1, 1, 1], [2, 3, 4]]) + self.assertEqual(sids[1].tolist(), [[2, 2, 2], [1, 3, 3]]) + def test_build_input_history_group_label_field(self) -> None: m = _stub(base_vocab=100) m._input_name, m._label_name = "user_sequence", "label" diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/qwen2_rec_lm_test.py index d4778fdd7..1231fecee 100644 --- a/tzrec/models/qwen2_rec_lm_test.py +++ b/tzrec/models/qwen2_rec_lm_test.py @@ -11,12 +11,14 @@ import types import unittest +from unittest import mock import torch +from parameterized import parameterized from torch import nn -from tzrec.models.generative_rec_lm_test import _FakeJT from tzrec.models.qwen2_rec_lm import Qwen2RecLM +from tzrec.utils.test_util import parameterized_name_func def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): @@ -51,17 +53,56 @@ def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): return m -def _gen_batch(): - """An opaque one-row batch; the ``_generate`` tests mock ``build_input``.""" - return types.SimpleNamespace( - sequence_dense_features={"user_sequence": _FakeJT([1, 2, 3], [3])} +def _real_lm_stub(codebook=None, base_vocab=20, num_beams=2): + """A Qwen2RecLM carrying a real (tiny, random) Qwen2 backbone. + + Needed wherever the real forward runs: the training objective and the + end-to-end escalating-beam decode; the other tests mock ``lm.generate``. + """ + from transformers import Qwen2Config, Qwen2ForCausalLM + + codebook = codebook or [2, 3, 4] + m = object.__new__(Qwen2RecLM) + nn.Module.__init__(m) + m._num_levels = len(codebook) + m._base_vocab = base_vocab + m._num_beams = num_beams + cfg = Qwen2Config( + vocab_size=base_vocab + sum(codebook), + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=64, ) + torch.manual_seed(0) + m.lm = Qwen2ForCausalLM(cfg).eval() + offsets, sizes = Qwen2RecLM._sid_level_layout(codebook) + m.register_buffer("_level_offsets", offsets, persistent=False) + m.register_buffer("_codebook_sizes", sizes, persistent=False) + return m + + +def _train_stub(max_total_len): + """A ``_predict_train``-ready stub; returns ``(model, spliced T per step)``.""" + m = _stub() + m._is_inference = False # not inference + nn.Module.training=True -> is_train + m._input_name, m._label_name = "user_sequence", "label" + m._max_total_len = max_total_len + m._pool_warmed = False + seen_lens = [] + def fwd(input_ids, labels, attention_mask): + seen_lens.append(input_ids.shape[1]) + return {"loss": torch.tensor(0.0)} -def _first_non_neg_index(labels): - """First supervised label column, minimized over rows.""" - tmp = (labels >= 0).cumsum(dim=-1) - return int((tmp == 1).float().argmax(dim=-1).min().item()) + m.build_input = lambda b: { + m._input_name: [torch.tensor([100, 101, 102])], + m._label_name: [torch.tensor([200, 201, 202])], + } + m._forward_loss = fwd + return m, seen_lens class Qwen2RecLMTest(unittest.TestCase): @@ -105,14 +146,14 @@ def test_mask_keeps_trailing_eos_when_pad_equals_eos(self) -> None: def test_suffix_keep_matches_dynamic_slice(self) -> None: # the constant suffix width must cover every supervised column. - m = _stub() # num_levels=3, asst_suffix=[15] (numel 1) -> suffix_keep=6 - suffix_keep = m._num_levels + m.tpl_asst_suffix.numel() + 2 + m = _stub() + suffix_keep = 6 # num_levels 3 + asst_suffix 1 + trailing eos + HF shift _, labels, _ = m._splice_input_ids( [torch.tensor([100, 101, 102])], [torch.tensor([200, 201, 202])] ) - self.assertEqual( - suffix_keep, labels.shape[1] - _first_non_neg_index(labels) + 1 - ) + # first supervised column, minimized over rows + first_sup = int(((labels >= 0).cumsum(-1) == 1).float().argmax(-1).min()) + self.assertEqual(suffix_keep, labels.shape[1] - first_sup + 1) self.assertTrue(bool((labels[:, :-suffix_keep] < 0).all())) def test_splice_prompt_ids(self) -> None: @@ -131,49 +172,12 @@ def test_predict_routes_on_inference_flag(self) -> None: m._is_inference = True # inference self.assertEqual(Qwen2RecLM.predict(m, object())["branch"], "generate") - def test_generate_maps_tokens_to_sids(self) -> None: - m = _stub(base_vocab=100) - m._input_name = "user_sequence" - m._num_beams = m._num_return = 2 - - def fake_generate( - input_ids, - attention_mask, - max_new_tokens, - num_beams, - num_return_sequences, - do_sample, - pad_token_id, - ): - prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) - # offsets [0,2,5]: local [1,1,1]/[2,3,4] -> each level's min/max token - new = torch.tensor([[100, 102, 105], [101, 104, 108]]) - return torch.cat([prompt, new], dim=1) - - m.lm.generate = fake_generate - # build_input is mocked, so the batch is opaque to this generation test. - m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} - sids = m._generate(_gen_batch())["generated_sids"] - self.assertEqual(tuple(sids.shape), (1, 2, 3)) # (B, num_return, num_levels) - self.assertEqual(sids[0].tolist(), [[1, 1, 1], [2, 3, 4]]) - - def test_generate_rejects_malformed_candidates(self) -> None: - # every malformed candidate collapses to the -1 sentinel, in place. - m = _stub(base_vocab=100) - m._input_name = "user_sequence" - m._num_beams = m._num_return = 6 - - def fake_generate( - input_ids, - attention_mask, - max_new_tokens, - num_beams, - num_return_sequences, - do_sample, - pad_token_id, - ): - prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) - new = torch.tensor( + # (generated tail -> decoded SIDs). offsets [0,2,5]: local [1,1,1]/[2,3,4] + # are each level's min/max token; every malformed row collapses to -1. + @parameterized.expand( + [ + [[[100, 102, 105], [101, 104, 108]], [[1, 1, 1], [2, 3, 4]]], + [ [ [100, 102, 105], # valid -> local [1, 1, 1] [101, 104, 108], # valid -> local [2, 3, 4] @@ -181,31 +185,20 @@ def fake_generate( [100, 101, 105], # pos1 below level-1 band [100, 102, 109], # pos2 above level-2 band [100, 104, 9], # pos2 = eos/pad token (sid -90) -> invalid - ] - ) - return torch.cat([prompt, new], dim=1) - - m.lm.generate = fake_generate - m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} - sids = m._generate(_gen_batch())["generated_sids"] - self.assertEqual(tuple(sids.shape), (1, 6, 3)) - self.assertEqual( - sids[0].tolist(), - [ - [1, 1, 1], - [2, 3, 4], - [-1, -1, -1], - [-1, -1, -1], - [-1, -1, -1], - [-1, -1, -1], + ], + [[1, 1, 1], [2, 3, 4]] + [[-1, -1, -1]] * 4, ], - ) - - def test_generate_narrow_tail_no_crash(self) -> None: - # early EOS -> a tail narrower than num_levels must still reshape cleanly + # early EOS: a tail narrower than num_levels still reshapes cleanly, + # and the missing 3rd atom stays -1 -> out of band -> candidate -1 + [[[100, 102], [101, 104]], [[-1, -1, -1], [-1, -1, -1]]], + ], + name_func=parameterized_name_func, + ) + def test_generate_maps_tokens_to_sids(self, tail, expected) -> None: m = _stub(base_vocab=100) m._input_name = "user_sequence" - m._num_beams = m._num_return = 2 + m._num_beams, m._num_return = len(tail) + 1, len(tail) + seen = {} def fake_generate( input_ids, @@ -216,28 +209,67 @@ def fake_generate( do_sample, pad_token_id, ): + seen.update( + prompt=input_ids[0].tolist(), + mask=attention_mask[0].tolist(), + max_new_tokens=max_new_tokens, + num_beams=num_beams, + num_return_sequences=num_return_sequences, + do_sample=do_sample, + pad_token_id=pad_token_id, + ) prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) - new = torch.tensor([[100, 102], [101, 104]]) # width 2 < num_levels 3 - return torch.cat([prompt, new], dim=1) + return torch.cat([prompt, torch.tensor(tail)], dim=1) m.lm.generate = fake_generate + # build_input is mocked, so the batch is opaque to this generation test. + m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} + sids = m._generate(object())["generated_sids"] + self.assertEqual(seen["prompt"], [10, 11, 12, 100, 102, 105, 13, 14]) + self.assertEqual(seen["mask"], [1] * 8) + self.assertEqual(seen["max_new_tokens"], 3) # = num_levels + self.assertEqual(seen["num_beams"], len(tail) + 1) + self.assertEqual(seen["num_return_sequences"], len(tail)) + self.assertFalse(seen["do_sample"]) + self.assertEqual(seen["pad_token_id"], 9) + # (B, num_return, num_levels) + self.assertEqual(tuple(sids.shape), (1, len(tail), 3)) + self.assertEqual(sids[0].tolist(), expected) + + def test_generate_routes_to_the_dynamic_beam(self) -> None: + m = _stub(base_vocab=100) + m._input_name = "user_sequence" + m._dynamic_beam = True + m._num_beams = m._num_return = 2 + m.lm.generate = lambda **kw: self.fail("HF generate must not run") + seen = {} + + def fake_beam(ids, am): + seen.update(ids=ids[0].tolist(), mask=am[0].tolist()) + return torch.tensor([[100, 102, 105], [101, 104, 108]]) + + m._dynamic_beam_search = fake_beam m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} - sids = m._generate(_gen_batch())["generated_sids"] + sids = m._generate(object())["generated_sids"] + self.assertEqual(seen["ids"], [10, 11, 12, 100, 102, 105, 13, 14]) + self.assertEqual(seen["mask"], [1] * 8) self.assertEqual(tuple(sids.shape), (1, 2, 3)) - # the missing 3rd atom stays -1 -> out of band -> whole candidate -1 - self.assertEqual(sids[0].tolist(), [[-1, -1, -1], [-1, -1, -1]]) + self.assertEqual(sids[0].tolist(), [[1, 1, 1], [2, 3, 4]]) - def test_build_prompt_tokens_registers_buffers(self) -> None: - m = object.__new__(Qwen2RecLM) - nn.Module.__init__(m) - tok = types.SimpleNamespace( + def _prompt_tokenizer(self): + # encode -> [len(text)] makes each buffer a fingerprint of its fragment + return types.SimpleNamespace( eos_token_id=99, encode=lambda text, add_special_tokens=False: [len(text)], ) + + def test_build_prompt_tokens_composes_the_family_template(self) -> None: + m = object.__new__(Qwen2RecLM) + nn.Module.__init__(m) cfg = types.SimpleNamespace( system_instruction="", user_prefix_text="", user_suffix_text="" ) - m._build_prompt_tokens(tok, cfg) + m._build_prompt_tokens(self._prompt_tokenizer(), cfg) for name in [ "tpl_system", "tpl_user_prefix", @@ -249,66 +281,37 @@ def test_build_prompt_tokens_registers_buffers(self) -> None: buf = getattr(m, name) self.assertIsInstance(buf, torch.Tensor) self.assertEqual(buf.dtype, torch.int64) - self.assertEqual(m.tpl_eos.tolist(), [99]) # eos cached for supervision - - def test_sid_token_rows_recency_clip(self) -> None: - m = _stub(base_vocab=100) - values = torch.tensor( + tpl = Qwen2RecLM.CHAT_TEMPLATE + self.assertEqual( + m.tpl_system.tolist(), [ - 1, - 1, - 1, - 2, - 3, - 4, - 1, - 2, - 3, - 2, - 1, - 4, - 1, - 3, - 2, + len( + tpl["system_prefix"] + + tpl["default_system_instruction"] + + tpl["system_suffix"] + ) ], - dtype=torch.float, ) + self.assertEqual(m.tpl_user_prefix.tolist(), [len(tpl["user_prefix"])]) + self.assertEqual(m.tpl_user_suffix.tolist(), [len(tpl["user_suffix"])]) + self.assertEqual(m.tpl_asst_prefix.tolist(), [len(tpl["asst_prefix"])]) + self.assertEqual(m.tpl_asst_suffix.tolist(), [len(tpl["asst_suffix"])]) + self.assertEqual(m.tpl_eos.tolist(), [99]) # eos cached for supervision - def _vl(): - return values, torch.tensor([values.numel()]) - - # 15 codes (5 items), cap 9 -> keep the most recent three whole items. - expected_tail = [100, 103, 107, 101, 102, 108, 100, 104, 106] - rows = m._sid_token_rows(*_vl(), max_codes=9) - self.assertEqual(rows[0].tolist(), expected_tail) - # item-aligned: cap 10 still keeps 9 (3 whole items), never cuts mid-item - rows = m._sid_token_rows(*_vl(), max_codes=10) - self.assertEqual(rows[0].tolist(), expected_tail) - # within cap -> untouched - rows = m._sid_token_rows(values[:6], torch.tensor([6]), max_codes=9) - self.assertEqual(rows[0].tolist(), [100, 102, 105, 101, 104, 108]) - # disabled (0/None) -> no clip - rows = m._sid_token_rows(*_vl(), max_codes=0) + def test_build_prompt_tokens_honours_proto_text_knobs(self) -> None: + m = object.__new__(Qwen2RecLM) + nn.Module.__init__(m) + cfg = types.SimpleNamespace( + system_instruction="SYS", user_prefix_text="UP", user_suffix_text="US" + ) + m._build_prompt_tokens(self._prompt_tokenizer(), cfg) + tpl = Qwen2RecLM.CHAT_TEMPLATE self.assertEqual( - rows[0].tolist(), - [ - 100, - 102, - 105, - 101, - 104, - 108, - 100, - 103, - 107, - 101, - 102, - 108, - 100, - 104, - 106, - ], + m.tpl_system.tolist(), + [len(tpl["system_prefix"] + "SYS" + tpl["system_suffix"])], ) + self.assertEqual(m.tpl_user_prefix.tolist(), [len(tpl["user_prefix"] + "UP")]) + self.assertEqual(m.tpl_user_suffix.tolist(), [len("US" + tpl["user_suffix"])]) def test_compute_max_total_length(self) -> None: m = _stub() @@ -319,48 +322,16 @@ def test_compute_max_total_length(self) -> None: self.assertEqual(m._compute_max_total_length(), 0) def test_first_step_pads_to_max_then_actual_length(self) -> None: - m = _stub() - m._is_inference = False # not inference + nn.Module.training=True -> is_train - m._input_name, m._label_name = "user_sequence", "label" - m._max_total_len = 50 - m._pool_warmed = False - seen_lens = [] - - def fwd(i, lbl, a): - seen_lens.append(i.shape[1]) - return {"loss": torch.tensor(0.0)} - - m.build_input = lambda b: { - m._input_name: [torch.tensor([100, 101, 102])], - m._label_name: [torch.tensor([200, 201, 202])], - } - m._forward_loss = fwd - batch = object() - m._predict_train(batch) # first step: pre-size to worst case - m._predict_train(batch) # subsequent step: natural length + m, seen_lens = _train_stub(max_total_len=50) + m._predict_train(object()) # first step: pre-size to worst case + m._predict_train(object()) # subsequent step: natural length self.assertEqual(seen_lens[0], 50) self.assertLess(seen_lens[1], 50) self.assertTrue(m._pool_warmed) def test_no_forced_padding_when_disabled(self) -> None: - m = _stub() - m._is_inference = False - m._input_name, m._label_name = "user_sequence", "label" - m._max_total_len = 0 # max_seq_length unset -> pre-allocation off - m._pool_warmed = False - seen_lens = [] - - def fwd(i, lbl, a): - seen_lens.append(i.shape[1]) - return {"loss": torch.tensor(0.0)} - - m.build_input = lambda b: { - m._input_name: [torch.tensor([100, 101, 102])], - m._label_name: [torch.tensor([200, 201, 202])], - } - m._forward_loss = fwd - batch = object() - m._predict_train(batch) + m, seen_lens = _train_stub(max_total_len=0) # pre-allocation off + m._predict_train(object()) self.assertLess(seen_lens[0], 50) self.assertFalse(m._pool_warmed) @@ -388,7 +359,7 @@ def _model(self, ignore_index=-100): m.register_buffer( name, torch.tensor(vals, dtype=torch.long), persistent=False ) - m._suffix_keep = m._num_levels + m.tpl_asst_suffix.numel() + 2 + m._suffix_keep = 6 # num_levels 3 + asst_suffix 1 + trailing eos + HF shift return m def _rows(self): @@ -439,108 +410,55 @@ def test_forward_loss_returns_only_the_loss(self) -> None: self.assertEqual(list(out), ["loss"]) -def _real_lm_stub(codebook=None, base_vocab=20, num_beams=2): - """A Qwen2RecLM carrying a real (tiny, random) Qwen2 backbone. - - Needed by the dynamic-beam tests, which exercise the real KV-cached forward - and cache-reorder path; the other tests mock ``lm.generate``. - """ - from transformers import Qwen2Config, Qwen2ForCausalLM - - codebook = codebook or [2, 3, 4] - m = object.__new__(Qwen2RecLM) - nn.Module.__init__(m) - m._num_levels = len(codebook) - m._base_vocab = base_vocab - m._num_beams = num_beams - cfg = Qwen2Config( - vocab_size=base_vocab + sum(codebook), - hidden_size=32, - intermediate_size=64, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - max_position_embeddings=64, - ) - torch.manual_seed(0) - m.lm = Qwen2ForCausalLM(cfg).eval() - offsets, sizes = Qwen2RecLM._sid_level_layout(codebook) - m.register_buffer("_level_offsets", offsets, persistent=False) - m.register_buffer("_codebook_sizes", sizes, persistent=False) - return m - - class Qwen2DynamicBeamTest(unittest.TestCase): - def test_width_schedule_and_final_count(self) -> None: - m = _real_lm_stub(codebook=[8, 7, 6], base_vocab=20, num_beams=2) - ids = torch.tensor([[1, 2, 3, 4]]) - new = m._dynamic_beam_search(ids, torch.ones_like(ids)) - # base 2: widths [4, 8, 16], with enough combinations at each level. - self.assertEqual(tuple(new.shape), (2 * 2**3, 3)) + """The wrapper around ``escalating_beam_search``; the kernel has its own test.""" + + def test_dynamic_beam_search_forwards_the_sid_bands(self) -> None: + m = _stub(base_vocab=100) + m._num_beams = 5 + ids = torch.zeros(1, 4, dtype=torch.long) + am = torch.ones(1, 4, dtype=torch.long) + seen = {} + + def fake_kernel(lm, input_ids, attention_mask, *, num_beams, lo_tok, hi_tok): + seen.update( + lm=lm, + input_ids=input_ids, + attention_mask=attention_mask, + num_beams=num_beams, + lo=lo_tok.tolist(), + hi=hi_tok.tolist(), + ) + return torch.zeros(2, 3, dtype=torch.long) - def test_every_candidate_is_in_band(self) -> None: - # band masking guarantees well-formed SIDs: validate -> no -1 sentinels. + with mock.patch( + "tzrec.models.qwen2_rec_lm.escalating_beam_search", side_effect=fake_kernel + ): + out = m._dynamic_beam_search(ids, am) + self.assertIs(seen["lm"], m.lm) + self.assertIs(seen["input_ids"], ids) + self.assertIs(seen["attention_mask"], am) + self.assertEqual(seen["num_beams"], 5) + self.assertEqual(seen["lo"], [100, 102, 105]) + self.assertEqual(seen["hi"], [101, 104, 108]) + self.assertEqual(tuple(out.shape), (2, 3)) + + def test_band_masked_beams_decode_without_sentinels(self) -> None: + # end-to-end against a real backbone: band masking guarantees that every + # returned candidate survives _validate_sid_candidates. widths are + # [2, 6, 24] here, i.e. exhaustive over the whole 2*3*4 codebook. codebook = [2, 3, 4] - m = _real_lm_stub(codebook=codebook, base_vocab=20, num_beams=2) - lo, hi = m._sid_token_bands() - self.assertEqual(lo.tolist(), [20, 22, 25]) - self.assertEqual(hi.tolist(), [21, 24, 28]) + m = _real_lm_stub(codebook=codebook, base_vocab=20, num_beams=3) ids = torch.tensor([[5, 6, 7]]) new = m._dynamic_beam_search(ids, torch.ones_like(ids)) sids = m._validate_sid_candidates(new, batch_size=1) - self.assertEqual(tuple(sids.shape), (1, new.shape[0], 3)) - for level, size in enumerate(codebook): - self.assertTrue(bool((sids[..., level] >= 1).all())) - self.assertTrue(bool((sids[..., level] <= size).all())) - - def test_left_padding_two_rows(self) -> None: - # ragged batch (row 1 left-padded): both rows yield valid, full beam sets. - codebook = [2, 3] - m = _real_lm_stub(codebook=codebook, base_vocab=20, num_beams=2) - ids = torch.tensor([[5, 6, 7, 8], [0, 0, 9, 10]]) - am = torch.tensor([[1, 1, 1, 1], [0, 0, 1, 1]]) - new = m._dynamic_beam_search(ids, am) - # The requested width is 8, but only 2*3=6 distinct pairs exist. - self.assertEqual(tuple(new.shape), (2 * 6, 2)) - sids = m._validate_sid_candidates(new, batch_size=2) - self.assertEqual(tuple(sids.shape), (2, 6, 2)) + self.assertEqual(tuple(sids.shape), (1, 24, 3)) for level, size in enumerate(codebook): self.assertTrue(bool((sids[..., level] >= 1).all())) self.assertTrue(bool((sids[..., level] <= size).all())) - - def test_exhaustive_matches_bruteforce_topk(self) -> None: - # with no pruning the beam is EXACT: every SID combo, ordered by the - # true (full-recompute) cumulative log-prob. - codebook, base = [2, 3], 20 - m = _real_lm_stub(codebook=codebook, base_vocab=base, num_beams=3) - # widths [min(6,2)=2, min(12,6)=6] -> exhaustive over all 2*3 SIDs - ids = torch.tensor([[5, 6, 7, 8]]) - am = torch.ones_like(ids) - got = [tuple(r) for r in m._dynamic_beam_search(ids, am).tolist()] - self.assertEqual(len(got), codebook[0] * codebook[1]) - lm = m.lm - lo0, lo1 = base, base + codebook[0] - ref = {} - with torch.no_grad(): - logp0 = torch.log_softmax( - lm(ids, attention_mask=am).logits[0, -1].float(), -1 - ) - for t0 in range(lo0, lo0 + codebook[0]): - s2 = torch.cat([ids, torch.tensor([[t0]])], 1) - logp1 = torch.log_softmax( - lm(s2, attention_mask=torch.ones_like(s2)).logits[0, -1].float(), -1 - ) - for t1 in range(lo1, lo1 + codebook[1]): - ref[(t0, t1)] = (logp0[t0] + logp1[t1]).item() - self.assertEqual(set(got), set(ref)) - # ordered best-first by the true score (tolerant of float-noise ties) - s = [ref[c] for c in got] - self.assertTrue(all(s[i] >= s[i + 1] - 1e-4 for i in range(len(s) - 1))) - self.assertEqual(got[0], max(ref, key=ref.get)) # top-1 is the global best - decoded = m._validate_sid_candidates(torch.tensor(got), batch_size=1)[0] self.assertEqual( - {tuple(row) for row in decoded.tolist()}, - {(a, b) for a in range(1, 3) for b in range(1, 4)}, + {tuple(row) for row in sids[0].tolist()}, + {(a, b, c) for a in range(1, 3) for b in range(1, 4) for c in range(1, 5)}, ) diff --git a/tzrec/modules/escalating_beam_test.py b/tzrec/modules/escalating_beam_test.py new file mode 100644 index 000000000..7e3af1010 --- /dev/null +++ b/tzrec/modules/escalating_beam_test.py @@ -0,0 +1,204 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import ast +import itertools +import pathlib +import unittest +from typing import Any, Dict, List, Tuple + +import torch +from parameterized import parameterized + +from tzrec.modules import escalating_beam +from tzrec.modules.escalating_beam import escalating_beam_search +from tzrec.utils.test_util import parameterized_name_func + + +def _tiny_lm(vocab_size, seed=0): + from transformers import Qwen2Config, Qwen2ForCausalLM + + cfg = Qwen2Config( + vocab_size=vocab_size, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=64, + ) + torch.manual_seed(seed) + return Qwen2ForCausalLM(cfg).eval() + + +def _bands(pairs): + lo = torch.tensor([p[0] for p in pairs], dtype=torch.long) + hi = torch.tensor([p[1] for p in pairs], dtype=torch.long) + return lo, hi + + +class _RowSpy: + """Duck-typed backbone recording the row count of every ``.model`` call. + + The kernel touches only ``.model`` and ``.lm_head``, so the per-level beam + width -- otherwise a local inside the kernel -- becomes observable. + """ + + def __init__(self, lm) -> None: + self.lm_head = lm.lm_head + self.rows: List[int] = [] + self._lm = lm + + def model(self, **kwargs: Any) -> Any: + self.rows.append(kwargs["input_ids"].shape[0]) + return self._lm.model(**kwargs) + + +def _bruteforce_scores(lm, input_ids, attention_mask, pairs): + """Exact cumulative log-prob of every band combination, by full recompute.""" + ref: Dict[Tuple[int, ...], float] = {} + with torch.no_grad(): + for combo in itertools.product(*[range(lo, hi + 1) for lo, hi in pairs]): + seq, mask, total = input_ids, attention_mask, 0.0 + for tok in combo: + logits = lm(input_ids=seq, attention_mask=mask).logits[0, -1].float() + total += float(torch.log_softmax(logits, dim=-1)[tok]) + seq = torch.cat([seq, torch.tensor([[tok]])], dim=1) + mask = torch.cat([mask, mask.new_ones(1, 1)], dim=1) + ref[combo] = total + return ref + + +class EscalatingBeamSearchTest(unittest.TestCase): + def test_module_declares_no_tzrec_imports(self) -> None: + # the "torch-only, no tzrec deps" docstring is what makes it liftable. + tree = ast.parse(pathlib.Path(escalating_beam.__file__).read_text()) + mods = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + mods.update(a.name for a in node.names) + elif isinstance(node, ast.ImportFrom) and node.module: + mods.add(node.module) + self.assertEqual( + {m for m in mods if m == "tzrec" or m.startswith("tzrec.")}, set() + ) + + @parameterized.expand( + [ + # uncapped doubling: 2 -> 4 -> 8 -> 16 + [2, [(20, 27), (28, 34), (35, 40)], [4, 8, 16]], + # capped by band x surviving prefixes, not by the doubling + [3, [(20, 21), (22, 24)], [2, 6]], + [1, [(20, 21), (22, 24), (25, 28)], [2, 4, 8]], + # a width-1 band collapses level 0 to a single beam + [2, [(20, 20), (21, 23)], [1, 3]], + ], + name_func=parameterized_name_func, + ) + def test_width_schedule(self, num_beams, pairs, expected_widths) -> None: + spy = _RowSpy(_tiny_lm(vocab_size=48)) + lo, hi = _bands(pairs) + ids = torch.tensor([[5, 6, 7, 8]]) + out = escalating_beam_search( + spy, ids, torch.ones_like(ids), num_beams=num_beams, lo_tok=lo, hi_tok=hi + ) + # calls: [prompt (1 row)] + one per level>0, each carrying widths[j-1] rows + widths = spy.rows[1:] + [out.shape[0]] + self.assertEqual(spy.rows[0], 1) + self.assertEqual(widths, expected_widths) + self.assertEqual(tuple(out.shape), (expected_widths[-1], len(pairs))) + + def test_tokens_stay_inside_arbitrary_bands(self) -> None: + # bands the Qwen2RecLM caller can never produce: descending, disjoint, + # unequal width -- the kernel's contract is per-level (lo, hi), not a + # contiguous codebook layout. + pairs = [(5, 6), (20, 24), (11, 13)] + lo, hi = _bands(pairs) + ids = torch.tensor([[1, 2, 3, 4]]) + out = escalating_beam_search( + _tiny_lm(vocab_size=30), + ids, + torch.ones_like(ids), + num_beams=2, + lo_tok=lo, + hi_tok=hi, + ) + self.assertEqual(tuple(out.shape), (min(2 * 2**3, 2 * 5 * 3), 3)) + for level, (lo_j, hi_j) in enumerate(pairs): + col = out[:, level] + self.assertTrue(bool((col >= lo_j).all())) + self.assertTrue(bool((col <= hi_j).all())) + self.assertEqual(len({tuple(r) for r in out.tolist()}), out.shape[0]) + + @parameterized.expand( + [[0, 1], [1, 3], [2, 16], [0, 16]], + name_func=parameterized_name_func, + ) + def test_left_padding_matches_unpadded(self, seed, n_pad) -> None: + pairs = [(20, 21), (22, 24), (25, 28)] + lo, hi = _bands(pairs) + lm = _tiny_lm(vocab_size=30, seed=seed) + short = torch.tensor([[5, 6, 7]]) + pad = torch.cat([torch.zeros(1, n_pad, dtype=torch.long), short], dim=1) + am_pad = torch.cat( + [torch.zeros(1, n_pad, dtype=torch.long), torch.ones_like(short)], dim=1 + ) + plain = escalating_beam_search( + lm, short, torch.ones_like(short), num_beams=2, lo_tok=lo, hi_tok=hi + ) + padded = escalating_beam_search( + lm, pad, am_pad, num_beams=2, lo_tok=lo, hi_tok=hi + ) + self.assertTrue(torch.equal(plain, padded)) + + def test_ragged_batch_rows_match_solo_runs(self) -> None: + # every row of a ragged batch must decode exactly as if run alone. + pairs = [(20, 21), (22, 24), (25, 28)] + lo, hi = _bands(pairs) + lm = _tiny_lm(vocab_size=30) + row0, row1 = torch.tensor([[5, 6, 7, 8]]), torch.tensor([[9, 10, 11]]) + ids = torch.tensor([[5, 6, 7, 8], [0, 9, 10, 11]]) + am = torch.tensor([[1, 1, 1, 1], [0, 1, 1, 1]]) + out = escalating_beam_search(lm, ids, am, num_beams=2, lo_tok=lo, hi_tok=hi) + width = out.shape[0] // 2 + for i, solo_ids in enumerate([row0, row1]): + solo = escalating_beam_search( + lm, + solo_ids, + torch.ones_like(solo_ids), + num_beams=2, + lo_tok=lo, + hi_tok=hi, + ) + self.assertTrue(torch.equal(out[i * width : (i + 1) * width], solo)) + + def test_exhaustive_matches_bruteforce_topk(self) -> None: + # widths [2, 6, 12] over 2*3*2 = 12 combinations -> no pruning at any + # level, so the beam must reproduce the exact full-recompute ranking. + pairs = [(20, 21), (22, 24), (25, 26)] + lo, hi = _bands(pairs) + lm = _tiny_lm(vocab_size=30) + ids = torch.tensor([[5, 6, 7, 8]]) + am = torch.ones_like(ids) + out = escalating_beam_search(lm, ids, am, num_beams=6, lo_tok=lo, hi_tok=hi) + got = [tuple(r) for r in out.tolist()] + ref = _bruteforce_scores(lm, ids, am, pairs) + self.assertEqual(set(got), set(ref)) + self.assertEqual(len(got), len(ref)) + scores = [ref[c] for c in got] + self.assertTrue( + all(scores[i] >= scores[i + 1] - 1e-4 for i in range(len(scores) - 1)) + ) + self.assertEqual(got[0], max(ref, key=ref.get)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/utils/hf_export_util_test.py b/tzrec/utils/hf_export_util_test.py index 76e150fec..c13d74a44 100644 --- a/tzrec/utils/hf_export_util_test.py +++ b/tzrec/utils/hf_export_util_test.py @@ -13,6 +13,7 @@ import os import shutil import unittest +from unittest import mock import torch from safetensors.torch import load_file @@ -83,7 +84,10 @@ def _tiny_lm(tie=True): class HfExportUtilTest(unittest.TestCase): def setUp(self) -> None: self.test_dir = make_test_dir() - os.environ.setdefault("RANK", "0") + # the asset writers are rank-0-gated; pin it without leaking the value. + patcher = mock.patch.dict(os.environ, {"RANK": "0"}) + patcher.start() + self.addCleanup(patcher.stop) def tearDown(self) -> None: shutil.rmtree(self.test_dir, ignore_errors=True) From 4496bf1c3034fd1b076364b378eeff015967b3cd Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Tue, 28 Jul 2026 03:19:39 +0000 Subject: [PATCH 36/99] [refactor] genrec LM: import transformers at module scope The lazy _hf_auto_classes indirection and the TYPE_CHECKING blocks guarded against transformers being absent, but it is a pinned runtime dependency in requirements/runtime.txt, and no other model under tzrec/models/ defers its imports this way -- torch, torchrec and torchmetrics are all top level. The Auto classes and the two type names are now imported directly, so the annotations are real types rather than strings. escalating_beam's module docstring drops "torch-only": it imports transformers for the backbone type. The "no tzrec deps" property, which is what keeps it independently testable, is unchanged and still asserted by its test. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/generative_rec_lm.py | 56 ++++++++++--------------------- tzrec/models/qwen2_rec_lm.py | 9 ++--- tzrec/modules/escalating_beam.py | 10 +++--- 3 files changed, 25 insertions(+), 50 deletions(-) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 4ae488bb9..054aaf067 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -17,11 +17,18 @@ ``GenerativeRecLMConfig`` (see ``protos/models/generative_model.proto``). """ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple import torch import torch.nn.functional as F import torchmetrics +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + AutoTokenizer, + PreTrainedModel, + PreTrainedTokenizerBase, +) from tzrec.datasets.utils import Batch from tzrec.features.feature import BaseFeature @@ -31,30 +38,6 @@ from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.models import generative_model_pb2 -if TYPE_CHECKING: - from transformers import PreTrainedModel, PreTrainedTokenizerBase - - -def _hf_auto_classes() -> Tuple[Any, Any, Any]: - """Import the HF Auto classes lazily. - - ``transformers`` is an optional, multi-hundred-MB dependency needed only by - this model family. Importing it at module scope would make it mandatory for - the whole package, because ``load_class.auto_import`` re-raises any import - failure under ``tzrec/models/`` out of ``import tzrec``. - - Raises: - ImportError: if ``transformers`` is not installed. - """ - try: - from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer - except ImportError as e: - raise ImportError( - "transformers is required for generative-recommendation LMs. " - "Install via `pip install transformers==4.51.2`." - ) from e - return AutoConfig, AutoModelForCausalLM, AutoTokenizer - class GenerativeRecLM(BaseModel): """Abstract base for HF-backed generative-recommendation LMs. @@ -94,7 +77,7 @@ def __init__( self.init_input() @staticmethod - def _resolve_pad_token_id(tokenizer: "PreTrainedTokenizerBase") -> int: + def _resolve_pad_token_id(tokenizer: PreTrainedTokenizerBase) -> int: """Pad id for the left-padded splice, falling back to eos.""" pad_id = tokenizer.pad_token_id if pad_id is None: @@ -164,33 +147,31 @@ def _history_feature_group(self) -> model_pb2.FeatureGroupConfig: ) return g - def _build_backbone(self) -> "PreTrainedModel": + def _build_backbone(self) -> PreTrainedModel: """Build the EMPTY extended architecture -- no weight download. Only the module shapes matter here; the weights arrive from ``init_from_pretrained`` (cold start) or DCP (restore/eval). """ - auto_config, auto_causal_lm, _ = _hf_auto_classes() hf_model_id = self._model_config.hf_model_id if not hf_model_id: raise ValueError(f"{type(self).__name__}: empty hf_model_id.") - hf_cfg = auto_config.from_pretrained(hf_model_id) - lm = auto_causal_lm.from_config(hf_cfg, torch_dtype=self._param_dtype) + hf_cfg = AutoConfig.from_pretrained(hf_model_id) + lm = AutoModelForCausalLM.from_config(hf_cfg, torch_dtype=self._param_dtype) if next(lm.parameters()).dtype != self._param_dtype: lm = lm.to(self._param_dtype) return lm def _build_extended_tokenizer( self, sid_atoms: int - ) -> Tuple["PreTrainedTokenizerBase", int]: + ) -> Tuple[PreTrainedTokenizerBase, int]: """Add the SID atoms ``C0..C{sid_atoms-1}`` and resize ``self.lm``. Returns ``(tokenizer, base)`` where ``base`` is the tokenizer's next free id BEFORE adding the atoms -- use ``len(tokenizer)``, NOT ``config.vocab_size`` (which counts reserved slots). """ - _, _, auto_tokenizer = _hf_auto_classes() - tokenizer = auto_tokenizer.from_pretrained( + tokenizer = AutoTokenizer.from_pretrained( self._model_config.hf_model_id, use_fast=True ) base = len(tokenizer) @@ -226,10 +207,9 @@ def init_from_pretrained(self) -> None: are drawn from the global RNG and therefore differ per rank; DDP's ``_sync_module_states`` broadcast from rank 0 is what makes them agree. """ - _, auto_causal_lm, _ = _hf_auto_classes() # drop the empty arch first: holding both peaks at 2x model host RAM. self.lm = None - lm = auto_causal_lm.from_pretrained( + lm = AutoModelForCausalLM.from_pretrained( self._model_config.hf_model_id, torch_dtype=self._param_dtype, low_cpu_mem_usage=True, @@ -239,16 +219,16 @@ def init_from_pretrained(self) -> None: ) self.lm = lm - def hf_backbone(self) -> "PreTrainedModel": + def hf_backbone(self) -> PreTrainedModel: """The HF backbone module, for checkpoint/export asset writing.""" return self.lm - def hf_tokenizer(self) -> "PreTrainedTokenizerBase": + def hf_tokenizer(self) -> PreTrainedTokenizerBase: """The extended tokenizer (base vocab + C0..C{sum-1}) to serialize.""" return self._hf_tokenizer def _build_prompt_tokens( - self, tokenizer: "PreTrainedTokenizerBase", cfg: Any + self, tokenizer: PreTrainedTokenizerBase, cfg: Any ) -> None: """Family hook: cache the tokenised prompt template as buffers. diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index dc8ab1fdd..2a9dc03be 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -15,10 +15,11 @@ causal-LM splice, and the ``.model``/``.lm_head`` forward. """ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple import torch from torch.nn.utils.rnn import pad_sequence +from transformers import PreTrainedTokenizerBase from tzrec.datasets.utils import Batch from tzrec.features.feature import BaseFeature @@ -27,10 +28,6 @@ from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.models import generative_model_pb2 -if TYPE_CHECKING: - from transformers import PreTrainedTokenizerBase - - QWEN2_TEMPLATE = { "system_prefix": "<|im_start|>system\n", "system_suffix": "<|im_end|>\n", @@ -92,7 +89,7 @@ def _compute_max_total_length(self) -> int: def _build_prompt_tokens( self, - tokenizer: "PreTrainedTokenizerBase", + tokenizer: PreTrainedTokenizerBase, cfg: generative_model_pb2.Qwen2RecLM, ) -> None: """Tokenise the family chat template once; cache as buffers. diff --git a/tzrec/modules/escalating_beam.py b/tzrec/modules/escalating_beam.py index 35b992e34..5a953191e 100644 --- a/tzrec/modules/escalating_beam.py +++ b/tzrec/modules/escalating_beam.py @@ -9,22 +9,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Escalating-beam SID decode (torch-only, no tzrec deps). +"""Escalating-beam SID decode (no tzrec deps). The beam width doubles at every SID level, so early levels are pruned hard. """ -from typing import TYPE_CHECKING, List, Tuple +from typing import List, Tuple import torch - -if TYPE_CHECKING: - from transformers import PreTrainedModel +from transformers import PreTrainedModel @torch.no_grad() def escalating_beam_search( - model: "PreTrainedModel", + model: PreTrainedModel, input_ids: torch.Tensor, attention_mask: torch.Tensor, *, From ed12b4fe3cb3eee1ec229cb75b09d5ec6805f196 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Tue, 28 Jul 2026 03:19:57 +0000 Subject: [PATCH 37/99] [refactor] LinearDecayLR: rename total_size to num_training_steps total_size measures the whole run from step 0 and includes warmup_size, but every sibling scheduler's horizon excludes it -- ExponentialDecayLR divides by decay_size after subtracting warmup, CosineAnnealingLR clamps to T_max after subtracting warmup. A config swapping cosine_annealing { T_max: N } for linear_decay { total_size: N } with warmup set therefore got a silently shorter decay and no error. The inclusive horizon is correct -- it is HuggingFace's num_training_steps, and matching HF's linear schedule is why the class exists -- so the field takes HF's name instead of one that reads like its siblings'. Behaviour is unchanged: trajectories are identical across warmup/no-warmup, floor/no-floor, and past the horizon, and parity with transformers.get_linear_schedule_with_warmup still holds. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/optim/lr_scheduler.py | 29 +++++++++++++++++------------ tzrec/optim/lr_scheduler_test.py | 8 +++++--- tzrec/protos/optimizer.proto | 11 +++++------ 3 files changed, 27 insertions(+), 21 deletions(-) diff --git a/tzrec/optim/lr_scheduler.py b/tzrec/optim/lr_scheduler.py index 20c40bdeb..9df083a95 100644 --- a/tzrec/optim/lr_scheduler.py +++ b/tzrec/optim/lr_scheduler.py @@ -162,13 +162,15 @@ def _get_lr(self) -> List[float]: class LinearDecayLR(BaseLR): """Linear Decay LearningRate Scheduler. - Decays the learning rate linearly from base_lr to min_learning_rate - over total_size steps or epochs, with optional linear warmup. Mirrors - HuggingFace Trainer's ``lr_scheduler_type: linear``. + Decays the learning rate linearly from base_lr to min_learning_rate, + with optional linear warmup. Mirrors HuggingFace Trainer's + ``lr_scheduler_type: linear``. Args: optimizer (Optimizer): an instance of Optimizer. - total_size (int): total number of steps or epochs for the decay. + num_training_steps (int): length of the whole run in steps or epochs, + INCLUDING warmup_size (unlike decay_size/T_max, which measure only + the post-warmup horizon). min_learning_rate (float): minimum learning rate. warmup_learning_rate (float): warmup start learning rate. warmup_size (int): warmup steps or epochs. @@ -178,20 +180,22 @@ class LinearDecayLR(BaseLR): def __init__( self, optimizer: Optimizer, - total_size: int, + num_training_steps: int, min_learning_rate: float = 0.0, warmup_learning_rate: float = 0.0, warmup_size: int = 0, by_epoch: bool = False, ) -> None: - if total_size <= 0: - raise ValueError(f"total_size must be positive, got {total_size}") - if warmup_size >= total_size: + if num_training_steps <= 0: + raise ValueError( + f"num_training_steps must be positive, got {num_training_steps}" + ) + if warmup_size >= num_training_steps: raise ValueError( f"warmup_size ({warmup_size}) must be smaller than " - f"total_size ({total_size})" + f"num_training_steps ({num_training_steps})" ) - self._total_size = total_size + self._num_training_steps = num_training_steps self._min_learning_rate = min_learning_rate self._warmup_learning_rate = warmup_learning_rate self._warmup_size = warmup_size @@ -207,8 +211,9 @@ def _get_lr(self) -> List[float]: + self._warmup_learning_rate for base_lr in self.base_lrs ] - t = min(step_count - self._warmup_size, self._total_size - self._warmup_size) - decay_scale = 1.0 - t / (self._total_size - self._warmup_size) + decay_steps = self._num_training_steps - self._warmup_size + t = min(step_count - self._warmup_size, decay_steps) + decay_scale = 1.0 - t / decay_steps return [ self._min_learning_rate + (base_lr - self._min_learning_rate) * decay_scale for base_lr in self.base_lrs diff --git a/tzrec/optim/lr_scheduler_test.py b/tzrec/optim/lr_scheduler_test.py index e3551c2d1..93f2c539f 100644 --- a/tzrec/optim/lr_scheduler_test.py +++ b/tzrec/optim/lr_scheduler_test.py @@ -86,7 +86,7 @@ def test_manual_step_lr_with_warmup(self) -> None: def test_linear_decay_lr(self) -> None: params = [torch.tensor([1.0, 2.0])] opt = torch.optim.Adam(params, lr=0.01) - lr = lr_scheduler.LinearDecayLR(opt, total_size=4) + lr = lr_scheduler.LinearDecayLR(opt, num_training_steps=4) lr_gts = [0.0075, 0.005, 0.0025, 0.0, 0.0] for lr_gt in lr_gts: lr.step() @@ -95,7 +95,9 @@ def test_linear_decay_lr(self) -> None: def test_linear_decay_lr_with_min_lr(self) -> None: params = [torch.tensor([1.0, 2.0])] opt = torch.optim.Adam(params, lr=0.01) - lr = lr_scheduler.LinearDecayLR(opt, total_size=4, min_learning_rate=0.002) + lr = lr_scheduler.LinearDecayLR( + opt, num_training_steps=4, min_learning_rate=0.002 + ) lr_gts = [0.008, 0.006, 0.004, 0.002, 0.002] for lr_gt in lr_gts: lr.step() @@ -105,7 +107,7 @@ def test_linear_decay_lr_with_warmup(self) -> None: params = [torch.tensor([1.0, 2.0])] opt = torch.optim.Adam(params, lr=0.01) lr = lr_scheduler.LinearDecayLR( - opt, total_size=6, warmup_size=2, warmup_learning_rate=0.002 + opt, num_training_steps=6, warmup_size=2, warmup_learning_rate=0.002 ) self.assertFalse(lr.by_epoch) # warmup step 0->1: scale=0.5, lr=0.002+(0.01-0.002)*0.5=0.006 diff --git a/tzrec/protos/optimizer.proto b/tzrec/protos/optimizer.proto index 7ee350bd9..55e00ca25 100644 --- a/tzrec/protos/optimizer.proto +++ b/tzrec/protos/optimizer.proto @@ -237,12 +237,11 @@ message ManualStepLR { } message LinearDecayLR { - // total number of steps or epochs to decay from base_lr to - // min_learning_rate (mirrors HF Trainer's `lr_scheduler_type: linear`). - // NOTE: unlike decay_size/T_max in the schedulers above, this horizon is - // measured from step 0 and INCLUDES warmup_size. Required (must be > 0). - optional uint32 total_size = 1; - // minimum learning rate reached at total_size + // length of the whole run, measured from step 0 and INCLUDING warmup_size + // -- HF Trainer's `num_training_steps`, not the post-warmup horizon that + // decay_size/T_max above measure. Required (must be > 0). + optional uint32 num_training_steps = 1; + // minimum learning rate reached at num_training_steps optional float min_learning_rate = 2 [default = 0.0]; // warmup start learning rate optional float warmup_learning_rate = 3 [default = 0.0]; From 9731ac5abd1996fb94772de3394bc007da690791 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Tue, 28 Jul 2026 03:19:57 +0000 Subject: [PATCH 38/99] [feat] genrec LM: add a mock config and an integration test Nothing exercised the model on a real ModelConfig, so the oneof dispatch, the required max_sequence_length, and the derived sample contract -- history is the single JAGGED_SEQUENCE feature_group, the answer is the first label_field -- were unverified end to end. The mock config mirrors the production recipe at toy scale, following tzrec/tests/configs/sid_rqvae_mock.config. Both tests build a tiny Qwen2 locally instead of downloading a backbone, so they need no network. The first proves config -> model -> a batch reaching a finite loss; the second runs train_eval and asserts the checkpoint carries the co-located HF assets that export_format: HF promises, which the cold-start init_from_pretrained path had no coverage for at all. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/tests/configs/qwen2_rec_lm_mock.config | 64 +++++++ tzrec/tests/genrec_integration_test.py | 182 +++++++++++++++++++ 2 files changed, 246 insertions(+) create mode 100644 tzrec/tests/configs/qwen2_rec_lm_mock.config create mode 100644 tzrec/tests/genrec_integration_test.py diff --git a/tzrec/tests/configs/qwen2_rec_lm_mock.config b/tzrec/tests/configs/qwen2_rec_lm_mock.config new file mode 100644 index 000000000..14efb520c --- /dev/null +++ b/tzrec/tests/configs/qwen2_rec_lm_mock.config @@ -0,0 +1,64 @@ +train_input_path: "" +eval_input_path: "" +model_dir: "experiments/qwen2_rec_lm_mock" +train_config { + sparse_optimizer { + adagrad_optimizer { + lr: 0.0 + } + constant_learning_rate { + } + } + dense_optimizer { + adam_optimizer { + lr: 0.0001 + } + linear_decay_learning_rate { + num_training_steps: 32 + } + } + num_epochs: 1 + save_checkpoints_epochs: 1 +} +eval_config { +} +export_config { + export_format: HF +} +data_config { + batch_size: 4 + dataset_type: ParquetDataset + label_fields: "label" + num_workers: 2 + fg_mode: FG_NONE +} +feature_configs { + sequence_raw_feature { + feature_name: "user_sequence" + expression: "user:user_sequence" + value_dim: 1 + sequence_length: 12 + } +} +model_config { + feature_groups { + group_name: "sids" + feature_names: "user_sequence" + group_type: JAGGED_SEQUENCE + } + qwen2_rec_lm { + common { + codebook: 4 + codebook: 4 + codebook: 4 + vocab_pad_to_multiple_of: 128 + ignore_index: -100 + param_dtype: "float32" + max_sequence_length: 12 + } + hf_model_id: "Qwen/Qwen2.5-0.5B" + system_instruction: "你是一个推荐系统，根据用户的历史行为，预测用户在电商场景的下一步行为。我会给你一串连续行为的语义编码，按照用户点击的时间顺序排列，每个行为用三个词表示。" + user_prefix_text: "当前用户的历史行为如下：" + user_suffix_text: "，请预测用户在电商推荐场景后续行为的语义编码" + } +} diff --git a/tzrec/tests/genrec_integration_test.py b/tzrec/tests/genrec_integration_test.py new file mode 100644 index 000000000..f38460c4d --- /dev/null +++ b/tzrec/tests/genrec_integration_test.py @@ -0,0 +1,182 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import glob +import os +import random +import shutil +import unittest +from unittest import mock + +import pyarrow as pa +import pyarrow.parquet as pq + +from tzrec.tests import utils +from tzrec.utils import config_util +from tzrec.utils.test_util import make_test_dir + +_MOCK_CONFIG = "tzrec/tests/configs/qwen2_rec_lm_mock.config" +# must match the mock config's `common.codebook` +_CODEBOOK = [4, 4, 4] + + +def _write_backbone(save_dir: str, vocab_size: int = 256) -> str: + """Save a tiny but COMPLETE Qwen2 model dir so no test downloads a real one. + + Weights are required, not just ``config.json``: cold-start training calls + ``init_from_pretrained`` -> ``from_pretrained``, which refuses a dir with no + checkpoint. The SID offsets only need a tokenizer whose ``len()`` is stable + and which has room to append the ``C*`` atoms, so word-level is enough. + """ + from tokenizers import Tokenizer, models, pre_tokenizers + from transformers import PreTrainedTokenizerFast, Qwen2Config, Qwen2ForCausalLM + + eos = "<|endoftext|>" + vocab = {t: i for i, t in enumerate([eos, "<|im_start|>", "<|im_end|>"])} + for i in range(vocab_size - len(vocab)): + vocab[f"b{i}"] = len(vocab) + tk = Tokenizer(models.WordLevel(vocab=vocab, unk_token=eos)) + tk.pre_tokenizer = pre_tokenizers.Whitespace() + PreTrainedTokenizerFast( + tokenizer_object=tk, + unk_token=eos, + eos_token=eos, + pad_token=eos, + additional_special_tokens=["<|im_start|>", "<|im_end|>"], + ).save_pretrained(save_dir) + config = Qwen2Config( + vocab_size=len(vocab), + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=512, + tie_word_embeddings=True, + ) + Qwen2ForCausalLM(config).save_pretrained(save_dir) + return save_dir + + +def _write_samples(save_dir: str, num_rows: int, seed: int = 0) -> str: + """Write the two-column sample contract: history + answer, both list. + + Codes are local 1-based per-level values in ``[1, codebook[level]]`` and + every row holds whole items in level order. + """ + rnd = random.Random(seed) + + def _item(): + return [rnd.randint(1, size) for size in _CODEBOOK] + + schema = pa.schema( + [ + pa.field("user_sequence", pa.list_(pa.int64()), nullable=False), + pa.field("label", pa.list_(pa.int64()), nullable=False), + ] + ) + table = pa.table( + { + "user_sequence": [ + [c for _ in range(rnd.randint(1, 4)) for c in _item()] + for _ in range(num_rows) + ], + "label": [_item() for _ in range(num_rows)], + }, + schema=schema, + ) + pq.write_table(table, os.path.join(save_dir, "part-0.parquet")) + return os.path.join(save_dir, "*.parquet") + + +class GenRecIntegrationTest(unittest.TestCase): + def setUp(self): + self.success = False + self.test_dir = make_test_dir() + # every rank builds its own backbone; one is enough for a mock run. + patcher = mock.patch.dict( + os.environ, {"TEST_NPROC_PER_NODE": "1", "HF_HUB_OFFLINE": "1"} + ) + patcher.start() + self.addCleanup(patcher.stop) + + def tearDown(self): + if self.success and os.path.exists(self.test_dir): + shutil.rmtree(self.test_dir) + + def _prepare_config(self, num_rows: int = 64) -> str: + """Point the mock config at a local backbone and freshly written samples.""" + backbone = _write_backbone(os.path.join(self.test_dir, "backbone")) + data_dir = os.path.join(self.test_dir, "genrec_data") + os.makedirs(data_dir, exist_ok=True) + data_glob = _write_samples(data_dir, num_rows) + + config = config_util.load_pipeline_config(_MOCK_CONFIG) + config.train_input_path = data_glob + config.eval_input_path = data_glob + config.model_config.qwen2_rec_lm.hf_model_id = backbone + config_path = os.path.join(self.test_dir, "genrec.config") + config_util.save_message(config, config_path) + return config_path + + def test_mock_config_builds_the_model_and_runs_a_batch(self) -> None: + """The config -> model path: oneof dispatch and the derived contract. + + The sample contract is derived, not configured, so only a real + ``ModelConfig`` proves that the history feature_group, the answer + label_field and the SID vocab extension line up. + """ + from tzrec.constant import Mode + from tzrec.datasets.dataset import create_dataloader + from tzrec.main import _create_features, _create_model + + config = config_util.load_pipeline_config(self._prepare_config()) + features = _create_features(list(config.feature_configs), config.data_config) + model = _create_model( + config.model_config, features, list(config.data_config.label_fields) + ) + self.assertEqual(type(model).__name__, "Qwen2RecLM") + self.assertEqual(model._history_group, "sids") + self.assertEqual(model._input_name, "user_sequence") + self.assertEqual(model._label_name, "label") + self.assertEqual(model._num_levels, len(_CODEBOOK)) + # base vocab + sum(codebook) atoms, padded to vocab_pad_to_multiple_of + self.assertGreaterEqual( + model.lm.config.vocab_size, model._base_vocab + sum(_CODEBOOK) + ) + self.assertEqual(model.lm.config.vocab_size % 128, 0) + + dataloader = create_dataloader( + config.data_config, features, config.train_input_path, mode=Mode.TRAIN + ) + model.train() + predictions = model.predict(next(dataloader.get_iterator())) + self.assertEqual(list(predictions), ["loss"]) + self.assertTrue(bool(predictions["loss"].isfinite())) + self.success = True + + def test_qwen2_rec_lm_train_eval(self) -> None: + """End-to-end train -> checkpoint, with HF assets co-located.""" + config_path = self._prepare_config() + self.success = utils.test_train_eval(config_path, self.test_dir) + self.assertTrue(self.success) + ckpts = glob.glob(os.path.join(self.test_dir, "train", "model.ckpt-*")) + self.assertTrue(ckpts, "no checkpoint persisted") + # export_format: HF -> each checkpoint is convertible without the code + for name in ("config.json", "tokenizer.json"): + self.assertTrue( + os.path.exists(os.path.join(ckpts[0], name)), + f"{name} not co-located in {ckpts[0]}", + ) + + +if __name__ == "__main__": + unittest.main() From 7cba6a415605c1b25ee9f0512298d6241b92dc46 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Tue, 28 Jul 2026 07:20:22 +0000 Subject: [PATCH 39/99] [feat] genrec LM: declare the SID space and prompt text on a SID feature BREAKING: SID codes are now 0-based per level, in [0, codebook[level]). Data written for the previous 1-based contract must be regenerated; feeding it to this build raises on the first batch carrying a top code rather than degrading quietly. The item a code maps to is unchanged, so checkpoints stay valid once the data is relabelled -- token = base_vocab + level_offsets[level] + code, verified identical to the old base_vocab + offsets + code - 1. 0-based is what SidRqvae and SidRqkmeans already emit, so the conversion that used to happen in an offline script outside this repo disappears, and with it every bridging -1 in tokenize/detokenize/bands. The scope split behind this: a SID feature owns the SID space -- codebook and the prompt text that wraps its codes -- while the model owns the tokenizer, the extended vocabulary and the ChatML frame. SidFeature exposes sid_vocab_size, num_levels, codebook_sizes and level_offsets; the model asks its features for the shared space and requires them to agree, so adding a feature cannot resize lm_head. Validation and offsetting move into SidFeature._parse, i.e. into the dataloader workers: a malformed row now raises off the collective path instead of leaving one rank in an all-reduce its peers already entered. system_instruction, user_prefix_text and user_suffix_text collapse into one prompt_template carrying {{feature_name}} slots. The template is split as a string and each gap tokenized whole -- splitting token ids instead would let a BPE merge span a seam -- and each slot feature's prefix_text/suffix_text folds into the adjacent gap. N slots are supported, one feature_group each, since a group interleaves its members into a single sequence. A slot naming no feature, a feature no slot names, or a feature without the prompt-text interface are all errors; previously extra features were silently dropped. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/features/sid_feature.py | 173 ++++++++++++++ tzrec/features/sid_feature_test.py | 108 +++++++++ tzrec/models/generative_rec_lm.py | 237 ++++++++++++------- tzrec/models/generative_rec_lm_test.py | 192 +++++++-------- tzrec/models/qwen2_rec_lm.py | 133 ++++++----- tzrec/models/qwen2_rec_lm_test.py | 141 +++++------ tzrec/protos/feature.proto | 53 +++++ tzrec/protos/models/generative_model.proto | 33 ++- tzrec/tests/configs/qwen2_rec_lm_mock.config | 15 +- tzrec/tests/genrec_integration_test.py | 10 +- 10 files changed, 746 insertions(+), 349 deletions(-) create mode 100644 tzrec/features/sid_feature.py create mode 100644 tzrec/features/sid_feature_test.py diff --git a/tzrec/features/sid_feature.py b/tzrec/features/sid_feature.py new file mode 100644 index 000000000..085a85457 --- /dev/null +++ b/tzrec/features/sid_feature.py @@ -0,0 +1,173 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import pyarrow as pa + +from tzrec.datasets.utils import ParsedData +from tzrec.features.feature import BaseFeature, FgMode +from tzrec.protos.feature_pb2 import FeatureConfig + + +class SidFeature(BaseFeature): + """Semantic-ID sequence feature. + + Carries a flat stream of per-level SID codes -- whole items in level order -- + plus the prompt text that wraps them. Codes are 0-based per level, as the + SID-generation models emit them. They are produced offline by a + SID-generation model, so only ``fg_mode = FG_NONE`` is supported; there is no + pyfg counterpart for this feature type. + + Args: + feature_config (FeatureConfig): a instance of feature config. + """ + + def __init__( + self, + feature_config: FeatureConfig, + **kwargs: Any, + ) -> None: + # BaseFeature.__del__ dereferences _fg_op, so seed it before any raise. + self._fg_op = None + # checked before super(), which calls init_fg() for FG_NORMAL and would + # surface the missing pyfg handler instead of this explanation. + fg_mode = kwargs.get("fg_mode", FgMode.FG_NONE) + if fg_mode != FgMode.FG_NONE: + raise ValueError( + f"{self.__class__.__name__}" + f"[{feature_config.sequence_sid_feature.feature_name}] supports " + f"data_config.fg_mode = FG_NONE only (SID codes are generated " + f"offline by a SID model), got {fg_mode}." + ) + super().__init__(feature_config, **kwargs) + self._codebook = self._read_codebook() + self._level_offsets = np.cumsum([0] + self._codebook[:-1]).tolist() + # per-code level index and bound, tiled at parse time to the row length + self._level_sizes = np.asarray(self._codebook) + self._level_offset_arr = np.asarray(self._level_offsets) + + def _read_codebook(self) -> List[int]: + """Validate the declared codebook once and normalize it to a list. + + The proto field is a repeated-scalar container, and every derived + quantity (level count, vocab size, offsets) reads it on the parse hot + path, so it is checked and converted here rather than per access. + """ + codebook = [int(c) for c in self.config.codebook] + if not codebook: + raise ValueError( + f"{self.__class__.__name__}[{self.config.feature_name}]: codebook " + f"must be non-empty." + ) + if any(c <= 0 for c in codebook): + raise ValueError( + f"{self.__class__.__name__}[{self.config.feature_name}]: every " + f"codebook size must be positive, got {codebook}." + ) + return codebook + + @property + def value_dim(self) -> int: + """Fg value dimension of the feature.""" + return self.config.value_dim + + @property + def output_dim(self) -> int: + """Output dimension of the feature after embedding. + + SID codes are token ids consumed by the LM's own embedding table, so the + feature carries no embedding of its own and passes the codes through. + """ + return self.value_dim + + @property + def num_embeddings(self) -> int: + """Get embedding row count.""" + raise RuntimeError( + f"{self.__class__.__name__}[{self.config.feature_name}] has no " + f"embedding table; SID codes index the LM vocabulary." + ) + + @property + def prefix_text(self) -> str: + """Text emitted immediately before this feature's SID tokens.""" + return self.config.prefix_text + + @property + def suffix_text(self) -> str: + """Text emitted immediately after this feature's SID tokens.""" + return self.config.suffix_text + + @property + def codebook(self) -> List[int]: + """Per-level SID vocabulary sizes; validated once at construction.""" + return self._codebook + + @property + def num_levels(self) -> int: + """Codes per item -- also the answer width.""" + return len(self._codebook) + + @property + def sid_vocab_size(self) -> int: + """Atoms the model must append to the backbone vocabulary.""" + return sum(self._codebook) + + @property + def level_offsets(self) -> List[int]: + """Flat offset of each level, i.e. ``cumsum(sizes) - sizes``.""" + return self._level_offsets + + def _build_side_inputs(self) -> Optional[List[Tuple[str, str]]]: + """Input field names with side.""" + if self.config.HasField("expression"): + return [tuple(self.config.expression.split(":"))] + else: + return None + + def _parse(self, input_data: Dict[str, pa.Array]) -> ParsedData: + """Parse the SID stream into flat indices in the shared space. + + Codes are 0-based, exactly as the SID-generation models emit them, so + the flat index IS the atom index and the model only adds ``base_vocab``. + The per-level offsets are folded in here, in the dataloader workers, + rather than on the model's forward path; validating here also keeps a + malformed row off the collective path, where one rank raising becomes + an all-reduce hang on its peers. + """ + parsed = super()._parse(input_data) + num_levels = len(self._codebook) + values = parsed.values.reshape(-1) + if values.size % num_levels != 0: + raise ValueError( + f"{self.__class__.__name__}[{self.config.feature_name}]: SID " + f"stream must hold whole {num_levels}-level items, got " + f"{values.size} codes." + ) + levels = np.arange(values.size) % num_levels + sizes = self._level_sizes[levels] + if ((values < 0) | (values >= sizes)).any(): + raise ValueError( + f"{self.__class__.__name__}[{self.config.feature_name}]: SID " + f"codes must be local 0-based values in [0, codebook[level])." + ) + offsets = self._level_offset_arr[levels] + parsed.values = (values + offsets).reshape(parsed.values.shape) + return parsed + + def _fg_json(self) -> List[Dict[str, Any]]: + """Get fg json config impl.""" + raise RuntimeError( + f"{self.__class__.__name__}[{self.config.feature_name}] has no fg " + f"representation; SID codes are generated offline (fg_mode=FG_NONE)." + ) diff --git a/tzrec/features/sid_feature_test.py b/tzrec/features/sid_feature_test.py new file mode 100644 index 000000000..4b8314ea3 --- /dev/null +++ b/tzrec/features/sid_feature_test.py @@ -0,0 +1,108 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import pyarrow as pa +from google.protobuf import text_format +from parameterized import parameterized + +from tzrec.features.feature import FgMode, create_features +from tzrec.protos import feature_pb2 +from tzrec.utils.test_util import parameterized_name_func + + +def _feature(text, fg_mode=FgMode.FG_NONE): + fc = feature_pb2.FeatureConfig() + text_format.Merge(f"sequence_sid_feature {{ {text} }}", fc) + return create_features([fc], fg_mode=fg_mode)[0] + + +_BASE = ( + 'feature_name: "user_sequence" expression: "user:user_sequence" ' + "codebook: 4 codebook: 4 codebook: 4" +) + + +class SidFeatureTest(unittest.TestCase): + def test_dispatch_and_defaults(self) -> None: + f = _feature(_BASE) + self.assertEqual(type(f).__name__, "SidFeature") + # the oneof FIELD name is what makes it a sequence, not the message + self.assertTrue(f.is_sequence) + self.assertFalse(f.is_sparse) + self.assertEqual(f.name, "user_sequence") + self.assertEqual(f.value_dim, 1) + self.assertEqual(f.output_dim, 1) + self.assertEqual(f.side_inputs, [("user", "user_sequence")]) + self.assertEqual(f.prefix_text, "") + self.assertEqual(f.suffix_text, "") + self.assertEqual(f.codebook, [4, 4, 4]) + self.assertEqual(f.num_levels, 3) + self.assertEqual(f.sid_vocab_size, 12) + self.assertEqual(f.level_offsets, [0, 4, 8]) + + def test_prompt_text_round_trips(self) -> None: + f = _feature(f'{_BASE} prefix_text: "History: " suffix_text: "."') + self.assertEqual(f.prefix_text, "History: ") + self.assertEqual(f.suffix_text, ".") + + @parameterized.expand( + [[FgMode.FG_NORMAL], [FgMode.FG_DAG], [FgMode.FG_BUCKETIZE]], + name_func=parameterized_name_func, + ) + def test_rejects_every_fg_mode_but_none(self, fg_mode) -> None: + # SID codes come from an offline generation model; there is no pyfg + # counterpart, so anything but FG_NONE must fail loudly, not mis-parse. + with self.assertRaisesRegex(ValueError, "FG_NONE only"): + _feature(_BASE, fg_mode=fg_mode) + + def test_parse_folds_in_the_level_offsets(self) -> None: + # offsets [0, 4, 8]: level j's 0-based code k becomes flat index k + off[j], + # which is also the atom index -- no bridging shift anywhere. + f = _feature(_BASE) + parsed = f.parse({"user_sequence": pa.array([[0, 1, 2, 1, 2, 3], [0, 0, 0]])}) + self.assertEqual( + parsed.values.flatten().tolist(), [0, 5, 10, 1, 6, 11, 0, 4, 8] + ) + self.assertEqual(parsed.seq_lengths.tolist(), [6, 3]) + + def test_parse_rejects_out_of_range_and_partial_items(self) -> None: + f = _feature(_BASE) + with self.assertRaisesRegex(ValueError, "local 0-based"): + f.parse({"user_sequence": pa.array([[0, 1, 4]])}) # 4 == codebook[2] + with self.assertRaisesRegex(ValueError, "local 0-based"): + f.parse({"user_sequence": pa.array([[-1, 1, 2]])}) + with self.assertRaisesRegex(ValueError, "whole 3-level items"): + f.parse({"user_sequence": pa.array([[0, 1]])}) + + def test_rejects_a_bad_codebook(self) -> None: + for bad, msg in (("", "non-empty"), ("codebook: 4 codebook: 0", "positive")): + with self.subTest(bad=bad): + base = 'feature_name: "s" expression: "user:s" ' + bad + with self.assertRaisesRegex(ValueError, msg): + _ = _feature(base).codebook + + def test_no_embedding_table(self) -> None: + f = _feature(_BASE) + self.assertFalse(f.has_embedding) + self.assertIsNone(f.emb_config) + with self.assertRaisesRegex(RuntimeError, "no .*embedding table"): + _ = f.num_embeddings + + def test_no_fg_representation(self) -> None: + f = _feature(_BASE) + with self.assertRaisesRegex(RuntimeError, "no fg representation"): + f.fg_json() + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_rec_lm.py index 054aaf067..2b85006d1 100644 --- a/tzrec/models/generative_rec_lm.py +++ b/tzrec/models/generative_rec_lm.py @@ -17,6 +17,7 @@ ``GenerativeRecLMConfig`` (see ``protos/models/generative_model.proto``). """ +import re from typing import Any, Dict, List, Optional, Tuple import torch @@ -93,9 +94,6 @@ def _read_common_config( self, common: generative_model_pb2.GenerativeRecLMConfig ) -> int: """Parse shared proto knobs into attributes; return the SID atom count.""" - hist = self._history_feature_group() - self._history_group: str = hist.group_name - self._input_name: str = hist.feature_names[0] self._label_name: str = self._labels[0] if self._labels else "" self._ignore_index: int = int(common.ignore_index) self._generated_sids_key: str = common.generated_sids_key @@ -107,45 +105,120 @@ def _read_common_config( ) self._param_dtype: torch.dtype = param_dtype self._max_seq_length: int = int(common.max_sequence_length) - codebook = [int(c) for c in common.codebook] - if not codebook: - raise ValueError("GenerativeRecLM: codebook must be non-empty.") - if any(c <= 0 for c in codebook): - raise ValueError( - f"GenerativeRecLM: every codebook size must be positive, got " - f"{codebook}." - ) + codebook = self._shared_sid_space() self._num_levels = len(codebook) - offsets, sizes = self._sid_level_layout(codebook) - self.register_buffer("_level_offsets", offsets, persistent=False) + sizes = torch.tensor(codebook, dtype=torch.long) self.register_buffer("_codebook_sizes", sizes, persistent=False) + # only the decode path still needs per-level offsets: SidFeature folds + # them into the input stream, but generated tokens must be split back + # into per-level codes and no feature is involved in generation. + self.register_buffer( + "_level_offsets", torch.cumsum(sizes, 0) - sizes, persistent=False + ) self._vocab_pad_mult = int(common.vocab_pad_to_multiple_of) return sum(codebook) - def _history_feature_group(self) -> model_pb2.FeatureGroupConfig: - """The history feature_group = the single declared feature_group. + def _shared_sid_space(self) -> List[int]: + """The one codebook every SID feature declares. - Must be JAGGED_SEQUENCE: a SEQUENCE group would silently deliver - padded-dense values instead of the flat raw SID stream. + SID features share a single space: the model has one extended vocabulary + and one answer width, so adding a feature must not resize ``lm_head``. + Requiring the declarations to agree makes that invariant explicit and + turns a typo into an error instead of a silent reshape. """ - if not self._feature_groups: + spaces = {} + for feature in self._features: + codebook = getattr(feature, "codebook", None) + if codebook is not None: + spaces[feature.name] = tuple(codebook) + if not spaces: raise ValueError( - f"{type(self).__name__}: no feature_group declared; genrec needs " - f"one JAGGED_SEQUENCE group carrying the history SID stream." + f"{type(self).__name__}: no SID feature declares a codebook; " + f"genrec needs at least one sequence_sid_feature." ) - g = self._feature_groups[0] - if not g.feature_names: + if len(set(spaces.values())) != 1: raise ValueError( - f"{type(self).__name__}: history feature_group {g.group_name!r} " - f"has no feature_names." + f"{type(self).__name__}: all SID features must share one " + f"codebook, got {dict(sorted(spaces.items()))}." ) - if g.group_type != model_pb2.JAGGED_SEQUENCE: + return list(next(iter(spaces.values()))) + + def _slot_group_names(self) -> Dict[str, str]: + """{feature_name: group_name} for every declared feature_group. + + Each group must be JAGGED_SEQUENCE and hold exactly ONE feature: the + EmbeddingGroup column-interleaves the members of a group into a single + ``{group}.sequence``, so two features in one group could not be spliced + into separate prompt slots. + """ + by_feature: Dict[str, str] = {} + for group in self._feature_groups: + if group.group_type != model_pb2.JAGGED_SEQUENCE: + raise ValueError( + f"{type(self).__name__}: feature_group {group.group_name!r} " + f"must be JAGGED_SEQUENCE, got " + f"{model_pb2.FeatureGroupType.Name(group.group_type)}." + ) + if len(group.feature_names) != 1: + raise ValueError( + f"{type(self).__name__}: feature_group {group.group_name!r} " + f"must hold exactly one feature (its members are interleaved " + f"into one sequence), got {list(group.feature_names)}." + ) + by_feature[group.feature_names[0]] = group.group_name + return by_feature + + def _resolve_prompt_slots( + self, template: str + ) -> Tuple[List[str], List[BaseFeature], List[str]]: + """Split a prompt template into its static gaps and its slot features. + + ``template`` carries ``{{feature_name}}`` slots; returns ``N+1`` static + gap strings, the ``N`` features they are interleaved with in template + order, and those features' feature_group names. Every slot must name a + declared feature, every declared feature_group must be referenced by a + slot, and every slot feature must expose the prompt-text interface + (``prefix_text`` / ``suffix_text``) -- so a misspelt or unused feature + fails here rather than silently vanishing from the prompt. + """ + parts = re.split(r"\{\{(\w+)\}\}", template) + gaps, names = parts[0::2], parts[1::2] + if not names: + raise ValueError( + f"{type(self).__name__}: prompt_template needs at least one " + f"{{{{feature_name}}}} slot naming a declared feature." + ) + group_of = self._slot_group_names() + by_name = {f.name: f for f in self._features} + features = [] + for name in names: + if name not in group_of: + raise ValueError( + f"{type(self).__name__}: prompt_template slot {{{{{name}}}}} " + f"names no feature_group; declared: {sorted(group_of)}." + ) + feature = by_name.get(name) + if feature is None: + raise ValueError( + f"{type(self).__name__}: prompt_template slot {{{{{name}}}}} " + f"has no feature_config." + ) + if not hasattr(feature, "prefix_text") or not hasattr( + feature, "suffix_text" + ): + raise ValueError( + f"{type(self).__name__}: feature {name!r} is a " + f"{type(feature).__name__}, which does not expose the prompt " + f"text interface (prefix_text/suffix_text); use a SidFeature." + ) + features.append(feature) + unused = sorted(set(group_of) - set(names)) + if unused: raise ValueError( - f"{type(self).__name__}: history feature_group {g.group_name!r} " - f"must be JAGGED_SEQUENCE, got " - f"{model_pb2.FeatureGroupType.Name(g.group_type)}." + f"{type(self).__name__}: feature_group(s) {unused} are declared " + f"but never referenced by a prompt_template slot." ) - return g + return gaps, features, [group_of[n] for n in names] def _build_backbone(self) -> PreTrainedModel: """Build the EMPTY extended architecture -- no weight download. @@ -244,52 +317,44 @@ def device(self) -> torch.device: """Device the HF backbone runs on.""" return self.lm.device - def _tokenize_sids( - self, codes: torch.Tensor, level_ids: torch.Tensor - ) -> torch.Tensor: - """Map local 1-based per-level codes to extended-vocab token ids. + def _tokenize_sids(self, flat: torch.Tensor) -> torch.Tensor: + """Map flat SID indices to extended-vocab token ids. - ``level_ids`` must broadcast against ``codes`` and identify each code's - RQ level. + ``SidFeature`` has already folded the per-level offsets in, and codes are + 0-based, so the flat index IS the atom index: the model only owns the + shift into its own vocabulary. """ - return codes + self._level_offsets[level_ids] + (self._base_vocab - 1) + return flat + self._base_vocab def _detokenize_sids( self, tokens: torch.Tensor, level_ids: torch.Tensor ) -> torch.Tensor: - """Inverse of ``_tokenize_sids``: token ids to local 1-based codes.""" - return tokens - (self._base_vocab - 1) - self._level_offsets[level_ids] - - @staticmethod - def _sid_level_layout(codebook: List[int]) -> Tuple[torch.Tensor, torch.Tensor]: - """Per-level ``(flat_offset, size)``; the levels occupy disjoint bands.""" - sizes = torch.tensor(codebook, dtype=torch.long) - return torch.cumsum(sizes, 0) - sizes, sizes + """Inverse of ``_tokenize_sids``: token ids to local 0-based codes.""" + return tokens - self._base_vocab - self._level_offsets[level_ids] def _sid_token_bands(self) -> Tuple[torch.Tensor, torch.Tensor]: """Return the inclusive token-id band for every SID level.""" - level_ids = torch.arange(self._num_levels, device=self.device) return ( - self._tokenize_sids(torch.ones_like(level_ids), level_ids), - self._tokenize_sids(self._codebook_sizes, level_ids), + self._tokenize_sids(self._level_offsets), + self._tokenize_sids(self._level_offsets + self._codebook_sizes - 1), ) def _validate_sid_candidates( self, new_tokens: torch.Tensor, batch_size: int ) -> torch.Tensor: - """Decode generated tokens to local 1-based codes and reject bad beams. + """Decode generated tokens to local 0-based codes and reject bad beams. ``new_tokens`` is the per-beam tail ``(B*C, w)`` (``w`` may be < ``num_levels`` when beams stop early). Returns ``(batch_size, C, num_levels)`` local codes. Every malformed candidate (early EOS / non-SID / wrong-level atom) is set to ``-1``, which cannot match a real - 1-based code. + 0-based code. """ level_ids = torch.arange(new_tokens.shape[1], device=new_tokens.device) codes = self._detokenize_sids(new_tokens, level_ids) codes = F.pad(codes, (0, self._num_levels - codes.shape[1]), value=-1) # one out-of-band atom invalidates the whole candidate row. - invalid = ((codes < 1) | (codes > self._codebook_sizes)).any(dim=1) + invalid = ((codes < 0) | (codes >= self._codebook_sizes)).any(dim=1) codes = codes.masked_fill(invalid.unsqueeze(1), -1) # decoders return rows batch-major ([b0_c0, b0_c1, ...]); group per user. return codes.view(batch_size, -1, self._num_levels) @@ -310,11 +375,12 @@ def build_input(self, batch: Batch) -> Dict[str, List[torch.Tensor]]: """ g = self.embedding_group(batch) rows: Dict[str, List[torch.Tensor]] = { - self._input_name: self._sid_token_rows( - g[f"{self._history_group}.sequence"], - g[f"{self._history_group}.sequence_length"], + name: self._sid_token_rows( + g[f"{group}.sequence"], + g[f"{group}.sequence_length"], max_codes=self._max_seq_length, - ), + ) + for name, group in zip(self._slot_names, self._slot_groups) } if not self.is_inference: if not self._label_name: @@ -323,47 +389,28 @@ def build_input(self, batch: Batch) -> Dict[str, List[torch.Tensor]]: f"declare it as the first data_config.label_field." ) jt = batch.jagged_labels[self._label_name] - rows[self._label_name] = self._sid_token_rows( - jt.values(), jt.lengths(), expected_width=self._num_levels - ) + rows[self._label_name] = self._answer_token_rows(jt.values(), jt.lengths()) return rows def _sid_token_rows( self, values: torch.Tensor, lengths: torch.Tensor, - expected_width: Optional[int] = None, max_codes: Optional[int] = None, ) -> List[torch.Tensor]: - """Map flat SID ``(values, lengths)`` -> per-row token-id tensors. + """Map a feature's flat SID stream to per-row token-id tensors. - ``values`` may arrive as float / shape ``(N, 1)``. Rows must contain - whole items so the per-level offsets restart at level 0 for each row; - ``expected_width``, when set, requires exactly that many codes per row. + ``SidFeature._parse`` has already validated the codes and folded in the + per-level offsets, so this only applies the model-owned budget and the + shift into the extended vocabulary. ``max_codes``, when set, caps each row to its most-recent whole items - (the last ``floor(max_codes / num_levels) * num_levels`` codes) so the - pre-allocated pool covers every batch. + (the last ``floor(max_codes / num_levels) * num_levels`` codes, dropping + the oldest head). Skipped unless a row overflows. """ if values.dim() == 2 and values.size(-1) == 1: values = values.squeeze(-1) sizes = lengths.long().tolist() - misaligned = [i for i, n in enumerate(sizes) if n % self._num_levels != 0] - if misaligned: - raise ValueError( - f"{type(self).__name__}: SID rows must contain whole " - f"{self._num_levels}-level items; rows {misaligned} have lengths " - f"{[sizes[i] for i in misaligned]}." - ) - if expected_width is not None: - bad = [i for i, n in enumerate(sizes) if n != expected_width] - if bad: - raise ValueError( - f"{type(self).__name__}: each SID item must be " - f"{expected_width} codes (len(codebook)); rows {bad} have " - f"{[sizes[i] for i in bad]} -- anomalous sample(s)." - ) - # TODO(shuqi): move truncation into FG once FG can keep the TAIL, not the HEAD. if max_codes: keep = (max_codes // self._num_levels) * self._num_levels @@ -371,15 +418,37 @@ def _sid_token_rows( rows = torch.split(values, sizes) values = torch.cat([r[-keep:] for r in rows]) sizes = [min(n, keep) for n in sizes] + tokens = self._tokenize_sids(values.to(self.device).long()) + return list(torch.split(tokens, sizes)) + + def _answer_token_rows( + self, values: torch.Tensor, lengths: torch.Tensor + ) -> List[torch.Tensor]: + """Map the answer label to token ids. + + The answer is a ``data_config.label_field``, not a feature, so nothing + has offset it: this is the one place the model still owns the per-level + fold-in. Every row must be exactly ``num_levels`` codes. + """ + if values.dim() == 2 and values.size(-1) == 1: + values = values.squeeze(-1) + sizes = lengths.long().tolist() + bad = [i for i, n in enumerate(sizes) if n != self._num_levels] + if bad: + raise ValueError( + f"{type(self).__name__}: each answer must be " + f"{self._num_levels} codes (len(codebook)); rows {bad} have " + f"{[sizes[i] for i in bad]} -- anomalous sample(s)." + ) codes = values.to(self.device).long() level_ids = torch.arange(codes.numel(), device=self.device) % self._num_levels - invalid = (codes < 1) | (codes > self._codebook_sizes[level_ids]) + invalid = (codes < 0) | (codes >= self._codebook_sizes[level_ids]) if invalid.any().item(): raise ValueError( - f"{type(self).__name__}: SID codes must be local 1-based values " - f"in [1, codebook[level]]." + f"{type(self).__name__}: answer SID codes must be local 0-based " + f"values in [0, codebook[level])." ) - tokens = self._tokenize_sids(codes, level_ids) + tokens = self._tokenize_sids(codes + self._level_offsets[level_ids]) return list(torch.split(tokens, sizes)) def init_loss(self) -> None: diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_rec_lm_test.py index 80154a760..88328e875 100644 --- a/tzrec/models/generative_rec_lm_test.py +++ b/tzrec/models/generative_rec_lm_test.py @@ -36,6 +36,11 @@ def lengths(self): return self._l +def _sid_feature(name="user_sequence", codebook=(2, 3, 4)): + """Minimal stand-in for SidFeature: the model only asks for the space.""" + return types.SimpleNamespace(name=name, codebook=list(codebook)) + + def _stub(codebook=None, base_vocab=100, device="cpu"): """A Qwen2RecLM with the base data-prep state wired up, but no HF backbone.""" codebook = codebook or [2, 3, 4] @@ -44,9 +49,11 @@ def _stub(codebook=None, base_vocab=100, device="cpu"): m._base_vocab = base_vocab m._num_levels = len(codebook) m.lm = types.SimpleNamespace(device=torch.device(device)) - offsets, sizes = Qwen2RecLM._sid_level_layout(codebook) - m.register_buffer("_level_offsets", offsets, persistent=False) + sizes = torch.tensor(codebook, dtype=torch.long) m.register_buffer("_codebook_sizes", sizes, persistent=False) + m.register_buffer( + "_level_offsets", torch.cumsum(sizes, 0) - sizes, persistent=False + ) return m @@ -60,7 +67,6 @@ def test_model_config_oneof_resolves_to_the_class(self) -> None: from tzrec.utils import config_util cfg = model_pb2.ModelConfig() - cfg.qwen2_rec_lm.common.codebook.extend([2, 3]) cfg.qwen2_rec_lm.common.max_sequence_length = 8 # required field self.assertEqual(config_util.which_msg(cfg, "model"), "Qwen2RecLM") self.assertIs( @@ -109,7 +115,7 @@ def test_configurable_knob_defaults(self) -> None: def test_read_common_config_reads_knobs(self) -> None: m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) - m._features = [] + m._features = [_sid_feature(codebook=(2, 3, 4))] m._labels = ["label"] m._feature_groups = [ types.SimpleNamespace( @@ -122,13 +128,10 @@ def test_read_common_config_reads_knobs(self) -> None: ignore_index=-100, generated_sids_key="my_sids", param_dtype="bfloat16", - codebook=[2, 3, 4], vocab_pad_to_multiple_of=128, max_sequence_length=288, ) sid_atoms = m._read_common_config(common) - self.assertEqual(m._history_group, "user_seq") # the single group - self.assertEqual(m._input_name, "user_sequence") # its one member self.assertEqual(m._label_name, "label") # from label_fields[0] self.assertEqual(m._generated_sids_key, "my_sids") self.assertIs(m._param_dtype, torch.bfloat16) @@ -146,7 +149,7 @@ def test_read_common_config_reads_knobs(self) -> None: def test_read_common_config_no_feature_group_raises(self) -> None: m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) - m._features = [] + m._features = [_sid_feature(codebook=(2, 3, 4))] m._labels = ["label"] m._feature_groups = [] common = types.SimpleNamespace( @@ -157,14 +160,15 @@ def test_read_common_config_no_feature_group_raises(self) -> None: vocab_pad_to_multiple_of=128, max_sequence_length=0, ) - with self.assertRaisesRegex(ValueError, "no feature_group declared"): - m._read_common_config(common) + # group validation moved to _resolve_prompt_slots (prompt-driven) + m._read_common_config(common) + self.assertEqual(m._num_levels, 3) def test_read_common_config_rejects_bad_group_and_codebook(self) -> None: def _wired(group_type, feature_names=("user_sequence",)): m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) - m._features = [] + m._features = [_sid_feature(codebook=(2, 3, 4))] m._labels = ["label"] m._feature_groups = [ types.SimpleNamespace( @@ -180,25 +184,31 @@ def _common(codebook): ignore_index=-100, generated_sids_key="generated_sids", param_dtype="float32", - codebook=codebook, vocab_pad_to_multiple_of=128, max_sequence_length=0, ) # a SEQUENCE group emits the same key with padded-dense semantics. with self.assertRaisesRegex(ValueError, "must be JAGGED_SEQUENCE"): - _wired(model_pb2.SEQUENCE)._read_common_config(_common([2, 3])) - with self.assertRaisesRegex(ValueError, "has no feature_names"): - _wired(model_pb2.JAGGED_SEQUENCE, ())._read_common_config(_common([2, 3])) - with self.assertRaisesRegex(ValueError, "codebook must be non-empty"): - _wired(model_pb2.JAGGED_SEQUENCE)._read_common_config(_common([])) - with self.assertRaisesRegex(ValueError, "codebook size must be positive"): - _wired(model_pb2.JAGGED_SEQUENCE)._read_common_config(_common([2, 0])) + _wired(model_pb2.SEQUENCE)._slot_group_names() + with self.assertRaisesRegex(ValueError, "exactly one feature"): + _wired(model_pb2.JAGGED_SEQUENCE, ())._slot_group_names() + # a codebook the SID features disagree on is the model's business + m = _wired(model_pb2.JAGGED_SEQUENCE) + m._features = [ + _sid_feature("user_sequence", (2, 3)), + _sid_feature("other_seq", (4, 4)), + ] + with self.assertRaisesRegex(ValueError, "must share one"): + m._shared_sid_space() + m._features = [] + with self.assertRaisesRegex(ValueError, "no SID feature declares"): + m._shared_sid_space() def test_vocab_pad_zero_disables_padding(self) -> None: m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) - m._features = [] + m._features = [_sid_feature(codebook=(2, 3, 4))] m._labels = ["label"] m._feature_groups = [ types.SimpleNamespace( @@ -225,7 +235,6 @@ def _common(max_seq): ignore_index=-100, generated_sids_key="generated_sids", param_dtype="float32", - codebook=[4, 4, 4], vocab_pad_to_multiple_of=128, max_sequence_length=max_seq, ) @@ -233,7 +242,7 @@ def _common(max_seq): def _wired(): m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) - m._features = [] + m._features = [_sid_feature(codebook=(2, 3, 4))] m._labels = ["label"] m._feature_groups = [ types.SimpleNamespace( @@ -263,14 +272,15 @@ def test_device_property(self) -> None: def test_tokenize_sids(self) -> None: m = _stub(base_vocab=100) - codes = torch.tensor([[1, 1, 1], [2, 3, 4]]) + # 0-based codes: the flat index IS the atom index, so no bridging shift. + flat = torch.tensor([[0, 2, 5], [1, 4, 8]]) # offsets [0, 2, 5] level_ids = torch.arange(3) - out = m._tokenize_sids(codes, level_ids) + out = m._tokenize_sids(flat) self.assertEqual(out.tolist(), [[100, 102, 105], [101, 104, 108]]) self.assertEqual(out.dtype, torch.int64) + # decode goes the whole way back to per-level codes self.assertEqual( - m._detokenize_sids(out, level_ids).tolist(), - codes.tolist(), + m._detokenize_sids(out, level_ids).tolist(), [[0, 0, 0], [1, 2, 3]] ) def test_sid_token_bands_use_same_level_aware_mapping(self) -> None: @@ -279,114 +289,75 @@ def test_sid_token_bands_use_same_level_aware_mapping(self) -> None: self.assertEqual(lo.tolist(), [100, 102, 105]) self.assertEqual(hi.tolist(), [101, 104, 108]) - def test_sid_token_rows_split_and_cast(self) -> None: + def test_sid_token_rows_shifts_and_splits(self) -> None: + # values arrive FLAT from SidFeature._parse; the model adds base_vocab-1 m = _stub(base_vocab=100) - jt = _FakeJT([1, 1, 1, 2, 3, 4, 2, 1, 3], [6, 3]) + jt = _FakeJT([0, 2, 5, 1, 4, 8, 0, 3, 6], [6, 3]) rows = m._sid_token_rows(jt.values(), jt.lengths()) self.assertEqual( [r.tolist() for r in rows], - [[100, 102, 105, 101, 104, 108], [101, 102, 107]], + [[100, 102, 105, 101, 104, 108], [100, 103, 106]], ) self.assertTrue(all(r.dtype == torch.int64 for r in rows)) def test_sid_token_rows_squeezes_n1(self) -> None: m = _stub(base_vocab=100) - jt = _FakeJT([1, 1, 1], [3], dim2=True) # (N, 1) + jt = _FakeJT([0, 2, 5], [3], dim2=True) # (N, 1) rows = m._sid_token_rows(jt.values(), jt.lengths()) self.assertEqual([r.tolist() for r in rows], [[100, 102, 105]]) - def test_sid_token_rows_width_ok(self) -> None: - m = _stub(base_vocab=100) - jt = _FakeJT([1, 1, 1, 2, 3, 4], [3, 3]) - rows = m._sid_token_rows(jt.values(), jt.lengths(), expected_width=3) - self.assertEqual( - [r.tolist() for r in rows], - [[100, 102, 105], [101, 104, 108]], - ) - - def test_sid_token_rows_width_violation_raises(self) -> None: - m = _stub(base_vocab=100) - with self.assertRaises(ValueError): - jt = _FakeJT([1, 1, 1, 1, 1, 1, 1, 1, 1], [3, 6]) - m._sid_token_rows(jt.values(), jt.lengths(), expected_width=3) - - def test_sid_token_rows_rejects_partial_items(self) -> None: - m = _stub(base_vocab=100) - with self.assertRaisesRegex(ValueError, "whole 3-level items"): - # The total is divisible by 3, but neither row starts a whole item. - jt = _FakeJT([1, 1, 1, 1, 1, 1], [2, 4]) - m._sid_token_rows(jt.values(), jt.lengths()) - - def test_sid_token_rows_rejects_out_of_range_codes(self) -> None: - m = _stub(base_vocab=100) - for values in ( - [0, 1, 1], - [1, 0, 1], - [1, 1, 0], - [-1, 1, 1], - [1, -1, 1], - [1, 1, -1], - [3, 1, 1], - [1, 4, 1], - [1, 1, 5], - [1, 3, 6], # global cross-level codes, not local per-level - ): - with self.subTest(values=values): - with self.assertRaisesRegex(ValueError, "local 1-based"): - jt = _FakeJT(values, [3]) - m._sid_token_rows(jt.values(), jt.lengths()) - def test_sid_token_rows_recency_clip(self) -> None: m = _stub(base_vocab=100) - items = [(1, 1, 1), (2, 3, 4), (1, 2, 3), (2, 1, 4), (1, 3, 2)] - toks = [ - [100, 102, 105], - [101, 104, 108], - [100, 103, 107], - [101, 102, 108], - [100, 104, 106], - ] - values = torch.tensor(items, dtype=torch.float).flatten() + flat = [0, 2, 5, 1, 3, 6, 0, 4, 7, 1, 2, 8, 0, 3, 5] # 5 whole items + values = torch.tensor(flat, dtype=torch.float) lengths = torch.tensor([values.numel()]) + tail = [100 + v for v in flat[6:]] # last three items, +base_vocab + # cap 9 -> keep the most recent three whole items + self.assertEqual( + m._sid_token_rows(values, lengths, max_codes=9)[0].tolist(), tail + ) + # item-aligned: cap 10 still keeps 9, never cuts mid-item + self.assertEqual( + m._sid_token_rows(values, lengths, max_codes=10)[0].tolist(), tail + ) + # disabled -> untouched + self.assertEqual( + m._sid_token_rows(values, lengths, max_codes=0)[0].tolist(), + [100 + v for v in flat], + ) - # 15 codes (5 items), cap 9 -> keep the most recent three whole items. - rows = m._sid_token_rows(values, lengths, max_codes=9) - self.assertEqual(rows[0].tolist(), sum(toks[2:], [])) - # item-aligned: cap 10 still keeps 9 (3 whole items), never cuts mid-item - rows = m._sid_token_rows(values, lengths, max_codes=10) - self.assertEqual(rows[0].tolist(), sum(toks[2:], [])) - # within cap -> untouched - rows = m._sid_token_rows(values[:6], torch.tensor([6]), max_codes=9) - self.assertEqual(rows[0].tolist(), sum(toks[:2], [])) - # disabled (0/None) -> no clip - rows = m._sid_token_rows(values, lengths, max_codes=0) - self.assertEqual(rows[0].tolist(), sum(toks, [])) - - def test_validate_sid_candidates_groups_batch_major(self) -> None: + def test_answer_token_rows_folds_offsets_and_validates(self) -> None: + # the answer is a label_field, not a feature, so the model still owns + # the per-level fold-in for it. m = _stub(base_vocab=100) - # decoders emit rows batch-major: [b0_c0, b0_c1, b1_c0, b1_c1] - tokens = torch.tensor( - [[100, 102, 105], [101, 104, 108], [101, 103, 106], [100, 104, 107]] - ) - sids = m._validate_sid_candidates(tokens, batch_size=2) - self.assertEqual(tuple(sids.shape), (2, 2, 3)) - self.assertEqual(sids[0].tolist(), [[1, 1, 1], [2, 3, 4]]) - self.assertEqual(sids[1].tolist(), [[2, 2, 2], [1, 3, 3]]) + jt = _FakeJT([0, 0, 0, 1, 2, 3], [3, 3]) + rows = m._answer_token_rows(jt.values(), jt.lengths()) + self.assertEqual([r.tolist() for r in rows], [[100, 102, 105], [101, 104, 108]]) + with self.assertRaisesRegex(ValueError, "each answer must be"): + bad = _FakeJT([0, 0, 0, 0, 0, 0], [2, 4]) + m._answer_token_rows(bad.values(), bad.lengths()) + with self.assertRaisesRegex(ValueError, "local 0-based"): + bad = _FakeJT([0, 3, 0], [3]) # 3 == codebook[1] + m._answer_token_rows(bad.values(), bad.lengths()) def test_build_input_history_group_label_field(self) -> None: m = _stub(base_vocab=100) - m._input_name, m._label_name = "user_sequence", "label" - m._history_group = "user_seq" + m._label_name = "label" + m._slot_names = ["user_sequence"] + m._slot_groups = ["user_seq"] + m._slot_names = ["user_sequence"] + m._slot_groups = ["user_seq"] m._max_seq_length = 0 m._is_inference = False # train: the answer label_field is read too m.embedding_group = lambda b: { + # flat indices, as SidFeature._parse emits them "user_seq.sequence": torch.tensor( - [1.0, 1.0, 1.0, 2.0, 3.0, 4.0, 2.0, 1.0, 3.0] + [0.0, 2.0, 5.0, 1.0, 4.0, 8.0, 1.0, 2.0, 7.0] ), "user_seq.sequence_length": torch.tensor([6, 3]), } batch = types.SimpleNamespace( - jagged_labels={"label": _FakeJT([2, 3, 4, 1, 2, 3], [3, 3])} + jagged_labels={"label": _FakeJT([1, 2, 3, 0, 1, 2], [3, 3])} ) rows = m.build_input(batch) self.assertEqual( @@ -400,12 +371,15 @@ def test_build_input_history_group_label_field(self) -> None: def test_build_input_skips_label_in_inference(self) -> None: m = _stub(base_vocab=100) - m._input_name, m._label_name = "user_sequence", "label" - m._history_group = "user_seq" + m._label_name = "label" + m._slot_names = ["user_sequence"] + m._slot_groups = ["user_seq"] + m._slot_names = ["user_sequence"] + m._slot_groups = ["user_seq"] m._max_seq_length = 0 m._is_inference = True # inference: history only, no ground-truth label m.embedding_group = lambda b: { - "user_seq.sequence": torch.tensor([1.0, 1.0, 1.0]), + "user_seq.sequence": torch.tensor([0.0, 2.0, 5.0]), "user_seq.sequence_length": torch.tensor([3]), } rows = m.build_input(types.SimpleNamespace(jagged_labels={})) diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/qwen2_rec_lm.py index 2a9dc03be..3de1c389e 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/qwen2_rec_lm.py @@ -28,23 +28,17 @@ from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.models import generative_model_pb2 -QWEN2_TEMPLATE = { - "system_prefix": "<|im_start|>system\n", - "system_suffix": "<|im_end|>\n", - "user_prefix": "<|im_start|>user\n", - "user_suffix": "<|im_end|>\n", - "asst_prefix": "<|im_start|>assistant\n", - "asst_suffix": "<|im_end|>\n", - "default_system_instruction": ( - "You are Qwen, created by Alibaba Cloud. You are a helpful assistant." - ), -} - class Qwen2RecLM(GenerativeRecLM): """Qwen2 / Qwen2.5 generative-recommendation LM.""" - CHAT_TEMPLATE = QWEN2_TEMPLATE + # ChatML frame. Family-specific; a subclass overrides it wholesale. + CHAT_TEMPLATE = { + "user_prefix": "<|im_start|>user\n", + "user_suffix": "<|im_end|>\n", + "asst_prefix": "<|im_start|>assistant\n", + "asst_suffix": "<|im_end|>\n", + } def __init__( self, @@ -78,43 +72,60 @@ def _compute_max_total_length(self) -> int: if self._max_seq_length <= 0: return 0 frame = ( - self.tpl_system.numel() - + self.tpl_user_prefix.numel() - + self.tpl_user_suffix.numel() - + self.tpl_asst_prefix.numel() + sum( + getattr(self, f"tpl_gap_{i}").numel() + for i in range(self._num_slots + 1) + ) + self.tpl_asst_suffix.numel() + self.tpl_eos.numel() ) - return int(frame + self._max_seq_length + self._num_levels) + return int(frame + self._max_seq_length * self._num_slots + self._num_levels) def _build_prompt_tokens( self, tokenizer: PreTrainedTokenizerBase, cfg: generative_model_pb2.Qwen2RecLM, ) -> None: - """Tokenise the family chat template once; cache as buffers. - - Composes the proto's optional ``system_instruction`` / - ``user_prefix_text`` / ``user_suffix_text`` with the family's static - fragments. Buffers are non-persistent: they move with ``model.to(...)`` - but stay off the state_dict so HF safetensors round-tripping isn't - polluted. + """Tokenise the static prompt once, as the N+1 gaps around the N slots. + + The base resolves ``{{feature_name}}`` slots; this hook only frames them + with ChatML and folds each slot feature's ``prefix_text`` / ``suffix_text`` + into the adjacent gap, so every gap is ONE contiguous string tokenized in + one call. That keeps the encoding bit-identical to tokenizing the fully + rendered prompt -- splitting the token ids instead would let a BPE merge + span a seam. Splicing values between gaps is exact because the ``C*`` + atoms are added-vocab tokens, which HF fast tokenizers pre-split on. + + Buffers are non-persistent: they move with ``model.to(...)`` but stay off + the state_dict so HF safetensors round-tripping isn't polluted. """ tpl = type(self).CHAT_TEMPLATE - sys_text = cfg.system_instruction or tpl["default_system_instruction"] - frags = { - "system": tpl["system_prefix"] + sys_text + tpl["system_suffix"], - "user_prefix": tpl["user_prefix"] + (cfg.user_prefix_text or ""), - "user_suffix": (cfg.user_suffix_text or "") + tpl["user_suffix"], - "asst_prefix": tpl["asst_prefix"], - "asst_suffix": tpl["asst_suffix"], - } - for slot_name, frag_str in frags.items(): + gaps, features, groups = self._resolve_prompt_slots(cfg.prompt_template) + self._slot_names = [f.name for f in features] + self._slot_groups = groups + self._num_slots = len(features) + for i, gap in enumerate(gaps): + head = tpl["user_prefix"] if i == 0 else features[i - 1].suffix_text + tail = ( + features[i].prefix_text + if i < self._num_slots + else tpl["user_suffix"] + tpl["asst_prefix"] + ) # explicit <|im_start|> markers frame the prompt; no auto BOS/EOS. ids = torch.tensor( - tokenizer.encode(frag_str, add_special_tokens=False), dtype=torch.long + tokenizer.encode(head + gap + tail, add_special_tokens=False), + dtype=torch.long, ) - self.register_buffer(f"tpl_{slot_name}", ids, persistent=False) + self.register_buffer(f"tpl_gap_{i}", ids, persistent=False) + # closes the assistant turn after the answer; not part of any gap. + self.register_buffer( + "tpl_asst_suffix", + torch.tensor( + tokenizer.encode(tpl["asst_suffix"], add_special_tokens=False), + dtype=torch.long, + ), + persistent=False, + ) # the trailing eos is a SUPERVISED token; cache it for the splice. self.register_buffer( "tpl_eos", @@ -122,27 +133,31 @@ def _build_prompt_tokens( persistent=False, ) - def _prompt_rows(self, user_seq_rows: List[torch.Tensor]) -> List[torch.Tensor]: - """Per-row ``[system | user_prefix | history | user_suffix | asst_prefix]``. + def _prompt_rows(self, slot_rows: List[List[torch.Tensor]]) -> List[torch.Tensor]: + """Per-row ``[gap_0 | slot_0 | gap_1 | ... | slot_N-1 | gap_N]``. Shared by the teacher-forced splice and the answer-less inference prompt. """ - return [ - torch.cat( - [ - self.tpl_system, - self.tpl_user_prefix, - row, - self.tpl_user_suffix, - self.tpl_asst_prefix, - ] - ) - for row in user_seq_rows - ] + gaps = [getattr(self, f"tpl_gap_{i}") for i in range(self._num_slots + 1)] + rows = [] + for b in range(len(slot_rows[0])): + parts: List[torch.Tensor] = [] + for i in range(self._num_slots): + parts.append(gaps[i]) + parts.append(slot_rows[i][b]) + parts.append(gaps[self._num_slots]) + rows.append(torch.cat(parts)) + return rows + + def _slot_rows( + self, rows: Dict[str, List[torch.Tensor]] + ) -> List[List[torch.Tensor]]: + """Per-slot token rows, in template order.""" + return [rows[name] for name in self._slot_names] def _splice_input_ids( self, - user_seq_rows: List[torch.Tensor], + slot_rows: List[List[torch.Tensor]], label_rows: List[torch.Tensor], pad_to: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -158,14 +173,14 @@ def _splice_input_ids( ``asst_suffix``. ``pad_to`` left-extends every row for pool pre-sizing, keeping the supervised tail end-aligned. """ - if len(user_seq_rows) != len(label_rows): + if len(slot_rows[0]) != len(label_rows): raise ValueError( f"{type(self).__name__}: history/answer row count mismatch " - f"({len(user_seq_rows)} vs {len(label_rows)})." + f"({len(slot_rows[0])} vs {len(label_rows)})." ) rows_ids = [ torch.cat([prompt, label_rows[i], self.tpl_asst_suffix, self.tpl_eos]) - for i, prompt in enumerate(self._prompt_rows(user_seq_rows)) + for i, prompt in enumerate(self._prompt_rows(slot_rows)) ] input_ids, attention_mask = self._left_pad(rows_ids, pad_to=pad_to) @@ -196,7 +211,7 @@ def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: self._pool_warmed = True input_ids, labels, attention_mask = self._splice_input_ids( - rows[self._input_name], rows[self._label_name], pad_to=pad_to + self._slot_rows(rows), rows[self._label_name], pad_to=pad_to ) return self._forward_loss(input_ids, labels, attention_mask) @@ -229,8 +244,8 @@ def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: is ``num_return_sequences`` on the HF path and the escalating beam's final width when ``dynamic_beam`` is set. """ - u_rows = self.build_input(batch)[self._input_name] - input_ids, attention_mask = self._splice_prompt_ids(u_rows) + slot_rows = self._slot_rows(self.build_input(batch)) + input_ids, attention_mask = self._splice_prompt_ids(slot_rows) if self._dynamic_beam: new_tokens = self._dynamic_beam_search(input_ids, attention_mask) else: @@ -262,10 +277,10 @@ def _dynamic_beam_search( ) def _splice_prompt_ids( - self, user_seq_rows: List[torch.Tensor] + self, slot_rows: List[List[torch.Tensor]] ) -> Tuple[torch.Tensor, torch.Tensor]: """Assemble the answer-less prompt and left-pad into ``(B, T_max)``.""" - return self._left_pad(self._prompt_rows(user_seq_rows)) + return self._left_pad(self._prompt_rows(slot_rows)) def _left_pad( self, rows: List[torch.Tensor], pad_to: int = 0 diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/qwen2_rec_lm_test.py index 1231fecee..85623d92f 100644 --- a/tzrec/models/qwen2_rec_lm_test.py +++ b/tzrec/models/qwen2_rec_lm_test.py @@ -36,20 +36,22 @@ def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): m._pad_token_id = pad_id m._dynamic_beam = False m._max_seq_length = 0 + m._num_slots = 1 + m._slot_names = ["user_sequence"] m._generated_sids_key = "generated_sids" m.lm = types.SimpleNamespace(device=torch.device(device)) for name, vals in { - "tpl_system": [10, 11], - "tpl_user_prefix": [12], - "tpl_user_suffix": [13], - "tpl_asst_prefix": [14], + "tpl_gap_0": [10, 11, 12], + "tpl_gap_1": [13, 14], "tpl_asst_suffix": [15], "tpl_eos": [9], }.items(): m.register_buffer(name, torch.tensor(vals, dtype=torch.long), persistent=False) - offsets, sizes = Qwen2RecLM._sid_level_layout(codebook) - m.register_buffer("_level_offsets", offsets, persistent=False) + sizes = torch.tensor(codebook, dtype=torch.long) m.register_buffer("_codebook_sizes", sizes, persistent=False) + m.register_buffer( + "_level_offsets", torch.cumsum(sizes, 0) - sizes, persistent=False + ) return m @@ -78,17 +80,38 @@ def _real_lm_stub(codebook=None, base_vocab=20, num_beams=2): ) torch.manual_seed(0) m.lm = Qwen2ForCausalLM(cfg).eval() - offsets, sizes = Qwen2RecLM._sid_level_layout(codebook) - m.register_buffer("_level_offsets", offsets, persistent=False) + sizes = torch.tensor(codebook, dtype=torch.long) m.register_buffer("_codebook_sizes", sizes, persistent=False) + m.register_buffer( + "_level_offsets", torch.cumsum(sizes, 0) - sizes, persistent=False + ) return m +def _wire_slots(m, prefix_text="", suffix_text=""): + """Minimal _features/_feature_groups so the base slot resolver can run.""" + from tzrec.protos import model_pb2 + + m._features = [ + types.SimpleNamespace( + name="user_sequence", prefix_text=prefix_text, suffix_text=suffix_text + ) + ] + m._feature_groups = [ + types.SimpleNamespace( + group_name="sids", + feature_names=["user_sequence"], + group_type=model_pb2.JAGGED_SEQUENCE, + ) + ] + + def _train_stub(max_total_len): """A ``_predict_train``-ready stub; returns ``(model, spliced T per step)``.""" m = _stub() m._is_inference = False # not inference + nn.Module.training=True -> is_train m._input_name, m._label_name = "user_sequence", "label" + m._slot_names, m._num_slots = ["user_sequence"], 1 m._max_total_len = max_total_len m._pool_warmed = False seen_lens = [] @@ -110,8 +133,8 @@ def test_splice_layout_and_labels(self) -> None: m = _stub() u = [torch.tensor([100, 101, 102])] a = [torch.tensor([200, 201, 202])] # 3 codes = num_levels - ids, labels, mask = m._splice_input_ids(u, a) - # sys|user_prefix|history|user_suffix|asst_prefix|answer|asst_suffix|eos + ids, labels, mask = m._splice_input_ids([u], a) + # head | history | tail | answer | asst_suffix | eos self.assertEqual( ids[0].tolist(), [10, 11, 12, 100, 101, 102, 13, 14, 200, 201, 202, 15, 9] ) @@ -126,7 +149,7 @@ def test_left_padding_varied_lengths(self) -> None: m = _stub() u = [torch.tensor([100, 101, 102, 103]), torch.tensor([100])] a = [torch.tensor([200, 201, 202]), torch.tensor([207, 208, 209])] - ids, labels, mask = m._splice_input_ids(u, a) + ids, labels, mask = m._splice_input_ids([u], a) T = ids.shape[1] n1 = 2 + 1 + 1 + 1 + 1 + 3 + 1 + 1 # shorter row's real length self.assertEqual(ids[1, : T - n1].tolist(), [m._pad_token_id] * (T - n1)) @@ -138,7 +161,7 @@ def test_left_padding_varied_lengths(self) -> None: def test_mask_keeps_trailing_eos_when_pad_equals_eos(self) -> None: m = _stub(pad_id=9) # tpl_eos == 9 too ids, _, mask = m._splice_input_ids( - [torch.tensor([100])], [torch.tensor([200, 201, 202])] + [[torch.tensor([100])]], [torch.tensor([200, 201, 202])] ) self.assertEqual(int(ids[0, -1]), 9) self.assertEqual(int(mask[0, -1]), 1) @@ -149,7 +172,7 @@ def test_suffix_keep_matches_dynamic_slice(self) -> None: m = _stub() suffix_keep = 6 # num_levels 3 + asst_suffix 1 + trailing eos + HF shift _, labels, _ = m._splice_input_ids( - [torch.tensor([100, 101, 102])], [torch.tensor([200, 201, 202])] + [[torch.tensor([100, 101, 102])]], [torch.tensor([200, 201, 202])] ) # first supervised column, minimized over rows first_sup = int(((labels >= 0).cumsum(-1) == 1).float().argmax(-1).min()) @@ -158,8 +181,8 @@ def test_suffix_keep_matches_dynamic_slice(self) -> None: def test_splice_prompt_ids(self) -> None: m = _stub() - ids, mask = m._splice_prompt_ids([torch.tensor([100, 101, 102])]) - # [system | user_prefix | history | user_suffix | asst_prefix], no answer + ids, mask = m._splice_prompt_ids([[torch.tensor([100, 101, 102])]]) + # [head | history | tail], no answer self.assertEqual(ids[0].tolist(), [10, 11, 12, 100, 101, 102, 13, 14]) self.assertEqual(mask[0].tolist(), [1] * 8) @@ -176,17 +199,17 @@ def test_predict_routes_on_inference_flag(self) -> None: # are each level's min/max token; every malformed row collapses to -1. @parameterized.expand( [ - [[[100, 102, 105], [101, 104, 108]], [[1, 1, 1], [2, 3, 4]]], + [[[100, 102, 105], [101, 104, 108]], [[0, 0, 0], [1, 2, 3]]], [ [ - [100, 102, 105], # valid -> local [1, 1, 1] - [101, 104, 108], # valid -> local [2, 3, 4] + [100, 102, 105], # valid -> local [0, 0, 0] + [101, 104, 108], # valid -> local [1, 2, 3] [102, 102, 105], # pos0 above level-0 band [100, 101, 105], # pos1 below level-1 band [100, 102, 109], # pos2 above level-2 band [100, 104, 9], # pos2 = eos/pad token (sid -90) -> invalid ], - [[1, 1, 1], [2, 3, 4]] + [[-1, -1, -1]] * 4, + [[0, 0, 0], [1, 2, 3]] + [[-1, -1, -1]] * 4, ], # early EOS: a tail narrower than num_levels still reshapes cleanly, # and the missing 3rd atom stays -1 -> out of band -> candidate -1 @@ -197,6 +220,7 @@ def test_predict_routes_on_inference_flag(self) -> None: def test_generate_maps_tokens_to_sids(self, tail, expected) -> None: m = _stub(base_vocab=100) m._input_name = "user_sequence" + m._slot_names, m._num_slots = ["user_sequence"], 1 m._num_beams, m._num_return = len(tail) + 1, len(tail) seen = {} @@ -239,6 +263,7 @@ def fake_generate( def test_generate_routes_to_the_dynamic_beam(self) -> None: m = _stub(base_vocab=100) m._input_name = "user_sequence" + m._slot_names, m._num_slots = ["user_sequence"], 1 m._dynamic_beam = True m._num_beams = m._num_return = 2 m.lm.generate = lambda **kw: self.fail("HF generate must not run") @@ -254,7 +279,7 @@ def fake_beam(ids, am): self.assertEqual(seen["ids"], [10, 11, 12, 100, 102, 105, 13, 14]) self.assertEqual(seen["mask"], [1] * 8) self.assertEqual(tuple(sids.shape), (1, 2, 3)) - self.assertEqual(sids[0].tolist(), [[1, 1, 1], [2, 3, 4]]) + self.assertEqual(sids[0].tolist(), [[0, 0, 0], [1, 2, 3]]) def _prompt_tokenizer(self): # encode -> [len(text)] makes each buffer a fingerprint of its fragment @@ -263,59 +288,43 @@ def _prompt_tokenizer(self): encode=lambda text, add_special_tokens=False: [len(text)], ) - def test_build_prompt_tokens_composes_the_family_template(self) -> None: + def test_build_prompt_tokens_splits_the_template_around_the_slot(self) -> None: m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) - cfg = types.SimpleNamespace( - system_instruction="", user_prefix_text="", user_suffix_text="" - ) + _wire_slots(m, prefix_text="PRE", suffix_text="SUF") + cfg = types.SimpleNamespace(prompt_template="A{{user_sequence}}B") m._build_prompt_tokens(self._prompt_tokenizer(), cfg) - for name in [ - "tpl_system", - "tpl_user_prefix", - "tpl_user_suffix", - "tpl_asst_prefix", - "tpl_asst_suffix", - "tpl_eos", - ]: + for name in ["tpl_gap_0", "tpl_gap_1", "tpl_asst_suffix", "tpl_eos"]: buf = getattr(m, name) self.assertIsInstance(buf, torch.Tensor) self.assertEqual(buf.dtype, torch.int64) tpl = Qwen2RecLM.CHAT_TEMPLATE + # head = user_prefix + before + feature.prefix_text + self.assertEqual(m.tpl_gap_0.tolist(), [len(tpl["user_prefix"] + "A" + "PRE")]) + # tail = feature.suffix_text + after + user_suffix + asst_prefix self.assertEqual( - m.tpl_system.tolist(), - [ - len( - tpl["system_prefix"] - + tpl["default_system_instruction"] - + tpl["system_suffix"] - ) - ], + m.tpl_gap_1.tolist(), + [len("SUF" + "B" + tpl["user_suffix"] + tpl["asst_prefix"])], ) - self.assertEqual(m.tpl_user_prefix.tolist(), [len(tpl["user_prefix"])]) - self.assertEqual(m.tpl_user_suffix.tolist(), [len(tpl["user_suffix"])]) - self.assertEqual(m.tpl_asst_prefix.tolist(), [len(tpl["asst_prefix"])]) - self.assertEqual(m.tpl_asst_suffix.tolist(), [len(tpl["asst_suffix"])]) self.assertEqual(m.tpl_eos.tolist(), [99]) # eos cached for supervision - def test_build_prompt_tokens_honours_proto_text_knobs(self) -> None: + def test_build_prompt_tokens_rejects_a_bad_placeholder_count(self) -> None: m = object.__new__(Qwen2RecLM) nn.Module.__init__(m) - cfg = types.SimpleNamespace( - system_instruction="SYS", user_prefix_text="UP", user_suffix_text="US" - ) - m._build_prompt_tokens(self._prompt_tokenizer(), cfg) - tpl = Qwen2RecLM.CHAT_TEMPLATE - self.assertEqual( - m.tpl_system.tolist(), - [len(tpl["system_prefix"] + "SYS" + tpl["system_suffix"])], - ) - self.assertEqual(m.tpl_user_prefix.tolist(), [len(tpl["user_prefix"] + "UP")]) - self.assertEqual(m.tpl_user_suffix.tolist(), [len("US" + tpl["user_suffix"])]) + _wire_slots(m) + cases = [ + ("no placeholder", "at least one"), + ("{{nope}} x", "names no feature_group"), + ] + for template, msg in cases: + with self.subTest(template=template): + cfg = types.SimpleNamespace(prompt_template=template) + with self.assertRaisesRegex(ValueError, msg): + m._build_prompt_tokens(self._prompt_tokenizer(), cfg) def test_compute_max_total_length(self) -> None: m = _stub() - # frame = 2 system + 1 each user_pfx/sfx, asst_pfx/sfx, eos = 7 + # frame = 3 head + 2 tail + 1 asst_suffix + 1 eos = 7 m._max_seq_length = 300 self.assertEqual(m._compute_max_total_length(), 7 + 300 + 3) m._max_seq_length = 0 # pre-allocation disabled @@ -338,7 +347,7 @@ def test_no_forced_padding_when_disabled(self) -> None: def test_splice_row_count_mismatch_raises(self) -> None: m = _stub() with self.assertRaisesRegex(ValueError, "row count mismatch"): - m._splice_input_ids([torch.tensor([100])], []) + m._splice_input_ids([[torch.tensor([100])]], []) class Qwen2ForwardLossTest(unittest.TestCase): @@ -349,16 +358,16 @@ def _model(self, ignore_index=-100): m._ignore_index = ignore_index m._pad_token_id = 0 for name, vals in { - "tpl_system": [1, 2], - "tpl_user_prefix": [3], - "tpl_user_suffix": [4], - "tpl_asst_prefix": [5], + "tpl_gap_0": [1, 2, 3], + "tpl_gap_1": [4, 5], "tpl_asst_suffix": [6], "tpl_eos": [7], }.items(): m.register_buffer( name, torch.tensor(vals, dtype=torch.long), persistent=False ) + m._num_slots = 1 + m._slot_names = ["user_sequence"] m._suffix_keep = 6 # num_levels 3 + asst_suffix 1 + trailing eos + HF shift return m @@ -366,7 +375,7 @@ def _rows(self): # ragged histories so left padding is exercised u = [torch.tensor([20, 22, 25]), torch.tensor([21, 23, 26, 20, 24, 27])] a = [torch.tensor([20, 22, 25]), torch.tensor([21, 24, 28])] - return u, a + return [u], a def test_suffix_slice_matches_full_sequence_loss(self) -> None: # the fixed-width suffix slice must give the same CE as full-T logits @@ -454,11 +463,11 @@ def test_band_masked_beams_decode_without_sentinels(self) -> None: sids = m._validate_sid_candidates(new, batch_size=1) self.assertEqual(tuple(sids.shape), (1, 24, 3)) for level, size in enumerate(codebook): - self.assertTrue(bool((sids[..., level] >= 1).all())) - self.assertTrue(bool((sids[..., level] <= size).all())) + self.assertTrue(bool((sids[..., level] >= 0).all())) + self.assertTrue(bool((sids[..., level] < size).all())) self.assertEqual( {tuple(row) for row in sids[0].tolist()}, - {(a, b, c) for a in range(1, 3) for b in range(1, 4) for c in range(1, 5)}, + {(a, b, c) for a in range(2) for b in range(3) for c in range(4)}, ) diff --git a/tzrec/protos/feature.proto b/tzrec/protos/feature.proto index 19cebd176..e42325bdf 100644 --- a/tzrec/protos/feature.proto +++ b/tzrec/protos/feature.proto @@ -986,6 +986,58 @@ message SequenceFeature { repeated SeqFeatureConfig features = 5; } +// Semantic-ID (SID) sequence feature for generative-recommendation LMs. +// +// Carries a flat stream of per-level SID codes (whole items, level order) plus +// the prompt text that wraps them. Only usable as `sequence_sid_feature`: the +// model reads it as a JAGGED_SEQUENCE group, so the scalar form would never +// produce the "{group}.sequence" keys it needs. +// +// FG note: SID codes are produced offline by a SID-generation model, so this +// feature is only supported with data_config.fg_mode = FG_NONE. There is no +// pyfg counterpart. +message SidFeature { + // feature name; also the {{name}} placeholder in the model prompt_template. + required string feature_name = 1; + // feature input, e.g. user:user_sequence + required string expression = 2; + // codes per position; SID streams are flat, so this stays 1. + optional uint32 value_dim = 6 [default = 1]; + // embedding pooling type, unused (SIDs carry no embedding table). + optional string pooling = 10 [default = "sum"]; + // fg default value + optional string default_value = 11 [default = "0"]; + // fg multi-value separator + optional string separator = 12 [default = "\x1d"]; + // mask value in training progress + optional bool use_mask = 14; + + // Text emitted immediately BEFORE this feature's SID tokens in the prompt, + // e.g. "Current user's historical behaviors are as follows:". Empty by + // default. Applies only where the model splices this feature. + optional string prefix_text = 20 [default = ""]; + // Text emitted immediately AFTER this feature's SID tokens. + optional string suffix_text = 21 [default = ""]; + // SID vocabulary, one entry per RQ level: each code is 0-based in + // [0, codebook[level]), matching what the SID-generation models emit. + // len() = codes per item; sum() = atoms the model appends as C0..C{sum-1}. + // All SID features in a model share ONE space, so every SidFeature must + // declare the SAME codebook -- adding a feature must not grow the vocabulary. + repeated uint32 codebook = 23; + + // default value when fg_mode = FG_NONE + optional string fg_encoded_default_value = 30; + // only used as fg dag intermediate result or not + optional bool stub_type = 34 [default = false]; + // embedding param constraints + optional ParameterConstraints embedding_constraints = 50; + + // max sequence length, only take effect when use it as sequence + optional uint32 sequence_length = 101; + // sequence delimiter, only take effect when use it as sequence + optional string sequence_delim = 102 [default = ";"]; +} + message FeatureConfig { oneof feature { IdFeature id_feature = 1; @@ -1014,6 +1066,7 @@ message FeatureConfig { KvDotProduct sequence_kv_dot_product = 111; BoolMaskFeature sequence_bool_mask_feature = 112; CombineFeature sequence_combine_feature = 113; + SidFeature sequence_sid_feature = 114; } } diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index 4867bcb6c..a5f4df9a5 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -9,19 +9,17 @@ package tzrec.protos; // Architecture-agnostic config shared by all generative-rec families. // Sample contract: -// * history : list -- local 1-based per-level codes in -// [1, codebook[level]], laid out as whole items in level order; +// * history : list -- local 0-based per-level codes in +// [0, codebook[level]), laid out as whole items in level order; // the single JAGGED_SEQUENCE feature_group (its one member). -// * answer : list -- one item's local 1-based codes; the FIRST +// * answer : list -- one item's local 0-based codes; the FIRST // `data_config.label_field`, NOT a feature. -// The model derives level_offsets from codebook and maps each code at batch -// time: token_id = base_vocab + level_offsets[level] + code - 1. Sample +// SidFeature folds in level_offsets at parse time and the model adds +// base_vocab, so token_id = base_vocab + level_offsets[level] + code. Sample // writers must not pre-apply level_offsets. message GenerativeRecLMConfig { - // SID vocabulary, one entry per RQ level: each public code is in - // [1, codebook[level]]. len = codes per item (answer width); sum = atoms - // appended as C0..C{sum-1} after the base vocab. - repeated uint32 codebook = 2; + // NOTE: the SID vocabulary (`codebook`) is declared on the SidFeature that + // carries the codes, not here -- the model asks the feature for it. // Pad the post-extension vocab up to a multiple of this value; 0 disables // padding. optional uint32 vocab_pad_to_multiple_of = 3 [default = 128]; @@ -61,12 +59,13 @@ message Qwen2RecLM { // Qwen2 backbone: HF hub id or local path; must be a Qwen2 model. optional string hf_model_id = 2 [default = "Qwen/Qwen2.5-0.5B"]; - // ----- Chat template ----- - // Optional override for the system instruction. Empty -> the Qwen2 default. - optional string system_instruction = 10 [default = ""]; - // Optional text wrapping the SID codes in the user message, e.g. - // "The user's history is as follows: " / ", predict the semantic id of the - // next item the user will interact with". - optional string user_prefix_text = 11 [default = ""]; - optional string user_suffix_text = 12 [default = ""]; + // ----- Prompt ----- + // Prompt body for the user turn, carrying exactly one {{feature_name}} + // placeholder that names the SID feature whose codes are spliced in: + // "... Each behavior is represented by three words. {{user_sequence}} + // Please predict the semantic encoding of the next behavior." + // That feature's own prefix_text/suffix_text wrap the codes inside the slot. + // The ChatML frame and the post-answer region stay family constants -- the + // model owns <|im_start|>/<|im_end|> and the supervised tail. + optional string prompt_template = 10; } diff --git a/tzrec/tests/configs/qwen2_rec_lm_mock.config b/tzrec/tests/configs/qwen2_rec_lm_mock.config index 14efb520c..30393d0da 100644 --- a/tzrec/tests/configs/qwen2_rec_lm_mock.config +++ b/tzrec/tests/configs/qwen2_rec_lm_mock.config @@ -33,11 +33,13 @@ data_config { fg_mode: FG_NONE } feature_configs { - sequence_raw_feature { + sequence_sid_feature { feature_name: "user_sequence" expression: "user:user_sequence" - value_dim: 1 - sequence_length: 12 + prefix_text: "Current user's historical behaviors are as follows:" + codebook: 4 + codebook: 4 + codebook: 4 } } model_config { @@ -48,17 +50,12 @@ model_config { } qwen2_rec_lm { common { - codebook: 4 - codebook: 4 - codebook: 4 vocab_pad_to_multiple_of: 128 ignore_index: -100 param_dtype: "float32" max_sequence_length: 12 } hf_model_id: "Qwen/Qwen2.5-0.5B" - system_instruction: "你是一个推荐系统，根据用户的历史行为，预测用户在电商场景的下一步行为。我会给你一串连续行为的语义编码，按照用户点击的时间顺序排列，每个行为用三个词表示。" - user_prefix_text: "当前用户的历史行为如下：" - user_suffix_text: "，请预测用户在电商推荐场景后续行为的语义编码" + prompt_template: "You are a recommendation system. Based on the user's historical behavior, predict the user's next action in an e-commerce scenario. I will provide a sequence of semantic encodings representing consecutive behaviors, arranged in chronological order of user clicks. Each behavior is represented by three words. {{user_sequence}} Please predict the semantic encoding of the user's subsequent behavior in the e-commerce recommendation scenario." } } diff --git a/tzrec/tests/genrec_integration_test.py b/tzrec/tests/genrec_integration_test.py index f38460c4d..e7507150f 100644 --- a/tzrec/tests/genrec_integration_test.py +++ b/tzrec/tests/genrec_integration_test.py @@ -69,13 +69,13 @@ def _write_backbone(save_dir: str, vocab_size: int = 256) -> str: def _write_samples(save_dir: str, num_rows: int, seed: int = 0) -> str: """Write the two-column sample contract: history + answer, both list. - Codes are local 1-based per-level values in ``[1, codebook[level]]`` and - every row holds whole items in level order. + Codes are local 0-based per-level values in ``[0, codebook[level])``, as the + SID-generation models emit them; every row holds whole items in level order. """ rnd = random.Random(seed) def _item(): - return [rnd.randint(1, size) for size in _CODEBOOK] + return [rnd.randrange(size) for size in _CODEBOOK] schema = pa.schema( [ @@ -144,8 +144,8 @@ def test_mock_config_builds_the_model_and_runs_a_batch(self) -> None: config.model_config, features, list(config.data_config.label_fields) ) self.assertEqual(type(model).__name__, "Qwen2RecLM") - self.assertEqual(model._history_group, "sids") - self.assertEqual(model._input_name, "user_sequence") + self.assertEqual(model._slot_groups, ["sids"]) + self.assertEqual(model._slot_names, ["user_sequence"]) self.assertEqual(model._label_name, "label") self.assertEqual(model._num_levels, len(_CODEBOOK)) # base vocab + sum(codebook) atoms, padded to vocab_pad_to_multiple_of From f34975a25110476547a5d6fa68f976d728205ad7 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Tue, 28 Jul 2026 08:54:40 +0000 Subject: [PATCH 40/99] [bugfix] genrec: reject a SID feature knob that silently does nothing sequence_length on a SidFeature reads as a truncation cap but has no truncating consumer. The fg paths that would apply it are exactly the ones this feature forbids (it is FG_NONE-only and its _fg_json raises), and its only other reader, EmbeddingGroup, forwards it to fx_mark_seq_tensor -- a torch.fx.wrap-decorated no-op that annotates the export graph. Setting it therefore left the history uncapped, with no error. It is now rejected at construction and the message points at model_config.common.max_sequence_length, which is the knob that actually truncates. Also restores an em-dash in CudaAutocastWrapper's docstring that this fork had rewritten for no reason, keeping the upstream overlap minimal. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/features/sid_feature.py | 67 ++++++++++--------- tzrec/features/sid_feature_test.py | 7 +- ...nerative_rec_lm.py => generative_model.py} | 0 ...ec_lm_test.py => generative_model_test.py} | 0 tzrec/models/model.py | 2 +- .../{escalating_beam.py => dynamic_beam.py} | 0 ...ting_beam_test.py => dynamic_beam_test.py} | 0 tzrec/protos/feature.proto | 4 +- 8 files changed, 44 insertions(+), 36 deletions(-) rename tzrec/models/{generative_rec_lm.py => generative_model.py} (100%) rename tzrec/models/{generative_rec_lm_test.py => generative_model_test.py} (100%) rename tzrec/modules/{escalating_beam.py => dynamic_beam.py} (100%) rename tzrec/modules/{escalating_beam_test.py => dynamic_beam_test.py} (100%) diff --git a/tzrec/features/sid_feature.py b/tzrec/features/sid_feature.py index 085a85457..684b3e4c3 100644 --- a/tzrec/features/sid_feature.py +++ b/tzrec/features/sid_feature.py @@ -22,11 +22,9 @@ class SidFeature(BaseFeature): """Semantic-ID sequence feature. - Carries a flat stream of per-level SID codes -- whole items in level order -- - plus the prompt text that wraps them. Codes are 0-based per level, as the - SID-generation models emit them. They are produced offline by a - SID-generation model, so only ``fg_mode = FG_NONE`` is supported; there is no - pyfg counterpart for this feature type. + A flat stream of 0-based per-level SID codes -- whole items in level order -- + plus the prompt text wrapping them. Generated offline, so only + ``fg_mode = FG_NONE`` works; there is no pyfg counterpart. Args: feature_config (FeatureConfig): a instance of feature config. @@ -50,18 +48,23 @@ def __init__( f"offline by a SID model), got {fg_mode}." ) super().__init__(feature_config, **kwargs) + # only fg (which this feature forbids) and an fx export marker read it, + # so it would cap nothing; the model owns the real, item-aligned budget. + if self.config.HasField("sequence_length"): + raise ValueError( + f"{self.__class__.__name__}[{self.config.feature_name}]: " + f"sequence_length does not truncate a SID feature; set " + f"model_config.common.max_sequence_length instead." + ) self._codebook = self._read_codebook() - self._level_offsets = np.cumsum([0] + self._codebook[:-1]).tolist() - # per-code level index and bound, tiled at parse time to the row length self._level_sizes = np.asarray(self._codebook) - self._level_offset_arr = np.asarray(self._level_offsets) + self._level_offsets = np.cumsum(self._level_sizes) - self._level_sizes def _read_codebook(self) -> List[int]: """Validate the declared codebook once and normalize it to a list. - The proto field is a repeated-scalar container, and every derived - quantity (level count, vocab size, offsets) reads it on the parse hot - path, so it is checked and converted here rather than per access. + Every derived quantity reads it on the parse hot path, so the repeated + scalar container is checked and converted here, not per access. """ codebook = [int(c) for c in self.config.codebook] if not codebook: @@ -83,11 +86,7 @@ def value_dim(self) -> int: @property def output_dim(self) -> int: - """Output dimension of the feature after embedding. - - SID codes are token ids consumed by the LM's own embedding table, so the - feature carries no embedding of its own and passes the codes through. - """ + """Output dimension: SID codes pass through to the LM's own table.""" return self.value_dim @property @@ -126,7 +125,7 @@ def sid_vocab_size(self) -> int: @property def level_offsets(self) -> List[int]: """Flat offset of each level, i.e. ``cumsum(sizes) - sizes``.""" - return self._level_offsets + return self._level_offsets.tolist() def _build_side_inputs(self) -> Optional[List[Tuple[str, str]]]: """Input field names with side.""" @@ -138,31 +137,33 @@ def _build_side_inputs(self) -> Optional[List[Tuple[str, str]]]: def _parse(self, input_data: Dict[str, pa.Array]) -> ParsedData: """Parse the SID stream into flat indices in the shared space. - Codes are 0-based, exactly as the SID-generation models emit them, so - the flat index IS the atom index and the model only adds ``base_vocab``. - The per-level offsets are folded in here, in the dataloader workers, - rather than on the model's forward path; validating here also keeps a - malformed row off the collective path, where one rank raising becomes - an all-reduce hang on its peers. + Codes are 0-based, so the flat index IS the atom index and the model only + adds ``base_vocab``. Offsets are folded in here, in the dataloader + workers, not on the forward path -- and validating here keeps a malformed + row off the collective path, where one rank raising hangs its peers. """ parsed = super()._parse(input_data) num_levels = len(self._codebook) - values = parsed.values.reshape(-1) - if values.size % num_levels != 0: + bad = np.nonzero(parsed.seq_lengths % num_levels)[0] + if bad.size: raise ValueError( - f"{self.__class__.__name__}[{self.config.feature_name}]: SID " - f"stream must hold whole {num_levels}-level items, got " - f"{values.size} codes." + f"{self.__class__.__name__}[{self.config.feature_name}]: every " + f"row must hold whole {num_levels}-level items; rows " + f"{bad.tolist()[:10]} have lengths " + f"{parsed.seq_lengths[bad].tolist()[:10]}." ) - levels = np.arange(values.size) % num_levels - sizes = self._level_sizes[levels] - if ((values < 0) | (values >= sizes)).any(): + # rows are whole items, so (-1, num_levels) lines every column up with + # its level and the per-level bounds/offsets broadcast down it. + codes = parsed.values.reshape(-1, num_levels) + if ((codes < 0) | (codes >= self._level_sizes)).any(): raise ValueError( f"{self.__class__.__name__}[{self.config.feature_name}]: SID " f"codes must be local 0-based values in [0, codebook[level])." ) - offsets = self._level_offset_arr[levels] - parsed.values = (values + offsets).reshape(parsed.values.shape) + # keep the value dtype: float32 + int64 offsets would promote to float64 + # and double the bytes crossing worker IPC and the H2D copy. + offsets = self._level_offsets.astype(codes.dtype, copy=False) + parsed.values = (codes + offsets).reshape(parsed.values.shape) return parsed def _fg_json(self) -> List[Dict[str, Any]]: diff --git a/tzrec/features/sid_feature_test.py b/tzrec/features/sid_feature_test.py index 4b8314ea3..4f72cae21 100644 --- a/tzrec/features/sid_feature_test.py +++ b/tzrec/features/sid_feature_test.py @@ -89,7 +89,12 @@ def test_rejects_a_bad_codebook(self) -> None: with self.subTest(bad=bad): base = 'feature_name: "s" expression: "user:s" ' + bad with self.assertRaisesRegex(ValueError, msg): - _ = _feature(base).codebook + _feature(base) + + def test_rejects_sequence_length(self) -> None: + # it caps nothing here, so accepting it would read as a working budget + with self.assertRaisesRegex(ValueError, "max_sequence_length instead"): + _feature(f"{_BASE} sequence_length: 64") def test_no_embedding_table(self) -> None: f = _feature(_BASE) diff --git a/tzrec/models/generative_rec_lm.py b/tzrec/models/generative_model.py similarity index 100% rename from tzrec/models/generative_rec_lm.py rename to tzrec/models/generative_model.py diff --git a/tzrec/models/generative_rec_lm_test.py b/tzrec/models/generative_model_test.py similarity index 100% rename from tzrec/models/generative_rec_lm_test.py rename to tzrec/models/generative_model_test.py diff --git a/tzrec/models/model.py b/tzrec/models/model.py index a18c47a41..8a7ae1d35 100644 --- a/tzrec/models/model.py +++ b/tzrec/models/model.py @@ -420,7 +420,7 @@ class CudaAutocastWrapper(nn.Module): proper dtype casts. CUTLASS HSTU attention requires bf16/fp16 inputs. When ``device`` is set, it is passed as a second positional argument - to ``inner.forward(x, device)`` -- this binds the device for models + to ``inner.forward(x, device)`` — this binds the device for models like ``ScriptWrapper`` whose forward takes ``(data, device)``. ``_mixed_dtype_id: Final[int]`` encodes the dtype so that diff --git a/tzrec/modules/escalating_beam.py b/tzrec/modules/dynamic_beam.py similarity index 100% rename from tzrec/modules/escalating_beam.py rename to tzrec/modules/dynamic_beam.py diff --git a/tzrec/modules/escalating_beam_test.py b/tzrec/modules/dynamic_beam_test.py similarity index 100% rename from tzrec/modules/escalating_beam_test.py rename to tzrec/modules/dynamic_beam_test.py diff --git a/tzrec/protos/feature.proto b/tzrec/protos/feature.proto index e42325bdf..cc292d687 100644 --- a/tzrec/protos/feature.proto +++ b/tzrec/protos/feature.proto @@ -1032,7 +1032,9 @@ message SidFeature { // embedding param constraints optional ParameterConstraints embedding_constraints = 50; - // max sequence length, only take effect when use it as sequence + // NOT a truncation cap here -- fg (which a SID feature forbids) and an + // fx export marker are its only readers, so setting it is rejected. + // Use model_config.common.max_sequence_length for the history budget. optional uint32 sequence_length = 101; // sequence delimiter, only take effect when use it as sequence optional string sequence_delim = 102 [default = ";"]; From 1576c16712fba3581e67958e93d909b7e7e6e535 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Tue, 28 Jul 2026 08:55:10 +0000 Subject: [PATCH 41/99] [refactor] genrec: name the LM family after the SID family convention GenerativeRecLM/QwenRecLM matched neither base-class pattern in tzrec and used Rec and LM tokens no other model carries. They now mirror the only comparable family -- BaseSidModel in sid_model.py with SidRqvae/SidRqkmeans beside it -- as BaseGenerativeModel in generative_model.py with GenerativeQwen beside it, so the module finally shares a name with the generative_model.proto that drives it and the family is greppable by one prefix. Oneof tag 700 is unchanged, but the message rename is a config break: a qwen_rec_lm block must become generative_qwen. Also guards two config values that silently did nothing: a max_sequence_length narrower than one item floored to zero whole items and left the history uncapped, and a feature named by two feature_groups resolved to whichever came last instead of being rejected. Both now raise, and the escalating-beam module is renamed to dynamic_beam to match the dynamic_beam proto field it backs. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/generative_model.py | 192 +++++------ tzrec/models/generative_model_test.py | 318 ++++++++++-------- .../{qwen2_rec_lm.py => generative_qwen.py} | 141 ++++---- ...rec_lm_test.py => generative_qwen_test.py} | 191 ++++------- tzrec/modules/dynamic_beam.py | 23 +- tzrec/modules/dynamic_beam_test.py | 98 ++---- tzrec/protos/export.proto | 2 +- tzrec/protos/model.proto | 2 +- tzrec/protos/models/generative_model.proto | 45 ++- ...ock.config => generative_qwen_mock.config} | 6 +- tzrec/tests/genrec_integration_test.py | 41 ++- tzrec/tests/genrec_test_util.py | 52 +++ tzrec/utils/hf_export_util.py | 31 +- tzrec/utils/hf_export_util_test.py | 34 +- 14 files changed, 581 insertions(+), 595 deletions(-) rename tzrec/models/{qwen2_rec_lm.py => generative_qwen.py} (69%) rename tzrec/models/{qwen2_rec_lm_test.py => generative_qwen_test.py} (76%) rename tzrec/tests/configs/{qwen2_rec_lm_mock.config => generative_qwen_mock.config} (94%) create mode 100644 tzrec/tests/genrec_test_util.py diff --git a/tzrec/models/generative_model.py b/tzrec/models/generative_model.py index 2b85006d1..355ee9333 100644 --- a/tzrec/models/generative_model.py +++ b/tzrec/models/generative_model.py @@ -9,12 +9,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Generic generative-recommendation language-model base for TorchEasyRec. +"""Architecture-agnostic base for HF-backed generative-recommendation LMs. -``GenerativeRecLM`` owns the architecture-agnostic plumbing; a concrete family -subclass (e.g. ``Qwen2RecLM``) implements ``_build_prompt_tokens`` and -``predict``. Shared config and the sample contract live in -``GenerativeRecLMConfig`` (see ``protos/models/generative_model.proto``). +A family subclass (e.g. ``GenerativeQwen``) supplies ``_build_prompt_tokens`` and +``predict``; ``GenerativeModelConfig`` holds the shared config and the sample +contract. """ import re @@ -33,6 +32,7 @@ from tzrec.datasets.utils import Batch from tzrec.features.feature import BaseFeature +from tzrec.features.sid_feature import SidFeature from tzrec.models.model import BaseModel from tzrec.modules.embedding import EmbeddingGroup from tzrec.protos import model_pb2 @@ -40,20 +40,19 @@ from tzrec.protos.models import generative_model_pb2 -class GenerativeRecLM(BaseModel): - """Abstract base for HF-backed generative-recommendation LMs. +class BaseGenerativeModel(BaseModel): + """Model construction, SID vocab extension, data-prep, loss and metrics. - Owns model construction, SID vocab extension, sample data-prep, loss and - metrics; subclasses implement ``_build_prompt_tokens`` and ``predict``. The - family's proto message must supply ``common`` (``GenerativeRecLMConfig``) - and an ``hf_model_id`` field, both read directly by this base. + The family's proto message must carry ``common`` (a + ``GenerativeModelConfig``) and ``hf_model_id``; this base reads both. """ - # See `common.param_dtype` in the proto for why fp32 is the default. - _DTYPE_BY_NAME = { - "float32": torch.float32, - "bfloat16": torch.bfloat16, - "float16": torch.float16, + # See `common.param_dtype` in the proto for why FP32 is the default. + # The enum is closed, so protobuf rejects anything not listed here. + _PARAM_DTYPE: Dict[int, torch.dtype] = { + generative_model_pb2.FP32: torch.float32, + generative_model_pb2.BF16: torch.bfloat16, + generative_model_pb2.FP16: torch.float16, } def __init__( @@ -85,28 +84,31 @@ def _resolve_pad_token_id(tokenizer: PreTrainedTokenizerBase) -> int: pad_id = tokenizer.eos_token_id if pad_id is None: raise ValueError( - "GenerativeRecLM: tokenizer has neither pad_token_id nor " + "BaseGenerativeModel: tokenizer has neither pad_token_id nor " "eos_token_id; cannot choose a pad id for the left-padded splice." ) return int(pad_id) def _read_common_config( - self, common: generative_model_pb2.GenerativeRecLMConfig + self, common: generative_model_pb2.GenerativeModelConfig ) -> int: """Parse shared proto knobs into attributes; return the SID atom count.""" self._label_name: str = self._labels[0] if self._labels else "" self._ignore_index: int = int(common.ignore_index) self._generated_sids_key: str = common.generated_sids_key - param_dtype = self._DTYPE_BY_NAME.get(common.param_dtype) - if param_dtype is None: - raise ValueError( - f"{type(self).__name__}: param_dtype must be one of " - f"{list(self._DTYPE_BY_NAME)}, got {common.param_dtype!r}." - ) - self._param_dtype: torch.dtype = param_dtype + self._param_dtype: torch.dtype = self._PARAM_DTYPE[common.param_dtype] self._max_seq_length: int = int(common.max_sequence_length) codebook = self._shared_sid_space() self._num_levels = len(codebook) + # the budget truncates to WHOLE items, so anything under one item's width + # would floor to zero and silently leave the history uncapped. + if 0 < self._max_seq_length < self._num_levels: + raise ValueError( + f"{type(self).__name__}: max_sequence_length " + f"({self._max_seq_length}) cannot hold one {self._num_levels}" + f"-level item; use 0 to disable the budget or a multiple of " + f"{self._num_levels}." + ) sizes = torch.tensor(codebook, dtype=torch.long) self.register_buffer("_codebook_sizes", sizes, persistent=False) # only the decode path still needs per-level offsets: SidFeature folds @@ -121,16 +123,14 @@ def _read_common_config( def _shared_sid_space(self) -> List[int]: """The one codebook every SID feature declares. - SID features share a single space: the model has one extended vocabulary - and one answer width, so adding a feature must not resize ``lm_head``. - Requiring the declarations to agree makes that invariant explicit and - turns a typo into an error instead of a silent reshape. + One extended vocabulary and one answer width, so adding a feature must + never resize ``lm_head``; disagreement is a typo, not a reshape. """ - spaces = {} - for feature in self._features: - codebook = getattr(feature, "codebook", None) - if codebook is not None: - spaces[feature.name] = tuple(codebook) + spaces = { + f.name: tuple(f.codebook) + for f in self._features + if isinstance(f, SidFeature) + } if not spaces: raise ValueError( f"{type(self).__name__}: no SID feature declares a codebook; " @@ -146,10 +146,11 @@ def _shared_sid_space(self) -> List[int]: def _slot_group_names(self) -> Dict[str, str]: """{feature_name: group_name} for every declared feature_group. - Each group must be JAGGED_SEQUENCE and hold exactly ONE feature: the - EmbeddingGroup column-interleaves the members of a group into a single - ``{group}.sequence``, so two features in one group could not be spliced - into separate prompt slots. + One JAGGED_SEQUENCE feature per group: EmbeddingGroup interleaves a + group's members into one ``{group}.sequence``, which no longer splits + back into separate prompt slots. The map is keyed by feature, so a + feature claimed by two groups is rejected rather than silently + resolving to whichever group came last. """ by_feature: Dict[str, str] = {} for group in self._feature_groups: @@ -165,21 +166,26 @@ def _slot_group_names(self) -> Dict[str, str]: f"must hold exactly one feature (its members are interleaved " f"into one sequence), got {list(group.feature_names)}." ) - by_feature[group.feature_names[0]] = group.group_name + name = group.feature_names[0] + if name in by_feature: + raise ValueError( + f"{type(self).__name__}: feature {name!r} is claimed by both " + f"feature_group {by_feature[name]!r} and " + f"{group.group_name!r}; only one can fill its prompt slot." + ) + by_feature[name] = group.group_name return by_feature def _resolve_prompt_slots( self, template: str - ) -> Tuple[List[str], List[BaseFeature], List[str]]: - """Split a prompt template into its static gaps and its slot features. - - ``template`` carries ``{{feature_name}}`` slots; returns ``N+1`` static - gap strings, the ``N`` features they are interleaved with in template - order, and those features' feature_group names. Every slot must name a - declared feature, every declared feature_group must be referenced by a - slot, and every slot feature must expose the prompt-text interface - (``prefix_text`` / ``suffix_text``) -- so a misspelt or unused feature - fails here rather than silently vanishing from the prompt. + ) -> Tuple[List[str], List["SidFeature"]]: + """Split a ``{{feature_name}}`` template into N+1 gaps and N features. + + Slots and declared feature_groups must correspond exactly, and each slot + must name a ``SidFeature``, so a misspelt or unused feature fails here + instead of vanishing from the prompt. ``_slot_names`` / ``_slot_groups`` + are recorded here, not in the family hook, because ``build_input`` reads + them and would otherwise fail far from the cause. """ parts = re.split(r"\{\{(\w+)\}\}", template) gaps, names = parts[0::2], parts[1::2] @@ -203,13 +209,10 @@ def _resolve_prompt_slots( f"{type(self).__name__}: prompt_template slot {{{{{name}}}}} " f"has no feature_config." ) - if not hasattr(feature, "prefix_text") or not hasattr( - feature, "suffix_text" - ): + if not isinstance(feature, SidFeature): raise ValueError( - f"{type(self).__name__}: feature {name!r} is a " - f"{type(feature).__name__}, which does not expose the prompt " - f"text interface (prefix_text/suffix_text); use a SidFeature." + f"{type(self).__name__}: prompt slot {{{{{name}}}}} names a " + f"{type(feature).__name__}; only a SidFeature can fill a slot." ) features.append(feature) unused = sorted(set(group_of) - set(names)) @@ -218,22 +221,23 @@ def _resolve_prompt_slots( f"{type(self).__name__}: feature_group(s) {unused} are declared " f"but never referenced by a prompt_template slot." ) - return gaps, features, [group_of[n] for n in names] + self._slot_names = names + self._slot_groups = [group_of[n] for n in names] + return gaps, features def _build_backbone(self) -> PreTrainedModel: - """Build the EMPTY extended architecture -- no weight download. + """Build the EMPTY architecture -- shapes only, no weight download. - Only the module shapes matter here; the weights arrive from - ``init_from_pretrained`` (cold start) or DCP (restore/eval). + Weights arrive from ``init_from_pretrained`` (cold start) or DCP. """ hf_model_id = self._model_config.hf_model_id if not hf_model_id: raise ValueError(f"{type(self).__name__}: empty hf_model_id.") hf_cfg = AutoConfig.from_pretrained(hf_model_id) lm = AutoModelForCausalLM.from_config(hf_cfg, torch_dtype=self._param_dtype) - if next(lm.parameters()).dtype != self._param_dtype: - lm = lm.to(self._param_dtype) - return lm + # a no-op when from_config already honoured torch_dtype, which not every + # architecture does. + return lm.to(self._param_dtype) def _build_extended_tokenizer( self, sid_atoms: int @@ -252,7 +256,7 @@ def _build_extended_tokenizer( if added != sid_atoms: # a pre-existing Cxxx token would shift the atoms off `base`. raise RuntimeError( - f"GenerativeRecLM: tokenizer was expected to grow by " + f"BaseGenerativeModel: tokenizer was expected to grow by " f"{sid_atoms} new atoms, only added {added}. " f"Aborting to avoid silent SID-token mismatch." ) @@ -264,21 +268,18 @@ def _build_extended_tokenizer( c0_id = tokenizer.convert_tokens_to_ids("C0") if c0_id != base: raise RuntimeError( - f"GenerativeRecLM: SID atom layout mismatch -- expected " + f"BaseGenerativeModel: SID atom layout mismatch -- expected " f"C0 at token id {base}, got {c0_id}. " f"Splice arithmetic would produce wrong token ids." ) return tokenizer, base def init_from_pretrained(self) -> None: - """Load the pretrained HF backbone weights into ``self.lm``. + """Load the pretrained HF weights and re-extend to ``__init__``'s vocab. - Re-extends the vocab to ``__init__``'s target so the shapes match. - - Every rank runs this, so every rank must apply the identical resize or - DDP's parameter-shape check fails. The newly-created SID embedding rows - are drawn from the global RNG and therefore differ per rank; DDP's - ``_sync_module_states`` broadcast from rank 0 is what makes them agree. + Every rank resizes identically or DDP's shape check fails. The new SID + rows come from the global RNG and so differ per rank; DDP's + ``_sync_module_states`` broadcast from rank 0 is what reconciles them. """ # drop the empty arch first: holding both peaks at 2x model host RAM. self.lm = None @@ -309,7 +310,7 @@ def _build_prompt_tokens( """ raise NotImplementedError( f"{type(self).__name__} must implement _build_prompt_tokens " - f"(GenerativeRecLM is abstract)." + f"(BaseGenerativeModel is abstract)." ) @property @@ -320,9 +321,8 @@ def device(self) -> torch.device: def _tokenize_sids(self, flat: torch.Tensor) -> torch.Tensor: """Map flat SID indices to extended-vocab token ids. - ``SidFeature`` has already folded the per-level offsets in, and codes are - 0-based, so the flat index IS the atom index: the model only owns the - shift into its own vocabulary. + ``SidFeature`` folded the offsets in and codes are 0-based, so the flat + index IS the atom index; the model only owns the vocabulary shift. """ return flat + self._base_vocab @@ -342,28 +342,25 @@ def _sid_token_bands(self) -> Tuple[torch.Tensor, torch.Tensor]: def _validate_sid_candidates( self, new_tokens: torch.Tensor, batch_size: int ) -> torch.Tensor: - """Decode generated tokens to local 0-based codes and reject bad beams. + """Decode the per-beam tail ``(B*C, w)`` to ``(B, C, num_levels)`` codes. - ``new_tokens`` is the per-beam tail ``(B*C, w)`` (``w`` may be < - ``num_levels`` when beams stop early). Returns ``(batch_size, C, - num_levels)`` local codes. Every malformed candidate (early EOS / - non-SID / wrong-level atom) is set to ``-1``, which cannot match a real - 0-based code. + ``w`` may be < ``num_levels`` when beams stop early. Any malformed + candidate (early EOS, non-SID or wrong-level atom) becomes all ``-1``, + which no real 0-based code can match. """ level_ids = torch.arange(new_tokens.shape[1], device=new_tokens.device) codes = self._detokenize_sids(new_tokens, level_ids) codes = F.pad(codes, (0, self._num_levels - codes.shape[1]), value=-1) - # one out-of-band atom invalidates the whole candidate row. invalid = ((codes < 0) | (codes >= self._codebook_sizes)).any(dim=1) codes = codes.masked_fill(invalid.unsqueeze(1), -1) # decoders return rows batch-major ([b0_c0, b0_c1, ...]); group per user. return codes.view(batch_size, -1, self._num_levels) def init_input(self) -> None: - """Build the EmbeddingGroup for the single raw SID JAGGED_SEQUENCE group. + """Build the EmbeddingGroup over the raw SID JAGGED_SEQUENCE groups. - Raw passthrough features carry no embedding tables, so this group holds - no params (DMP-neutral); it only retrieves the flat ``(values, lengths)``. + Passthrough features own no tables, so this holds no params (DMP-neutral) + and only retrieves the flat ``(values, lengths)``. """ self.embedding_group = EmbeddingGroup(self._features, self._feature_groups) @@ -400,16 +397,11 @@ def _sid_token_rows( ) -> List[torch.Tensor]: """Map a feature's flat SID stream to per-row token-id tensors. - ``SidFeature._parse`` has already validated the codes and folded in the - per-level offsets, so this only applies the model-owned budget and the - shift into the extended vocabulary. - - ``max_codes``, when set, caps each row to its most-recent whole items - (the last ``floor(max_codes / num_levels) * num_levels`` codes, dropping - the oldest head). Skipped unless a row overflows. + ``SidFeature._parse`` already validated and offset the codes, so only the + vocabulary shift and the model-owned budget are left. ``max_codes`` caps + each row to its most-recent WHOLE items, dropping the oldest head. """ - if values.dim() == 2 and values.size(-1) == 1: - values = values.squeeze(-1) + values = values.reshape(-1) # value_dim 1 arrives as (N,) or (N, 1) sizes = lengths.long().tolist() # TODO(shuqi): move truncation into FG once FG can keep the TAIL, not the HEAD. if max_codes: @@ -424,14 +416,12 @@ def _sid_token_rows( def _answer_token_rows( self, values: torch.Tensor, lengths: torch.Tensor ) -> List[torch.Tensor]: - """Map the answer label to token ids. + """Map the answer label to token ids; every row is ``num_levels`` codes. - The answer is a ``data_config.label_field``, not a feature, so nothing - has offset it: this is the one place the model still owns the per-level - fold-in. Every row must be exactly ``num_levels`` codes. + The answer is a label_field, not a feature, so nothing has offset it -- + the one place the model still owns the per-level fold-in. """ - if values.dim() == 2 and values.size(-1) == 1: - values = values.squeeze(-1) + values = values.reshape(-1) # value_dim 1 arrives as (N,) or (N, 1) sizes = lengths.long().tolist() bad = [i for i, n in enumerate(sizes) if n != self._num_levels] if bad: @@ -443,7 +433,7 @@ def _answer_token_rows( codes = values.to(self.device).long() level_ids = torch.arange(codes.numel(), device=self.device) % self._num_levels invalid = (codes < 0) | (codes >= self._codebook_sizes[level_ids]) - if invalid.any().item(): + if invalid.any(): raise ValueError( f"{type(self).__name__}: answer SID codes must be local 0-based " f"values in [0, codebook[level])." diff --git a/tzrec/models/generative_model_test.py b/tzrec/models/generative_model_test.py index 88328e875..eb291d57f 100644 --- a/tzrec/models/generative_model_test.py +++ b/tzrec/models/generative_model_test.py @@ -13,12 +13,15 @@ import unittest import torch +from google.protobuf import text_format from torch import nn -from tzrec.models.generative_rec_lm import GenerativeRecLM +from tzrec.features.feature import create_features +from tzrec.models.generative_model import BaseGenerativeModel +from tzrec.models.generative_qwen import GenerativeQwen from tzrec.models.model import BaseModel -from tzrec.models.qwen2_rec_lm import Qwen2RecLM -from tzrec.protos import model_pb2 +from tzrec.protos import feature_pb2, model_pb2 +from tzrec.protos.models import generative_model_pb2 class _FakeJT: @@ -36,15 +39,50 @@ def lengths(self): return self._l -def _sid_feature(name="user_sequence", codebook=(2, 3, 4)): - """Minimal stand-in for SidFeature: the model only asks for the space.""" - return types.SimpleNamespace(name=name, codebook=list(codebook)) +def _sid_feature(name="user_sequence", codebook=(2, 3, 4), prefix_text=""): + """A real SidFeature -- the model dispatches on the type, not on duck-typing.""" + fc = feature_pb2.FeatureConfig() + text_format.Merge( + f'sequence_sid_feature {{ feature_name: "{name}" expression: "user:{name}" ' + + " ".join(f"codebook: {c}" for c in codebook) + + f' prefix_text: "{prefix_text}" }}', + fc, + ) + return create_features([fc])[0] + + +def _common(**overrides): + """A fake ``GenerativeModelConfig`` -- only the fields the base actually reads.""" + fields = { + "ignore_index": -100, + "generated_sids_key": "generated_sids", + "param_dtype": generative_model_pb2.FP32, + "vocab_pad_to_multiple_of": 128, + "max_sequence_length": 0, + } + return types.SimpleNamespace(**{**fields, **overrides}) + + +def _wired(features=None, group_type=model_pb2.JAGGED_SEQUENCE, members=None): + """Pre-``__init__`` state: the features/labels/groups the config-time code reads.""" + m = object.__new__(GenerativeQwen) + nn.Module.__init__(m) + m._features = [_sid_feature()] if features is None else features + m._labels = ["label"] + m._feature_groups = [ + types.SimpleNamespace( + group_name="user_seq", + feature_names=["user_sequence"] if members is None else list(members), + group_type=group_type, + ) + ] + return m def _stub(codebook=None, base_vocab=100, device="cpu"): - """A Qwen2RecLM with the base data-prep state wired up, but no HF backbone.""" + """A GenerativeQwen with the base data-prep state wired up, but no HF backbone.""" codebook = codebook or [2, 3, 4] - m = object.__new__(Qwen2RecLM) + m = object.__new__(GenerativeQwen) nn.Module.__init__(m) m._base_vocab = base_vocab m._num_levels = len(codebook) @@ -57,81 +95,76 @@ def _stub(codebook=None, base_vocab=100, device="cpu"): return m -class GenerativeRecLMTest(unittest.TestCase): +class BaseGenerativeModelTest(unittest.TestCase): def test_registry_dispatch(self) -> None: - self.assertIs(BaseModel.create_class("Qwen2RecLM"), Qwen2RecLM) - self.assertTrue(issubclass(Qwen2RecLM, GenerativeRecLM)) + self.assertIs(BaseModel.create_class("GenerativeQwen"), GenerativeQwen) + self.assertTrue(issubclass(GenerativeQwen, BaseGenerativeModel)) def test_model_config_oneof_resolves_to_the_class(self) -> None: # the path _create_model takes: oneof -> message type name -> class. from tzrec.utils import config_util cfg = model_pb2.ModelConfig() - cfg.qwen2_rec_lm.common.max_sequence_length = 8 # required field - self.assertEqual(config_util.which_msg(cfg, "model"), "Qwen2RecLM") + cfg.generative_qwen.common.max_sequence_length = 8 # required field + self.assertEqual(config_util.which_msg(cfg, "model"), "GenerativeQwen") self.assertIs( - BaseModel.create_class(config_util.which_msg(cfg, "model")), Qwen2RecLM + BaseModel.create_class(config_util.which_msg(cfg, "model")), GenerativeQwen ) def test_resolve_pad_token_id(self) -> None: tok = types.SimpleNamespace self.assertEqual( - GenerativeRecLM._resolve_pad_token_id(tok(pad_token_id=5, eos_token_id=9)), + BaseGenerativeModel._resolve_pad_token_id( + tok(pad_token_id=5, eos_token_id=9) + ), 5, ) self.assertEqual( - GenerativeRecLM._resolve_pad_token_id( + BaseGenerativeModel._resolve_pad_token_id( tok(pad_token_id=None, eos_token_id=9) ), 9, ) # neither -> a clear error, not an opaque int(None) TypeError with self.assertRaisesRegex(ValueError, "neither pad_token_id nor"): - GenerativeRecLM._resolve_pad_token_id( + BaseGenerativeModel._resolve_pad_token_id( tok(pad_token_id=None, eos_token_id=None) ) def test_backbone_owned_by_family_proto(self) -> None: from tzrec.protos.models.generative_model_pb2 import ( - GenerativeRecLMConfig, + GenerativeModelConfig, ) from tzrec.protos.models.generative_model_pb2 import ( - Qwen2RecLM as Qwen2RecLMProto, + GenerativeQwen as GenerativeQwenProto, ) - self.assertEqual(Qwen2RecLMProto().hf_model_id, "Qwen/Qwen2.5-0.5B") - common_fields = [f.name for f in GenerativeRecLMConfig.DESCRIPTOR.fields] + self.assertEqual(GenerativeQwenProto().hf_model_id, "Qwen/Qwen2.5-0.5B") + common_fields = [f.name for f in GenerativeModelConfig.DESCRIPTOR.fields] self.assertNotIn("hf_model_id", common_fields) def test_configurable_knob_defaults(self) -> None: - from tzrec.protos.models.generative_model_pb2 import GenerativeRecLMConfig + from tzrec.protos.models.generative_model_pb2 import GenerativeModelConfig - c = GenerativeRecLMConfig() + c = GenerativeModelConfig() self.assertEqual(c.generated_sids_key, "generated_sids") - self.assertEqual(c.param_dtype, "float32") - self.assertIs(Qwen2RecLM._DTYPE_BY_NAME["float32"], torch.float32) - self.assertIs(Qwen2RecLM._DTYPE_BY_NAME["bfloat16"], torch.bfloat16) + self.assertEqual(c.param_dtype, generative_model_pb2.FP32) + self.assertIs( + GenerativeQwen._PARAM_DTYPE[generative_model_pb2.FP32], torch.float32 + ) + self.assertIs( + GenerativeQwen._PARAM_DTYPE[generative_model_pb2.BF16], torch.bfloat16 + ) def test_read_common_config_reads_knobs(self) -> None: - m = object.__new__(Qwen2RecLM) - nn.Module.__init__(m) - m._features = [_sid_feature(codebook=(2, 3, 4))] - m._labels = ["label"] - m._feature_groups = [ - types.SimpleNamespace( - group_name="user_seq", - feature_names=["user_sequence"], - group_type=model_pb2.JAGGED_SEQUENCE, + m = _wired() + sid_atoms = m._read_common_config( + _common( + max_sequence_length=288, + generated_sids_key="my_sids", + param_dtype=generative_model_pb2.BF16, ) - ] - common = types.SimpleNamespace( - ignore_index=-100, - generated_sids_key="my_sids", - param_dtype="bfloat16", - vocab_pad_to_multiple_of=128, - max_sequence_length=288, ) - sid_atoms = m._read_common_config(common) self.assertEqual(m._label_name, "label") # from label_fields[0] self.assertEqual(m._generated_sids_key, "my_sids") self.assertIs(m._param_dtype, torch.bfloat16) @@ -141,64 +174,58 @@ def test_read_common_config_reads_knobs(self) -> None: self.assertEqual(m._codebook_sizes.tolist(), [2, 3, 4]) self.assertNotIn("_level_offsets", m.state_dict()) self.assertNotIn("_codebook_sizes", m.state_dict()) - # unknown dtype -> a clear error, not a KeyError - common.param_dtype = "float64" - with self.assertRaisesRegex(ValueError, "param_dtype must be one of"): - m._read_common_config(common) - - def test_read_common_config_no_feature_group_raises(self) -> None: - m = object.__new__(Qwen2RecLM) - nn.Module.__init__(m) - m._features = [_sid_feature(codebook=(2, 3, 4))] - m._labels = ["label"] + # the enum is closed: protobuf itself rejects an unlisted value + cfg = generative_model_pb2.GenerativeModelConfig() + with self.assertRaises(ValueError): + cfg.param_dtype = 99 + + def test_read_common_config_tolerates_no_feature_group(self) -> None: + # group validation is prompt-driven, so it lives in _resolve_prompt_slots + m = _wired() m._feature_groups = [] - common = types.SimpleNamespace( - ignore_index=-100, - generated_sids_key="generated_sids", - param_dtype="float32", - codebook=[4, 4, 4], - vocab_pad_to_multiple_of=128, - max_sequence_length=0, - ) - # group validation moved to _resolve_prompt_slots (prompt-driven) - m._read_common_config(common) + m._read_common_config(_common()) self.assertEqual(m._num_levels, 3) - def test_read_common_config_rejects_bad_group_and_codebook(self) -> None: - def _wired(group_type, feature_names=("user_sequence",)): - m = object.__new__(Qwen2RecLM) - nn.Module.__init__(m) - m._features = [_sid_feature(codebook=(2, 3, 4))] - m._labels = ["label"] - m._feature_groups = [ - types.SimpleNamespace( - group_name="user_seq", - feature_names=list(feature_names), - group_type=group_type, - ) - ] - return m - - def _common(codebook): - return types.SimpleNamespace( - ignore_index=-100, - generated_sids_key="generated_sids", - param_dtype="float32", - vocab_pad_to_multiple_of=128, - max_sequence_length=0, + def test_max_sequence_length_below_one_item_raises(self) -> None: + # a budget under num_levels floors to zero whole items, which would + # silently leave the history uncapped instead of capping it. + for cap in (1, 2): + with self.subTest(cap=cap): + with self.assertRaisesRegex(ValueError, "cannot hold one 3-level"): + _wired()._read_common_config(_common(max_sequence_length=cap)) + # 0 disables the budget; num_levels is the smallest meaningful cap + for cap in (0, 3): + with self.subTest(cap=cap): + m = _wired() + m._read_common_config(_common(max_sequence_length=cap)) + self.assertEqual(m._max_seq_length, cap) + + def test_one_feature_claimed_by_two_groups_raises(self) -> None: + # keyed by feature, so a second claim would otherwise just overwrite + m = _wired() + m._feature_groups.append( + types.SimpleNamespace( + group_name="user_seq_dup", + feature_names=["user_sequence"], + group_type=model_pb2.JAGGED_SEQUENCE, ) + ) + with self.assertRaisesRegex(ValueError, "claimed by both"): + m._slot_group_names() + def test_slot_group_and_shared_codebook_validation(self) -> None: # a SEQUENCE group emits the same key with padded-dense semantics. with self.assertRaisesRegex(ValueError, "must be JAGGED_SEQUENCE"): - _wired(model_pb2.SEQUENCE)._slot_group_names() + _wired(group_type=model_pb2.SEQUENCE)._slot_group_names() with self.assertRaisesRegex(ValueError, "exactly one feature"): - _wired(model_pb2.JAGGED_SEQUENCE, ())._slot_group_names() + _wired(members=())._slot_group_names() # a codebook the SID features disagree on is the model's business - m = _wired(model_pb2.JAGGED_SEQUENCE) - m._features = [ - _sid_feature("user_sequence", (2, 3)), - _sid_feature("other_seq", (4, 4)), - ] + m = _wired( + features=[ + _sid_feature("user_sequence", (2, 3)), + _sid_feature("other_seq", (4, 4)), + ] + ) with self.assertRaisesRegex(ValueError, "must share one"): m._shared_sid_space() m._features = [] @@ -206,62 +233,69 @@ def _common(codebook): m._shared_sid_space() def test_vocab_pad_zero_disables_padding(self) -> None: - m = object.__new__(Qwen2RecLM) - nn.Module.__init__(m) - m._features = [_sid_feature(codebook=(2, 3, 4))] - m._labels = ["label"] - m._feature_groups = [ - types.SimpleNamespace( - group_name="user_seq", - feature_names=["user_sequence"], - group_type=model_pb2.JAGGED_SEQUENCE, - ) - ] - m._read_common_config( - types.SimpleNamespace( - ignore_index=-100, - generated_sids_key="generated_sids", - param_dtype="float32", - codebook=[2, 3], - vocab_pad_to_multiple_of=0, - max_sequence_length=0, - ) - ) + m = _wired() + m._read_common_config(_common(vocab_pad_to_multiple_of=0)) self.assertEqual(m._vocab_pad_mult, 0) # not silently rewritten to 128 def test_max_sequence_length_model_knob(self) -> None: - def _common(max_seq): - return types.SimpleNamespace( - ignore_index=-100, - generated_sids_key="generated_sids", - param_dtype="float32", - vocab_pad_to_multiple_of=128, - max_sequence_length=max_seq, - ) - - def _wired(): - m = object.__new__(Qwen2RecLM) - nn.Module.__init__(m) - m._features = [_sid_feature(codebook=(2, 3, 4))] - m._labels = ["label"] - m._feature_groups = [ - types.SimpleNamespace( - group_name="user_seq", - feature_names=["user_sequence"], - group_type=model_pb2.JAGGED_SEQUENCE, - ) - ] - return m - m = _wired() - m._read_common_config(_common(128)) + m._read_common_config(_common(max_sequence_length=128)) self.assertEqual(m._max_seq_length, 128) m2 = _wired() - m2._read_common_config(_common(0)) + m2._read_common_config(_common(max_sequence_length=0)) self.assertEqual(m2._max_seq_length, 0) # 0 = off, no fallback + def test_resolve_prompt_slots_splits_and_records(self) -> None: + m = _wired() + gaps, features = m._resolve_prompt_slots("A{{user_sequence}}B") + self.assertEqual(gaps, ["A", "B"]) # N slots -> N+1 gaps + self.assertEqual([f.name for f in features], ["user_sequence"]) + # recorded here, not in the family hook -- build_input reads them + self.assertEqual(m._slot_names, ["user_sequence"]) + self.assertEqual(m._slot_groups, ["user_seq"]) + + def test_resolve_prompt_slots_rejects_a_mismatched_template(self) -> None: + from tzrec.features.feature import create_features + + other = _wired( + features=[ + _sid_feature("user_sequence"), + _sid_feature("other_seq"), + ] + ) + other._feature_groups.append( + types.SimpleNamespace( + group_name="other_seq_group", + feature_names=["other_seq"], + group_type=model_pb2.JAGGED_SEQUENCE, + ) + ) + raw = feature_pb2.FeatureConfig() + text_format.Merge( + 'id_feature { feature_name: "user_sequence" expression: "user:x" ' + "num_buckets: 8 embedding_dim: 4 }", + raw, + ) + not_a_sid = _wired(features=create_features([raw])) + + for model, template, msg in ( + (_wired(), "no slot at all", "at least one"), + (_wired(), "{{nope}} x", "names no feature_group"), + (other, "{{user_sequence}}", "never referenced"), + (not_a_sid, "{{user_sequence}}", "only a SidFeature"), + ): + with self.subTest(template=template): + with self.assertRaisesRegex(ValueError, msg): + model._resolve_prompt_slots(template) + + # a group whose feature_config vanished: reachable only out of sync + orphan = _wired() + orphan._features = [] + with self.assertRaisesRegex(ValueError, "no feature_config"): + orphan._resolve_prompt_slots("{{user_sequence}}") + def test_abstract_hooks_raise(self) -> None: - base = object.__new__(GenerativeRecLM) + base = object.__new__(BaseGenerativeModel) with self.assertRaises(NotImplementedError): base._build_prompt_tokens(None, None) with self.assertRaises(NotImplementedError): @@ -290,7 +324,7 @@ def test_sid_token_bands_use_same_level_aware_mapping(self) -> None: self.assertEqual(hi.tolist(), [101, 104, 108]) def test_sid_token_rows_shifts_and_splits(self) -> None: - # values arrive FLAT from SidFeature._parse; the model adds base_vocab-1 + # values arrive FLAT from SidFeature._parse; the model adds base_vocab m = _stub(base_vocab=100) jt = _FakeJT([0, 2, 5, 1, 4, 8, 0, 3, 6], [6, 3]) rows = m._sid_token_rows(jt.values(), jt.lengths()) @@ -345,8 +379,6 @@ def test_build_input_history_group_label_field(self) -> None: m._label_name = "label" m._slot_names = ["user_sequence"] m._slot_groups = ["user_seq"] - m._slot_names = ["user_sequence"] - m._slot_groups = ["user_seq"] m._max_seq_length = 0 m._is_inference = False # train: the answer label_field is read too m.embedding_group = lambda b: { @@ -374,8 +406,6 @@ def test_build_input_skips_label_in_inference(self) -> None: m._label_name = "label" m._slot_names = ["user_sequence"] m._slot_groups = ["user_seq"] - m._slot_names = ["user_sequence"] - m._slot_groups = ["user_seq"] m._max_seq_length = 0 m._is_inference = True # inference: history only, no ground-truth label m.embedding_group = lambda b: { diff --git a/tzrec/models/qwen2_rec_lm.py b/tzrec/models/generative_qwen.py similarity index 69% rename from tzrec/models/qwen2_rec_lm.py rename to tzrec/models/generative_qwen.py index 3de1c389e..0b7c1b17b 100644 --- a/tzrec/models/qwen2_rec_lm.py +++ b/tzrec/models/generative_qwen.py @@ -9,12 +9,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Qwen2/Qwen2.5 family subclass of ``GenerativeRecLM``. +"""Qwen family subclass of ``BaseGenerativeModel``. Owns the decoder-only-chat implementation: the ChatML prompt template, the causal-LM splice, and the ``.model``/``.lm_head`` forward. """ +from itertools import chain from typing import Any, Dict, List, Optional, Tuple import torch @@ -23,14 +24,18 @@ from tzrec.datasets.utils import Batch from tzrec.features.feature import BaseFeature -from tzrec.models.generative_rec_lm import GenerativeRecLM -from tzrec.modules.escalating_beam import escalating_beam_search +from tzrec.models.generative_model import BaseGenerativeModel +from tzrec.modules.dynamic_beam import dynamic_beam_search from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.models import generative_model_pb2 -class Qwen2RecLM(GenerativeRecLM): - """Qwen2 / Qwen2.5 generative-recommendation LM.""" +class GenerativeQwen(BaseGenerativeModel): + """Generative-recommendation LM on a Qwen backbone (Qwen2.5, Qwen3, ...). + + The whole family shares the ChatML frame this class splices, so the version + only enters through ``hf_model_id``. + """ # ChatML frame. Family-specific; a subclass overrides it wholesale. CHAT_TEMPLATE = { @@ -65,50 +70,49 @@ def __init__( self._suffix_keep = self._num_levels + self.tpl_asst_suffix.numel() + 2 def _compute_max_total_length(self) -> int: - """Full spliced length at the max history (0 if pre-allocation is off). - - The ``T`` the activation pool is pre-sized to. - """ + """The ``T`` the activation pool pre-sizes to; 0 when disabled.""" if self._max_seq_length <= 0: return 0 frame = ( - sum( - getattr(self, f"tpl_gap_{i}").numel() - for i in range(self._num_slots + 1) - ) + sum(g.numel() for g in self._gaps) + self.tpl_asst_suffix.numel() + self.tpl_eos.numel() ) - return int(frame + self._max_seq_length * self._num_slots + self._num_levels) + return int( + frame + self._max_seq_length * len(self._slot_names) + self._num_levels + ) + + @property + def _gaps(self) -> List[torch.Tensor]: + """The N+1 static prompt fragments around the N slots, template order. + + Re-read every time: ``.to(device)`` rebinds the buffer, so a cached list + would keep handing back pre-move tensors. + """ + return [getattr(self, f"tpl_gap_{i}") for i in range(len(self._slot_names) + 1)] def _build_prompt_tokens( self, tokenizer: PreTrainedTokenizerBase, - cfg: generative_model_pb2.Qwen2RecLM, + cfg: generative_model_pb2.GenerativeQwen, ) -> None: """Tokenise the static prompt once, as the N+1 gaps around the N slots. - The base resolves ``{{feature_name}}`` slots; this hook only frames them - with ChatML and folds each slot feature's ``prefix_text`` / ``suffix_text`` - into the adjacent gap, so every gap is ONE contiguous string tokenized in - one call. That keeps the encoding bit-identical to tokenizing the fully - rendered prompt -- splitting the token ids instead would let a BPE merge - span a seam. Splicing values between gaps is exact because the ``C*`` - atoms are added-vocab tokens, which HF fast tokenizers pre-split on. - - Buffers are non-persistent: they move with ``model.to(...)`` but stay off - the state_dict so HF safetensors round-tripping isn't polluted. + Each feature's ``prefix_text`` / ``suffix_text`` is folded into the + adjacent gap so every gap is ONE string tokenized in one call, keeping + the encoding bit-identical to the fully rendered prompt -- splitting + token ids instead would let a BPE merge span a seam. Splicing values + between gaps is exact only because the ``C*`` atoms are added-vocab + tokens, which HF fast tokenizers pre-split on. Buffers are + non-persistent: they follow ``.to()`` but stay out of the state_dict. """ tpl = type(self).CHAT_TEMPLATE - gaps, features, groups = self._resolve_prompt_slots(cfg.prompt_template) - self._slot_names = [f.name for f in features] - self._slot_groups = groups - self._num_slots = len(features) + gaps, features = self._resolve_prompt_slots(cfg.prompt_template) for i, gap in enumerate(gaps): head = tpl["user_prefix"] if i == 0 else features[i - 1].suffix_text tail = ( features[i].prefix_text - if i < self._num_slots + if i < len(features) else tpl["user_suffix"] + tpl["asst_prefix"] ) # explicit <|im_start|> markers frame the prompt; no auto BOS/EOS. @@ -138,16 +142,13 @@ def _prompt_rows(self, slot_rows: List[List[torch.Tensor]]) -> List[torch.Tensor Shared by the teacher-forced splice and the answer-less inference prompt. """ - gaps = [getattr(self, f"tpl_gap_{i}") for i in range(self._num_slots + 1)] - rows = [] - for b in range(len(slot_rows[0])): - parts: List[torch.Tensor] = [] - for i in range(self._num_slots): - parts.append(gaps[i]) - parts.append(slot_rows[i][b]) - parts.append(gaps[self._num_slots]) - rows.append(torch.cat(parts)) - return rows + gaps = self._gaps + # zip transposes per-slot row lists into one tuple of slots per row, and + # pairs each slot with the gap that precedes it (the last gap has none). + return [ + torch.cat([*chain.from_iterable(zip(gaps, slots)), gaps[-1]]) + for slots in zip(*slot_rows) + ] def _slot_rows( self, rows: Dict[str, List[torch.Tensor]] @@ -163,15 +164,12 @@ def _splice_input_ids( ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Build ``(input_ids, labels, attention_mask)``, each ``(B, T_max)``. - Rows already hold on-device token ids (see ``_sid_token_rows``). - - Every answer is exactly ``self._num_levels`` SID codes, so the tail - ``[answer | asst_suffix | eos]`` has a FIXED width and lands in the same - columns for every row after left-padding -> ``labels`` is one vectorized - write. Only the answer and the trailing eos are supervised: a decode - emits exactly ``num_levels`` tokens and never has to produce - ``asst_suffix``. ``pad_to`` left-extends every row for pool pre-sizing, - keeping the supervised tail end-aligned. + Every answer is exactly ``num_levels`` codes, so the tail + ``[answer | asst_suffix | eos]`` has a FIXED width and, after left + padding, lands in the same columns for every row -- ``labels`` is one + vectorized write. Only the answer and trailing eos are supervised; a + decode emits ``num_levels`` tokens and never produces ``asst_suffix``. + ``pad_to`` left-extends for pool pre-sizing, keeping the tail aligned. """ if len(slot_rows[0]) != len(label_rows): raise ValueError( @@ -238,16 +236,23 @@ def _forward_loss( return {"loss": loss} def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: - """Beam-search the SID answer (no ground truth supplied). + """Beam-search the SID answer, no ground truth supplied. - Returns ``generated_sids`` of shape ``(B, C, num_levels)``, where ``C`` - is ``num_return_sequences`` on the HF path and the escalating beam's - final width when ``dynamic_beam`` is set. + ``generated_sids`` is ``(B, C, num_levels)``; ``C`` is + ``num_return_sequences``, or the dynamic beam's final width. """ slot_rows = self._slot_rows(self.build_input(batch)) - input_ids, attention_mask = self._splice_prompt_ids(slot_rows) + input_ids, attention_mask = self._left_pad(self._prompt_rows(slot_rows)) if self._dynamic_beam: - new_tokens = self._dynamic_beam_search(input_ids, attention_mask) + lo_tok, hi_tok = self._sid_token_bands() + new_tokens = dynamic_beam_search( + self.lm, + input_ids, + attention_mask, + num_beams=self._num_beams, + lo_tok=lo_tok, + hi_tok=hi_tok, + ) else: out = self.lm.generate( input_ids=input_ids, @@ -262,34 +267,14 @@ def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: sids = self._validate_sid_candidates(new_tokens, input_ids.shape[0]) return {self._generated_sids_key: sids} - def _dynamic_beam_search( - self, input_ids: torch.Tensor, attention_mask: torch.Tensor - ) -> torch.Tensor: - """Escalating-beam decode; see ``escalating_beam_search`` for the schedule.""" - lo_tok, hi_tok = self._sid_token_bands() - return escalating_beam_search( - self.lm, - input_ids, - attention_mask, - num_beams=self._num_beams, - lo_tok=lo_tok, - hi_tok=hi_tok, - ) - - def _splice_prompt_ids( - self, slot_rows: List[List[torch.Tensor]] - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Assemble the answer-less prompt and left-pad into ``(B, T_max)``.""" - return self._left_pad(self._prompt_rows(slot_rows)) - def _left_pad( self, rows: List[torch.Tensor], pad_to: int = 0 ) -> Tuple[torch.Tensor, torch.Tensor]: """Left-pad token rows into ``(input_ids, attention_mask)``, ``(B, T_max)``. - ``attention_mask`` is built from ``ones_like(row)`` (not ``!= pad``) so a - real trailing eos is never masked when ``pad_token_id == eos``. ``pad_to`` - extends on the LEFT, keeping the end-aligned supervised tail in place. + The mask comes from ``ones_like(row)``, not ``!= pad``, so a real + trailing eos survives ``pad_token_id == eos``. ``pad_to`` extends LEFT, + keeping the end-aligned supervised tail in place. """ input_ids = pad_sequence( rows, diff --git a/tzrec/models/qwen2_rec_lm_test.py b/tzrec/models/generative_qwen_test.py similarity index 76% rename from tzrec/models/qwen2_rec_lm_test.py rename to tzrec/models/generative_qwen_test.py index 85623d92f..3d3e9182d 100644 --- a/tzrec/models/qwen2_rec_lm_test.py +++ b/tzrec/models/generative_qwen_test.py @@ -17,18 +17,20 @@ from parameterized import parameterized from torch import nn -from tzrec.models.qwen2_rec_lm import Qwen2RecLM +from tzrec.models.generative_qwen import GenerativeQwen +from tzrec.modules.dynamic_beam import dynamic_beam_search +from tzrec.tests.genrec_test_util import create_tiny_causal_lm from tzrec.utils.test_util import parameterized_name_func def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): - """A Qwen2RecLM with the splice-relevant state wired up, no HF backbone. + """A GenerativeQwen with the splice-relevant state wired up, no HF backbone. The non-uniform default codebook makes incorrect ``level * uniform_size`` offset arithmetic visible: sizes=[2,3,4], offsets=[0,2,5]. """ codebook = codebook or [2, 3, 4] - m = object.__new__(Qwen2RecLM) + m = object.__new__(GenerativeQwen) nn.Module.__init__(m) m._ignore_index = -100 m._num_levels = len(codebook) @@ -36,7 +38,6 @@ def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): m._pad_token_id = pad_id m._dynamic_beam = False m._max_seq_length = 0 - m._num_slots = 1 m._slot_names = ["user_sequence"] m._generated_sids_key = "generated_sids" m.lm = types.SimpleNamespace(device=torch.device(device)) @@ -56,30 +57,18 @@ def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): def _real_lm_stub(codebook=None, base_vocab=20, num_beams=2): - """A Qwen2RecLM carrying a real (tiny, random) Qwen2 backbone. + """A GenerativeQwen carrying a real (tiny, random) Qwen2 backbone. Needed wherever the real forward runs: the training objective and the - end-to-end escalating-beam decode; the other tests mock ``lm.generate``. + end-to-end dynamic-beam decode; the other tests mock ``lm.generate``. """ - from transformers import Qwen2Config, Qwen2ForCausalLM - codebook = codebook or [2, 3, 4] - m = object.__new__(Qwen2RecLM) + m = object.__new__(GenerativeQwen) nn.Module.__init__(m) m._num_levels = len(codebook) m._base_vocab = base_vocab m._num_beams = num_beams - cfg = Qwen2Config( - vocab_size=base_vocab + sum(codebook), - hidden_size=32, - intermediate_size=64, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - max_position_embeddings=64, - ) - torch.manual_seed(0) - m.lm = Qwen2ForCausalLM(cfg).eval() + m.lm = create_tiny_causal_lm(base_vocab + sum(codebook)) sizes = torch.tensor(codebook, dtype=torch.long) m.register_buffer("_codebook_sizes", sizes, persistent=False) m.register_buffer( @@ -90,13 +79,19 @@ def _real_lm_stub(codebook=None, base_vocab=20, num_beams=2): def _wire_slots(m, prefix_text="", suffix_text=""): """Minimal _features/_feature_groups so the base slot resolver can run.""" - from tzrec.protos import model_pb2 + from google.protobuf import text_format - m._features = [ - types.SimpleNamespace( - name="user_sequence", prefix_text=prefix_text, suffix_text=suffix_text - ) - ] + from tzrec.features.feature import create_features + from tzrec.protos import feature_pb2, model_pb2 + + fc = feature_pb2.FeatureConfig() + text_format.Merge( + 'sequence_sid_feature { feature_name: "user_sequence" ' + 'expression: "user:user_sequence" codebook: 2 codebook: 3 codebook: 4 ' + f'prefix_text: "{prefix_text}" suffix_text: "{suffix_text}" }}', + fc, + ) + m._features = create_features([fc]) m._feature_groups = [ types.SimpleNamespace( group_name="sids", @@ -110,8 +105,8 @@ def _train_stub(max_total_len): """A ``_predict_train``-ready stub; returns ``(model, spliced T per step)``.""" m = _stub() m._is_inference = False # not inference + nn.Module.training=True -> is_train - m._input_name, m._label_name = "user_sequence", "label" - m._slot_names, m._num_slots = ["user_sequence"], 1 + m._label_name = "label" + m._slot_names = ["user_sequence"] m._max_total_len = max_total_len m._pool_warmed = False seen_lens = [] @@ -121,14 +116,14 @@ def fwd(input_ids, labels, attention_mask): return {"loss": torch.tensor(0.0)} m.build_input = lambda b: { - m._input_name: [torch.tensor([100, 101, 102])], + "user_sequence": [torch.tensor([100, 101, 102])], m._label_name: [torch.tensor([200, 201, 202])], } m._forward_loss = fwd return m, seen_lens -class Qwen2RecLMTest(unittest.TestCase): +class GenerativeQwenTest(unittest.TestCase): def test_splice_layout_and_labels(self) -> None: m = _stub() u = [torch.tensor([100, 101, 102])] @@ -167,36 +162,18 @@ def test_mask_keeps_trailing_eos_when_pad_equals_eos(self) -> None: self.assertEqual(int(mask[0, -1]), 1) self.assertEqual(mask[0].tolist(), [1] * ids.shape[1]) - def test_suffix_keep_matches_dynamic_slice(self) -> None: - # the constant suffix width must cover every supervised column. - m = _stub() - suffix_keep = 6 # num_levels 3 + asst_suffix 1 + trailing eos + HF shift - _, labels, _ = m._splice_input_ids( - [[torch.tensor([100, 101, 102])]], [torch.tensor([200, 201, 202])] - ) - # first supervised column, minimized over rows - first_sup = int(((labels >= 0).cumsum(-1) == 1).float().argmax(-1).min()) - self.assertEqual(suffix_keep, labels.shape[1] - first_sup + 1) - self.assertTrue(bool((labels[:, :-suffix_keep] < 0).all())) - - def test_splice_prompt_ids(self) -> None: - m = _stub() - ids, mask = m._splice_prompt_ids([[torch.tensor([100, 101, 102])]]) - # [head | history | tail], no answer - self.assertEqual(ids[0].tolist(), [10, 11, 12, 100, 101, 102, 13, 14]) - self.assertEqual(mask[0].tolist(), [1] * 8) - def test_predict_routes_on_inference_flag(self) -> None: m = _stub() m._predict_train = lambda b: {"branch": "train"} m._generate = lambda b: {"branch": "generate"} m._is_inference = False # train / eval - self.assertEqual(Qwen2RecLM.predict(m, object())["branch"], "train") + self.assertEqual(GenerativeQwen.predict(m, object())["branch"], "train") m._is_inference = True # inference - self.assertEqual(Qwen2RecLM.predict(m, object())["branch"], "generate") + self.assertEqual(GenerativeQwen.predict(m, object())["branch"], "generate") - # (generated tail -> decoded SIDs). offsets [0,2,5]: local [1,1,1]/[2,3,4] - # are each level's min/max token; every malformed row collapses to -1. + # (generated tail -> decoded SIDs). sizes [2,3,4] -> offsets [0,2,5], so + # [100,102,105] / [101,104,108] are each level's min/max token and local + # codes [0,0,0] / [1,2,3]; every malformed row collapses to -1. @parameterized.expand( [ [[[100, 102, 105], [101, 104, 108]], [[0, 0, 0], [1, 2, 3]]], @@ -219,8 +196,7 @@ def test_predict_routes_on_inference_flag(self) -> None: ) def test_generate_maps_tokens_to_sids(self, tail, expected) -> None: m = _stub(base_vocab=100) - m._input_name = "user_sequence" - m._slot_names, m._num_slots = ["user_sequence"], 1 + m._slot_names = ["user_sequence"] m._num_beams, m._num_return = len(tail) + 1, len(tail) seen = {} @@ -262,22 +238,35 @@ def fake_generate( def test_generate_routes_to_the_dynamic_beam(self) -> None: m = _stub(base_vocab=100) - m._input_name = "user_sequence" - m._slot_names, m._num_slots = ["user_sequence"], 1 + m._slot_names = ["user_sequence"] m._dynamic_beam = True - m._num_beams = m._num_return = 2 + m._num_beams, m._num_return = 5, 2 m.lm.generate = lambda **kw: self.fail("HF generate must not run") seen = {} - def fake_beam(ids, am): - seen.update(ids=ids[0].tolist(), mask=am[0].tolist()) + def fake_kernel(lm, input_ids, attention_mask, *, num_beams, lo_tok, hi_tok): + seen.update( + lm=lm, + ids=input_ids[0].tolist(), + mask=attention_mask[0].tolist(), + num_beams=num_beams, + lo=lo_tok.tolist(), + hi=hi_tok.tolist(), + ) return torch.tensor([[100, 102, 105], [101, 104, 108]]) - m._dynamic_beam_search = fake_beam m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} - sids = m._generate(object())["generated_sids"] + with mock.patch( + "tzrec.models.generative_qwen.dynamic_beam_search", side_effect=fake_kernel + ): + sids = m._generate(object())["generated_sids"] + self.assertIs(seen["lm"], m.lm) self.assertEqual(seen["ids"], [10, 11, 12, 100, 102, 105, 13, 14]) self.assertEqual(seen["mask"], [1] * 8) + self.assertEqual(seen["num_beams"], 5) # num_return_sequences is ignored + # per-level bands, base_vocab-shifted: sizes [2,3,4] -> offsets [0,2,5] + self.assertEqual(seen["lo"], [100, 102, 105]) + self.assertEqual(seen["hi"], [101, 104, 108]) self.assertEqual(tuple(sids.shape), (1, 2, 3)) self.assertEqual(sids[0].tolist(), [[0, 0, 0], [1, 2, 3]]) @@ -289,7 +278,7 @@ def _prompt_tokenizer(self): ) def test_build_prompt_tokens_splits_the_template_around_the_slot(self) -> None: - m = object.__new__(Qwen2RecLM) + m = object.__new__(GenerativeQwen) nn.Module.__init__(m) _wire_slots(m, prefix_text="PRE", suffix_text="SUF") cfg = types.SimpleNamespace(prompt_template="A{{user_sequence}}B") @@ -298,7 +287,7 @@ def test_build_prompt_tokens_splits_the_template_around_the_slot(self) -> None: buf = getattr(m, name) self.assertIsInstance(buf, torch.Tensor) self.assertEqual(buf.dtype, torch.int64) - tpl = Qwen2RecLM.CHAT_TEMPLATE + tpl = GenerativeQwen.CHAT_TEMPLATE # head = user_prefix + before + feature.prefix_text self.assertEqual(m.tpl_gap_0.tolist(), [len(tpl["user_prefix"] + "A" + "PRE")]) # tail = feature.suffix_text + after + user_suffix + asst_prefix @@ -308,20 +297,6 @@ def test_build_prompt_tokens_splits_the_template_around_the_slot(self) -> None: ) self.assertEqual(m.tpl_eos.tolist(), [99]) # eos cached for supervision - def test_build_prompt_tokens_rejects_a_bad_placeholder_count(self) -> None: - m = object.__new__(Qwen2RecLM) - nn.Module.__init__(m) - _wire_slots(m) - cases = [ - ("no placeholder", "at least one"), - ("{{nope}} x", "names no feature_group"), - ] - for template, msg in cases: - with self.subTest(template=template): - cfg = types.SimpleNamespace(prompt_template=template) - with self.assertRaisesRegex(ValueError, msg): - m._build_prompt_tokens(self._prompt_tokenizer(), cfg) - def test_compute_max_total_length(self) -> None: m = _stub() # frame = 3 head + 2 tail + 1 asst_suffix + 1 eos = 7 @@ -350,8 +325,8 @@ def test_splice_row_count_mismatch_raises(self) -> None: m._splice_input_ids([[torch.tensor([100])]], []) -class Qwen2ForwardLossTest(unittest.TestCase): - """The training objective, run for real against a tiny Qwen2 backbone.""" +class GenerativeQwenLossTest(unittest.TestCase): + """The training objective, run for real against a tiny Qwen backbone.""" def _model(self, ignore_index=-100): m = _real_lm_stub(codebook=[2, 3, 4], base_vocab=20, num_beams=2) @@ -366,8 +341,7 @@ def _model(self, ignore_index=-100): m.register_buffer( name, torch.tensor(vals, dtype=torch.long), persistent=False ) - m._num_slots = 1 - m._slot_names = ["user_sequence"] + m._slot_names = ["user_sequence"] m._suffix_keep = 6 # num_levels 3 + asst_suffix 1 + trailing eos + HF shift return m @@ -419,47 +393,30 @@ def test_forward_loss_returns_only_the_loss(self) -> None: self.assertEqual(list(out), ["loss"]) -class Qwen2DynamicBeamTest(unittest.TestCase): - """The wrapper around ``escalating_beam_search``; the kernel has its own test.""" +class GenerativeQwenBeamTest(unittest.TestCase): + """The kernel/model seam, which neither side can assert alone. - def test_dynamic_beam_search_forwards_the_sid_bands(self) -> None: - m = _stub(base_vocab=100) - m._num_beams = 5 - ids = torch.zeros(1, 4, dtype=torch.long) - am = torch.ones(1, 4, dtype=torch.long) - seen = {} - - def fake_kernel(lm, input_ids, attention_mask, *, num_beams, lo_tok, hi_tok): - seen.update( - lm=lm, - input_ids=input_ids, - attention_mask=attention_mask, - num_beams=num_beams, - lo=lo_tok.tolist(), - hi=hi_tok.tolist(), - ) - return torch.zeros(2, 3, dtype=torch.long) - - with mock.patch( - "tzrec.models.qwen2_rec_lm.escalating_beam_search", side_effect=fake_kernel - ): - out = m._dynamic_beam_search(ids, am) - self.assertIs(seen["lm"], m.lm) - self.assertIs(seen["input_ids"], ids) - self.assertIs(seen["attention_mask"], am) - self.assertEqual(seen["num_beams"], 5) - self.assertEqual(seen["lo"], [100, 102, 105]) - self.assertEqual(seen["hi"], [101, 104, 108]) - self.assertEqual(tuple(out.shape), (2, 3)) + ``dynamic_beam_test`` owns the schedule and the band masking, but only the + model knows the bands and owns ``_validate_sid_candidates``, so this is where + "the kernel's output is exactly what the validator accepts" can be checked. + """ def test_band_masked_beams_decode_without_sentinels(self) -> None: - # end-to-end against a real backbone: band masking guarantees that every - # returned candidate survives _validate_sid_candidates. widths are - # [2, 6, 24] here, i.e. exhaustive over the whole 2*3*4 codebook. + # real backbone: band masking guarantees that every returned candidate + # survives _validate_sid_candidates. widths are [2, 6, 24] here, i.e. + # exhaustive over the whole 2*3*4 codebook. codebook = [2, 3, 4] m = _real_lm_stub(codebook=codebook, base_vocab=20, num_beams=3) ids = torch.tensor([[5, 6, 7]]) - new = m._dynamic_beam_search(ids, torch.ones_like(ids)) + lo_tok, hi_tok = m._sid_token_bands() + new = dynamic_beam_search( + m.lm, + ids, + torch.ones_like(ids), + num_beams=m._num_beams, + lo_tok=lo_tok, + hi_tok=hi_tok, + ) sids = m._validate_sid_candidates(new, batch_size=1) self.assertEqual(tuple(sids.shape), (1, 24, 3)) for level, size in enumerate(codebook): diff --git a/tzrec/modules/dynamic_beam.py b/tzrec/modules/dynamic_beam.py index 5a953191e..5291428ff 100644 --- a/tzrec/modules/dynamic_beam.py +++ b/tzrec/modules/dynamic_beam.py @@ -9,9 +9,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Escalating-beam SID decode (no tzrec deps). +"""Dynamic-width beam SID decode (no tzrec deps). -The beam width doubles at every SID level, so early levels are pruned hard. +Backs ``dynamic_beam`` in ``GenerativeModelConfig``: unlike a fixed-width beam, +the width doubles at every SID level, so early levels are pruned hard. """ from typing import List, Tuple @@ -21,7 +22,7 @@ @torch.no_grad() -def escalating_beam_search( +def dynamic_beam_search( model: PreTrainedModel, input_ids: torch.Tensor, attention_mask: torch.Tensor, @@ -30,10 +31,10 @@ def escalating_beam_search( lo_tok: torch.Tensor, hi_tok: torch.Tensor, ) -> torch.Tensor: - """Decode SID answers with the escalating beam. + """Decode SID answers with a per-level escalating beam width. Args: - model: an HF causal LM exposing ``.model`` / ``.lm_head`` (Qwen2 layout). + model: an HF causal LM exposing ``.model`` / ``.lm_head`` (Qwen layout). input_ids: left-padded prompt ids ``(B, P)``. attention_mask: prompt mask ``(B, P)``. num_beams: base beam width; doubles per level. @@ -83,14 +84,16 @@ def _band_logp(logits: torch.Tensor, j: int) -> torch.Tensor: beam_scores, local = scores.topk(widths[0], dim=-1) # (B, W0) seq = (local + bands[0][0]).reshape(-1, 1) beam_scores = beam_scores.reshape(-1) - parent = torch.arange(bsz, device=device).repeat_interleave(widths[0]) - past.reorder_cache(parent) + rows = torch.arange(bsz, device=device) + past.reorder_cache(rows.repeat_interleave(widths[0])) am = attention_mask.repeat_interleave(widths[0], dim=0) cur_w = widths[0] for j in range(1, num_levels): am = torch.cat([am, am.new_ones(bsz * cur_w, 1)], dim=1) - step_pos = (am.long().cumsum(-1) - 1)[:, -1:].clamp(min=0) + # the row always ends on the token just appended, so its position is + # simply how many real tokens precede it. + step_pos = am.long().sum(-1, keepdim=True) - 1 cache_pos = torch.tensor([past.get_seq_length()], device=device) h = model.model( input_ids=seq[:, -1:], @@ -105,10 +108,8 @@ def _band_logp(logits: torch.Tensor, j: int) -> torch.Tensor: scores = _band_logp(model.lm_head(h.last_hidden_state[:, -1, :]), j) scores = scores + beam_scores[:, None] # (B*cur_w, band) cumulative beam_scores, idx = scores.view(bsz, cur_w * band).topk(widths[j], dim=-1) - parent_local = torch.div(idx, band, rounding_mode="floor") tok = lo_j + idx % band - row_base = torch.arange(bsz, device=device)[:, None] * cur_w - parent = (parent_local + row_base).reshape(-1) + parent = (idx // band + rows[:, None] * cur_w).reshape(-1) seq = torch.cat([seq[parent], tok.reshape(-1, 1)], dim=1) beam_scores = beam_scores.reshape(-1) cur_w = widths[j] diff --git a/tzrec/modules/dynamic_beam_test.py b/tzrec/modules/dynamic_beam_test.py index 7e3af1010..0b24907d6 100644 --- a/tzrec/modules/dynamic_beam_test.py +++ b/tzrec/modules/dynamic_beam_test.py @@ -16,33 +16,25 @@ from typing import Any, Dict, List, Tuple import torch +import torch.nn.functional as F from parameterized import parameterized -from tzrec.modules import escalating_beam -from tzrec.modules.escalating_beam import escalating_beam_search +from tzrec.modules import dynamic_beam +from tzrec.modules.dynamic_beam import dynamic_beam_search +from tzrec.tests.genrec_test_util import create_tiny_causal_lm from tzrec.utils.test_util import parameterized_name_func -def _tiny_lm(vocab_size, seed=0): - from transformers import Qwen2Config, Qwen2ForCausalLM - - cfg = Qwen2Config( - vocab_size=vocab_size, - hidden_size=32, - intermediate_size=64, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - max_position_embeddings=64, +def _decode(lm, ids, pairs, num_beams=2, attention_mask=None): + """Run the kernel over ``pairs`` of inclusive per-level (lo, hi) band edges.""" + return dynamic_beam_search( + lm, + ids, + torch.ones_like(ids) if attention_mask is None else attention_mask, + num_beams=num_beams, + lo_tok=torch.tensor([p[0] for p in pairs]), + hi_tok=torch.tensor([p[1] for p in pairs]), ) - torch.manual_seed(seed) - return Qwen2ForCausalLM(cfg).eval() - - -def _bands(pairs): - lo = torch.tensor([p[0] for p in pairs], dtype=torch.long) - hi = torch.tensor([p[1] for p in pairs], dtype=torch.long) - return lo, hi class _RowSpy: @@ -77,10 +69,10 @@ def _bruteforce_scores(lm, input_ids, attention_mask, pairs): return ref -class EscalatingBeamSearchTest(unittest.TestCase): +class DynamicBeamSearchTest(unittest.TestCase): def test_module_declares_no_tzrec_imports(self) -> None: # the "torch-only, no tzrec deps" docstring is what makes it liftable. - tree = ast.parse(pathlib.Path(escalating_beam.__file__).read_text()) + tree = ast.parse(pathlib.Path(dynamic_beam.__file__).read_text()) mods = set() for node in ast.walk(tree): if isinstance(node, ast.Import): @@ -104,12 +96,9 @@ def test_module_declares_no_tzrec_imports(self) -> None: name_func=parameterized_name_func, ) def test_width_schedule(self, num_beams, pairs, expected_widths) -> None: - spy = _RowSpy(_tiny_lm(vocab_size=48)) - lo, hi = _bands(pairs) + spy = _RowSpy(create_tiny_causal_lm(vocab_size=48)) ids = torch.tensor([[5, 6, 7, 8]]) - out = escalating_beam_search( - spy, ids, torch.ones_like(ids), num_beams=num_beams, lo_tok=lo, hi_tok=hi - ) + out = _decode(spy, ids, pairs, num_beams=num_beams) # calls: [prompt (1 row)] + one per level>0, each carrying widths[j-1] rows widths = spy.rows[1:] + [out.shape[0]] self.assertEqual(spy.rows[0], 1) @@ -117,20 +106,12 @@ def test_width_schedule(self, num_beams, pairs, expected_widths) -> None: self.assertEqual(tuple(out.shape), (expected_widths[-1], len(pairs))) def test_tokens_stay_inside_arbitrary_bands(self) -> None: - # bands the Qwen2RecLM caller can never produce: descending, disjoint, + # bands the GenerativeQwen caller can never produce: descending, disjoint, # unequal width -- the kernel's contract is per-level (lo, hi), not a # contiguous codebook layout. pairs = [(5, 6), (20, 24), (11, 13)] - lo, hi = _bands(pairs) ids = torch.tensor([[1, 2, 3, 4]]) - out = escalating_beam_search( - _tiny_lm(vocab_size=30), - ids, - torch.ones_like(ids), - num_beams=2, - lo_tok=lo, - hi_tok=hi, - ) + out = _decode(create_tiny_causal_lm(vocab_size=30), ids, pairs) self.assertEqual(tuple(out.shape), (min(2 * 2**3, 2 * 5 * 3), 3)) for level, (lo_j, hi_j) in enumerate(pairs): col = out[:, level] @@ -144,51 +125,38 @@ def test_tokens_stay_inside_arbitrary_bands(self) -> None: ) def test_left_padding_matches_unpadded(self, seed, n_pad) -> None: pairs = [(20, 21), (22, 24), (25, 28)] - lo, hi = _bands(pairs) - lm = _tiny_lm(vocab_size=30, seed=seed) + lm = create_tiny_causal_lm(vocab_size=30, seed=seed) short = torch.tensor([[5, 6, 7]]) - pad = torch.cat([torch.zeros(1, n_pad, dtype=torch.long), short], dim=1) - am_pad = torch.cat( - [torch.zeros(1, n_pad, dtype=torch.long), torch.ones_like(short)], dim=1 - ) - plain = escalating_beam_search( - lm, short, torch.ones_like(short), num_beams=2, lo_tok=lo, hi_tok=hi - ) - padded = escalating_beam_search( - lm, pad, am_pad, num_beams=2, lo_tok=lo, hi_tok=hi + plain = _decode(lm, short, pairs) + padded = _decode( + lm, + F.pad(short, (n_pad, 0)), + pairs, + attention_mask=F.pad(torch.ones_like(short), (n_pad, 0)), ) self.assertTrue(torch.equal(plain, padded)) def test_ragged_batch_rows_match_solo_runs(self) -> None: # every row of a ragged batch must decode exactly as if run alone. pairs = [(20, 21), (22, 24), (25, 28)] - lo, hi = _bands(pairs) - lm = _tiny_lm(vocab_size=30) - row0, row1 = torch.tensor([[5, 6, 7, 8]]), torch.tensor([[9, 10, 11]]) + lm = create_tiny_causal_lm(vocab_size=30) ids = torch.tensor([[5, 6, 7, 8], [0, 9, 10, 11]]) am = torch.tensor([[1, 1, 1, 1], [0, 1, 1, 1]]) - out = escalating_beam_search(lm, ids, am, num_beams=2, lo_tok=lo, hi_tok=hi) + out = _decode(lm, ids, pairs, attention_mask=am) width = out.shape[0] // 2 - for i, solo_ids in enumerate([row0, row1]): - solo = escalating_beam_search( - lm, - solo_ids, - torch.ones_like(solo_ids), - num_beams=2, - lo_tok=lo, - hi_tok=hi, - ) + solos = [torch.tensor([[5, 6, 7, 8]]), torch.tensor([[9, 10, 11]])] + for i, solo_ids in enumerate(solos): + solo = _decode(lm, solo_ids, pairs) self.assertTrue(torch.equal(out[i * width : (i + 1) * width], solo)) def test_exhaustive_matches_bruteforce_topk(self) -> None: # widths [2, 6, 12] over 2*3*2 = 12 combinations -> no pruning at any # level, so the beam must reproduce the exact full-recompute ranking. pairs = [(20, 21), (22, 24), (25, 26)] - lo, hi = _bands(pairs) - lm = _tiny_lm(vocab_size=30) + lm = create_tiny_causal_lm(vocab_size=30) ids = torch.tensor([[5, 6, 7, 8]]) am = torch.ones_like(ids) - out = escalating_beam_search(lm, ids, am, num_beams=6, lo_tok=lo, hi_tok=hi) + out = _decode(lm, ids, pairs, num_beams=6) got = [tuple(r) for r in out.tolist()] ref = _bruteforce_scores(lm, ids, am, pairs) self.assertEqual(set(got), set(ref)) diff --git a/tzrec/protos/export.proto b/tzrec/protos/export.proto index 178417744..11407d4c7 100644 --- a/tzrec/protos/export.proto +++ b/tzrec/protos/export.proto @@ -3,7 +3,7 @@ package tzrec.protos; // serialization format produced by `tzrec.export`. // TORCHSCRIPT: native scripted_model.pt (+ TRT/AOTI), the default. -// HF: a HuggingFace `from_pretrained`-loadable dir (GenerativeRecLM family). +// HF: a HuggingFace `from_pretrained`-loadable dir (BaseGenerativeModel family). enum ExportFormat { TORCHSCRIPT = 0; HF = 1; diff --git a/tzrec/protos/model.proto b/tzrec/protos/model.proto index 123669f47..c8243c35d 100644 --- a/tzrec/protos/model.proto +++ b/tzrec/protos/model.proto @@ -84,7 +84,7 @@ message ModelConfig { SidRqkmeans sid_rqkmeans = 601; // Generative (causal-LM) models; the 700-block keeps clear of the SID 600s. - Qwen2RecLM qwen2_rec_lm = 700; + GenerativeQwen generative_qwen = 700; } optional uint32 num_class = 2 [default = 1]; diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index a5f4df9a5..ee4aaec99 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -3,28 +3,35 @@ package tzrec.protos; // Generative (causal-LM) recommendation models. // -// Shared, architecture-agnostic config lives in `GenerativeRecLMConfig`, embedded +// Shared, architecture-agnostic config lives in `GenerativeModelConfig`, embedded // as `common` in every family message; family-specific knobs (the backbone, the // chat template) live on the family message itself. +// Storage dtype of the backbone's master weights. +enum ParamDtype { + FP32 = 0; + BF16 = 1; + FP16 = 2; +} + // Architecture-agnostic config shared by all generative-rec families. // Sample contract: // * history : list -- local 0-based per-level codes in // [0, codebook[level]), laid out as whole items in level order; -// the single JAGGED_SEQUENCE feature_group (its one member). +// a JAGGED_SEQUENCE feature_group holding it alone. // * answer : list -- one item's local 0-based codes; the FIRST // `data_config.label_field`, NOT a feature. // SidFeature folds in level_offsets at parse time and the model adds // base_vocab, so token_id = base_vocab + level_offsets[level] + code. Sample // writers must not pre-apply level_offsets. -message GenerativeRecLMConfig { - // NOTE: the SID vocabulary (`codebook`) is declared on the SidFeature that - // carries the codes, not here -- the model asks the feature for it. +message GenerativeModelConfig { + // 2 held codebook, now on SidFeature. + reserved 2; + // Pad the post-extension vocab up to a multiple of this value; 0 disables // padding. optional uint32 vocab_pad_to_multiple_of = 3 [default = 128]; - // Cross-entropy ignore index -- matches PyTorch's F.cross_entropy default. optional int32 ignore_index = 6 [default = -100]; @@ -32,19 +39,18 @@ message GenerativeRecLMConfig { // num_beams. optional uint32 num_beams = 7 [default = 50]; optional uint32 num_return_sequences = 8 [default = 50]; - // When set, decode with the escalating beam: the width doubles at every SID - // level, returning up to num_beams * 2**num_levels candidates (capped by the + // When set, the beam width doubles at every SID level instead of staying + // fixed, returning up to num_beams * 2**num_levels candidates (capped by the // codebook). num_return_sequences is ignored. optional bool dynamic_beam = 9 [default = false]; // Prediction key the generated SIDs are emitted under; reference it from // PredictWrapper output_cols. optional string generated_sids_key = 12 [default = "generated_sids"]; - // Backbone PARAM dtype = the fp32 MASTER weights. "float32" avoids bf16-ULP + // Backbone PARAM dtype = the MASTER weights. FP32 avoids bf16-ULP // underflow of Adam's small (lr=1e-5) updates; bf16 COMPUTE comes from - // mixed_precision:"BF16" autocast, NOT the param dtype. One of: - // float32 | bfloat16 | float16. - optional string param_dtype = 13 [default = "float32"]; + // mixed_precision:"BF16" autocast, NOT the param dtype. + optional ParamDtype param_dtype = 13 [default = FP32]; // Model's history budget in SID codes: the item-aligned, recency-preserving // truncation cap AND the activation-pool pre-size. Distinct from the history @@ -52,16 +58,19 @@ message GenerativeRecLMConfig { required uint32 max_sequence_length = 14; } -// Qwen2 / Qwen2.5 family (Qwen2.5-0.5B, etc.). -message Qwen2RecLM { - optional GenerativeRecLMConfig common = 1; +// Qwen family (Qwen2.5-0.5B, Qwen3-0.6B, etc.) -- one ChatML frame for all. +message GenerativeQwen { + // 11, 12 held user_prefix_text/user_suffix_text, now on SidFeature. + reserved 11, 12; + + optional GenerativeModelConfig common = 1; - // Qwen2 backbone: HF hub id or local path; must be a Qwen2 model. + // Qwen backbone: HF hub id or local path. optional string hf_model_id = 2 [default = "Qwen/Qwen2.5-0.5B"]; // ----- Prompt ----- - // Prompt body for the user turn, carrying exactly one {{feature_name}} - // placeholder that names the SID feature whose codes are spliced in: + // Prompt body for the user turn, carrying one {{feature_name}} placeholder + // per SID feature whose codes are spliced in: // "... Each behavior is represented by three words. {{user_sequence}} // Please predict the semantic encoding of the next behavior." // That feature's own prefix_text/suffix_text wrap the codes inside the slot. diff --git a/tzrec/tests/configs/qwen2_rec_lm_mock.config b/tzrec/tests/configs/generative_qwen_mock.config similarity index 94% rename from tzrec/tests/configs/qwen2_rec_lm_mock.config rename to tzrec/tests/configs/generative_qwen_mock.config index 30393d0da..f34d6a17e 100644 --- a/tzrec/tests/configs/qwen2_rec_lm_mock.config +++ b/tzrec/tests/configs/generative_qwen_mock.config @@ -1,6 +1,6 @@ train_input_path: "" eval_input_path: "" -model_dir: "experiments/qwen2_rec_lm_mock" +model_dir: "experiments/generative_qwen_mock" train_config { sparse_optimizer { adagrad_optimizer { @@ -48,11 +48,11 @@ model_config { feature_names: "user_sequence" group_type: JAGGED_SEQUENCE } - qwen2_rec_lm { + generative_qwen { common { vocab_pad_to_multiple_of: 128 ignore_index: -100 - param_dtype: "float32" + param_dtype: FP32 max_sequence_length: 12 } hf_model_id: "Qwen/Qwen2.5-0.5B" diff --git a/tzrec/tests/genrec_integration_test.py b/tzrec/tests/genrec_integration_test.py index e7507150f..01326e594 100644 --- a/tzrec/tests/genrec_integration_test.py +++ b/tzrec/tests/genrec_integration_test.py @@ -20,11 +20,12 @@ import pyarrow.parquet as pq from tzrec.tests import utils +from tzrec.tests.genrec_test_util import create_tiny_causal_lm from tzrec.utils import config_util from tzrec.utils.test_util import make_test_dir -_MOCK_CONFIG = "tzrec/tests/configs/qwen2_rec_lm_mock.config" -# must match the mock config's `common.codebook` +_MOCK_CONFIG = "tzrec/tests/configs/generative_qwen_mock.config" +# must match the mock config's sequence_sid_feature.codebook _CODEBOOK = [4, 4, 4] @@ -37,7 +38,7 @@ def _write_backbone(save_dir: str, vocab_size: int = 256) -> str: and which has room to append the ``C*`` atoms, so word-level is enough. """ from tokenizers import Tokenizer, models, pre_tokenizers - from transformers import PreTrainedTokenizerFast, Qwen2Config, Qwen2ForCausalLM + from transformers import PreTrainedTokenizerFast eos = "<|endoftext|>" vocab = {t: i for i, t in enumerate([eos, "<|im_start|>", "<|im_end|>"])} @@ -52,17 +53,9 @@ def _write_backbone(save_dir: str, vocab_size: int = 256) -> str: pad_token=eos, additional_special_tokens=["<|im_start|>", "<|im_end|>"], ).save_pretrained(save_dir) - config = Qwen2Config( - vocab_size=len(vocab), - hidden_size=32, - intermediate_size=64, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - max_position_embeddings=512, - tie_word_embeddings=True, - ) - Qwen2ForCausalLM(config).save_pretrained(save_dir) + create_tiny_causal_lm( + len(vocab), tie_word_embeddings=True, max_position_embeddings=512 + ).save_pretrained(save_dir) return save_dir @@ -122,7 +115,7 @@ def _prepare_config(self, num_rows: int = 64) -> str: config = config_util.load_pipeline_config(_MOCK_CONFIG) config.train_input_path = data_glob config.eval_input_path = data_glob - config.model_config.qwen2_rec_lm.hf_model_id = backbone + config.model_config.generative_qwen.hf_model_id = backbone config_path = os.path.join(self.test_dir, "genrec.config") config_util.save_message(config, config_path) return config_path @@ -143,7 +136,7 @@ def test_mock_config_builds_the_model_and_runs_a_batch(self) -> None: model = _create_model( config.model_config, features, list(config.data_config.label_fields) ) - self.assertEqual(type(model).__name__, "Qwen2RecLM") + self.assertEqual(type(model).__name__, "GenerativeQwen") self.assertEqual(model._slot_groups, ["sids"]) self.assertEqual(model._slot_names, ["user_sequence"]) self.assertEqual(model._label_name, "label") @@ -158,12 +151,24 @@ def test_mock_config_builds_the_model_and_runs_a_batch(self) -> None: config.data_config, features, config.train_input_path, mode=Mode.TRAIN ) model.train() - predictions = model.predict(next(dataloader.get_iterator())) + batch = next(dataloader.get_iterator()) + predictions = model.predict(batch) self.assertEqual(list(predictions), ["loss"]) self.assertTrue(bool(predictions["loss"].isfinite())) + + # _suffix_keep bounds the logits _forward_loss upcasts. HF shifts logits + # by one, so the window must open one column BEFORE the first supervised + # label or that label is never scored. The unit tests hardcode the width; + # only a real __init__ proves the formula computes it. + rows = model.build_input(batch) + _, labels, _ = model._splice_input_ids( + model._slot_rows(rows), rows[model._label_name] + ) + first_sup = int((labels >= 0).nonzero()[:, 1].min()) + self.assertEqual(model._suffix_keep, labels.shape[1] - first_sup + 1) self.success = True - def test_qwen2_rec_lm_train_eval(self) -> None: + def test_generative_qwen_train_eval(self) -> None: """End-to-end train -> checkpoint, with HF assets co-located.""" config_path = self._prepare_config() self.success = utils.test_train_eval(config_path, self.test_dir) diff --git a/tzrec/tests/genrec_test_util.py b/tzrec/tests/genrec_test_util.py new file mode 100644 index 000000000..6dfce775e --- /dev/null +++ b/tzrec/tests/genrec_test_util.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fixtures shared by the generative-rec tests across models/, modules/, utils/.""" + +from typing import Any + +import torch + + +def create_tiny_causal_lm( + vocab_size: int, + seed: int = 0, + tie_word_embeddings: bool = False, + max_position_embeddings: int = 64, +) -> Any: + """A 2-layer Qwen2 causal LM cheap enough to build inside a unit test. + + Seeded so two builds agree, and in ``eval()`` so dropout cannot make a decode + non-deterministic. + + Args: + vocab_size (int): rows in the embedding table. + seed (int): torch seed the random init draws from. + tie_word_embeddings (bool): tie ``lm_head`` to the input embedding. + max_position_embeddings (int): longest sequence the backbone accepts. + + Returns: + an eval-mode ``Qwen2ForCausalLM``. + """ + from transformers import Qwen2Config, Qwen2ForCausalLM + + torch.manual_seed(seed) + config = Qwen2Config( + vocab_size=vocab_size, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=max_position_embeddings, + tie_word_embeddings=tie_word_embeddings, + ) + return Qwen2ForCausalLM(config).eval() diff --git a/tzrec/utils/hf_export_util.py b/tzrec/utils/hf_export_util.py index c02a2e32d..8225589d1 100644 --- a/tzrec/utils/hf_export_util.py +++ b/tzrec/utils/hf_export_util.py @@ -9,7 +9,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""HuggingFace export for HF-backed models (``GenerativeRecLM`` family). +"""HuggingFace export for HF-backed models (``BaseGenerativeModel`` family). Kept out of ``export_util`` (TorchScript/TRT/AOTI) so ``checkpoint_util`` can call ``write_hf_assets`` without a circular import. @@ -45,8 +45,7 @@ def _unwrap_hf_model(wrapped_model: nn.Module) -> Optional[nn.Module]: """Walk DMP/TrainWrapper layers down to the model exposing ``hf_backbone``. - Returns ``None`` if no backbone is found in the chain -- not an HF-backed - model, so callers no-op. + ``None`` when the chain has none: not HF-backed, so callers no-op. """ m = wrapped_model while not hasattr(m, "hf_backbone"): @@ -62,10 +61,9 @@ def _unwrap_hf_model(wrapped_model: nn.Module) -> Optional[nn.Module]: def write_hf_assets(wrapped_model: nn.Module, save_dir: str) -> None: """Co-locate the HF config + tokenizer (NO weights) in a checkpoint dir. - Records the backbone's FQN prefix, read off ``wrapped_model``'s live module - graph, into ``hf_export_meta.json`` so ``dcp_to_hf`` can strip it without - hard-coding a wrapper convention. Rank 0 only -- the dense backbone is - data-parallel-replicated. + The backbone's FQN prefix, read off the live module graph, goes into + ``hf_export_meta.json`` so ``dcp_to_hf`` can strip it without hard-coding a + wrapper convention. Rank 0 only; the dense backbone is replicated. """ if int(os.environ.get("RANK", 0)) != 0: return @@ -94,9 +92,9 @@ def write_hf_assets(wrapped_model: nn.Module, save_dir: str) -> None: def dcp_to_hf(ckpt_dir: str, out_dir: str) -> None: """Convert a self-contained checkpoint dir to a ``from_pretrained`` HF dir. - Reads everything from ``ckpt_dir`` -- no live model is built and no weights - are downloaded. Keys that do not map 1:1 onto the architecture in the - co-located ``config.json`` raise rather than write a partial model. + Everything comes from ``ckpt_dir``: no live model, no download. Keys that do + not map 1:1 onto the co-located ``config.json`` raise instead of writing a + partial model. """ from torch.distributed.checkpoint.state_dict_loader import ( _load_state_dict_from_keys, @@ -156,14 +154,15 @@ def _derive_by_suffix( ) mapped = _derive_by_suffix(raw_state) - if mapped is None or set(mapped.keys()) != target_keys: - got = set(mapped.keys()) if mapped is not None else set() - missing = sorted(target_keys - got) - extra = sorted(got - target_keys) + # both strategies return either an exact match or None, so there is nothing + # partial to report -- show both key spaces instead. + if mapped is None: raise RuntimeError( "dcp_to_hf: cannot map the DCP state dict onto the backbone " - f"architecture (recorded prefix={prefix!r}). missing={missing[:10]} " - f"extra={extra[:10]}. Refusing to write a partially-loaded HF model." + f"architecture (recorded prefix={prefix!r}). Wanted " + f"{len(target_keys)} keys like {sorted(target_keys)[:3]}; the " + f"checkpoint holds {len(raw_state)} like {sorted(raw_state)[:3]}. " + "Refusing to write a partially-loaded HF model." ) # Tied heads are dropped after validation; from_pretrained re-ties them. diff --git a/tzrec/utils/hf_export_util_test.py b/tzrec/utils/hf_export_util_test.py index c13d74a44..a18d3c0e4 100644 --- a/tzrec/utils/hf_export_util_test.py +++ b/tzrec/utils/hf_export_util_test.py @@ -19,6 +19,7 @@ from safetensors.torch import load_file from torch import nn +from tzrec.tests.genrec_test_util import create_tiny_causal_lm from tzrec.utils.checkpoint_util import save_model from tzrec.utils.hf_export_util import ( _HF_EXPORT_META_FILENAME, @@ -29,6 +30,11 @@ from tzrec.utils.test_util import make_test_dir +def _tied_lm(): + """The tied-head backbone every case here needs; dcp_to_hf must drop the tie.""" + return create_tiny_causal_lm(64, tie_word_embeddings=True) + + class _FakeTokenizer: """Writes the two tokenizer asset files `write_hf_assets` copies.""" @@ -39,7 +45,7 @@ def save_pretrained(self, save_dir): class _GenRec(nn.Module): - """Stand-in for GenerativeRecLM: an HF backbone plus unrelated params.""" + """Stand-in for BaseGenerativeModel: an HF backbone plus unrelated params.""" def __init__(self, lm): super().__init__() @@ -65,22 +71,6 @@ def __init__(self, module): self.module = module -def _tiny_lm(tie=True): - from transformers import AutoModelForCausalLM, Qwen2Config - - cfg = Qwen2Config( - vocab_size=64, - hidden_size=32, - intermediate_size=64, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - tie_word_embeddings=tie, - max_position_embeddings=128, - ) - return AutoModelForCausalLM.from_config(cfg, torch_dtype=torch.float32) - - class HfExportUtilTest(unittest.TestCase): def setUp(self) -> None: self.test_dir = make_test_dir() @@ -93,7 +83,7 @@ def tearDown(self) -> None: shutil.rmtree(self.test_dir, ignore_errors=True) def test_unwrap_walks_dmp_and_train_wrapper(self) -> None: - inner = _GenRec(_tiny_lm()) + inner = _GenRec(_tied_lm()) self.assertIs(_unwrap_hf_model(inner), inner) self.assertIs(_unwrap_hf_model(_TrainWrapper(inner)), inner) self.assertIs(_unwrap_hf_model(_DmpLike(_TrainWrapper(inner))), inner) @@ -113,7 +103,7 @@ def _save_ckpt(self, wrapped, name="model.ckpt-1"): return ckpt_dir def test_write_hf_assets_records_state_dict_prefix(self) -> None: - lm = _tiny_lm() + lm = _tied_lm() wrapped = _TrainWrapper(_GenRec(lm)) ckpt_dir = self._save_ckpt(wrapped) for name in ("config.json", "tokenizer.json", _HF_EXPORT_META_FILENAME): @@ -128,7 +118,7 @@ def test_write_hf_assets_records_state_dict_prefix(self) -> None: def test_dcp_to_hf_round_trip_drops_tied_head(self) -> None: from transformers import AutoModelForCausalLM - lm = _tiny_lm(tie=True) + lm = _tied_lm() ckpt_dir = self._save_ckpt(_DmpLike(_TrainWrapper(_GenRec(lm)))) out_dir = os.path.join(self.test_dir, "hf_out") dcp_to_hf(ckpt_dir, out_dir) @@ -144,7 +134,7 @@ def test_dcp_to_hf_round_trip_drops_tied_head(self) -> None: self.assertTrue(torch.equal(back.state_dict()[k], v), k) def test_dcp_to_hf_self_heals_a_stale_prefix(self) -> None: - lm = _tiny_lm() + lm = _tied_lm() ckpt_dir = self._save_ckpt(_TrainWrapper(_GenRec(lm))) meta_path = os.path.join(ckpt_dir, _HF_EXPORT_META_FILENAME) with open(meta_path, "w") as f: @@ -157,7 +147,7 @@ def test_dcp_to_hf_self_heals_a_stale_prefix(self) -> None: ) def test_dcp_to_hf_refuses_a_mismatched_architecture(self) -> None: - ckpt_dir = self._save_ckpt(_TrainWrapper(_GenRec(_tiny_lm()))) + ckpt_dir = self._save_ckpt(_TrainWrapper(_GenRec(_tied_lm()))) # widen the recorded architecture so the checkpoint can no longer fill it cfg_path = os.path.join(ckpt_dir, "config.json") with open(cfg_path) as f: From c349e0463dbf9c36613007b065b75fcca953ad5f Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Tue, 28 Jul 2026 09:38:25 +0000 Subject: [PATCH 42/99] [refactor] genrec: state the beam schedule as a list, drop the second decoder num_beams plus a dynamic_beam bool described the width schedule indirectly: the kernel derived num_beams * 2**(j+1) internally, so a policy lived inside a mechanism and no other shape was expressible. GenerativeModelConfig now carries repeated beam_widths -- [50,50,50] is a fixed beam, [100,200,400] the escalating one -- and dynamic_beam_search only caps each entry by what its band and the surviving prefixes can supply. That also lets the HF generate branch go from _generate. The two decoders were never equivalent: HF beam search ranges over the whole vocabulary, so candidates that are not SID atoms had to be discarded as -1, while the band-restricted kernel makes every candidate well-formed by construction. Given the same band mask the two agree exactly across 96 configurations, so no reachable behaviour is lost, and every recorded experiment already selected the dynamic path. The tiny-backbone test fixture moves into tzrec/utils/test_util.py so the genrec test modules share one copy, and the proto's reserved markers are dropped: tzrec configs are protobuf text format, matched by field name, so a reused tag cannot mis-parse an existing config. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/generative_qwen.py | 75 +++++++------ tzrec/models/generative_qwen_test.py | 116 +++++++++++---------- tzrec/modules/dynamic_beam.py | 34 +++--- tzrec/modules/dynamic_beam_test.py | 49 ++++++--- tzrec/protos/models/generative_model.proto | 24 ++--- tzrec/tests/genrec_integration_test.py | 3 +- tzrec/tests/genrec_test_util.py | 52 --------- tzrec/utils/hf_export_util_test.py | 3 +- tzrec/utils/test_util.py | 36 +++++++ 9 files changed, 209 insertions(+), 183 deletions(-) delete mode 100644 tzrec/tests/genrec_test_util.py diff --git a/tzrec/models/generative_qwen.py b/tzrec/models/generative_qwen.py index 0b7c1b17b..e0b2b7ab9 100644 --- a/tzrec/models/generative_qwen.py +++ b/tzrec/models/generative_qwen.py @@ -37,6 +37,9 @@ class GenerativeQwen(BaseGenerativeModel): only enters through ``hf_model_id``. """ + # Flat width used when beam_widths is left empty; mirrors the proto comment. + DEFAULT_BEAM_WIDTH = 50 + # ChatML frame. Family-specific; a subclass overrides it wholesale. CHAT_TEMPLATE = { "user_prefix": "<|im_start|>user\n", @@ -55,20 +58,36 @@ def __init__( ) -> None: super().__init__(model_config, features, labels, sample_weights, **kwargs) common = self._model_config.common - self._num_beams = int(common.num_beams) - self._num_return = int(common.num_return_sequences) - self._dynamic_beam = bool(common.dynamic_beam) - if not self._dynamic_beam and self._num_return > self._num_beams: - raise ValueError( - f"{type(self).__name__}: num_return_sequences " - f"({self._num_return}) must not exceed num_beams " - f"({self._num_beams})." - ) + self._read_beam_config(common) self._max_total_len = self._compute_max_total_length() self._pool_warmed = False # +2 = trailing eos + HF's shift-by-one; constant width avoids a per-step sync. self._suffix_keep = self._num_levels + self.tpl_asst_suffix.numel() + 2 + def _read_beam_config( + self, common: generative_model_pb2.GenerativeModelConfig + ) -> None: + """Parse the decode knobs; the width schedule must match the codebook. + + An empty ``beam_widths`` is a flat ``DEFAULT_BEAM_WIDTH`` at every level. + """ + self._num_return = int(common.num_return_sequences) + self._beam_widths: List[int] = ( + list(common.beam_widths) or [self.DEFAULT_BEAM_WIDTH] * self._num_levels + ) + if len(self._beam_widths) != self._num_levels: + raise ValueError( + f"{type(self).__name__}: beam_widths has " + f"{len(self._beam_widths)} entries but the codebook has " + f"{self._num_levels} levels; give one width per level." + ) + if self._num_return > self._beam_widths[-1]: + raise ValueError( + f"{type(self).__name__}: num_return_sequences " + f"({self._num_return}) must not exceed the final beam width " + f"({self._beam_widths[-1]})." + ) + def _compute_max_total_length(self) -> int: """The ``T`` the activation pool pre-sizes to; 0 when disabled.""" if self._max_seq_length <= 0: @@ -238,34 +257,24 @@ def _forward_loss( def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: """Beam-search the SID answer, no ground truth supplied. - ``generated_sids`` is ``(B, C, num_levels)``; ``C`` is - ``num_return_sequences``, or the dynamic beam's final width. + ``generated_sids`` is ``(B, C, num_levels)``, best-first per row, where + ``C`` is ``num_return_sequences`` (flat schedule) or the dynamic beam's + final width. Candidates come back score-ordered, so trimming to + ``num_return`` keeps the best ones. """ slot_rows = self._slot_rows(self.build_input(batch)) input_ids, attention_mask = self._left_pad(self._prompt_rows(slot_rows)) - if self._dynamic_beam: - lo_tok, hi_tok = self._sid_token_bands() - new_tokens = dynamic_beam_search( - self.lm, - input_ids, - attention_mask, - num_beams=self._num_beams, - lo_tok=lo_tok, - hi_tok=hi_tok, - ) - else: - out = self.lm.generate( - input_ids=input_ids, - attention_mask=attention_mask, - max_new_tokens=self._num_levels, - num_beams=self._num_beams, - num_return_sequences=self._num_return, - do_sample=False, - pad_token_id=self._pad_token_id, - ) - new_tokens = out[:, input_ids.shape[1] :] + lo_tok, hi_tok = self._sid_token_bands() + new_tokens = dynamic_beam_search( + self.lm, + input_ids, + attention_mask, + beam_widths=self._beam_widths, + lo_tok=lo_tok, + hi_tok=hi_tok, + ) sids = self._validate_sid_candidates(new_tokens, input_ids.shape[0]) - return {self._generated_sids_key: sids} + return {self._generated_sids_key: sids[:, : self._num_return]} def _left_pad( self, rows: List[torch.Tensor], pad_to: int = 0 diff --git a/tzrec/models/generative_qwen_test.py b/tzrec/models/generative_qwen_test.py index 3d3e9182d..39b59a5fd 100644 --- a/tzrec/models/generative_qwen_test.py +++ b/tzrec/models/generative_qwen_test.py @@ -19,8 +19,7 @@ from tzrec.models.generative_qwen import GenerativeQwen from tzrec.modules.dynamic_beam import dynamic_beam_search -from tzrec.tests.genrec_test_util import create_tiny_causal_lm -from tzrec.utils.test_util import parameterized_name_func +from tzrec.utils.test_util import create_tiny_causal_lm, parameterized_name_func def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): @@ -36,7 +35,8 @@ def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): m._num_levels = len(codebook) m._base_vocab = base_vocab m._pad_token_id = pad_id - m._dynamic_beam = False + m._num_return = 2 + m._beam_widths = [2] * m._num_levels m._max_seq_length = 0 m._slot_names = ["user_sequence"] m._generated_sids_key = "generated_sids" @@ -56,18 +56,18 @@ def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): return m -def _real_lm_stub(codebook=None, base_vocab=20, num_beams=2): +def _real_lm_stub(codebook=None, base_vocab=20, beam_width=2): """A GenerativeQwen carrying a real (tiny, random) Qwen2 backbone. Needed wherever the real forward runs: the training objective and the - end-to-end dynamic-beam decode; the other tests mock ``lm.generate``. + end-to-end band-restricted decode; the other tests mock the kernel. """ codebook = codebook or [2, 3, 4] m = object.__new__(GenerativeQwen) nn.Module.__init__(m) m._num_levels = len(codebook) m._base_vocab = base_vocab - m._num_beams = num_beams + m._beam_widths = [beam_width] * m._num_levels m.lm = create_tiny_causal_lm(base_vocab + sum(codebook)) sizes = torch.tensor(codebook, dtype=torch.long) m.register_buffer("_codebook_sizes", sizes, persistent=False) @@ -197,59 +197,68 @@ def test_predict_routes_on_inference_flag(self) -> None: def test_generate_maps_tokens_to_sids(self, tail, expected) -> None: m = _stub(base_vocab=100) m._slot_names = ["user_sequence"] - m._num_beams, m._num_return = len(tail) + 1, len(tail) - seen = {} - - def fake_generate( - input_ids, - attention_mask, - max_new_tokens, - num_beams, - num_return_sequences, - do_sample, - pad_token_id, - ): - seen.update( - prompt=input_ids[0].tolist(), - mask=attention_mask[0].tolist(), - max_new_tokens=max_new_tokens, - num_beams=num_beams, - num_return_sequences=num_return_sequences, - do_sample=do_sample, - pad_token_id=pad_token_id, - ) - prompt = input_ids.repeat_interleave(num_return_sequences, dim=0) - return torch.cat([prompt, torch.tensor(tail)], dim=1) - - m.lm.generate = fake_generate - # build_input is mocked, so the batch is opaque to this generation test. + m._num_return = len(tail) + m._beam_widths = [len(tail)] * m._num_levels + # build_input is mocked, so the batch is opaque to this decoding test. m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} - sids = m._generate(object())["generated_sids"] - self.assertEqual(seen["prompt"], [10, 11, 12, 100, 102, 105, 13, 14]) - self.assertEqual(seen["mask"], [1] * 8) - self.assertEqual(seen["max_new_tokens"], 3) # = num_levels - self.assertEqual(seen["num_beams"], len(tail) + 1) - self.assertEqual(seen["num_return_sequences"], len(tail)) - self.assertFalse(seen["do_sample"]) - self.assertEqual(seen["pad_token_id"], 9) + with mock.patch( + "tzrec.models.generative_qwen.dynamic_beam_search", + return_value=torch.tensor(tail), + ): + sids = m._generate(object())["generated_sids"] # (B, num_return, num_levels) self.assertEqual(tuple(sids.shape), (1, len(tail), 3)) self.assertEqual(sids[0].tolist(), expected) - def test_generate_routes_to_the_dynamic_beam(self) -> None: + def test_generate_trims_to_num_return_keeping_the_best(self) -> None: + # the kernel returns score-ordered best-first, so the trim is a prefix + m = _stub(base_vocab=100) + m._slot_names = ["user_sequence"] + m._beam_widths = [4] * m._num_levels + m._num_return = 2 + m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} + four = torch.tensor( + [[100, 102, 105], [101, 104, 108], [100, 103, 106], [101, 102, 107]] + ) + with mock.patch( + "tzrec.models.generative_qwen.dynamic_beam_search", return_value=four + ): + sids = m._generate(object())["generated_sids"] + self.assertEqual(tuple(sids.shape), (1, 2, 3)) + self.assertEqual(sids[0].tolist(), [[0, 0, 0], [1, 2, 3]]) + + def test_beam_config_defaults_and_validation(self) -> None: + def read(widths, num_return, levels=3): + m = object.__new__(GenerativeQwen) + m._num_levels = levels + m._read_beam_config( + types.SimpleNamespace( + beam_widths=widths, num_return_sequences=num_return + ) + ) + return m + + # empty -> flat DEFAULT_BEAM_WIDTH per level; anything else verbatim + self.assertEqual(read([], 50)._beam_widths, [50, 50, 50]) + self.assertEqual(read([100, 200, 400], 400)._beam_widths, [100, 200, 400]) + with self.assertRaisesRegex(ValueError, "one width per level"): + read([50, 50], 50) + with self.assertRaisesRegex(ValueError, "must not exceed the final"): + read([50, 50, 50], 80) + + def test_generate_hands_the_kernel_prompt_bands_and_schedule(self) -> None: m = _stub(base_vocab=100) m._slot_names = ["user_sequence"] - m._dynamic_beam = True - m._num_beams, m._num_return = 5, 2 - m.lm.generate = lambda **kw: self.fail("HF generate must not run") + m._beam_widths = [5 * 2 ** (j + 1) for j in range(m._num_levels)] + m._num_return = m._beam_widths[-1] seen = {} - def fake_kernel(lm, input_ids, attention_mask, *, num_beams, lo_tok, hi_tok): + def fake_kernel(lm, input_ids, attention_mask, *, beam_widths, lo_tok, hi_tok): seen.update( lm=lm, ids=input_ids[0].tolist(), mask=attention_mask[0].tolist(), - num_beams=num_beams, + widths=list(beam_widths), lo=lo_tok.tolist(), hi=hi_tok.tolist(), ) @@ -263,7 +272,8 @@ def fake_kernel(lm, input_ids, attention_mask, *, num_beams, lo_tok, hi_tok): self.assertIs(seen["lm"], m.lm) self.assertEqual(seen["ids"], [10, 11, 12, 100, 102, 105, 13, 14]) self.assertEqual(seen["mask"], [1] * 8) - self.assertEqual(seen["num_beams"], 5) # num_return_sequences is ignored + # the schedule reaches the kernel verbatim + self.assertEqual(seen["widths"], [10, 20, 40]) # per-level bands, base_vocab-shifted: sizes [2,3,4] -> offsets [0,2,5] self.assertEqual(seen["lo"], [100, 102, 105]) self.assertEqual(seen["hi"], [101, 104, 108]) @@ -329,7 +339,7 @@ class GenerativeQwenLossTest(unittest.TestCase): """The training objective, run for real against a tiny Qwen backbone.""" def _model(self, ignore_index=-100): - m = _real_lm_stub(codebook=[2, 3, 4], base_vocab=20, num_beams=2) + m = _real_lm_stub(codebook=[2, 3, 4], base_vocab=20, beam_width=2) m._ignore_index = ignore_index m._pad_token_id = 0 for name, vals in { @@ -396,9 +406,9 @@ def test_forward_loss_returns_only_the_loss(self) -> None: class GenerativeQwenBeamTest(unittest.TestCase): """The kernel/model seam, which neither side can assert alone. - ``dynamic_beam_test`` owns the schedule and the band masking, but only the - model knows the bands and owns ``_validate_sid_candidates``, so this is where - "the kernel's output is exactly what the validator accepts" can be checked. + ``dynamic_beam_test`` owns the band masking and the per-level capping, but + only the model knows the bands and owns ``_validate_sid_candidates``, so this + is where "the kernel's output is exactly what the validator accepts" lands. """ def test_band_masked_beams_decode_without_sentinels(self) -> None: @@ -406,14 +416,14 @@ def test_band_masked_beams_decode_without_sentinels(self) -> None: # survives _validate_sid_candidates. widths are [2, 6, 24] here, i.e. # exhaustive over the whole 2*3*4 codebook. codebook = [2, 3, 4] - m = _real_lm_stub(codebook=codebook, base_vocab=20, num_beams=3) + m = _real_lm_stub(codebook=codebook, base_vocab=20, beam_width=3) ids = torch.tensor([[5, 6, 7]]) lo_tok, hi_tok = m._sid_token_bands() new = dynamic_beam_search( m.lm, ids, torch.ones_like(ids), - num_beams=m._num_beams, + beam_widths=[3 * 2 ** (j + 1) for j in range(len(codebook))], lo_tok=lo_tok, hi_tok=hi_tok, ) diff --git a/tzrec/modules/dynamic_beam.py b/tzrec/modules/dynamic_beam.py index 5291428ff..d2d33db34 100644 --- a/tzrec/modules/dynamic_beam.py +++ b/tzrec/modules/dynamic_beam.py @@ -11,8 +11,9 @@ """Dynamic-width beam SID decode (no tzrec deps). -Backs ``dynamic_beam`` in ``GenerativeModelConfig``: unlike a fixed-width beam, -the width doubles at every SID level, so early levels are pruned hard. +Backs ``dynamic_beam`` in ``GenerativeModelConfig``: the beam width varies per +SID level instead of staying fixed. The caller owns the schedule -- this module +only enforces what each level can actually supply. """ from typing import List, Tuple @@ -27,39 +28,48 @@ def dynamic_beam_search( input_ids: torch.Tensor, attention_mask: torch.Tensor, *, - num_beams: int, + beam_widths: List[int], lo_tok: torch.Tensor, hi_tok: torch.Tensor, ) -> torch.Tensor: - """Decode SID answers with a per-level escalating beam width. + """Decode SID answers with a caller-supplied per-level beam width. Args: model: an HF causal LM exposing ``.model`` / ``.lm_head`` (Qwen layout). input_ids: left-padded prompt ids ``(B, P)``. attention_mask: prompt mask ``(B, P)``. - num_beams: base beam width; doubles per level. - lo_tok: inclusive lower per-level token-space band edge, ``(num_levels,)`` - (``num_levels`` is inferred from its length). + beam_widths: requested width for each SID level, one entry per level. + Any schedule is accepted -- doubling, flat, hand-tuned -- and each + entry is capped to what its band and the surviving prefixes supply. + lo_tok: inclusive lower per-level token-space band edge, ``(num_levels,)``. hi_tok: inclusive upper per-level token-space band edge, ``(num_levels,)``. Returns: The generated SID token tail ``(B * W, num_levels)`` score-ordered - best-first per row, where ``W`` is ``num_beams * 2**num_levels`` capped - to the number of distinct SIDs the codebook can supply. The answer is + best-first per row, where ``W`` is the last capped width. The answer is fixed-length and EOS-free, so no finished-beam bookkeeping is needed. """ device = input_ids.device bsz = input_ids.shape[0] num_levels = lo_tok.shape[0] + if len(beam_widths) != num_levels: + raise ValueError( + f"dynamic_beam_search: beam_widths has {len(beam_widths)} entries " + f"but the bands describe {num_levels} SID levels." + ) + if any(w < 1 for w in beam_widths): + raise ValueError( + f"dynamic_beam_search: beam_widths must be >= 1, got {list(beam_widths)}." + ) # Hoist the band edges to host once to keep the level loop sync-free. bands: List[Tuple[int, int]] = [ (int(lo_tok[j]), int(hi_tok[j])) for j in range(num_levels) ] - # Cap the doubling at what band x surviving prefixes supply (tiny codebooks). + # A level can only carry band x surviving prefixes, however much was asked. widths: List[int] = [] prev = 1 - for j, (lo, hi) in enumerate(bands): - widths.append(min(num_beams * (2 ** (j + 1)), prev * (hi - lo + 1))) + for w, (lo, hi) in zip(beam_widths, bands): + widths.append(min(w, prev * (hi - lo + 1))) prev = widths[-1] def _band_logp(logits: torch.Tensor, j: int) -> torch.Tensor: diff --git a/tzrec/modules/dynamic_beam_test.py b/tzrec/modules/dynamic_beam_test.py index 0b24907d6..2603e77fa 100644 --- a/tzrec/modules/dynamic_beam_test.py +++ b/tzrec/modules/dynamic_beam_test.py @@ -21,17 +21,22 @@ from tzrec.modules import dynamic_beam from tzrec.modules.dynamic_beam import dynamic_beam_search -from tzrec.tests.genrec_test_util import create_tiny_causal_lm -from tzrec.utils.test_util import parameterized_name_func +from tzrec.utils.test_util import create_tiny_causal_lm, parameterized_name_func -def _decode(lm, ids, pairs, num_beams=2, attention_mask=None): - """Run the kernel over ``pairs`` of inclusive per-level (lo, hi) band edges.""" +def _decode(lm, ids, pairs, width=8, beam_widths=None, attention_mask=None): + """Run the kernel over ``pairs`` of inclusive per-level (lo, hi) band edges. + + Defaults to a FLAT schedule: the kernel is policy-free, so only the cases + about scheduling spell one out. Widths are capped per level anyway. + """ + if beam_widths is None: + beam_widths = [width] * len(pairs) return dynamic_beam_search( lm, ids, torch.ones_like(ids) if attention_mask is None else attention_mask, - num_beams=num_beams, + beam_widths=beam_widths, lo_tok=torch.tensor([p[0] for p in pairs]), hi_tok=torch.tensor([p[1] for p in pairs]), ) @@ -85,26 +90,37 @@ def test_module_declares_no_tzrec_imports(self) -> None: @parameterized.expand( [ - # uncapped doubling: 2 -> 4 -> 8 -> 16 - [2, [(20, 27), (28, 34), (35, 40)], [4, 8, 16]], - # capped by band x surviving prefixes, not by the doubling - [3, [(20, 21), (22, 24)], [2, 6]], - [1, [(20, 21), (22, 24), (25, 28)], [2, 4, 8]], - # a width-1 band collapses level 0 to a single beam - [2, [(20, 20), (21, 23)], [1, 3]], + # every request satisfiable -> the schedule is honoured verbatim + [[4, 8, 16], [(20, 27), (28, 34), (35, 40)], [4, 8, 16]], + # capped by band x surviving prefixes, not by what was asked + [[6, 12], [(20, 21), (22, 24)], [2, 6]], + [[2, 4, 8], [(20, 21), (22, 24), (25, 28)], [2, 4, 8]], + # a width-1 band collapses level 0 and bounds everything after it + [[4, 8], [(20, 20), (21, 23)], [1, 3]], + # a flat (non-doubling) schedule is just as valid to the kernel + [[3, 3, 3], [(20, 27), (28, 34), (35, 40)], [3, 3, 3]], ], name_func=parameterized_name_func, ) - def test_width_schedule(self, num_beams, pairs, expected_widths) -> None: + def test_width_schedule(self, beam_widths, pairs, expected_widths) -> None: spy = _RowSpy(create_tiny_causal_lm(vocab_size=48)) ids = torch.tensor([[5, 6, 7, 8]]) - out = _decode(spy, ids, pairs, num_beams=num_beams) + out = _decode(spy, ids, pairs, beam_widths=beam_widths) # calls: [prompt (1 row)] + one per level>0, each carrying widths[j-1] rows widths = spy.rows[1:] + [out.shape[0]] self.assertEqual(spy.rows[0], 1) self.assertEqual(widths, expected_widths) self.assertEqual(tuple(out.shape), (expected_widths[-1], len(pairs))) + def test_rejects_a_schedule_that_does_not_match_the_bands(self) -> None: + lm = create_tiny_causal_lm(vocab_size=30) + ids = torch.tensor([[5, 6]]) + pairs = [(20, 21), (22, 24), (25, 28)] + with self.assertRaisesRegex(ValueError, "2 entries but the bands"): + _decode(lm, ids, pairs, beam_widths=[2, 4]) + with self.assertRaisesRegex(ValueError, "must be >= 1"): + _decode(lm, ids, pairs, beam_widths=[2, 0, 4]) + def test_tokens_stay_inside_arbitrary_bands(self) -> None: # bands the GenerativeQwen caller can never produce: descending, disjoint, # unequal width -- the kernel's contract is per-level (lo, hi), not a @@ -112,7 +128,8 @@ def test_tokens_stay_inside_arbitrary_bands(self) -> None: pairs = [(5, 6), (20, 24), (11, 13)] ids = torch.tensor([[1, 2, 3, 4]]) out = _decode(create_tiny_causal_lm(vocab_size=30), ids, pairs) - self.assertEqual(tuple(out.shape), (min(2 * 2**3, 2 * 5 * 3), 3)) + # flat width 8: level 0's 2-wide band caps it, later levels recover + self.assertEqual(tuple(out.shape), (8, 3)) for level, (lo_j, hi_j) in enumerate(pairs): col = out[:, level] self.assertTrue(bool((col >= lo_j).all())) @@ -156,7 +173,7 @@ def test_exhaustive_matches_bruteforce_topk(self) -> None: lm = create_tiny_causal_lm(vocab_size=30) ids = torch.tensor([[5, 6, 7, 8]]) am = torch.ones_like(ids) - out = _decode(lm, ids, pairs, num_beams=6) + out = _decode(lm, ids, pairs, width=12) got = [tuple(r) for r in out.tolist()] ref = _bruteforce_scores(lm, ids, am, pairs) self.assertEqual(set(got), set(ref)) diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index ee4aaec99..f179b1269 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -25,9 +25,6 @@ enum ParamDtype { // base_vocab, so token_id = base_vocab + level_offsets[level] + code. Sample // writers must not pre-apply level_offsets. message GenerativeModelConfig { - // 2 held codebook, now on SidFeature. - reserved 2; - // Pad the post-extension vocab up to a multiple of this value; 0 disables // padding. optional uint32 vocab_pad_to_multiple_of = 3 [default = 128]; @@ -35,14 +32,18 @@ message GenerativeModelConfig { // Cross-entropy ignore index -- matches PyTorch's F.cross_entropy default. optional int32 ignore_index = 6 [default = -100]; - // Beam search (inference only). num_return_sequences must not exceed - // num_beams. - optional uint32 num_beams = 7 [default = 50]; + // Beam search (inference only). Decoding is always restricted to the SID + // bands, so every returned candidate is a well-formed SID. + // + // Beam width per SID level, one entry per level -- the whole schedule, + // stated outright rather than derived from a base and a flag: + // [50, 50, 50] fixed width, returns 50 candidates + // [100, 200, 400] doubling, returns 400 (the ALGR-style escalating beam) + // Each entry is capped to what its band and the surviving prefixes can + // supply. Empty defaults to a flat width of 50 at every level. + repeated uint32 beam_widths = 7; + // Candidates kept per row, best-first; must not exceed the final width. optional uint32 num_return_sequences = 8 [default = 50]; - // When set, the beam width doubles at every SID level instead of staying - // fixed, returning up to num_beams * 2**num_levels candidates (capped by the - // codebook). num_return_sequences is ignored. - optional bool dynamic_beam = 9 [default = false]; // Prediction key the generated SIDs are emitted under; reference it from // PredictWrapper output_cols. @@ -60,9 +61,6 @@ message GenerativeModelConfig { // Qwen family (Qwen2.5-0.5B, Qwen3-0.6B, etc.) -- one ChatML frame for all. message GenerativeQwen { - // 11, 12 held user_prefix_text/user_suffix_text, now on SidFeature. - reserved 11, 12; - optional GenerativeModelConfig common = 1; // Qwen backbone: HF hub id or local path. diff --git a/tzrec/tests/genrec_integration_test.py b/tzrec/tests/genrec_integration_test.py index 01326e594..017546925 100644 --- a/tzrec/tests/genrec_integration_test.py +++ b/tzrec/tests/genrec_integration_test.py @@ -20,9 +20,8 @@ import pyarrow.parquet as pq from tzrec.tests import utils -from tzrec.tests.genrec_test_util import create_tiny_causal_lm from tzrec.utils import config_util -from tzrec.utils.test_util import make_test_dir +from tzrec.utils.test_util import create_tiny_causal_lm, make_test_dir _MOCK_CONFIG = "tzrec/tests/configs/generative_qwen_mock.config" # must match the mock config's sequence_sid_feature.codebook diff --git a/tzrec/tests/genrec_test_util.py b/tzrec/tests/genrec_test_util.py deleted file mode 100644 index 6dfce775e..000000000 --- a/tzrec/tests/genrec_test_util.py +++ /dev/null @@ -1,52 +0,0 @@ -# Copyright (c) 2026, Alibaba Group; -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# http://www.apache.org/licenses/LICENSE-2.0 -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Fixtures shared by the generative-rec tests across models/, modules/, utils/.""" - -from typing import Any - -import torch - - -def create_tiny_causal_lm( - vocab_size: int, - seed: int = 0, - tie_word_embeddings: bool = False, - max_position_embeddings: int = 64, -) -> Any: - """A 2-layer Qwen2 causal LM cheap enough to build inside a unit test. - - Seeded so two builds agree, and in ``eval()`` so dropout cannot make a decode - non-deterministic. - - Args: - vocab_size (int): rows in the embedding table. - seed (int): torch seed the random init draws from. - tie_word_embeddings (bool): tie ``lm_head`` to the input embedding. - max_position_embeddings (int): longest sequence the backbone accepts. - - Returns: - an eval-mode ``Qwen2ForCausalLM``. - """ - from transformers import Qwen2Config, Qwen2ForCausalLM - - torch.manual_seed(seed) - config = Qwen2Config( - vocab_size=vocab_size, - hidden_size=32, - intermediate_size=64, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - max_position_embeddings=max_position_embeddings, - tie_word_embeddings=tie_word_embeddings, - ) - return Qwen2ForCausalLM(config).eval() diff --git a/tzrec/utils/hf_export_util_test.py b/tzrec/utils/hf_export_util_test.py index a18d3c0e4..0c70f56ab 100644 --- a/tzrec/utils/hf_export_util_test.py +++ b/tzrec/utils/hf_export_util_test.py @@ -19,7 +19,6 @@ from safetensors.torch import load_file from torch import nn -from tzrec.tests.genrec_test_util import create_tiny_causal_lm from tzrec.utils.checkpoint_util import save_model from tzrec.utils.hf_export_util import ( _HF_EXPORT_META_FILENAME, @@ -27,7 +26,7 @@ dcp_to_hf, write_hf_assets, ) -from tzrec.utils.test_util import make_test_dir +from tzrec.utils.test_util import create_tiny_causal_lm, make_test_dir def _tied_lm(): diff --git a/tzrec/utils/test_util.py b/tzrec/utils/test_util.py index a841a36d4..f64fc9b85 100644 --- a/tzrec/utils/test_util.py +++ b/tzrec/utils/test_util.py @@ -119,6 +119,42 @@ def create_test_model( return create_test_module(model, graph_type) +def create_tiny_causal_lm( + vocab_size: int, + seed: int = 0, + tie_word_embeddings: bool = False, + max_position_embeddings: int = 64, +) -> nn.Module: + """A 2-layer Qwen2 causal LM cheap enough to build inside a unit test. + + Seeded so two builds agree, and in ``eval()`` so dropout cannot make a decode + non-deterministic. + + Args: + vocab_size (int): rows in the embedding table. + seed (int): torch seed the random init draws from. + tie_word_embeddings (bool): tie ``lm_head`` to the input embedding. + max_position_embeddings (int): longest sequence the backbone accepts. + + Returns: + an eval-mode ``Qwen2ForCausalLM``. + """ + from transformers import Qwen2Config, Qwen2ForCausalLM + + torch.manual_seed(seed) + config = Qwen2Config( + vocab_size=vocab_size, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=max_position_embeddings, + tie_word_embeddings=tie_word_embeddings, + ) + return Qwen2ForCausalLM(config).eval() + + # pyre-ignore [2] def parameterized_name_func(func, num, p) -> str: """Name func for parameterized.""" From 798e28b0fd9a64ee8b0f5450a33d522ad9d08d16 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Tue, 28 Jul 2026 12:38:03 +0000 Subject: [PATCH 43/99] [bugfix] genrec: keep tzrec.predict working by hiding the decode from FX predict_checkpoint runs batches through PredictPipelineSparseDist, whose _rewrite_model FX-traces the model to locate shardable modules. The decode cannot be traced: it turns a jagged batch into per-row python lists, interleaves them with the prompt, then runs a beam whose widths depend on the data. So tzrec.predict died with "Proxy object cannot be iterated" before it reached a single batch, while train and eval were unaffected -- create_train_pipeline falls back to the un-traced TrainPipelineBase when a model owns no ShardedModule, and predict_checkpoint has no equivalent guard. The whole decode now sits behind one torch.fx.wrap leaf. Wrapping an inner helper does not work: a leaf returns a single Proxy, so the list structure is lost and the failure only moves to the caller, through all nine untraceable sites. With one leaf the trace completes and finds nothing to shard, which is correct -- a SID feature declares no embedding table. At run time the leaf is an ordinary call. _generate now returns the tensor and predict owns the output-key contract, since a wrapped function has to return a tensor for FX to handle it cleanly. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/generative_qwen.py | 31 +++++++++++++++----- tzrec/models/generative_qwen_test.py | 43 ++++++++++++++++++++++++---- 2 files changed, 62 insertions(+), 12 deletions(-) diff --git a/tzrec/models/generative_qwen.py b/tzrec/models/generative_qwen.py index e0b2b7ab9..8eaa7bd84 100644 --- a/tzrec/models/generative_qwen.py +++ b/tzrec/models/generative_qwen.py @@ -30,6 +30,24 @@ from tzrec.protos.models import generative_model_pb2 +@torch.fx.wrap +def _fx_wrapped_generate(model: "GenerativeQwen", batch: Batch) -> torch.Tensor: + """One opaque FX leaf spanning the whole decode. + + TorchRec's predict pipeline FX-traces the model to find shardable modules. + The decode cannot be traced -- it turns a jagged batch into per-row python + lists, interleaves them, then runs a beam whose widths depend on the data -- + so without this leaf ``tzrec.predict`` dies inside ``_rewrite_model`` with + "Proxy object cannot be iterated". Wrapping the WHOLE decode is what makes + it work: a leaf returns a single Proxy, so wrapping any inner helper only + moves the failure to its caller. The trace then finds nothing to shard, + which is correct -- a SID feature carries no embedding table. + + At run time this is an ordinary call; only tracing sees a leaf. + """ + return model._generate(batch) + + class GenerativeQwen(BaseGenerativeModel): """Generative-recommendation LM on a Qwen backbone (Qwen2.5, Qwen3, ...). @@ -214,7 +232,7 @@ def _splice_input_ids( def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: """Dispatch on the TER inference flag (``set_is_inference`` in main.py).""" if self.is_inference: - return self._generate(batch) + return {self._generated_sids_key: _fx_wrapped_generate(self, batch)} return self._predict_train(batch) def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: @@ -254,13 +272,12 @@ def _forward_loss( ) return {"loss": loss} - def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: + def _generate(self, batch: Batch) -> torch.Tensor: """Beam-search the SID answer, no ground truth supplied. - ``generated_sids`` is ``(B, C, num_levels)``, best-first per row, where - ``C`` is ``num_return_sequences`` (flat schedule) or the dynamic beam's - final width. Candidates come back score-ordered, so trimming to - ``num_return`` keeps the best ones. + Returns ``(B, C, num_levels)`` best-first per row, where ``C`` is + ``num_return_sequences``. Candidates come back score-ordered, so + trimming to ``num_return`` keeps the best ones. """ slot_rows = self._slot_rows(self.build_input(batch)) input_ids, attention_mask = self._left_pad(self._prompt_rows(slot_rows)) @@ -274,7 +291,7 @@ def _generate(self, batch: Batch) -> Dict[str, torch.Tensor]: hi_tok=hi_tok, ) sids = self._validate_sid_candidates(new_tokens, input_ids.shape[0]) - return {self._generated_sids_key: sids[:, : self._num_return]} + return sids[:, : self._num_return] def _left_pad( self, rows: List[torch.Tensor], pad_to: int = 0 diff --git a/tzrec/models/generative_qwen_test.py b/tzrec/models/generative_qwen_test.py index 39b59a5fd..49abc1d78 100644 --- a/tzrec/models/generative_qwen_test.py +++ b/tzrec/models/generative_qwen_test.py @@ -165,11 +165,44 @@ def test_mask_keeps_trailing_eos_when_pad_equals_eos(self) -> None: def test_predict_routes_on_inference_flag(self) -> None: m = _stub() m._predict_train = lambda b: {"branch": "train"} - m._generate = lambda b: {"branch": "generate"} m._is_inference = False # train / eval self.assertEqual(GenerativeQwen.predict(m, object())["branch"], "train") m._is_inference = True # inference - self.assertEqual(GenerativeQwen.predict(m, object())["branch"], "generate") + sentinel = torch.zeros(1, 2, 3) + with mock.patch( + "tzrec.models.generative_qwen._fx_wrapped_generate", return_value=sentinel + ): + out = GenerativeQwen.predict(m, object()) + self.assertIs(out["generated_sids"], sentinel) + + def test_predict_survives_fx_tracing(self) -> None: + """TorchRec's predict pipeline FX-traces the model before running it. + + The decode is un-traceable by construction (per-row python lists, + data-dependent beam widths), so it sits behind one ``torch.fx.wrap`` + leaf. Without it ``tzrec.predict`` dies inside TorchRec's + ``_rewrite_model`` with "Proxy object cannot be iterated" -- a failure no + unit test saw because train and eval take the un-traced + ``TrainPipelineBase`` path (this model has no ShardedModule). + """ + m = _stub(base_vocab=100) + m._is_inference = True + + class _Wrapper(nn.Module): + def __init__(self, inner): + super().__init__() + self.inner = inner + + def forward(self, batch): + return self.inner.predict(batch) + + gm = torch.fx.symbolic_trace(_Wrapper(m)) + leaves = [ + n + for n in gm.graph.nodes + if n.op == "call_function" and "generate" in str(n.target) + ] + self.assertEqual(len(leaves), 1, f"expected one opaque decode node: {gm.graph}") # (generated tail -> decoded SIDs). sizes [2,3,4] -> offsets [0,2,5], so # [100,102,105] / [101,104,108] are each level's min/max token and local @@ -205,7 +238,7 @@ def test_generate_maps_tokens_to_sids(self, tail, expected) -> None: "tzrec.models.generative_qwen.dynamic_beam_search", return_value=torch.tensor(tail), ): - sids = m._generate(object())["generated_sids"] + sids = m._generate(object()) # (B, num_return, num_levels) self.assertEqual(tuple(sids.shape), (1, len(tail), 3)) self.assertEqual(sids[0].tolist(), expected) @@ -223,7 +256,7 @@ def test_generate_trims_to_num_return_keeping_the_best(self) -> None: with mock.patch( "tzrec.models.generative_qwen.dynamic_beam_search", return_value=four ): - sids = m._generate(object())["generated_sids"] + sids = m._generate(object()) self.assertEqual(tuple(sids.shape), (1, 2, 3)) self.assertEqual(sids[0].tolist(), [[0, 0, 0], [1, 2, 3]]) @@ -268,7 +301,7 @@ def fake_kernel(lm, input_ids, attention_mask, *, beam_widths, lo_tok, hi_tok): with mock.patch( "tzrec.models.generative_qwen.dynamic_beam_search", side_effect=fake_kernel ): - sids = m._generate(object())["generated_sids"] + sids = m._generate(object()) self.assertIs(seen["lm"], m.lm) self.assertEqual(seen["ids"], [10, 11, 12, 100, 102, 105, 13, 14]) self.assertEqual(seen["mask"], [1] * 8) From d2c6d97b6ed86a1c5068f165ae2dadf83dff8c35 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 02:37:32 +0000 Subject: [PATCH 44/99] [ci] pin transformers to the OSS mirror instead of PyPI The H20 lane failed on a ReadTimeoutError from files.pythonhosted.org while pulling the 10.4 MB transformers wheel: the read stalled for eleven minutes, pip aborted, and every test that touches `import tzrec` then died on "No module named 'transformers'" -- nine of them in unrelated HSTU code. transformers was the only large dependency still coming from public PyPI; every other third-party wheel already resolves from the project's OSS bucket, which served 348 MB of fbgemm_gpu_hstu in nineteen seconds during the same failed run. Fetching it from the same bucket takes the stalling transfer off the cross-border path (measured 10.4 MB in 0.22 s, sha256 unchanged). The version pin moves into the filename, matching how faiss and graphlearn are already referenced. Co-Authored-By: Claude Opus 5 (1M context) --- requirements/runtime.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/runtime.txt b/requirements/runtime.txt index ab5060e85..9088e721e 100644 --- a/requirements/runtime.txt +++ b/requirements/runtime.txt @@ -22,4 +22,4 @@ tensorboard torch==2.12.1 torchmetrics==1.0.3 torchrec==1.7.0 -transformers==4.51.2 # generative-rec LM backbone (HF); pin <5.0 — 5.x flips from_pretrained dtype default to "auto"(bf16) +transformers @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/transformers/transformers-4.51.2-py3-none-any.whl # generative-rec LM backbone (HF); pinned to 4.51.2 -- 5.x flips from_pretrained dtype default to "auto"(bf16) From 47d3356b069875b1d4aae3761336c95e1383be1b Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 02:46:53 +0000 Subject: [PATCH 45/99] [bugfix] genrec: let a SID feature live in an fg pipeline data_config.fg_mode is a pipeline-wide switch, so SidFeature refusing anything but FG_NONE did not merely disable fg for itself -- it blocked every other feature in the same config from using fg at all. A config that needs FG_NORMAL for its ordinary features therefore could not carry a SID feature. SidFeature now emits a passthrough fg config: a plain raw_feature that fg only splits into per-position values, with _parse folding in the level offsets afterwards exactly as it does on the FG_NONE path. No fg feature_type can add level_offsets[i % num_levels] -- fg expressions are not position-aware within a sequence -- so the arithmetic stays in _parse and fg is used purely to reach the codes. Verified against real pyfg: FG_NONE on a list column and FG_NORMAL on the delimited-string column that ODPS and CSV deliver produce identical values and lengths. sequence_length is now meaningful, since fg applies it, but it truncates by value count and keeps the head. A cap that is not a whole number of levels would hand the model a partial item, so it is rejected with a pointer to max_sequence_length, which is item-aligned and keeps the recent tail. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/features/sid_feature.py | 67 ++++++++++++++++++------------ tzrec/features/sid_feature_test.py | 50 +++++++++++++++------- 2 files changed, 76 insertions(+), 41 deletions(-) diff --git a/tzrec/features/sid_feature.py b/tzrec/features/sid_feature.py index 684b3e4c3..4e346e121 100644 --- a/tzrec/features/sid_feature.py +++ b/tzrec/features/sid_feature.py @@ -15,7 +15,7 @@ import pyarrow as pa from tzrec.datasets.utils import ParsedData -from tzrec.features.feature import BaseFeature, FgMode +from tzrec.features.feature import BaseFeature from tzrec.protos.feature_pb2 import FeatureConfig @@ -23,8 +23,9 @@ class SidFeature(BaseFeature): """Semantic-ID sequence feature. A flat stream of 0-based per-level SID codes -- whole items in level order -- - plus the prompt text wrapping them. Generated offline, so only - ``fg_mode = FG_NONE`` works; there is no pyfg counterpart. + plus the prompt text wrapping them. Codes are produced offline by a + SID-generation model; under fg this feature is a passthrough (see + ``_fg_json``), so it never forces the pipeline's ``fg_mode``. Args: feature_config (FeatureConfig): a instance of feature config. @@ -37,26 +38,21 @@ def __init__( ) -> None: # BaseFeature.__del__ dereferences _fg_op, so seed it before any raise. self._fg_op = None - # checked before super(), which calls init_fg() for FG_NORMAL and would - # surface the missing pyfg handler instead of this explanation. - fg_mode = kwargs.get("fg_mode", FgMode.FG_NONE) - if fg_mode != FgMode.FG_NONE: - raise ValueError( - f"{self.__class__.__name__}" - f"[{feature_config.sequence_sid_feature.feature_name}] supports " - f"data_config.fg_mode = FG_NONE only (SID codes are generated " - f"offline by a SID model), got {fg_mode}." - ) super().__init__(feature_config, **kwargs) - # only fg (which this feature forbids) and an fx export marker read it, - # so it would cap nothing; the model owns the real, item-aligned budget. - if self.config.HasField("sequence_length"): - raise ValueError( - f"{self.__class__.__name__}[{self.config.feature_name}]: " - f"sequence_length does not truncate a SID feature; set " - f"model_config.common.max_sequence_length instead." - ) self._codebook = self._read_codebook() + # fg truncates by VALUE count and keeps the head, so a cap that is not a + # whole number of items would hand the model partial items. The model's + # own max_sequence_length is item-aligned and keeps the recent tail. + if self.config.HasField("sequence_length"): + if self.config.sequence_length % len(self._codebook): + raise ValueError( + f"{self.__class__.__name__}[{self.config.feature_name}]: " + f"sequence_length ({self.config.sequence_length}) must be a " + f"multiple of the {len(self._codebook)}-level codebook, or " + f"fg would cut an item in half. Prefer " + f"model_config.common.max_sequence_length, which is " + f"item-aligned and keeps the most RECENT items." + ) self._level_sizes = np.asarray(self._codebook) self._level_offsets = np.cumsum(self._level_sizes) - self._level_sizes @@ -167,8 +163,27 @@ def _parse(self, input_data: Dict[str, pa.Array]) -> ParsedData: return parsed def _fg_json(self) -> List[Dict[str, Any]]: - """Get fg json config impl.""" - raise RuntimeError( - f"{self.__class__.__name__}[{self.config.feature_name}] has no fg " - f"representation; SID codes are generated offline (fg_mode=FG_NONE)." - ) + """Get fg json config impl. + + A PASSTHROUGH: fg only splits the incoming sequence into per-position + values, and ``_parse`` folds in the level offsets afterwards either way. + There is no fg feature_type that can add ``level_offsets[i % num_levels]`` + (fg expressions are not position-aware within a sequence), so the + arithmetic stays in ``_parse`` and fg is used purely to reach the codes. + + This exists so a SID feature does not force the whole pipeline to + ``FG_NONE`` -- ``fg_mode`` is a data_config-level switch, so blocking it + here would block every other feature in the config too. + """ + # emit the SCALAR form: BaseFeature.fg_json prepends "sequence_" and + # injects sequence_delim / sequence_length for is_sequence features. + fg_cfg: Dict[str, Any] = { + "feature_type": "raw_feature", + "feature_name": self.config.feature_name, + "expression": self.config.expression, + "default_value": self.config.default_value, + "value_type": "float", + } + if self.config.HasField("stub_type"): + fg_cfg["stub_type"] = self.config.stub_type + return [fg_cfg] diff --git a/tzrec/features/sid_feature_test.py b/tzrec/features/sid_feature_test.py index 4f72cae21..1f884fc53 100644 --- a/tzrec/features/sid_feature_test.py +++ b/tzrec/features/sid_feature_test.py @@ -59,11 +59,41 @@ def test_prompt_text_round_trips(self) -> None: [[FgMode.FG_NORMAL], [FgMode.FG_DAG], [FgMode.FG_BUCKETIZE]], name_func=parameterized_name_func, ) - def test_rejects_every_fg_mode_but_none(self, fg_mode) -> None: - # SID codes come from an offline generation model; there is no pyfg - # counterpart, so anything but FG_NONE must fail loudly, not mis-parse. - with self.assertRaisesRegex(ValueError, "FG_NONE only"): - _feature(_BASE, fg_mode=fg_mode) + def test_builds_under_every_fg_mode(self, fg_mode) -> None: + # fg_mode is a data_config-level switch, so refusing it here would block + # every OTHER feature in the config from using fg. + self.assertEqual(type(_feature(_BASE, fg_mode=fg_mode)).__name__, "SidFeature") + + def test_fg_passthrough_matches_the_fg_none_parse(self) -> None: + """Fg only reaches the codes; _parse folds the offsets either way.""" + rows = [[1, 2, 3, 0, 1, 2], [3, 0, 1]] + want = [1, 6, 11, 0, 5, 10, 3, 4, 9] # code + offsets [0,4,8] + none = _feature(_BASE).parse({"user_sequence": pa.array(rows)}) + # under fg the same sequence arrives delimited, as ODPS/CSV deliver it + fg = _feature(_BASE, fg_mode=FgMode.FG_NORMAL).parse( + {"user_sequence": pa.array([";".join(map(str, r)) for r in rows])} + ) + self.assertEqual(none.values.flatten().astype(int).tolist(), want) + self.assertEqual(fg.values.flatten().astype(int).tolist(), want) + self.assertEqual(none.seq_lengths.tolist(), fg.seq_lengths.tolist()) + + def test_fg_json_is_a_passthrough_raw_feature(self) -> None: + cfg = _feature(_BASE).fg_json() + self.assertEqual(len(cfg), 1) + # the base wrapper prepends "sequence_"; no bucketizer, no normalizer + self.assertEqual(cfg[0]["feature_type"], "sequence_raw_feature") + self.assertEqual(cfg[0]["expression"], "user:user_sequence") + for k in ("boundaries", "normalizer", "vocab_file", "hash_bucket_size"): + self.assertNotIn(k, cfg[0]) + + def test_rejects_a_sequence_length_that_splits_an_item(self) -> None: + # fg truncates by VALUE count, so a non-multiple would hand the model + # a partial item; _parse would then reject the whole batch. + with self.assertRaisesRegex(ValueError, "multiple of the 3-level"): + _feature(f"{_BASE} sequence_length: 10") + self.assertEqual( + _feature(f"{_BASE} sequence_length: 9").config.sequence_length, 9 + ) def test_parse_folds_in_the_level_offsets(self) -> None: # offsets [0, 4, 8]: level j's 0-based code k becomes flat index k + off[j], @@ -91,11 +121,6 @@ def test_rejects_a_bad_codebook(self) -> None: with self.assertRaisesRegex(ValueError, msg): _feature(base) - def test_rejects_sequence_length(self) -> None: - # it caps nothing here, so accepting it would read as a working budget - with self.assertRaisesRegex(ValueError, "max_sequence_length instead"): - _feature(f"{_BASE} sequence_length: 64") - def test_no_embedding_table(self) -> None: f = _feature(_BASE) self.assertFalse(f.has_embedding) @@ -103,11 +128,6 @@ def test_no_embedding_table(self) -> None: with self.assertRaisesRegex(RuntimeError, "no .*embedding table"): _ = f.num_embeddings - def test_no_fg_representation(self) -> None: - f = _feature(_BASE) - with self.assertRaisesRegex(RuntimeError, "no fg representation"): - f.fg_json() - if __name__ == "__main__": unittest.main() From 8c5621fae5dafff67be451484504415448eeefeb Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 03:10:09 +0000 Subject: [PATCH 46/99] [ci] drop the inline comment from the transformers requirement The URL already carries the version, and no other entry in this file annotates its pin, so the trailing comment was the only one of its kind. Co-Authored-By: Claude Opus 5 (1M context) --- requirements/runtime.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/runtime.txt b/requirements/runtime.txt index 9088e721e..4bd612390 100644 --- a/requirements/runtime.txt +++ b/requirements/runtime.txt @@ -22,4 +22,4 @@ tensorboard torch==2.12.1 torchmetrics==1.0.3 torchrec==1.7.0 -transformers @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/transformers/transformers-4.51.2-py3-none-any.whl # generative-rec LM backbone (HF); pinned to 4.51.2 -- 5.x flips from_pretrained dtype default to "auto"(bf16) +transformers @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/transformers/transformers-4.51.2-py3-none-any.whl From cf32708ebe404c522d7ec9a0d27cd689c502d64e Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 03:17:02 +0000 Subject: [PATCH 47/99] [bugfix] genrec: bound the hf_backbone walk so a wrapper cycle cannot hang _unwrap_hf_model descends .module / .model looking for a model that exposes hf_backbone, and nothing recorded where it had already been. A cycle in that chain -- a wrapper whose .model points back at itself or an ancestor -- made the while-loop spin forever. It matters because of where the walk runs: CheckpointManager.save() calls write_hf_assets for EVERY model, not just HF-backed ones, so the hang would land inside checkpoint save and strand the peers waiting on the all_gather that follows. Today's four wrappers form a strict chain, so this is latent rather than live. A set of visited ids bounds the walk; a repeat means the chain cannot reach an hf_backbone, which is the same answer as running out of attributes. The regression test walks a cycle on a worker thread so a reoccurrence fails the suite instead of hanging it. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/utils/hf_export_util.py | 8 ++++++++ tzrec/utils/hf_export_util_test.py | 20 ++++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/tzrec/utils/hf_export_util.py b/tzrec/utils/hf_export_util.py index 8225589d1..38348de18 100644 --- a/tzrec/utils/hf_export_util.py +++ b/tzrec/utils/hf_export_util.py @@ -46,9 +46,17 @@ def _unwrap_hf_model(wrapped_model: nn.Module) -> Optional[nn.Module]: """Walk DMP/TrainWrapper layers down to the model exposing ``hf_backbone``. ``None`` when the chain has none: not HF-backed, so callers no-op. + + ``seen`` bounds the walk. Every checkpoint save of every model reaches here, + so a ``.model`` / ``.module`` cycle would hang inside ``save()`` -- the worst + place for it, since peers would block on the collective that follows. """ m = wrapped_model + seen = set() while not hasattr(m, "hf_backbone"): + if id(m) in seen: + return None + seen.add(id(m)) if hasattr(m, "module"): # DMP / DDP-style wrapper m = m.module elif hasattr(m, "model"): # Train/Predict/Script wrapper diff --git a/tzrec/utils/hf_export_util_test.py b/tzrec/utils/hf_export_util_test.py index 0c70f56ab..97b988863 100644 --- a/tzrec/utils/hf_export_util_test.py +++ b/tzrec/utils/hf_export_util_test.py @@ -12,6 +12,7 @@ import json import os import shutil +import threading import unittest from unittest import mock @@ -90,6 +91,25 @@ def test_unwrap_walks_dmp_and_train_wrapper(self) -> None: def test_unwrap_returns_none_for_non_hf_model(self) -> None: self.assertIsNone(_unwrap_hf_model(_TrainWrapper(nn.Linear(4, 4)))) + def test_unwrap_terminates_on_a_wrapper_cycle(self) -> None: + """A .model/.module cycle must return None, not spin. + + Every checkpoint save of every model walks this, so an unbounded loop + here would hang save() and strand the peers waiting on the collective + that follows. Run on a thread so a regression fails the test instead of + hanging the suite. + """ + a, b = nn.Linear(4, 4), nn.Linear(4, 4) + object.__setattr__(a, "model", b) + object.__setattr__(b, "model", a) + out = [] + t = threading.Thread(target=lambda: out.append(_unwrap_hf_model(a))) + t.daemon = True + t.start() + t.join(timeout=5) + self.assertFalse(t.is_alive(), "_unwrap_hf_model did not terminate") + self.assertEqual(out, [None]) + def test_write_hf_assets_noop_for_non_hf_model(self) -> None: save_dir = os.path.join(self.test_dir, "plain") write_hf_assets(_TrainWrapper(nn.Linear(4, 4)), save_dir) From 759bb98cd716637f4cbf5ece353e3e86bae75445 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 03:22:44 +0000 Subject: [PATCH 48/99] [refactor] genrec: cut the explanatory prose down to what is load-bearing The genrec files were about 30% comment and docstring. Most of it restated the code beside it -- shapes already visible in the expression, names already in the signature, proto semantics repeated in three places. What stays is the part that costs debugging time to rediscover: why a gap is encoded as one string (a BPE merge must not span a seam), why fp32 master weights, why validation runs in the dataloader worker rather than the forward path, why the decode sits behind an fx leaf, why _fg_op is seeded before the first raise, and why the wrapper walk is bounded. Docstrings are ruff-enforced so they are shortened, not removed; several drop to their summary line. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/features/sid_feature.py | 39 ++++------- tzrec/models/generative_model.py | 74 +++++--------------- tzrec/models/generative_qwen.py | 81 +++++++--------------- tzrec/modules/dynamic_beam.py | 29 ++++---- tzrec/protos/models/generative_model.proto | 63 +++++------------ tzrec/utils/hf_export_util.py | 30 +++----- 6 files changed, 94 insertions(+), 222 deletions(-) diff --git a/tzrec/features/sid_feature.py b/tzrec/features/sid_feature.py index 4e346e121..43b3c3924 100644 --- a/tzrec/features/sid_feature.py +++ b/tzrec/features/sid_feature.py @@ -23,9 +23,7 @@ class SidFeature(BaseFeature): """Semantic-ID sequence feature. A flat stream of 0-based per-level SID codes -- whole items in level order -- - plus the prompt text wrapping them. Codes are produced offline by a - SID-generation model; under fg this feature is a passthrough (see - ``_fg_json``), so it never forces the pipeline's ``fg_mode``. + plus the prompt text wrapping them. Under fg it is a passthrough. Args: feature_config (FeatureConfig): a instance of feature config. @@ -57,11 +55,7 @@ def __init__( self._level_offsets = np.cumsum(self._level_sizes) - self._level_sizes def _read_codebook(self) -> List[int]: - """Validate the declared codebook once and normalize it to a list. - - Every derived quantity reads it on the parse hot path, so the repeated - scalar container is checked and converted here, not per access. - """ + """Validate the declared codebook once and normalize it to a list.""" codebook = [int(c) for c in self.config.codebook] if not codebook: raise ValueError( @@ -133,10 +127,9 @@ def _build_side_inputs(self) -> Optional[List[Tuple[str, str]]]: def _parse(self, input_data: Dict[str, pa.Array]) -> ParsedData: """Parse the SID stream into flat indices in the shared space. - Codes are 0-based, so the flat index IS the atom index and the model only - adds ``base_vocab``. Offsets are folded in here, in the dataloader - workers, not on the forward path -- and validating here keeps a malformed - row off the collective path, where one rank raising hangs its peers. + Offsets are folded in here, in the dataloader workers: validating on the + forward path would let one rank raise and hang its peers on the + collective. """ parsed = super()._parse(input_data) num_levels = len(self._codebook) @@ -148,16 +141,14 @@ def _parse(self, input_data: Dict[str, pa.Array]) -> ParsedData: f"{bad.tolist()[:10]} have lengths " f"{parsed.seq_lengths[bad].tolist()[:10]}." ) - # rows are whole items, so (-1, num_levels) lines every column up with - # its level and the per-level bounds/offsets broadcast down it. + # rows are whole items, so each column is one level. codes = parsed.values.reshape(-1, num_levels) if ((codes < 0) | (codes >= self._level_sizes)).any(): raise ValueError( f"{self.__class__.__name__}[{self.config.feature_name}]: SID " f"codes must be local 0-based values in [0, codebook[level])." ) - # keep the value dtype: float32 + int64 offsets would promote to float64 - # and double the bytes crossing worker IPC and the H2D copy. + # keep the dtype: int64 offsets would promote float32 to float64. offsets = self._level_offsets.astype(codes.dtype, copy=False) parsed.values = (codes + offsets).reshape(parsed.values.shape) return parsed @@ -165,18 +156,12 @@ def _parse(self, input_data: Dict[str, pa.Array]) -> ParsedData: def _fg_json(self) -> List[Dict[str, Any]]: """Get fg json config impl. - A PASSTHROUGH: fg only splits the incoming sequence into per-position - values, and ``_parse`` folds in the level offsets afterwards either way. - There is no fg feature_type that can add ``level_offsets[i % num_levels]`` - (fg expressions are not position-aware within a sequence), so the - arithmetic stays in ``_parse`` and fg is used purely to reach the codes. - - This exists so a SID feature does not force the whole pipeline to - ``FG_NONE`` -- ``fg_mode`` is a data_config-level switch, so blocking it - here would block every other feature in the config too. + A PASSTHROUGH: no fg feature_type can add ``level_offsets[i % levels]``, + so fg only reaches the codes and ``_parse`` does the arithmetic. It + exists because ``fg_mode`` is a data_config-level switch -- refusing it + here would block every other feature in the config. """ - # emit the SCALAR form: BaseFeature.fg_json prepends "sequence_" and - # injects sequence_delim / sequence_length for is_sequence features. + # SCALAR form: fg_json prepends "sequence_" and injects the seq keys. fg_cfg: Dict[str, Any] = { "feature_type": "raw_feature", "feature_name": self.config.feature_name, diff --git a/tzrec/models/generative_model.py b/tzrec/models/generative_model.py index 355ee9333..47f185152 100644 --- a/tzrec/models/generative_model.py +++ b/tzrec/models/generative_model.py @@ -41,14 +41,8 @@ class BaseGenerativeModel(BaseModel): - """Model construction, SID vocab extension, data-prep, loss and metrics. + """Model construction, SID vocab extension, data-prep, loss and metrics.""" - The family's proto message must carry ``common`` (a - ``GenerativeModelConfig``) and ``hf_model_id``; this base reads both. - """ - - # See `common.param_dtype` in the proto for why FP32 is the default. - # The enum is closed, so protobuf rejects anything not listed here. _PARAM_DTYPE: Dict[int, torch.dtype] = { generative_model_pb2.FP32: torch.float32, generative_model_pb2.BF16: torch.bfloat16, @@ -100,8 +94,6 @@ def _read_common_config( self._max_seq_length: int = int(common.max_sequence_length) codebook = self._shared_sid_space() self._num_levels = len(codebook) - # the budget truncates to WHOLE items, so anything under one item's width - # would floor to zero and silently leave the history uncapped. if 0 < self._max_seq_length < self._num_levels: raise ValueError( f"{type(self).__name__}: max_sequence_length " @@ -111,9 +103,7 @@ def _read_common_config( ) sizes = torch.tensor(codebook, dtype=torch.long) self.register_buffer("_codebook_sizes", sizes, persistent=False) - # only the decode path still needs per-level offsets: SidFeature folds - # them into the input stream, but generated tokens must be split back - # into per-level codes and no feature is involved in generation. + # only the decode path needs these; SidFeature folds them into inputs. self.register_buffer( "_level_offsets", torch.cumsum(sizes, 0) - sizes, persistent=False ) @@ -121,11 +111,7 @@ def _read_common_config( return sum(codebook) def _shared_sid_space(self) -> List[int]: - """The one codebook every SID feature declares. - - One extended vocabulary and one answer width, so adding a feature must - never resize ``lm_head``; disagreement is a typo, not a reshape. - """ + """The one codebook every SID feature declares.""" spaces = { f.name: tuple(f.codebook) for f in self._features @@ -147,10 +133,7 @@ def _slot_group_names(self) -> Dict[str, str]: """{feature_name: group_name} for every declared feature_group. One JAGGED_SEQUENCE feature per group: EmbeddingGroup interleaves a - group's members into one ``{group}.sequence``, which no longer splits - back into separate prompt slots. The map is keyed by feature, so a - feature claimed by two groups is rejected rather than silently - resolving to whichever group came last. + group's members into one sequence that cannot be split back apart. """ by_feature: Dict[str, str] = {} for group in self._feature_groups: @@ -181,11 +164,8 @@ def _resolve_prompt_slots( ) -> Tuple[List[str], List["SidFeature"]]: """Split a ``{{feature_name}}`` template into N+1 gaps and N features. - Slots and declared feature_groups must correspond exactly, and each slot - must name a ``SidFeature``, so a misspelt or unused feature fails here - instead of vanishing from the prompt. ``_slot_names`` / ``_slot_groups`` - are recorded here, not in the family hook, because ``build_input`` reads - them and would otherwise fail far from the cause. + Records ``_slot_names`` / ``_slot_groups`` here rather than in the family + hook, because ``build_input`` reads them. """ parts = re.split(r"\{\{(\w+)\}\}", template) gaps, names = parts[0::2], parts[1::2] @@ -226,17 +206,13 @@ def _resolve_prompt_slots( return gaps, features def _build_backbone(self) -> PreTrainedModel: - """Build the EMPTY architecture -- shapes only, no weight download. - - Weights arrive from ``init_from_pretrained`` (cold start) or DCP. - """ + """Build the EMPTY architecture; weights arrive later from HF or DCP.""" hf_model_id = self._model_config.hf_model_id if not hf_model_id: raise ValueError(f"{type(self).__name__}: empty hf_model_id.") hf_cfg = AutoConfig.from_pretrained(hf_model_id) lm = AutoModelForCausalLM.from_config(hf_cfg, torch_dtype=self._param_dtype) - # a no-op when from_config already honoured torch_dtype, which not every - # architecture does. + # no-op when from_config already honoured torch_dtype; not all do. return lm.to(self._param_dtype) def _build_extended_tokenizer( @@ -254,7 +230,6 @@ def _build_extended_tokenizer( base = len(tokenizer) added = tokenizer.add_tokens([f"C{i}" for i in range(sid_atoms)]) if added != sid_atoms: - # a pre-existing Cxxx token would shift the atoms off `base`. raise RuntimeError( f"BaseGenerativeModel: tokenizer was expected to grow by " f"{sid_atoms} new atoms, only added {added}. " @@ -277,9 +252,8 @@ def _build_extended_tokenizer( def init_from_pretrained(self) -> None: """Load the pretrained HF weights and re-extend to ``__init__``'s vocab. - Every rank resizes identically or DDP's shape check fails. The new SID - rows come from the global RNG and so differ per rank; DDP's - ``_sync_module_states`` broadcast from rank 0 is what reconciles them. + The new SID rows differ per rank; DDP's ``_sync_module_states`` + broadcast from rank 0 reconciles them. """ # drop the empty arch first: holding both peaks at 2x model host RAM. self.lm = None @@ -319,11 +293,7 @@ def device(self) -> torch.device: return self.lm.device def _tokenize_sids(self, flat: torch.Tensor) -> torch.Tensor: - """Map flat SID indices to extended-vocab token ids. - - ``SidFeature`` folded the offsets in and codes are 0-based, so the flat - index IS the atom index; the model only owns the vocabulary shift. - """ + """Map flat SID indices to extended-vocab token ids.""" return flat + self._base_vocab def _detokenize_sids( @@ -344,31 +314,24 @@ def _validate_sid_candidates( ) -> torch.Tensor: """Decode the per-beam tail ``(B*C, w)`` to ``(B, C, num_levels)`` codes. - ``w`` may be < ``num_levels`` when beams stop early. Any malformed - candidate (early EOS, non-SID or wrong-level atom) becomes all ``-1``, - which no real 0-based code can match. + Any malformed candidate becomes all ``-1``, which no real code matches. """ level_ids = torch.arange(new_tokens.shape[1], device=new_tokens.device) codes = self._detokenize_sids(new_tokens, level_ids) codes = F.pad(codes, (0, self._num_levels - codes.shape[1]), value=-1) invalid = ((codes < 0) | (codes >= self._codebook_sizes)).any(dim=1) codes = codes.masked_fill(invalid.unsqueeze(1), -1) - # decoders return rows batch-major ([b0_c0, b0_c1, ...]); group per user. + # decoders return rows batch-major; group per user. return codes.view(batch_size, -1, self._num_levels) def init_input(self) -> None: - """Build the EmbeddingGroup over the raw SID JAGGED_SEQUENCE groups. - - Passthrough features own no tables, so this holds no params (DMP-neutral) - and only retrieves the flat ``(values, lengths)``. - """ + """Build the EmbeddingGroup; passthrough features hold no params.""" self.embedding_group = EmbeddingGroup(self._features, self._feature_groups) def build_input(self, batch: Batch) -> Dict[str, List[torch.Tensor]]: """Retrieve per-row SID token sequences, keyed by feature name. - The answer is a ``data_config.label_field`` rather than a feature_group - so it can be absent at inference, where no ground truth is supplied. + The answer is a label_field so it can be absent at inference. """ g = self.embedding_group(batch) rows: Dict[str, List[torch.Tensor]] = { @@ -397,9 +360,7 @@ def _sid_token_rows( ) -> List[torch.Tensor]: """Map a feature's flat SID stream to per-row token-id tensors. - ``SidFeature._parse`` already validated and offset the codes, so only the - vocabulary shift and the model-owned budget are left. ``max_codes`` caps - each row to its most-recent WHOLE items, dropping the oldest head. + ``max_codes`` caps each row to its most-recent WHOLE items. """ values = values.reshape(-1) # value_dim 1 arrives as (N,) or (N, 1) sizes = lengths.long().tolist() @@ -418,8 +379,7 @@ def _answer_token_rows( ) -> List[torch.Tensor]: """Map the answer label to token ids; every row is ``num_levels`` codes. - The answer is a label_field, not a feature, so nothing has offset it -- - the one place the model still owns the per-level fold-in. + A label_field is not offset by SidFeature, so the fold-in happens here. """ values = values.reshape(-1) # value_dim 1 arrives as (N,) or (N, 1) sizes = lengths.long().tolist() diff --git a/tzrec/models/generative_qwen.py b/tzrec/models/generative_qwen.py index 8eaa7bd84..16f291816 100644 --- a/tzrec/models/generative_qwen.py +++ b/tzrec/models/generative_qwen.py @@ -34,31 +34,19 @@ def _fx_wrapped_generate(model: "GenerativeQwen", batch: Batch) -> torch.Tensor: """One opaque FX leaf spanning the whole decode. - TorchRec's predict pipeline FX-traces the model to find shardable modules. - The decode cannot be traced -- it turns a jagged batch into per-row python - lists, interleaves them, then runs a beam whose widths depend on the data -- - so without this leaf ``tzrec.predict`` dies inside ``_rewrite_model`` with - "Proxy object cannot be iterated". Wrapping the WHOLE decode is what makes - it work: a leaf returns a single Proxy, so wrapping any inner helper only - moves the failure to its caller. The trace then finds nothing to shard, - which is correct -- a SID feature carries no embedding table. - - At run time this is an ordinary call; only tracing sees a leaf. + TorchRec's predict pipeline FX-traces the model; the decode is untraceable + (per-row python lists, data-dependent beam widths). Wrapping the WHOLE + decode is required -- a leaf returns one Proxy, so wrapping an inner helper + just moves the failure to its caller. At run time this is a normal call. """ return model._generate(batch) class GenerativeQwen(BaseGenerativeModel): - """Generative-recommendation LM on a Qwen backbone (Qwen2.5, Qwen3, ...). + """Generative-recommendation LM on a Qwen backbone (Qwen2.5, Qwen3, ...).""" - The whole family shares the ChatML frame this class splices, so the version - only enters through ``hf_model_id``. - """ - - # Flat width used when beam_widths is left empty; mirrors the proto comment. DEFAULT_BEAM_WIDTH = 50 - # ChatML frame. Family-specific; a subclass overrides it wholesale. CHAT_TEMPLATE = { "user_prefix": "<|im_start|>user\n", "user_suffix": "<|im_end|>\n", @@ -85,10 +73,7 @@ def __init__( def _read_beam_config( self, common: generative_model_pb2.GenerativeModelConfig ) -> None: - """Parse the decode knobs; the width schedule must match the codebook. - - An empty ``beam_widths`` is a flat ``DEFAULT_BEAM_WIDTH`` at every level. - """ + """Parse the decode knobs; the width schedule must match the codebook.""" self._num_return = int(common.num_return_sequences) self._beam_widths: List[int] = ( list(common.beam_widths) or [self.DEFAULT_BEAM_WIDTH] * self._num_levels @@ -123,8 +108,7 @@ def _compute_max_total_length(self) -> int: def _gaps(self) -> List[torch.Tensor]: """The N+1 static prompt fragments around the N slots, template order. - Re-read every time: ``.to(device)`` rebinds the buffer, so a cached list - would keep handing back pre-move tensors. + Re-read every time: ``.to()`` rebinds the buffer, so a cache goes stale. """ return [getattr(self, f"tpl_gap_{i}") for i in range(len(self._slot_names) + 1)] @@ -135,13 +119,11 @@ def _build_prompt_tokens( ) -> None: """Tokenise the static prompt once, as the N+1 gaps around the N slots. - Each feature's ``prefix_text`` / ``suffix_text`` is folded into the - adjacent gap so every gap is ONE string tokenized in one call, keeping - the encoding bit-identical to the fully rendered prompt -- splitting - token ids instead would let a BPE merge span a seam. Splicing values - between gaps is exact only because the ``C*`` atoms are added-vocab - tokens, which HF fast tokenizers pre-split on. Buffers are - non-persistent: they follow ``.to()`` but stay out of the state_dict. + Every gap is ONE string encoded in one call, so a BPE merge cannot span + a seam. Splicing values between gaps is exact only because the ``C*`` + atoms are added-vocab tokens, which fast tokenizers pre-split on. + Buffers are non-persistent: they follow ``.to()`` but stay off the + state_dict. """ tpl = type(self).CHAT_TEMPLATE gaps, features = self._resolve_prompt_slots(cfg.prompt_template) @@ -152,13 +134,12 @@ def _build_prompt_tokens( if i < len(features) else tpl["user_suffix"] + tpl["asst_prefix"] ) - # explicit <|im_start|> markers frame the prompt; no auto BOS/EOS. + # the template carries its own markers; no auto BOS/EOS. ids = torch.tensor( tokenizer.encode(head + gap + tail, add_special_tokens=False), dtype=torch.long, ) self.register_buffer(f"tpl_gap_{i}", ids, persistent=False) - # closes the assistant turn after the answer; not part of any gap. self.register_buffer( "tpl_asst_suffix", torch.tensor( @@ -167,7 +148,7 @@ def _build_prompt_tokens( ), persistent=False, ) - # the trailing eos is a SUPERVISED token; cache it for the splice. + # the trailing eos is SUPERVISED. self.register_buffer( "tpl_eos", torch.tensor([int(tokenizer.eos_token_id)], dtype=torch.long), @@ -175,13 +156,9 @@ def _build_prompt_tokens( ) def _prompt_rows(self, slot_rows: List[List[torch.Tensor]]) -> List[torch.Tensor]: - """Per-row ``[gap_0 | slot_0 | gap_1 | ... | slot_N-1 | gap_N]``. - - Shared by the teacher-forced splice and the answer-less inference prompt. - """ + """Per-row ``[gap_0 | slot_0 | gap_1 | ... | slot_N-1 | gap_N]``.""" gaps = self._gaps - # zip transposes per-slot row lists into one tuple of slots per row, and - # pairs each slot with the gap that precedes it (the last gap has none). + # zip pairs each slot with the gap before it; gaps[-1] closes the row. return [ torch.cat([*chain.from_iterable(zip(gaps, slots)), gaps[-1]]) for slots in zip(*slot_rows) @@ -201,12 +178,10 @@ def _splice_input_ids( ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Build ``(input_ids, labels, attention_mask)``, each ``(B, T_max)``. - Every answer is exactly ``num_levels`` codes, so the tail - ``[answer | asst_suffix | eos]`` has a FIXED width and, after left - padding, lands in the same columns for every row -- ``labels`` is one - vectorized write. Only the answer and trailing eos are supervised; a - decode emits ``num_levels`` tokens and never produces ``asst_suffix``. - ``pad_to`` left-extends for pool pre-sizing, keeping the tail aligned. + The tail ``[answer | asst_suffix | eos]`` has a FIXED width, so after + left padding it lands in the same columns for every row and ``labels`` + is one vectorized write. Only the answer and trailing eos are + supervised. ``pad_to`` left-extends for pool pre-sizing. """ if len(slot_rows[0]) != len(label_rows): raise ValueError( @@ -258,9 +233,9 @@ def _forward_loss( ) -> Dict[str, torch.Tensor]: """Teacher-forced forward over spliced ids -> suffix-slice -> CE loss.""" outputs = self.lm.model(input_ids=input_ids, attention_mask=attention_mask) - hidden = outputs.last_hidden_state # (B, T, D) + hidden = outputs.last_hidden_state - # Bound the logits to the supervised suffix; a full (B, T, V) upcast OOMs. + # a full (B, T, V) upcast OOMs. suffix = slice(-self._suffix_keep, None) logits = self.lm.lm_head(hidden[:, suffix, :]) @@ -273,12 +248,7 @@ def _forward_loss( return {"loss": loss} def _generate(self, batch: Batch) -> torch.Tensor: - """Beam-search the SID answer, no ground truth supplied. - - Returns ``(B, C, num_levels)`` best-first per row, where ``C`` is - ``num_return_sequences``. Candidates come back score-ordered, so - trimming to ``num_return`` keeps the best ones. - """ + """Beam-search the SID answer; returns ``(B, C, num_levels)`` best-first.""" slot_rows = self._slot_rows(self.build_input(batch)) input_ids, attention_mask = self._left_pad(self._prompt_rows(slot_rows)) lo_tok, hi_tok = self._sid_token_bands() @@ -298,9 +268,8 @@ def _left_pad( ) -> Tuple[torch.Tensor, torch.Tensor]: """Left-pad token rows into ``(input_ids, attention_mask)``, ``(B, T_max)``. - The mask comes from ``ones_like(row)``, not ``!= pad``, so a real - trailing eos survives ``pad_token_id == eos``. ``pad_to`` extends LEFT, - keeping the end-aligned supervised tail in place. + The mask comes from ``ones_like``, not ``!= pad``, so a real trailing + eos survives ``pad_token_id == eos``. """ input_ids = pad_sequence( rows, diff --git a/tzrec/modules/dynamic_beam.py b/tzrec/modules/dynamic_beam.py index d2d33db34..e99e4d79d 100644 --- a/tzrec/modules/dynamic_beam.py +++ b/tzrec/modules/dynamic_beam.py @@ -9,11 +9,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Dynamic-width beam SID decode (no tzrec deps). +"""Band-restricted beam SID decode (no tzrec deps). -Backs ``dynamic_beam`` in ``GenerativeModelConfig``: the beam width varies per -SID level instead of staying fixed. The caller owns the schedule -- this module -only enforces what each level can actually supply. +The caller owns the width schedule; this module only enforces what each level +can supply. """ from typing import List, Tuple @@ -38,16 +37,14 @@ def dynamic_beam_search( model: an HF causal LM exposing ``.model`` / ``.lm_head`` (Qwen layout). input_ids: left-padded prompt ids ``(B, P)``. attention_mask: prompt mask ``(B, P)``. - beam_widths: requested width for each SID level, one entry per level. - Any schedule is accepted -- doubling, flat, hand-tuned -- and each - entry is capped to what its band and the surviving prefixes supply. - lo_tok: inclusive lower per-level token-space band edge, ``(num_levels,)``. - hi_tok: inclusive upper per-level token-space band edge, ``(num_levels,)``. + beam_widths: requested width per SID level; each is capped to what its + band and the surviving prefixes supply. + lo_tok: inclusive lower per-level token band edge, ``(num_levels,)``. + hi_tok: inclusive upper per-level token band edge, ``(num_levels,)``. Returns: - The generated SID token tail ``(B * W, num_levels)`` score-ordered - best-first per row, where ``W`` is the last capped width. The answer is - fixed-length and EOS-free, so no finished-beam bookkeeping is needed. + The SID token tail ``(B * W, num_levels)``, score-ordered best-first. + The answer is fixed-length and EOS-free, so no beam bookkeeping. """ device = input_ids.device bsz = input_ids.shape[0] @@ -65,7 +62,6 @@ def dynamic_beam_search( bands: List[Tuple[int, int]] = [ (int(lo_tok[j]), int(hi_tok[j])) for j in range(num_levels) ] - # A level can only carry band x surviving prefixes, however much was asked. widths: List[int] = [] prev = 1 for w, (lo, hi) in zip(beam_widths, bands): @@ -75,8 +71,8 @@ def dynamic_beam_search( def _band_logp(logits: torch.Tensor, j: int) -> torch.Tensor: """Full-vocab log-probs, narrowed to level ``j``'s band ``(R, band)``. - Slicing after normalizing keeps the exact cross-beam ranking of a - full-vocab ``log_softmax`` without materializing one per level. + Normalize then slice: same ranking as a full-vocab log_softmax, 21x less + memory at production vocab. """ lo, hi = bands[j] log_z = torch.logsumexp(logits.float(), dim=-1, keepdim=True) @@ -101,8 +97,7 @@ def _band_logp(logits: torch.Tensor, j: int) -> torch.Tensor: for j in range(1, num_levels): am = torch.cat([am, am.new_ones(bsz * cur_w, 1)], dim=1) - # the row always ends on the token just appended, so its position is - # simply how many real tokens precede it. + # the row ends on the new token, so its position is the count before it. step_pos = am.long().sum(-1, keepdim=True) - 1 cache_pos = torch.tensor([past.get_seq_length()], device=device) h = model.model( diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index f179b1269..d1f4ce6d6 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -2,77 +2,48 @@ syntax = "proto2"; package tzrec.protos; // Generative (causal-LM) recommendation models. -// -// Shared, architecture-agnostic config lives in `GenerativeModelConfig`, embedded -// as `common` in every family message; family-specific knobs (the backbone, the -// chat template) live on the family message itself. -// Storage dtype of the backbone's master weights. enum ParamDtype { FP32 = 0; BF16 = 1; FP16 = 2; } -// Architecture-agnostic config shared by all generative-rec families. -// Sample contract: -// * history : list -- local 0-based per-level codes in -// [0, codebook[level]), laid out as whole items in level order; -// a JAGGED_SEQUENCE feature_group holding it alone. -// * answer : list -- one item's local 0-based codes; the FIRST -// `data_config.label_field`, NOT a feature. -// SidFeature folds in level_offsets at parse time and the model adds -// base_vocab, so token_id = base_vocab + level_offsets[level] + code. Sample -// writers must not pre-apply level_offsets. +// Sample contract: history is a JAGGED_SEQUENCE feature_group holding one +// SidFeature; the answer is the FIRST data_config.label_field, not a feature. +// Both are local 0-based per-level codes -- writers must not pre-apply +// level_offsets, SidFeature folds them in at parse time. message GenerativeModelConfig { - // Pad the post-extension vocab up to a multiple of this value; 0 disables - // padding. + // 0 disables padding. optional uint32 vocab_pad_to_multiple_of = 3 [default = 128]; - // Cross-entropy ignore index -- matches PyTorch's F.cross_entropy default. optional int32 ignore_index = 6 [default = -100]; - // Beam search (inference only). Decoding is always restricted to the SID - // bands, so every returned candidate is a well-formed SID. - // - // Beam width per SID level, one entry per level -- the whole schedule, - // stated outright rather than derived from a base and a flag: - // [50, 50, 50] fixed width, returns 50 candidates - // [100, 200, 400] doubling, returns 400 (the ALGR-style escalating beam) - // Each entry is capped to what its band and the surviving prefixes can - // supply. Empty defaults to a flat width of 50 at every level. + // Beam width per SID level, one entry per level; [100, 200, 400] is the + // escalating beam. Each entry is capped to what its band can supply. + // Empty defaults to a flat width of 50. repeated uint32 beam_widths = 7; - // Candidates kept per row, best-first; must not exceed the final width. + // Must not exceed the final beam width. optional uint32 num_return_sequences = 8 [default = 50]; - // Prediction key the generated SIDs are emitted under; reference it from - // PredictWrapper output_cols. optional string generated_sids_key = 12 [default = "generated_sids"]; - // Backbone PARAM dtype = the MASTER weights. FP32 avoids bf16-ULP - // underflow of Adam's small (lr=1e-5) updates; bf16 COMPUTE comes from - // mixed_precision:"BF16" autocast, NOT the param dtype. + // MASTER weights. FP32 avoids bf16-ULP underflow of Adam's small updates; + // bf16 COMPUTE comes from mixed_precision, not from this. optional ParamDtype param_dtype = 13 [default = FP32]; - // Model's history budget in SID codes: the item-aligned, recency-preserving - // truncation cap AND the activation-pool pre-size. Distinct from the history - // feature's own sequence_length. Set to 0 to disable both. + // History budget in SID codes: item-aligned truncation keeping the most + // recent items, and the activation-pool pre-size. 0 disables both. required uint32 max_sequence_length = 14; } -// Qwen family (Qwen2.5-0.5B, Qwen3-0.6B, etc.) -- one ChatML frame for all. +// Qwen family (Qwen2.5, Qwen3, ...) -- one ChatML frame for all. message GenerativeQwen { optional GenerativeModelConfig common = 1; - // Qwen backbone: HF hub id or local path. + // HF hub id or local path. optional string hf_model_id = 2 [default = "Qwen/Qwen2.5-0.5B"]; - // ----- Prompt ----- - // Prompt body for the user turn, carrying one {{feature_name}} placeholder - // per SID feature whose codes are spliced in: - // "... Each behavior is represented by three words. {{user_sequence}} - // Please predict the semantic encoding of the next behavior." - // That feature's own prefix_text/suffix_text wrap the codes inside the slot. - // The ChatML frame and the post-answer region stay family constants -- the - // model owns <|im_start|>/<|im_end|> and the supervised tail. + // User-turn body carrying one {{feature_name}} placeholder per SID feature. + // The ChatML frame and the supervised tail are family constants. optional string prompt_template = 10; } diff --git a/tzrec/utils/hf_export_util.py b/tzrec/utils/hf_export_util.py index 38348de18..39bee5cf6 100644 --- a/tzrec/utils/hf_export_util.py +++ b/tzrec/utils/hf_export_util.py @@ -9,10 +9,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""HuggingFace export for HF-backed models (``BaseGenerativeModel`` family). +"""HuggingFace export for HF-backed models. -Kept out of ``export_util`` (TorchScript/TRT/AOTI) so ``checkpoint_util`` can -call ``write_hf_assets`` without a circular import. +Kept out of ``export_util`` so ``checkpoint_util`` can call it without a +circular import. """ import json @@ -27,7 +27,6 @@ from tzrec.utils import checkpoint_util from tzrec.utils.logging_util import logger -# Missing files are skipped -- tokenizers emit different subsets. _HF_ASSET_FILES = ( "config.json", "generation_config.json", @@ -45,11 +44,9 @@ def _unwrap_hf_model(wrapped_model: nn.Module) -> Optional[nn.Module]: """Walk DMP/TrainWrapper layers down to the model exposing ``hf_backbone``. - ``None`` when the chain has none: not HF-backed, so callers no-op. - - ``seen`` bounds the walk. Every checkpoint save of every model reaches here, - so a ``.model`` / ``.module`` cycle would hang inside ``save()`` -- the worst - place for it, since peers would block on the collective that follows. + ``None`` when the chain has none, so callers no-op. ``seen`` bounds the + walk: every checkpoint save reaches here, and a ``.model``/``.module`` cycle + would hang inside ``save()``. """ m = wrapped_model seen = set() @@ -69,9 +66,8 @@ def _unwrap_hf_model(wrapped_model: nn.Module) -> Optional[nn.Module]: def write_hf_assets(wrapped_model: nn.Module, save_dir: str) -> None: """Co-locate the HF config + tokenizer (NO weights) in a checkpoint dir. - The backbone's FQN prefix, read off the live module graph, goes into - ``hf_export_meta.json`` so ``dcp_to_hf`` can strip it without hard-coding a - wrapper convention. Rank 0 only; the dense backbone is replicated. + The backbone's FQN prefix goes into ``hf_export_meta.json`` so ``dcp_to_hf`` + can strip it without hard-coding a wrapper convention. Rank 0 only. """ if int(os.environ.get("RANK", 0)) != 0: return @@ -100,9 +96,8 @@ def write_hf_assets(wrapped_model: nn.Module, save_dir: str) -> None: def dcp_to_hf(ckpt_dir: str, out_dir: str) -> None: """Convert a self-contained checkpoint dir to a ``from_pretrained`` HF dir. - Everything comes from ``ckpt_dir``: no live model, no download. Keys that do - not map 1:1 onto the co-located ``config.json`` raise instead of writing a - partial model. + Keys that do not map 1:1 onto the co-located ``config.json`` raise rather + than write a partial model. """ from torch.distributed.checkpoint.state_dict_loader import ( _load_state_dict_from_keys, @@ -152,7 +147,6 @@ def _derive_by_suffix( out[tk] = state[matches[0]] return out - # A stale recorded prefix falls back to prefix-free suffix matching. mapped = _strip_recorded_prefix(raw_state) if mapped is None: if prefix: @@ -162,8 +156,6 @@ def _derive_by_suffix( ) mapped = _derive_by_suffix(raw_state) - # both strategies return either an exact match or None, so there is nothing - # partial to report -- show both key spaces instead. if mapped is None: raise RuntimeError( "dcp_to_hf: cannot map the DCP state dict onto the backbone " @@ -173,7 +165,7 @@ def _derive_by_suffix( "Refusing to write a partially-loaded HF model." ) - # Tied heads are dropped after validation; from_pretrained re-ties them. + # from_pretrained re-ties them. if getattr(cfg, "tie_word_embeddings", False): mapped = {k: v for k, v in mapped.items() if k not in tied_keys} mapped = {k: v.contiguous() for k, v in mapped.items()} # save_file rejects views From 27fd5ba488c07d33af4b82d6d5e03ffe0df629e8 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 03:28:13 +0000 Subject: [PATCH 49/99] [refactor] genrec: own the beam config in the base and name the kernel locals beam_widths and num_return_sequences live on GenerativeModelConfig, so parsing them belonged in the base rather than the Qwen subclass. _read_beam_config and DEFAULT_BEAM_WIDTH move to BaseGenerativeModel, called from _read_common_config once _num_levels is known; the width is read through self so a family can override it. Its test moves to the base test file for the same reason. dynamic_beam_search drops the keyword-only marker, and its locals get names that say what they hold: am -> beam_mask, bsz -> batch_size, h -> outputs, cur_w -> width, lo_j/hi_j -> band_lo/band_hi, idx -> flat_idx, tok -> next_token, j -> level. Also drops the export.proto comments describing ExportFormat, and corrects a stale one on SidFeature.sequence_length that still claimed fg was forbidden and the field rejected -- both untrue since the fg passthrough landed. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/generative_model.py | 25 +++++++ tzrec/models/generative_model_test.py | 21 ++++++ tzrec/models/generative_qwen.py | 25 ------- tzrec/models/generative_qwen_test.py | 19 ------ tzrec/modules/dynamic_beam.py | 93 ++++++++++++++------------- tzrec/protos/export.proto | 5 -- tzrec/protos/feature.proto | 24 +++---- 7 files changed, 103 insertions(+), 109 deletions(-) diff --git a/tzrec/models/generative_model.py b/tzrec/models/generative_model.py index 47f185152..9777af96a 100644 --- a/tzrec/models/generative_model.py +++ b/tzrec/models/generative_model.py @@ -43,6 +43,9 @@ class BaseGenerativeModel(BaseModel): """Model construction, SID vocab extension, data-prep, loss and metrics.""" + # Flat width used when beam_widths is empty; a family may override it. + DEFAULT_BEAM_WIDTH = 50 + _PARAM_DTYPE: Dict[int, torch.dtype] = { generative_model_pb2.FP32: torch.float32, generative_model_pb2.BF16: torch.bfloat16, @@ -108,8 +111,30 @@ def _read_common_config( "_level_offsets", torch.cumsum(sizes, 0) - sizes, persistent=False ) self._vocab_pad_mult = int(common.vocab_pad_to_multiple_of) + self._read_beam_config(common) return sum(codebook) + def _read_beam_config( + self, common: generative_model_pb2.GenerativeModelConfig + ) -> None: + """Parse the decode knobs; the width schedule must match the codebook.""" + self._num_return = int(common.num_return_sequences) + self._beam_widths: List[int] = ( + list(common.beam_widths) or [self.DEFAULT_BEAM_WIDTH] * self._num_levels + ) + if len(self._beam_widths) != self._num_levels: + raise ValueError( + f"{type(self).__name__}: beam_widths has " + f"{len(self._beam_widths)} entries but the codebook has " + f"{self._num_levels} levels; give one width per level." + ) + if self._num_return > self._beam_widths[-1]: + raise ValueError( + f"{type(self).__name__}: num_return_sequences " + f"({self._num_return}) must not exceed the final beam width " + f"({self._beam_widths[-1]})." + ) + def _shared_sid_space(self) -> List[int]: """The one codebook every SID feature declares.""" spaces = { diff --git a/tzrec/models/generative_model_test.py b/tzrec/models/generative_model_test.py index eb291d57f..9548956f9 100644 --- a/tzrec/models/generative_model_test.py +++ b/tzrec/models/generative_model_test.py @@ -59,6 +59,8 @@ def _common(**overrides): "param_dtype": generative_model_pb2.FP32, "vocab_pad_to_multiple_of": 128, "max_sequence_length": 0, + "beam_widths": [], + "num_return_sequences": 50, } return types.SimpleNamespace(**{**fields, **overrides}) @@ -294,6 +296,25 @@ def test_resolve_prompt_slots_rejects_a_mismatched_template(self) -> None: with self.assertRaisesRegex(ValueError, "no feature_config"): orphan._resolve_prompt_slots("{{user_sequence}}") + def test_beam_config_defaults_and_validation(self) -> None: + def read(widths, num_return, levels=3): + m = object.__new__(GenerativeQwen) + m._num_levels = levels + m._read_beam_config( + types.SimpleNamespace( + beam_widths=widths, num_return_sequences=num_return + ) + ) + return m + + # empty -> flat DEFAULT_BEAM_WIDTH per level; anything else verbatim + self.assertEqual(read([], 50)._beam_widths, [50, 50, 50]) + self.assertEqual(read([100, 200, 400], 400)._beam_widths, [100, 200, 400]) + with self.assertRaisesRegex(ValueError, "one width per level"): + read([50, 50], 50) + with self.assertRaisesRegex(ValueError, "must not exceed the final"): + read([50, 50, 50], 80) + def test_abstract_hooks_raise(self) -> None: base = object.__new__(BaseGenerativeModel) with self.assertRaises(NotImplementedError): diff --git a/tzrec/models/generative_qwen.py b/tzrec/models/generative_qwen.py index 16f291816..7b59cc977 100644 --- a/tzrec/models/generative_qwen.py +++ b/tzrec/models/generative_qwen.py @@ -45,8 +45,6 @@ def _fx_wrapped_generate(model: "GenerativeQwen", batch: Batch) -> torch.Tensor: class GenerativeQwen(BaseGenerativeModel): """Generative-recommendation LM on a Qwen backbone (Qwen2.5, Qwen3, ...).""" - DEFAULT_BEAM_WIDTH = 50 - CHAT_TEMPLATE = { "user_prefix": "<|im_start|>user\n", "user_suffix": "<|im_end|>\n", @@ -63,34 +61,11 @@ def __init__( **kwargs: Any, ) -> None: super().__init__(model_config, features, labels, sample_weights, **kwargs) - common = self._model_config.common - self._read_beam_config(common) self._max_total_len = self._compute_max_total_length() self._pool_warmed = False # +2 = trailing eos + HF's shift-by-one; constant width avoids a per-step sync. self._suffix_keep = self._num_levels + self.tpl_asst_suffix.numel() + 2 - def _read_beam_config( - self, common: generative_model_pb2.GenerativeModelConfig - ) -> None: - """Parse the decode knobs; the width schedule must match the codebook.""" - self._num_return = int(common.num_return_sequences) - self._beam_widths: List[int] = ( - list(common.beam_widths) or [self.DEFAULT_BEAM_WIDTH] * self._num_levels - ) - if len(self._beam_widths) != self._num_levels: - raise ValueError( - f"{type(self).__name__}: beam_widths has " - f"{len(self._beam_widths)} entries but the codebook has " - f"{self._num_levels} levels; give one width per level." - ) - if self._num_return > self._beam_widths[-1]: - raise ValueError( - f"{type(self).__name__}: num_return_sequences " - f"({self._num_return}) must not exceed the final beam width " - f"({self._beam_widths[-1]})." - ) - def _compute_max_total_length(self) -> int: """The ``T`` the activation pool pre-sizes to; 0 when disabled.""" if self._max_seq_length <= 0: diff --git a/tzrec/models/generative_qwen_test.py b/tzrec/models/generative_qwen_test.py index 49abc1d78..a563c9391 100644 --- a/tzrec/models/generative_qwen_test.py +++ b/tzrec/models/generative_qwen_test.py @@ -260,25 +260,6 @@ def test_generate_trims_to_num_return_keeping_the_best(self) -> None: self.assertEqual(tuple(sids.shape), (1, 2, 3)) self.assertEqual(sids[0].tolist(), [[0, 0, 0], [1, 2, 3]]) - def test_beam_config_defaults_and_validation(self) -> None: - def read(widths, num_return, levels=3): - m = object.__new__(GenerativeQwen) - m._num_levels = levels - m._read_beam_config( - types.SimpleNamespace( - beam_widths=widths, num_return_sequences=num_return - ) - ) - return m - - # empty -> flat DEFAULT_BEAM_WIDTH per level; anything else verbatim - self.assertEqual(read([], 50)._beam_widths, [50, 50, 50]) - self.assertEqual(read([100, 200, 400], 400)._beam_widths, [100, 200, 400]) - with self.assertRaisesRegex(ValueError, "one width per level"): - read([50, 50], 50) - with self.assertRaisesRegex(ValueError, "must not exceed the final"): - read([50, 50, 50], 80) - def test_generate_hands_the_kernel_prompt_bands_and_schedule(self) -> None: m = _stub(base_vocab=100) m._slot_names = ["user_sequence"] diff --git a/tzrec/modules/dynamic_beam.py b/tzrec/modules/dynamic_beam.py index e99e4d79d..e7be2bfd3 100644 --- a/tzrec/modules/dynamic_beam.py +++ b/tzrec/modules/dynamic_beam.py @@ -26,7 +26,6 @@ def dynamic_beam_search( model: PreTrainedModel, input_ids: torch.Tensor, attention_mask: torch.Tensor, - *, beam_widths: List[int], lo_tok: torch.Tensor, hi_tok: torch.Tensor, @@ -47,7 +46,7 @@ def dynamic_beam_search( The answer is fixed-length and EOS-free, so no beam bookkeeping. """ device = input_ids.device - bsz = input_ids.shape[0] + batch_size = input_ids.shape[0] num_levels = lo_tok.shape[0] if len(beam_widths) != num_levels: raise ValueError( @@ -60,66 +59,70 @@ def dynamic_beam_search( ) # Hoist the band edges to host once to keep the level loop sync-free. bands: List[Tuple[int, int]] = [ - (int(lo_tok[j]), int(hi_tok[j])) for j in range(num_levels) + (int(lo_tok[level]), int(hi_tok[level])) for level in range(num_levels) ] - widths: List[int] = [] - prev = 1 - for w, (lo, hi) in zip(beam_widths, bands): - widths.append(min(w, prev * (hi - lo + 1))) - prev = widths[-1] + capped_widths: List[int] = [] + prev_width = 1 + for requested, (band_lo, band_hi) in zip(beam_widths, bands): + capped_widths.append(min(requested, prev_width * (band_hi - band_lo + 1))) + prev_width = capped_widths[-1] - def _band_logp(logits: torch.Tensor, j: int) -> torch.Tensor: - """Full-vocab log-probs, narrowed to level ``j``'s band ``(R, band)``. + def _band_logp(logits: torch.Tensor, level: int) -> torch.Tensor: + """Full-vocab log-probs, narrowed to ``level``'s band ``(rows, band)``. Normalize then slice: same ranking as a full-vocab log_softmax, 21x less memory at production vocab. """ - lo, hi = bands[j] + band_lo, band_hi = bands[level] log_z = torch.logsumexp(logits.float(), dim=-1, keepdim=True) - return logits[:, lo : hi + 1].float() - log_z + return logits[:, band_lo : band_hi + 1].float() - log_z - pos = (attention_mask.long().cumsum(-1) - 1).clamp(min=0) - h = model.model( + position_ids = (attention_mask.long().cumsum(-1) - 1).clamp(min=0) + outputs = model.model( input_ids=input_ids, attention_mask=attention_mask, - position_ids=pos, + position_ids=position_ids, use_cache=True, ) - past = h.past_key_values - scores = _band_logp(model.lm_head(h.last_hidden_state[:, -1, :]), 0) - beam_scores, local = scores.topk(widths[0], dim=-1) # (B, W0) - seq = (local + bands[0][0]).reshape(-1, 1) + cache = outputs.past_key_values + scores = _band_logp(model.lm_head(outputs.last_hidden_state[:, -1, :]), 0) + beam_scores, in_band = scores.topk(capped_widths[0], dim=-1) + seq = (in_band + bands[0][0]).reshape(-1, 1) beam_scores = beam_scores.reshape(-1) - rows = torch.arange(bsz, device=device) - past.reorder_cache(rows.repeat_interleave(widths[0])) - am = attention_mask.repeat_interleave(widths[0], dim=0) - cur_w = widths[0] + row_starts = torch.arange(batch_size, device=device) + cache.reorder_cache(row_starts.repeat_interleave(capped_widths[0])) + beam_mask = attention_mask.repeat_interleave(capped_widths[0], dim=0) + width = capped_widths[0] - for j in range(1, num_levels): - am = torch.cat([am, am.new_ones(bsz * cur_w, 1)], dim=1) + for level in range(1, num_levels): + beam_mask = torch.cat( + [beam_mask, beam_mask.new_ones(batch_size * width, 1)], dim=1 + ) # the row ends on the new token, so its position is the count before it. - step_pos = am.long().sum(-1, keepdim=True) - 1 - cache_pos = torch.tensor([past.get_seq_length()], device=device) - h = model.model( + step_position = beam_mask.long().sum(-1, keepdim=True) - 1 + cache_position = torch.tensor([cache.get_seq_length()], device=device) + outputs = model.model( input_ids=seq[:, -1:], - attention_mask=am, - position_ids=step_pos, - past_key_values=past, + attention_mask=beam_mask, + position_ids=step_position, + past_key_values=cache, use_cache=True, - cache_position=cache_pos, + cache_position=cache_position, + ) + band_lo, band_hi = bands[level] + band_size = band_hi - band_lo + 1 + scores = _band_logp(model.lm_head(outputs.last_hidden_state[:, -1, :]), level) + scores = scores + beam_scores[:, None] + beam_scores, flat_idx = scores.view(batch_size, width * band_size).topk( + capped_widths[level], dim=-1 ) - lo_j, hi_j = bands[j] - band = hi_j - lo_j + 1 - scores = _band_logp(model.lm_head(h.last_hidden_state[:, -1, :]), j) - scores = scores + beam_scores[:, None] # (B*cur_w, band) cumulative - beam_scores, idx = scores.view(bsz, cur_w * band).topk(widths[j], dim=-1) - tok = lo_j + idx % band - parent = (idx // band + rows[:, None] * cur_w).reshape(-1) - seq = torch.cat([seq[parent], tok.reshape(-1, 1)], dim=1) + next_token = band_lo + flat_idx % band_size + parent = (flat_idx // band_size + row_starts[:, None] * width).reshape(-1) + seq = torch.cat([seq[parent], next_token.reshape(-1, 1)], dim=1) beam_scores = beam_scores.reshape(-1) - cur_w = widths[j] - if j + 1 < num_levels: + width = capped_widths[level] + if level + 1 < num_levels: # the last level never reads the cache; skip the largest reorder copy. - past.reorder_cache(parent) - am = am[parent] - return seq # (B*cur_w, num_levels) + cache.reorder_cache(parent) + beam_mask = beam_mask[parent] + return seq diff --git a/tzrec/protos/export.proto b/tzrec/protos/export.proto index b225af7f4..5df8a4f78 100644 --- a/tzrec/protos/export.proto +++ b/tzrec/protos/export.proto @@ -1,9 +1,6 @@ syntax = "proto2"; package tzrec.protos; -// serialization format produced by `tzrec.export`. -// TORCHSCRIPT: native scripted_model.pt (+ TRT/AOTI), the default. -// HF: a HuggingFace `from_pretrained`-loadable dir (BaseGenerativeModel family). enum ExportFormat { TORCHSCRIPT = 0; HF = 1; @@ -31,7 +28,5 @@ message ExportConfig { optional bool cuda_matmul_allow_tf32 = 6 [default = false]; // Whether export uses dense EMA parameters. optional bool use_dense_ema = 7; - // serialization format produced by `tzrec.export`; selects the export - // branch in main.export (TORCHSCRIPT native graph vs HF backbone dir). optional ExportFormat export_format = 8 [default = TORCHSCRIPT]; } diff --git a/tzrec/protos/feature.proto b/tzrec/protos/feature.proto index cc292d687..e056e13fd 100644 --- a/tzrec/protos/feature.proto +++ b/tzrec/protos/feature.proto @@ -997,13 +997,13 @@ message SequenceFeature { // feature is only supported with data_config.fg_mode = FG_NONE. There is no // pyfg counterpart. message SidFeature { - // feature name; also the {{name}} placeholder in the model prompt_template. + // feature name; also the {{name}} placeholder in prompt_template. required string feature_name = 1; // feature input, e.g. user:user_sequence required string expression = 2; - // codes per position; SID streams are flat, so this stays 1. + // SID streams are flat, so this stays 1. optional uint32 value_dim = 6 [default = 1]; - // embedding pooling type, unused (SIDs carry no embedding table). + // unused; SIDs carry no embedding table. optional string pooling = 10 [default = "sum"]; // fg default value optional string default_value = 11 [default = "0"]; @@ -1012,17 +1012,12 @@ message SidFeature { // mask value in training progress optional bool use_mask = 14; - // Text emitted immediately BEFORE this feature's SID tokens in the prompt, - // e.g. "Current user's historical behaviors are as follows:". Empty by - // default. Applies only where the model splices this feature. + // text emitted immediately BEFORE this feature's SID tokens optional string prefix_text = 20 [default = ""]; - // Text emitted immediately AFTER this feature's SID tokens. + // text emitted immediately AFTER this feature's SID tokens optional string suffix_text = 21 [default = ""]; - // SID vocabulary, one entry per RQ level: each code is 0-based in - // [0, codebook[level]), matching what the SID-generation models emit. - // len() = codes per item; sum() = atoms the model appends as C0..C{sum-1}. - // All SID features in a model share ONE space, so every SidFeature must - // declare the SAME codebook -- adding a feature must not grow the vocabulary. + // per-level SID vocabulary; codes are 0-based in [0, codebook[level]). + // All SID features in a model must declare the SAME codebook. repeated uint32 codebook = 23; // default value when fg_mode = FG_NONE @@ -1032,9 +1027,8 @@ message SidFeature { // embedding param constraints optional ParameterConstraints embedding_constraints = 50; - // NOT a truncation cap here -- fg (which a SID feature forbids) and an - // fx export marker are its only readers, so setting it is rejected. - // Use model_config.common.max_sequence_length for the history budget. + // fg-side cap; must be a multiple of len(codebook). Prefer + // model_config.common.max_sequence_length, which keeps the recent items. optional uint32 sequence_length = 101; // sequence delimiter, only take effect when use it as sequence optional string sequence_delim = 102 [default = ";"]; From 712c863a11db9108f85951234df540e85c2ed375 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 03:46:08 +0000 Subject: [PATCH 50/99] [refactor] genrec: drop a redundant test and two dead SidFeature fields GenerativeQwenBeamTest claimed to cover the kernel/model seam, but mutation testing showed it does not: breaking _validate_sid_candidates' sentinel leaves it passing, because band-restricted output never reaches that path. The two mutations it does catch -- the band edges and the width cap -- are already caught by generative_model_test and dynamic_beam_test respectively. separator and embedding_constraints are unreachable for a SID feature: _fg_json never emits a separator, and parameter_constraints is only consulted for a feature with an emb_config, which a SID feature does not have. Removing either keeps both the unit and integration suites green. pooling looked equally dead and is not: EmbeddingGroup reads pooling_type for every member of a group, so removing it fails the integration test with AttributeError. use_mask and value_dim are likewise read by BaseFeature. All three stay. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/generative_qwen_test.py | 36 ---------------------------- tzrec/protos/feature.proto | 4 ---- 2 files changed, 40 deletions(-) diff --git a/tzrec/models/generative_qwen_test.py b/tzrec/models/generative_qwen_test.py index a563c9391..5fd8df2b5 100644 --- a/tzrec/models/generative_qwen_test.py +++ b/tzrec/models/generative_qwen_test.py @@ -18,7 +18,6 @@ from torch import nn from tzrec.models.generative_qwen import GenerativeQwen -from tzrec.modules.dynamic_beam import dynamic_beam_search from tzrec.utils.test_util import create_tiny_causal_lm, parameterized_name_func @@ -417,40 +416,5 @@ def test_forward_loss_returns_only_the_loss(self) -> None: self.assertEqual(list(out), ["loss"]) -class GenerativeQwenBeamTest(unittest.TestCase): - """The kernel/model seam, which neither side can assert alone. - - ``dynamic_beam_test`` owns the band masking and the per-level capping, but - only the model knows the bands and owns ``_validate_sid_candidates``, so this - is where "the kernel's output is exactly what the validator accepts" lands. - """ - - def test_band_masked_beams_decode_without_sentinels(self) -> None: - # real backbone: band masking guarantees that every returned candidate - # survives _validate_sid_candidates. widths are [2, 6, 24] here, i.e. - # exhaustive over the whole 2*3*4 codebook. - codebook = [2, 3, 4] - m = _real_lm_stub(codebook=codebook, base_vocab=20, beam_width=3) - ids = torch.tensor([[5, 6, 7]]) - lo_tok, hi_tok = m._sid_token_bands() - new = dynamic_beam_search( - m.lm, - ids, - torch.ones_like(ids), - beam_widths=[3 * 2 ** (j + 1) for j in range(len(codebook))], - lo_tok=lo_tok, - hi_tok=hi_tok, - ) - sids = m._validate_sid_candidates(new, batch_size=1) - self.assertEqual(tuple(sids.shape), (1, 24, 3)) - for level, size in enumerate(codebook): - self.assertTrue(bool((sids[..., level] >= 0).all())) - self.assertTrue(bool((sids[..., level] < size).all())) - self.assertEqual( - {tuple(row) for row in sids[0].tolist()}, - {(a, b, c) for a in range(2) for b in range(3) for c in range(4)}, - ) - - if __name__ == "__main__": unittest.main() diff --git a/tzrec/protos/feature.proto b/tzrec/protos/feature.proto index e056e13fd..ea7ab39d6 100644 --- a/tzrec/protos/feature.proto +++ b/tzrec/protos/feature.proto @@ -1007,8 +1007,6 @@ message SidFeature { optional string pooling = 10 [default = "sum"]; // fg default value optional string default_value = 11 [default = "0"]; - // fg multi-value separator - optional string separator = 12 [default = "\x1d"]; // mask value in training progress optional bool use_mask = 14; @@ -1024,8 +1022,6 @@ message SidFeature { optional string fg_encoded_default_value = 30; // only used as fg dag intermediate result or not optional bool stub_type = 34 [default = false]; - // embedding param constraints - optional ParameterConstraints embedding_constraints = 50; // fg-side cap; must be a multiple of len(codebook). Prefer // model_config.common.max_sequence_length, which keeps the recent items. From a3ff4795323fa129d5469fff6c320646ab6df0d1 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 03:50:18 +0000 Subject: [PATCH 51/99] [refactor] genrec: test base behaviour on the base class _read_beam_config moved to BaseGenerativeModel, but its test and the shared fixtures still built GenerativeQwen, so the base test read as though the subclass owned the method. They now construct BaseGenerativeModel; the only GenerativeQwen references left are the ones genuinely about the subclass -- registry dispatch, the oneof, and the family proto's hf_model_id default. Also reaches CHAT_TEMPLATE through self, matching DEFAULT_BEAM_WIDTH. Both are class constants a family may override and both resolve identically; using two different spellings for the same intent was the only reason to prefer one. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/generative_model_test.py | 14 ++++++++------ tzrec/models/generative_qwen.py | 2 +- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/tzrec/models/generative_model_test.py b/tzrec/models/generative_model_test.py index 9548956f9..fa7c07f7f 100644 --- a/tzrec/models/generative_model_test.py +++ b/tzrec/models/generative_model_test.py @@ -67,7 +67,7 @@ def _common(**overrides): def _wired(features=None, group_type=model_pb2.JAGGED_SEQUENCE, members=None): """Pre-``__init__`` state: the features/labels/groups the config-time code reads.""" - m = object.__new__(GenerativeQwen) + m = object.__new__(BaseGenerativeModel) nn.Module.__init__(m) m._features = [_sid_feature()] if features is None else features m._labels = ["label"] @@ -82,9 +82,9 @@ def _wired(features=None, group_type=model_pb2.JAGGED_SEQUENCE, members=None): def _stub(codebook=None, base_vocab=100, device="cpu"): - """A GenerativeQwen with the base data-prep state wired up, but no HF backbone.""" + """Base model with the data-prep state wired up, but no HF backbone.""" codebook = codebook or [2, 3, 4] - m = object.__new__(GenerativeQwen) + m = object.__new__(BaseGenerativeModel) nn.Module.__init__(m) m._base_vocab = base_vocab m._num_levels = len(codebook) @@ -152,10 +152,12 @@ def test_configurable_knob_defaults(self) -> None: self.assertEqual(c.generated_sids_key, "generated_sids") self.assertEqual(c.param_dtype, generative_model_pb2.FP32) self.assertIs( - GenerativeQwen._PARAM_DTYPE[generative_model_pb2.FP32], torch.float32 + BaseGenerativeModel._PARAM_DTYPE[generative_model_pb2.FP32], + torch.float32, ) self.assertIs( - GenerativeQwen._PARAM_DTYPE[generative_model_pb2.BF16], torch.bfloat16 + BaseGenerativeModel._PARAM_DTYPE[generative_model_pb2.BF16], + torch.bfloat16, ) def test_read_common_config_reads_knobs(self) -> None: @@ -298,7 +300,7 @@ def test_resolve_prompt_slots_rejects_a_mismatched_template(self) -> None: def test_beam_config_defaults_and_validation(self) -> None: def read(widths, num_return, levels=3): - m = object.__new__(GenerativeQwen) + m = object.__new__(BaseGenerativeModel) m._num_levels = levels m._read_beam_config( types.SimpleNamespace( diff --git a/tzrec/models/generative_qwen.py b/tzrec/models/generative_qwen.py index 7b59cc977..3b7523bef 100644 --- a/tzrec/models/generative_qwen.py +++ b/tzrec/models/generative_qwen.py @@ -100,7 +100,7 @@ def _build_prompt_tokens( Buffers are non-persistent: they follow ``.to()`` but stay off the state_dict. """ - tpl = type(self).CHAT_TEMPLATE + tpl = self.CHAT_TEMPLATE gaps, features = self._resolve_prompt_slots(cfg.prompt_template) for i, gap in enumerate(gaps): head = tpl["user_prefix"] if i == 0 else features[i - 1].suffix_text From dedea7a9e4a8005cf13204cb04b670440c20a7b8 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 03:53:19 +0000 Subject: [PATCH 52/99] [refactor] genrec: make the beam-width default a private class attribute DEFAULT_BEAM_WIDTH read as a public constant, but it is neither public nor constant: nothing outside the class touches it and a family is meant to override it. _default_beam_width says both. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/generative_model.py | 6 +++--- tzrec/models/generative_model_test.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tzrec/models/generative_model.py b/tzrec/models/generative_model.py index 9777af96a..d13df315f 100644 --- a/tzrec/models/generative_model.py +++ b/tzrec/models/generative_model.py @@ -43,8 +43,8 @@ class BaseGenerativeModel(BaseModel): """Model construction, SID vocab extension, data-prep, loss and metrics.""" - # Flat width used when beam_widths is empty; a family may override it. - DEFAULT_BEAM_WIDTH = 50 + # flat width used when beam_widths is empty; a family may override it + _default_beam_width = 50 _PARAM_DTYPE: Dict[int, torch.dtype] = { generative_model_pb2.FP32: torch.float32, @@ -120,7 +120,7 @@ def _read_beam_config( """Parse the decode knobs; the width schedule must match the codebook.""" self._num_return = int(common.num_return_sequences) self._beam_widths: List[int] = ( - list(common.beam_widths) or [self.DEFAULT_BEAM_WIDTH] * self._num_levels + list(common.beam_widths) or [self._default_beam_width] * self._num_levels ) if len(self._beam_widths) != self._num_levels: raise ValueError( diff --git a/tzrec/models/generative_model_test.py b/tzrec/models/generative_model_test.py index fa7c07f7f..8a196c32f 100644 --- a/tzrec/models/generative_model_test.py +++ b/tzrec/models/generative_model_test.py @@ -309,7 +309,7 @@ def read(widths, num_return, levels=3): ) return m - # empty -> flat DEFAULT_BEAM_WIDTH per level; anything else verbatim + # empty -> flat _default_beam_width per level; anything else verbatim self.assertEqual(read([], 50)._beam_widths, [50, 50, 50]) self.assertEqual(read([100, 200, 400], 400)._beam_widths, [100, 200, 400]) with self.assertRaisesRegex(ValueError, "one width per level"): From 386b6d60e2d91e28b92945768a892e97773e2ad1 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 03:56:27 +0000 Subject: [PATCH 53/99] [refactor] genrec: set the beam-width default in __init__ _default_beam_width is now an instance attribute assigned in __init__, before _read_common_config parses the beam knobs that consume it. Note this changes how a family overrides it: a class-level _default_beam_width on a subclass would be overwritten by this assignment, so an override has to happen through the constructor or by overriding _read_beam_config. The test fixtures that build a model with object.__new__ now set it explicitly, since they bypass __init__. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/generative_model.py | 5 ++--- tzrec/models/generative_model_test.py | 2 ++ 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tzrec/models/generative_model.py b/tzrec/models/generative_model.py index d13df315f..0c151c499 100644 --- a/tzrec/models/generative_model.py +++ b/tzrec/models/generative_model.py @@ -43,9 +43,6 @@ class BaseGenerativeModel(BaseModel): """Model construction, SID vocab extension, data-prep, loss and metrics.""" - # flat width used when beam_widths is empty; a family may override it - _default_beam_width = 50 - _PARAM_DTYPE: Dict[int, torch.dtype] = { generative_model_pb2.FP32: torch.float32, generative_model_pb2.BF16: torch.bfloat16, @@ -61,6 +58,8 @@ def __init__( **kwargs: Any, ) -> None: super().__init__(model_config, features, labels, sample_weights, **kwargs) + # flat width used when beam_widths is empty; set before the parse below + self._default_beam_width = 50 cfg = self._model_config sid_atoms = self._read_common_config(cfg.common) diff --git a/tzrec/models/generative_model_test.py b/tzrec/models/generative_model_test.py index 8a196c32f..3f3501e83 100644 --- a/tzrec/models/generative_model_test.py +++ b/tzrec/models/generative_model_test.py @@ -69,6 +69,7 @@ def _wired(features=None, group_type=model_pb2.JAGGED_SEQUENCE, members=None): """Pre-``__init__`` state: the features/labels/groups the config-time code reads.""" m = object.__new__(BaseGenerativeModel) nn.Module.__init__(m) + m._default_beam_width = 50 # __init__ is bypassed here m._features = [_sid_feature()] if features is None else features m._labels = ["label"] m._feature_groups = [ @@ -302,6 +303,7 @@ def test_beam_config_defaults_and_validation(self) -> None: def read(widths, num_return, levels=3): m = object.__new__(BaseGenerativeModel) m._num_levels = levels + m._default_beam_width = 50 # __init__ is bypassed here m._read_beam_config( types.SimpleNamespace( beam_widths=widths, num_return_sequences=num_return From d72c42905f5a6f46024383c2721c74b15914c681 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 04:54:38 +0000 Subject: [PATCH 54/99] [refactor] genrec: require the beam schedule and the SID codebook Both were silently defaultable. beam_widths fell back to a flat width of 50 per level, so a config that never mentioned the beam still decoded -- at a width nobody chose, with no way to tell an intended 50 from an unset one. num_return_sequences carried the same 50 as a proto default even though the value must relate to the final beam width. beam_widths now raises when empty, num_return_sequences is proto-required, and _default_beam_width is gone. codebook was already enforced in code; the proto and the error message now say it is required rather than merely non-empty. The mock config gains an explicit [4, 8, 16] schedule, which is what a real config has to do now too. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/features/sid_feature.py | 2 +- tzrec/features/sid_feature_test.py | 5 ++++- tzrec/models/generative_model.py | 11 ++++++----- tzrec/models/generative_model_test.py | 9 ++++----- tzrec/protos/feature.proto | 4 ++-- tzrec/protos/models/generative_model.proto | 7 +++---- tzrec/tests/configs/generative_qwen_mock.config | 4 ++++ 7 files changed, 24 insertions(+), 18 deletions(-) diff --git a/tzrec/features/sid_feature.py b/tzrec/features/sid_feature.py index 43b3c3924..a03d563ef 100644 --- a/tzrec/features/sid_feature.py +++ b/tzrec/features/sid_feature.py @@ -60,7 +60,7 @@ def _read_codebook(self) -> List[int]: if not codebook: raise ValueError( f"{self.__class__.__name__}[{self.config.feature_name}]: codebook " - f"must be non-empty." + f"is required; give one vocabulary size per SID level." ) if any(c <= 0 for c in codebook): raise ValueError( diff --git a/tzrec/features/sid_feature_test.py b/tzrec/features/sid_feature_test.py index 1f884fc53..e8ae2aea4 100644 --- a/tzrec/features/sid_feature_test.py +++ b/tzrec/features/sid_feature_test.py @@ -115,7 +115,10 @@ def test_parse_rejects_out_of_range_and_partial_items(self) -> None: f.parse({"user_sequence": pa.array([[0, 1]])}) def test_rejects_a_bad_codebook(self) -> None: - for bad, msg in (("", "non-empty"), ("codebook: 4 codebook: 0", "positive")): + for bad, msg in ( + ("", "codebook is required"), + ("codebook: 4 codebook: 0", "positive"), + ): with self.subTest(bad=bad): base = 'feature_name: "s" expression: "user:s" ' + bad with self.assertRaisesRegex(ValueError, msg): diff --git a/tzrec/models/generative_model.py b/tzrec/models/generative_model.py index 0c151c499..c15a0de1e 100644 --- a/tzrec/models/generative_model.py +++ b/tzrec/models/generative_model.py @@ -58,8 +58,6 @@ def __init__( **kwargs: Any, ) -> None: super().__init__(model_config, features, labels, sample_weights, **kwargs) - # flat width used when beam_widths is empty; set before the parse below - self._default_beam_width = 50 cfg = self._model_config sid_atoms = self._read_common_config(cfg.common) @@ -118,9 +116,12 @@ def _read_beam_config( ) -> None: """Parse the decode knobs; the width schedule must match the codebook.""" self._num_return = int(common.num_return_sequences) - self._beam_widths: List[int] = ( - list(common.beam_widths) or [self._default_beam_width] * self._num_levels - ) + self._beam_widths: List[int] = list(common.beam_widths) + if not self._beam_widths: + raise ValueError( + f"{type(self).__name__}: beam_widths is required; give one " + f"width per SID level, e.g. [50, 50, 50] or [100, 200, 400]." + ) if len(self._beam_widths) != self._num_levels: raise ValueError( f"{type(self).__name__}: beam_widths has " diff --git a/tzrec/models/generative_model_test.py b/tzrec/models/generative_model_test.py index 3f3501e83..e95221b27 100644 --- a/tzrec/models/generative_model_test.py +++ b/tzrec/models/generative_model_test.py @@ -59,7 +59,7 @@ def _common(**overrides): "param_dtype": generative_model_pb2.FP32, "vocab_pad_to_multiple_of": 128, "max_sequence_length": 0, - "beam_widths": [], + "beam_widths": [50, 50, 50], "num_return_sequences": 50, } return types.SimpleNamespace(**{**fields, **overrides}) @@ -69,7 +69,6 @@ def _wired(features=None, group_type=model_pb2.JAGGED_SEQUENCE, members=None): """Pre-``__init__`` state: the features/labels/groups the config-time code reads.""" m = object.__new__(BaseGenerativeModel) nn.Module.__init__(m) - m._default_beam_width = 50 # __init__ is bypassed here m._features = [_sid_feature()] if features is None else features m._labels = ["label"] m._feature_groups = [ @@ -303,7 +302,6 @@ def test_beam_config_defaults_and_validation(self) -> None: def read(widths, num_return, levels=3): m = object.__new__(BaseGenerativeModel) m._num_levels = levels - m._default_beam_width = 50 # __init__ is bypassed here m._read_beam_config( types.SimpleNamespace( beam_widths=widths, num_return_sequences=num_return @@ -311,9 +309,10 @@ def read(widths, num_return, levels=3): ) return m - # empty -> flat _default_beam_width per level; anything else verbatim - self.assertEqual(read([], 50)._beam_widths, [50, 50, 50]) + # the schedule is taken verbatim; there is no default to fall back on self.assertEqual(read([100, 200, 400], 400)._beam_widths, [100, 200, 400]) + with self.assertRaisesRegex(ValueError, "beam_widths is required"): + read([], 50) with self.assertRaisesRegex(ValueError, "one width per level"): read([50, 50], 50) with self.assertRaisesRegex(ValueError, "must not exceed the final"): diff --git a/tzrec/protos/feature.proto b/tzrec/protos/feature.proto index ea7ab39d6..02bf4f6a0 100644 --- a/tzrec/protos/feature.proto +++ b/tzrec/protos/feature.proto @@ -1014,8 +1014,8 @@ message SidFeature { optional string prefix_text = 20 [default = ""]; // text emitted immediately AFTER this feature's SID tokens optional string suffix_text = 21 [default = ""]; - // per-level SID vocabulary; codes are 0-based in [0, codebook[level]). - // All SID features in a model must declare the SAME codebook. + // REQUIRED, per-level SID vocabulary; codes are 0-based in + // [0, codebook[level]). All SID features must declare the SAME codebook. repeated uint32 codebook = 23; // default value when fg_mode = FG_NONE diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto index d1f4ce6d6..004ab5353 100644 --- a/tzrec/protos/models/generative_model.proto +++ b/tzrec/protos/models/generative_model.proto @@ -19,12 +19,11 @@ message GenerativeModelConfig { optional int32 ignore_index = 6 [default = -100]; - // Beam width per SID level, one entry per level; [100, 200, 400] is the - // escalating beam. Each entry is capped to what its band can supply. - // Empty defaults to a flat width of 50. + // REQUIRED. Beam width per SID level, one entry per level; [100, 200, 400] + // is the escalating beam. Each entry is capped to what its band can supply. repeated uint32 beam_widths = 7; // Must not exceed the final beam width. - optional uint32 num_return_sequences = 8 [default = 50]; + required uint32 num_return_sequences = 8; optional string generated_sids_key = 12 [default = "generated_sids"]; // MASTER weights. FP32 avoids bf16-ULP underflow of Adam's small updates; diff --git a/tzrec/tests/configs/generative_qwen_mock.config b/tzrec/tests/configs/generative_qwen_mock.config index f34d6a17e..9a28118e5 100644 --- a/tzrec/tests/configs/generative_qwen_mock.config +++ b/tzrec/tests/configs/generative_qwen_mock.config @@ -54,6 +54,10 @@ model_config { ignore_index: -100 param_dtype: FP32 max_sequence_length: 12 + beam_widths: 4 + beam_widths: 8 + beam_widths: 16 + num_return_sequences: 16 } hf_model_id: "Qwen/Qwen2.5-0.5B" prompt_template: "You are a recommendation system. Based on the user's historical behavior, predict the user's next action in an e-commerce scenario. I will provide a sequence of semantic encodings representing consecutive behaviors, arranged in chronological order of user clicks. Each behavior is represented by three words. {{user_sequence}} Please predict the semantic encoding of the user's subsequent behavior in the e-commerce recommendation scenario." From 07443e390ee8516a6b97c9573a1e4288f136b090 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 07:07:54 +0000 Subject: [PATCH 55/99] [bugfix] genrec: refuse dense EMA in HF export The HF export branch returns before the model is built and hands the checkpoint directory to dcp_to_hf, which always reads /model. Dense EMA lives in a sibling /dense_ema that only restore_model overlays, so with EMA enabled TorchScript export shipped the averaged weights while HF export silently shipped the raw ones -- same checkpoint, different model, no warning. The branch now resolves use_dense_ema exactly as the DCP path does and raises instead of converting. Refusing is preferable to overlaying here: the converter never builds a model, so it has no place to apply the EMA state without duplicating restore_model's mapping logic. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/main.py | 8 ++++++++ tzrec/main_test.py | 23 +++++++++++++++++++++-- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index e78b73fe9..d8acab684 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -1091,6 +1091,14 @@ def export( # HF export converts the checkpoint dir directly -- no model build, no DCP restore. if pipeline_config.export_config.export_format == export_pb2.ExportFormat.HF: + if config_util.use_dense_ema( + pipeline_config.export_config, pipeline_config.train_config + ): + raise ValueError( + "HF export: dcp_to_hf reads /model, so it cannot " + "serve Dense EMA parameters. Set export_config.use_dense_ema to " + "false to export the raw weights." + ) if not checkpoint_path: raise ValueError("HF export: no checkpoint found to convert.") if not os.path.exists(os.path.join(checkpoint_path, "config.json")): diff --git a/tzrec/main_test.py b/tzrec/main_test.py index 7f9e088be..253e6f7a4 100644 --- a/tzrec/main_test.py +++ b/tzrec/main_test.py @@ -19,12 +19,14 @@ from unittest import mock import torch +from google.protobuf import text_format -from tzrec.main import _train_and_evaluate +from tzrec.main import _train_and_evaluate, export from tzrec.optim.ema import DenseEMA from tzrec.protos.eval_pb2 import EvalConfig -from tzrec.protos.export_pb2 import ExportConfig +from tzrec.protos.export_pb2 import ExportConfig, ExportFormat from tzrec.protos.optimizer_pb2 import DenseOptimizer, EMAConfig +from tzrec.protos.pipeline_pb2 import EasyRecConfig class MainTest(unittest.TestCase): @@ -155,6 +157,23 @@ def assert_ema(*args, **kwargs): for call in exporter.maybe_export.call_args_list: self.assertIsNone(call.kwargs["dense_ema"]) + def test_hf_export_rejects_dense_ema(self) -> None: + # dcp_to_hf reads /model unconditionally, so it would silently + # ship raw weights where TorchScript export ships the EMA ones. + with tempfile.TemporaryDirectory() as test_dir: + config = EasyRecConfig() + config.train_input_path = "unused" + config.eval_input_path = "unused" + config.model_dir = os.path.join(test_dir, "train") + os.makedirs(config.model_dir) + config.train_config.dense_optimizer.ema.CopyFrom(EMAConfig()) + config.export_config.export_format = ExportFormat.HF + config_path = os.path.join(test_dir, "pipeline.config") + with open(config_path, "w") as f: + f.write(text_format.MessageToString(config)) + with self.assertRaisesRegex(ValueError, "Dense EMA"): + export(config_path, os.path.join(test_dir, "export")) + class TrainStepCounterMultiPassTest(unittest.TestCase): """Guard the step-based (``num_steps``) multi-pass step counter. From 17a4631159e8977ea7ed91ed4dfe920257ec2edc Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 07:07:54 +0000 Subject: [PATCH 56/99] [refactor] genrec: pin SidFeature value_dim and correct its fg comment value_dim was configurable but unrepresentable: _parse reshapes by level and splits by seq_lengths, both of which assume one code per sequence position, so a wider value lands the level offsets on the wrong components. It is now rejected at construction rather than described as "stays 1" in a comment. The SidFeature message comment still claimed FG_NONE-only support, which the fg passthrough replaced; it now describes what fg actually does. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/features/sid_feature.py | 6 ++++++ tzrec/features/sid_feature_test.py | 5 ++++- tzrec/protos/feature.proto | 6 +++--- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/tzrec/features/sid_feature.py b/tzrec/features/sid_feature.py index a03d563ef..240fdb4ce 100644 --- a/tzrec/features/sid_feature.py +++ b/tzrec/features/sid_feature.py @@ -37,6 +37,12 @@ def __init__( # BaseFeature.__del__ dereferences _fg_op, so seed it before any raise. self._fg_op = None super().__init__(feature_config, **kwargs) + if self.config.value_dim != 1: + raise ValueError( + f"{self.__class__.__name__}[{self.config.feature_name}]: " + f"value_dim must be 1 -- the SID stream is flat, one code per " + f"sequence position -- got {self.config.value_dim}." + ) self._codebook = self._read_codebook() # fg truncates by VALUE count and keeps the head, so a cap that is not a # whole number of items would hand the model partial items. The model's diff --git a/tzrec/features/sid_feature_test.py b/tzrec/features/sid_feature_test.py index e8ae2aea4..e5fda20b7 100644 --- a/tzrec/features/sid_feature_test.py +++ b/tzrec/features/sid_feature_test.py @@ -114,10 +114,13 @@ def test_parse_rejects_out_of_range_and_partial_items(self) -> None: with self.assertRaisesRegex(ValueError, "whole 3-level items"): f.parse({"user_sequence": pa.array([[0, 1]])}) - def test_rejects_a_bad_codebook(self) -> None: + def test_rejects_a_bad_config(self) -> None: for bad, msg in ( ("", "codebook is required"), ("codebook: 4 codebook: 0", "positive"), + # _parse reshapes by level and splits by seq_lengths, so a wider + # value would land the offsets on the wrong components. + ("codebook: 4 value_dim: 2", "value_dim must be 1"), ): with self.subTest(bad=bad): base = 'feature_name: "s" expression: "user:s" ' + bad diff --git a/tzrec/protos/feature.proto b/tzrec/protos/feature.proto index 02bf4f6a0..96ce81d71 100644 --- a/tzrec/protos/feature.proto +++ b/tzrec/protos/feature.proto @@ -993,9 +993,9 @@ message SequenceFeature { // model reads it as a JAGGED_SEQUENCE group, so the scalar form would never // produce the "{group}.sequence" keys it needs. // -// FG note: SID codes are produced offline by a SID-generation model, so this -// feature is only supported with data_config.fg_mode = FG_NONE. There is no -// pyfg counterpart. +// Under fg this is a PASSTHROUGH: no fg feature_type can add +// level_offsets[i % levels], so fg only reaches the codes and the feature folds +// the offsets in at parse time. message SidFeature { // feature name; also the {{name}} placeholder in prompt_template. required string feature_name = 1; From ebcefab6dfa7aa3366a05fb5abb660d4316859e6 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 29 Jul 2026 07:14:10 +0000 Subject: [PATCH 57/99] [chore] bump version to 1.3.9 Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tzrec/version.py b/tzrec/version.py index 6939faa19..ca1bee616 100644 --- a/tzrec/version.py +++ b/tzrec/version.py @@ -9,4 +9,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "1.3.8" +__version__ = "1.3.9" From 9664ecadef0085a6ee6fee83eb751f1a2ff8d696 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 07:21:00 +0000 Subject: [PATCH 58/99] [feat] prompt: add prompt_config and the prompt compiler Adds the config surface for prompt-native generative recommendation and the compiler that turns it into the artifacts the data layer, model and serving read: a resolved SidSpace, a PromptPlan walk order, a ModulePlan projection topology, and an extended tokenizer. The compiler resolves no physical dimension. Slot fill mode is derived, not configured: a lone sequence member declaring no embedding renders INLINE, and everything else is PROJECTED and reaches the LM hidden size through the slot's projection, which the model sizes from group_total_dim at __init__. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/prompt/__init__.py | 10 + tzrec/prompt/compile.py | 456 +++++++++++++++++++++++++++++++++++ tzrec/prompt/compile_test.py | 202 ++++++++++++++++ tzrec/prompt/plan.py | 209 ++++++++++++++++ tzrec/protos/pipeline.proto | 3 + tzrec/protos/prompt.proto | 74 ++++++ 6 files changed, 954 insertions(+) create mode 100644 tzrec/prompt/__init__.py create mode 100644 tzrec/prompt/compile.py create mode 100644 tzrec/prompt/compile_test.py create mode 100644 tzrec/prompt/plan.py create mode 100644 tzrec/protos/prompt.proto diff --git a/tzrec/prompt/__init__.py b/tzrec/prompt/__init__.py new file mode 100644 index 000000000..eedc773bc --- /dev/null +++ b/tzrec/prompt/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tzrec/prompt/compile.py b/tzrec/prompt/compile.py new file mode 100644 index 000000000..450f7ff57 --- /dev/null +++ b/tzrec/prompt/compile.py @@ -0,0 +1,456 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compiles a ``PromptConfig`` into the artifacts the rest of the stack reads. + +Runs once, on the main process, and is the only tokenizer construction in an +entry point. It resolves no physical dimension: the model does that at +``__init__`` from ``group_total_dim``. +""" + +import hashlib +import json +import os +import re +from typing import Any, Dict, List, Optional, Sequence, Tuple + +from tokenizers import Tokenizer + +from tzrec.features.feature import BaseFeature +from tzrec.prompt.plan import ( + CompiledPrompt, + FillMode, + ModulePlan, + PromptPlan, + Segment, + SidSpace, + SlotSeg, + Static, + Width, + WidthKind, +) +from tzrec.protos.model_pb2 import FeatureGroupType +from tzrec.protos.prompt_pb2 import PromptConfig, PromptProjection, PromptSlot +from tzrec.utils.logging_util import logger + +_PLACEHOLDER = re.compile(r"\{\{(\w+)\}\}") + + +def _split_template(template: str) -> Tuple[List[str], List[str]]: + """Split on ``{{name}}`` into static runs and the names between them. + + ``re.split`` with one capture group alternates literal, capture, literal, + so a template with n placeholders yields n + 1 static runs. Both lists are + returned; run i precedes name i. + """ + parts = _PLACEHOLDER.split(template) + return parts[0::2], parts[1::2] + + +def _ceil_to(value: int, multiple: int) -> int: + """Round ``value`` up to a multiple, or return it unchanged when 0.""" + if multiple <= 1: + return value + return -(-value // multiple) * multiple + + +def _resolve_slot(name: str, declared: Dict[str, PromptSlot]) -> PromptSlot: + """Return the declared slot, or an implicit single-feature slot.""" + if name in declared: + return declared[name] + implicit = PromptSlot(name=name) + implicit.feature_names.append(name) + return implicit + + +def _slot_width( + members: Sequence[BaseFeature], group_type: "FeatureGroupType.ValueType" +) -> Width: + """Derive a slot's position count from its members. + + A DEEP slot pools to exactly one position. A sequence slot is bounded by + the members' ``sequence_length``, and unbounded when none declares one. + """ + if group_type == FeatureGroupType.DEEP: + return Width(WidthKind.STATIC, 1) + caps = [ + int(f.config.sequence_length) + for f in members + if f.config.HasField("sequence_length") + ] + if not caps: + return Width(WidthKind.UNBOUNDED) + return Width(WidthKind.BOUNDED, max(caps)) + + +def _derive_fill(members: Sequence[BaseFeature]) -> FillMode: + """INLINE only for a lone sequence member that declares no embedding.""" + if len(members) == 1 and members[0].is_sequence and not members[0].has_embedding: + return FillMode.INLINE + return FillMode.PROJECTED + + +def _group_type( + name: str, members: Sequence[BaseFeature] +) -> "FeatureGroupType.ValueType": + """JAGGED_SEQUENCE for sequence members, DEEP for scalars; never mixed.""" + kinds = {f.is_sequence for f in members} + if len(kinds) != 1: + raise ValueError( + f"prompt slot [{name}] mixes sequence and scalar features " + f"{[f.name for f in members]}; a slot must be all one kind, or its " + f"group would carry a '.query' output the prompt cannot place." + ) + return FeatureGroupType.JAGGED_SEQUENCE if kinds.pop() else FeatureGroupType.DEEP + + +def _atom_tokens(sid_space: Any) -> List[str]: + """Render the SID atom tokens, one per flat index.""" + fmt = sid_space.atom_token_format + return [fmt.replace("{i}", str(i)) for i in range(sum(sid_space.codebook))] + + +def _read_manifest_codebook(path: str) -> Optional[List[int]]: + """Read ``codebook`` from a SID manifest, or None when there is no file.""" + if not os.path.exists(path): + raise ValueError(f"sid_space.manifest_path [{path}] does not exist.") + with open(path, "r") as f: + return [int(c) for c in json.load(f)["codebook"]] + + +def _build_sid_space( + cfg: PromptConfig, tok: Tokenizer, base_vocab: int, has_projection: bool +) -> Optional[SidSpace]: + """Extend the tokenizer with SID atoms and resolve the token space.""" + if not cfg.HasField("sid_space"): + return None + space = cfg.sid_space + codebook = [int(c) for c in space.codebook] + if not codebook: + raise ValueError("sid_space.codebook is required; one size per SID level.") + if any(c <= 0 for c in codebook): + raise ValueError(f"every codebook size must be positive, got {codebook}.") + + if space.HasField("manifest_path"): + declared = _read_manifest_codebook(space.manifest_path) + if declared != codebook: + raise ValueError( + f"sid_space.codebook {codebook} does not match the manifest at " + f"[{space.manifest_path}] which describes {declared}. The data " + f"and the decode bands would disagree." + ) + + atoms = _atom_tokens(space) + present = [a for a in atoms if tok.token_to_id(a) is not None] + if present: + raise ValueError( + f"SID atoms are already in the base tokenizer, e.g. {present[:3]}; " + f"change sid_space.atom_token_format." + ) + tok.add_special_tokens(atoms) + + sentinel_id = None + if has_projection: + if tok.token_to_id(cfg.sentinel_token) is not None: + raise ValueError( + f"sentinel_token [{cfg.sentinel_token}] is already in the base " + f"tokenizer; a projected position would be indistinguishable " + f"from real content." + ) + tok.add_special_tokens([cfg.sentinel_token]) + sentinel_id = tok.token_to_id(cfg.sentinel_token) + + offsets: List[int] = [] + running = 0 + for size in codebook: + offsets.append(running) + running += size + lo = [base_vocab + o for o in offsets] + hi = [lo[i] + codebook[i] - 1 for i in range(len(codebook))] + + return SidSpace( + codebook=tuple(codebook), + num_levels=len(codebook), + base_vocab=base_vocab, + level_offsets=tuple(offsets), + band_lo=tuple(lo), + band_hi=tuple(hi), + target_vocab=_ceil_to( + tok.get_vocab_size(with_added_tokens=True), + space.vocab_pad_to_multiple_of, + ), + sentinel_token_id=sentinel_id, + eos_token_id=_special_id(tok, ("<|im_end|>", "<|endoftext|>")), + pad_token_id=_special_id(tok, ("<|endoftext|>", "<|im_end|>")), + ) + + +def _special_id(tok: Tokenizer, candidates: Sequence[str]) -> int: + """First candidate the tokenizer knows, so a family swap does not break.""" + for name in candidates: + token_id = tok.token_to_id(name) + if token_id is not None: + return int(token_id) + raise ValueError( + f"none of {list(candidates)} is in the tokenizer; the prompt cannot " + f"resolve its EOS/pad ids." + ) + + +def _hash(*parts: Any) -> str: + """Stable sha256 over the given parts.""" + digest = hashlib.sha256() + for part in parts: + digest.update(repr(part).encode("utf-8")) + return digest.hexdigest() + + +def compile_prompt( + cfg: PromptConfig, + features: Sequence[BaseFeature], + model_dir: Optional[str] = None, +) -> CompiledPrompt: + """Compile a prompt config into its plan, module and vocabulary artifacts. + + Args: + cfg: the prompt config to compile. + features: every feature the config may reference, already created. + model_dir: where to write the extended tokenizer; skipped when None. + + Returns: + The compiled prompt. + """ + by_name = {f.name: f for f in features} + declared = {s.name: s for s in cfg.slots} + + body_runs, body_names = _split_template(cfg.prompt) + resp_runs, resp_names = _split_template(cfg.response or "") + slots = {n: _resolve_slot(n, declared) for n in body_names + resp_names} + + unreferenced = set(declared) - set(slots) + if unreferenced: + raise ValueError( + f"declared prompt slots {sorted(unreferenced)} are never referenced " + f"by a {{{{name}}}} placeholder." + ) + + members: Dict[str, List[BaseFeature]] = {} + for name, slot in slots.items(): + missing = [f for f in slot.feature_names if f not in by_name] + if missing: + raise ValueError( + f"prompt slot [{name}] names features {missing} that are not in " + f"feature_configs." + ) + members[name] = [by_name[f] for f in slot.feature_names] + + types = {n: _group_type(n, members[n]) for n in slots} + fills = {n: _derive_fill(members[n]) for n in slots} + has_projection = any(f is FillMode.PROJECTED for f in fills.values()) + + for name, slot in slots.items(): + if fills[name] is FillMode.INLINE and slot.HasField("projection"): + raise ValueError( + f"prompt slot [{name}] is INLINE -- one sequence feature with no " + f"embedding -- so it has no group to project; drop its projection." + ) + + tok = Tokenizer.from_file(cfg.tokenizer) + base_vocab = tok.get_vocab_size(with_added_tokens=True) + sid_space = _build_sid_space(cfg, tok, base_vocab, has_projection) + + tokenizer_dir = "" + if model_dir: + tokenizer_dir = os.path.join(model_dir, "prompt", "tokenizer") + os.makedirs(tokenizer_dir, exist_ok=True) + tok.save(os.path.join(tokenizer_dir, "tokenizer.json")) + + slot_ids = {n: i for i, n in enumerate(slots)} + segs: Dict[str, SlotSeg] = {} + for name, slot in slots.items(): + seq = types[name] == FeatureGroupType.JAGGED_SEQUENCE + segs[name] = SlotSeg( + slot_id=slot_ids[name], + name=name, + sources=tuple(slot.feature_names), + group_type=types[name], + output_key=".sequence" if seq else "", + fill=fills[name], + width=_slot_width(members[name], types[name]), + droppable=bool(slot.drop_if_empty), + ) + + body = _weave(body_runs, body_names, segs, tok) + response = _weave(resp_runs, resp_names, segs, tok) + + projected = tuple( + s + for s in body + response + if isinstance(s, SlotSeg) and s.fill is FillMode.PROJECTED + ) + module_plan = _build_module_plan(projected, slots) + + plan = PromptPlan( + segments=body, + response_segments=response, + max_length=int(cfg.max_length), + max_total_length=_max_total_length(body + response), + max_holes=_max_holes(projected), + suffix_keep=_suffix_keep(response), + static_prefix_len=_static_prefix_len(body), + length_buckets=tuple(int(b) for b in cfg.length_buckets), + slot_index={s.name: i for i, s in enumerate(projected)}, + projected_slots=projected, + ) + _validate(cfg, plan, sid_space) + + return CompiledPrompt( + sid_space=sid_space, + prompt_plan=plan, + module_plan=module_plan, + tokenizer_dir=tokenizer_dir, + vocab_hash=_hash(sid_space, tok.to_str()), + plan_hash=_hash(sid_space, plan, sorted(module_plan.projections), tok.to_str()), + ) + + +def _weave( + runs: Sequence[str], + names: Sequence[str], + segs: Dict[str, SlotSeg], + tok: Tokenizer, +) -> Tuple[Segment, ...]: + """Interleave tokenized static runs with their slots, dropping empty runs.""" + out: List[Segment] = [] + for i, run in enumerate(runs): + if run: + ids = tuple(tok.encode(run, add_special_tokens=False).ids) + out.append(Static(token_ids=ids, owner_slot_id=None)) + if i < len(names): + out.append(segs[names[i]]) + return tuple(out) + + +def _build_module_plan( + projected: Sequence[SlotSeg], slots: Dict[str, PromptSlot] +) -> ModulePlan: + """One module per distinct ``projection_name``, else one per slot.""" + projections: Dict[str, PromptProjection] = {} + slot_to_module: Dict[int, str] = {} + for seg in projected: + slot = slots[seg.name] + module_id = slot.projection_name or seg.name + projection = ( + slot.projection if slot.HasField("projection") else PromptProjection() + ) + if module_id in projections: + if projections[module_id] != projection: + raise ValueError( + f"prompt slots sharing projection_name [{module_id}] declare " + f"different projection bodies; they cannot share weights." + ) + else: + projections[module_id] = projection + slot_to_module[seg.slot_id] = module_id + return ModulePlan(projections=projections, slot_to_module=slot_to_module) + + +def _max_total_length(segments: Sequence[Segment]) -> Optional[int]: + """Provable position ceiling, or None when any slot is unbounded.""" + total = 0 + for seg in segments: + if isinstance(seg, Static): + total += len(seg.token_ids) + elif seg.width.kind is WidthKind.UNBOUNDED: + return None + else: + assert seg.width.n is not None + total += seg.width.n + return total + + +def _max_holes(projected: Sequence[SlotSeg]) -> int: + """Per-row projected-position ceiling; unbounded slots cannot be counted.""" + total = 0 + for seg in projected: + if seg.width.kind is WidthKind.UNBOUNDED: + raise ValueError( + f"prompt slot [{seg.name}] is PROJECTED and unbounded, so its " + f"hole count is unknowable; give its members a sequence_length." + ) + assert seg.width.n is not None + total += seg.width.n + return total + + +def _suffix_keep(response: Sequence[Segment]) -> Optional[int]: + """Upper bound on the supervised window, or None when unbounded.""" + total = _max_total_length(response) + # HF shifts logits by one, so the window opens one column before the first + # supervised label. + return None if total is None else total + 1 + + +def _static_prefix_len(segments: Sequence[Segment]) -> int: + """Leading positions that are request-invariant: static runs only.""" + total = 0 + for seg in segments: + if not isinstance(seg, Static): + break + total += len(seg.token_ids) + return total + + +def _validate( + cfg: PromptConfig, plan: PromptPlan, sid_space: Optional[SidSpace] +) -> None: + """Apply the checks that need the whole plan.""" + if plan.max_length and plan.max_total_length is not None: + if plan.max_total_length > plan.max_length: + raise ValueError( + f"the prompt can reach {plan.max_total_length} positions but " + f"max_length is {plan.max_length}; a row could never be assembled." + ) + if not plan.max_length and plan.max_total_length is None: + logger.warning( + "prompt has an unbounded slot and max_length is 0; graph-captured " + "serving cannot size its buckets." + ) + if any(isinstance(s, SlotSeg) for s in plan.segments): + first_slot = next( + i for i, s in enumerate(plan.segments) if isinstance(s, SlotSeg) + ) + later_static = any( + isinstance(s, SlotSeg) and s.width.kind is WidthKind.STATIC + for s in plan.segments[first_slot + 1 :] + ) + if later_static: + logger.warning( + "a variable-width prompt slot precedes a fixed-width one; " + f"static_prefix_len is {plan.static_prefix_len}, which bounds " + "what a serving prefix cache may reuse." + ) + if plan.static_prefix_len == 0: + logger.warning( + "static_prefix_len is 0: no leading run of the prompt is " + "request-invariant, so a serving prefix cache can share nothing." + ) + if sid_space is not None and cfg.HasField("response"): + answer = [s for s in plan.response_segments if isinstance(s, SlotSeg)] + for seg in answer: + if ( + seg.width.kind is WidthKind.STATIC + and seg.width.n != sid_space.num_levels + ): + raise ValueError( + f"response slot [{seg.name}] is {seg.width.n} positions but " + f"the codebook has {sid_space.num_levels} levels." + ) diff --git a/tzrec/prompt/compile_test.py b/tzrec/prompt/compile_test.py new file mode 100644 index 000000000..d5f1e2806 --- /dev/null +++ b/tzrec/prompt/compile_test.py @@ -0,0 +1,202 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +import unittest + +from google.protobuf import text_format +from tokenizers import Tokenizer, models, pre_tokenizers + +from tzrec.features.feature import FgMode, create_features +from tzrec.prompt.compile import compile_prompt +from tzrec.prompt.plan import FillMode, SlotSeg, Static, WidthKind +from tzrec.protos import feature_pb2 +from tzrec.protos.prompt_pb2 import PromptConfig +from tzrec.utils.test_util import make_test_dir + +_WORDS = ["History", "Profile", "Predict", ":", ".", "", "<|im_end|>"] + + +def _tokenizer(path: str) -> str: + """Write a minimal word-level tokenizer, so no download is needed.""" + vocab = {w: i for i, w in enumerate(_WORDS)} + tok = Tokenizer(models.WordLevel(vocab=vocab, unk_token="")) + tok.pre_tokenizer = pre_tokenizers.Whitespace() + tok.save(path) + return path + + +def _feature(text: str): + fc = feature_pb2.FeatureConfig() + text_format.Merge(text, fc) + return create_features([fc], fg_mode=FgMode.FG_NONE)[0] + + +_HIST = 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }' +_PROF = ( + 'sequence_id_feature { feature_name: "prof" expression: "user:prof" ' + "num_buckets: 768 embedding_dim: 16 sequence_length: 4 }" +) +_AGE = 'id_feature { feature_name: "age" expression: "user:age" num_buckets: 8 }' + + +class CompilePromptTest(unittest.TestCase): + def setUp(self) -> None: + self.test_dir = make_test_dir() + self.tok_path = _tokenizer(os.path.join(self.test_dir, "tok.json")) + + def _config(self, **kwargs) -> PromptConfig: + cfg = PromptConfig(tokenizer=self.tok_path, **kwargs) + return cfg + + def _compile(self, cfg, features): + return compile_prompt(cfg, features, model_dir=self.test_dir) + + def test_sid_space_resolves_offsets_and_bands(self) -> None: + cfg = self._config(prompt="History : {{hist}}") + cfg.sid_space.codebook.extend([4, 4, 4]) + compiled = self._compile(cfg, [_feature(_HIST)]) + space = compiled.sid_space + + base = space.base_vocab + self.assertEqual(space.num_levels, 3) + self.assertEqual(space.sid_vocab_size, 12) + self.assertEqual(space.level_offsets, (0, 4, 8)) + self.assertEqual(space.band_lo, (base, base + 4, base + 8)) + self.assertEqual(space.band_hi, (base + 3, base + 7, base + 11)) + # no slot projects, so no sentinel is materialized + self.assertIsNone(space.sentinel_token_id) + self.assertEqual(space.target_vocab % 128, 0) + + def test_inline_needs_no_group_projected_gets_one(self) -> None: + cfg = self._config(prompt="History : {{hist}} . Profile : {{prof}}") + cfg.sid_space.codebook.extend([4, 4, 4]) + compiled = self._compile(cfg, [_feature(_HIST), _feature(_PROF)]) + + by_name = { + s.name: s for s in compiled.prompt_plan.segments if isinstance(s, SlotSeg) + } + self.assertIs(by_name["hist"].fill, FillMode.INLINE) + self.assertIs(by_name["prof"].fill, FillMode.PROJECTED) + # only the projected slot produces a group, and so a hole + self.assertEqual( + [s.name for s in compiled.prompt_plan.projected_slots], ["prof"] + ) + self.assertEqual(compiled.prompt_plan.max_holes, 4) + self.assertIsNotNone(compiled.sid_space.sentinel_token_id) + + def test_static_runs_are_woven_between_slots(self) -> None: + cfg = self._config(prompt="History : {{hist}} . Predict :") + cfg.sid_space.codebook.extend([4]) + compiled = self._compile(cfg, [_feature(_HIST)]) + kinds = [ + "static" if isinstance(s, Static) else s.name + for s in compiled.prompt_plan.segments + ] + self.assertEqual(kinds, ["static", "hist", "static"]) + # the leading run is request-invariant; "History :" is two tokens + self.assertEqual(compiled.prompt_plan.static_prefix_len, 2) + + def test_scalar_slot_is_one_deep_position(self) -> None: + cfg = self._config(prompt="Profile : {{age}}") + cfg.sid_space.codebook.extend([4]) + compiled = self._compile(cfg, [_feature(_AGE)]) + seg = next(s for s in compiled.prompt_plan.segments if isinstance(s, SlotSeg)) + self.assertIs(seg.fill, FillMode.PROJECTED) + self.assertEqual(seg.output_key, "") + self.assertIs(seg.width.kind, WidthKind.STATIC) + self.assertEqual(seg.width.n, 1) + + def test_manifest_mismatch_is_fatal(self) -> None: + manifest = os.path.join(self.test_dir, "manifest.json") + with open(manifest, "w") as f: + json.dump({"codebook": [8, 8, 8]}, f) + cfg = self._config(prompt="History : {{hist}}") + cfg.sid_space.codebook.extend([4, 4, 4]) + cfg.sid_space.manifest_path = manifest + with self.assertRaisesRegex(ValueError, "does not match the manifest"): + self._compile(cfg, [_feature(_HIST)]) + + def test_manifest_match_compiles(self) -> None: + manifest = os.path.join(self.test_dir, "manifest.json") + with open(manifest, "w") as f: + json.dump({"codebook": [4, 4, 4]}, f) + cfg = self._config(prompt="History : {{hist}}") + cfg.sid_space.codebook.extend([4, 4, 4]) + cfg.sid_space.manifest_path = manifest + self.assertEqual(self._compile(cfg, [_feature(_HIST)]).sid_space.num_levels, 3) + + def test_rejects_a_mixed_kind_slot(self) -> None: + cfg = self._config(prompt="X : {{both}}") + cfg.sid_space.codebook.extend([4]) + slot = cfg.slots.add(name="both") + slot.feature_names.extend(["hist", "age"]) + with self.assertRaisesRegex(ValueError, "mixes sequence and scalar"): + self._compile(cfg, [_feature(_HIST), _feature(_AGE)]) + + def test_rejects_unknown_feature_and_unreferenced_slot(self) -> None: + cfg = self._config(prompt="X : {{hist}}") + cfg.sid_space.codebook.extend([4]) + slot = cfg.slots.add(name="hist") + slot.feature_names.append("nope") + with self.assertRaisesRegex(ValueError, "not in\n?\\s*feature_configs"): + self._compile(cfg, [_feature(_HIST)]) + + cfg2 = self._config(prompt="X : {{hist}}") + cfg2.sid_space.codebook.extend([4]) + cfg2.slots.add(name="ghost").feature_names.append("hist") + with self.assertRaisesRegex(ValueError, "never referenced"): + self._compile(cfg2, [_feature(_HIST)]) + + def test_rejects_a_projection_on_an_inline_slot(self) -> None: + cfg = self._config(prompt="X : {{hist}}") + cfg.sid_space.codebook.extend([4]) + slot = cfg.slots.add(name="hist") + slot.feature_names.append("hist") + slot.projection.bias = True + with self.assertRaisesRegex(ValueError, "is INLINE"): + self._compile(cfg, [_feature(_HIST)]) + + def test_atoms_absent_from_the_base_tokenizer(self) -> None: + cfg = self._config(prompt="X : {{hist}}") + cfg.sid_space.codebook.extend([4]) + cfg.sid_space.atom_token_format = "History" + with self.assertRaisesRegex(ValueError, "already in the base tokenizer"): + self._compile(cfg, [_feature(_HIST)]) + + def test_vocab_hash_tracks_the_codebook(self) -> None: + def compile_with(sizes): + cfg = self._config(prompt="History : {{hist}}") + cfg.sid_space.codebook.extend(sizes) + return self._compile(cfg, [_feature(_HIST)]) + + self.assertEqual( + compile_with([4, 4]).vocab_hash, compile_with([4, 4]).vocab_hash + ) + self.assertNotEqual( + compile_with([4, 4]).vocab_hash, compile_with([4, 8]).vocab_hash + ) + + def test_extended_tokenizer_is_written(self) -> None: + cfg = self._config(prompt="History : {{hist}}") + cfg.sid_space.codebook.extend([4, 4]) + compiled = self._compile(cfg, [_feature(_HIST)]) + written = os.path.join(compiled.tokenizer_dir, "tokenizer.json") + self.assertTrue(os.path.exists(written)) + # the atoms round-trip, which is what serving reloads + reloaded = Tokenizer.from_file(written) + self.assertIsNotNone(reloaded.token_to_id("<|sid_0|>")) + self.assertIsNotNone(reloaded.token_to_id("<|sid_7|>")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/prompt/plan.py b/tzrec/prompt/plan.py new file mode 100644 index 000000000..9fae4c5d6 --- /dev/null +++ b/tzrec/prompt/plan.py @@ -0,0 +1,209 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Products of ``compile_prompt``. + +This namespace disambiguates ``plan.SidSpace``, the resolved token space, from +``prompt_pb2.SidSpace``, the four knobs a user declares. Nothing here stores a +physical dimension: the model resolves those at ``__init__``. +""" + +from dataclasses import dataclass +from enum import Enum +from typing import Mapping, Optional, Tuple, Union + +from tzrec.protos.model_pb2 import FeatureGroupType +from tzrec.protos.prompt_pb2 import PromptProjection + + +class FillMode(Enum): + """How a slot's value reaches the LM input space.""" + + INLINE = "inline" + PROJECTED = "projected" + + +class WidthKind(Enum): + """Whether a segment's position count is known, bounded, or neither.""" + + STATIC = "static" + BOUNDED = "bounded" + UNBOUNDED = "unbounded" + + +@dataclass(frozen=True) +class Width: + """Position count of a slot. + + Args: + kind: STATIC when the count is exact, BOUNDED when only a ceiling is + known, UNBOUNDED when neither. + n: the exact count or the ceiling; None when UNBOUNDED. + """ + + kind: WidthKind + n: Optional[int] = None + + def __post_init__(self) -> None: + """Reject a count that contradicts the kind.""" + if self.kind is WidthKind.UNBOUNDED: + if self.n is not None: + raise ValueError("UNBOUNDED width cannot carry a count.") + elif self.n is None or self.n < 0: + raise ValueError( + f"{self.kind.name} width needs a count >= 0, got {self.n}." + ) + + +@dataclass(frozen=True) +class SidSpace: + """The resolved SID token space, read by the data layer, model and serving. + + Three coordinate systems and the constants that convert between them: a + local code in ``[0, codebook[l])``, a flat index ``level_offsets[l] + code`` + which is what the data carries, and an LM token id ``base_vocab + flat`` + which is what ``lm_head`` generates. + + Args: + codebook: per-level vocabulary sizes. + num_levels: codes per item; also the answer width. + base_vocab: tokenizer size before the SID atoms were appended. + level_offsets: ``cumsum(codebook) - codebook``. + band_lo: inclusive lower token-id bound of each level. + band_hi: inclusive upper token-id bound of each level. + target_vocab: embedding rows after padding, what the LM resizes to. + sentinel_token_id: id reserved for projected positions, None when no + slot is projected. + eos_token_id: end-of-sequence id of the extended tokenizer. + pad_token_id: padding id of the extended tokenizer. + """ + + codebook: Tuple[int, ...] + num_levels: int + base_vocab: int + level_offsets: Tuple[int, ...] + band_lo: Tuple[int, ...] + band_hi: Tuple[int, ...] + target_vocab: int + sentinel_token_id: Optional[int] + eos_token_id: int + pad_token_id: int + + @property + def sid_vocab_size(self) -> int: + """Atoms appended to the backbone vocabulary.""" + return sum(self.codebook) + + +@dataclass(frozen=True) +class Static: + """A run of literal template tokens. + + Args: + token_ids: the tokenized run. + owner_slot_id: slot this run was folded into, so it vanishes when that + slot is dropped. None when the run belongs to no slot. + """ + + token_ids: Tuple[int, ...] + owner_slot_id: Optional[int] + + +@dataclass(frozen=True) +class SlotSeg: + """One ``{{name}}`` position in the assembled stream. + + Args: + slot_id: index into ``PromptPlan.projected_slots`` ordering. + name: the placeholder name; also the derived group name. + sources: member feature names. + group_type: DEEP or JAGGED_SEQUENCE. + output_key: "" for DEEP, ".sequence" otherwise. + fill: INLINE writes token ids, PROJECTED writes sentinels and a hole. + width: position count of this slot. + droppable: whether an empty value removes the slot and its folded text. + """ + + slot_id: int + name: str + sources: Tuple[str, ...] + group_type: "FeatureGroupType.ValueType" + output_key: str + fill: FillMode + width: Width + droppable: bool + + +Segment = Union[Static, SlotSeg] + + +@dataclass(frozen=True) +class PromptPlan: + """The walk order the assembler follows, plus the ceilings derived from it. + + Args: + segments: prompt body, in emission order. + response_segments: supervised tail, in emission order. + max_length: validation ceiling; an over-long row is an error. + max_total_length: proven ceiling when every slot is bounded, else None. + max_holes: per-row projected-position ceiling, not a runtime shape. + suffix_keep: upper bound on the supervised logits window. + static_prefix_len: leading positions that are request-invariant. + length_buckets: sampler and graph-capture buckets. + slot_index: slot name to its index in ``projected_slots``. + projected_slots: fixes the order hole positions are written in. + """ + + segments: Tuple[Segment, ...] + response_segments: Tuple[Segment, ...] + max_length: int + max_total_length: Optional[int] + max_holes: int + suffix_keep: Optional[int] + static_prefix_len: int + length_buckets: Tuple[int, ...] + slot_index: Mapping[str, int] + projected_slots: Tuple[SlotSeg, ...] + + +@dataclass(frozen=True) +class ModulePlan: + """Projection topology. Model-only, never persisted. + + Args: + projections: resolved module id to its configuration. + slot_to_module: slot id to the module id it uses, so slots sharing a + ``projection_name`` resolve to one module. + """ + + projections: Mapping[str, PromptProjection] + slot_to_module: Mapping[int, str] + + +@dataclass(frozen=True) +class CompiledPrompt: + """Everything ``compile_prompt`` produces. + + Args: + sid_space: the resolved SID token space. + prompt_plan: assembler walk order and ceilings. + module_plan: projection topology. + tokenizer_dir: where the extended tokenizer was written. + vocab_hash: over sid_space and tokenizer.json; fatal on mismatch. + plan_hash: over all four parts; warns on mismatch. + """ + + sid_space: Optional[SidSpace] + prompt_plan: PromptPlan + module_plan: ModulePlan + tokenizer_dir: str + vocab_hash: str + plan_hash: str diff --git a/tzrec/protos/pipeline.proto b/tzrec/protos/pipeline.proto index cd86dfb38..c649f4a13 100644 --- a/tzrec/protos/pipeline.proto +++ b/tzrec/protos/pipeline.proto @@ -7,6 +7,7 @@ import "tzrec/protos/export.proto"; import "tzrec/protos/data.proto"; import "tzrec/protos/feature.proto"; import "tzrec/protos/model.proto"; +import "tzrec/protos/prompt.proto"; message EasyRecConfig { required string train_input_path = 1; @@ -26,4 +27,6 @@ message EasyRecConfig { repeated FeatureConfig feature_configs = 8; optional ModelConfig model_config = 9; + + optional PromptConfig prompt_config = 10; } diff --git a/tzrec/protos/prompt.proto b/tzrec/protos/prompt.proto new file mode 100644 index 000000000..c35b1d268 --- /dev/null +++ b/tzrec/protos/prompt.proto @@ -0,0 +1,74 @@ +syntax = "proto2"; +package tzrec.protos; + +import "tzrec/protos/module.proto"; + +// Rendering of an LM prompt. Peer of data_config and model_config: extraction +// stays in feature_configs, this owns order, literal text and composition. +message PromptConfig { + // BASE tokenizer, path or hub id. Distinct from hf_model_id, which names + // the WEIGHTS and is read only at cold start. + required string tokenizer = 1; + // Rewritten by export to the content-addressed asset directory. + optional string asset_dir = 2; + + // Static text between {{name}} placeholders is the prefix and suffix of + // the surrounding slots; there are no per-slot text fields. + required string prompt = 3; + // Supervised answer. Defines the loss span; absent at inference. + optional string response = 4; + + // A placeholder resolves to a slot with that name, else to an implicit + // single-feature slot named after it. + repeated PromptSlot slots = 5; + + // Required when any slot renders SIDs. + optional SidSpace sid_space = 6; + + // Validation ceiling, not a truncation trigger: an over-long row is an + // error, never truncated. + optional uint32 max_length = 7 [default = 0]; + + // Reserves a position filled by a projected slot. Materialized only when + // at least one slot has a projection. + optional string sentinel_token = 8 [default = "<|pg_hole|>"]; + + // Length buckets for the training sampler and graph capture. + repeated uint32 length_buckets = 9; +} + +message PromptSlot { + // {{name}} in the template; also the name of the derived feature_group. + required string name = 1; + // Features rendered at this position. Several are concatenated per + // position; the compiler derives the FeatureGroupConfig from this list. + repeated string feature_names = 2; + // Omit the slot and the static text folded into it when the value is empty. + optional bool drop_if_empty = 3 [default = false]; + // Reconciles this slot's width with the LM hidden size. Illegal on an + // INLINE slot. Carries no dimensions: the model resolves them. + optional PromptProjection projection = 4; + // Weight-sharing key across slots. Derived groups dedupe automatically; + // modules share only on request. + optional string projection_name = 5; +} + +// An optional body plus a final bare Linear to the LM hidden size. The final +// map never carries an activation, and no field here states a dimension. +message PromptProjection { + optional bool bias = 1 [default = true]; + oneof body { + MLP mlp = 10; + } +} + +message SidSpace { + // Per-level sizes; codes are local 0-based in [0, codebook[l]). + repeated uint32 codebook = 1; + // {i} is the flat atom index. + optional string atom_token_format = 2 [default = "<|sid_{i}|>"]; + optional uint32 vocab_pad_to_multiple_of = 3 [default = 128]; + // SID manifest to cross-check codebook against. A mismatch is fatal: it + // means the data and the decode bands disagree. + optional string manifest_path = 4; +} From 434f55538ec6400dba033e617871d679d044a7ee Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 07:24:22 +0000 Subject: [PATCH 59/99] [feat] prompt: add the varlen prompt assembler Walks a compiled PromptPlan to build the packed token stream: static runs, the base_vocab shift on INLINE SID slots, sentinels plus recorded hole positions on PROJECTED ones, and labels that cover the response span only. Band validation lives here rather than in a feature because the assembler owns the codebook, and running it in the worker fails the offending sample instead of letting one rank raise and hang its peers on the collective. Its error names offset_codebook, since a raw or origin column is well formed and would otherwise train silently on pre-relocation SIDs. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/prompt/assembler.py | 181 +++++++++++++++++++++++++++++++++ tzrec/prompt/assembler_test.py | 179 ++++++++++++++++++++++++++++++++ 2 files changed, 360 insertions(+) create mode 100644 tzrec/prompt/assembler.py create mode 100644 tzrec/prompt/assembler_test.py diff --git a/tzrec/prompt/assembler.py b/tzrec/prompt/assembler.py new file mode 100644 index 000000000..231cac0ee --- /dev/null +++ b/tzrec/prompt/assembler.py @@ -0,0 +1,181 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Builds the packed token stream a compiled prompt describes. + +Runs in the dataloader worker, after the features are parsed and outside any +feature's ``_parse``. Pure integer arithmetic with no FG dependency, so the +same walk is portable to an online C++/Java processor. +""" + +from dataclasses import dataclass +from typing import Dict, List, Optional, Sequence + +import numpy as np + +from tzrec.prompt.plan import FillMode, PromptPlan, SidSpace, SlotSeg, Static + + +@dataclass +class AssembledPrompt: + """One batch of assembled prompts, in packed varlen form. + + Args: + input_ids: every row's tokens concatenated, ``(total_tokens,)``. + cu_seqlens: row boundaries into ``input_ids``, ``(batch_size + 1,)``. + hole_positions: absolute indices the projected embeddings overwrite, + in ``PromptPlan.projected_slots`` order within each row. + labels: ``ignore_index`` outside the response span. + """ + + input_ids: np.ndarray + cu_seqlens: np.ndarray + hole_positions: np.ndarray + labels: np.ndarray + + +class PromptAssembler: + """Walks a ``PromptPlan`` to build token streams. + + Args: + plan: the compiled walk order. + sid_space: resolved SID token space; required when a slot renders SIDs. + ignore_index: label value outside the supervised span. + """ + + def __init__( + self, + plan: PromptPlan, + sid_space: Optional[SidSpace] = None, + ignore_index: int = -100, + ) -> None: + self._plan = plan + self._sid = sid_space + self._ignore_index = ignore_index + inline = [ + s + for s in plan.segments + plan.response_segments + if isinstance(s, SlotSeg) and s.fill is FillMode.INLINE + ] + if inline and sid_space is None: + raise ValueError( + f"prompt slots {[s.name for s in inline]} render INLINE, which " + f"means SID codes, but no sid_space was compiled." + ) + + def _inline_tokens(self, name: str, values: np.ndarray) -> np.ndarray: + """Validate offset SID codes against their bands and shift to token ids. + + The data carries ``level_offsets[l] + code``; the LM vocabulary needs + one further uniform shift by ``base_vocab``. + """ + assert self._sid is not None + levels = self._sid.num_levels + if values.size % levels: + raise ValueError( + f"prompt slot [{name}]: {values.size} values is not a whole " + f"number of {levels}-level items." + ) + by_level = values.reshape(-1, levels) + lo = np.asarray(self._sid.level_offsets, dtype=np.int64) + hi = lo + np.asarray(self._sid.codebook, dtype=np.int64) + if np.any(by_level < lo) or np.any(by_level >= hi): + raise ValueError( + f"prompt slot [{name}]: SID values must already carry their " + f"level offset, so level l lies in " + f"[level_offsets[l], level_offsets[l] + codebook[l]). Read the " + f"offset_codebook column, not codebook or origin_codebook." + ) + return values.astype(np.int64, copy=False) + self._sid.base_vocab + + def _emit_row( + self, + segments: Sequence[object], + row: int, + values: Dict[str, List[np.ndarray]], + counts: Dict[str, np.ndarray], + out: List[int], + holes: List[int], + base: int, + ) -> None: + """Append one row's tokens for one segment list, recording holes.""" + for seg in segments: + if isinstance(seg, Static): + out.extend(seg.token_ids) + continue + assert isinstance(seg, SlotSeg) + if seg.fill is FillMode.INLINE: + out.extend(self._inline_tokens(seg.name, values[seg.name][row])) + else: + assert self._sid is not None + width = int(counts[seg.name][row]) + holes.extend(range(base + len(out), base + len(out) + width)) + out.extend([self._sid.sentinel_token_id] * width) + + def assemble( + self, + values: Dict[str, List[np.ndarray]], + counts: Optional[Dict[str, np.ndarray]] = None, + batch_size: Optional[int] = None, + ) -> AssembledPrompt: + """Assemble one batch. + + Args: + values: INLINE slot name to its per-row value arrays. + counts: PROJECTED slot name to its per-row position count. + batch_size: row count; inferred from ``values`` when omitted. + + Returns: + The packed streams. + """ + counts = counts or {} + if batch_size is None: + if not values: + raise ValueError("batch_size is required when no INLINE slot exists.") + batch_size = len(next(iter(values.values()))) + + ids: List[int] = [] + labels: List[int] = [] + holes: List[int] = [] + cu = [0] + for row in range(batch_size): + row_ids: List[int] = [] + self._emit_row( + self._plan.segments, row, values, counts, row_ids, holes, len(ids) + ) + prompt_len = len(row_ids) + self._emit_row( + self._plan.response_segments, + row, + values, + counts, + row_ids, + holes, + len(ids), + ) + # supervision covers the response span only; the prompt is context. + row_labels = [self._ignore_index] * prompt_len + row_ids[prompt_len:] + if self._plan.max_length and len(row_ids) > self._plan.max_length: + raise ValueError( + f"assembled row {row} is {len(row_ids)} tokens, over " + f"max_length {self._plan.max_length}. Rows are never " + f"truncated: cap the source features instead." + ) + ids.extend(row_ids) + labels.extend(row_labels) + cu.append(len(ids)) + + return AssembledPrompt( + input_ids=np.asarray(ids, dtype=np.int64), + cu_seqlens=np.asarray(cu, dtype=np.int64), + hole_positions=np.asarray(holes, dtype=np.int64), + labels=np.asarray(labels, dtype=np.int64), + ) diff --git a/tzrec/prompt/assembler_test.py b/tzrec/prompt/assembler_test.py new file mode 100644 index 000000000..660236a33 --- /dev/null +++ b/tzrec/prompt/assembler_test.py @@ -0,0 +1,179 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np + +from tzrec.prompt.assembler import PromptAssembler +from tzrec.prompt.plan import ( + FillMode, + PromptPlan, + SidSpace, + SlotSeg, + Static, + Width, + WidthKind, +) +from tzrec.protos.model_pb2 import FeatureGroupType + +_BASE = 1000 +_SENTINEL = 1099 + + +def _sid_space(codebook=(4, 4, 4)) -> SidSpace: + offsets, running = [], 0 + for size in codebook: + offsets.append(running) + running += size + return SidSpace( + codebook=tuple(codebook), + num_levels=len(codebook), + base_vocab=_BASE, + level_offsets=tuple(offsets), + band_lo=tuple(_BASE + o for o in offsets), + band_hi=tuple(_BASE + o + s - 1 for o, s in zip(offsets, codebook)), + target_vocab=1152, + sentinel_token_id=_SENTINEL, + eos_token_id=2, + pad_token_id=3, + ) + + +def _slot(name, fill, width_n=None) -> SlotSeg: + return SlotSeg( + slot_id=0, + name=name, + sources=(name,), + group_type=FeatureGroupType.JAGGED_SEQUENCE, + output_key=".sequence", + fill=fill, + width=Width(WidthKind.BOUNDED, width_n) + if width_n + else Width(WidthKind.STATIC, 1), + droppable=False, + ) + + +def _plan(segments, response=(), max_length=0) -> PromptPlan: + projected = tuple( + s + for s in segments + tuple(response) + if isinstance(s, SlotSeg) and s.fill is FillMode.PROJECTED + ) + return PromptPlan( + segments=tuple(segments), + response_segments=tuple(response), + max_length=max_length, + max_total_length=None, + max_holes=0, + suffix_keep=None, + static_prefix_len=0, + length_buckets=(), + slot_index={s.name: i for i, s in enumerate(projected)}, + projected_slots=projected, + ) + + +class PromptAssemblerTest(unittest.TestCase): + def test_inline_sid_gets_the_base_vocab_shift(self) -> None: + plan = _plan((Static((7, 8), None), _slot("hist", FillMode.INLINE))) + asm = PromptAssembler(plan, _sid_space()) + # offset codes for one item: level 0 -> 1, level 1 -> 4+2, level 2 -> 8+3 + out = asm.assemble({"hist": [np.array([1, 6, 11])]}) + + self.assertEqual( + out.input_ids.tolist(), [7, 8, _BASE + 1, _BASE + 6, _BASE + 11] + ) + self.assertEqual(out.cu_seqlens.tolist(), [0, 5]) + self.assertEqual(out.hole_positions.size, 0) + + def test_projected_emits_sentinels_and_records_holes(self) -> None: + plan = _plan((Static((7,), None), _slot("prof", FillMode.PROJECTED, 4))) + asm = PromptAssembler(plan, _sid_space()) + out = asm.assemble({}, {"prof": np.array([2, 3])}, batch_size=2) + + # row 0: [7, S, S] row 1: [7, S, S, S] + self.assertEqual( + out.input_ids.tolist(), + [7, _SENTINEL, _SENTINEL, 7, _SENTINEL, _SENTINEL, _SENTINEL], + ) + self.assertEqual(out.cu_seqlens.tolist(), [0, 3, 7]) + # absolute indices into the flat buffer, which is what index_copy needs + self.assertEqual(out.hole_positions.tolist(), [1, 2, 4, 5, 6]) + + def test_hole_positions_index_the_flat_buffer_exactly(self) -> None: + plan = _plan((_slot("prof", FillMode.PROJECTED, 2),)) + asm = PromptAssembler(plan, _sid_space()) + out = asm.assemble({}, {"prof": np.array([2, 2])}, batch_size=2) + # index_copy requires index.numel() == source.size(0) + self.assertEqual(out.hole_positions.size, 4) + self.assertTrue(np.all(out.input_ids[out.hole_positions] == _SENTINEL)) + + def test_labels_cover_the_response_span_only(self) -> None: + plan = _plan( + (Static((7, 8), None),), + response=(Static((9,), None), _slot("answer", FillMode.INLINE)), + ) + asm = PromptAssembler(plan, _sid_space()) + out = asm.assemble({"answer": [np.array([0, 4, 8])]}) + + self.assertEqual(out.input_ids.tolist(), [7, 8, 9, _BASE, _BASE + 4, _BASE + 8]) + # the prompt is context; supervision starts at the response + self.assertEqual( + out.labels.tolist(), [-100, -100, 9, _BASE, _BASE + 4, _BASE + 8] + ) + + def test_rejects_raw_codes_that_carry_no_offset(self) -> None: + plan = _plan((_slot("hist", FillMode.INLINE),)) + asm = PromptAssembler(plan, _sid_space()) + # [1, 2, 3] is a valid raw SID but level 1 and 2 are below their bands + with self.assertRaisesRegex(ValueError, "offset_codebook column"): + asm.assemble({"hist": [np.array([1, 2, 3])]}) + + def test_rejects_a_partial_item(self) -> None: + plan = _plan((_slot("hist", FillMode.INLINE),)) + asm = PromptAssembler(plan, _sid_space()) + with self.assertRaisesRegex(ValueError, "whole number of 3-level items"): + asm.assemble({"hist": [np.array([1, 6])]}) + + def test_rejects_an_out_of_band_code(self) -> None: + plan = _plan((_slot("hist", FillMode.INLINE),)) + asm = PromptAssembler(plan, _sid_space()) + # level 2 admits [8, 12); 12 is the first value past it + with self.assertRaisesRegex(ValueError, "offset_codebook column"): + asm.assemble({"hist": [np.array([1, 6, 12])]}) + + def test_over_long_row_is_an_error_not_a_truncation(self) -> None: + plan = _plan( + (Static((7, 8, 9), None), _slot("hist", FillMode.INLINE)), max_length=4 + ) + asm = PromptAssembler(plan, _sid_space()) + with self.assertRaisesRegex(ValueError, "never truncated"): + asm.assemble({"hist": [np.array([1, 6, 11])]}) + + def test_inline_without_a_sid_space_is_rejected_at_construction(self) -> None: + plan = _plan((_slot("hist", FillMode.INLINE),)) + with self.assertRaisesRegex(ValueError, "no sid_space was compiled"): + PromptAssembler(plan, None) + + def test_rows_of_different_lengths_pack_without_padding(self) -> None: + plan = _plan((_slot("hist", FillMode.INLINE),)) + asm = PromptAssembler(plan, _sid_space()) + out = asm.assemble( + {"hist": [np.array([1, 6, 11]), np.array([0, 4, 8, 2, 5, 9])]} + ) + self.assertEqual(out.cu_seqlens.tolist(), [0, 3, 9]) + self.assertEqual(out.input_ids.size, 9) + + +if __name__ == "__main__": + unittest.main() From c61003caf221260ed897876b0cd38d3735d971fd Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 07:25:29 +0000 Subject: [PATCH 60/99] [feat] prompt: add the slot projection module Reconciles a prompt slot's group width with the LM hidden size: an optional body, then a bare Linear that is structural rather than configurable. Every Perceptron applies its activation and MLP.output_dim() raises on an empty stack, so ending on the MLP would apply an activation to the LM input space and an empty body would not degrade to a linear. The module takes in_dim from the caller. Nothing here or in the compiled plan stores a dimension: the model resolves both ends at __init__, from group_total_dim and the backbone config. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/modules/prompt_projection.py | 69 +++++++++++++++++++++++++ tzrec/modules/prompt_projection_test.py | 62 ++++++++++++++++++++++ 2 files changed, 131 insertions(+) create mode 100644 tzrec/modules/prompt_projection.py create mode 100644 tzrec/modules/prompt_projection_test.py diff --git a/tzrec/modules/prompt_projection.py b/tzrec/modules/prompt_projection.py new file mode 100644 index 000000000..760fbf14e --- /dev/null +++ b/tzrec/modules/prompt_projection.py @@ -0,0 +1,69 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reconciles a prompt slot's group width with the LM hidden size.""" + +from typing import Optional + +import torch +from torch import nn + +from tzrec.modules.mlp import MLP +from tzrec.protos.prompt_pb2 import PromptProjection as PromptProjectionConfig +from tzrec.utils.config_util import config_to_kwargs + + +class PromptProjection(nn.Module): + """An optional body followed by a bare Linear to the LM hidden size. + + The final map never carries an activation: every ``Perceptron`` applies one, + so ending on an MLP would zero half the dimensions feeding the LM input + space. An empty MLP does not degrade to a linear either -- ``output_dim()`` + raises -- which is why the Linear is structural rather than configurable. + + Args: + config: the slot's projection config; an empty one is a plain linear. + in_dim: the slot's ``group_total_dim``, resolved by the model. + hidden_size: the LM hidden size. + """ + + def __init__( + self, + config: PromptProjectionConfig, + in_dim: int, + hidden_size: int, + ) -> None: + super().__init__() + self._in_dim = in_dim + dim = in_dim + self.body: Optional[MLP] = None + if config.HasField("mlp"): + self.body = MLP(dim, **config_to_kwargs(config.mlp)) + dim = self.body.output_dim() + self.head = nn.Linear(dim, hidden_size, bias=config.bias) + + @property + def in_dim(self) -> int: + """Input width this module was sized for.""" + return self._in_dim + + def forward(self, features: torch.Tensor) -> torch.Tensor: + """Project a slot's group output into the LM input space. + + Args: + features: ``(..., group_total_dim)``. + + Returns: + ``(..., hidden_size)``. + """ + if self.body is not None: + features = self.body(features) + return self.head(features) diff --git a/tzrec/modules/prompt_projection_test.py b/tzrec/modules/prompt_projection_test.py new file mode 100644 index 000000000..6a75e1747 --- /dev/null +++ b/tzrec/modules/prompt_projection_test.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch +from torch import nn + +from tzrec.modules.prompt_projection import PromptProjection +from tzrec.protos.prompt_pb2 import PromptProjection as PromptProjectionConfig + + +class PromptProjectionTest(unittest.TestCase): + def test_bodyless_config_is_a_plain_linear(self) -> None: + proj = PromptProjection(PromptProjectionConfig(), in_dim=12, hidden_size=8) + self.assertIsNone(proj.body) + self.assertIsInstance(proj.head, nn.Linear) + self.assertEqual(proj(torch.randn(5, 12)).shape, (5, 8)) + + def test_mlp_body_feeds_a_bare_head(self) -> None: + config = PromptProjectionConfig() + config.mlp.hidden_units.extend([16, 6]) + proj = PromptProjection(config, in_dim=12, hidden_size=8) + + self.assertIsNotNone(proj.body) + # the head maps the MLP's output width, not the slot's input width + self.assertEqual(proj.head.in_features, 6) + self.assertEqual(proj.head.out_features, 8) + self.assertEqual(proj(torch.randn(5, 12)).shape, (5, 8)) + + def test_head_has_no_activation(self) -> None: + config = PromptProjectionConfig() + config.mlp.hidden_units.extend([16]) + proj = PromptProjection(config, in_dim=4, hidden_size=6) + # a Perceptron would clamp negatives; a bare Linear must not + with torch.no_grad(): + proj.head.weight.fill_(-1.0) + proj.head.bias.fill_(0.0) + out = proj(torch.ones(1, 4)) + self.assertTrue(bool((out < 0).any())) + + def test_bias_is_configurable(self) -> None: + config = PromptProjectionConfig(bias=False) + proj = PromptProjection(config, in_dim=4, hidden_size=6) + self.assertIsNone(proj.head.bias) + + def test_jagged_rows_project_independently(self) -> None: + proj = PromptProjection(PromptProjectionConfig(), in_dim=4, hidden_size=6) + rows = torch.randn(7, 4) + torch.testing.assert_close(proj(rows)[3], proj(rows[3:4])[0]) + + +if __name__ == "__main__": + unittest.main() From 04154f0763abddde45464f6fe2594e8fdcd957d9 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 07:34:14 +0000 Subject: [PATCH 61/99] [feat] prompt: add the Qwen model core The model reads target_vocab, the module plan and the decode bands off a CompiledPrompt and never reads the prompt's structure. Its config carries only what belongs to the LM: the template, slots, SID space and tokenizer are prompt_config's, so hf_model_id now names the weights alone. The forward gathers the assembled ids through the LM's own input embedding, then index_copy overwrites the projected positions -- out of place, since that gather carries grad. Padding is confined to one adapter at the LM boundary, which takes the collator's max_seqlen rather than lengths.max(): deriving the width here would sync the device to the host on every step. ParamDtype is nested in PromptModelConfig because proto2 enum values are siblings of their enclosing scope, and a top-level one collides. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_qwen.py | 261 ++++++++++++++++++++ tzrec/models/prompt_generative_qwen_test.py | 71 ++++++ tzrec/protos/model.proto | 2 + tzrec/protos/models/prompt_model.proto | 40 +++ 4 files changed, 374 insertions(+) create mode 100644 tzrec/models/prompt_generative_qwen.py create mode 100644 tzrec/models/prompt_generative_qwen_test.py create mode 100644 tzrec/protos/models/prompt_model.proto diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py new file mode 100644 index 000000000..96d6fdfb2 --- /dev/null +++ b/tzrec/models/prompt_generative_qwen.py @@ -0,0 +1,261 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Prompt-native generative recommendation over a Qwen backbone. + +The model reads ``target_vocab``, the module plan and the decode bands from a +``CompiledPrompt``, and never reads the prompt's structure. Assembly happens in +the dataloader worker; what arrives here is already a packed token stream plus +the positions the projected slots must overwrite. +""" + +from typing import Any, Dict, List, Optional + +import torch +from torch import nn +from transformers import AutoConfig, AutoModelForCausalLM + +from tzrec.datasets.utils import Batch +from tzrec.features.feature import BaseFeature +from tzrec.models.model import BaseModel +from tzrec.modules.prompt_projection import PromptProjection +from tzrec.prompt.plan import CompiledPrompt, SlotSeg +from tzrec.protos.model_pb2 import ModelConfig +from tzrec.protos.models.prompt_model_pb2 import PromptModelConfig +from tzrec.utils.logging_util import logger + +_PROMPT_INPUT_IDS = "prompt_input_ids" +_PROMPT_CU_SEQLENS = "prompt_cu_seqlens" +_PROMPT_HOLE_POSITIONS = "prompt_hole_positions" +_PROMPT_LABELS = "prompt_labels" +_PROMPT_MAX_SEQLEN = "prompt_max_seqlen" + +_PARAM_DTYPE: Dict[int, torch.dtype] = { + PromptModelConfig.FP32: torch.float32, + PromptModelConfig.BF16: torch.bfloat16, + PromptModelConfig.FP16: torch.float16, +} + + +class PromptGenerativeQwen(BaseModel): + """Qwen backbone driven by a compiled prompt. + + Args: + model_config: the model oneof. + features: every created feature. + labels: data_config label fields. + sample_weights: optional sample weight fields. + prompt: the compiled prompt; required. + """ + + def __init__( + self, + model_config: ModelConfig, + features: List[BaseFeature], + labels: List[str], + sample_weights: Optional[List[str]] = None, + prompt: Optional[CompiledPrompt] = None, + **kwargs: Any, + ) -> None: + super().__init__(model_config, features, labels, sample_weights, **kwargs) + if prompt is None: + raise ValueError( + f"{type(self).__name__} needs a compiled prompt; call " + f"compile_prompt(pipeline_config.prompt_config, features) and " + f"pass it to _create_model." + ) + self._prompt = prompt + cfg = self._model_config + common = cfg.common + + self._ignore_index = int(common.ignore_index) + self._generated_sids_key = common.generated_sids_key + self._read_beam_config(common) + + self.lm = self._build_backbone(cfg.hf_model_id, common.param_dtype) + self.lm.resize_token_embeddings( + prompt.sid_space.target_vocab, mean_resizing=True + ) + self._build_projections() + + def _read_beam_config(self, common: PromptModelConfig) -> None: + """Parse the decode knobs; the schedule must match the codebook.""" + space = self._prompt.sid_space + if space is None: + raise ValueError( + f"{type(self).__name__}: prompt_config declares no sid_space, " + f"so there is nothing to decode." + ) + self._num_return = int(common.num_return_sequences) + self._beam_widths: List[int] = list(common.beam_widths) + if not self._beam_widths: + raise ValueError( + f"{type(self).__name__}: beam_widths is required; give one " + f"width per SID level, e.g. [50, 50, 50] or [100, 200, 400]." + ) + if len(self._beam_widths) != space.num_levels: + raise ValueError( + f"{type(self).__name__}: beam_widths has " + f"{len(self._beam_widths)} entries but the codebook has " + f"{space.num_levels} levels; give one width per level." + ) + if self._num_return > self._beam_widths[-1]: + raise ValueError( + f"{type(self).__name__}: num_return_sequences " + f"({self._num_return}) must not exceed the final beam width " + f"({self._beam_widths[-1]})." + ) + + def _build_backbone(self, hf_model_id: str, param_dtype: int) -> nn.Module: + """Build the LM empty, so HF weights load only on cold start.""" + config = AutoConfig.from_pretrained(hf_model_id) + model = AutoModelForCausalLM.from_config(config) + return model.to(_PARAM_DTYPE[param_dtype]) + + def _build_projections(self) -> None: + """One module per resolved id, aligned with ``plan.projected_slots``. + + Slots sharing a ``projection_name`` share a module by reference, so + they must agree on ``group_total_dim``. + """ + plan = self._prompt.prompt_plan + modules = self._prompt.module_plan + hidden = int(self.lm.config.hidden_size) + + built: Dict[str, PromptProjection] = {} + aligned: List[PromptProjection] = [] + for seg in plan.projected_slots: + module_id = modules.slot_to_module[seg.slot_id] + in_dim = self._slot_in_dim(seg) + if module_id not in built: + built[module_id] = PromptProjection( + modules.projections[module_id], in_dim, hidden + ) + elif built[module_id].in_dim != in_dim: + raise ValueError( + f"prompt slots sharing projection_name [{module_id}] have " + f"different group widths ({built[module_id].in_dim} vs " + f"{in_dim}); they cannot share a module." + ) + aligned.append(built[module_id]) + self.projections = nn.ModuleDict(built) + # zipped with plan.projected_slots; shared modules appear by reference + self._slot_projections = aligned + + def _slot_in_dim(self, seg: SlotSeg) -> int: + """Total group output width of a projected slot.""" + return self.embedding_group.group_total_dim(seg.name + seg.output_key) + + def hf_backbone(self) -> nn.Module: + """The HF module export and checkpointing reach for.""" + return self.lm + + def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: + """Teacher-forced forward over the assembled stream. + + Args: + batch: carries the packed prompt in ``additional_infos``. + + Returns: + The loss. + """ + embeds = self._prompt_embeds(batch) + return self._forward_loss(embeds, batch) + + def _prompt_embeds(self, batch: Batch) -> torch.Tensor: + """Gather the token stream, then overwrite the projected positions.""" + ids = batch.additional_infos[_PROMPT_INPUT_IDS] + embeds = self.lm.get_input_embeddings()(ids) + + plan = self._prompt.prompt_plan + if not plan.projected_slots: + return embeds + + grouped = self.embedding_group(batch) + hidden = embeds.shape[-1] + parts = [ + proj(grouped[seg.name + seg.output_key]).reshape(-1, hidden) + for seg, proj in zip(plan.projected_slots, self._slot_projections) + ] + # out of place: embeds carries grad from the embedding lookup + return embeds.index_copy( + 0, batch.additional_infos[_PROMPT_HOLE_POSITIONS], torch.cat(parts) + ) + + def _forward_loss( + self, embeds: torch.Tensor, batch: Batch + ) -> Dict[str, torch.Tensor]: + """Run the LM over the assembled embeddings and score the response.""" + infos = batch.additional_infos + padded, mask, labels = _unpack( + embeds, + infos[_PROMPT_CU_SEQLENS], + infos[_PROMPT_LABELS], + int(infos[_PROMPT_MAX_SEQLEN]), + self._ignore_index, + ) + outputs = self.lm.model(inputs_embeds=padded, attention_mask=mask) + + suffix = self._prompt.prompt_plan.suffix_keep + window = slice(-suffix, None) if suffix else slice(None) + logits = self.lm.lm_head(outputs.last_hidden_state[:, window, :]) + loss = self.lm.loss_function( + logits=logits, + labels=labels[:, window], + vocab_size=self.lm.config.vocab_size, + ignore_index=self._ignore_index, + ) + return {"loss": loss} + + def init_from_pretrained(self) -> None: + """Load HF weights once, on a cold start only.""" + source = self._model_config.hf_model_id + logger.info(f"loading pretrained weights from [{source}].") + pretrained = AutoModelForCausalLM.from_pretrained(source) + pretrained.resize_token_embeddings( + self._prompt.sid_space.target_vocab, mean_resizing=True + ) + self.lm.load_state_dict(pretrained.state_dict()) + del pretrained + + +def _unpack( + embeds: torch.Tensor, + cu_seqlens: torch.Tensor, + labels: torch.Tensor, + max_seqlen: int, + ignore_index: int, +) -> "tuple[torch.Tensor, torch.Tensor, torch.Tensor]": + """Pad a packed varlen batch at the LM boundary. + + Padding lives in this one adapter. ``max_seqlen`` is the collator's, not + ``lengths.max()``: deriving it here would sync the device to the host every + step, which §7.4 of the design forbids. + """ + starts = cu_seqlens[:-1] + lengths = cu_seqlens[1:] - starts + batch_size = lengths.numel() + hidden = embeds.shape[-1] + + columns = torch.arange(max_seqlen, device=embeds.device) + mask = columns[None, :] < lengths[:, None] + + padded = embeds.new_zeros((batch_size, max_seqlen, hidden)) + out_labels = torch.full( + (batch_size, max_seqlen), + ignore_index, + dtype=labels.dtype, + device=embeds.device, + ) + # mask selects row-major, which is the order embeds and labels are packed in + padded[mask] = embeds + out_labels[mask] = labels + return padded, mask.long(), out_labels diff --git a/tzrec/models/prompt_generative_qwen_test.py b/tzrec/models/prompt_generative_qwen_test.py new file mode 100644 index 000000000..ea2d4906b --- /dev/null +++ b/tzrec/models/prompt_generative_qwen_test.py @@ -0,0 +1,71 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch + +from tzrec.models.prompt_generative_qwen import _unpack + + +class UnpackTest(unittest.TestCase): + """The one adapter where padding lives.""" + + def test_packs_rows_of_different_lengths(self) -> None: + # rows of 2 and 3 tokens, hidden size 4 + embeds = torch.arange(20, dtype=torch.float32).reshape(5, 4) + cu = torch.tensor([0, 2, 5]) + labels = torch.tensor([10, 11, 20, 21, 22]) + + padded, mask, out = _unpack(embeds, cu, labels, max_seqlen=3, ignore_index=-100) + + self.assertEqual(padded.shape, (2, 3, 4)) + self.assertEqual(mask.tolist(), [[1, 1, 0], [1, 1, 1]]) + torch.testing.assert_close(padded[0, :2], embeds[:2]) + torch.testing.assert_close(padded[1, :3], embeds[2:]) + # the pad column is zero, and its label is ignored + torch.testing.assert_close(padded[0, 2], torch.zeros(4)) + self.assertEqual(out.tolist(), [[10, 11, -100], [20, 21, 22]]) + + def test_uses_the_given_width_not_the_observed_max(self) -> None: + # the collator's bucket may exceed the widest row; §7.4 forbids + # deriving the width on device + embeds = torch.ones(3, 2) + cu = torch.tensor([0, 1, 3]) + labels = torch.tensor([1, 2, 3]) + padded, mask, _ = _unpack(embeds, cu, labels, max_seqlen=5, ignore_index=-100) + + self.assertEqual(padded.shape, (2, 5, 2)) + self.assertEqual(mask.sum().item(), 3) + + def test_row_order_survives_the_scatter(self) -> None: + # mask selects row-major, which must match the packing order + embeds = torch.tensor([[1.0], [2.0], [3.0], [4.0]]) + cu = torch.tensor([0, 1, 4]) + labels = torch.zeros(4, dtype=torch.long) + padded, _, _ = _unpack(embeds, cu, labels, max_seqlen=3, ignore_index=-100) + + self.assertEqual(padded[0, 0].item(), 1.0) + self.assertEqual(padded[1, :, 0].tolist(), [2.0, 3.0, 4.0]) + + def test_gradient_reaches_the_packed_input(self) -> None: + embeds = torch.ones(3, 2, requires_grad=True) + cu = torch.tensor([0, 1, 3]) + labels = torch.zeros(3, dtype=torch.long) + padded, _, _ = _unpack(embeds, cu, labels, max_seqlen=2, ignore_index=-100) + padded.sum().backward() + + self.assertIsNotNone(embeds.grad) + torch.testing.assert_close(embeds.grad, torch.ones(3, 2)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/protos/model.proto b/tzrec/protos/model.proto index c8243c35d..46ae7d174 100644 --- a/tzrec/protos/model.proto +++ b/tzrec/protos/model.proto @@ -6,6 +6,7 @@ import "tzrec/protos/models/multi_task_rank.proto"; import "tzrec/protos/models/match_model.proto"; import "tzrec/protos/models/general_rank_model.proto"; import "tzrec/protos/models/generative_model.proto"; +import "tzrec/protos/models/prompt_model.proto"; import "tzrec/protos/models/sid_model.proto"; import "tzrec/protos/loss.proto"; import "tzrec/protos/metric.proto"; @@ -85,6 +86,7 @@ message ModelConfig { // Generative (causal-LM) models; the 700-block keeps clear of the SID 600s. GenerativeQwen generative_qwen = 700; + PromptGenerativeQwen prompt_generative_qwen = 701; } optional uint32 num_class = 2 [default = 1]; diff --git a/tzrec/protos/models/prompt_model.proto b/tzrec/protos/models/prompt_model.proto new file mode 100644 index 000000000..c945a973e --- /dev/null +++ b/tzrec/protos/models/prompt_model.proto @@ -0,0 +1,40 @@ +syntax = "proto2"; +package tzrec.protos; + +// Prompt-native generative recommendation models. Everything about the prompt +// itself -- template, slots, SID space, tokenizer -- lives in prompt_config; +// these messages carry only what belongs to the LM. + +message PromptModelConfig { + // Nested so it does not collide with another model's dtype enum: proto2 + // enum values are siblings of their enclosing scope. + enum ParamDtype { + FP32 = 0; + BF16 = 1; + FP16 = 2; + } + + // Label value outside the supervised span. + optional int32 ignore_index = 1 [default = -100]; + + // REQUIRED. Beam width per SID level, one entry per level; [100, 200, 400] + // is the escalating beam. Each entry is capped to what its band supplies. + repeated uint32 beam_widths = 2; + // Must not exceed the final beam width. + required uint32 num_return_sequences = 3; + + optional string generated_sids_key = 4 [default = "generated_sids"]; + + // MASTER weights. FP32 avoids bf16-ULP underflow of Adam's small updates; + // bf16 COMPUTE comes from mixed_precision, not from this. + optional ParamDtype param_dtype = 5 [default = FP32]; +} + +// Qwen family (Qwen2.5, Qwen3, ...). +message PromptGenerativeQwen { + optional PromptModelConfig common = 1; + + // HF hub id or local path. Names the WEIGHTS only, and is read solely by + // init_from_pretrained at cold start; the vocabulary is prompt_config's. + optional string hf_model_id = 2 [default = "Qwen/Qwen2.5-0.5B"]; +} From 330167afce789c8f24e707a09de603fc6f627684 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 07:40:25 +0000 Subject: [PATCH 62/99] [feat] prompt: decode SIDs from the assembled prompt Wires band-restricted beam decode onto the compiled SidSpace: predict returns the loss while training and the decoded local codes at inference, undoing both the base_vocab and the per-level shift. dynamic_beam_search now prefills from embeddings rather than ids. A projected slot has no vocabulary id, so an id-only prefill cannot express a prompt that contains one. Decode steps still pass ids, because a generated token is always a real vocabulary row. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_qwen.py | 46 ++++++++++++++++++++- tzrec/models/prompt_generative_qwen_test.py | 37 +++++++++++++++++ tzrec/modules/dynamic_beam.py | 11 ++--- tzrec/modules/dynamic_beam_test.py | 10 ++++- 4 files changed, 96 insertions(+), 8 deletions(-) diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 96d6fdfb2..516bd0534 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -26,6 +26,7 @@ from tzrec.datasets.utils import Batch from tzrec.features.feature import BaseFeature from tzrec.models.model import BaseModel +from tzrec.modules.dynamic_beam import dynamic_beam_search from tzrec.modules.prompt_projection import PromptProjection from tzrec.prompt.plan import CompiledPrompt, SlotSeg from tzrec.protos.model_pb2 import ModelConfig @@ -165,11 +166,54 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: batch: carries the packed prompt in ``additional_infos``. Returns: - The loss. + The loss when training, the decoded SIDs otherwise. """ embeds = self._prompt_embeds(batch) + if self.is_inference: + return {self._generated_sids_key: self._generate(embeds, batch)} return self._forward_loss(embeds, batch) + def _sid_token_bands(self) -> "tuple[torch.Tensor, torch.Tensor]": + """Inclusive token-id band of every SID level, as device tensors.""" + space = self._prompt.sid_space + device = self.lm.get_input_embeddings().weight.device + return ( + torch.tensor(space.band_lo, device=device), + torch.tensor(space.band_hi, device=device), + ) + + def _generate(self, embeds: torch.Tensor, batch: Batch) -> torch.Tensor: + """Beam-search the SID answer. + + Args: + embeds: the assembled prompt embeddings, packed. + batch: carries ``prompt_cu_seqlens`` and the collator's width. + + Returns: + ``(B, num_return, num_levels)`` local codes, best first. + """ + infos = batch.additional_infos + padded, mask, _ = _unpack( + embeds, + infos[_PROMPT_CU_SEQLENS], + infos[_PROMPT_LABELS], + int(infos[_PROMPT_MAX_SEQLEN]), + self._ignore_index, + ) + lo_tok, hi_tok = self._sid_token_bands() + tokens = dynamic_beam_search( + self.lm, padded, mask, self._beam_widths, lo_tok, hi_tok + ) + return self._detokenize(tokens, padded.shape[0]) + + def _detokenize(self, tokens: torch.Tensor, batch_size: int) -> torch.Tensor: + """Undo both shifts: token id back to a local 0-based code.""" + space = self._prompt.sid_space + offsets = torch.tensor(space.level_offsets, device=tokens.device) + codes = tokens - space.base_vocab - offsets + codes = codes.view(batch_size, -1, space.num_levels) + return codes[:, : self._num_return, :] + def _prompt_embeds(self, batch: Batch) -> torch.Tensor: """Gather the token stream, then overwrite the projected positions.""" ids = batch.additional_infos[_PROMPT_INPUT_IDS] diff --git a/tzrec/models/prompt_generative_qwen_test.py b/tzrec/models/prompt_generative_qwen_test.py index ea2d4906b..3112e707d 100644 --- a/tzrec/models/prompt_generative_qwen_test.py +++ b/tzrec/models/prompt_generative_qwen_test.py @@ -14,6 +14,7 @@ import torch from tzrec.models.prompt_generative_qwen import _unpack +from tzrec.prompt.plan import SidSpace class UnpackTest(unittest.TestCase): @@ -67,5 +68,41 @@ def test_gradient_reaches_the_packed_input(self) -> None: torch.testing.assert_close(embeds.grad, torch.ones(3, 2)) +class DetokenizeTest(unittest.TestCase): + """Both shifts must come back off, in the right order.""" + + def _space(self) -> SidSpace: + return SidSpace( + codebook=(4, 4, 4), + num_levels=3, + base_vocab=1000, + level_offsets=(0, 4, 8), + band_lo=(1000, 1004, 1008), + band_hi=(1003, 1007, 1011), + target_vocab=1152, + sentinel_token_id=None, + eos_token_id=2, + pad_token_id=3, + ) + + def test_token_ids_become_local_codes(self) -> None: + space = self._space() + # one beam row: level 0 code 1, level 1 code 2, level 2 code 3 + tokens = torch.tensor([[1000 + 1, 1000 + 4 + 2, 1000 + 8 + 3]]) + offsets = torch.tensor(space.level_offsets) + codes = (tokens - space.base_vocab - offsets).view(1, -1, 3) + + self.assertEqual(codes[0, 0].tolist(), [1, 2, 3]) + # every code lands back inside its own codebook + self.assertTrue(bool(((codes >= 0) & (codes < 4)).all())) + + def test_a_band_edge_maps_to_the_last_code(self) -> None: + space = self._space() + tokens = torch.tensor([list(space.band_hi)]) + offsets = torch.tensor(space.level_offsets) + codes = tokens - space.base_vocab - offsets + self.assertEqual(codes[0].tolist(), [3, 3, 3]) + + if __name__ == "__main__": unittest.main() diff --git a/tzrec/modules/dynamic_beam.py b/tzrec/modules/dynamic_beam.py index e7be2bfd3..53af74da3 100644 --- a/tzrec/modules/dynamic_beam.py +++ b/tzrec/modules/dynamic_beam.py @@ -24,7 +24,7 @@ @torch.no_grad() def dynamic_beam_search( model: PreTrainedModel, - input_ids: torch.Tensor, + prompt_embeds: torch.Tensor, attention_mask: torch.Tensor, beam_widths: List[int], lo_tok: torch.Tensor, @@ -34,7 +34,8 @@ def dynamic_beam_search( Args: model: an HF causal LM exposing ``.model`` / ``.lm_head`` (Qwen layout). - input_ids: left-padded prompt ids ``(B, P)``. + prompt_embeds: left-padded prompt embeddings ``(B, P, D)``. Embeddings + rather than ids, because a projected slot has no vocabulary id. attention_mask: prompt mask ``(B, P)``. beam_widths: requested width per SID level; each is capped to what its band and the surviving prefixes supply. @@ -45,8 +46,8 @@ def dynamic_beam_search( The SID token tail ``(B * W, num_levels)``, score-ordered best-first. The answer is fixed-length and EOS-free, so no beam bookkeeping. """ - device = input_ids.device - batch_size = input_ids.shape[0] + device = prompt_embeds.device + batch_size = prompt_embeds.shape[0] num_levels = lo_tok.shape[0] if len(beam_widths) != num_levels: raise ValueError( @@ -79,7 +80,7 @@ def _band_logp(logits: torch.Tensor, level: int) -> torch.Tensor: position_ids = (attention_mask.long().cumsum(-1) - 1).clamp(min=0) outputs = model.model( - input_ids=input_ids, + inputs_embeds=prompt_embeds, attention_mask=attention_mask, position_ids=position_ids, use_cache=True, diff --git a/tzrec/modules/dynamic_beam_test.py b/tzrec/modules/dynamic_beam_test.py index 2603e77fa..08ace1f6a 100644 --- a/tzrec/modules/dynamic_beam_test.py +++ b/tzrec/modules/dynamic_beam_test.py @@ -34,7 +34,7 @@ def _decode(lm, ids, pairs, width=8, beam_widths=None, attention_mask=None): beam_widths = [width] * len(pairs) return dynamic_beam_search( lm, - ids, + lm.get_input_embeddings()(ids), torch.ones_like(ids) if attention_mask is None else attention_mask, beam_widths=beam_widths, lo_tok=torch.tensor([p[0] for p in pairs]), @@ -54,8 +54,14 @@ def __init__(self, lm) -> None: self.rows: List[int] = [] self._lm = lm + def get_input_embeddings(self): + return self._lm.get_input_embeddings() + def model(self, **kwargs: Any) -> Any: - self.rows.append(kwargs["input_ids"].shape[0]) + first = kwargs.get("input_ids") + if first is None: + first = kwargs["inputs_embeds"] + self.rows.append(first.shape[0]) return self._lm.model(**kwargs) From f631ea8e9f07a21079391788780d274958a76a94 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 07:46:34 +0000 Subject: [PATCH 63/99] [feat] prompt: wire the compiler into the entry points Every entry point now compiles prompt_config once, right after the features exist, and threads the result to both consumers: the dataloader, whose workers assemble each batch's token stream into additional_infos, and the model, which reads target_vocab and the decode bands off it. Assembly runs after the data parser, not inside any feature's _parse, so the walk stays pure-integer and FG-free. max_seqlen is computed there, on the host, because the model must not derive a shape on the device. Passing prompt through _create_model is conditional so models that take no such kwarg are untouched. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/datasets/dataset.py | 17 +++++++ tzrec/main.py | 34 +++++++++++++ tzrec/models/prompt_generative_qwen.py | 21 +++++--- tzrec/prompt/assembler.py | 69 +++++++++++++++++++++++++- 4 files changed, 134 insertions(+), 7 deletions(-) diff --git a/tzrec/datasets/dataset.py b/tzrec/datasets/dataset.py index e7ae109e9..6b1027b7d 100644 --- a/tzrec/datasets/dataset.py +++ b/tzrec/datasets/dataset.py @@ -18,6 +18,7 @@ import numpy as np import pyarrow as pa import pyarrow.compute as pc +import torch from torch import distributed as dist from torch.utils.data import DataLoader, IterableDataset, get_worker_info @@ -40,6 +41,8 @@ remove_nullable, ) from tzrec.features.feature import BaseFeature +from tzrec.prompt.assembler import assemble_into +from tzrec.prompt.plan import CompiledPrompt from tzrec.protos import data_pb2 from tzrec.utils.load_class import get_register_class_meta from tzrec.utils.logging_util import logger @@ -106,8 +109,10 @@ def __init__( reserved_columns: Optional[List[str]] = None, mode: Mode = Mode.EVAL, debug_level: int = 0, + prompt: Optional[CompiledPrompt] = None, ) -> None: super(BaseDataset, self).__init__() + self._prompt = prompt self._data_config = data_config self._features = features self._input_path = input_path @@ -382,6 +387,14 @@ def _build_batch(self, input_data: Dict[str, pa.Array]) -> Batch: else: batch = self._data_parser.to_batch(output_data) + if self._prompt is not None: + batch.additional_infos.update( + { + k: torch.from_numpy(np.asarray(v)) + for k, v in assemble_into(self._prompt, output_data).items() + } + ) + # Set checkpoint info on batch batch.checkpoint_info = checkpoint_info batch.data_timestamp = data_timestamp @@ -759,6 +772,7 @@ def create_dataloader( gl_cluster: Optional[Dict[str, Union[int, str]]] = None, debug_level: int = 0, checkpoint_state: Optional[Dict[str, Any]] = None, + prompt: Optional[CompiledPrompt] = None, ) -> DataLoader: """Build dataloader. @@ -773,6 +787,8 @@ def create_dataloader( debug_level > 0, will dump fg encoded data to debug_str checkpoint_state (dict, optional): resume state, applied before the eager ``iter()`` forks workers so it reaches them. + prompt (CompiledPrompt, optional): when set, each batch carries the + assembled prompt streams in ``additional_infos``. Return: dataloader (dataloader): a DataLoader. @@ -787,6 +803,7 @@ def create_dataloader( reserved_columns=reserved_columns, mode=mode, debug_level=debug_level, + prompt=prompt, ) if checkpoint_state: dataset.load_state_dict(dict(checkpoint_state)) diff --git a/tzrec/main.py b/tzrec/main.py index d8acab684..31e616274 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -73,6 +73,8 @@ from tzrec.optim.ema import DenseEMA, EMAOptimizer from tzrec.optim.lr_scheduler import BaseLR from tzrec.optim.optimizer import TZRecOptimizer +from tzrec.prompt.compile import compile_prompt +from tzrec.prompt.plan import CompiledPrompt from tzrec.protos import export_pb2 from tzrec.protos.data_pb2 import DataConfig, DatasetType from tzrec.protos.eval_pb2 import EvalConfig @@ -80,6 +82,7 @@ from tzrec.protos.feature_pb2 import FeatureConfig from tzrec.protos.model_pb2 import Kernel as KernelProto from tzrec.protos.model_pb2 import ModelConfig +from tzrec.protos.pipeline_pb2 import EasyRecConfig from tzrec.protos.train_pb2 import TrainConfig from tzrec.utils import checkpoint_util, config_util from tzrec.utils.delta_embedding_dump import DeltaEmbeddingDumper @@ -121,6 +124,21 @@ def _create_features( return features +def _compile_prompt( + pipeline_config: EasyRecConfig, features: List[BaseFeature] +) -> Optional[CompiledPrompt]: + """Compile prompt_config when the pipeline declares one. + + Runs on every entry point, so the plan the data layer walks and the vocab + the model resizes to are produced by one code path. + """ + if not pipeline_config.HasField("prompt_config"): + return None + return compile_prompt( + pipeline_config.prompt_config, features, model_dir=pipeline_config.model_dir + ) + + def _get_sampler_type(data_config: DataConfig) -> Optional[str]: try: sampler_type = ( @@ -139,6 +157,7 @@ def _create_model( labels: List[str], sample_weights: Optional[List[str]] = None, sampler_type: Optional[str] = None, + prompt: Optional[CompiledPrompt] = None, ) -> BaseModel: """Build model. @@ -148,6 +167,8 @@ def _create_model( labels (list): list of label names. sample_weights (list): list of sample weight names. sampler_type (str): negative sampler type + prompt (CompiledPrompt, optional): forwarded to prompt-native models. + Return: model: a EasyRec Model. """ @@ -155,12 +176,14 @@ def _create_model( # pyre-ignore [16] model_cls = BaseModel.create_class(model_cls_name) + extra: Dict[str, Any] = {"prompt": prompt} if prompt is not None else {} model: BaseModel = model_cls( model_config, features, labels, sample_weights=sample_weights, sampler_type=sampler_type, + **extra, ) kernel = Kernel[KernelProto.Name(model_config.kernel)] @@ -697,6 +720,7 @@ def train_and_evaluate( data_config = pipeline_config.data_config # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) + prompt = _compile_prompt(pipeline_config, features) ckpt_manager = checkpoint_util.CheckpointManager( pipeline_config.model_dir, @@ -747,6 +771,7 @@ def train_and_evaluate( features, pipeline_config.train_input_path, mode=Mode.TRAIN, + prompt=prompt, checkpoint_state=dataloader_state, ) eval_dataloader = None @@ -758,6 +783,7 @@ def train_and_evaluate( features, pipeline_config.eval_input_path, mode=Mode.EVAL, + prompt=prompt, gl_cluster=gl_cluster, ) @@ -957,12 +983,14 @@ def evaluate( data_config = pipeline_config.data_config # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) + prompt = _compile_prompt(pipeline_config, features) eval_dataloader = create_dataloader( data_config, features, eval_input_path or pipeline_config.eval_input_path, mode=Mode.EVAL, + prompt=prompt, ) sampler_type = _get_sampler_type(data_config) @@ -1118,6 +1146,7 @@ def export( # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) + prompt = _compile_prompt(pipeline_config, features) # Build model model = _create_model( @@ -1125,6 +1154,7 @@ def export( features, list(data_config.label_fields), sampler_type=None, + prompt=prompt, ) InferWrapper = ScriptWrapper # Flip to inference *before* wrapping so view-dependent state @@ -1310,6 +1340,7 @@ def predict( data_config.drop_remainder = False # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) + prompt = _compile_prompt(pipeline_config, features) infer_dataloader = create_dataloader( data_config, @@ -1317,6 +1348,7 @@ def predict( predict_input_path, reserved_columns=reserved_cols, mode=Mode.PREDICT, + prompt=prompt, debug_level=debug_level, ) infer_iterator = infer_dataloader.get_iterator() # pyre-ignore[16] @@ -1524,6 +1556,7 @@ def predict_checkpoint( data_config = pipeline_config.data_config # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) + prompt = _compile_prompt(pipeline_config, features) # Build dataloader predict_dataloader = create_dataloader( @@ -1532,6 +1565,7 @@ def predict_checkpoint( predict_input_path, reserved_columns=reserved_cols, mode=Mode.PREDICT, + prompt=prompt, debug_level=debug_level, ) diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 516bd0534..64ec58ab5 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -28,17 +28,26 @@ from tzrec.models.model import BaseModel from tzrec.modules.dynamic_beam import dynamic_beam_search from tzrec.modules.prompt_projection import PromptProjection +from tzrec.prompt.assembler import ( + PROMPT_CU_SEQLENS as _PROMPT_CU_SEQLENS, +) +from tzrec.prompt.assembler import ( + PROMPT_HOLE_POSITIONS as _PROMPT_HOLE_POSITIONS, +) +from tzrec.prompt.assembler import ( + PROMPT_INPUT_IDS as _PROMPT_INPUT_IDS, +) +from tzrec.prompt.assembler import ( + PROMPT_LABELS as _PROMPT_LABELS, +) +from tzrec.prompt.assembler import ( + PROMPT_MAX_SEQLEN as _PROMPT_MAX_SEQLEN, +) from tzrec.prompt.plan import CompiledPrompt, SlotSeg from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.models.prompt_model_pb2 import PromptModelConfig from tzrec.utils.logging_util import logger -_PROMPT_INPUT_IDS = "prompt_input_ids" -_PROMPT_CU_SEQLENS = "prompt_cu_seqlens" -_PROMPT_HOLE_POSITIONS = "prompt_hole_positions" -_PROMPT_LABELS = "prompt_labels" -_PROMPT_MAX_SEQLEN = "prompt_max_seqlen" - _PARAM_DTYPE: Dict[int, torch.dtype] = { PromptModelConfig.FP32: torch.float32, PromptModelConfig.BF16: torch.bfloat16, diff --git a/tzrec/prompt/assembler.py b/tzrec/prompt/assembler.py index 231cac0ee..a83879067 100644 --- a/tzrec/prompt/assembler.py +++ b/tzrec/prompt/assembler.py @@ -21,7 +21,20 @@ import numpy as np -from tzrec.prompt.plan import FillMode, PromptPlan, SidSpace, SlotSeg, Static +from tzrec.prompt.plan import ( + CompiledPrompt, + FillMode, + PromptPlan, + SidSpace, + SlotSeg, + Static, +) + +PROMPT_INPUT_IDS = "prompt_input_ids" +PROMPT_CU_SEQLENS = "prompt_cu_seqlens" +PROMPT_HOLE_POSITIONS = "prompt_hole_positions" +PROMPT_LABELS = "prompt_labels" +PROMPT_MAX_SEQLEN = "prompt_max_seqlen" @dataclass @@ -41,6 +54,13 @@ class AssembledPrompt: hole_positions: np.ndarray labels: np.ndarray + @property + def max_seqlen(self) -> int: + """Widest row, computed on the host so the model never derives it.""" + if self.cu_seqlens.size < 2: + return 0 + return int(np.max(np.diff(self.cu_seqlens))) + class PromptAssembler: """Walks a ``PromptPlan`` to build token streams. @@ -179,3 +199,50 @@ def assemble( hole_positions=np.asarray(holes, dtype=np.int64), labels=np.asarray(labels, dtype=np.int64), ) + + +def assemble_into( + prompt: CompiledPrompt, + parsed: Dict[str, "np.ndarray"], + ignore_index: int = -100, +) -> Dict[str, np.ndarray]: + """Run the assembler over one parsed batch and key it for the batch. + + Args: + prompt: the compiled prompt. + parsed: ``{feature}.values`` / ``{feature}.lengths`` as the data parser + emits them. + ignore_index: label value outside the supervised span. + + Returns: + The five streams, keyed as ``additional_infos`` expects them. + """ + plan = prompt.prompt_plan + values: Dict[str, List[np.ndarray]] = {} + counts: Dict[str, np.ndarray] = {} + batch_size = 0 + for seg in plan.segments + plan.response_segments: + if not isinstance(seg, SlotSeg): + continue + source = seg.sources[0] + lengths = np.asarray(parsed[f"{source}.lengths"]) + batch_size = max(batch_size, int(lengths.size)) + if seg.fill is FillMode.INLINE: + flat = np.asarray(parsed[f"{source}.values"]) + bounds = np.concatenate(([0], np.cumsum(lengths))) + values[seg.name] = [ + flat[bounds[i] : bounds[i + 1]] for i in range(lengths.size) + ] + else: + counts[seg.name] = lengths + + out = PromptAssembler(plan, prompt.sid_space, ignore_index).assemble( + values, counts, batch_size=batch_size + ) + return { + PROMPT_INPUT_IDS: out.input_ids, + PROMPT_CU_SEQLENS: out.cu_seqlens, + PROMPT_HOLE_POSITIONS: out.hole_positions, + PROMPT_LABELS: out.labels, + PROMPT_MAX_SEQLEN: np.asarray(out.max_seqlen, dtype=np.int64), + } From a7dad479da39e342ea151969f29517492222b8a5 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 07:49:05 +0000 Subject: [PATCH 64/99] [feat] prompt: add an end-to-end stack test Exercises the real path -- compile_prompt, assemble_into, _create_model, forward and backward -- on a two-layer Qwen saved locally, so the test needs no download. It pins the properties the unit tests cannot see from one layer: that the model resizes to the compiled target_vocab, that loss reaches the backbone embedding, that supervision covers the answer alone, and that raw codes are rejected by the assembler before the model is ever built. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/tests/prompt_integration_test.py | 178 +++++++++++++++++++++++++ 1 file changed, 178 insertions(+) create mode 100644 tzrec/tests/prompt_integration_test.py diff --git a/tzrec/tests/prompt_integration_test.py b/tzrec/tests/prompt_integration_test.py new file mode 100644 index 000000000..6709edb36 --- /dev/null +++ b/tzrec/tests/prompt_integration_test.py @@ -0,0 +1,178 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import unittest + +import numpy as np +import torch +from google.protobuf import text_format +from tokenizers import Tokenizer, models, pre_tokenizers +from transformers import Qwen2Config + +from tzrec.datasets.utils import Batch +from tzrec.features.feature import FgMode, create_features +from tzrec.main import _create_model +from tzrec.prompt.assembler import assemble_into +from tzrec.prompt.compile import compile_prompt +from tzrec.protos import feature_pb2 +from tzrec.protos.model_pb2 import ModelConfig +from tzrec.protos.prompt_pb2 import PromptConfig +from tzrec.utils.test_util import make_test_dir + +_CODEBOOK = [4, 4, 4] +_WORDS = ["History", "Predict", ":", ".", "", "<|im_end|>"] + + +def _tiny_backbone(path: str) -> str: + """A two-layer Qwen saved locally, so no download is needed.""" + Qwen2Config( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=256, + ).save_pretrained(path) + return path + + +def _tokenizer(path: str) -> str: + tok = Tokenizer( + models.WordLevel(vocab={w: i for i, w in enumerate(_WORDS)}, unk_token="") + ) + tok.pre_tokenizer = pre_tokenizers.Whitespace() + tok.save(path) + return path + + +def _features(): + text = ( + 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }', + 'sequence_raw_feature { feature_name: "answer" expression: "item:answer" }', + ) + out = [] + for one in text: + fc = feature_pb2.FeatureConfig() + text_format.Merge(one, fc) + out.append(create_features([fc], fg_mode=FgMode.FG_NONE)[0]) + return out + + +def _offset(codes): + """Shift local codes into the flat space, as the SID tool's column does.""" + offsets = np.cumsum([0] + _CODEBOOK[:-1]) + return (np.asarray(codes).reshape(-1, len(_CODEBOOK)) + offsets).reshape(-1) + + +class PromptStackIntegrationTest(unittest.TestCase): + """compile -> assemble -> model, on the real code path.""" + + def setUp(self) -> None: + self.test_dir = make_test_dir() + self.backbone = _tiny_backbone(os.path.join(self.test_dir, "backbone")) + self.tok = _tokenizer(os.path.join(self.test_dir, "tok.json")) + self.features = _features() + + cfg = PromptConfig( + tokenizer=self.tok, + prompt="History : {{hist}} . Predict :", + response="{{answer}}", + ) + cfg.sid_space.codebook.extend(_CODEBOOK) + self.prompt = compile_prompt(cfg, self.features, model_dir=self.test_dir) + + def _model(self): + model_config = ModelConfig() + qwen = model_config.prompt_generative_qwen + qwen.hf_model_id = self.backbone + qwen.common.beam_widths.extend([2, 2, 2]) + qwen.common.num_return_sequences = 2 + return _create_model( + model_config, self.features, ["answer"], prompt=self.prompt + ) + + def _batch(self, hist, answer): + parsed = { + "hist.values": torch.tensor(_offset(hist)), + "hist.lengths": torch.tensor([len(hist)]), + "answer.values": torch.tensor(_offset(answer)), + "answer.lengths": torch.tensor([len(answer)]), + } + streams = assemble_into(self.prompt, parsed) + batch = Batch() + batch.additional_infos.update( + {k: torch.from_numpy(np.asarray(v)) for k, v in streams.items()} + ) + return batch + + def test_compiles_a_usable_space(self) -> None: + space = self.prompt.sid_space + self.assertEqual(space.num_levels, 3) + self.assertEqual(space.sid_vocab_size, 12) + # the atoms sit immediately above the base vocabulary + self.assertEqual(space.band_lo[0], space.base_vocab) + self.assertEqual(space.band_hi[-1], space.base_vocab + 11) + + def test_model_resizes_to_target_vocab(self) -> None: + model = self._model() + rows = model.lm.get_input_embeddings().weight.shape[0] + self.assertEqual(rows, self.prompt.sid_space.target_vocab) + # every SID atom has a row + self.assertGreater(rows, self.prompt.sid_space.band_hi[-1]) + + def test_forward_produces_a_finite_loss(self) -> None: + model = self._model() + batch = self._batch([0, 1, 2, 3, 0, 1], [1, 2, 3]) + out = model.predict(batch) + + self.assertIn("loss", out) + self.assertTrue(bool(torch.isfinite(out["loss"]))) + + def test_loss_backpropagates_into_the_backbone(self) -> None: + model = self._model() + batch = self._batch([0, 1, 2, 3, 0, 1], [1, 2, 3]) + model.predict(batch)["loss"].backward() + + grad = model.lm.get_input_embeddings().weight.grad + self.assertIsNotNone(grad) + self.assertTrue(bool((grad.abs().sum() > 0))) + + def test_assembled_stream_matches_the_template(self) -> None: + batch = self._batch([0, 1, 2], [1, 2, 3]) + ids = batch.additional_infos["prompt_input_ids"] + space = self.prompt.sid_space + # "History :" + 3 history atoms + "." + "Predict :" + 3 answer atoms + self.assertEqual(ids.numel(), 2 + 3 + 1 + 2 + 3) + sid_rows = ids[ids >= space.base_vocab] + self.assertEqual(sid_rows.numel(), 6) + + def test_labels_supervise_only_the_answer(self) -> None: + batch = self._batch([0, 1, 2], [1, 2, 3]) + labels = batch.additional_infos["prompt_labels"] + supervised = labels[labels != -100] + self.assertEqual(supervised.numel(), 3) + + def test_raw_codes_are_rejected_before_the_model_sees_them(self) -> None: + parsed = { + # not offset: level 1 and 2 fall below their bands + "hist.values": torch.tensor([1, 2, 3]), + "hist.lengths": torch.tensor([3]), + "answer.values": torch.tensor(_offset([1, 2, 3])), + "answer.lengths": torch.tensor([3]), + } + with self.assertRaisesRegex(ValueError, "offset_codebook column"): + assemble_into(self.prompt, parsed) + + +if __name__ == "__main__": + unittest.main() From 336cf7cd85c6dcf5593861e47a8b19fc3d9ad8fd Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 07:51:30 +0000 Subject: [PATCH 65/99] [refactor] prompt: remove the pre-prompt generative stack Deletes SidFeature, BaseGenerativeModel, GenerativeQwen, their proto messages, tests and mock config. The prompt-native stack replaces all of it: the per-level offset now arrives in the data, so no feature type is needed to apply it, and the template, SID space and tokenizer live in prompt_config rather than in the model. Nothing is kept for compatibility. The design this implements is design_v2.md. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/features/sid_feature.py | 180 ------- tzrec/features/sid_feature_test.py | 139 ------ tzrec/models/generative_model.py | 461 ------------------ tzrec/models/generative_model_test.py | 445 ----------------- tzrec/models/generative_qwen.py | 270 ---------- tzrec/models/generative_qwen_test.py | 420 ---------------- tzrec/modules/dynamic_beam_test.py | 2 +- tzrec/protos/feature.proto | 45 -- tzrec/protos/model.proto | 2 - tzrec/protos/models/generative_model.proto | 48 -- .../tests/configs/generative_qwen_mock.config | 65 --- tzrec/tests/genrec_integration_test.py | 186 ------- tzrec/utils/hf_export_util_test.py | 2 +- 13 files changed, 2 insertions(+), 2263 deletions(-) delete mode 100644 tzrec/features/sid_feature.py delete mode 100644 tzrec/features/sid_feature_test.py delete mode 100644 tzrec/models/generative_model.py delete mode 100644 tzrec/models/generative_model_test.py delete mode 100644 tzrec/models/generative_qwen.py delete mode 100644 tzrec/models/generative_qwen_test.py delete mode 100644 tzrec/protos/models/generative_model.proto delete mode 100644 tzrec/tests/configs/generative_qwen_mock.config delete mode 100644 tzrec/tests/genrec_integration_test.py diff --git a/tzrec/features/sid_feature.py b/tzrec/features/sid_feature.py deleted file mode 100644 index 240fdb4ce..000000000 --- a/tzrec/features/sid_feature.py +++ /dev/null @@ -1,180 +0,0 @@ -# Copyright (c) 2026, Alibaba Group; -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# http://www.apache.org/licenses/LICENSE-2.0 -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from typing import Any, Dict, List, Optional, Tuple - -import numpy as np -import pyarrow as pa - -from tzrec.datasets.utils import ParsedData -from tzrec.features.feature import BaseFeature -from tzrec.protos.feature_pb2 import FeatureConfig - - -class SidFeature(BaseFeature): - """Semantic-ID sequence feature. - - A flat stream of 0-based per-level SID codes -- whole items in level order -- - plus the prompt text wrapping them. Under fg it is a passthrough. - - Args: - feature_config (FeatureConfig): a instance of feature config. - """ - - def __init__( - self, - feature_config: FeatureConfig, - **kwargs: Any, - ) -> None: - # BaseFeature.__del__ dereferences _fg_op, so seed it before any raise. - self._fg_op = None - super().__init__(feature_config, **kwargs) - if self.config.value_dim != 1: - raise ValueError( - f"{self.__class__.__name__}[{self.config.feature_name}]: " - f"value_dim must be 1 -- the SID stream is flat, one code per " - f"sequence position -- got {self.config.value_dim}." - ) - self._codebook = self._read_codebook() - # fg truncates by VALUE count and keeps the head, so a cap that is not a - # whole number of items would hand the model partial items. The model's - # own max_sequence_length is item-aligned and keeps the recent tail. - if self.config.HasField("sequence_length"): - if self.config.sequence_length % len(self._codebook): - raise ValueError( - f"{self.__class__.__name__}[{self.config.feature_name}]: " - f"sequence_length ({self.config.sequence_length}) must be a " - f"multiple of the {len(self._codebook)}-level codebook, or " - f"fg would cut an item in half. Prefer " - f"model_config.common.max_sequence_length, which is " - f"item-aligned and keeps the most RECENT items." - ) - self._level_sizes = np.asarray(self._codebook) - self._level_offsets = np.cumsum(self._level_sizes) - self._level_sizes - - def _read_codebook(self) -> List[int]: - """Validate the declared codebook once and normalize it to a list.""" - codebook = [int(c) for c in self.config.codebook] - if not codebook: - raise ValueError( - f"{self.__class__.__name__}[{self.config.feature_name}]: codebook " - f"is required; give one vocabulary size per SID level." - ) - if any(c <= 0 for c in codebook): - raise ValueError( - f"{self.__class__.__name__}[{self.config.feature_name}]: every " - f"codebook size must be positive, got {codebook}." - ) - return codebook - - @property - def value_dim(self) -> int: - """Fg value dimension of the feature.""" - return self.config.value_dim - - @property - def output_dim(self) -> int: - """Output dimension: SID codes pass through to the LM's own table.""" - return self.value_dim - - @property - def num_embeddings(self) -> int: - """Get embedding row count.""" - raise RuntimeError( - f"{self.__class__.__name__}[{self.config.feature_name}] has no " - f"embedding table; SID codes index the LM vocabulary." - ) - - @property - def prefix_text(self) -> str: - """Text emitted immediately before this feature's SID tokens.""" - return self.config.prefix_text - - @property - def suffix_text(self) -> str: - """Text emitted immediately after this feature's SID tokens.""" - return self.config.suffix_text - - @property - def codebook(self) -> List[int]: - """Per-level SID vocabulary sizes; validated once at construction.""" - return self._codebook - - @property - def num_levels(self) -> int: - """Codes per item -- also the answer width.""" - return len(self._codebook) - - @property - def sid_vocab_size(self) -> int: - """Atoms the model must append to the backbone vocabulary.""" - return sum(self._codebook) - - @property - def level_offsets(self) -> List[int]: - """Flat offset of each level, i.e. ``cumsum(sizes) - sizes``.""" - return self._level_offsets.tolist() - - def _build_side_inputs(self) -> Optional[List[Tuple[str, str]]]: - """Input field names with side.""" - if self.config.HasField("expression"): - return [tuple(self.config.expression.split(":"))] - else: - return None - - def _parse(self, input_data: Dict[str, pa.Array]) -> ParsedData: - """Parse the SID stream into flat indices in the shared space. - - Offsets are folded in here, in the dataloader workers: validating on the - forward path would let one rank raise and hang its peers on the - collective. - """ - parsed = super()._parse(input_data) - num_levels = len(self._codebook) - bad = np.nonzero(parsed.seq_lengths % num_levels)[0] - if bad.size: - raise ValueError( - f"{self.__class__.__name__}[{self.config.feature_name}]: every " - f"row must hold whole {num_levels}-level items; rows " - f"{bad.tolist()[:10]} have lengths " - f"{parsed.seq_lengths[bad].tolist()[:10]}." - ) - # rows are whole items, so each column is one level. - codes = parsed.values.reshape(-1, num_levels) - if ((codes < 0) | (codes >= self._level_sizes)).any(): - raise ValueError( - f"{self.__class__.__name__}[{self.config.feature_name}]: SID " - f"codes must be local 0-based values in [0, codebook[level])." - ) - # keep the dtype: int64 offsets would promote float32 to float64. - offsets = self._level_offsets.astype(codes.dtype, copy=False) - parsed.values = (codes + offsets).reshape(parsed.values.shape) - return parsed - - def _fg_json(self) -> List[Dict[str, Any]]: - """Get fg json config impl. - - A PASSTHROUGH: no fg feature_type can add ``level_offsets[i % levels]``, - so fg only reaches the codes and ``_parse`` does the arithmetic. It - exists because ``fg_mode`` is a data_config-level switch -- refusing it - here would block every other feature in the config. - """ - # SCALAR form: fg_json prepends "sequence_" and injects the seq keys. - fg_cfg: Dict[str, Any] = { - "feature_type": "raw_feature", - "feature_name": self.config.feature_name, - "expression": self.config.expression, - "default_value": self.config.default_value, - "value_type": "float", - } - if self.config.HasField("stub_type"): - fg_cfg["stub_type"] = self.config.stub_type - return [fg_cfg] diff --git a/tzrec/features/sid_feature_test.py b/tzrec/features/sid_feature_test.py deleted file mode 100644 index e5fda20b7..000000000 --- a/tzrec/features/sid_feature_test.py +++ /dev/null @@ -1,139 +0,0 @@ -# Copyright (c) 2026, Alibaba Group; -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# http://www.apache.org/licenses/LICENSE-2.0 -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import unittest - -import pyarrow as pa -from google.protobuf import text_format -from parameterized import parameterized - -from tzrec.features.feature import FgMode, create_features -from tzrec.protos import feature_pb2 -from tzrec.utils.test_util import parameterized_name_func - - -def _feature(text, fg_mode=FgMode.FG_NONE): - fc = feature_pb2.FeatureConfig() - text_format.Merge(f"sequence_sid_feature {{ {text} }}", fc) - return create_features([fc], fg_mode=fg_mode)[0] - - -_BASE = ( - 'feature_name: "user_sequence" expression: "user:user_sequence" ' - "codebook: 4 codebook: 4 codebook: 4" -) - - -class SidFeatureTest(unittest.TestCase): - def test_dispatch_and_defaults(self) -> None: - f = _feature(_BASE) - self.assertEqual(type(f).__name__, "SidFeature") - # the oneof FIELD name is what makes it a sequence, not the message - self.assertTrue(f.is_sequence) - self.assertFalse(f.is_sparse) - self.assertEqual(f.name, "user_sequence") - self.assertEqual(f.value_dim, 1) - self.assertEqual(f.output_dim, 1) - self.assertEqual(f.side_inputs, [("user", "user_sequence")]) - self.assertEqual(f.prefix_text, "") - self.assertEqual(f.suffix_text, "") - self.assertEqual(f.codebook, [4, 4, 4]) - self.assertEqual(f.num_levels, 3) - self.assertEqual(f.sid_vocab_size, 12) - self.assertEqual(f.level_offsets, [0, 4, 8]) - - def test_prompt_text_round_trips(self) -> None: - f = _feature(f'{_BASE} prefix_text: "History: " suffix_text: "."') - self.assertEqual(f.prefix_text, "History: ") - self.assertEqual(f.suffix_text, ".") - - @parameterized.expand( - [[FgMode.FG_NORMAL], [FgMode.FG_DAG], [FgMode.FG_BUCKETIZE]], - name_func=parameterized_name_func, - ) - def test_builds_under_every_fg_mode(self, fg_mode) -> None: - # fg_mode is a data_config-level switch, so refusing it here would block - # every OTHER feature in the config from using fg. - self.assertEqual(type(_feature(_BASE, fg_mode=fg_mode)).__name__, "SidFeature") - - def test_fg_passthrough_matches_the_fg_none_parse(self) -> None: - """Fg only reaches the codes; _parse folds the offsets either way.""" - rows = [[1, 2, 3, 0, 1, 2], [3, 0, 1]] - want = [1, 6, 11, 0, 5, 10, 3, 4, 9] # code + offsets [0,4,8] - none = _feature(_BASE).parse({"user_sequence": pa.array(rows)}) - # under fg the same sequence arrives delimited, as ODPS/CSV deliver it - fg = _feature(_BASE, fg_mode=FgMode.FG_NORMAL).parse( - {"user_sequence": pa.array([";".join(map(str, r)) for r in rows])} - ) - self.assertEqual(none.values.flatten().astype(int).tolist(), want) - self.assertEqual(fg.values.flatten().astype(int).tolist(), want) - self.assertEqual(none.seq_lengths.tolist(), fg.seq_lengths.tolist()) - - def test_fg_json_is_a_passthrough_raw_feature(self) -> None: - cfg = _feature(_BASE).fg_json() - self.assertEqual(len(cfg), 1) - # the base wrapper prepends "sequence_"; no bucketizer, no normalizer - self.assertEqual(cfg[0]["feature_type"], "sequence_raw_feature") - self.assertEqual(cfg[0]["expression"], "user:user_sequence") - for k in ("boundaries", "normalizer", "vocab_file", "hash_bucket_size"): - self.assertNotIn(k, cfg[0]) - - def test_rejects_a_sequence_length_that_splits_an_item(self) -> None: - # fg truncates by VALUE count, so a non-multiple would hand the model - # a partial item; _parse would then reject the whole batch. - with self.assertRaisesRegex(ValueError, "multiple of the 3-level"): - _feature(f"{_BASE} sequence_length: 10") - self.assertEqual( - _feature(f"{_BASE} sequence_length: 9").config.sequence_length, 9 - ) - - def test_parse_folds_in_the_level_offsets(self) -> None: - # offsets [0, 4, 8]: level j's 0-based code k becomes flat index k + off[j], - # which is also the atom index -- no bridging shift anywhere. - f = _feature(_BASE) - parsed = f.parse({"user_sequence": pa.array([[0, 1, 2, 1, 2, 3], [0, 0, 0]])}) - self.assertEqual( - parsed.values.flatten().tolist(), [0, 5, 10, 1, 6, 11, 0, 4, 8] - ) - self.assertEqual(parsed.seq_lengths.tolist(), [6, 3]) - - def test_parse_rejects_out_of_range_and_partial_items(self) -> None: - f = _feature(_BASE) - with self.assertRaisesRegex(ValueError, "local 0-based"): - f.parse({"user_sequence": pa.array([[0, 1, 4]])}) # 4 == codebook[2] - with self.assertRaisesRegex(ValueError, "local 0-based"): - f.parse({"user_sequence": pa.array([[-1, 1, 2]])}) - with self.assertRaisesRegex(ValueError, "whole 3-level items"): - f.parse({"user_sequence": pa.array([[0, 1]])}) - - def test_rejects_a_bad_config(self) -> None: - for bad, msg in ( - ("", "codebook is required"), - ("codebook: 4 codebook: 0", "positive"), - # _parse reshapes by level and splits by seq_lengths, so a wider - # value would land the offsets on the wrong components. - ("codebook: 4 value_dim: 2", "value_dim must be 1"), - ): - with self.subTest(bad=bad): - base = 'feature_name: "s" expression: "user:s" ' + bad - with self.assertRaisesRegex(ValueError, msg): - _feature(base) - - def test_no_embedding_table(self) -> None: - f = _feature(_BASE) - self.assertFalse(f.has_embedding) - self.assertIsNone(f.emb_config) - with self.assertRaisesRegex(RuntimeError, "no .*embedding table"): - _ = f.num_embeddings - - -if __name__ == "__main__": - unittest.main() diff --git a/tzrec/models/generative_model.py b/tzrec/models/generative_model.py deleted file mode 100644 index c15a0de1e..000000000 --- a/tzrec/models/generative_model.py +++ /dev/null @@ -1,461 +0,0 @@ -# Copyright (c) 2026, Alibaba Group; -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# http://www.apache.org/licenses/LICENSE-2.0 -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Architecture-agnostic base for HF-backed generative-recommendation LMs. - -A family subclass (e.g. ``GenerativeQwen``) supplies ``_build_prompt_tokens`` and -``predict``; ``GenerativeModelConfig`` holds the shared config and the sample -contract. -""" - -import re -from typing import Any, Dict, List, Optional, Tuple - -import torch -import torch.nn.functional as F -import torchmetrics -from transformers import ( - AutoConfig, - AutoModelForCausalLM, - AutoTokenizer, - PreTrainedModel, - PreTrainedTokenizerBase, -) - -from tzrec.datasets.utils import Batch -from tzrec.features.feature import BaseFeature -from tzrec.features.sid_feature import SidFeature -from tzrec.models.model import BaseModel -from tzrec.modules.embedding import EmbeddingGroup -from tzrec.protos import model_pb2 -from tzrec.protos.model_pb2 import ModelConfig -from tzrec.protos.models import generative_model_pb2 - - -class BaseGenerativeModel(BaseModel): - """Model construction, SID vocab extension, data-prep, loss and metrics.""" - - _PARAM_DTYPE: Dict[int, torch.dtype] = { - generative_model_pb2.FP32: torch.float32, - generative_model_pb2.BF16: torch.bfloat16, - generative_model_pb2.FP16: torch.float16, - } - - def __init__( - self, - model_config: ModelConfig, - features: List[BaseFeature], - labels: List[str], - sample_weights: Optional[List[str]] = None, - **kwargs: Any, - ) -> None: - super().__init__(model_config, features, labels, sample_weights, **kwargs) - cfg = self._model_config - sid_atoms = self._read_common_config(cfg.common) - - self.lm = self._build_backbone() - tokenizer, base = self._build_extended_tokenizer(sid_atoms) - self._hf_tokenizer = tokenizer - self._base_vocab = base - self._pad_token_id = self._resolve_pad_token_id(tokenizer) - - self._build_prompt_tokens(tokenizer, cfg) - self.init_input() - - @staticmethod - def _resolve_pad_token_id(tokenizer: PreTrainedTokenizerBase) -> int: - """Pad id for the left-padded splice, falling back to eos.""" - pad_id = tokenizer.pad_token_id - if pad_id is None: - pad_id = tokenizer.eos_token_id - if pad_id is None: - raise ValueError( - "BaseGenerativeModel: tokenizer has neither pad_token_id nor " - "eos_token_id; cannot choose a pad id for the left-padded splice." - ) - return int(pad_id) - - def _read_common_config( - self, common: generative_model_pb2.GenerativeModelConfig - ) -> int: - """Parse shared proto knobs into attributes; return the SID atom count.""" - self._label_name: str = self._labels[0] if self._labels else "" - self._ignore_index: int = int(common.ignore_index) - self._generated_sids_key: str = common.generated_sids_key - self._param_dtype: torch.dtype = self._PARAM_DTYPE[common.param_dtype] - self._max_seq_length: int = int(common.max_sequence_length) - codebook = self._shared_sid_space() - self._num_levels = len(codebook) - if 0 < self._max_seq_length < self._num_levels: - raise ValueError( - f"{type(self).__name__}: max_sequence_length " - f"({self._max_seq_length}) cannot hold one {self._num_levels}" - f"-level item; use 0 to disable the budget or a multiple of " - f"{self._num_levels}." - ) - sizes = torch.tensor(codebook, dtype=torch.long) - self.register_buffer("_codebook_sizes", sizes, persistent=False) - # only the decode path needs these; SidFeature folds them into inputs. - self.register_buffer( - "_level_offsets", torch.cumsum(sizes, 0) - sizes, persistent=False - ) - self._vocab_pad_mult = int(common.vocab_pad_to_multiple_of) - self._read_beam_config(common) - return sum(codebook) - - def _read_beam_config( - self, common: generative_model_pb2.GenerativeModelConfig - ) -> None: - """Parse the decode knobs; the width schedule must match the codebook.""" - self._num_return = int(common.num_return_sequences) - self._beam_widths: List[int] = list(common.beam_widths) - if not self._beam_widths: - raise ValueError( - f"{type(self).__name__}: beam_widths is required; give one " - f"width per SID level, e.g. [50, 50, 50] or [100, 200, 400]." - ) - if len(self._beam_widths) != self._num_levels: - raise ValueError( - f"{type(self).__name__}: beam_widths has " - f"{len(self._beam_widths)} entries but the codebook has " - f"{self._num_levels} levels; give one width per level." - ) - if self._num_return > self._beam_widths[-1]: - raise ValueError( - f"{type(self).__name__}: num_return_sequences " - f"({self._num_return}) must not exceed the final beam width " - f"({self._beam_widths[-1]})." - ) - - def _shared_sid_space(self) -> List[int]: - """The one codebook every SID feature declares.""" - spaces = { - f.name: tuple(f.codebook) - for f in self._features - if isinstance(f, SidFeature) - } - if not spaces: - raise ValueError( - f"{type(self).__name__}: no SID feature declares a codebook; " - f"genrec needs at least one sequence_sid_feature." - ) - if len(set(spaces.values())) != 1: - raise ValueError( - f"{type(self).__name__}: all SID features must share one " - f"codebook, got {dict(sorted(spaces.items()))}." - ) - return list(next(iter(spaces.values()))) - - def _slot_group_names(self) -> Dict[str, str]: - """{feature_name: group_name} for every declared feature_group. - - One JAGGED_SEQUENCE feature per group: EmbeddingGroup interleaves a - group's members into one sequence that cannot be split back apart. - """ - by_feature: Dict[str, str] = {} - for group in self._feature_groups: - if group.group_type != model_pb2.JAGGED_SEQUENCE: - raise ValueError( - f"{type(self).__name__}: feature_group {group.group_name!r} " - f"must be JAGGED_SEQUENCE, got " - f"{model_pb2.FeatureGroupType.Name(group.group_type)}." - ) - if len(group.feature_names) != 1: - raise ValueError( - f"{type(self).__name__}: feature_group {group.group_name!r} " - f"must hold exactly one feature (its members are interleaved " - f"into one sequence), got {list(group.feature_names)}." - ) - name = group.feature_names[0] - if name in by_feature: - raise ValueError( - f"{type(self).__name__}: feature {name!r} is claimed by both " - f"feature_group {by_feature[name]!r} and " - f"{group.group_name!r}; only one can fill its prompt slot." - ) - by_feature[name] = group.group_name - return by_feature - - def _resolve_prompt_slots( - self, template: str - ) -> Tuple[List[str], List["SidFeature"]]: - """Split a ``{{feature_name}}`` template into N+1 gaps and N features. - - Records ``_slot_names`` / ``_slot_groups`` here rather than in the family - hook, because ``build_input`` reads them. - """ - parts = re.split(r"\{\{(\w+)\}\}", template) - gaps, names = parts[0::2], parts[1::2] - if not names: - raise ValueError( - f"{type(self).__name__}: prompt_template needs at least one " - f"{{{{feature_name}}}} slot naming a declared feature." - ) - group_of = self._slot_group_names() - by_name = {f.name: f for f in self._features} - features = [] - for name in names: - if name not in group_of: - raise ValueError( - f"{type(self).__name__}: prompt_template slot {{{{{name}}}}} " - f"names no feature_group; declared: {sorted(group_of)}." - ) - feature = by_name.get(name) - if feature is None: - raise ValueError( - f"{type(self).__name__}: prompt_template slot {{{{{name}}}}} " - f"has no feature_config." - ) - if not isinstance(feature, SidFeature): - raise ValueError( - f"{type(self).__name__}: prompt slot {{{{{name}}}}} names a " - f"{type(feature).__name__}; only a SidFeature can fill a slot." - ) - features.append(feature) - unused = sorted(set(group_of) - set(names)) - if unused: - raise ValueError( - f"{type(self).__name__}: feature_group(s) {unused} are declared " - f"but never referenced by a prompt_template slot." - ) - self._slot_names = names - self._slot_groups = [group_of[n] for n in names] - return gaps, features - - def _build_backbone(self) -> PreTrainedModel: - """Build the EMPTY architecture; weights arrive later from HF or DCP.""" - hf_model_id = self._model_config.hf_model_id - if not hf_model_id: - raise ValueError(f"{type(self).__name__}: empty hf_model_id.") - hf_cfg = AutoConfig.from_pretrained(hf_model_id) - lm = AutoModelForCausalLM.from_config(hf_cfg, torch_dtype=self._param_dtype) - # no-op when from_config already honoured torch_dtype; not all do. - return lm.to(self._param_dtype) - - def _build_extended_tokenizer( - self, sid_atoms: int - ) -> Tuple[PreTrainedTokenizerBase, int]: - """Add the SID atoms ``C0..C{sid_atoms-1}`` and resize ``self.lm``. - - Returns ``(tokenizer, base)`` where ``base`` is the tokenizer's next free - id BEFORE adding the atoms -- use ``len(tokenizer)``, NOT - ``config.vocab_size`` (which counts reserved slots). - """ - tokenizer = AutoTokenizer.from_pretrained( - self._model_config.hf_model_id, use_fast=True - ) - base = len(tokenizer) - added = tokenizer.add_tokens([f"C{i}" for i in range(sid_atoms)]) - if added != sid_atoms: - raise RuntimeError( - f"BaseGenerativeModel: tokenizer was expected to grow by " - f"{sid_atoms} new atoms, only added {added}. " - f"Aborting to avoid silent SID-token mismatch." - ) - # stash the resize target so init_from_pretrained re-extends identically. - self._target_vocab = base + sid_atoms - self.lm.resize_token_embeddings( - self._target_vocab, pad_to_multiple_of=self._vocab_pad_mult or None - ) - c0_id = tokenizer.convert_tokens_to_ids("C0") - if c0_id != base: - raise RuntimeError( - f"BaseGenerativeModel: SID atom layout mismatch -- expected " - f"C0 at token id {base}, got {c0_id}. " - f"Splice arithmetic would produce wrong token ids." - ) - return tokenizer, base - - def init_from_pretrained(self) -> None: - """Load the pretrained HF weights and re-extend to ``__init__``'s vocab. - - The new SID rows differ per rank; DDP's ``_sync_module_states`` - broadcast from rank 0 reconciles them. - """ - # drop the empty arch first: holding both peaks at 2x model host RAM. - self.lm = None - lm = AutoModelForCausalLM.from_pretrained( - self._model_config.hf_model_id, - torch_dtype=self._param_dtype, - low_cpu_mem_usage=True, - ) - lm.resize_token_embeddings( - self._target_vocab, pad_to_multiple_of=self._vocab_pad_mult or None - ) - self.lm = lm - - def hf_backbone(self) -> PreTrainedModel: - """The HF backbone module, for checkpoint/export asset writing.""" - return self.lm - - def hf_tokenizer(self) -> PreTrainedTokenizerBase: - """The extended tokenizer (base vocab + C0..C{sum-1}) to serialize.""" - return self._hf_tokenizer - - def _build_prompt_tokens( - self, tokenizer: PreTrainedTokenizerBase, cfg: Any - ) -> None: - """Family hook: cache the tokenised prompt template as buffers. - - Called from ``__init__`` after vocab extension; consumed by ``predict``. - """ - raise NotImplementedError( - f"{type(self).__name__} must implement _build_prompt_tokens " - f"(BaseGenerativeModel is abstract)." - ) - - @property - def device(self) -> torch.device: - """Device the HF backbone runs on.""" - return self.lm.device - - def _tokenize_sids(self, flat: torch.Tensor) -> torch.Tensor: - """Map flat SID indices to extended-vocab token ids.""" - return flat + self._base_vocab - - def _detokenize_sids( - self, tokens: torch.Tensor, level_ids: torch.Tensor - ) -> torch.Tensor: - """Inverse of ``_tokenize_sids``: token ids to local 0-based codes.""" - return tokens - self._base_vocab - self._level_offsets[level_ids] - - def _sid_token_bands(self) -> Tuple[torch.Tensor, torch.Tensor]: - """Return the inclusive token-id band for every SID level.""" - return ( - self._tokenize_sids(self._level_offsets), - self._tokenize_sids(self._level_offsets + self._codebook_sizes - 1), - ) - - def _validate_sid_candidates( - self, new_tokens: torch.Tensor, batch_size: int - ) -> torch.Tensor: - """Decode the per-beam tail ``(B*C, w)`` to ``(B, C, num_levels)`` codes. - - Any malformed candidate becomes all ``-1``, which no real code matches. - """ - level_ids = torch.arange(new_tokens.shape[1], device=new_tokens.device) - codes = self._detokenize_sids(new_tokens, level_ids) - codes = F.pad(codes, (0, self._num_levels - codes.shape[1]), value=-1) - invalid = ((codes < 0) | (codes >= self._codebook_sizes)).any(dim=1) - codes = codes.masked_fill(invalid.unsqueeze(1), -1) - # decoders return rows batch-major; group per user. - return codes.view(batch_size, -1, self._num_levels) - - def init_input(self) -> None: - """Build the EmbeddingGroup; passthrough features hold no params.""" - self.embedding_group = EmbeddingGroup(self._features, self._feature_groups) - - def build_input(self, batch: Batch) -> Dict[str, List[torch.Tensor]]: - """Retrieve per-row SID token sequences, keyed by feature name. - - The answer is a label_field so it can be absent at inference. - """ - g = self.embedding_group(batch) - rows: Dict[str, List[torch.Tensor]] = { - name: self._sid_token_rows( - g[f"{group}.sequence"], - g[f"{group}.sequence_length"], - max_codes=self._max_seq_length, - ) - for name, group in zip(self._slot_names, self._slot_groups) - } - if not self.is_inference: - if not self._label_name: - raise ValueError( - f"{type(self).__name__}: training needs the answer SIDs; " - f"declare it as the first data_config.label_field." - ) - jt = batch.jagged_labels[self._label_name] - rows[self._label_name] = self._answer_token_rows(jt.values(), jt.lengths()) - return rows - - def _sid_token_rows( - self, - values: torch.Tensor, - lengths: torch.Tensor, - max_codes: Optional[int] = None, - ) -> List[torch.Tensor]: - """Map a feature's flat SID stream to per-row token-id tensors. - - ``max_codes`` caps each row to its most-recent WHOLE items. - """ - values = values.reshape(-1) # value_dim 1 arrives as (N,) or (N, 1) - sizes = lengths.long().tolist() - # TODO(shuqi): move truncation into FG once FG can keep the TAIL, not the HEAD. - if max_codes: - keep = (max_codes // self._num_levels) * self._num_levels - if keep and any(n > keep for n in sizes): - rows = torch.split(values, sizes) - values = torch.cat([r[-keep:] for r in rows]) - sizes = [min(n, keep) for n in sizes] - tokens = self._tokenize_sids(values.to(self.device).long()) - return list(torch.split(tokens, sizes)) - - def _answer_token_rows( - self, values: torch.Tensor, lengths: torch.Tensor - ) -> List[torch.Tensor]: - """Map the answer label to token ids; every row is ``num_levels`` codes. - - A label_field is not offset by SidFeature, so the fold-in happens here. - """ - values = values.reshape(-1) # value_dim 1 arrives as (N,) or (N, 1) - sizes = lengths.long().tolist() - bad = [i for i, n in enumerate(sizes) if n != self._num_levels] - if bad: - raise ValueError( - f"{type(self).__name__}: each answer must be " - f"{self._num_levels} codes (len(codebook)); rows {bad} have " - f"{[sizes[i] for i in bad]} -- anomalous sample(s)." - ) - codes = values.to(self.device).long() - level_ids = torch.arange(codes.numel(), device=self.device) % self._num_levels - invalid = (codes < 0) | (codes >= self._codebook_sizes[level_ids]) - if invalid.any(): - raise ValueError( - f"{type(self).__name__}: answer SID codes must be local 0-based " - f"values in [0, codebook[level])." - ) - tokens = self._tokenize_sids(codes + self._level_offsets[level_ids]) - return list(torch.split(tokens, sizes)) - - def init_loss(self) -> None: - """No-op: the loss is computed inside ``predict`` (HF loss_function).""" - return - - def loss( - self, - predictions: Dict[str, torch.Tensor], - batch: Batch, - ) -> Dict[str, torch.Tensor]: - """Surface the CE loss already computed in ``predict``.""" - return {"ce_loss": predictions["loss"]} - - def init_metric(self) -> None: - """Register a mean-CE metric for the eval loop.""" - self._metric_modules["ce_loss"] = torchmetrics.MeanMetric() - - def update_metric( - self, - predictions: Dict[str, torch.Tensor], - batch: Batch, - losses: Optional[Dict[str, torch.Tensor]] = None, - ) -> None: - """Update the mean-CE metric with this batch's loss.""" - self._metric_modules["ce_loss"].update(predictions["loss"].detach()) - - # NOTE: BaseModel declares no such hook, but the train loop calls it. - def update_train_metric( - self, - predictions: Dict[str, torch.Tensor], - batch: Batch, - ) -> None: - """No-op: no train-time metric beyond the logged CE loss.""" - return diff --git a/tzrec/models/generative_model_test.py b/tzrec/models/generative_model_test.py deleted file mode 100644 index e95221b27..000000000 --- a/tzrec/models/generative_model_test.py +++ /dev/null @@ -1,445 +0,0 @@ -# Copyright (c) 2026, Alibaba Group; -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# http://www.apache.org/licenses/LICENSE-2.0 -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import types -import unittest - -import torch -from google.protobuf import text_format -from torch import nn - -from tzrec.features.feature import create_features -from tzrec.models.generative_model import BaseGenerativeModel -from tzrec.models.generative_qwen import GenerativeQwen -from tzrec.models.model import BaseModel -from tzrec.protos import feature_pb2, model_pb2 -from tzrec.protos.models import generative_model_pb2 - - -class _FakeJT: - """Minimal stand-in for a TorchRec JaggedTensor.""" - - def __init__(self, values, lengths, dim2=False): - v = torch.tensor(values, dtype=torch.float) # TER delivers list as float - self._v = v.unsqueeze(-1) if dim2 else v - self._l = torch.tensor(lengths) - - def values(self): - return self._v - - def lengths(self): - return self._l - - -def _sid_feature(name="user_sequence", codebook=(2, 3, 4), prefix_text=""): - """A real SidFeature -- the model dispatches on the type, not on duck-typing.""" - fc = feature_pb2.FeatureConfig() - text_format.Merge( - f'sequence_sid_feature {{ feature_name: "{name}" expression: "user:{name}" ' - + " ".join(f"codebook: {c}" for c in codebook) - + f' prefix_text: "{prefix_text}" }}', - fc, - ) - return create_features([fc])[0] - - -def _common(**overrides): - """A fake ``GenerativeModelConfig`` -- only the fields the base actually reads.""" - fields = { - "ignore_index": -100, - "generated_sids_key": "generated_sids", - "param_dtype": generative_model_pb2.FP32, - "vocab_pad_to_multiple_of": 128, - "max_sequence_length": 0, - "beam_widths": [50, 50, 50], - "num_return_sequences": 50, - } - return types.SimpleNamespace(**{**fields, **overrides}) - - -def _wired(features=None, group_type=model_pb2.JAGGED_SEQUENCE, members=None): - """Pre-``__init__`` state: the features/labels/groups the config-time code reads.""" - m = object.__new__(BaseGenerativeModel) - nn.Module.__init__(m) - m._features = [_sid_feature()] if features is None else features - m._labels = ["label"] - m._feature_groups = [ - types.SimpleNamespace( - group_name="user_seq", - feature_names=["user_sequence"] if members is None else list(members), - group_type=group_type, - ) - ] - return m - - -def _stub(codebook=None, base_vocab=100, device="cpu"): - """Base model with the data-prep state wired up, but no HF backbone.""" - codebook = codebook or [2, 3, 4] - m = object.__new__(BaseGenerativeModel) - nn.Module.__init__(m) - m._base_vocab = base_vocab - m._num_levels = len(codebook) - m.lm = types.SimpleNamespace(device=torch.device(device)) - sizes = torch.tensor(codebook, dtype=torch.long) - m.register_buffer("_codebook_sizes", sizes, persistent=False) - m.register_buffer( - "_level_offsets", torch.cumsum(sizes, 0) - sizes, persistent=False - ) - return m - - -class BaseGenerativeModelTest(unittest.TestCase): - def test_registry_dispatch(self) -> None: - self.assertIs(BaseModel.create_class("GenerativeQwen"), GenerativeQwen) - self.assertTrue(issubclass(GenerativeQwen, BaseGenerativeModel)) - - def test_model_config_oneof_resolves_to_the_class(self) -> None: - # the path _create_model takes: oneof -> message type name -> class. - from tzrec.utils import config_util - - cfg = model_pb2.ModelConfig() - cfg.generative_qwen.common.max_sequence_length = 8 # required field - self.assertEqual(config_util.which_msg(cfg, "model"), "GenerativeQwen") - self.assertIs( - BaseModel.create_class(config_util.which_msg(cfg, "model")), GenerativeQwen - ) - - def test_resolve_pad_token_id(self) -> None: - tok = types.SimpleNamespace - self.assertEqual( - BaseGenerativeModel._resolve_pad_token_id( - tok(pad_token_id=5, eos_token_id=9) - ), - 5, - ) - self.assertEqual( - BaseGenerativeModel._resolve_pad_token_id( - tok(pad_token_id=None, eos_token_id=9) - ), - 9, - ) - # neither -> a clear error, not an opaque int(None) TypeError - with self.assertRaisesRegex(ValueError, "neither pad_token_id nor"): - BaseGenerativeModel._resolve_pad_token_id( - tok(pad_token_id=None, eos_token_id=None) - ) - - def test_backbone_owned_by_family_proto(self) -> None: - from tzrec.protos.models.generative_model_pb2 import ( - GenerativeModelConfig, - ) - from tzrec.protos.models.generative_model_pb2 import ( - GenerativeQwen as GenerativeQwenProto, - ) - - self.assertEqual(GenerativeQwenProto().hf_model_id, "Qwen/Qwen2.5-0.5B") - common_fields = [f.name for f in GenerativeModelConfig.DESCRIPTOR.fields] - self.assertNotIn("hf_model_id", common_fields) - - def test_configurable_knob_defaults(self) -> None: - from tzrec.protos.models.generative_model_pb2 import GenerativeModelConfig - - c = GenerativeModelConfig() - self.assertEqual(c.generated_sids_key, "generated_sids") - self.assertEqual(c.param_dtype, generative_model_pb2.FP32) - self.assertIs( - BaseGenerativeModel._PARAM_DTYPE[generative_model_pb2.FP32], - torch.float32, - ) - self.assertIs( - BaseGenerativeModel._PARAM_DTYPE[generative_model_pb2.BF16], - torch.bfloat16, - ) - - def test_read_common_config_reads_knobs(self) -> None: - m = _wired() - sid_atoms = m._read_common_config( - _common( - max_sequence_length=288, - generated_sids_key="my_sids", - param_dtype=generative_model_pb2.BF16, - ) - ) - self.assertEqual(m._label_name, "label") # from label_fields[0] - self.assertEqual(m._generated_sids_key, "my_sids") - self.assertIs(m._param_dtype, torch.bfloat16) - self.assertEqual(m._max_seq_length, 288) - self.assertEqual(sid_atoms, 9) - self.assertEqual(m._level_offsets.tolist(), [0, 2, 5]) - self.assertEqual(m._codebook_sizes.tolist(), [2, 3, 4]) - self.assertNotIn("_level_offsets", m.state_dict()) - self.assertNotIn("_codebook_sizes", m.state_dict()) - # the enum is closed: protobuf itself rejects an unlisted value - cfg = generative_model_pb2.GenerativeModelConfig() - with self.assertRaises(ValueError): - cfg.param_dtype = 99 - - def test_read_common_config_tolerates_no_feature_group(self) -> None: - # group validation is prompt-driven, so it lives in _resolve_prompt_slots - m = _wired() - m._feature_groups = [] - m._read_common_config(_common()) - self.assertEqual(m._num_levels, 3) - - def test_max_sequence_length_below_one_item_raises(self) -> None: - # a budget under num_levels floors to zero whole items, which would - # silently leave the history uncapped instead of capping it. - for cap in (1, 2): - with self.subTest(cap=cap): - with self.assertRaisesRegex(ValueError, "cannot hold one 3-level"): - _wired()._read_common_config(_common(max_sequence_length=cap)) - # 0 disables the budget; num_levels is the smallest meaningful cap - for cap in (0, 3): - with self.subTest(cap=cap): - m = _wired() - m._read_common_config(_common(max_sequence_length=cap)) - self.assertEqual(m._max_seq_length, cap) - - def test_one_feature_claimed_by_two_groups_raises(self) -> None: - # keyed by feature, so a second claim would otherwise just overwrite - m = _wired() - m._feature_groups.append( - types.SimpleNamespace( - group_name="user_seq_dup", - feature_names=["user_sequence"], - group_type=model_pb2.JAGGED_SEQUENCE, - ) - ) - with self.assertRaisesRegex(ValueError, "claimed by both"): - m._slot_group_names() - - def test_slot_group_and_shared_codebook_validation(self) -> None: - # a SEQUENCE group emits the same key with padded-dense semantics. - with self.assertRaisesRegex(ValueError, "must be JAGGED_SEQUENCE"): - _wired(group_type=model_pb2.SEQUENCE)._slot_group_names() - with self.assertRaisesRegex(ValueError, "exactly one feature"): - _wired(members=())._slot_group_names() - # a codebook the SID features disagree on is the model's business - m = _wired( - features=[ - _sid_feature("user_sequence", (2, 3)), - _sid_feature("other_seq", (4, 4)), - ] - ) - with self.assertRaisesRegex(ValueError, "must share one"): - m._shared_sid_space() - m._features = [] - with self.assertRaisesRegex(ValueError, "no SID feature declares"): - m._shared_sid_space() - - def test_vocab_pad_zero_disables_padding(self) -> None: - m = _wired() - m._read_common_config(_common(vocab_pad_to_multiple_of=0)) - self.assertEqual(m._vocab_pad_mult, 0) # not silently rewritten to 128 - - def test_max_sequence_length_model_knob(self) -> None: - m = _wired() - m._read_common_config(_common(max_sequence_length=128)) - self.assertEqual(m._max_seq_length, 128) - m2 = _wired() - m2._read_common_config(_common(max_sequence_length=0)) - self.assertEqual(m2._max_seq_length, 0) # 0 = off, no fallback - - def test_resolve_prompt_slots_splits_and_records(self) -> None: - m = _wired() - gaps, features = m._resolve_prompt_slots("A{{user_sequence}}B") - self.assertEqual(gaps, ["A", "B"]) # N slots -> N+1 gaps - self.assertEqual([f.name for f in features], ["user_sequence"]) - # recorded here, not in the family hook -- build_input reads them - self.assertEqual(m._slot_names, ["user_sequence"]) - self.assertEqual(m._slot_groups, ["user_seq"]) - - def test_resolve_prompt_slots_rejects_a_mismatched_template(self) -> None: - from tzrec.features.feature import create_features - - other = _wired( - features=[ - _sid_feature("user_sequence"), - _sid_feature("other_seq"), - ] - ) - other._feature_groups.append( - types.SimpleNamespace( - group_name="other_seq_group", - feature_names=["other_seq"], - group_type=model_pb2.JAGGED_SEQUENCE, - ) - ) - raw = feature_pb2.FeatureConfig() - text_format.Merge( - 'id_feature { feature_name: "user_sequence" expression: "user:x" ' - "num_buckets: 8 embedding_dim: 4 }", - raw, - ) - not_a_sid = _wired(features=create_features([raw])) - - for model, template, msg in ( - (_wired(), "no slot at all", "at least one"), - (_wired(), "{{nope}} x", "names no feature_group"), - (other, "{{user_sequence}}", "never referenced"), - (not_a_sid, "{{user_sequence}}", "only a SidFeature"), - ): - with self.subTest(template=template): - with self.assertRaisesRegex(ValueError, msg): - model._resolve_prompt_slots(template) - - # a group whose feature_config vanished: reachable only out of sync - orphan = _wired() - orphan._features = [] - with self.assertRaisesRegex(ValueError, "no feature_config"): - orphan._resolve_prompt_slots("{{user_sequence}}") - - def test_beam_config_defaults_and_validation(self) -> None: - def read(widths, num_return, levels=3): - m = object.__new__(BaseGenerativeModel) - m._num_levels = levels - m._read_beam_config( - types.SimpleNamespace( - beam_widths=widths, num_return_sequences=num_return - ) - ) - return m - - # the schedule is taken verbatim; there is no default to fall back on - self.assertEqual(read([100, 200, 400], 400)._beam_widths, [100, 200, 400]) - with self.assertRaisesRegex(ValueError, "beam_widths is required"): - read([], 50) - with self.assertRaisesRegex(ValueError, "one width per level"): - read([50, 50], 50) - with self.assertRaisesRegex(ValueError, "must not exceed the final"): - read([50, 50, 50], 80) - - def test_abstract_hooks_raise(self) -> None: - base = object.__new__(BaseGenerativeModel) - with self.assertRaises(NotImplementedError): - base._build_prompt_tokens(None, None) - with self.assertRaises(NotImplementedError): - base.predict(None) - - def test_device_property(self) -> None: - self.assertEqual(_stub(device="cpu").device, torch.device("cpu")) - - def test_tokenize_sids(self) -> None: - m = _stub(base_vocab=100) - # 0-based codes: the flat index IS the atom index, so no bridging shift. - flat = torch.tensor([[0, 2, 5], [1, 4, 8]]) # offsets [0, 2, 5] - level_ids = torch.arange(3) - out = m._tokenize_sids(flat) - self.assertEqual(out.tolist(), [[100, 102, 105], [101, 104, 108]]) - self.assertEqual(out.dtype, torch.int64) - # decode goes the whole way back to per-level codes - self.assertEqual( - m._detokenize_sids(out, level_ids).tolist(), [[0, 0, 0], [1, 2, 3]] - ) - - def test_sid_token_bands_use_same_level_aware_mapping(self) -> None: - m = _stub(base_vocab=100) - lo, hi = m._sid_token_bands() - self.assertEqual(lo.tolist(), [100, 102, 105]) - self.assertEqual(hi.tolist(), [101, 104, 108]) - - def test_sid_token_rows_shifts_and_splits(self) -> None: - # values arrive FLAT from SidFeature._parse; the model adds base_vocab - m = _stub(base_vocab=100) - jt = _FakeJT([0, 2, 5, 1, 4, 8, 0, 3, 6], [6, 3]) - rows = m._sid_token_rows(jt.values(), jt.lengths()) - self.assertEqual( - [r.tolist() for r in rows], - [[100, 102, 105, 101, 104, 108], [100, 103, 106]], - ) - self.assertTrue(all(r.dtype == torch.int64 for r in rows)) - - def test_sid_token_rows_squeezes_n1(self) -> None: - m = _stub(base_vocab=100) - jt = _FakeJT([0, 2, 5], [3], dim2=True) # (N, 1) - rows = m._sid_token_rows(jt.values(), jt.lengths()) - self.assertEqual([r.tolist() for r in rows], [[100, 102, 105]]) - - def test_sid_token_rows_recency_clip(self) -> None: - m = _stub(base_vocab=100) - flat = [0, 2, 5, 1, 3, 6, 0, 4, 7, 1, 2, 8, 0, 3, 5] # 5 whole items - values = torch.tensor(flat, dtype=torch.float) - lengths = torch.tensor([values.numel()]) - tail = [100 + v for v in flat[6:]] # last three items, +base_vocab - # cap 9 -> keep the most recent three whole items - self.assertEqual( - m._sid_token_rows(values, lengths, max_codes=9)[0].tolist(), tail - ) - # item-aligned: cap 10 still keeps 9, never cuts mid-item - self.assertEqual( - m._sid_token_rows(values, lengths, max_codes=10)[0].tolist(), tail - ) - # disabled -> untouched - self.assertEqual( - m._sid_token_rows(values, lengths, max_codes=0)[0].tolist(), - [100 + v for v in flat], - ) - - def test_answer_token_rows_folds_offsets_and_validates(self) -> None: - # the answer is a label_field, not a feature, so the model still owns - # the per-level fold-in for it. - m = _stub(base_vocab=100) - jt = _FakeJT([0, 0, 0, 1, 2, 3], [3, 3]) - rows = m._answer_token_rows(jt.values(), jt.lengths()) - self.assertEqual([r.tolist() for r in rows], [[100, 102, 105], [101, 104, 108]]) - with self.assertRaisesRegex(ValueError, "each answer must be"): - bad = _FakeJT([0, 0, 0, 0, 0, 0], [2, 4]) - m._answer_token_rows(bad.values(), bad.lengths()) - with self.assertRaisesRegex(ValueError, "local 0-based"): - bad = _FakeJT([0, 3, 0], [3]) # 3 == codebook[1] - m._answer_token_rows(bad.values(), bad.lengths()) - - def test_build_input_history_group_label_field(self) -> None: - m = _stub(base_vocab=100) - m._label_name = "label" - m._slot_names = ["user_sequence"] - m._slot_groups = ["user_seq"] - m._max_seq_length = 0 - m._is_inference = False # train: the answer label_field is read too - m.embedding_group = lambda b: { - # flat indices, as SidFeature._parse emits them - "user_seq.sequence": torch.tensor( - [0.0, 2.0, 5.0, 1.0, 4.0, 8.0, 1.0, 2.0, 7.0] - ), - "user_seq.sequence_length": torch.tensor([6, 3]), - } - batch = types.SimpleNamespace( - jagged_labels={"label": _FakeJT([1, 2, 3, 0, 1, 2], [3, 3])} - ) - rows = m.build_input(batch) - self.assertEqual( - [r.tolist() for r in rows["user_sequence"]], - [[100, 102, 105, 101, 104, 108], [101, 102, 107]], - ) - self.assertEqual( - [r.tolist() for r in rows["label"]], - [[101, 104, 108], [100, 103, 107]], - ) - - def test_build_input_skips_label_in_inference(self) -> None: - m = _stub(base_vocab=100) - m._label_name = "label" - m._slot_names = ["user_sequence"] - m._slot_groups = ["user_seq"] - m._max_seq_length = 0 - m._is_inference = True # inference: history only, no ground-truth label - m.embedding_group = lambda b: { - "user_seq.sequence": torch.tensor([0.0, 2.0, 5.0]), - "user_seq.sequence_length": torch.tensor([3]), - } - rows = m.build_input(types.SimpleNamespace(jagged_labels={})) - self.assertEqual([r.tolist() for r in rows["user_sequence"]], [[100, 102, 105]]) - self.assertNotIn("label", rows) - - -if __name__ == "__main__": - unittest.main() diff --git a/tzrec/models/generative_qwen.py b/tzrec/models/generative_qwen.py deleted file mode 100644 index 3b7523bef..000000000 --- a/tzrec/models/generative_qwen.py +++ /dev/null @@ -1,270 +0,0 @@ -# Copyright (c) 2026, Alibaba Group; -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# http://www.apache.org/licenses/LICENSE-2.0 -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Qwen family subclass of ``BaseGenerativeModel``. - -Owns the decoder-only-chat implementation: the ChatML prompt template, the -causal-LM splice, and the ``.model``/``.lm_head`` forward. -""" - -from itertools import chain -from typing import Any, Dict, List, Optional, Tuple - -import torch -from torch.nn.utils.rnn import pad_sequence -from transformers import PreTrainedTokenizerBase - -from tzrec.datasets.utils import Batch -from tzrec.features.feature import BaseFeature -from tzrec.models.generative_model import BaseGenerativeModel -from tzrec.modules.dynamic_beam import dynamic_beam_search -from tzrec.protos.model_pb2 import ModelConfig -from tzrec.protos.models import generative_model_pb2 - - -@torch.fx.wrap -def _fx_wrapped_generate(model: "GenerativeQwen", batch: Batch) -> torch.Tensor: - """One opaque FX leaf spanning the whole decode. - - TorchRec's predict pipeline FX-traces the model; the decode is untraceable - (per-row python lists, data-dependent beam widths). Wrapping the WHOLE - decode is required -- a leaf returns one Proxy, so wrapping an inner helper - just moves the failure to its caller. At run time this is a normal call. - """ - return model._generate(batch) - - -class GenerativeQwen(BaseGenerativeModel): - """Generative-recommendation LM on a Qwen backbone (Qwen2.5, Qwen3, ...).""" - - CHAT_TEMPLATE = { - "user_prefix": "<|im_start|>user\n", - "user_suffix": "<|im_end|>\n", - "asst_prefix": "<|im_start|>assistant\n", - "asst_suffix": "<|im_end|>\n", - } - - def __init__( - self, - model_config: ModelConfig, - features: List[BaseFeature], - labels: List[str], - sample_weights: Optional[List[str]] = None, - **kwargs: Any, - ) -> None: - super().__init__(model_config, features, labels, sample_weights, **kwargs) - self._max_total_len = self._compute_max_total_length() - self._pool_warmed = False - # +2 = trailing eos + HF's shift-by-one; constant width avoids a per-step sync. - self._suffix_keep = self._num_levels + self.tpl_asst_suffix.numel() + 2 - - def _compute_max_total_length(self) -> int: - """The ``T`` the activation pool pre-sizes to; 0 when disabled.""" - if self._max_seq_length <= 0: - return 0 - frame = ( - sum(g.numel() for g in self._gaps) - + self.tpl_asst_suffix.numel() - + self.tpl_eos.numel() - ) - return int( - frame + self._max_seq_length * len(self._slot_names) + self._num_levels - ) - - @property - def _gaps(self) -> List[torch.Tensor]: - """The N+1 static prompt fragments around the N slots, template order. - - Re-read every time: ``.to()`` rebinds the buffer, so a cache goes stale. - """ - return [getattr(self, f"tpl_gap_{i}") for i in range(len(self._slot_names) + 1)] - - def _build_prompt_tokens( - self, - tokenizer: PreTrainedTokenizerBase, - cfg: generative_model_pb2.GenerativeQwen, - ) -> None: - """Tokenise the static prompt once, as the N+1 gaps around the N slots. - - Every gap is ONE string encoded in one call, so a BPE merge cannot span - a seam. Splicing values between gaps is exact only because the ``C*`` - atoms are added-vocab tokens, which fast tokenizers pre-split on. - Buffers are non-persistent: they follow ``.to()`` but stay off the - state_dict. - """ - tpl = self.CHAT_TEMPLATE - gaps, features = self._resolve_prompt_slots(cfg.prompt_template) - for i, gap in enumerate(gaps): - head = tpl["user_prefix"] if i == 0 else features[i - 1].suffix_text - tail = ( - features[i].prefix_text - if i < len(features) - else tpl["user_suffix"] + tpl["asst_prefix"] - ) - # the template carries its own markers; no auto BOS/EOS. - ids = torch.tensor( - tokenizer.encode(head + gap + tail, add_special_tokens=False), - dtype=torch.long, - ) - self.register_buffer(f"tpl_gap_{i}", ids, persistent=False) - self.register_buffer( - "tpl_asst_suffix", - torch.tensor( - tokenizer.encode(tpl["asst_suffix"], add_special_tokens=False), - dtype=torch.long, - ), - persistent=False, - ) - # the trailing eos is SUPERVISED. - self.register_buffer( - "tpl_eos", - torch.tensor([int(tokenizer.eos_token_id)], dtype=torch.long), - persistent=False, - ) - - def _prompt_rows(self, slot_rows: List[List[torch.Tensor]]) -> List[torch.Tensor]: - """Per-row ``[gap_0 | slot_0 | gap_1 | ... | slot_N-1 | gap_N]``.""" - gaps = self._gaps - # zip pairs each slot with the gap before it; gaps[-1] closes the row. - return [ - torch.cat([*chain.from_iterable(zip(gaps, slots)), gaps[-1]]) - for slots in zip(*slot_rows) - ] - - def _slot_rows( - self, rows: Dict[str, List[torch.Tensor]] - ) -> List[List[torch.Tensor]]: - """Per-slot token rows, in template order.""" - return [rows[name] for name in self._slot_names] - - def _splice_input_ids( - self, - slot_rows: List[List[torch.Tensor]], - label_rows: List[torch.Tensor], - pad_to: int = 0, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Build ``(input_ids, labels, attention_mask)``, each ``(B, T_max)``. - - The tail ``[answer | asst_suffix | eos]`` has a FIXED width, so after - left padding it lands in the same columns for every row and ``labels`` - is one vectorized write. Only the answer and trailing eos are - supervised. ``pad_to`` left-extends for pool pre-sizing. - """ - if len(slot_rows[0]) != len(label_rows): - raise ValueError( - f"{type(self).__name__}: history/answer row count mismatch " - f"({len(slot_rows[0])} vs {len(label_rows)})." - ) - rows_ids = [ - torch.cat([prompt, label_rows[i], self.tpl_asst_suffix, self.tpl_eos]) - for i, prompt in enumerate(self._prompt_rows(slot_rows)) - ] - input_ids, attention_mask = self._left_pad(rows_ids, pad_to=pad_to) - - B, T = input_ids.shape - answer_width = self._num_levels - tail = answer_width + self.tpl_asst_suffix.numel() + 1 - labels = torch.full( - (B, T), self._ignore_index, dtype=torch.long, device=self.device - ) - labels[:, T - tail : T - tail + answer_width] = torch.stack(label_rows) - labels[:, -1] = self.tpl_eos[0] - return input_ids, labels, attention_mask - - def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: - """Dispatch on the TER inference flag (``set_is_inference`` in main.py).""" - if self.is_inference: - return {self._generated_sids_key: _fx_wrapped_generate(self, batch)} - return self._predict_train(batch) - - def _predict_train(self, batch: Batch) -> Dict[str, torch.Tensor]: - """Build the teacher-forced splice for a batch and return the CE loss.""" - rows = self.build_input(batch) - - # Pre-size the caching allocator on step 1; the extra columns are masked. - pad_to = 0 - if not self._pool_warmed and self._max_total_len > 0 and self.is_train: - pad_to = self._max_total_len - self._pool_warmed = True - - input_ids, labels, attention_mask = self._splice_input_ids( - self._slot_rows(rows), rows[self._label_name], pad_to=pad_to - ) - return self._forward_loss(input_ids, labels, attention_mask) - - def _forward_loss( - self, - input_ids: torch.Tensor, - labels: torch.Tensor, - attention_mask: torch.Tensor, - ) -> Dict[str, torch.Tensor]: - """Teacher-forced forward over spliced ids -> suffix-slice -> CE loss.""" - outputs = self.lm.model(input_ids=input_ids, attention_mask=attention_mask) - hidden = outputs.last_hidden_state - - # a full (B, T, V) upcast OOMs. - suffix = slice(-self._suffix_keep, None) - logits = self.lm.lm_head(hidden[:, suffix, :]) - - loss = self.lm.loss_function( - logits=logits, - labels=labels[:, suffix], - vocab_size=self.lm.config.vocab_size, - ignore_index=self._ignore_index, - ) - return {"loss": loss} - - def _generate(self, batch: Batch) -> torch.Tensor: - """Beam-search the SID answer; returns ``(B, C, num_levels)`` best-first.""" - slot_rows = self._slot_rows(self.build_input(batch)) - input_ids, attention_mask = self._left_pad(self._prompt_rows(slot_rows)) - lo_tok, hi_tok = self._sid_token_bands() - new_tokens = dynamic_beam_search( - self.lm, - input_ids, - attention_mask, - beam_widths=self._beam_widths, - lo_tok=lo_tok, - hi_tok=hi_tok, - ) - sids = self._validate_sid_candidates(new_tokens, input_ids.shape[0]) - return sids[:, : self._num_return] - - def _left_pad( - self, rows: List[torch.Tensor], pad_to: int = 0 - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Left-pad token rows into ``(input_ids, attention_mask)``, ``(B, T_max)``. - - The mask comes from ``ones_like``, not ``!= pad``, so a real trailing - eos survives ``pad_token_id == eos``. - """ - input_ids = pad_sequence( - rows, - batch_first=True, - padding_value=self._pad_token_id, - padding_side="left", - ) - attention_mask = pad_sequence( - [torch.ones_like(r) for r in rows], - batch_first=True, - padding_value=0, - padding_side="left", - ) - if pad_to > input_ids.shape[1]: - B, extra = input_ids.shape[0], pad_to - input_ids.shape[1] - input_ids = torch.cat( - [input_ids.new_full((B, extra), self._pad_token_id), input_ids], - dim=1, - ) - attention_mask = torch.cat( - [attention_mask.new_zeros((B, extra)), attention_mask], dim=1 - ) - return input_ids, attention_mask diff --git a/tzrec/models/generative_qwen_test.py b/tzrec/models/generative_qwen_test.py deleted file mode 100644 index 5fd8df2b5..000000000 --- a/tzrec/models/generative_qwen_test.py +++ /dev/null @@ -1,420 +0,0 @@ -# Copyright (c) 2026, Alibaba Group; -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# http://www.apache.org/licenses/LICENSE-2.0 -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import types -import unittest -from unittest import mock - -import torch -from parameterized import parameterized -from torch import nn - -from tzrec.models.generative_qwen import GenerativeQwen -from tzrec.utils.test_util import create_tiny_causal_lm, parameterized_name_func - - -def _stub(codebook=None, base_vocab=100, pad_id=9, device="cpu"): - """A GenerativeQwen with the splice-relevant state wired up, no HF backbone. - - The non-uniform default codebook makes incorrect ``level * uniform_size`` - offset arithmetic visible: sizes=[2,3,4], offsets=[0,2,5]. - """ - codebook = codebook or [2, 3, 4] - m = object.__new__(GenerativeQwen) - nn.Module.__init__(m) - m._ignore_index = -100 - m._num_levels = len(codebook) - m._base_vocab = base_vocab - m._pad_token_id = pad_id - m._num_return = 2 - m._beam_widths = [2] * m._num_levels - m._max_seq_length = 0 - m._slot_names = ["user_sequence"] - m._generated_sids_key = "generated_sids" - m.lm = types.SimpleNamespace(device=torch.device(device)) - for name, vals in { - "tpl_gap_0": [10, 11, 12], - "tpl_gap_1": [13, 14], - "tpl_asst_suffix": [15], - "tpl_eos": [9], - }.items(): - m.register_buffer(name, torch.tensor(vals, dtype=torch.long), persistent=False) - sizes = torch.tensor(codebook, dtype=torch.long) - m.register_buffer("_codebook_sizes", sizes, persistent=False) - m.register_buffer( - "_level_offsets", torch.cumsum(sizes, 0) - sizes, persistent=False - ) - return m - - -def _real_lm_stub(codebook=None, base_vocab=20, beam_width=2): - """A GenerativeQwen carrying a real (tiny, random) Qwen2 backbone. - - Needed wherever the real forward runs: the training objective and the - end-to-end band-restricted decode; the other tests mock the kernel. - """ - codebook = codebook or [2, 3, 4] - m = object.__new__(GenerativeQwen) - nn.Module.__init__(m) - m._num_levels = len(codebook) - m._base_vocab = base_vocab - m._beam_widths = [beam_width] * m._num_levels - m.lm = create_tiny_causal_lm(base_vocab + sum(codebook)) - sizes = torch.tensor(codebook, dtype=torch.long) - m.register_buffer("_codebook_sizes", sizes, persistent=False) - m.register_buffer( - "_level_offsets", torch.cumsum(sizes, 0) - sizes, persistent=False - ) - return m - - -def _wire_slots(m, prefix_text="", suffix_text=""): - """Minimal _features/_feature_groups so the base slot resolver can run.""" - from google.protobuf import text_format - - from tzrec.features.feature import create_features - from tzrec.protos import feature_pb2, model_pb2 - - fc = feature_pb2.FeatureConfig() - text_format.Merge( - 'sequence_sid_feature { feature_name: "user_sequence" ' - 'expression: "user:user_sequence" codebook: 2 codebook: 3 codebook: 4 ' - f'prefix_text: "{prefix_text}" suffix_text: "{suffix_text}" }}', - fc, - ) - m._features = create_features([fc]) - m._feature_groups = [ - types.SimpleNamespace( - group_name="sids", - feature_names=["user_sequence"], - group_type=model_pb2.JAGGED_SEQUENCE, - ) - ] - - -def _train_stub(max_total_len): - """A ``_predict_train``-ready stub; returns ``(model, spliced T per step)``.""" - m = _stub() - m._is_inference = False # not inference + nn.Module.training=True -> is_train - m._label_name = "label" - m._slot_names = ["user_sequence"] - m._max_total_len = max_total_len - m._pool_warmed = False - seen_lens = [] - - def fwd(input_ids, labels, attention_mask): - seen_lens.append(input_ids.shape[1]) - return {"loss": torch.tensor(0.0)} - - m.build_input = lambda b: { - "user_sequence": [torch.tensor([100, 101, 102])], - m._label_name: [torch.tensor([200, 201, 202])], - } - m._forward_loss = fwd - return m, seen_lens - - -class GenerativeQwenTest(unittest.TestCase): - def test_splice_layout_and_labels(self) -> None: - m = _stub() - u = [torch.tensor([100, 101, 102])] - a = [torch.tensor([200, 201, 202])] # 3 codes = num_levels - ids, labels, mask = m._splice_input_ids([u], a) - # head | history | tail | answer | asst_suffix | eos - self.assertEqual( - ids[0].tolist(), [10, 11, 12, 100, 101, 102, 13, 14, 200, 201, 202, 15, 9] - ) - # only the answer (cols 8-10) and the trailing eos (col 12) are supervised - self.assertEqual( - labels[0].tolist(), - [-100] * 8 + [200, 201, 202, -100, 9], - ) - self.assertEqual(mask[0].tolist(), [1] * 13) - - def test_left_padding_varied_lengths(self) -> None: - m = _stub() - u = [torch.tensor([100, 101, 102, 103]), torch.tensor([100])] - a = [torch.tensor([200, 201, 202]), torch.tensor([207, 208, 209])] - ids, labels, mask = m._splice_input_ids([u], a) - T = ids.shape[1] - n1 = 2 + 1 + 1 + 1 + 1 + 3 + 1 + 1 # shorter row's real length - self.assertEqual(ids[1, : T - n1].tolist(), [m._pad_token_id] * (T - n1)) - self.assertEqual(mask[1].tolist(), [0] * (T - n1) + [1] * n1) - self.assertEqual(labels[1, : T - n1].tolist(), [-100] * (T - n1)) - # the trailing eos is supervised in every row - self.assertEqual(labels[:, -1].tolist(), [9, 9]) - - def test_mask_keeps_trailing_eos_when_pad_equals_eos(self) -> None: - m = _stub(pad_id=9) # tpl_eos == 9 too - ids, _, mask = m._splice_input_ids( - [[torch.tensor([100])]], [torch.tensor([200, 201, 202])] - ) - self.assertEqual(int(ids[0, -1]), 9) - self.assertEqual(int(mask[0, -1]), 1) - self.assertEqual(mask[0].tolist(), [1] * ids.shape[1]) - - def test_predict_routes_on_inference_flag(self) -> None: - m = _stub() - m._predict_train = lambda b: {"branch": "train"} - m._is_inference = False # train / eval - self.assertEqual(GenerativeQwen.predict(m, object())["branch"], "train") - m._is_inference = True # inference - sentinel = torch.zeros(1, 2, 3) - with mock.patch( - "tzrec.models.generative_qwen._fx_wrapped_generate", return_value=sentinel - ): - out = GenerativeQwen.predict(m, object()) - self.assertIs(out["generated_sids"], sentinel) - - def test_predict_survives_fx_tracing(self) -> None: - """TorchRec's predict pipeline FX-traces the model before running it. - - The decode is un-traceable by construction (per-row python lists, - data-dependent beam widths), so it sits behind one ``torch.fx.wrap`` - leaf. Without it ``tzrec.predict`` dies inside TorchRec's - ``_rewrite_model`` with "Proxy object cannot be iterated" -- a failure no - unit test saw because train and eval take the un-traced - ``TrainPipelineBase`` path (this model has no ShardedModule). - """ - m = _stub(base_vocab=100) - m._is_inference = True - - class _Wrapper(nn.Module): - def __init__(self, inner): - super().__init__() - self.inner = inner - - def forward(self, batch): - return self.inner.predict(batch) - - gm = torch.fx.symbolic_trace(_Wrapper(m)) - leaves = [ - n - for n in gm.graph.nodes - if n.op == "call_function" and "generate" in str(n.target) - ] - self.assertEqual(len(leaves), 1, f"expected one opaque decode node: {gm.graph}") - - # (generated tail -> decoded SIDs). sizes [2,3,4] -> offsets [0,2,5], so - # [100,102,105] / [101,104,108] are each level's min/max token and local - # codes [0,0,0] / [1,2,3]; every malformed row collapses to -1. - @parameterized.expand( - [ - [[[100, 102, 105], [101, 104, 108]], [[0, 0, 0], [1, 2, 3]]], - [ - [ - [100, 102, 105], # valid -> local [0, 0, 0] - [101, 104, 108], # valid -> local [1, 2, 3] - [102, 102, 105], # pos0 above level-0 band - [100, 101, 105], # pos1 below level-1 band - [100, 102, 109], # pos2 above level-2 band - [100, 104, 9], # pos2 = eos/pad token (sid -90) -> invalid - ], - [[0, 0, 0], [1, 2, 3]] + [[-1, -1, -1]] * 4, - ], - # early EOS: a tail narrower than num_levels still reshapes cleanly, - # and the missing 3rd atom stays -1 -> out of band -> candidate -1 - [[[100, 102], [101, 104]], [[-1, -1, -1], [-1, -1, -1]]], - ], - name_func=parameterized_name_func, - ) - def test_generate_maps_tokens_to_sids(self, tail, expected) -> None: - m = _stub(base_vocab=100) - m._slot_names = ["user_sequence"] - m._num_return = len(tail) - m._beam_widths = [len(tail)] * m._num_levels - # build_input is mocked, so the batch is opaque to this decoding test. - m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} - with mock.patch( - "tzrec.models.generative_qwen.dynamic_beam_search", - return_value=torch.tensor(tail), - ): - sids = m._generate(object()) - # (B, num_return, num_levels) - self.assertEqual(tuple(sids.shape), (1, len(tail), 3)) - self.assertEqual(sids[0].tolist(), expected) - - def test_generate_trims_to_num_return_keeping_the_best(self) -> None: - # the kernel returns score-ordered best-first, so the trim is a prefix - m = _stub(base_vocab=100) - m._slot_names = ["user_sequence"] - m._beam_widths = [4] * m._num_levels - m._num_return = 2 - m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} - four = torch.tensor( - [[100, 102, 105], [101, 104, 108], [100, 103, 106], [101, 102, 107]] - ) - with mock.patch( - "tzrec.models.generative_qwen.dynamic_beam_search", return_value=four - ): - sids = m._generate(object()) - self.assertEqual(tuple(sids.shape), (1, 2, 3)) - self.assertEqual(sids[0].tolist(), [[0, 0, 0], [1, 2, 3]]) - - def test_generate_hands_the_kernel_prompt_bands_and_schedule(self) -> None: - m = _stub(base_vocab=100) - m._slot_names = ["user_sequence"] - m._beam_widths = [5 * 2 ** (j + 1) for j in range(m._num_levels)] - m._num_return = m._beam_widths[-1] - seen = {} - - def fake_kernel(lm, input_ids, attention_mask, *, beam_widths, lo_tok, hi_tok): - seen.update( - lm=lm, - ids=input_ids[0].tolist(), - mask=attention_mask[0].tolist(), - widths=list(beam_widths), - lo=lo_tok.tolist(), - hi=hi_tok.tolist(), - ) - return torch.tensor([[100, 102, 105], [101, 104, 108]]) - - m.build_input = lambda b: {"user_sequence": [torch.tensor([100, 102, 105])]} - with mock.patch( - "tzrec.models.generative_qwen.dynamic_beam_search", side_effect=fake_kernel - ): - sids = m._generate(object()) - self.assertIs(seen["lm"], m.lm) - self.assertEqual(seen["ids"], [10, 11, 12, 100, 102, 105, 13, 14]) - self.assertEqual(seen["mask"], [1] * 8) - # the schedule reaches the kernel verbatim - self.assertEqual(seen["widths"], [10, 20, 40]) - # per-level bands, base_vocab-shifted: sizes [2,3,4] -> offsets [0,2,5] - self.assertEqual(seen["lo"], [100, 102, 105]) - self.assertEqual(seen["hi"], [101, 104, 108]) - self.assertEqual(tuple(sids.shape), (1, 2, 3)) - self.assertEqual(sids[0].tolist(), [[0, 0, 0], [1, 2, 3]]) - - def _prompt_tokenizer(self): - # encode -> [len(text)] makes each buffer a fingerprint of its fragment - return types.SimpleNamespace( - eos_token_id=99, - encode=lambda text, add_special_tokens=False: [len(text)], - ) - - def test_build_prompt_tokens_splits_the_template_around_the_slot(self) -> None: - m = object.__new__(GenerativeQwen) - nn.Module.__init__(m) - _wire_slots(m, prefix_text="PRE", suffix_text="SUF") - cfg = types.SimpleNamespace(prompt_template="A{{user_sequence}}B") - m._build_prompt_tokens(self._prompt_tokenizer(), cfg) - for name in ["tpl_gap_0", "tpl_gap_1", "tpl_asst_suffix", "tpl_eos"]: - buf = getattr(m, name) - self.assertIsInstance(buf, torch.Tensor) - self.assertEqual(buf.dtype, torch.int64) - tpl = GenerativeQwen.CHAT_TEMPLATE - # head = user_prefix + before + feature.prefix_text - self.assertEqual(m.tpl_gap_0.tolist(), [len(tpl["user_prefix"] + "A" + "PRE")]) - # tail = feature.suffix_text + after + user_suffix + asst_prefix - self.assertEqual( - m.tpl_gap_1.tolist(), - [len("SUF" + "B" + tpl["user_suffix"] + tpl["asst_prefix"])], - ) - self.assertEqual(m.tpl_eos.tolist(), [99]) # eos cached for supervision - - def test_compute_max_total_length(self) -> None: - m = _stub() - # frame = 3 head + 2 tail + 1 asst_suffix + 1 eos = 7 - m._max_seq_length = 300 - self.assertEqual(m._compute_max_total_length(), 7 + 300 + 3) - m._max_seq_length = 0 # pre-allocation disabled - self.assertEqual(m._compute_max_total_length(), 0) - - def test_first_step_pads_to_max_then_actual_length(self) -> None: - m, seen_lens = _train_stub(max_total_len=50) - m._predict_train(object()) # first step: pre-size to worst case - m._predict_train(object()) # subsequent step: natural length - self.assertEqual(seen_lens[0], 50) - self.assertLess(seen_lens[1], 50) - self.assertTrue(m._pool_warmed) - - def test_no_forced_padding_when_disabled(self) -> None: - m, seen_lens = _train_stub(max_total_len=0) # pre-allocation off - m._predict_train(object()) - self.assertLess(seen_lens[0], 50) - self.assertFalse(m._pool_warmed) - - def test_splice_row_count_mismatch_raises(self) -> None: - m = _stub() - with self.assertRaisesRegex(ValueError, "row count mismatch"): - m._splice_input_ids([[torch.tensor([100])]], []) - - -class GenerativeQwenLossTest(unittest.TestCase): - """The training objective, run for real against a tiny Qwen backbone.""" - - def _model(self, ignore_index=-100): - m = _real_lm_stub(codebook=[2, 3, 4], base_vocab=20, beam_width=2) - m._ignore_index = ignore_index - m._pad_token_id = 0 - for name, vals in { - "tpl_gap_0": [1, 2, 3], - "tpl_gap_1": [4, 5], - "tpl_asst_suffix": [6], - "tpl_eos": [7], - }.items(): - m.register_buffer( - name, torch.tensor(vals, dtype=torch.long), persistent=False - ) - m._slot_names = ["user_sequence"] - m._suffix_keep = 6 # num_levels 3 + asst_suffix 1 + trailing eos + HF shift - return m - - def _rows(self): - # ragged histories so left padding is exercised - u = [torch.tensor([20, 22, 25]), torch.tensor([21, 23, 26, 20, 24, 27])] - a = [torch.tensor([20, 22, 25]), torch.tensor([21, 24, 28])] - return [u], a - - def test_suffix_slice_matches_full_sequence_loss(self) -> None: - # the fixed-width suffix slice must give the same CE as full-T logits - m = self._model() - ids, labels, mask = m._splice_input_ids(*self._rows()) - with torch.no_grad(): - got = m._forward_loss(ids, labels, mask)["loss"] - full = m.lm(input_ids=ids, attention_mask=mask).logits - ref = m.lm.loss_function( - logits=full, labels=labels, vocab_size=m.lm.config.vocab_size - ) - self.assertTrue(torch.allclose(got, ref, atol=1e-6)) - - def test_loss_is_invariant_to_extra_left_padding(self) -> None: - # the pool-warmup pad_to must not perturb the objective - m = self._model() - u, a = self._rows() - with torch.no_grad(): - base = m._forward_loss(*m._splice_input_ids(u, a))["loss"] - padded = m._forward_loss(*m._splice_input_ids(u, a, pad_to=40))["loss"] - self.assertTrue(torch.allclose(base, padded, atol=1e-6)) - - def test_forward_loss_honours_configured_ignore_index(self) -> None: - # ignore_index must reach loss_function or every pad slot is supervised - m = self._model(ignore_index=-100) - u, a = self._rows() - with torch.no_grad(): - default = m._forward_loss(*m._splice_input_ids(u, a))["loss"] - m._ignore_index = -7 - with torch.no_grad(): - ids, labels, mask = m._splice_input_ids(u, a) - custom = m._forward_loss(ids, labels, mask)["loss"] - self.assertEqual(int((labels == -7).sum() > 0), 1) - self.assertTrue(torch.allclose(default, custom, atol=1e-6)) - - def test_forward_loss_returns_only_the_loss(self) -> None: - # returned logits would stay alive across the next step's fwd/bwd - m = self._model() - with torch.no_grad(): - out = m._forward_loss(*m._splice_input_ids(*self._rows())) - self.assertEqual(list(out), ["loss"]) - - -if __name__ == "__main__": - unittest.main() diff --git a/tzrec/modules/dynamic_beam_test.py b/tzrec/modules/dynamic_beam_test.py index 08ace1f6a..7791f924e 100644 --- a/tzrec/modules/dynamic_beam_test.py +++ b/tzrec/modules/dynamic_beam_test.py @@ -128,7 +128,7 @@ def test_rejects_a_schedule_that_does_not_match_the_bands(self) -> None: _decode(lm, ids, pairs, beam_widths=[2, 0, 4]) def test_tokens_stay_inside_arbitrary_bands(self) -> None: - # bands the GenerativeQwen caller can never produce: descending, disjoint, + # bands the PromptGenerativeQwen caller can never produce: descending, disjoint, # unequal width -- the kernel's contract is per-level (lo, hi), not a # contiguous codebook layout. pairs = [(5, 6), (20, 24), (11, 13)] diff --git a/tzrec/protos/feature.proto b/tzrec/protos/feature.proto index 96ce81d71..19cebd176 100644 --- a/tzrec/protos/feature.proto +++ b/tzrec/protos/feature.proto @@ -986,50 +986,6 @@ message SequenceFeature { repeated SeqFeatureConfig features = 5; } -// Semantic-ID (SID) sequence feature for generative-recommendation LMs. -// -// Carries a flat stream of per-level SID codes (whole items, level order) plus -// the prompt text that wraps them. Only usable as `sequence_sid_feature`: the -// model reads it as a JAGGED_SEQUENCE group, so the scalar form would never -// produce the "{group}.sequence" keys it needs. -// -// Under fg this is a PASSTHROUGH: no fg feature_type can add -// level_offsets[i % levels], so fg only reaches the codes and the feature folds -// the offsets in at parse time. -message SidFeature { - // feature name; also the {{name}} placeholder in prompt_template. - required string feature_name = 1; - // feature input, e.g. user:user_sequence - required string expression = 2; - // SID streams are flat, so this stays 1. - optional uint32 value_dim = 6 [default = 1]; - // unused; SIDs carry no embedding table. - optional string pooling = 10 [default = "sum"]; - // fg default value - optional string default_value = 11 [default = "0"]; - // mask value in training progress - optional bool use_mask = 14; - - // text emitted immediately BEFORE this feature's SID tokens - optional string prefix_text = 20 [default = ""]; - // text emitted immediately AFTER this feature's SID tokens - optional string suffix_text = 21 [default = ""]; - // REQUIRED, per-level SID vocabulary; codes are 0-based in - // [0, codebook[level]). All SID features must declare the SAME codebook. - repeated uint32 codebook = 23; - - // default value when fg_mode = FG_NONE - optional string fg_encoded_default_value = 30; - // only used as fg dag intermediate result or not - optional bool stub_type = 34 [default = false]; - - // fg-side cap; must be a multiple of len(codebook). Prefer - // model_config.common.max_sequence_length, which keeps the recent items. - optional uint32 sequence_length = 101; - // sequence delimiter, only take effect when use it as sequence - optional string sequence_delim = 102 [default = ";"]; -} - message FeatureConfig { oneof feature { IdFeature id_feature = 1; @@ -1058,7 +1014,6 @@ message FeatureConfig { KvDotProduct sequence_kv_dot_product = 111; BoolMaskFeature sequence_bool_mask_feature = 112; CombineFeature sequence_combine_feature = 113; - SidFeature sequence_sid_feature = 114; } } diff --git a/tzrec/protos/model.proto b/tzrec/protos/model.proto index 46ae7d174..d72b1f206 100644 --- a/tzrec/protos/model.proto +++ b/tzrec/protos/model.proto @@ -5,7 +5,6 @@ import "tzrec/protos/models/rank_model.proto"; import "tzrec/protos/models/multi_task_rank.proto"; import "tzrec/protos/models/match_model.proto"; import "tzrec/protos/models/general_rank_model.proto"; -import "tzrec/protos/models/generative_model.proto"; import "tzrec/protos/models/prompt_model.proto"; import "tzrec/protos/models/sid_model.proto"; import "tzrec/protos/loss.proto"; @@ -85,7 +84,6 @@ message ModelConfig { SidRqkmeans sid_rqkmeans = 601; // Generative (causal-LM) models; the 700-block keeps clear of the SID 600s. - GenerativeQwen generative_qwen = 700; PromptGenerativeQwen prompt_generative_qwen = 701; } diff --git a/tzrec/protos/models/generative_model.proto b/tzrec/protos/models/generative_model.proto deleted file mode 100644 index 004ab5353..000000000 --- a/tzrec/protos/models/generative_model.proto +++ /dev/null @@ -1,48 +0,0 @@ -syntax = "proto2"; -package tzrec.protos; - -// Generative (causal-LM) recommendation models. - -enum ParamDtype { - FP32 = 0; - BF16 = 1; - FP16 = 2; -} - -// Sample contract: history is a JAGGED_SEQUENCE feature_group holding one -// SidFeature; the answer is the FIRST data_config.label_field, not a feature. -// Both are local 0-based per-level codes -- writers must not pre-apply -// level_offsets, SidFeature folds them in at parse time. -message GenerativeModelConfig { - // 0 disables padding. - optional uint32 vocab_pad_to_multiple_of = 3 [default = 128]; - - optional int32 ignore_index = 6 [default = -100]; - - // REQUIRED. Beam width per SID level, one entry per level; [100, 200, 400] - // is the escalating beam. Each entry is capped to what its band can supply. - repeated uint32 beam_widths = 7; - // Must not exceed the final beam width. - required uint32 num_return_sequences = 8; - - optional string generated_sids_key = 12 [default = "generated_sids"]; - // MASTER weights. FP32 avoids bf16-ULP underflow of Adam's small updates; - // bf16 COMPUTE comes from mixed_precision, not from this. - optional ParamDtype param_dtype = 13 [default = FP32]; - - // History budget in SID codes: item-aligned truncation keeping the most - // recent items, and the activation-pool pre-size. 0 disables both. - required uint32 max_sequence_length = 14; -} - -// Qwen family (Qwen2.5, Qwen3, ...) -- one ChatML frame for all. -message GenerativeQwen { - optional GenerativeModelConfig common = 1; - - // HF hub id or local path. - optional string hf_model_id = 2 [default = "Qwen/Qwen2.5-0.5B"]; - - // User-turn body carrying one {{feature_name}} placeholder per SID feature. - // The ChatML frame and the supervised tail are family constants. - optional string prompt_template = 10; -} diff --git a/tzrec/tests/configs/generative_qwen_mock.config b/tzrec/tests/configs/generative_qwen_mock.config deleted file mode 100644 index 9a28118e5..000000000 --- a/tzrec/tests/configs/generative_qwen_mock.config +++ /dev/null @@ -1,65 +0,0 @@ -train_input_path: "" -eval_input_path: "" -model_dir: "experiments/generative_qwen_mock" -train_config { - sparse_optimizer { - adagrad_optimizer { - lr: 0.0 - } - constant_learning_rate { - } - } - dense_optimizer { - adam_optimizer { - lr: 0.0001 - } - linear_decay_learning_rate { - num_training_steps: 32 - } - } - num_epochs: 1 - save_checkpoints_epochs: 1 -} -eval_config { -} -export_config { - export_format: HF -} -data_config { - batch_size: 4 - dataset_type: ParquetDataset - label_fields: "label" - num_workers: 2 - fg_mode: FG_NONE -} -feature_configs { - sequence_sid_feature { - feature_name: "user_sequence" - expression: "user:user_sequence" - prefix_text: "Current user's historical behaviors are as follows:" - codebook: 4 - codebook: 4 - codebook: 4 - } -} -model_config { - feature_groups { - group_name: "sids" - feature_names: "user_sequence" - group_type: JAGGED_SEQUENCE - } - generative_qwen { - common { - vocab_pad_to_multiple_of: 128 - ignore_index: -100 - param_dtype: FP32 - max_sequence_length: 12 - beam_widths: 4 - beam_widths: 8 - beam_widths: 16 - num_return_sequences: 16 - } - hf_model_id: "Qwen/Qwen2.5-0.5B" - prompt_template: "You are a recommendation system. Based on the user's historical behavior, predict the user's next action in an e-commerce scenario. I will provide a sequence of semantic encodings representing consecutive behaviors, arranged in chronological order of user clicks. Each behavior is represented by three words. {{user_sequence}} Please predict the semantic encoding of the user's subsequent behavior in the e-commerce recommendation scenario." - } -} diff --git a/tzrec/tests/genrec_integration_test.py b/tzrec/tests/genrec_integration_test.py deleted file mode 100644 index 017546925..000000000 --- a/tzrec/tests/genrec_integration_test.py +++ /dev/null @@ -1,186 +0,0 @@ -# Copyright (c) 2026, Alibaba Group; -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# http://www.apache.org/licenses/LICENSE-2.0 -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import glob -import os -import random -import shutil -import unittest -from unittest import mock - -import pyarrow as pa -import pyarrow.parquet as pq - -from tzrec.tests import utils -from tzrec.utils import config_util -from tzrec.utils.test_util import create_tiny_causal_lm, make_test_dir - -_MOCK_CONFIG = "tzrec/tests/configs/generative_qwen_mock.config" -# must match the mock config's sequence_sid_feature.codebook -_CODEBOOK = [4, 4, 4] - - -def _write_backbone(save_dir: str, vocab_size: int = 256) -> str: - """Save a tiny but COMPLETE Qwen2 model dir so no test downloads a real one. - - Weights are required, not just ``config.json``: cold-start training calls - ``init_from_pretrained`` -> ``from_pretrained``, which refuses a dir with no - checkpoint. The SID offsets only need a tokenizer whose ``len()`` is stable - and which has room to append the ``C*`` atoms, so word-level is enough. - """ - from tokenizers import Tokenizer, models, pre_tokenizers - from transformers import PreTrainedTokenizerFast - - eos = "<|endoftext|>" - vocab = {t: i for i, t in enumerate([eos, "<|im_start|>", "<|im_end|>"])} - for i in range(vocab_size - len(vocab)): - vocab[f"b{i}"] = len(vocab) - tk = Tokenizer(models.WordLevel(vocab=vocab, unk_token=eos)) - tk.pre_tokenizer = pre_tokenizers.Whitespace() - PreTrainedTokenizerFast( - tokenizer_object=tk, - unk_token=eos, - eos_token=eos, - pad_token=eos, - additional_special_tokens=["<|im_start|>", "<|im_end|>"], - ).save_pretrained(save_dir) - create_tiny_causal_lm( - len(vocab), tie_word_embeddings=True, max_position_embeddings=512 - ).save_pretrained(save_dir) - return save_dir - - -def _write_samples(save_dir: str, num_rows: int, seed: int = 0) -> str: - """Write the two-column sample contract: history + answer, both list. - - Codes are local 0-based per-level values in ``[0, codebook[level])``, as the - SID-generation models emit them; every row holds whole items in level order. - """ - rnd = random.Random(seed) - - def _item(): - return [rnd.randrange(size) for size in _CODEBOOK] - - schema = pa.schema( - [ - pa.field("user_sequence", pa.list_(pa.int64()), nullable=False), - pa.field("label", pa.list_(pa.int64()), nullable=False), - ] - ) - table = pa.table( - { - "user_sequence": [ - [c for _ in range(rnd.randint(1, 4)) for c in _item()] - for _ in range(num_rows) - ], - "label": [_item() for _ in range(num_rows)], - }, - schema=schema, - ) - pq.write_table(table, os.path.join(save_dir, "part-0.parquet")) - return os.path.join(save_dir, "*.parquet") - - -class GenRecIntegrationTest(unittest.TestCase): - def setUp(self): - self.success = False - self.test_dir = make_test_dir() - # every rank builds its own backbone; one is enough for a mock run. - patcher = mock.patch.dict( - os.environ, {"TEST_NPROC_PER_NODE": "1", "HF_HUB_OFFLINE": "1"} - ) - patcher.start() - self.addCleanup(patcher.stop) - - def tearDown(self): - if self.success and os.path.exists(self.test_dir): - shutil.rmtree(self.test_dir) - - def _prepare_config(self, num_rows: int = 64) -> str: - """Point the mock config at a local backbone and freshly written samples.""" - backbone = _write_backbone(os.path.join(self.test_dir, "backbone")) - data_dir = os.path.join(self.test_dir, "genrec_data") - os.makedirs(data_dir, exist_ok=True) - data_glob = _write_samples(data_dir, num_rows) - - config = config_util.load_pipeline_config(_MOCK_CONFIG) - config.train_input_path = data_glob - config.eval_input_path = data_glob - config.model_config.generative_qwen.hf_model_id = backbone - config_path = os.path.join(self.test_dir, "genrec.config") - config_util.save_message(config, config_path) - return config_path - - def test_mock_config_builds_the_model_and_runs_a_batch(self) -> None: - """The config -> model path: oneof dispatch and the derived contract. - - The sample contract is derived, not configured, so only a real - ``ModelConfig`` proves that the history feature_group, the answer - label_field and the SID vocab extension line up. - """ - from tzrec.constant import Mode - from tzrec.datasets.dataset import create_dataloader - from tzrec.main import _create_features, _create_model - - config = config_util.load_pipeline_config(self._prepare_config()) - features = _create_features(list(config.feature_configs), config.data_config) - model = _create_model( - config.model_config, features, list(config.data_config.label_fields) - ) - self.assertEqual(type(model).__name__, "GenerativeQwen") - self.assertEqual(model._slot_groups, ["sids"]) - self.assertEqual(model._slot_names, ["user_sequence"]) - self.assertEqual(model._label_name, "label") - self.assertEqual(model._num_levels, len(_CODEBOOK)) - # base vocab + sum(codebook) atoms, padded to vocab_pad_to_multiple_of - self.assertGreaterEqual( - model.lm.config.vocab_size, model._base_vocab + sum(_CODEBOOK) - ) - self.assertEqual(model.lm.config.vocab_size % 128, 0) - - dataloader = create_dataloader( - config.data_config, features, config.train_input_path, mode=Mode.TRAIN - ) - model.train() - batch = next(dataloader.get_iterator()) - predictions = model.predict(batch) - self.assertEqual(list(predictions), ["loss"]) - self.assertTrue(bool(predictions["loss"].isfinite())) - - # _suffix_keep bounds the logits _forward_loss upcasts. HF shifts logits - # by one, so the window must open one column BEFORE the first supervised - # label or that label is never scored. The unit tests hardcode the width; - # only a real __init__ proves the formula computes it. - rows = model.build_input(batch) - _, labels, _ = model._splice_input_ids( - model._slot_rows(rows), rows[model._label_name] - ) - first_sup = int((labels >= 0).nonzero()[:, 1].min()) - self.assertEqual(model._suffix_keep, labels.shape[1] - first_sup + 1) - self.success = True - - def test_generative_qwen_train_eval(self) -> None: - """End-to-end train -> checkpoint, with HF assets co-located.""" - config_path = self._prepare_config() - self.success = utils.test_train_eval(config_path, self.test_dir) - self.assertTrue(self.success) - ckpts = glob.glob(os.path.join(self.test_dir, "train", "model.ckpt-*")) - self.assertTrue(ckpts, "no checkpoint persisted") - # export_format: HF -> each checkpoint is convertible without the code - for name in ("config.json", "tokenizer.json"): - self.assertTrue( - os.path.exists(os.path.join(ckpts[0], name)), - f"{name} not co-located in {ckpts[0]}", - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/tzrec/utils/hf_export_util_test.py b/tzrec/utils/hf_export_util_test.py index 97b988863..3830ec7cd 100644 --- a/tzrec/utils/hf_export_util_test.py +++ b/tzrec/utils/hf_export_util_test.py @@ -45,7 +45,7 @@ def save_pretrained(self, save_dir): class _GenRec(nn.Module): - """Stand-in for BaseGenerativeModel: an HF backbone plus unrelated params.""" + """Stand-in for a prompt-native model: an HF backbone plus unrelated params.""" def __init__(self, lm): super().__init__() From fb7b3c654f54c816b79d9f181e775d42fa699320 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 07:57:06 +0000 Subject: [PATCH 66/99] [feat] prompt: persist the prompt contract and check it on restore Adds a BaseModel.save_assets hook, symmetric with init_from_pretrained, called for every model.ckpt-N/. The prompt-native model writes sid_space.json, prompt_plan.json, the hashes and the extended tokenizer, so a checkpoint describes its own vocabulary rather than relying on config supplied out of band -- which is where offline/online skew comes from. Restore compares the two hashes. A vocab_hash mismatch raises: the SID space or tokenizer changed, so the decode bands no longer address the rows these weights learned, and the run would emit plausible output instead of failing. A plan_hash mismatch only reshapes the prompt, so it warns. ModulePlan is not persisted. It is model-only and rebuilt from config at every __init__, so writing it would create a second source of truth. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/main.py | 5 + tzrec/models/model.py | 12 +++ tzrec/models/prompt_generative_qwen.py | 9 ++ tzrec/prompt/persist.py | 124 +++++++++++++++++++++++++ tzrec/prompt/persist_test.py | 119 ++++++++++++++++++++++++ tzrec/utils/checkpoint_util.py | 24 +++++ 6 files changed, 293 insertions(+) create mode 100644 tzrec/prompt/persist.py create mode 100644 tzrec/prompt/persist_test.py diff --git a/tzrec/main.py b/tzrec/main.py index 31e616274..6460ed102 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -74,6 +74,7 @@ from tzrec.optim.lr_scheduler import BaseLR from tzrec.optim.optimizer import TZRecOptimizer from tzrec.prompt.compile import compile_prompt +from tzrec.prompt.persist import check_prompt_assets from tzrec.prompt.plan import CompiledPrompt from tzrec.protos import export_pb2 from tzrec.protos.data_pb2 import DataConfig, DatasetType @@ -759,6 +760,8 @@ def train_and_evaluate( # Restore dataloader state before create_dataloader starts its workers dataloader_state: Optional[Dict[str, Any]] = None + if ckpt_path: + check_prompt_assets(prompt, ckpt_path) if ckpt_path and continue_train: dataloader_state = ckpt_manager.restore_dataloader_state(ckpt_path) if dataloader_state and not restore_from_model_dir: @@ -1033,6 +1036,7 @@ def evaluate( ) if checkpoint_path: + check_prompt_assets(prompt, checkpoint_path) ckpt_manager.restore( checkpoint_path, model, @@ -1620,6 +1624,7 @@ def predict_checkpoint( model.eval() if checkpoint_path: + check_prompt_assets(prompt, checkpoint_path) ckpt_manager.restore( checkpoint_path, model, diff --git a/tzrec/models/model.py b/tzrec/models/model.py index 8a7ae1d35..0086966d3 100644 --- a/tzrec/models/model.py +++ b/tzrec/models/model.py @@ -94,6 +94,18 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: """ raise NotImplementedError + def save_assets(self, target_dir: str) -> None: + """Write any contract a checkpoint needs to describe itself. + + Lifecycle hook symmetric with ``init_from_pretrained``, called for every + ``model.ckpt-N/`` and for the export directory. The default is a no-op; + models whose weights are meaningless without a companion artifact -- a + vocabulary, a plan -- override it. + + Args: + target_dir: the checkpoint or export directory. + """ + def init_from_pretrained(self) -> None: """Load pretrained weights at cold start (no checkpoint to restore). diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 64ec58ab5..3910e3f9e 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -43,6 +43,7 @@ from tzrec.prompt.assembler import ( PROMPT_MAX_SEQLEN as _PROMPT_MAX_SEQLEN, ) +from tzrec.prompt.persist import save_prompt_assets from tzrec.prompt.plan import CompiledPrompt, SlotSeg from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.models.prompt_model_pb2 import PromptModelConfig @@ -268,6 +269,14 @@ def _forward_loss( ) return {"loss": loss} + def save_assets(self, target_dir: str) -> None: + """Co-locate the prompt contract, so the checkpoint is self-describing. + + Args: + target_dir: the checkpoint or export directory. + """ + save_prompt_assets(self._prompt, target_dir) + def init_from_pretrained(self) -> None: """Load HF weights once, on a cold start only.""" source = self._model_config.hf_model_id diff --git a/tzrec/prompt/persist.py b/tzrec/prompt/persist.py new file mode 100644 index 000000000..296783e87 --- /dev/null +++ b/tzrec/prompt/persist.py @@ -0,0 +1,124 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Writes the prompt contract beside the weights, and checks it on restore. + +A checkpoint that cannot describe its own vocabulary is a checkpoint serving +has to be told about out of band, which is where offline/online skew comes +from. ``ModulePlan`` is deliberately absent: it is model-only and rebuilt from +config at every ``__init__``. +""" + +import dataclasses +import json +import os +import shutil +from enum import Enum +from typing import Any, Dict, Optional + +from tzrec.prompt.plan import CompiledPrompt +from tzrec.utils.logging_util import logger + +PROMPT_DIR = "prompt" +_SID_SPACE = "sid_space.json" +_PROMPT_PLAN = "prompt_plan.json" +_HASHES = "prompt_hashes.json" +_TOKENIZER = "tokenizer" + + +def _plain(value: Any) -> Any: + """Render a compiled artifact as JSON-safe values.""" + if isinstance(value, Enum): + return value.value + if dataclasses.is_dataclass(value) and not isinstance(value, type): + return { + f.name: _plain(getattr(value, f.name)) for f in dataclasses.fields(value) + } + if isinstance(value, (list, tuple)): + return [_plain(v) for v in value] + if isinstance(value, dict): + return {str(k): _plain(v) for k, v in value.items()} + return value + + +def save_prompt_assets(prompt: CompiledPrompt, target_dir: str) -> None: + """Write the prompt contract into a checkpoint or export directory. + + Args: + prompt: the compiled prompt. + target_dir: the checkpoint or export directory. + """ + out = os.path.join(target_dir, PROMPT_DIR) + os.makedirs(out, exist_ok=True) + + with open(os.path.join(out, _SID_SPACE), "w") as f: + json.dump(_plain(prompt.sid_space), f, indent=2) + with open(os.path.join(out, _PROMPT_PLAN), "w") as f: + json.dump(_plain(prompt.prompt_plan), f, indent=2) + with open(os.path.join(out, _HASHES), "w") as f: + json.dump( + {"vocab_hash": prompt.vocab_hash, "plan_hash": prompt.plan_hash}, + f, + indent=2, + ) + + if prompt.tokenizer_dir and os.path.isdir(prompt.tokenizer_dir): + shutil.copytree( + prompt.tokenizer_dir, os.path.join(out, _TOKENIZER), dirs_exist_ok=True + ) + + +def read_prompt_hashes(source_dir: str) -> Optional[Dict[str, str]]: + """Read the hashes a checkpoint recorded, or None when it has none.""" + path = os.path.join(source_dir, PROMPT_DIR, _HASHES) + if not os.path.exists(path): + return None + with open(path, "r") as f: + return json.load(f) + + +def check_prompt_assets(prompt: Optional[CompiledPrompt], ckpt_dir: str) -> None: + """Compare a compiled prompt against what a checkpoint recorded. + + A ``vocab_hash`` mismatch is fatal: the decode bands would point at token + ranges the weights never learned, which produces plausible output rather + than an error. A ``plan_hash`` mismatch only reshapes the prompt, so it + warns. + + Args: + prompt: the freshly compiled prompt, or None when the pipeline declares + no prompt_config. + ckpt_dir: the checkpoint being restored. + """ + if prompt is None: + return + recorded = read_prompt_hashes(ckpt_dir) + if recorded is None: + logger.warning( + f"checkpoint [{ckpt_dir}] records no prompt assets, so its " + f"vocabulary cannot be checked against the current prompt_config." + ) + return + + if recorded.get("vocab_hash") != prompt.vocab_hash: + raise ValueError( + f"prompt vocabulary does not match checkpoint [{ckpt_dir}]: the " + f"checkpoint was trained against {recorded.get('vocab_hash')} but " + f"prompt_config now compiles to {prompt.vocab_hash}. The SID space " + f"or the tokenizer changed, so the decode bands no longer address " + f"the rows these weights learned." + ) + if recorded.get("plan_hash") != prompt.plan_hash: + logger.warning( + f"prompt plan differs from checkpoint [{ckpt_dir}]: the vocabulary " + f"matches, so the weights are usable, but the template, slots or " + f"projections changed." + ) diff --git a/tzrec/prompt/persist_test.py b/tzrec/prompt/persist_test.py new file mode 100644 index 000000000..a8e2a2b49 --- /dev/null +++ b/tzrec/prompt/persist_test.py @@ -0,0 +1,119 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import dataclasses +import json +import os +import unittest + +from google.protobuf import text_format +from tokenizers import Tokenizer, models, pre_tokenizers + +from tzrec.features.feature import FgMode, create_features +from tzrec.prompt.compile import compile_prompt +from tzrec.prompt.persist import ( + PROMPT_DIR, + check_prompt_assets, + read_prompt_hashes, + save_prompt_assets, +) +from tzrec.protos import feature_pb2 +from tzrec.protos.prompt_pb2 import PromptConfig +from tzrec.utils.test_util import make_test_dir + +_WORDS = ["History", "Predict", ":", "", "<|im_end|>"] + + +class PromptPersistTest(unittest.TestCase): + def setUp(self) -> None: + self.test_dir = make_test_dir() + tok_path = os.path.join(self.test_dir, "tok.json") + tok = Tokenizer( + models.WordLevel( + vocab={w: i for i, w in enumerate(_WORDS)}, unk_token="" + ) + ) + tok.pre_tokenizer = pre_tokenizers.Whitespace() + tok.save(tok_path) + self.tok_path = tok_path + + fc = feature_pb2.FeatureConfig() + text_format.Merge( + 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }', fc + ) + self.features = create_features([fc], fg_mode=FgMode.FG_NONE) + + def _compile(self, codebook=(4, 4, 4), prompt="History : {{hist}}"): + cfg = PromptConfig(tokenizer=self.tok_path, prompt=prompt) + cfg.sid_space.codebook.extend(codebook) + return compile_prompt(cfg, self.features, model_dir=self.test_dir) + + def test_writes_a_self_describing_directory(self) -> None: + prompt = self._compile() + ckpt = os.path.join(self.test_dir, "model.ckpt-1") + save_prompt_assets(prompt, ckpt) + + out = os.path.join(ckpt, PROMPT_DIR) + for name in ("sid_space.json", "prompt_plan.json", "prompt_hashes.json"): + self.assertTrue(os.path.exists(os.path.join(out, name)), name) + # serving reloads the extended tokenizer from the checkpoint + self.assertTrue( + os.path.exists(os.path.join(out, "tokenizer", "tokenizer.json")) + ) + + def test_sid_space_round_trips_as_plain_json(self) -> None: + prompt = self._compile() + ckpt = os.path.join(self.test_dir, "model.ckpt-1") + save_prompt_assets(prompt, ckpt) + + with open(os.path.join(ckpt, PROMPT_DIR, "sid_space.json")) as f: + space = json.load(f) + self.assertEqual(space["codebook"], [4, 4, 4]) + self.assertEqual(space["level_offsets"], [0, 4, 8]) + self.assertEqual(space["band_lo"][0], prompt.sid_space.base_vocab) + # every declared field survives, so serving needs no tzrec code + self.assertEqual( + set(space), {f.name for f in dataclasses.fields(prompt.sid_space)} + ) + + def test_matching_prompt_passes(self) -> None: + prompt = self._compile() + ckpt = os.path.join(self.test_dir, "model.ckpt-1") + save_prompt_assets(prompt, ckpt) + check_prompt_assets(self._compile(), ckpt) + + def test_a_changed_codebook_is_fatal(self) -> None: + ckpt = os.path.join(self.test_dir, "model.ckpt-1") + save_prompt_assets(self._compile(codebook=(4, 4, 4)), ckpt) + with self.assertRaisesRegex(ValueError, "does not match checkpoint"): + check_prompt_assets(self._compile(codebook=(8, 8, 8)), ckpt) + + def test_a_changed_template_only_warns(self) -> None: + ckpt = os.path.join(self.test_dir, "model.ckpt-1") + save_prompt_assets(self._compile(), ckpt) + moved = self._compile(prompt="Predict : {{hist}}") + # the vocabulary is untouched, so the weights are still usable + self.assertEqual(moved.vocab_hash, read_prompt_hashes(ckpt)["vocab_hash"]) + self.assertNotEqual(moved.plan_hash, read_prompt_hashes(ckpt)["plan_hash"]) + check_prompt_assets(moved, ckpt) + + def test_a_checkpoint_without_assets_only_warns(self) -> None: + bare = os.path.join(self.test_dir, "model.ckpt-bare") + os.makedirs(bare, exist_ok=True) + self.assertIsNone(read_prompt_hashes(bare)) + check_prompt_assets(self._compile(), bare) + + def test_no_prompt_config_is_a_no_op(self) -> None: + check_prompt_assets(None, os.path.join(self.test_dir, "nowhere")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/utils/checkpoint_util.py b/tzrec/utils/checkpoint_util.py index f26b7a6f3..c8a11017a 100644 --- a/tzrec/utils/checkpoint_util.py +++ b/tzrec/utils/checkpoint_util.py @@ -332,6 +332,23 @@ def best_checkpoint( return latest_checkpoint(model_dir) +def _unwrap_model(model: nn.Module) -> nn.Module: + """Walk DMP/TrainWrapper layers down to the model that owns the hooks.""" + inner = model + seen = set() + while not hasattr(inner, "save_assets"): + if id(inner) in seen: + return model + seen.add(id(inner)) + if hasattr(inner, "module"): + inner = inner.module + elif hasattr(inner, "model"): + inner = inner.model + else: + return model + return inner + + class CheckpointManager: """Saves training checkpoints and prunes old ones asynchronously. @@ -411,6 +428,13 @@ def save( f"write_hf_assets failed for {ckpt_dir}: {e} -- checkpoint " f"weights are saved; skipping HF assets." ) + try: + _unwrap_model(model).save_assets(ckpt_dir) + except Exception as e: # noqa: BLE001 + logger.warning( + f"save_assets failed for {ckpt_dir}: {e} -- checkpoint weights " + f"are saved; skipping model assets." + ) if dataloader_state is not None: save_dataloader_state(ckpt_dir, dataloader_state) self._last_ckpt_dir = ckpt_dir From c6ec0e2ac47a39425270e91336c3cd6cdf08868c Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 08:07:09 +0000 Subject: [PATCH 67/99] [bugfix] prompt: fix four gaps a real training run exposed Unit tests passed while the pipeline could not run. Running it found: train_and_evaluate never passed prompt= to _create_model, so the model raised at construction. The earlier wiring matched other call sites by shape and missed this one. assemble_into read a dense sequence feature's values as flat, but the parser emits (total, value_dim); the (n, 1) slices made a ragged list. It now flattens, with a test that feeds the parser's real shape. The model implemented none of init_loss, loss, init_metric, update_metric or update_train_metric, which TrainWrapper and the eval loop all call. write_hf_assets required an hf_tokenizer. A prompt-native model owns none: its extended vocabulary is a separately versioned artifact that save_assets writes to prompt/tokenizer, so the tokenizer step is now optional. Verified end to end: three steps train, eval reports ce_loss, the checkpoint carries prompt/ assets, resume works, and a changed codebook is refused at restore. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/main.py | 2 ++ tzrec/models/prompt_generative_qwen.py | 49 ++++++++++++++++++++++++++ tzrec/prompt/assembler.py | 4 ++- tzrec/prompt/assembler_test.py | 22 ++++++++++++ tzrec/utils/hf_export_util.py | 6 +++- 5 files changed, 81 insertions(+), 2 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index 6460ed102..a0fd841d6 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -799,6 +799,7 @@ def train_and_evaluate( list(data_config.label_fields), sample_weights=list(data_config.sample_weight_fields), sampler_type=sampler_type, + prompt=prompt, ) # Cold start only; a resumed or fine-tuned run gets its weights from DCP. if ckpt_path is None: @@ -1005,6 +1006,7 @@ def evaluate( list(data_config.label_fields), sample_weights=list(data_config.sample_weight_fields), sampler_type=sampler_type, + prompt=prompt, ) model = TrainWrapper( model, device=device, mixed_precision=train_config.mixed_precision diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 3910e3f9e..02c0b289a 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -20,6 +20,7 @@ from typing import Any, Dict, List, Optional import torch +import torchmetrics from torch import nn from transformers import AutoConfig, AutoModelForCausalLM @@ -269,6 +270,54 @@ def _forward_loss( ) return {"loss": loss} + def init_loss(self) -> None: + """No-op: the LM computes its own CE inside ``predict``.""" + return + + def loss( + self, predictions: Dict[str, torch.Tensor], batch: Batch + ) -> Dict[str, torch.Tensor]: + """Surface the CE already computed in ``predict``. + + Args: + predictions: what ``predict`` returned. + batch: the batch, unused. + + Returns: + The named loss. + """ + return {"ce_loss": predictions["loss"]} + + def init_metric(self) -> None: + """Register a mean-CE metric for the eval loop.""" + self._metric_modules["ce_loss"] = torchmetrics.MeanMetric() + + def update_metric( + self, + predictions: Dict[str, torch.Tensor], + batch: Batch, + losses: Optional[Dict[str, torch.Tensor]] = None, + ) -> None: + """Update the mean-CE metric with this batch's loss. + + Args: + predictions: what ``predict`` returned. + batch: the batch, unused. + losses: the named losses, unused. + """ + self._metric_modules["ce_loss"].update(predictions["loss"].detach()) + + def update_train_metric( + self, predictions: Dict[str, torch.Tensor], batch: Batch + ) -> None: + """No-op: nothing beyond the logged CE. + + Args: + predictions: what ``predict`` returned. + batch: the batch, unused. + """ + return + def save_assets(self, target_dir: str) -> None: """Co-locate the prompt contract, so the checkpoint is self-describing. diff --git a/tzrec/prompt/assembler.py b/tzrec/prompt/assembler.py index a83879067..d4917ebd1 100644 --- a/tzrec/prompt/assembler.py +++ b/tzrec/prompt/assembler.py @@ -228,7 +228,9 @@ def assemble_into( lengths = np.asarray(parsed[f"{source}.lengths"]) batch_size = max(batch_size, int(lengths.size)) if seg.fill is FillMode.INLINE: - flat = np.asarray(parsed[f"{source}.values"]) + # a dense sequence feature emits (total, value_dim); the stream is + # one code per position, so value_dim is always 1 here + flat = np.asarray(parsed[f"{source}.values"]).reshape(-1) bounds = np.concatenate(([0], np.cumsum(lengths))) values[seg.name] = [ flat[bounds[i] : bounds[i + 1]] for i in range(lengths.size) diff --git a/tzrec/prompt/assembler_test.py b/tzrec/prompt/assembler_test.py index 660236a33..ca32eb1ef 100644 --- a/tzrec/prompt/assembler_test.py +++ b/tzrec/prompt/assembler_test.py @@ -165,6 +165,28 @@ def test_inline_without_a_sid_space_is_rejected_at_construction(self) -> None: with self.assertRaisesRegex(ValueError, "no sid_space was compiled"): PromptAssembler(plan, None) + def test_column_shaped_values_are_flattened(self) -> None: + # the data parser emits (total, value_dim) for a dense sequence feature + from tzrec.prompt.assembler import assemble_into + from tzrec.prompt.plan import CompiledPrompt, ModulePlan + + plan = _plan((_slot("hist", FillMode.INLINE),)) + prompt = CompiledPrompt( + sid_space=_sid_space(), + prompt_plan=plan, + module_plan=ModulePlan(projections={}, slot_to_module={}), + tokenizer_dir="", + vocab_hash="v", + plan_hash="p", + ) + parsed = { + "hist.values": np.array([[1], [6], [11], [0], [4], [8]]), + "hist.lengths": np.array([3, 3]), + } + out = assemble_into(prompt, parsed) + self.assertEqual(out["prompt_cu_seqlens"].tolist(), [0, 3, 6]) + self.assertEqual(out["prompt_input_ids"].tolist()[0], _BASE + 1) + def test_rows_of_different_lengths_pack_without_padding(self) -> None: plan = _plan((_slot("hist", FillMode.INLINE),)) asm = PromptAssembler(plan, _sid_space()) diff --git a/tzrec/utils/hf_export_util.py b/tzrec/utils/hf_export_util.py index 39bee5cf6..f46eb9f9e 100644 --- a/tzrec/utils/hf_export_util.py +++ b/tzrec/utils/hf_export_util.py @@ -81,7 +81,11 @@ def write_hf_assets(wrapped_model: nn.Module, save_dir: str) -> None: gen_cfg = getattr(backbone, "generation_config", None) if gen_cfg is not None: gen_cfg.save_pretrained(save_dir) - inner.hf_tokenizer().save_pretrained(save_dir) + # A prompt-native model owns no HF tokenizer: its extended vocabulary is a + # separately versioned artifact that save_assets writes to prompt/tokenizer. + tokenizer = getattr(inner, "hf_tokenizer", None) + if tokenizer is not None: + tokenizer().save_pretrained(save_dir) # named_modules() FQNs carry the DMP prefix that state_dict() strips. raw_prefix = next( From a589b7fa589fee92b0ac008802266b659b19c9ae Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 08:17:25 +0000 Subject: [PATCH 68/99] [bugfix] prompt: derive the feature groups a projected slot needs A PROJECTED slot reaches the model through EmbeddingGroup, but nothing built one: the compiler produced a ModulePlan without the FeatureGroupConfig it implies, so the model raised on embedding_group at construction. Step 12 of the compile algorithm was specified and never implemented; no unit test caught it because every earlier test used an INLINE-only prompt. The compiler now derives one group per projected slot and the model builds its EmbeddingGroup from them. Derived rather than declared, because a prompt group is never shared with a model tower and most of FeatureGroupConfig is meaningless here. Verified by training both shapes end to end. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_qwen.py | 4 ++ tzrec/prompt/compile.py | 17 ++++++- tzrec/prompt/plan.py | 6 ++- tzrec/tests/prompt_integration_test.py | 70 +++++++++++++++++++++++++- 4 files changed, 92 insertions(+), 5 deletions(-) diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 02c0b289a..aa2286251 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -28,6 +28,7 @@ from tzrec.features.feature import BaseFeature from tzrec.models.model import BaseModel from tzrec.modules.dynamic_beam import dynamic_beam_search +from tzrec.modules.embedding import EmbeddingGroup from tzrec.modules.prompt_projection import PromptProjection from tzrec.prompt.assembler import ( PROMPT_CU_SEQLENS as _PROMPT_CU_SEQLENS, @@ -96,6 +97,9 @@ def __init__( self.lm.resize_token_embeddings( prompt.sid_space.target_vocab, mean_resizing=True ) + self.embedding_group = EmbeddingGroup( + self._features, list(self._prompt.module_plan.feature_groups) + ) self._build_projections() def _read_beam_config(self, common: PromptModelConfig) -> None: diff --git a/tzrec/prompt/compile.py b/tzrec/prompt/compile.py index 450f7ff57..8b9b0f374 100644 --- a/tzrec/prompt/compile.py +++ b/tzrec/prompt/compile.py @@ -37,7 +37,7 @@ Width, WidthKind, ) -from tzrec.protos.model_pb2 import FeatureGroupType +from tzrec.protos.model_pb2 import FeatureGroupConfig, FeatureGroupType from tzrec.protos.prompt_pb2 import PromptConfig, PromptProjection, PromptSlot from tzrec.utils.logging_util import logger @@ -360,7 +360,20 @@ def _build_module_plan( else: projections[module_id] = projection slot_to_module[seg.slot_id] = module_id - return ModulePlan(projections=projections, slot_to_module=slot_to_module) + + groups = tuple( + FeatureGroupConfig( + group_name=seg.name, + feature_names=list(seg.sources), + group_type=seg.group_type, + ) + for seg in projected + ) + return ModulePlan( + projections=projections, + slot_to_module=slot_to_module, + feature_groups=groups, + ) def _max_total_length(segments: Sequence[Segment]) -> Optional[int]: diff --git a/tzrec/prompt/plan.py b/tzrec/prompt/plan.py index 9fae4c5d6..fe40ce0c8 100644 --- a/tzrec/prompt/plan.py +++ b/tzrec/prompt/plan.py @@ -20,7 +20,7 @@ from enum import Enum from typing import Mapping, Optional, Tuple, Union -from tzrec.protos.model_pb2 import FeatureGroupType +from tzrec.protos.model_pb2 import FeatureGroupConfig, FeatureGroupType from tzrec.protos.prompt_pb2 import PromptProjection @@ -182,10 +182,14 @@ class ModulePlan: projections: resolved module id to its configuration. slot_to_module: slot id to the module id it uses, so slots sharing a ``projection_name`` resolve to one module. + feature_groups: one derived group per PROJECTED slot. Derived rather + than declared: a prompt group is never shared with a model tower, + and four of FeatureGroupConfig's six fields are meaningless here. """ projections: Mapping[str, PromptProjection] slot_to_module: Mapping[int, str] + feature_groups: Tuple[FeatureGroupConfig, ...] = () @dataclass(frozen=True) diff --git a/tzrec/tests/prompt_integration_test.py b/tzrec/tests/prompt_integration_test.py index 6709edb36..93f1af10e 100644 --- a/tzrec/tests/prompt_integration_test.py +++ b/tzrec/tests/prompt_integration_test.py @@ -55,11 +55,17 @@ def _tokenizer(path: str) -> str: return path -def _features(): +_PROF = ( + 'sequence_id_feature { feature_name: "prof" expression: "user:prof" ' + "num_buckets: 32 embedding_dim: 8 sequence_length: 4 }" +) + + +def _features(extra=()): text = ( 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }', 'sequence_raw_feature { feature_name: "answer" expression: "item:answer" }', - ) + ) + tuple(extra) out = [] for one in text: fc = feature_pb2.FeatureConfig() @@ -174,5 +180,65 @@ def test_raw_codes_are_rejected_before_the_model_sees_them(self) -> None: assemble_into(self.prompt, parsed) +class ProjectedSlotTest(unittest.TestCase): + """A slot whose value reaches the LM through a table and a projection.""" + + def setUp(self) -> None: + self.test_dir = make_test_dir() + self.backbone = _tiny_backbone(os.path.join(self.test_dir, "backbone")) + self.tok = _tokenizer(os.path.join(self.test_dir, "tok.json")) + self.features = _features(extra=(_PROF,)) + + cfg = PromptConfig( + tokenizer=self.tok, + prompt="History : {{hist}} . Predict {{prof}} :", + response="{{answer}}", + ) + cfg.sid_space.codebook.extend(_CODEBOOK) + self.prompt = compile_prompt(cfg, self.features, model_dir=self.test_dir) + + def test_compiler_derives_one_group_per_projected_slot(self) -> None: + groups = self.prompt.module_plan.feature_groups + self.assertEqual([g.group_name for g in groups], ["prof"]) + self.assertEqual(list(groups[0].feature_names), ["prof"]) + # the INLINE slot produces none: its tokens are already in the stream + self.assertEqual( + [s.name for s in self.prompt.prompt_plan.projected_slots], ["prof"] + ) + + def test_sentinel_is_materialized_and_holes_recorded(self) -> None: + space = self.prompt.sid_space + self.assertIsNotNone(space.sentinel_token_id) + + parsed = { + "hist.values": torch.tensor(_offset([0, 1, 2])).reshape(-1, 1), + "hist.lengths": torch.tensor([3]), + "answer.values": torch.tensor(_offset([1, 2, 3])), + "answer.lengths": torch.tensor([3]), + "prof.values": torch.tensor([5, 9]), + "prof.lengths": torch.tensor([2]), + } + streams = assemble_into(self.prompt, parsed) + # two profile items -> two sentinels -> two holes + self.assertEqual(streams["prompt_hole_positions"].tolist(), [7, 8]) + ids = streams["prompt_input_ids"] + self.assertTrue(all(ids[p] == space.sentinel_token_id for p in [7, 8])) + + def test_projection_receives_gradient(self) -> None: + model_config = ModelConfig() + qwen = model_config.prompt_generative_qwen + qwen.hf_model_id = self.backbone + qwen.common.beam_widths.extend([2, 2, 2]) + qwen.common.num_return_sequences = 2 + model = _create_model( + model_config, self.features, ["answer"], prompt=self.prompt + ) + self.assertEqual(len(model.projections), 1) + + # the scatter is what puts the projection on the autograd path at all + proj = next(iter(model.projections.values())) + self.assertIsNone(proj.head.weight.grad) + + if __name__ == "__main__": unittest.main() From 8bcf3d68463522765db2485fcbf406557c19f574 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 08:21:19 +0000 Subject: [PATCH 69/99] [bugfix] prompt: make predict work Two gaps, both found by running tzrec.predict rather than by a unit test. predict_checkpoint never passed prompt= to _create_model, so the model raised at construction. That is the third call site with its own shape; all four are now explicit. PredictPipelineSparseDist FX-traces the model, and beam decode reads host ints and branches on them, so tracing died on a Proxy. The decode loop is now a single torch.fx.wrap leaf. One leaf rather than a wrapped inner helper, because every host read inside the loop is untraceable and wrapping one only moves the failure to the next. Verified on both checkpoint shapes: 64 rows in, 64 out, generated_sids of shape (num_return, num_levels), every code inside [0, codebook) -- so the bands held and detokenize inverted both shifts. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/main.py | 1 + tzrec/models/prompt_generative_qwen.py | 22 ++++++++++++++++------ 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index a0fd841d6..7adfb795e 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -1591,6 +1591,7 @@ def predict_checkpoint( pipeline_config.model_config, features, [], + prompt=prompt, ) model.set_is_inference(True) model = PredictWrapper( diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index aa2286251..8a24b8aaa 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -183,10 +183,9 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: Returns: The loss when training, the decoded SIDs otherwise. """ - embeds = self._prompt_embeds(batch) if self.is_inference: - return {self._generated_sids_key: self._generate(embeds, batch)} - return self._forward_loss(embeds, batch) + return {self._generated_sids_key: _fx_wrapped_generate(self, batch)} + return self._forward_loss(self._prompt_embeds(batch), batch) def _sid_token_bands(self) -> "tuple[torch.Tensor, torch.Tensor]": """Inclusive token-id band of every SID level, as device tensors.""" @@ -197,16 +196,16 @@ def _sid_token_bands(self) -> "tuple[torch.Tensor, torch.Tensor]": torch.tensor(space.band_hi, device=device), ) - def _generate(self, embeds: torch.Tensor, batch: Batch) -> torch.Tensor: + def _generate(self, batch: Batch) -> torch.Tensor: """Beam-search the SID answer. Args: - embeds: the assembled prompt embeddings, packed. - batch: carries ``prompt_cu_seqlens`` and the collator's width. + batch: carries the packed prompt and the collator's width. Returns: ``(B, num_return, num_levels)`` local codes, best first. """ + embeds = self._prompt_embeds(batch) infos = batch.additional_infos padded, mask, _ = _unpack( embeds, @@ -374,3 +373,14 @@ def _unpack( padded[mask] = embeds out_labels[mask] = labels return padded, mask.long(), out_labels + + +@torch.fx.wrap +def _fx_wrapped_generate(model: "PromptGenerativeQwen", batch: Batch) -> torch.Tensor: + """Hide the decode loop from FX. + + ``PredictPipelineSparseDist`` FX-traces the model, and beam decode reads + host ints and branches on them. Wrapping an inner helper only moves the + failure to the next such read, so the whole loop is one leaf. + """ + return model._generate(batch) From 6149ffb2b7cf37086fb3b0d6ccb48988a8fd19be Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 08:26:08 +0000 Subject: [PATCH 70/99] [bugfix] prompt: make export self-describing and refuse TorchScript The HF branch returns before a model exists, so save_assets never ran and the exported directory carried weights with no vocabulary: serving would have had to reach back into model_dir for the SID space and tokenizer. The branch now copies the checkpoint's prompt/ forward, so one directory is the whole contract. TorchScript export of a prompt-native model died on a missing prompt_input_ids. That is architectural, not a bug: the model's input is an assembled token stream the dataloader builds, and export has no dataloader. The design exports the LM as a HuggingFace directory and reserves TorchScript for the prompt front end, so the format is now refused up front with a message that says which knob to set. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/main.py | 15 ++++++++++++++- tzrec/prompt/persist.py | 17 +++++++++++++++++ tzrec/prompt/persist_test.py | 22 ++++++++++++++++++++++ 3 files changed, 53 insertions(+), 1 deletion(-) diff --git a/tzrec/main.py b/tzrec/main.py index 7adfb795e..12ccf16cf 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -74,7 +74,7 @@ from tzrec.optim.lr_scheduler import BaseLR from tzrec.optim.optimizer import TZRecOptimizer from tzrec.prompt.compile import compile_prompt -from tzrec.prompt.persist import check_prompt_assets +from tzrec.prompt.persist import check_prompt_assets, copy_prompt_assets from tzrec.prompt.plan import CompiledPrompt from tzrec.protos import export_pb2 from tzrec.protos.data_pb2 import DataConfig, DatasetType @@ -1124,6 +1124,16 @@ def export( checkpoint_path, _ = ckpt_manager.latest_checkpoint() # HF export converts the checkpoint dir directly -- no model build, no DCP restore. + if pipeline_config.HasField("prompt_config") and ( + pipeline_config.export_config.export_format != export_pb2.ExportFormat.HF + ): + raise ValueError( + "a prompt-native model exports to a HuggingFace directory, not " + "TorchScript: its input is an assembled token stream the dataloader " + "builds, which an export-time dummy batch cannot supply. Set " + "export_config.export_format to HF." + ) + if pipeline_config.export_config.export_format == export_pb2.ExportFormat.HF: if config_util.use_dense_ema( pipeline_config.export_config, pipeline_config.train_config @@ -1146,6 +1156,9 @@ def export( from tzrec.utils.hf_export_util import dcp_to_hf dcp_to_hf(checkpoint_path, export_dir) + # this branch never builds a model, so save_assets cannot run; the + # checkpoint already carries the contract, so copy it forward + copy_prompt_assets(checkpoint_path, export_dir) return data_config = pipeline_config.data_config diff --git a/tzrec/prompt/persist.py b/tzrec/prompt/persist.py index 296783e87..cf53b8ebf 100644 --- a/tzrec/prompt/persist.py +++ b/tzrec/prompt/persist.py @@ -122,3 +122,20 @@ def check_prompt_assets(prompt: Optional[CompiledPrompt], ckpt_dir: str) -> None f"matches, so the weights are usable, but the template, slots or " f"projections changed." ) + + +def copy_prompt_assets(source_dir: str, target_dir: str) -> None: + """Carry a checkpoint's prompt contract into an export directory. + + Args: + source_dir: the checkpoint being exported. + target_dir: the export directory. + """ + source = os.path.join(source_dir, PROMPT_DIR) + if not os.path.isdir(source): + logger.warning( + f"checkpoint [{source_dir}] carries no prompt assets, so the export " + f"will not describe its own vocabulary." + ) + return + shutil.copytree(source, os.path.join(target_dir, PROMPT_DIR), dirs_exist_ok=True) diff --git a/tzrec/prompt/persist_test.py b/tzrec/prompt/persist_test.py index a8e2a2b49..6cd0fdc20 100644 --- a/tzrec/prompt/persist_test.py +++ b/tzrec/prompt/persist_test.py @@ -22,6 +22,7 @@ from tzrec.prompt.persist import ( PROMPT_DIR, check_prompt_assets, + copy_prompt_assets, read_prompt_hashes, save_prompt_assets, ) @@ -114,6 +115,27 @@ def test_a_checkpoint_without_assets_only_warns(self) -> None: def test_no_prompt_config_is_a_no_op(self) -> None: check_prompt_assets(None, os.path.join(self.test_dir, "nowhere")) + def test_export_carries_the_contract_forward(self) -> None: + ckpt = os.path.join(self.test_dir, "model.ckpt-1") + save_prompt_assets(self._compile(), ckpt) + export = os.path.join(self.test_dir, "export") + copy_prompt_assets(ckpt, export) + + # the HF branch never builds a model, so save_assets cannot run there + self.assertEqual(read_prompt_hashes(export), read_prompt_hashes(ckpt)) + self.assertTrue( + os.path.exists( + os.path.join(export, PROMPT_DIR, "tokenizer", "tokenizer.json") + ) + ) + + def test_copying_from_a_bare_checkpoint_only_warns(self) -> None: + bare = os.path.join(self.test_dir, "bare") + os.makedirs(bare, exist_ok=True) + export = os.path.join(self.test_dir, "export_bare") + copy_prompt_assets(bare, export) + self.assertIsNone(read_prompt_hashes(export)) + if __name__ == "__main__": unittest.main() From 21c70304e7a653df98a1a3b3db01a2762ca48296 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 08:33:58 +0000 Subject: [PATCH 71/99] [doc] add the prompt-native generative recommendation manual Covers what an operator has to get right: reading offset_codebook rather than codebook or origin_codebook, the config surface, how a slot's fill mode is derived rather than configured, the artifacts a checkpoint carries, and why export is HuggingFace-only. The troubleshooting section is keyed on the exact error strings the code raises, including the two that are fatal by design -- a wrong SID column and a vocabulary that no longer matches the checkpoint. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/models/generative.rst | 1 + docs/source/models/prompt_generative_qwen.md | 176 +++++++++++++++++++ 2 files changed, 177 insertions(+) create mode 100644 docs/source/models/prompt_generative_qwen.md diff --git a/docs/source/models/generative.rst b/docs/source/models/generative.rst index 9e0aa39b2..e92016094 100644 --- a/docs/source/models/generative.rst +++ b/docs/source/models/generative.rst @@ -5,5 +5,6 @@ :maxdepth: 2 dlrm_hstu + prompt_generative_qwen ultra_hstu hstu_match diff --git a/docs/source/models/prompt_generative_qwen.md b/docs/source/models/prompt_generative_qwen.md new file mode 100644 index 000000000..97ff74b5f --- /dev/null +++ b/docs/source/models/prompt_generative_qwen.md @@ -0,0 +1,176 @@ +# Prompt 原生生成式推荐（PromptGenerativeQwen） + +以 Qwen 为骨干，把用户历史的语义 ID（SID）拼进 prompt，让模型直接生成下一个物品的 SID。 + +prompt 的模板、槽位、SID 空间与词表都由新的 `prompt_config` 描述，`model_config` 只保留属于 LM 的部分。 + +## 1. 数据准备 + +SID 必须以 **offset 形式**进入 tzrec，即 SID 生成工具 `resolve_sid_collisions` 输出的 `offset_codebook` 列： + +``` +第 l 层的取值 = level_offsets[l] + code code 属于 [0, codebook[l]) +``` + +以 `codebook: 4 4 4` 为例，`level_offsets` 为 `[0, 4, 8]`，因此一个 item 的三层取值分别落在 `[0,4)`、`[4,8)`、`[8,12)`。 + +```{warning} +只读 `offset_codebook`。`codebook` 与 `origin_codebook` 两列同样格式合法，但前者未加 offset、后者是冲突解析**之前**的 SID；误用不会报格式错，而是训练出静默错误的模型。assembler 的 band 校验能挡住未加 offset 的列，但挡不住手工对 `origin_codebook` 施加 offset 得到的流。 +``` + +一行历史是若干个 item 的三层 code 依次拼平，长度必须是层数的整数倍。 + +## 2. 配置 + +一个最小可运行的配置： + +``` +data_config { + batch_size: 4 + dataset_type: ParquetDataset + fg_mode: FG_NONE + label_fields: "answer" +} + +feature_configs { + sequence_raw_feature { feature_name: "hist" expression: "user:hist" } +} +feature_configs { + sequence_raw_feature { feature_name: "answer" expression: "item:answer" } +} + +prompt_config { + tokenizer: "path/to/tokenizer.json" + prompt: "用户历史行为为：{{hist}}。请预测下一个商品：" + response: "{{answer}}" + sid_space { codebook: 256 codebook: 256 codebook: 256 } + max_length: 4096 +} + +model_config { + prompt_generative_qwen { + hf_model_id: "Qwen/Qwen2.5-0.5B" + common { + beam_widths: 100 + beam_widths: 200 + beam_widths: 400 + num_return_sequences: 50 + } + } +} +``` + +### prompt_config + +| 字段 | 说明 | +| ------------------------- | ------------------------------------------------------------------------------------------------------- | +| `tokenizer` | **基础** tokenizer 的路径或 hub id。注意它与 `hf_model_id` 不同：后者只表示权重，且只在冷启动时读取一次 | +| `prompt` | 模板。`{{name}}` 之间的静态文本自动成为相邻槽位的前后缀，无需逐槽位配置 | +| `response` | 监督目标。定义 loss 覆盖的范围；推理时不生成该段 | +| `sid_space.codebook` | 每层的 SID 词表大小 | +| `sid_space.manifest_path` | 可选。指向 SID manifest，编译期与 `codebook` 逐元素比对，不一致直接报错 | +| `max_length` | 校验上限，**不是**截断开关：超长的行会报错，不会被截断 | + +### 槽位如何被推导 + +`{{name}}` 默认解析为同名特征。槽位的填充方式不需要配置，由成员特征推导： + +| 槽位成员 | 填充方式 | 说明 | +| -------------------------------------------------------- | --------- | --------------------------------------------------- | +| 单个序列特征且不声明 embedding（`sequence_raw_feature`） | INLINE | SID 直接进入 token 流，与答案共享 embedding | +| 其他情形（如 `sequence_id_feature`、标量特征、多成员） | PROJECTED | 走自己的 embedding 表，再经一次投影抵达 LM 隐层维度 | + +PROJECTED 槽位在 token 流中占位为 sentinel，真实取值在前向时写入对应位置。 + +### model_config + +| 字段 | 说明 | +| ---------------------- | ----------------------------------------------------------- | +| `hf_model_id` | 预训练权重的 hub id 或本地目录 | +| `beam_widths` | 每层一个宽度，长度必须等于 `codebook` 的层数 | +| `num_return_sequences` | 不得超过最后一层的宽度 | +| `param_dtype` | 主权重精度，默认 FP32。bf16 会让 Adam 的小更新在 ULP 下丢失 | + +## 3. 训练 + +```bash +torchrun --master_addr=localhost --master_port=32555 --nnodes=1 --nproc-per-node=2 --node_rank=0 \ + -m tzrec.train_eval --pipeline_config_path prompt_qwen.config +``` + +续跑加 `--continue_train`。 + +每个 `model.ckpt-N/` 除权重外还会写出 `prompt/` 目录： + +``` +model.ckpt-N/prompt/ + sid_space.json 解析后的 SID 空间：codebook、level_offsets、band、target_vocab + prompt_plan.json assembler 的遍历顺序与各项上界 + prompt_hashes.json vocab_hash 与 plan_hash + tokenizer/ 扩展后的 tokenizer（含 SID atom） +``` + +即 checkpoint 自带词表契约，服务端无需另行配置。 + +## 4. 预测 + +```bash +torchrun --master_addr=localhost --master_port=32555 --nnodes=1 --nproc-per-node=1 --node_rank=0 \ + -m tzrec.predict --pipeline_config_path experiments/run/pipeline.config \ + --predict_input_path 'data/*.parquet' --predict_output_path out +``` + +输出列 `generated_sids`，形状为 `(num_return_sequences, 层数)`，取值是**局部 0-based** code，可直接与 SID 映射表的 `codebook` 列对齐。 + +## 5. 导出 + +只支持导出为 HuggingFace 目录： + +``` +export_config { export_format: HF } +``` + +```bash +torchrun ... -m tzrec.export --pipeline_config_path experiments/run/pipeline.config \ + --export_dir exported +``` + +产出目录同时包含权重与 prompt 契约，可直接被 `AutoModelForCausalLM.from_pretrained` 加载： + +``` +exported/ + config.json generation_config.json model.safetensors + prompt/ sid_space.json prompt_plan.json prompt_hashes.json tokenizer/ +``` + +```{note} +本模型不支持 TorchScript 导出。它的输入是 dataloader 组装出的 token 流，而导出期的伪造 batch 无法提供。配置为默认的 TORCHSCRIPT 时会直接报错并提示改用 HF。 +``` + +## 6. 常见问题 + +**`SID values must already carry their level offset ... Read the offset_codebook column`** + +读错了列。改用 `offset_codebook`，见第 1 节。 + +**`prompt vocabulary does not match checkpoint`** + +`codebook`、`atom_token_format` 或 tokenizer 变了，与该 checkpoint 训练时的词表不一致。这是硬失败：解码 band 会指向这批权重从未学过的行，继续跑只会产出看似合理的错误结果。要么改回原配置，要么从头训练。 + +若只是模板或槽位变了（`plan_hash` 不同、`vocab_hash` 相同），只会告警，权重仍可用。 + +**`beam_widths has N entries but the codebook has M levels`** + +每层一个宽度，两者长度必须相等。 + +**`assembled row X is N tokens, over max_length`** + +超长的行不会被截断。请在特征上用 `sequence_length` 限制历史长度，而不是调大 `max_length`。 + +**`static_prefix_len is 0`（告警）** + +模板开头就是一个槽位，导致服务端前缀缓存无内容可共享。把静态指令文本放在最前、变长槽位放在最后即可。 + +**`a prompt-native model exports to a HuggingFace directory, not TorchScript`** + +见第 5 节，设置 `export_config.export_format: HF`。 From dce8ed17bd935dd210f11fa1bc77db276be03a74 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 10:29:36 +0000 Subject: [PATCH 72/99] [bugfix] prompt: write checkpoint assets from rank 0 only save_assets ran on every rank, so a multi-rank save had all of them doing json.dump and copytree to the same paths at once. The files survived a 2-rank local run, but concurrent writes to one path can interleave into a truncated artifact, and a checkpoint whose sid_space.json is truncated is one that restore cannot validate. Guarded in save_prompt_assets, symmetric with write_hf_assets. copy_prompt_assets needs none: the export path already calls it inside is_rank_zero. Verified at --nproc-per-node=2 for both the INLINE-only and projected configs, train and predict. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/prompt/persist.py | 5 +++++ tzrec/prompt/persist_test.py | 12 ++++++++++++ 2 files changed, 17 insertions(+) diff --git a/tzrec/prompt/persist.py b/tzrec/prompt/persist.py index cf53b8ebf..06184e2a1 100644 --- a/tzrec/prompt/persist.py +++ b/tzrec/prompt/persist.py @@ -52,10 +52,15 @@ def _plain(value: Any) -> Any: def save_prompt_assets(prompt: CompiledPrompt, target_dir: str) -> None: """Write the prompt contract into a checkpoint or export directory. + Rank 0 only: every rank reaches here on save, and concurrent json.dump and + copytree to one path can interleave into a truncated file. + Args: prompt: the compiled prompt. target_dir: the checkpoint or export directory. """ + if int(os.environ.get("RANK", 0)) != 0: + return out = os.path.join(target_dir, PROMPT_DIR) os.makedirs(out, exist_ok=True) diff --git a/tzrec/prompt/persist_test.py b/tzrec/prompt/persist_test.py index 6cd0fdc20..159249f57 100644 --- a/tzrec/prompt/persist_test.py +++ b/tzrec/prompt/persist_test.py @@ -13,6 +13,7 @@ import json import os import unittest +from unittest import mock from google.protobuf import text_format from tokenizers import Tokenizer, models, pre_tokenizers @@ -136,6 +137,17 @@ def test_copying_from_a_bare_checkpoint_only_warns(self) -> None: copy_prompt_assets(bare, export) self.assertIsNone(read_prompt_hashes(export)) + def test_only_rank_zero_writes(self) -> None: + ckpt = os.path.join(self.test_dir, "model.ckpt-rank") + with mock.patch.dict(os.environ, {"RANK": "1"}): + save_prompt_assets(self._compile(), ckpt) + # a non-zero rank must not race rank 0's json.dump and copytree + self.assertFalse(os.path.exists(os.path.join(ckpt, PROMPT_DIR))) + + with mock.patch.dict(os.environ, {"RANK": "0"}): + save_prompt_assets(self._compile(), ckpt) + self.assertIsNotNone(read_prompt_hashes(ckpt)) + if __name__ == "__main__": unittest.main() From 53e1062a721b83db02caa506826efa3cd66b6eaf Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 5 Aug 2026 10:34:10 +0000 Subject: [PATCH 73/99] [bugfix] prompt: bound the supervised logits window suffix_keep was always None, so the model scored logits over every position instead of the answer. It came out None because the answer slot's width was derived from sequence_length, which the answer feature does not declare and should not have to: the answer is exactly one SID item, so its width is the codebook depth. At the toy vocab the smoke runs use, the difference is invisible. At a real vocabulary it is not: full-length logits are batch x length x vocab in fp32, which is terabytes where the window is megabytes, so the first training step would have failed to allocate. The width now comes from the codebook, and an unbounded response is rejected at compile rather than silently falling back to scoring everything. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/prompt/compile.py | 29 +++++++++++++++++++++++++---- tzrec/prompt/compile_test.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 4 deletions(-) diff --git a/tzrec/prompt/compile.py b/tzrec/prompt/compile.py index 8b9b0f374..6b82b68ea 100644 --- a/tzrec/prompt/compile.py +++ b/tzrec/prompt/compile.py @@ -72,13 +72,19 @@ def _resolve_slot(name: str, declared: Dict[str, PromptSlot]) -> PromptSlot: def _slot_width( - members: Sequence[BaseFeature], group_type: "FeatureGroupType.ValueType" + members: Sequence[BaseFeature], + group_type: "FeatureGroupType.ValueType", + answer_levels: Optional[int] = None, ) -> Width: """Derive a slot's position count from its members. - A DEEP slot pools to exactly one position. A sequence slot is bounded by - the members' ``sequence_length``, and unbounded when none declares one. + A DEEP slot pools to exactly one position. The answer is exactly one SID + item, so its width is the codebook depth and needs no sequence_length. Any + other sequence slot is bounded by its members' sequence_length, and + unbounded when none declares one. """ + if answer_levels is not None: + return Width(WidthKind.STATIC, answer_levels) if group_type == FeatureGroupType.DEEP: return Width(WidthKind.STATIC, 1) caps = [ @@ -274,9 +280,17 @@ def compile_prompt( tok.save(os.path.join(tokenizer_dir, "tokenizer.json")) slot_ids = {n: i for i, n in enumerate(slots)} + answer_names = set(resp_names) segs: Dict[str, SlotSeg] = {} for name, slot in slots.items(): seq = types[name] == FeatureGroupType.JAGGED_SEQUENCE + levels = ( + sid_space.num_levels + if sid_space is not None + and name in answer_names + and fills[name] is FillMode.INLINE + else None + ) segs[name] = SlotSeg( slot_id=slot_ids[name], name=name, @@ -284,7 +298,7 @@ def compile_prompt( group_type=types[name], output_key=".sequence" if seq else "", fill=fills[name], - width=_slot_width(members[name], types[name]), + width=_slot_width(members[name], types[name], levels), droppable=bool(slot.drop_if_empty), ) @@ -451,6 +465,13 @@ def _validate( f"static_prefix_len is {plan.static_prefix_len}, which bounds " "what a serving prefix cache may reuse." ) + if plan.response_segments and plan.suffix_keep is None: + raise ValueError( + "the response has an unbounded slot, so the supervised logits " + "window cannot be bounded. A decoder-only model would then " + "materialize logits for every position, which is (batch x length x " + "vocab) and will not fit. Give the response slot a fixed width." + ) if plan.static_prefix_len == 0: logger.warning( "static_prefix_len is 0: no leading run of the prompt is " diff --git a/tzrec/prompt/compile_test.py b/tzrec/prompt/compile_test.py index d5f1e2806..f6c65bca3 100644 --- a/tzrec/prompt/compile_test.py +++ b/tzrec/prompt/compile_test.py @@ -197,6 +197,35 @@ def test_extended_tokenizer_is_written(self) -> None: self.assertIsNotNone(reloaded.token_to_id("<|sid_0|>")) self.assertIsNotNone(reloaded.token_to_id("<|sid_7|>")) + def test_answer_width_comes_from_the_codebook(self) -> None: + cfg = self._config(prompt="History : {{hist}}", response="{{answer}}") + cfg.sid_space.codebook.extend([4, 4, 4]) + answer = _feature( + 'sequence_raw_feature { feature_name: "answer" expression: "item:answer" }' + ) + compiled = self._compile(cfg, [_feature(_HIST), answer]) + + seg = next( + s for s in compiled.prompt_plan.response_segments if isinstance(s, SlotSeg) + ) + # the answer is one SID item, so its width needs no sequence_length + self.assertIs(seg.width.kind, WidthKind.STATIC) + self.assertEqual(seg.width.n, 3) + # +1 because HF shifts logits: the window opens one column before the + # first supervised label + self.assertEqual(compiled.prompt_plan.suffix_keep, 4) + + def test_unbounded_response_is_rejected(self) -> None: + # with no sid_space the response has no codebook-derived width, so the + # supervised window is unbounded and the logits would cover every + # position + cfg = self._config(prompt="History : {{hist}}", response="{{answer}}") + answer = _feature( + 'sequence_raw_feature { feature_name: "answer" expression: "item:answer" }' + ) + with self.assertRaisesRegex(ValueError, "window cannot be bounded"): + self._compile(cfg, [_feature(_HIST), answer]) + if __name__ == "__main__": unittest.main() From ea45b36c584921cb65379b932df491421ad637f4 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Thu, 6 Aug 2026 11:31:27 +0000 Subject: [PATCH 74/99] [refactor] prompt: pass the compiled prompt directly _create_model built a kwargs dict to avoid handing prompt= to models that do not declare it. BaseModule already absorbs unknown kwargs and forwards none to nn.Module, so a None prompt reaching an unrelated model was never a problem and the conditional guarded nothing. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/main.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index 9b6bb63f3..f1e4fdc74 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -176,14 +176,13 @@ def _create_model( # pyre-ignore [16] model_cls = BaseModel.create_class(model_cls_name) - extra: Dict[str, Any] = {"prompt": prompt} if prompt is not None else {} model: BaseModel = model_cls( model_config, features, labels, sample_weights=sample_weights, sampler_type=sampler_type, - **extra, + prompt=prompt, ) kernel = Kernel[KernelProto.Name(model_config.kernel)] From 16e4a5c90e992ac2750783be50e8e48ba87b4b8e Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Thu, 6 Aug 2026 11:51:19 +0000 Subject: [PATCH 75/99] [refactor] prompt: split out BasePromptGenerativeModel Everything independent of how a family runs its transformer moves to the base: building the LM empty, resizing to the compiled vocabulary, wiring slot projections, the SID coordinate conversions, and the loss, metric and checkpoint hooks. The subclass keeps predict, the teacher-forced loss, the decode loop and the parameters only those need -- ignore_index, generated_sids_key and the beam schedule. Those differ irreducibly rather than by configuration: a decoder-only family reaches past lm(...) into body and head so logits cover a suffix window, while an encoder-decoder would pass labels and get a loss back. Making them abstract states that, where a parameterized hook would imply the difference is one of degree. Llama and Mistral share Qwen's attribute layout exactly, so they need no subclass of their own; the split earns its keep only when a family with a different forward arrives. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_model.py | 284 ++++++++++++++++++++++ tzrec/models/prompt_generative_qwen.py | 303 ++++++------------------ 2 files changed, 357 insertions(+), 230 deletions(-) create mode 100644 tzrec/models/prompt_generative_model.py diff --git a/tzrec/models/prompt_generative_model.py b/tzrec/models/prompt_generative_model.py new file mode 100644 index 000000000..5591c6246 --- /dev/null +++ b/tzrec/models/prompt_generative_model.py @@ -0,0 +1,284 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Backbone-agnostic half of a prompt-native generative model. + +Everything here is independent of how a family runs its transformer: building +the LM empty, resizing to the compiled vocabulary, wiring slot projections, +converting between the SID coordinate systems, and the checkpoint hooks. + +What a subclass owns is the forward: ``predict``, the teacher-forced loss and +the decode loop. Those differ irreducibly -- a decoder-only model reaches past +``lm(...)`` into body and head so it can score a suffix window, while an +encoder-decoder passes labels and gets a loss back -- so they are abstract here +rather than parameterized. +""" + +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torchmetrics +from torch import nn +from transformers import AutoConfig, AutoModelForCausalLM + +from tzrec.datasets.utils import Batch +from tzrec.features.feature import BaseFeature +from tzrec.models.model import BaseModel +from tzrec.modules.embedding import EmbeddingGroup +from tzrec.modules.prompt_projection import PromptProjection +from tzrec.prompt.assembler import ( + PROMPT_HOLE_POSITIONS, + PROMPT_INPUT_IDS, +) +from tzrec.prompt.persist import save_prompt_assets +from tzrec.prompt.plan import CompiledPrompt, SlotSeg +from tzrec.protos.model_pb2 import ModelConfig +from tzrec.protos.models.prompt_model_pb2 import PromptModelConfig +from tzrec.utils.logging_util import logger + +_PARAM_DTYPE: Dict[int, torch.dtype] = { + PromptModelConfig.FP32: torch.float32, + PromptModelConfig.BF16: torch.bfloat16, + PromptModelConfig.FP16: torch.float16, +} + + +class BasePromptGenerativeModel(BaseModel): + """An HF backbone driven by a compiled prompt. + + Args: + model_config: the model oneof. + features: every created feature. + labels: data_config label fields. + sample_weights: optional sample weight fields. + prompt: the compiled prompt; required. + """ + + def __init__( + self, + model_config: ModelConfig, + features: List[BaseFeature], + labels: List[str], + sample_weights: Optional[List[str]] = None, + prompt: Optional[CompiledPrompt] = None, + **kwargs: Any, + ) -> None: + super().__init__(model_config, features, labels, sample_weights, **kwargs) + if prompt is None: + raise ValueError( + f"{type(self).__name__} needs a compiled prompt; call " + f"compile_prompt(pipeline_config.prompt_config, features) and " + f"pass it to _create_model." + ) + if prompt.sid_space is None: + raise ValueError( + f"{type(self).__name__}: prompt_config declares no sid_space, " + f"so there is no SID vocabulary to extend or decode." + ) + self._prompt = prompt + cfg = self._model_config + + self.lm = self._build_backbone(cfg.hf_model_id, cfg.common.param_dtype) + self.lm.resize_token_embeddings( + prompt.sid_space.target_vocab, mean_resizing=True + ) + self.embedding_group = EmbeddingGroup( + self._features, list(self._prompt.module_plan.feature_groups) + ) + self._build_projections() + + def _build_backbone(self, hf_model_id: str, param_dtype: int) -> nn.Module: + """Build the LM empty, so HF weights load only on cold start. + + Args: + hf_model_id: hub id or local directory naming the weights. + param_dtype: master-weight dtype. + + Returns: + The uninitialized backbone. + """ + config = AutoConfig.from_pretrained(hf_model_id) + model = AutoModelForCausalLM.from_config(config) + return model.to(_PARAM_DTYPE[param_dtype]) + + def _build_projections(self) -> None: + """One module per resolved id, aligned with ``plan.projected_slots``. + + Slots sharing a ``projection_name`` share a module by reference, so + they must agree on ``group_total_dim``. + """ + plan = self._prompt.prompt_plan + modules = self._prompt.module_plan + hidden = int(self.lm.config.hidden_size) + + built: Dict[str, PromptProjection] = {} + aligned: List[PromptProjection] = [] + for seg in plan.projected_slots: + module_id = modules.slot_to_module[seg.slot_id] + in_dim = self._slot_in_dim(seg) + if module_id not in built: + built[module_id] = PromptProjection( + modules.projections[module_id], in_dim, hidden + ) + elif built[module_id].in_dim != in_dim: + raise ValueError( + f"prompt slots sharing projection_name [{module_id}] have " + f"different group widths ({built[module_id].in_dim} vs " + f"{in_dim}); they cannot share a module." + ) + aligned.append(built[module_id]) + self.projections = nn.ModuleDict(built) + # zipped with plan.projected_slots; shared modules appear by reference + self._slot_projections = aligned + + def _slot_in_dim(self, seg: SlotSeg) -> int: + """Total group output width of a projected slot. + + Args: + seg: the projected slot. + + Returns: + Its ``group_total_dim``. + """ + return self.embedding_group.group_total_dim(seg.name + seg.output_key) + + def hf_backbone(self) -> nn.Module: + """The HF module export and checkpointing reach for.""" + return self.lm + + def _prompt_embeds(self, batch: Batch) -> torch.Tensor: + """Gather the token stream, then overwrite the projected positions. + + Args: + batch: carries the packed prompt in ``additional_infos``. + + Returns: + ``(total_tokens, hidden)``. + """ + ids = batch.additional_infos[PROMPT_INPUT_IDS] + embeds = self.lm.get_input_embeddings()(ids) + + plan = self._prompt.prompt_plan + if not plan.projected_slots: + return embeds + + grouped = self.embedding_group(batch) + hidden = embeds.shape[-1] + parts = [ + proj(grouped[seg.name + seg.output_key]).reshape(-1, hidden) + for seg, proj in zip(plan.projected_slots, self._slot_projections) + ] + # out of place: embeds carries grad from the embedding lookup + return embeds.index_copy( + 0, batch.additional_infos[PROMPT_HOLE_POSITIONS], torch.cat(parts) + ) + + def _sid_token_bands(self) -> Tuple[torch.Tensor, torch.Tensor]: + """Inclusive token-id band of every SID level, as device tensors.""" + space = self._prompt.sid_space + device = self.lm.get_input_embeddings().weight.device + return ( + torch.tensor(space.band_lo, device=device), + torch.tensor(space.band_hi, device=device), + ) + + def _detokenize(self, tokens: torch.Tensor, batch_size: int) -> torch.Tensor: + """Undo both shifts: token id back to a local 0-based code. + + Args: + tokens: generated token ids, ``(batch_size * beams, num_levels)``. + batch_size: rows in the batch. + + Returns: + ``(batch_size, beams, num_levels)`` local codes. + """ + space = self._prompt.sid_space + offsets = torch.tensor(space.level_offsets, device=tokens.device) + codes = tokens - space.base_vocab - offsets + return codes.view(batch_size, -1, space.num_levels) + + def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: + """Run the model over an assembled prompt. + + Args: + batch: carries the packed prompt in ``additional_infos``. + + Returns: + The loss when training, the decoded SIDs otherwise. + """ + raise NotImplementedError + + def init_loss(self) -> None: + """No-op: an LM computes its own CE inside ``predict``.""" + return + + def loss( + self, predictions: Dict[str, torch.Tensor], batch: Batch + ) -> Dict[str, torch.Tensor]: + """Surface the CE already computed in ``predict``. + + Args: + predictions: what ``predict`` returned. + batch: the batch, unused. + + Returns: + The named loss. + """ + return {"ce_loss": predictions["loss"]} + + def init_metric(self) -> None: + """Register a mean-CE metric for the eval loop.""" + self._metric_modules["ce_loss"] = torchmetrics.MeanMetric() + + def update_metric( + self, + predictions: Dict[str, torch.Tensor], + batch: Batch, + losses: Optional[Dict[str, torch.Tensor]] = None, + ) -> None: + """Update the mean-CE metric with this batch's loss. + + Args: + predictions: what ``predict`` returned. + batch: the batch, unused. + losses: the named losses, unused. + """ + self._metric_modules["ce_loss"].update(predictions["loss"].detach()) + + def update_train_metric( + self, predictions: Dict[str, torch.Tensor], batch: Batch + ) -> None: + """No-op: nothing beyond the logged CE. + + Args: + predictions: what ``predict`` returned. + batch: the batch, unused. + """ + return + + def save_assets(self, target_dir: str) -> None: + """Co-locate the prompt contract, so the checkpoint is self-describing. + + Args: + target_dir: the checkpoint or export directory. + """ + save_prompt_assets(self._prompt, target_dir) + + def init_from_pretrained(self) -> None: + """Load HF weights once, on a cold start only.""" + source = self._model_config.hf_model_id + logger.info(f"loading pretrained weights from [{source}].") + pretrained = AutoModelForCausalLM.from_pretrained(source) + pretrained.resize_token_embeddings( + self._prompt.sid_space.target_vocab, mean_resizing=True + ) + self.lm.load_state_dict(pretrained.state_dict()) + del pretrained diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 8a24b8aaa..237f32984 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -9,57 +9,36 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Prompt-native generative recommendation over a Qwen backbone. +"""Decoder-only forward and decode over a Qwen backbone. -The model reads ``target_vocab``, the module plan and the decode bands from a -``CompiledPrompt``, and never reads the prompt's structure. Assembly happens in -the dataloader worker; what arrives here is already a packed token stream plus -the positions the projected slots must overwrite. +Everything backbone-agnostic is in ``BasePromptGenerativeModel``. What is here +is what a decoder-only family does differently: it reaches past ``lm(...)`` into +``lm.model`` and ``lm.lm_head`` so logits are materialized for a suffix window +only, and it decodes by prefilling once and stepping a self-attention cache. + +The layout is Qwen's, and Llama and Mistral share it exactly. """ -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple import torch -import torchmetrics -from torch import nn -from transformers import AutoConfig, AutoModelForCausalLM from tzrec.datasets.utils import Batch from tzrec.features.feature import BaseFeature -from tzrec.models.model import BaseModel +from tzrec.models.prompt_generative_model import BasePromptGenerativeModel from tzrec.modules.dynamic_beam import dynamic_beam_search -from tzrec.modules.embedding import EmbeddingGroup -from tzrec.modules.prompt_projection import PromptProjection -from tzrec.prompt.assembler import ( - PROMPT_CU_SEQLENS as _PROMPT_CU_SEQLENS, -) -from tzrec.prompt.assembler import ( - PROMPT_HOLE_POSITIONS as _PROMPT_HOLE_POSITIONS, -) -from tzrec.prompt.assembler import ( - PROMPT_INPUT_IDS as _PROMPT_INPUT_IDS, -) -from tzrec.prompt.assembler import ( - PROMPT_LABELS as _PROMPT_LABELS, -) from tzrec.prompt.assembler import ( - PROMPT_MAX_SEQLEN as _PROMPT_MAX_SEQLEN, + PROMPT_CU_SEQLENS, + PROMPT_LABELS, + PROMPT_MAX_SEQLEN, ) -from tzrec.prompt.persist import save_prompt_assets -from tzrec.prompt.plan import CompiledPrompt, SlotSeg +from tzrec.prompt.plan import CompiledPrompt from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.models.prompt_model_pb2 import PromptModelConfig -from tzrec.utils.logging_util import logger -_PARAM_DTYPE: Dict[int, torch.dtype] = { - PromptModelConfig.FP32: torch.float32, - PromptModelConfig.BF16: torch.bfloat16, - PromptModelConfig.FP16: torch.float16, -} - -class PromptGenerativeQwen(BaseModel): - """Qwen backbone driven by a compiled prompt. +class PromptGenerativeQwen(BasePromptGenerativeModel): + """Qwen family (Qwen2.5, Qwen3, ...) driven by a compiled prompt. Args: model_config: the model oneof. @@ -78,38 +57,21 @@ def __init__( prompt: Optional[CompiledPrompt] = None, **kwargs: Any, ) -> None: - super().__init__(model_config, features, labels, sample_weights, **kwargs) - if prompt is None: - raise ValueError( - f"{type(self).__name__} needs a compiled prompt; call " - f"compile_prompt(pipeline_config.prompt_config, features) and " - f"pass it to _create_model." - ) - self._prompt = prompt - cfg = self._model_config - common = cfg.common - + super().__init__( + model_config, features, labels, sample_weights, prompt, **kwargs + ) + common = self._model_config.common self._ignore_index = int(common.ignore_index) self._generated_sids_key = common.generated_sids_key self._read_beam_config(common) - self.lm = self._build_backbone(cfg.hf_model_id, common.param_dtype) - self.lm.resize_token_embeddings( - prompt.sid_space.target_vocab, mean_resizing=True - ) - self.embedding_group = EmbeddingGroup( - self._features, list(self._prompt.module_plan.feature_groups) - ) - self._build_projections() - def _read_beam_config(self, common: PromptModelConfig) -> None: - """Parse the decode knobs; the schedule must match the codebook.""" + """Parse the decode knobs; the schedule must match the codebook. + + Args: + common: the shared model config. + """ space = self._prompt.sid_space - if space is None: - raise ValueError( - f"{type(self).__name__}: prompt_config declares no sid_space, " - f"so there is nothing to decode." - ) self._num_return = int(common.num_return_sequences) self._beam_widths: List[int] = list(common.beam_widths) if not self._beam_widths: @@ -130,50 +92,6 @@ def _read_beam_config(self, common: PromptModelConfig) -> None: f"({self._beam_widths[-1]})." ) - def _build_backbone(self, hf_model_id: str, param_dtype: int) -> nn.Module: - """Build the LM empty, so HF weights load only on cold start.""" - config = AutoConfig.from_pretrained(hf_model_id) - model = AutoModelForCausalLM.from_config(config) - return model.to(_PARAM_DTYPE[param_dtype]) - - def _build_projections(self) -> None: - """One module per resolved id, aligned with ``plan.projected_slots``. - - Slots sharing a ``projection_name`` share a module by reference, so - they must agree on ``group_total_dim``. - """ - plan = self._prompt.prompt_plan - modules = self._prompt.module_plan - hidden = int(self.lm.config.hidden_size) - - built: Dict[str, PromptProjection] = {} - aligned: List[PromptProjection] = [] - for seg in plan.projected_slots: - module_id = modules.slot_to_module[seg.slot_id] - in_dim = self._slot_in_dim(seg) - if module_id not in built: - built[module_id] = PromptProjection( - modules.projections[module_id], in_dim, hidden - ) - elif built[module_id].in_dim != in_dim: - raise ValueError( - f"prompt slots sharing projection_name [{module_id}] have " - f"different group widths ({built[module_id].in_dim} vs " - f"{in_dim}); they cannot share a module." - ) - aligned.append(built[module_id]) - self.projections = nn.ModuleDict(built) - # zipped with plan.projected_slots; shared modules appear by reference - self._slot_projections = aligned - - def _slot_in_dim(self, seg: SlotSeg) -> int: - """Total group output width of a projected slot.""" - return self.embedding_group.group_total_dim(seg.name + seg.output_key) - - def hf_backbone(self) -> nn.Module: - """The HF module export and checkpointing reach for.""" - return self.lm - def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: """Teacher-forced forward over the assembled stream. @@ -187,77 +105,27 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: return {self._generated_sids_key: _fx_wrapped_generate(self, batch)} return self._forward_loss(self._prompt_embeds(batch), batch) - def _sid_token_bands(self) -> "tuple[torch.Tensor, torch.Tensor]": - """Inclusive token-id band of every SID level, as device tensors.""" - space = self._prompt.sid_space - device = self.lm.get_input_embeddings().weight.device - return ( - torch.tensor(space.band_lo, device=device), - torch.tensor(space.band_hi, device=device), - ) + def _forward_loss( + self, embeds: torch.Tensor, batch: Batch + ) -> Dict[str, torch.Tensor]: + """Run the LM over the assembled embeddings and score the response. - def _generate(self, batch: Batch) -> torch.Tensor: - """Beam-search the SID answer. + Body and head are called separately so logits cover the supervised + window only: a full (batch, length, vocab) upcast does not fit. Args: - batch: carries the packed prompt and the collator's width. + embeds: the assembled prompt embeddings, packed. + batch: carries the row boundaries and the collator's width. Returns: - ``(B, num_return, num_levels)`` local codes, best first. + The loss. """ - embeds = self._prompt_embeds(batch) - infos = batch.additional_infos - padded, mask, _ = _unpack( - embeds, - infos[_PROMPT_CU_SEQLENS], - infos[_PROMPT_LABELS], - int(infos[_PROMPT_MAX_SEQLEN]), - self._ignore_index, - ) - lo_tok, hi_tok = self._sid_token_bands() - tokens = dynamic_beam_search( - self.lm, padded, mask, self._beam_widths, lo_tok, hi_tok - ) - return self._detokenize(tokens, padded.shape[0]) - - def _detokenize(self, tokens: torch.Tensor, batch_size: int) -> torch.Tensor: - """Undo both shifts: token id back to a local 0-based code.""" - space = self._prompt.sid_space - offsets = torch.tensor(space.level_offsets, device=tokens.device) - codes = tokens - space.base_vocab - offsets - codes = codes.view(batch_size, -1, space.num_levels) - return codes[:, : self._num_return, :] - - def _prompt_embeds(self, batch: Batch) -> torch.Tensor: - """Gather the token stream, then overwrite the projected positions.""" - ids = batch.additional_infos[_PROMPT_INPUT_IDS] - embeds = self.lm.get_input_embeddings()(ids) - - plan = self._prompt.prompt_plan - if not plan.projected_slots: - return embeds - - grouped = self.embedding_group(batch) - hidden = embeds.shape[-1] - parts = [ - proj(grouped[seg.name + seg.output_key]).reshape(-1, hidden) - for seg, proj in zip(plan.projected_slots, self._slot_projections) - ] - # out of place: embeds carries grad from the embedding lookup - return embeds.index_copy( - 0, batch.additional_infos[_PROMPT_HOLE_POSITIONS], torch.cat(parts) - ) - - def _forward_loss( - self, embeds: torch.Tensor, batch: Batch - ) -> Dict[str, torch.Tensor]: - """Run the LM over the assembled embeddings and score the response.""" infos = batch.additional_infos padded, mask, labels = _unpack( embeds, - infos[_PROMPT_CU_SEQLENS], - infos[_PROMPT_LABELS], - int(infos[_PROMPT_MAX_SEQLEN]), + infos[PROMPT_CU_SEQLENS], + infos[PROMPT_LABELS], + int(infos[PROMPT_MAX_SEQLEN]), self._ignore_index, ) outputs = self.lm.model(inputs_embeds=padded, attention_mask=mask) @@ -273,72 +141,30 @@ def _forward_loss( ) return {"loss": loss} - def init_loss(self) -> None: - """No-op: the LM computes its own CE inside ``predict``.""" - return - - def loss( - self, predictions: Dict[str, torch.Tensor], batch: Batch - ) -> Dict[str, torch.Tensor]: - """Surface the CE already computed in ``predict``. + def _generate(self, batch: Batch) -> torch.Tensor: + """Beam-search the SID answer. Args: - predictions: what ``predict`` returned. - batch: the batch, unused. + batch: carries the packed prompt and the collator's width. Returns: - The named loss. - """ - return {"ce_loss": predictions["loss"]} - - def init_metric(self) -> None: - """Register a mean-CE metric for the eval loop.""" - self._metric_modules["ce_loss"] = torchmetrics.MeanMetric() - - def update_metric( - self, - predictions: Dict[str, torch.Tensor], - batch: Batch, - losses: Optional[Dict[str, torch.Tensor]] = None, - ) -> None: - """Update the mean-CE metric with this batch's loss. - - Args: - predictions: what ``predict`` returned. - batch: the batch, unused. - losses: the named losses, unused. - """ - self._metric_modules["ce_loss"].update(predictions["loss"].detach()) - - def update_train_metric( - self, predictions: Dict[str, torch.Tensor], batch: Batch - ) -> None: - """No-op: nothing beyond the logged CE. - - Args: - predictions: what ``predict`` returned. - batch: the batch, unused. - """ - return - - def save_assets(self, target_dir: str) -> None: - """Co-locate the prompt contract, so the checkpoint is self-describing. - - Args: - target_dir: the checkpoint or export directory. + ``(B, num_return, num_levels)`` local codes, best first. """ - save_prompt_assets(self._prompt, target_dir) - - def init_from_pretrained(self) -> None: - """Load HF weights once, on a cold start only.""" - source = self._model_config.hf_model_id - logger.info(f"loading pretrained weights from [{source}].") - pretrained = AutoModelForCausalLM.from_pretrained(source) - pretrained.resize_token_embeddings( - self._prompt.sid_space.target_vocab, mean_resizing=True + embeds = self._prompt_embeds(batch) + infos = batch.additional_infos + padded, mask, _ = _unpack( + embeds, + infos[PROMPT_CU_SEQLENS], + infos[PROMPT_LABELS], + int(infos[PROMPT_MAX_SEQLEN]), + self._ignore_index, + ) + lo_tok, hi_tok = self._sid_token_bands() + tokens = dynamic_beam_search( + self.lm, padded, mask, self._beam_widths, lo_tok, hi_tok ) - self.lm.load_state_dict(pretrained.state_dict()) - del pretrained + codes = self._detokenize(tokens, padded.shape[0]) + return codes[:, : self._num_return, :] def _unpack( @@ -347,12 +173,22 @@ def _unpack( labels: torch.Tensor, max_seqlen: int, ignore_index: int, -) -> "tuple[torch.Tensor, torch.Tensor, torch.Tensor]": +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Pad a packed varlen batch at the LM boundary. Padding lives in this one adapter. ``max_seqlen`` is the collator's, not ``lengths.max()``: deriving it here would sync the device to the host every - step, which §7.4 of the design forbids. + step, which the design forbids. + + Args: + embeds: packed embeddings, ``(total_tokens, hidden)``. + cu_seqlens: row boundaries, ``(batch_size + 1,)``. + labels: packed labels, ``(total_tokens,)``. + max_seqlen: the collator's padded width. + ignore_index: label value for padding. + + Returns: + Padded embeddings, attention mask and labels. """ starts = cu_seqlens[:-1] lengths = cu_seqlens[1:] - starts @@ -382,5 +218,12 @@ def _fx_wrapped_generate(model: "PromptGenerativeQwen", batch: Batch) -> torch.T ``PredictPipelineSparseDist`` FX-traces the model, and beam decode reads host ints and branches on them. Wrapping an inner helper only moves the failure to the next such read, so the whole loop is one leaf. + + Args: + model: the model whose decode loop to run. + batch: the batch to decode. + + Returns: + The decoded local codes. """ return model._generate(batch) From f9b1af67fa39a01b8c20ffcd651b99b5eae610b6 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Thu, 6 Aug 2026 12:06:17 +0000 Subject: [PATCH 76/99] [refactor] prompt: cleanup pass over the new stack Collapses the two identical wrapper-unwind walks into one checkpoint_util.unwrap_to(model, attr); hf_export_util had a byte-for-byte copy differing only in the marker attribute. Stops round-tripping constants through the device. The SID band edges were built as tensors and immediately read back with int(), which is 2*num_levels host syncs per decode to recover values the caller already had; they now stay Python ints. The level offsets decode subtracts every step become a buffer, and the assembler's per-level bounds are hoisted out of the per-row path. Drops what nothing reads: PromptPlan.slot_index, Static.owner_slot_id (always None, so its documented drop behaviour never existed), PromptProjection.in_dim and the one-line _slot_in_dim wrapper. Inference no longer builds and scatters a label tensor it discards. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_model.py | 41 +++++++++---------------- tzrec/models/prompt_generative_qwen.py | 23 +++++++++----- tzrec/modules/dynamic_beam.py | 15 +++------ tzrec/modules/dynamic_beam_test.py | 3 +- tzrec/modules/prompt_projection.py | 6 ---- tzrec/prompt/assembler.py | 7 +++-- tzrec/prompt/assembler_test.py | 13 +++----- tzrec/prompt/compile.py | 3 +- tzrec/prompt/plan.py | 5 --- tzrec/tests/prompt_integration_test.py | 15 ++------- tzrec/utils/checkpoint_util.py | 25 +++++++++++---- tzrec/utils/hf_export_util.py | 24 +-------------- tzrec/utils/hf_export_util_test.py | 18 ++++++----- 13 files changed, 79 insertions(+), 119 deletions(-) diff --git a/tzrec/models/prompt_generative_model.py b/tzrec/models/prompt_generative_model.py index 5591c6246..6a6f5c0c0 100644 --- a/tzrec/models/prompt_generative_model.py +++ b/tzrec/models/prompt_generative_model.py @@ -22,7 +22,7 @@ rather than parameterized. """ -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional import torch import torchmetrics @@ -39,7 +39,7 @@ PROMPT_INPUT_IDS, ) from tzrec.prompt.persist import save_prompt_assets -from tzrec.prompt.plan import CompiledPrompt, SlotSeg +from tzrec.prompt.plan import CompiledPrompt from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.models.prompt_model_pb2 import PromptModelConfig from tzrec.utils.logging_util import logger @@ -94,6 +94,12 @@ def __init__( self._features, list(self._prompt.module_plan.feature_groups) ) self._build_projections() + # decode subtracts these every step; a buffer follows the module's device + self.register_buffer( + "_level_offsets", + torch.tensor(prompt.sid_space.level_offsets), + persistent=False, + ) def _build_backbone(self, hf_model_id: str, param_dtype: int) -> nn.Module: """Build the LM empty, so HF weights load only on cold start. @@ -120,18 +126,20 @@ def _build_projections(self) -> None: hidden = int(self.lm.config.hidden_size) built: Dict[str, PromptProjection] = {} + widths: Dict[str, int] = {} aligned: List[PromptProjection] = [] for seg in plan.projected_slots: module_id = modules.slot_to_module[seg.slot_id] - in_dim = self._slot_in_dim(seg) + in_dim = self.embedding_group.group_total_dim(seg.name + seg.output_key) if module_id not in built: built[module_id] = PromptProjection( modules.projections[module_id], in_dim, hidden ) - elif built[module_id].in_dim != in_dim: + widths[module_id] = in_dim + elif widths[module_id] != in_dim: raise ValueError( f"prompt slots sharing projection_name [{module_id}] have " - f"different group widths ({built[module_id].in_dim} vs " + f"different group widths ({widths[module_id]} vs " f"{in_dim}); they cannot share a module." ) aligned.append(built[module_id]) @@ -139,17 +147,6 @@ def _build_projections(self) -> None: # zipped with plan.projected_slots; shared modules appear by reference self._slot_projections = aligned - def _slot_in_dim(self, seg: SlotSeg) -> int: - """Total group output width of a projected slot. - - Args: - seg: the projected slot. - - Returns: - Its ``group_total_dim``. - """ - return self.embedding_group.group_total_dim(seg.name + seg.output_key) - def hf_backbone(self) -> nn.Module: """The HF module export and checkpointing reach for.""" return self.lm @@ -181,15 +178,6 @@ def _prompt_embeds(self, batch: Batch) -> torch.Tensor: 0, batch.additional_infos[PROMPT_HOLE_POSITIONS], torch.cat(parts) ) - def _sid_token_bands(self) -> Tuple[torch.Tensor, torch.Tensor]: - """Inclusive token-id band of every SID level, as device tensors.""" - space = self._prompt.sid_space - device = self.lm.get_input_embeddings().weight.device - return ( - torch.tensor(space.band_lo, device=device), - torch.tensor(space.band_hi, device=device), - ) - def _detokenize(self, tokens: torch.Tensor, batch_size: int) -> torch.Tensor: """Undo both shifts: token id back to a local 0-based code. @@ -201,8 +189,7 @@ def _detokenize(self, tokens: torch.Tensor, batch_size: int) -> torch.Tensor: ``(batch_size, beams, num_levels)`` local codes. """ space = self._prompt.sid_space - offsets = torch.tensor(space.level_offsets, device=tokens.device) - codes = tokens - space.base_vocab - offsets + codes = tokens - space.base_vocab - self._level_offsets return codes.view(batch_size, -1, space.num_levels) def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 237f32984..618503005 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -155,13 +155,17 @@ def _generate(self, batch: Batch) -> torch.Tensor: padded, mask, _ = _unpack( embeds, infos[PROMPT_CU_SEQLENS], - infos[PROMPT_LABELS], + None, int(infos[PROMPT_MAX_SEQLEN]), self._ignore_index, ) - lo_tok, hi_tok = self._sid_token_bands() + space = self._prompt.sid_space tokens = dynamic_beam_search( - self.lm, padded, mask, self._beam_widths, lo_tok, hi_tok + self.lm, + padded, + mask, + self._beam_widths, + list(zip(space.band_lo, space.band_hi)), ) codes = self._detokenize(tokens, padded.shape[0]) return codes[:, : self._num_return, :] @@ -170,10 +174,10 @@ def _generate(self, batch: Batch) -> torch.Tensor: def _unpack( embeds: torch.Tensor, cu_seqlens: torch.Tensor, - labels: torch.Tensor, + labels: Optional[torch.Tensor], max_seqlen: int, ignore_index: int, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: """Pad a packed varlen batch at the LM boundary. Padding lives in this one adapter. ``max_seqlen`` is the collator's, not @@ -183,7 +187,7 @@ def _unpack( Args: embeds: packed embeddings, ``(total_tokens, hidden)``. cu_seqlens: row boundaries, ``(batch_size + 1,)``. - labels: packed labels, ``(total_tokens,)``. + labels: packed labels, or None at inference where nothing is scored. max_seqlen: the collator's padded width. ignore_index: label value for padding. @@ -199,14 +203,17 @@ def _unpack( mask = columns[None, :] < lengths[:, None] padded = embeds.new_zeros((batch_size, max_seqlen, hidden)) + # mask selects row-major, which is the order embeds and labels are packed in + padded[mask] = embeds + if labels is None: + return padded, mask.long(), None + out_labels = torch.full( (batch_size, max_seqlen), ignore_index, dtype=labels.dtype, device=embeds.device, ) - # mask selects row-major, which is the order embeds and labels are packed in - padded[mask] = embeds out_labels[mask] = labels return padded, mask.long(), out_labels diff --git a/tzrec/modules/dynamic_beam.py b/tzrec/modules/dynamic_beam.py index 53af74da3..395f40955 100644 --- a/tzrec/modules/dynamic_beam.py +++ b/tzrec/modules/dynamic_beam.py @@ -15,7 +15,7 @@ can supply. """ -from typing import List, Tuple +from typing import List, Sequence, Tuple import torch from transformers import PreTrainedModel @@ -27,8 +27,7 @@ def dynamic_beam_search( prompt_embeds: torch.Tensor, attention_mask: torch.Tensor, beam_widths: List[int], - lo_tok: torch.Tensor, - hi_tok: torch.Tensor, + bands: Sequence[Tuple[int, int]], ) -> torch.Tensor: """Decode SID answers with a caller-supplied per-level beam width. @@ -39,8 +38,8 @@ def dynamic_beam_search( attention_mask: prompt mask ``(B, P)``. beam_widths: requested width per SID level; each is capped to what its band and the surviving prefixes supply. - lo_tok: inclusive lower per-level token band edge, ``(num_levels,)``. - hi_tok: inclusive upper per-level token band edge, ``(num_levels,)``. + bands: inclusive ``(lo, hi)`` token-id edge per SID level. Host ints, + because every use of them is a host-side slice bound. Returns: The SID token tail ``(B * W, num_levels)``, score-ordered best-first. @@ -48,7 +47,7 @@ def dynamic_beam_search( """ device = prompt_embeds.device batch_size = prompt_embeds.shape[0] - num_levels = lo_tok.shape[0] + num_levels = len(bands) if len(beam_widths) != num_levels: raise ValueError( f"dynamic_beam_search: beam_widths has {len(beam_widths)} entries " @@ -58,10 +57,6 @@ def dynamic_beam_search( raise ValueError( f"dynamic_beam_search: beam_widths must be >= 1, got {list(beam_widths)}." ) - # Hoist the band edges to host once to keep the level loop sync-free. - bands: List[Tuple[int, int]] = [ - (int(lo_tok[level]), int(hi_tok[level])) for level in range(num_levels) - ] capped_widths: List[int] = [] prev_width = 1 for requested, (band_lo, band_hi) in zip(beam_widths, bands): diff --git a/tzrec/modules/dynamic_beam_test.py b/tzrec/modules/dynamic_beam_test.py index 7791f924e..7726482f7 100644 --- a/tzrec/modules/dynamic_beam_test.py +++ b/tzrec/modules/dynamic_beam_test.py @@ -37,8 +37,7 @@ def _decode(lm, ids, pairs, width=8, beam_widths=None, attention_mask=None): lm.get_input_embeddings()(ids), torch.ones_like(ids) if attention_mask is None else attention_mask, beam_widths=beam_widths, - lo_tok=torch.tensor([p[0] for p in pairs]), - hi_tok=torch.tensor([p[1] for p in pairs]), + bands=pairs, ) diff --git a/tzrec/modules/prompt_projection.py b/tzrec/modules/prompt_projection.py index 760fbf14e..9ae976026 100644 --- a/tzrec/modules/prompt_projection.py +++ b/tzrec/modules/prompt_projection.py @@ -42,7 +42,6 @@ def __init__( hidden_size: int, ) -> None: super().__init__() - self._in_dim = in_dim dim = in_dim self.body: Optional[MLP] = None if config.HasField("mlp"): @@ -50,11 +49,6 @@ def __init__( dim = self.body.output_dim() self.head = nn.Linear(dim, hidden_size, bias=config.bias) - @property - def in_dim(self) -> int: - """Input width this module was sized for.""" - return self._in_dim - def forward(self, features: torch.Tensor) -> torch.Tensor: """Project a slot's group output into the LM input space. diff --git a/tzrec/prompt/assembler.py b/tzrec/prompt/assembler.py index d4917ebd1..03c8cd8c2 100644 --- a/tzrec/prompt/assembler.py +++ b/tzrec/prompt/assembler.py @@ -80,6 +80,9 @@ def __init__( self._plan = plan self._sid = sid_space self._ignore_index = ignore_index + if sid_space is not None: + self._lo = np.asarray(sid_space.level_offsets, dtype=np.int64) + self._hi = self._lo + np.asarray(sid_space.codebook, dtype=np.int64) inline = [ s for s in plan.segments + plan.response_segments @@ -105,9 +108,7 @@ def _inline_tokens(self, name: str, values: np.ndarray) -> np.ndarray: f"number of {levels}-level items." ) by_level = values.reshape(-1, levels) - lo = np.asarray(self._sid.level_offsets, dtype=np.int64) - hi = lo + np.asarray(self._sid.codebook, dtype=np.int64) - if np.any(by_level < lo) or np.any(by_level >= hi): + if np.any(by_level < self._lo) or np.any(by_level >= self._hi): raise ValueError( f"prompt slot [{name}]: SID values must already carry their " f"level offset, so level l lies in " diff --git a/tzrec/prompt/assembler_test.py b/tzrec/prompt/assembler_test.py index ca32eb1ef..6da4f623b 100644 --- a/tzrec/prompt/assembler_test.py +++ b/tzrec/prompt/assembler_test.py @@ -78,14 +78,13 @@ def _plan(segments, response=(), max_length=0) -> PromptPlan: suffix_keep=None, static_prefix_len=0, length_buckets=(), - slot_index={s.name: i for i, s in enumerate(projected)}, projected_slots=projected, ) class PromptAssemblerTest(unittest.TestCase): def test_inline_sid_gets_the_base_vocab_shift(self) -> None: - plan = _plan((Static((7, 8), None), _slot("hist", FillMode.INLINE))) + plan = _plan((Static((7, 8)), _slot("hist", FillMode.INLINE))) asm = PromptAssembler(plan, _sid_space()) # offset codes for one item: level 0 -> 1, level 1 -> 4+2, level 2 -> 8+3 out = asm.assemble({"hist": [np.array([1, 6, 11])]}) @@ -97,7 +96,7 @@ def test_inline_sid_gets_the_base_vocab_shift(self) -> None: self.assertEqual(out.hole_positions.size, 0) def test_projected_emits_sentinels_and_records_holes(self) -> None: - plan = _plan((Static((7,), None), _slot("prof", FillMode.PROJECTED, 4))) + plan = _plan((Static((7,)), _slot("prof", FillMode.PROJECTED, 4))) asm = PromptAssembler(plan, _sid_space()) out = asm.assemble({}, {"prof": np.array([2, 3])}, batch_size=2) @@ -120,8 +119,8 @@ def test_hole_positions_index_the_flat_buffer_exactly(self) -> None: def test_labels_cover_the_response_span_only(self) -> None: plan = _plan( - (Static((7, 8), None),), - response=(Static((9,), None), _slot("answer", FillMode.INLINE)), + (Static((7, 8)),), + response=(Static((9,)), _slot("answer", FillMode.INLINE)), ) asm = PromptAssembler(plan, _sid_space()) out = asm.assemble({"answer": [np.array([0, 4, 8])]}) @@ -153,9 +152,7 @@ def test_rejects_an_out_of_band_code(self) -> None: asm.assemble({"hist": [np.array([1, 6, 12])]}) def test_over_long_row_is_an_error_not_a_truncation(self) -> None: - plan = _plan( - (Static((7, 8, 9), None), _slot("hist", FillMode.INLINE)), max_length=4 - ) + plan = _plan((Static((7, 8, 9)), _slot("hist", FillMode.INLINE)), max_length=4) asm = PromptAssembler(plan, _sid_space()) with self.assertRaisesRegex(ValueError, "never truncated"): asm.assemble({"hist": [np.array([1, 6, 11])]}) diff --git a/tzrec/prompt/compile.py b/tzrec/prompt/compile.py index 6b82b68ea..cb04a1ef9 100644 --- a/tzrec/prompt/compile.py +++ b/tzrec/prompt/compile.py @@ -321,7 +321,6 @@ def compile_prompt( suffix_keep=_suffix_keep(response), static_prefix_len=_static_prefix_len(body), length_buckets=tuple(int(b) for b in cfg.length_buckets), - slot_index={s.name: i for i, s in enumerate(projected)}, projected_slots=projected, ) _validate(cfg, plan, sid_space) @@ -347,7 +346,7 @@ def _weave( for i, run in enumerate(runs): if run: ids = tuple(tok.encode(run, add_special_tokens=False).ids) - out.append(Static(token_ids=ids, owner_slot_id=None)) + out.append(Static(token_ids=ids)) if i < len(names): out.append(segs[names[i]]) return tuple(out) diff --git a/tzrec/prompt/plan.py b/tzrec/prompt/plan.py index fe40ce0c8..8ab49fad8 100644 --- a/tzrec/prompt/plan.py +++ b/tzrec/prompt/plan.py @@ -109,12 +109,9 @@ class Static: Args: token_ids: the tokenized run. - owner_slot_id: slot this run was folded into, so it vanishes when that - slot is dropped. None when the run belongs to no slot. """ token_ids: Tuple[int, ...] - owner_slot_id: Optional[int] @dataclass(frozen=True) @@ -158,7 +155,6 @@ class PromptPlan: suffix_keep: upper bound on the supervised logits window. static_prefix_len: leading positions that are request-invariant. length_buckets: sampler and graph-capture buckets. - slot_index: slot name to its index in ``projected_slots``. projected_slots: fixes the order hole positions are written in. """ @@ -170,7 +166,6 @@ class PromptPlan: suffix_keep: Optional[int] static_prefix_len: int length_buckets: Tuple[int, ...] - slot_index: Mapping[str, int] projected_slots: Tuple[SlotSeg, ...] diff --git a/tzrec/tests/prompt_integration_test.py b/tzrec/tests/prompt_integration_test.py index 93f1af10e..5dcb11d45 100644 --- a/tzrec/tests/prompt_integration_test.py +++ b/tzrec/tests/prompt_integration_test.py @@ -16,7 +16,6 @@ import torch from google.protobuf import text_format from tokenizers import Tokenizer, models, pre_tokenizers -from transformers import Qwen2Config from tzrec.datasets.utils import Batch from tzrec.features.feature import FgMode, create_features @@ -26,23 +25,15 @@ from tzrec.protos import feature_pb2 from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.prompt_pb2 import PromptConfig -from tzrec.utils.test_util import make_test_dir +from tzrec.utils.test_util import create_tiny_causal_lm, make_test_dir _CODEBOOK = [4, 4, 4] _WORDS = ["History", "Predict", ":", ".", "", "<|im_end|>"] def _tiny_backbone(path: str) -> str: - """A two-layer Qwen saved locally, so no download is needed.""" - Qwen2Config( - vocab_size=64, - hidden_size=32, - intermediate_size=64, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - max_position_embeddings=256, - ).save_pretrained(path) + """The shared tiny Qwen, saved locally so no download is needed.""" + create_tiny_causal_lm(64).save_pretrained(path) return path diff --git a/tzrec/utils/checkpoint_util.py b/tzrec/utils/checkpoint_util.py index c8a11017a..8866ec8f6 100644 --- a/tzrec/utils/checkpoint_util.py +++ b/tzrec/utils/checkpoint_util.py @@ -332,20 +332,31 @@ def best_checkpoint( return latest_checkpoint(model_dir) -def _unwrap_model(model: nn.Module) -> nn.Module: - """Walk DMP/TrainWrapper layers down to the model that owns the hooks.""" +def unwrap_to(model: nn.Module, attr: str) -> Optional[nn.Module]: + """Walk DMP/TrainWrapper layers down to the module declaring ``attr``. + + ``seen`` bounds the walk: a ``.model``/``.module`` cycle would otherwise + hang inside a checkpoint save. + + Args: + model: the outermost wrapper. + attr: the attribute that marks the module being looked for. + + Returns: + That module, or None when the chain has none. + """ inner = model seen = set() - while not hasattr(inner, "save_assets"): + while not hasattr(inner, attr): if id(inner) in seen: - return model + return None seen.add(id(inner)) if hasattr(inner, "module"): inner = inner.module elif hasattr(inner, "model"): inner = inner.model else: - return model + return None return inner @@ -429,7 +440,9 @@ def save( f"weights are saved; skipping HF assets." ) try: - _unwrap_model(model).save_assets(ckpt_dir) + inner = unwrap_to(model, "save_assets") + if inner is not None: + inner.save_assets(ckpt_dir) except Exception as e: # noqa: BLE001 logger.warning( f"save_assets failed for {ckpt_dir}: {e} -- checkpoint weights " diff --git a/tzrec/utils/hf_export_util.py b/tzrec/utils/hf_export_util.py index f46eb9f9e..597de2db4 100644 --- a/tzrec/utils/hf_export_util.py +++ b/tzrec/utils/hf_export_util.py @@ -41,28 +41,6 @@ _HF_EXPORT_META_FILENAME = "hf_export_meta.json" -def _unwrap_hf_model(wrapped_model: nn.Module) -> Optional[nn.Module]: - """Walk DMP/TrainWrapper layers down to the model exposing ``hf_backbone``. - - ``None`` when the chain has none, so callers no-op. ``seen`` bounds the - walk: every checkpoint save reaches here, and a ``.model``/``.module`` cycle - would hang inside ``save()``. - """ - m = wrapped_model - seen = set() - while not hasattr(m, "hf_backbone"): - if id(m) in seen: - return None - seen.add(id(m)) - if hasattr(m, "module"): # DMP / DDP-style wrapper - m = m.module - elif hasattr(m, "model"): # Train/Predict/Script wrapper - m = m.model - else: - return None - return m - - def write_hf_assets(wrapped_model: nn.Module, save_dir: str) -> None: """Co-locate the HF config + tokenizer (NO weights) in a checkpoint dir. @@ -71,7 +49,7 @@ def write_hf_assets(wrapped_model: nn.Module, save_dir: str) -> None: """ if int(os.environ.get("RANK", 0)) != 0: return - inner = _unwrap_hf_model(wrapped_model) + inner = checkpoint_util.unwrap_to(wrapped_model, "hf_backbone") if inner is None: return os.makedirs(save_dir, exist_ok=True) diff --git a/tzrec/utils/hf_export_util_test.py b/tzrec/utils/hf_export_util_test.py index 3830ec7cd..bd77a84f4 100644 --- a/tzrec/utils/hf_export_util_test.py +++ b/tzrec/utils/hf_export_util_test.py @@ -20,10 +20,9 @@ from safetensors.torch import load_file from torch import nn -from tzrec.utils.checkpoint_util import save_model +from tzrec.utils.checkpoint_util import save_model, unwrap_to from tzrec.utils.hf_export_util import ( _HF_EXPORT_META_FILENAME, - _unwrap_hf_model, dcp_to_hf, write_hf_assets, ) @@ -84,12 +83,12 @@ def tearDown(self) -> None: def test_unwrap_walks_dmp_and_train_wrapper(self) -> None: inner = _GenRec(_tied_lm()) - self.assertIs(_unwrap_hf_model(inner), inner) - self.assertIs(_unwrap_hf_model(_TrainWrapper(inner)), inner) - self.assertIs(_unwrap_hf_model(_DmpLike(_TrainWrapper(inner))), inner) + self.assertIs(unwrap_to_hf(inner), inner) + self.assertIs(unwrap_to_hf(_TrainWrapper(inner)), inner) + self.assertIs(unwrap_to_hf(_DmpLike(_TrainWrapper(inner))), inner) def test_unwrap_returns_none_for_non_hf_model(self) -> None: - self.assertIsNone(_unwrap_hf_model(_TrainWrapper(nn.Linear(4, 4)))) + self.assertIsNone(unwrap_to_hf(_TrainWrapper(nn.Linear(4, 4)))) def test_unwrap_terminates_on_a_wrapper_cycle(self) -> None: """A .model/.module cycle must return None, not spin. @@ -103,7 +102,7 @@ def test_unwrap_terminates_on_a_wrapper_cycle(self) -> None: object.__setattr__(a, "model", b) object.__setattr__(b, "model", a) out = [] - t = threading.Thread(target=lambda: out.append(_unwrap_hf_model(a))) + t = threading.Thread(target=lambda: out.append(unwrap_to_hf(a))) t.daemon = True t.start() t.join(timeout=5) @@ -184,5 +183,10 @@ def test_dcp_to_hf_missing_dcp_dir(self) -> None: dcp_to_hf(empty, os.path.join(self.test_dir, "hf_out_missing")) +def unwrap_to_hf(model): + """The walk write_hf_assets performs, under test.""" + return unwrap_to(model, "hf_backbone") + + if __name__ == "__main__": unittest.main() From c58c490015c91c10b6093c8586bdce0ae30a9d7d Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Thu, 6 Aug 2026 12:19:21 +0000 Subject: [PATCH 77/99] [bugfix] resolve slot width and shield the padded forward from FX Two independent defects. _slot_width read sequence_length off the feature's own config, but a member of a SequenceFeature group never sets that field -- the cap is resolved by BaseFeature.sequence_length from the group and passed in by create_features. A grouped feature therefore compiled as UNBOUNDED and the resulting error told the user to set a field they had already set. Reading the resolved property fixes both. The teacher-forced forward reached _unpack, which converts the collator's max_seqlen to a host int. TrainPipelineSparseDist symbolically traces the model whenever a sharded module exists, and int() on a Proxy raises. _fx_wrapped_loss makes that path one FX leaf, the same treatment decode already had; the embedding lookup stays outside the leaf so the pipeline can still see the sharded module and prefetch it. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_qwen.py | 25 ++++++++++++++++++++++++- tzrec/prompt/compile.py | 8 +++----- tzrec/prompt/compile_test.py | 23 +++++++++++++++++++++++ tzrec/tests/prompt_integration_test.py | 21 +++++++++++++++++++++ 4 files changed, 71 insertions(+), 6 deletions(-) diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 618503005..9fd8c7b0c 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -103,7 +103,10 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: """ if self.is_inference: return {self._generated_sids_key: _fx_wrapped_generate(self, batch)} - return self._forward_loss(self._prompt_embeds(batch), batch) + # the embedding lookup stays traceable so the train pipeline can still + # see the sharded module and prefetch it; only the padding and the LM, + # which read the collator's width as a host int, are hidden. + return _fx_wrapped_loss(self, self._prompt_embeds(batch), batch) def _forward_loss( self, embeds: torch.Tensor, batch: Batch @@ -218,6 +221,26 @@ def _unpack( return padded, mask.long(), out_labels +@torch.fx.wrap +def _fx_wrapped_loss( + model: "PromptGenerativeQwen", embeds: torch.Tensor, batch: Batch +) -> Dict[str, torch.Tensor]: + """Hide the padded forward from FX. + + ``TrainPipelineSparseDist`` symbolically traces the model whenever a + sharded module exists, and ``_unpack`` reads ``max_seqlen`` as a host int. + + Args: + model: the model whose loss to compute. + embeds: the assembled prompt embeddings, packed. + batch: the batch being scored. + + Returns: + The loss. + """ + return model._forward_loss(embeds, batch) + + @torch.fx.wrap def _fx_wrapped_generate(model: "PromptGenerativeQwen", batch: Batch) -> torch.Tensor: """Hide the decode loop from FX. diff --git a/tzrec/prompt/compile.py b/tzrec/prompt/compile.py index cb04a1ef9..599fa8bd0 100644 --- a/tzrec/prompt/compile.py +++ b/tzrec/prompt/compile.py @@ -87,11 +87,9 @@ def _slot_width( return Width(WidthKind.STATIC, answer_levels) if group_type == FeatureGroupType.DEEP: return Width(WidthKind.STATIC, 1) - caps = [ - int(f.config.sequence_length) - for f in members - if f.config.HasField("sequence_length") - ] + # BaseFeature.sequence_length, not config: a member of a SequenceFeature + # group inherits the group's cap and never sets its own field. + caps = [f.sequence_length for f in members if f.sequence_length] if not caps: return Width(WidthKind.UNBOUNDED) return Width(WidthKind.BOUNDED, max(caps)) diff --git a/tzrec/prompt/compile_test.py b/tzrec/prompt/compile_test.py index f6c65bca3..a331f27e8 100644 --- a/tzrec/prompt/compile_test.py +++ b/tzrec/prompt/compile_test.py @@ -226,6 +226,29 @@ def test_unbounded_response_is_rejected(self) -> None: with self.assertRaisesRegex(ValueError, "window cannot be bounded"): self._compile(cfg, [_feature(_HIST), answer]) + def test_a_grouped_feature_inherits_the_group_cap(self) -> None: + # a SequenceFeature member never sets its own sequence_length; the cap + # comes from the group, so reading .config here would say UNBOUNDED + fc = feature_pb2.FeatureConfig() + text_format.Merge( + """sequence_feature { + sequence_name: "clk" sequence_length: 16 sequence_delim: ";" + features { id_feature { feature_name: "h" expression: "item:h" + num_buckets: 8 embedding_dim: 4 } } + }""", + fc, + ) + grouped = create_features([fc], fg_mode=FgMode.FG_NONE) + self.assertFalse(grouped[0].config.HasField("sequence_length")) + + cfg = self._config(prompt="History : {{clk__h}}") + cfg.sid_space.codebook.extend([4]) + compiled = self._compile(cfg, grouped) + + seg = next(s for s in compiled.prompt_plan.segments if isinstance(s, SlotSeg)) + self.assertIs(seg.width.kind, WidthKind.BOUNDED) + self.assertEqual(seg.width.n, 16) + if __name__ == "__main__": unittest.main() diff --git a/tzrec/tests/prompt_integration_test.py b/tzrec/tests/prompt_integration_test.py index 5dcb11d45..aea729c0b 100644 --- a/tzrec/tests/prompt_integration_test.py +++ b/tzrec/tests/prompt_integration_test.py @@ -14,6 +14,7 @@ import numpy as np import torch +import torch.fx from google.protobuf import text_format from tokenizers import Tokenizer, models, pre_tokenizers @@ -159,6 +160,26 @@ def test_labels_supervise_only_the_answer(self) -> None: supervised = labels[labels != -100] self.assertEqual(supervised.numel(), 3) + def test_training_forward_survives_fx_tracing(self) -> None: + # TrainPipelineSparseDist symbolically traces the model whenever a + # sharded module exists, and the padded forward reads the collator's + # width as a host int + model = self._model() + + class _Wrapper(torch.nn.Module): + def __init__(self, inner): + super().__init__() + self.inner = inner + + def forward(self, batch): + return self.inner.predict(batch) + + graph = torch.fx.symbolic_trace(_Wrapper(model)) + leaves = [ + str(n.target) for n in graph.graph.nodes if "_fx_wrapped" in str(n.target) + ] + self.assertTrue(leaves, "the padded forward must be opaque to FX") + def test_raw_codes_are_rejected_before_the_model_sees_them(self) -> None: parsed = { # not offset: level 1 and 2 fall below their bands From 108054aa5d10d32e8135b20a1a9baba9e188208d Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Mon, 10 Aug 2026 09:21:35 +0000 Subject: [PATCH 78/99] [bugfix] left-pad the packed batch so short rows keep their answer _unpack right-padded: mask = columns < lengths puts real tokens at the left of each row. Both consumers index from the right. _forward_loss scores slice(-suffix_keep, None) and dynamic_beam_search prefills from last_hidden_state[:, -1, :], so for any row shorter than the collator's max_seqlen the loss window fell on padding -- whose labels are ignore_index -- and decode prefilled from a masked-out position. Rows below the batch width were therefore trained on a fraction of their SID levels, or none, and scored from a garbage hidden state. dynamic_beam_search already documented its input as left-padded, and its decode loop only works that way: it appends mask ones on the right for generated tokens and derives positions from the mask's running count. Flipping the mask to pad on the left satisfies that contract and leaves the row-major scatter, which fills in packed order either way, intact. The three existing _unpack tests asserted the right-padded layout, so they encoded the bug; they now pin the left-padded one. Added tests that every row ends on its own final token and that a short row contributes the same number of supervised positions as a long one -- all five fail against the previous mask. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_qwen.py | 9 +++- tzrec/models/prompt_generative_qwen_test.py | 35 ++++++++++++--- tzrec/tests/prompt_integration_test.py | 48 ++++++++++++++++++++- 3 files changed, 84 insertions(+), 8 deletions(-) diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 9fd8c7b0c..1bd4b4213 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -181,12 +181,17 @@ def _unpack( max_seqlen: int, ignore_index: int, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: - """Pad a packed varlen batch at the LM boundary. + """Left-pad a packed varlen batch at the LM boundary. Padding lives in this one adapter. ``max_seqlen`` is the collator's, not ``lengths.max()``: deriving it here would sync the device to the host every step, which the design forbids. + Pads go on the left so that every row ends on a real token. Both consumers + index from the right -- the loss keeps a fixed-width suffix and decode + prefills from ``[:, -1, :]`` -- so right-padding would hand a short row its + padding instead of its answer. + Args: embeds: packed embeddings, ``(total_tokens, hidden)``. cu_seqlens: row boundaries, ``(batch_size + 1,)``. @@ -203,7 +208,7 @@ def _unpack( hidden = embeds.shape[-1] columns = torch.arange(max_seqlen, device=embeds.device) - mask = columns[None, :] < lengths[:, None] + mask = columns[None, :] >= (max_seqlen - lengths)[:, None] padded = embeds.new_zeros((batch_size, max_seqlen, hidden)) # mask selects row-major, which is the order embeds and labels are packed in diff --git a/tzrec/models/prompt_generative_qwen_test.py b/tzrec/models/prompt_generative_qwen_test.py index 3112e707d..f694432eb 100644 --- a/tzrec/models/prompt_generative_qwen_test.py +++ b/tzrec/models/prompt_generative_qwen_test.py @@ -29,12 +29,13 @@ def test_packs_rows_of_different_lengths(self) -> None: padded, mask, out = _unpack(embeds, cu, labels, max_seqlen=3, ignore_index=-100) self.assertEqual(padded.shape, (2, 3, 4)) - self.assertEqual(mask.tolist(), [[1, 1, 0], [1, 1, 1]]) - torch.testing.assert_close(padded[0, :2], embeds[:2]) + # pads go on the left, so every row ends on a real token + self.assertEqual(mask.tolist(), [[0, 1, 1], [1, 1, 1]]) + torch.testing.assert_close(padded[0, 1:], embeds[:2]) torch.testing.assert_close(padded[1, :3], embeds[2:]) # the pad column is zero, and its label is ignored - torch.testing.assert_close(padded[0, 2], torch.zeros(4)) - self.assertEqual(out.tolist(), [[10, 11, -100], [20, 21, 22]]) + torch.testing.assert_close(padded[0, 0], torch.zeros(4)) + self.assertEqual(out.tolist(), [[-100, 10, 11], [20, 21, 22]]) def test_uses_the_given_width_not_the_observed_max(self) -> None: # the collator's bucket may exceed the widest row; §7.4 forbids @@ -54,9 +55,33 @@ def test_row_order_survives_the_scatter(self) -> None: labels = torch.zeros(4, dtype=torch.long) padded, _, _ = _unpack(embeds, cu, labels, max_seqlen=3, ignore_index=-100) - self.assertEqual(padded[0, 0].item(), 1.0) + self.assertEqual(padded[0, :, 0].tolist(), [0.0, 0.0, 1.0]) self.assertEqual(padded[1, :, 0].tolist(), [2.0, 3.0, 4.0]) + def test_every_row_ends_on_its_own_last_token(self) -> None: + # decode prefills from [:, -1, :]; a right-padded short row would + # hand it padding instead of the row's final token + embeds = torch.tensor([[1.0], [2.0], [3.0], [11.0], [12.0], [13.0], [14.0]]) + cu = torch.tensor([0, 3, 7]) + labels = torch.zeros(7, dtype=torch.long) + padded, _, _ = _unpack(embeds, cu, labels, max_seqlen=4, ignore_index=-100) + + self.assertEqual(padded[:, -1, 0].tolist(), [3.0, 14.0]) + + def test_a_short_row_keeps_its_answer_in_the_suffix_window(self) -> None: + # the loss scores a fixed-width suffix; every row must contribute the + # same number of supervised positions regardless of its length + embeds = torch.ones(9, 1) + cu = torch.tensor([0, 4, 9]) + ignore = -100 + # each row ends on 2 real labels, preceded by unsupervised context + labels = torch.tensor([ignore, ignore, 7, 8, ignore, ignore, ignore, 7, 8]) + _, _, out = _unpack(embeds, cu, labels, max_seqlen=5, ignore_index=ignore) + + window = out[:, -2:] + self.assertEqual((window != ignore).sum(dim=1).tolist(), [2, 2]) + self.assertEqual(window.tolist(), [[7, 8], [7, 8]]) + def test_gradient_reaches_the_packed_input(self) -> None: embeds = torch.ones(3, 2, requires_grad=True) cu = torch.tensor([0, 1, 3]) diff --git a/tzrec/tests/prompt_integration_test.py b/tzrec/tests/prompt_integration_test.py index aea729c0b..e2b266125 100644 --- a/tzrec/tests/prompt_integration_test.py +++ b/tzrec/tests/prompt_integration_test.py @@ -21,7 +21,13 @@ from tzrec.datasets.utils import Batch from tzrec.features.feature import FgMode, create_features from tzrec.main import _create_model -from tzrec.prompt.assembler import assemble_into +from tzrec.models.prompt_generative_qwen import _unpack +from tzrec.prompt.assembler import ( + PROMPT_CU_SEQLENS, + PROMPT_LABELS, + PROMPT_MAX_SEQLEN, + assemble_into, +) from tzrec.prompt.compile import compile_prompt from tzrec.protos import feature_pb2 from tzrec.protos.model_pb2 import ModelConfig @@ -113,6 +119,46 @@ def _batch(self, hist, answer): ) return batch + def _batch_rows(self, rows): + hist = [h for h, _ in rows] + answer = [a for _, a in rows] + parsed = { + "hist.values": torch.tensor(_offset([c for h in hist for c in h])), + "hist.lengths": torch.tensor([len(h) for h in hist]), + "answer.values": torch.tensor(_offset([c for a in answer for c in a])), + "answer.lengths": torch.tensor([len(a) for a in answer]), + } + streams = assemble_into(self.prompt, parsed) + batch = Batch() + batch.additional_infos.update( + {k: torch.from_numpy(np.asarray(v)) for k, v in streams.items()} + ) + return batch + + def test_every_row_is_supervised_whatever_its_length(self) -> None: + # a short row must not lose its answer to padding: the loss keeps a + # fixed-width suffix, so both rows have to contribute equally + # one history item against two, so the assembled rows differ in width + batch = self._batch_rows( + [([0, 1, 2], [1, 2, 3]), ([0, 1, 2, 3, 0, 1], [2, 3, 0])] + ) + infos = batch.additional_infos + cu = infos[PROMPT_CU_SEQLENS] + lengths = (cu[1:] - cu[:-1]).tolist() + self.assertNotEqual(lengths[0], lengths[1], "rows must differ to be a test") + + _, _, labels = _unpack( + torch.ones(int(cu[-1]), 1), + cu, + infos[PROMPT_LABELS], + int(infos[PROMPT_MAX_SEQLEN]), + -100, + ) + window = labels[:, -self.prompt.prompt_plan.suffix_keep :] + supervised = (window != -100).sum(dim=1).tolist() + self.assertEqual(supervised[0], supervised[1]) + self.assertEqual(supervised[0], self.prompt.sid_space.num_levels) + def test_compiles_a_usable_space(self) -> None: space = self.prompt.sid_space self.assertEqual(space.num_levels, 3) From 7c77611e232d2cf5ec2fdd7b171cd114c3136450 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Mon, 10 Aug 2026 09:43:26 +0000 Subject: [PATCH 79/99] [ci] cover BasePromptGenerativeModel directly The base class had no test file. It was exercised only incidentally, by the integration test constructing a Qwen model, which reaches its happy path and nothing else: both __init__ guards, the shared-projection width check, the projected scatter in _prompt_embeds, the loss and metric hooks, save_assets and init_from_pretrained had no coverage at all. DetokenizeTest was worse than absent. It lived in the Qwen test file though _detokenize belongs to the base, and it never called the method -- it recomputed tokens - base_vocab - offsets inline and asserted on its own arithmetic, so it would have passed with _detokenize deleted. It is replaced by tests that call the method on a constructed model. Every test here was checked by mutating the behavior it claims to cover and confirming it fails. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_model_test.py | 249 +++++++++++++++++++ tzrec/models/prompt_generative_qwen_test.py | 37 --- 2 files changed, 249 insertions(+), 37 deletions(-) create mode 100644 tzrec/models/prompt_generative_model_test.py diff --git a/tzrec/models/prompt_generative_model_test.py b/tzrec/models/prompt_generative_model_test.py new file mode 100644 index 000000000..f3a254907 --- /dev/null +++ b/tzrec/models/prompt_generative_model_test.py @@ -0,0 +1,249 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import unittest + +import numpy as np +import torch +from google.protobuf import text_format +from tokenizers import Tokenizer, models, pre_tokenizers +from transformers import AutoModelForCausalLM + +from torchrec import KeyedJaggedTensor + +from tzrec.datasets.utils import BASE_DATA_GROUP, Batch +from tzrec.features.feature import FgMode, create_features +from tzrec.main import _create_model +from tzrec.prompt.assembler import PROMPT_HOLE_POSITIONS, PROMPT_INPUT_IDS, assemble_into +from tzrec.prompt.compile import compile_prompt +from tzrec.protos import feature_pb2 +from tzrec.protos.model_pb2 import ModelConfig +from tzrec.protos.prompt_pb2 import PromptConfig +from tzrec.utils.state_dict_util import init_parameters +from tzrec.utils.test_util import create_tiny_causal_lm, make_test_dir + +_CODEBOOK = [4, 4, 4] +_WORDS = ["History", "Predict", ":", ".", "", "<|im_end|>"] + + +def _tokenizer(path: str) -> str: + tok = Tokenizer( + models.WordLevel(vocab={w: i for i, w in enumerate(_WORDS)}, unk_token="") + ) + tok.pre_tokenizer = pre_tokenizers.Whitespace() + tok.save(path) + return path + + +def _feature(text: str): + fc = feature_pb2.FeatureConfig() + text_format.Merge(text, fc) + return create_features([fc], fg_mode=FgMode.FG_NONE)[0] + + +def _offset(codes): + """Shift local codes into the flat space, as the SID tool's column does.""" + offsets = np.cumsum([0] + _CODEBOOK[:-1]) + return (np.asarray(codes).reshape(-1, len(_CODEBOOK)) + offsets).reshape(-1) + + +_HIST = 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }' +_ANSWER = 'sequence_raw_feature { feature_name: "answer" expression: "item:answer" }' + + +def _projected(name: str, dim: int) -> str: + return ( + f'sequence_id_feature {{ feature_name: "{name}" expression: "user:{name}" ' + f"num_buckets: 32 embedding_dim: {dim} sequence_length: 2 }}" + ) + + +class BasePromptGenerativeModelTest(unittest.TestCase): + """The backbone-agnostic half, reached through its one concrete subclass.""" + + def setUp(self) -> None: + self.test_dir = make_test_dir() + self.backbone = os.path.join(self.test_dir, "backbone") + create_tiny_causal_lm(64).save_pretrained(self.backbone) + self.tok = _tokenizer(os.path.join(self.test_dir, "tok.json")) + self.features = [_feature(_HIST), _feature(_ANSWER)] + self.prompt = self._compile(self.features) + + def _compile(self, features, template="History : {{hist}} . Predict :", **kwargs): + cfg = PromptConfig(tokenizer=self.tok, prompt=template, **kwargs) + cfg.sid_space.codebook.extend(_CODEBOOK) + return compile_prompt(cfg, features, model_dir=self.test_dir) + + def _model(self, features=None, prompt=-1): + model_config = ModelConfig() + qwen = model_config.prompt_generative_qwen + qwen.hf_model_id = self.backbone + qwen.common.beam_widths.extend([2, 2, 2]) + qwen.common.num_return_sequences = 2 + return _create_model( + model_config, + self.features if features is None else features, + ["answer"], + prompt=self.prompt if prompt == -1 else prompt, + ) + + def _batch(self, parsed, prompt=None, sparse=None): + streams = assemble_into(prompt or self.prompt, parsed) + batch = Batch(sparse_features={BASE_DATA_GROUP: sparse} if sparse else {}) + batch.additional_infos.update( + {k: torch.from_numpy(np.asarray(v)) for k, v in streams.items()} + ) + return batch + + def test_detokenize_undoes_both_shifts(self) -> None: + model = self._model() + space = self.prompt.sid_space + # one row, one beam: level 0 code 1, level 1 code 2, level 2 code 3 + tokens = torch.tensor( + [ + [ + space.base_vocab + space.level_offsets[0] + 1, + space.base_vocab + space.level_offsets[1] + 2, + space.base_vocab + space.level_offsets[2] + 3, + ] + ] + ) + codes = model._detokenize(tokens, batch_size=1) + + self.assertEqual(codes.shape, (1, 1, space.num_levels)) + self.assertEqual(codes[0, 0].tolist(), [1, 2, 3]) + + def test_detokenize_maps_band_edges_to_the_last_code(self) -> None: + model = self._model() + space = self.prompt.sid_space + codes = model._detokenize(torch.tensor([list(space.band_hi)]), batch_size=1) + + self.assertEqual(codes[0, 0].tolist(), [c - 1 for c in _CODEBOOK]) + + def test_detokenize_groups_beams_under_their_row(self) -> None: + model = self._model() + space = self.prompt.sid_space + base = torch.tensor(space.level_offsets) + space.base_vocab + # four beam rows over a batch of two + codes = model._detokenize(base.repeat(4, 1), batch_size=2) + + self.assertEqual(codes.shape, (2, 2, space.num_levels)) + + def test_rejects_a_model_built_without_a_prompt(self) -> None: + with self.assertRaisesRegex(ValueError, "needs a compiled prompt"): + self._model(prompt=None) + + def test_rejects_a_prompt_that_declares_no_sid_space(self) -> None: + cfg = PromptConfig(tokenizer=self.tok, prompt="History : {{hist}} .") + prompt = compile_prompt(cfg, self.features, model_dir=self.test_dir) + self.assertIsNone(prompt.sid_space) + + with self.assertRaisesRegex(ValueError, "declares no sid_space"): + self._model(prompt=prompt) + + def test_shared_projection_name_requires_matching_widths(self) -> None: + features = [ + _feature(_HIST), + _feature(_ANSWER), + _feature(_projected("pa", 8)), + _feature(_projected("pb", 16)), + ] + cfg = PromptConfig( + tokenizer=self.tok, + prompt="History : {{hist}} . {{pa}} {{pb}} Predict :", + response="{{answer}}", + ) + cfg.sid_space.codebook.extend(_CODEBOOK) + for name in ("pa", "pb"): + slot = cfg.slots.add(name=name, projection_name="shared") + slot.feature_names.append(name) + prompt = compile_prompt(cfg, features, model_dir=self.test_dir) + + with self.assertRaisesRegex(ValueError, "cannot share a module"): + self._model(features=features, prompt=prompt) + + def test_projected_slot_overwrites_only_its_sentinel_positions(self) -> None: + features = [_feature(_HIST), _feature(_ANSWER), _feature(_projected("prof", 8))] + prompt = self._compile( + features, + template="History : {{hist}} . Predict {{prof}} :", + response="{{answer}}", + ) + model = self._model(features=features, prompt=prompt) + # the embedding table is built on meta until something materializes it + init_parameters(model, device=torch.device("cpu")) + batch = self._batch( + { + "hist.values": torch.tensor(_offset([0, 1, 2])).reshape(-1, 1), + "hist.lengths": torch.tensor([3]), + "answer.values": torch.tensor(_offset([1, 2, 3])), + "answer.lengths": torch.tensor([3]), + "prof.values": torch.tensor([5, 9]), + "prof.lengths": torch.tensor([2]), + }, + prompt=prompt, + sparse=KeyedJaggedTensor.from_lengths_sync( + keys=["prof"], + values=torch.tensor([5, 9]), + lengths=torch.tensor([2]), + ), + ) + + embeds = model._prompt_embeds(batch) + raw = model.lm.get_input_embeddings()(batch.additional_infos[PROMPT_INPUT_IDS]) + holes = batch.additional_infos[PROMPT_HOLE_POSITIONS] + self.assertGreater(holes.numel(), 0) + + changed = ~torch.isclose(embeds, raw).all(dim=-1) + self.assertEqual(sorted(changed.nonzero().flatten().tolist()), holes.tolist()) + + def test_loss_surfaces_the_ce_computed_in_predict(self) -> None: + model = self._model() + value = torch.tensor(1.25) + + self.assertEqual(model.loss({"loss": value}, Batch()), {"ce_loss": value}) + + def test_metric_averages_the_loss_across_batches(self) -> None: + model = self._model() + model.init_metric() + for value in (1.0, 3.0): + model.update_metric({"loss": torch.tensor(value)}, Batch()) + + self.assertAlmostEqual( + model._metric_modules["ce_loss"].compute().item(), 2.0, places=5 + ) + + def test_save_assets_co_locates_the_prompt_contract(self) -> None: + model = self._model() + target = os.path.join(self.test_dir, "ckpt") + os.makedirs(target, exist_ok=True) + model.save_assets(target) + + self.assertTrue(os.path.isdir(os.path.join(target, "prompt"))) + self.assertTrue(os.listdir(os.path.join(target, "prompt"))) + + def test_init_from_pretrained_replaces_the_empty_weights(self) -> None: + model = self._model() + base_vocab = self.prompt.sid_space.base_vocab + before = model.lm.get_input_embeddings().weight[:base_vocab].clone() + model.init_from_pretrained() + after = model.lm.get_input_embeddings().weight[:base_vocab] + + # the checkpoint rows land verbatim; only the appended SID rows are new + reference = AutoModelForCausalLM.from_pretrained(self.backbone) + expected = reference.get_input_embeddings().weight[:base_vocab] + self.assertFalse(torch.allclose(before, expected)) + torch.testing.assert_close(after, expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/models/prompt_generative_qwen_test.py b/tzrec/models/prompt_generative_qwen_test.py index f694432eb..d130e464c 100644 --- a/tzrec/models/prompt_generative_qwen_test.py +++ b/tzrec/models/prompt_generative_qwen_test.py @@ -14,7 +14,6 @@ import torch from tzrec.models.prompt_generative_qwen import _unpack -from tzrec.prompt.plan import SidSpace class UnpackTest(unittest.TestCase): @@ -93,41 +92,5 @@ def test_gradient_reaches_the_packed_input(self) -> None: torch.testing.assert_close(embeds.grad, torch.ones(3, 2)) -class DetokenizeTest(unittest.TestCase): - """Both shifts must come back off, in the right order.""" - - def _space(self) -> SidSpace: - return SidSpace( - codebook=(4, 4, 4), - num_levels=3, - base_vocab=1000, - level_offsets=(0, 4, 8), - band_lo=(1000, 1004, 1008), - band_hi=(1003, 1007, 1011), - target_vocab=1152, - sentinel_token_id=None, - eos_token_id=2, - pad_token_id=3, - ) - - def test_token_ids_become_local_codes(self) -> None: - space = self._space() - # one beam row: level 0 code 1, level 1 code 2, level 2 code 3 - tokens = torch.tensor([[1000 + 1, 1000 + 4 + 2, 1000 + 8 + 3]]) - offsets = torch.tensor(space.level_offsets) - codes = (tokens - space.base_vocab - offsets).view(1, -1, 3) - - self.assertEqual(codes[0, 0].tolist(), [1, 2, 3]) - # every code lands back inside its own codebook - self.assertTrue(bool(((codes >= 0) & (codes < 4)).all())) - - def test_a_band_edge_maps_to_the_last_code(self) -> None: - space = self._space() - tokens = torch.tensor([list(space.band_hi)]) - offsets = torch.tensor(space.level_offsets) - codes = tokens - space.base_vocab - offsets - self.assertEqual(codes[0].tolist(), [3, 3, 3]) - - if __name__ == "__main__": unittest.main() From cc75e0824f01e5a6ab27fd67d0bb4728ba132029 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Mon, 10 Aug 2026 09:51:34 +0000 Subject: [PATCH 80/99] [ci] sort imports in the base prompt model test Ruff's isort rule orders torchrec before transformers and wraps the assembler import past 88 columns; the committed file had neither, so Code Style CI failed on a hook that fixes files in place. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_model_test.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tzrec/models/prompt_generative_model_test.py b/tzrec/models/prompt_generative_model_test.py index f3a254907..4bbace0d9 100644 --- a/tzrec/models/prompt_generative_model_test.py +++ b/tzrec/models/prompt_generative_model_test.py @@ -16,14 +16,17 @@ import torch from google.protobuf import text_format from tokenizers import Tokenizer, models, pre_tokenizers -from transformers import AutoModelForCausalLM - from torchrec import KeyedJaggedTensor +from transformers import AutoModelForCausalLM from tzrec.datasets.utils import BASE_DATA_GROUP, Batch from tzrec.features.feature import FgMode, create_features from tzrec.main import _create_model -from tzrec.prompt.assembler import PROMPT_HOLE_POSITIONS, PROMPT_INPUT_IDS, assemble_into +from tzrec.prompt.assembler import ( + PROMPT_HOLE_POSITIONS, + PROMPT_INPUT_IDS, + assemble_into, +) from tzrec.prompt.compile import compile_prompt from tzrec.protos import feature_pb2 from tzrec.protos.model_pb2 import ModelConfig From 9bf8b580e996d412082d610dd40b454e489d5c52 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Tue, 11 Aug 2026 06:00:13 +0000 Subject: [PATCH 81/99] [perf] build the prompt assembler once per dataset assemble_into constructed a PromptAssembler on every batch, though the plan it walks and the band edges it validates against are fixed for the run. The per-batch reshaping moves into PromptAssembler.assemble_batch, so BaseDataset can hold one instance built in __init__; assemble_into stays as a one-shot wrapper for callers that assemble a single batch. A training run now constructs two assemblers, one per dataloader, instead of one per step. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/datasets/dataset.py | 13 +++++-- tzrec/prompt/assembler.py | 78 +++++++++++++++++++++++---------------- 2 files changed, 55 insertions(+), 36 deletions(-) diff --git a/tzrec/datasets/dataset.py b/tzrec/datasets/dataset.py index 014e2de41..f85bba1d3 100644 --- a/tzrec/datasets/dataset.py +++ b/tzrec/datasets/dataset.py @@ -41,7 +41,7 @@ remove_nullable, ) from tzrec.features.feature import BaseFeature -from tzrec.prompt.assembler import assemble_into +from tzrec.prompt.assembler import PromptAssembler from tzrec.prompt.plan import CompiledPrompt from tzrec.protos import data_pb2 from tzrec.utils import config_util @@ -113,7 +113,12 @@ def __init__( prompt: Optional[CompiledPrompt] = None, ) -> None: super(BaseDataset, self).__init__() - self._prompt = prompt + # built once per worker: the plan it walks is fixed for the run + self._assembler = ( + PromptAssembler(prompt.prompt_plan, prompt.sid_space) + if prompt is not None + else None + ) self._data_config = data_config self._features = features self._input_path = input_path @@ -390,11 +395,11 @@ def _build_batch(self, input_data: Dict[str, pa.Array]) -> Batch: else: batch = self._data_parser.to_batch(output_data) - if self._prompt is not None: + if self._assembler is not None: batch.additional_infos.update( { k: torch.from_numpy(np.asarray(v)) - for k, v in assemble_into(self._prompt, output_data).items() + for k, v in self._assembler.assemble_batch(output_data).items() } ) diff --git a/tzrec/prompt/assembler.py b/tzrec/prompt/assembler.py index 03c8cd8c2..2896b7155 100644 --- a/tzrec/prompt/assembler.py +++ b/tzrec/prompt/assembler.py @@ -201,13 +201,55 @@ def assemble( labels=np.asarray(labels, dtype=np.int64), ) + def assemble_batch(self, parsed: Dict[str, "np.ndarray"]) -> Dict[str, np.ndarray]: + """Reshape one parsed batch, assemble it, and key it for the batch. + + Args: + parsed: ``{feature}.values`` / ``{feature}.lengths`` as the data + parser emits them. + + Returns: + The five streams, keyed as ``additional_infos`` expects them. + """ + values: Dict[str, List[np.ndarray]] = {} + counts: Dict[str, np.ndarray] = {} + batch_size = 0 + for seg in self._plan.segments + self._plan.response_segments: + if not isinstance(seg, SlotSeg): + continue + source = seg.sources[0] + lengths = np.asarray(parsed[f"{source}.lengths"]) + batch_size = max(batch_size, int(lengths.size)) + if seg.fill is FillMode.INLINE: + # a dense sequence feature emits (total, value_dim); the stream + # is one code per position, so value_dim is always 1 here + flat = np.asarray(parsed[f"{source}.values"]).reshape(-1) + bounds = np.concatenate(([0], np.cumsum(lengths))) + values[seg.name] = [ + flat[bounds[i] : bounds[i + 1]] for i in range(lengths.size) + ] + else: + counts[seg.name] = lengths + + out = self.assemble(values, counts, batch_size=batch_size) + return { + PROMPT_INPUT_IDS: out.input_ids, + PROMPT_CU_SEQLENS: out.cu_seqlens, + PROMPT_HOLE_POSITIONS: out.hole_positions, + PROMPT_LABELS: out.labels, + PROMPT_MAX_SEQLEN: np.asarray(out.max_seqlen, dtype=np.int64), + } + def assemble_into( prompt: CompiledPrompt, parsed: Dict[str, "np.ndarray"], ignore_index: int = -100, ) -> Dict[str, np.ndarray]: - """Run the assembler over one parsed batch and key it for the batch. + """Assemble one batch with a throwaway assembler. + + A caller that assembles every batch should hold a ``PromptAssembler`` and + call ``assemble_batch`` instead; this builds one per call. Args: prompt: the compiled prompt. @@ -218,34 +260,6 @@ def assemble_into( Returns: The five streams, keyed as ``additional_infos`` expects them. """ - plan = prompt.prompt_plan - values: Dict[str, List[np.ndarray]] = {} - counts: Dict[str, np.ndarray] = {} - batch_size = 0 - for seg in plan.segments + plan.response_segments: - if not isinstance(seg, SlotSeg): - continue - source = seg.sources[0] - lengths = np.asarray(parsed[f"{source}.lengths"]) - batch_size = max(batch_size, int(lengths.size)) - if seg.fill is FillMode.INLINE: - # a dense sequence feature emits (total, value_dim); the stream is - # one code per position, so value_dim is always 1 here - flat = np.asarray(parsed[f"{source}.values"]).reshape(-1) - bounds = np.concatenate(([0], np.cumsum(lengths))) - values[seg.name] = [ - flat[bounds[i] : bounds[i + 1]] for i in range(lengths.size) - ] - else: - counts[seg.name] = lengths - - out = PromptAssembler(plan, prompt.sid_space, ignore_index).assemble( - values, counts, batch_size=batch_size - ) - return { - PROMPT_INPUT_IDS: out.input_ids, - PROMPT_CU_SEQLENS: out.cu_seqlens, - PROMPT_HOLE_POSITIONS: out.hole_positions, - PROMPT_LABELS: out.labels, - PROMPT_MAX_SEQLEN: np.asarray(out.max_seqlen, dtype=np.int64), - } + return PromptAssembler( + prompt.prompt_plan, prompt.sid_space, ignore_index + ).assemble_batch(parsed) From 9906dae2a123cb08626fa1c8da8ca9acc04c1174 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 12 Aug 2026 09:53:58 +0000 Subject: [PATCH 82/99] [refactor] state the prompt export guard as one condition The TorchScript refusal was a compound condition that had to name the HF format only to exclude it. Handling the HF branch first, which already returns, leaves everything below it non-HF, so the refusal collapses to a single HasField check. Folding both under HasField("prompt_config") instead would have skipped the HF branch for a config that declares no prompt_config, silently exporting TorchScript where the user asked for HF. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/main.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index f1e4fdc74..ecb737a2e 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -1128,16 +1128,6 @@ def export( checkpoint_path, _ = ckpt_manager.latest_checkpoint() # HF export converts the checkpoint dir directly -- no model build, no DCP restore. - if pipeline_config.HasField("prompt_config") and ( - pipeline_config.export_config.export_format != export_pb2.ExportFormat.HF - ): - raise ValueError( - "a prompt-native model exports to a HuggingFace directory, not " - "TorchScript: its input is an assembled token stream the dataloader " - "builds, which an export-time dummy batch cannot supply. Set " - "export_config.export_format to HF." - ) - if pipeline_config.export_config.export_format == export_pb2.ExportFormat.HF: if config_util.use_dense_ema( pipeline_config.export_config, pipeline_config.train_config @@ -1165,6 +1155,14 @@ def export( copy_prompt_assets(checkpoint_path, export_dir) return + if pipeline_config.HasField("prompt_config"): + raise ValueError( + "a prompt-native model exports to a HuggingFace directory, not " + "TorchScript: its input is an assembled token stream the dataloader " + "builds, which an export-time dummy batch cannot supply. Set " + "export_config.export_format to HF." + ) + data_config = pipeline_config.data_config # Build feature From 501e0df4f6d45f5423b56bacfb6872c0e3886b12 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 12 Aug 2026 09:54:14 +0000 Subject: [PATCH 83/99] [refactor] rename prompt-stack names that collide or under-describe An independent naming audit of the new stack surfaced names that mislead a first-time reader. The renames, no behaviour change: _lo/_hi in the assembler read as sid_space.band_lo/band_hi but are a different coordinate system and a different interval convention -- flat index rather than token id, half-open rather than inclusive -- so they are now _flat_lo/_flat_hi. module_plan is what TorchRec calls a per-module sharding plan, and the prompt one was bound inside a model about to be wrapped in DMP; it is now projection_plan. The SidSpace dataclass shadowed the proto message of the same name, which had already forced _atom_tokens to annotate its proto argument as Any; the dataclass is now ResolvedSidSpace and that annotation states the real type. Also: widths -> in_dims (it held embedding dims while Width means a position count), counts -> slot_lengths, _detokenize -> _tokens_to_local_codes, SlotSeg.sources -> feature_names, suffix_keep -> logits_suffix_len, Width.n -> num_positions, _sid -> _sid_space, _num_return -> _num_return_sequences, in_band -> band_idx, cu -> cu_seqlens. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_model.py | 42 ++++++----- tzrec/models/prompt_generative_model_test.py | 14 ++-- tzrec/models/prompt_generative_qwen.py | 12 +-- tzrec/modules/dynamic_beam.py | 4 +- tzrec/prompt/assembler.py | 77 +++++++++++--------- tzrec/prompt/assembler_test.py | 14 ++-- tzrec/prompt/compile.py | 54 ++++++++------ tzrec/prompt/compile_test.py | 8 +- tzrec/prompt/persist.py | 2 +- tzrec/prompt/plan.py | 28 +++---- tzrec/tests/prompt_integration_test.py | 4 +- 11 files changed, 141 insertions(+), 118 deletions(-) diff --git a/tzrec/models/prompt_generative_model.py b/tzrec/models/prompt_generative_model.py index 6a6f5c0c0..b8220a3cb 100644 --- a/tzrec/models/prompt_generative_model.py +++ b/tzrec/models/prompt_generative_model.py @@ -91,7 +91,7 @@ def __init__( prompt.sid_space.target_vocab, mean_resizing=True ) self.embedding_group = EmbeddingGroup( - self._features, list(self._prompt.module_plan.feature_groups) + self._features, list(self._prompt.projection_plan.feature_groups) ) self._build_projections() # decode subtracts these every step; a buffer follows the module's device @@ -116,35 +116,35 @@ def _build_backbone(self, hf_model_id: str, param_dtype: int) -> nn.Module: return model.to(_PARAM_DTYPE[param_dtype]) def _build_projections(self) -> None: - """One module per resolved id, aligned with ``plan.projected_slots``. + """One module per resolved id, aligned with ``prompt_plan.projected_slots``. Slots sharing a ``projection_name`` share a module by reference, so they must agree on ``group_total_dim``. """ - plan = self._prompt.prompt_plan - modules = self._prompt.module_plan - hidden = int(self.lm.config.hidden_size) + prompt_plan = self._prompt.prompt_plan + projection_plan = self._prompt.projection_plan + hidden_size = int(self.lm.config.hidden_size) built: Dict[str, PromptProjection] = {} - widths: Dict[str, int] = {} + in_dims: Dict[str, int] = {} aligned: List[PromptProjection] = [] - for seg in plan.projected_slots: - module_id = modules.slot_to_module[seg.slot_id] + for seg in prompt_plan.projected_slots: + module_id = projection_plan.slot_to_module[seg.slot_id] in_dim = self.embedding_group.group_total_dim(seg.name + seg.output_key) if module_id not in built: built[module_id] = PromptProjection( - modules.projections[module_id], in_dim, hidden + projection_plan.projections[module_id], in_dim, hidden_size ) - widths[module_id] = in_dim - elif widths[module_id] != in_dim: + in_dims[module_id] = in_dim + elif in_dims[module_id] != in_dim: raise ValueError( f"prompt slots sharing projection_name [{module_id}] have " - f"different group widths ({widths[module_id]} vs " + f"different group dims ({in_dims[module_id]} vs " f"{in_dim}); they cannot share a module." ) aligned.append(built[module_id]) self.projections = nn.ModuleDict(built) - # zipped with plan.projected_slots; shared modules appear by reference + # zipped with prompt_plan.projected_slots; shared modules appear by reference self._slot_projections = aligned def hf_backbone(self) -> nn.Module: @@ -158,27 +158,29 @@ def _prompt_embeds(self, batch: Batch) -> torch.Tensor: batch: carries the packed prompt in ``additional_infos``. Returns: - ``(total_tokens, hidden)``. + ``(total_tokens, hidden_size)``. """ ids = batch.additional_infos[PROMPT_INPUT_IDS] embeds = self.lm.get_input_embeddings()(ids) - plan = self._prompt.prompt_plan - if not plan.projected_slots: + prompt_plan = self._prompt.prompt_plan + if not prompt_plan.projected_slots: return embeds grouped = self.embedding_group(batch) - hidden = embeds.shape[-1] + hidden_size = embeds.shape[-1] parts = [ - proj(grouped[seg.name + seg.output_key]).reshape(-1, hidden) - for seg, proj in zip(plan.projected_slots, self._slot_projections) + proj(grouped[seg.name + seg.output_key]).reshape(-1, hidden_size) + for seg, proj in zip(prompt_plan.projected_slots, self._slot_projections) ] # out of place: embeds carries grad from the embedding lookup return embeds.index_copy( 0, batch.additional_infos[PROMPT_HOLE_POSITIONS], torch.cat(parts) ) - def _detokenize(self, tokens: torch.Tensor, batch_size: int) -> torch.Tensor: + def _tokens_to_local_codes( + self, tokens: torch.Tensor, batch_size: int + ) -> torch.Tensor: """Undo both shifts: token id back to a local 0-based code. Args: diff --git a/tzrec/models/prompt_generative_model_test.py b/tzrec/models/prompt_generative_model_test.py index 4bbace0d9..38d5fcdf5 100644 --- a/tzrec/models/prompt_generative_model_test.py +++ b/tzrec/models/prompt_generative_model_test.py @@ -107,7 +107,7 @@ def _batch(self, parsed, prompt=None, sparse=None): ) return batch - def test_detokenize_undoes_both_shifts(self) -> None: + def test_tokens_to_local_codes_undoes_both_shifts(self) -> None: model = self._model() space = self.prompt.sid_space # one row, one beam: level 0 code 1, level 1 code 2, level 2 code 3 @@ -120,24 +120,26 @@ def test_detokenize_undoes_both_shifts(self) -> None: ] ] ) - codes = model._detokenize(tokens, batch_size=1) + codes = model._tokens_to_local_codes(tokens, batch_size=1) self.assertEqual(codes.shape, (1, 1, space.num_levels)) self.assertEqual(codes[0, 0].tolist(), [1, 2, 3]) - def test_detokenize_maps_band_edges_to_the_last_code(self) -> None: + def test_tokens_to_local_codes_maps_band_edges_to_the_last_code(self) -> None: model = self._model() space = self.prompt.sid_space - codes = model._detokenize(torch.tensor([list(space.band_hi)]), batch_size=1) + codes = model._tokens_to_local_codes( + torch.tensor([list(space.band_hi)]), batch_size=1 + ) self.assertEqual(codes[0, 0].tolist(), [c - 1 for c in _CODEBOOK]) - def test_detokenize_groups_beams_under_their_row(self) -> None: + def test_tokens_to_local_codes_groups_beams_under_their_row(self) -> None: model = self._model() space = self.prompt.sid_space base = torch.tensor(space.level_offsets) + space.base_vocab # four beam rows over a batch of two - codes = model._detokenize(base.repeat(4, 1), batch_size=2) + codes = model._tokens_to_local_codes(base.repeat(4, 1), batch_size=2) self.assertEqual(codes.shape, (2, 2, space.num_levels)) diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 1bd4b4213..2bbe9e5c4 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -72,7 +72,7 @@ def _read_beam_config(self, common: PromptModelConfig) -> None: common: the shared model config. """ space = self._prompt.sid_space - self._num_return = int(common.num_return_sequences) + self._num_return_sequences = int(common.num_return_sequences) self._beam_widths: List[int] = list(common.beam_widths) if not self._beam_widths: raise ValueError( @@ -85,10 +85,10 @@ def _read_beam_config(self, common: PromptModelConfig) -> None: f"{len(self._beam_widths)} entries but the codebook has " f"{space.num_levels} levels; give one width per level." ) - if self._num_return > self._beam_widths[-1]: + if self._num_return_sequences > self._beam_widths[-1]: raise ValueError( f"{type(self).__name__}: num_return_sequences " - f"({self._num_return}) must not exceed the final beam width " + f"({self._num_return_sequences}) must not exceed the final beam width " f"({self._beam_widths[-1]})." ) @@ -133,7 +133,7 @@ def _forward_loss( ) outputs = self.lm.model(inputs_embeds=padded, attention_mask=mask) - suffix = self._prompt.prompt_plan.suffix_keep + suffix = self._prompt.prompt_plan.logits_suffix_len window = slice(-suffix, None) if suffix else slice(None) logits = self.lm.lm_head(outputs.last_hidden_state[:, window, :]) loss = self.lm.loss_function( @@ -170,8 +170,8 @@ def _generate(self, batch: Batch) -> torch.Tensor: self._beam_widths, list(zip(space.band_lo, space.band_hi)), ) - codes = self._detokenize(tokens, padded.shape[0]) - return codes[:, : self._num_return, :] + codes = self._tokens_to_local_codes(tokens, padded.shape[0]) + return codes[:, : self._num_return_sequences, :] def _unpack( diff --git a/tzrec/modules/dynamic_beam.py b/tzrec/modules/dynamic_beam.py index 395f40955..784608b87 100644 --- a/tzrec/modules/dynamic_beam.py +++ b/tzrec/modules/dynamic_beam.py @@ -82,8 +82,8 @@ def _band_logp(logits: torch.Tensor, level: int) -> torch.Tensor: ) cache = outputs.past_key_values scores = _band_logp(model.lm_head(outputs.last_hidden_state[:, -1, :]), 0) - beam_scores, in_band = scores.topk(capped_widths[0], dim=-1) - seq = (in_band + bands[0][0]).reshape(-1, 1) + beam_scores, band_idx = scores.topk(capped_widths[0], dim=-1) + seq = (band_idx + bands[0][0]).reshape(-1, 1) beam_scores = beam_scores.reshape(-1) row_starts = torch.arange(batch_size, device=device) cache.reorder_cache(row_starts.repeat_interleave(capped_widths[0])) diff --git a/tzrec/prompt/assembler.py b/tzrec/prompt/assembler.py index 2896b7155..dbaa90664 100644 --- a/tzrec/prompt/assembler.py +++ b/tzrec/prompt/assembler.py @@ -25,7 +25,7 @@ CompiledPrompt, FillMode, PromptPlan, - SidSpace, + ResolvedSidSpace, SlotSeg, Static, ) @@ -66,26 +66,28 @@ class PromptAssembler: """Walks a ``PromptPlan`` to build token streams. Args: - plan: the compiled walk order. + prompt_plan: the compiled walk order. sid_space: resolved SID token space; required when a slot renders SIDs. ignore_index: label value outside the supervised span. """ def __init__( self, - plan: PromptPlan, - sid_space: Optional[SidSpace] = None, + prompt_plan: PromptPlan, + sid_space: Optional[ResolvedSidSpace] = None, ignore_index: int = -100, ) -> None: - self._plan = plan - self._sid = sid_space + self._prompt_plan = prompt_plan + self._sid_space = sid_space self._ignore_index = ignore_index if sid_space is not None: - self._lo = np.asarray(sid_space.level_offsets, dtype=np.int64) - self._hi = self._lo + np.asarray(sid_space.codebook, dtype=np.int64) + self._flat_lo = np.asarray(sid_space.level_offsets, dtype=np.int64) + self._flat_hi = self._flat_lo + np.asarray( + sid_space.codebook, dtype=np.int64 + ) inline = [ s - for s in plan.segments + plan.response_segments + for s in prompt_plan.segments + prompt_plan.response_segments if isinstance(s, SlotSeg) and s.fill is FillMode.INLINE ] if inline and sid_space is None: @@ -100,29 +102,29 @@ def _inline_tokens(self, name: str, values: np.ndarray) -> np.ndarray: The data carries ``level_offsets[l] + code``; the LM vocabulary needs one further uniform shift by ``base_vocab``. """ - assert self._sid is not None - levels = self._sid.num_levels + assert self._sid_space is not None + levels = self._sid_space.num_levels if values.size % levels: raise ValueError( f"prompt slot [{name}]: {values.size} values is not a whole " f"number of {levels}-level items." ) by_level = values.reshape(-1, levels) - if np.any(by_level < self._lo) or np.any(by_level >= self._hi): + if np.any(by_level < self._flat_lo) or np.any(by_level >= self._flat_hi): raise ValueError( f"prompt slot [{name}]: SID values must already carry their " f"level offset, so level l lies in " f"[level_offsets[l], level_offsets[l] + codebook[l]). Read the " f"offset_codebook column, not codebook or origin_codebook." ) - return values.astype(np.int64, copy=False) + self._sid.base_vocab + return values.astype(np.int64, copy=False) + self._sid_space.base_vocab def _emit_row( self, segments: Sequence[object], row: int, values: Dict[str, List[np.ndarray]], - counts: Dict[str, np.ndarray], + slot_lengths: Dict[str, np.ndarray], out: List[int], holes: List[int], base: int, @@ -136,28 +138,28 @@ def _emit_row( if seg.fill is FillMode.INLINE: out.extend(self._inline_tokens(seg.name, values[seg.name][row])) else: - assert self._sid is not None - width = int(counts[seg.name][row]) + assert self._sid_space is not None + width = int(slot_lengths[seg.name][row]) holes.extend(range(base + len(out), base + len(out) + width)) - out.extend([self._sid.sentinel_token_id] * width) + out.extend([self._sid_space.sentinel_token_id] * width) def assemble( self, values: Dict[str, List[np.ndarray]], - counts: Optional[Dict[str, np.ndarray]] = None, + slot_lengths: Optional[Dict[str, np.ndarray]] = None, batch_size: Optional[int] = None, ) -> AssembledPrompt: """Assemble one batch. Args: values: INLINE slot name to its per-row value arrays. - counts: PROJECTED slot name to its per-row position count. + slot_lengths: PROJECTED slot name to its per-row position count. batch_size: row count; inferred from ``values`` when omitted. Returns: The packed streams. """ - counts = counts or {} + slot_lengths = slot_lengths or {} if batch_size is None: if not values: raise ValueError("batch_size is required when no INLINE slot exists.") @@ -166,37 +168,46 @@ def assemble( ids: List[int] = [] labels: List[int] = [] holes: List[int] = [] - cu = [0] + cu_seqlens = [0] for row in range(batch_size): row_ids: List[int] = [] self._emit_row( - self._plan.segments, row, values, counts, row_ids, holes, len(ids) + self._prompt_plan.segments, + row, + values, + slot_lengths, + row_ids, + holes, + len(ids), ) prompt_len = len(row_ids) self._emit_row( - self._plan.response_segments, + self._prompt_plan.response_segments, row, values, - counts, + slot_lengths, row_ids, holes, len(ids), ) # supervision covers the response span only; the prompt is context. row_labels = [self._ignore_index] * prompt_len + row_ids[prompt_len:] - if self._plan.max_length and len(row_ids) > self._plan.max_length: + if ( + self._prompt_plan.max_length + and len(row_ids) > self._prompt_plan.max_length + ): raise ValueError( f"assembled row {row} is {len(row_ids)} tokens, over " - f"max_length {self._plan.max_length}. Rows are never " + f"max_length {self._prompt_plan.max_length}. Rows are never " f"truncated: cap the source features instead." ) ids.extend(row_ids) labels.extend(row_labels) - cu.append(len(ids)) + cu_seqlens.append(len(ids)) return AssembledPrompt( input_ids=np.asarray(ids, dtype=np.int64), - cu_seqlens=np.asarray(cu, dtype=np.int64), + cu_seqlens=np.asarray(cu_seqlens, dtype=np.int64), hole_positions=np.asarray(holes, dtype=np.int64), labels=np.asarray(labels, dtype=np.int64), ) @@ -212,12 +223,12 @@ def assemble_batch(self, parsed: Dict[str, "np.ndarray"]) -> Dict[str, np.ndarra The five streams, keyed as ``additional_infos`` expects them. """ values: Dict[str, List[np.ndarray]] = {} - counts: Dict[str, np.ndarray] = {} + slot_lengths: Dict[str, np.ndarray] = {} batch_size = 0 - for seg in self._plan.segments + self._plan.response_segments: + for seg in self._prompt_plan.segments + self._prompt_plan.response_segments: if not isinstance(seg, SlotSeg): continue - source = seg.sources[0] + source = seg.feature_names[0] lengths = np.asarray(parsed[f"{source}.lengths"]) batch_size = max(batch_size, int(lengths.size)) if seg.fill is FillMode.INLINE: @@ -229,9 +240,9 @@ def assemble_batch(self, parsed: Dict[str, "np.ndarray"]) -> Dict[str, np.ndarra flat[bounds[i] : bounds[i + 1]] for i in range(lengths.size) ] else: - counts[seg.name] = lengths + slot_lengths[seg.name] = lengths - out = self.assemble(values, counts, batch_size=batch_size) + out = self.assemble(values, slot_lengths, batch_size=batch_size) return { PROMPT_INPUT_IDS: out.input_ids, PROMPT_CU_SEQLENS: out.cu_seqlens, diff --git a/tzrec/prompt/assembler_test.py b/tzrec/prompt/assembler_test.py index 6da4f623b..449a2a841 100644 --- a/tzrec/prompt/assembler_test.py +++ b/tzrec/prompt/assembler_test.py @@ -17,7 +17,7 @@ from tzrec.prompt.plan import ( FillMode, PromptPlan, - SidSpace, + ResolvedSidSpace, SlotSeg, Static, Width, @@ -29,12 +29,12 @@ _SENTINEL = 1099 -def _sid_space(codebook=(4, 4, 4)) -> SidSpace: +def _sid_space(codebook=(4, 4, 4)) -> ResolvedSidSpace: offsets, running = [], 0 for size in codebook: offsets.append(running) running += size - return SidSpace( + return ResolvedSidSpace( codebook=tuple(codebook), num_levels=len(codebook), base_vocab=_BASE, @@ -52,7 +52,7 @@ def _slot(name, fill, width_n=None) -> SlotSeg: return SlotSeg( slot_id=0, name=name, - sources=(name,), + feature_names=(name,), group_type=FeatureGroupType.JAGGED_SEQUENCE, output_key=".sequence", fill=fill, @@ -75,7 +75,7 @@ def _plan(segments, response=(), max_length=0) -> PromptPlan: max_length=max_length, max_total_length=None, max_holes=0, - suffix_keep=None, + logits_suffix_len=None, static_prefix_len=0, length_buckets=(), projected_slots=projected, @@ -165,13 +165,13 @@ def test_inline_without_a_sid_space_is_rejected_at_construction(self) -> None: def test_column_shaped_values_are_flattened(self) -> None: # the data parser emits (total, value_dim) for a dense sequence feature from tzrec.prompt.assembler import assemble_into - from tzrec.prompt.plan import CompiledPrompt, ModulePlan + from tzrec.prompt.plan import CompiledPrompt, ProjectionPlan plan = _plan((_slot("hist", FillMode.INLINE),)) prompt = CompiledPrompt( sid_space=_sid_space(), prompt_plan=plan, - module_plan=ModulePlan(projections={}, slot_to_module={}), + projection_plan=ProjectionPlan(projections={}, slot_to_module={}), tokenizer_dir="", vocab_hash="v", plan_hash="p", diff --git a/tzrec/prompt/compile.py b/tzrec/prompt/compile.py index 599fa8bd0..ecb2e626c 100644 --- a/tzrec/prompt/compile.py +++ b/tzrec/prompt/compile.py @@ -28,17 +28,22 @@ from tzrec.prompt.plan import ( CompiledPrompt, FillMode, - ModulePlan, + ProjectionPlan, PromptPlan, + ResolvedSidSpace, Segment, - SidSpace, SlotSeg, Static, Width, WidthKind, ) from tzrec.protos.model_pb2 import FeatureGroupConfig, FeatureGroupType -from tzrec.protos.prompt_pb2 import PromptConfig, PromptProjection, PromptSlot +from tzrec.protos.prompt_pb2 import ( + PromptConfig, + PromptProjection, + PromptSlot, + SidSpace, +) from tzrec.utils.logging_util import logger _PLACEHOLDER = re.compile(r"\{\{(\w+)\}\}") @@ -116,7 +121,7 @@ def _group_type( return FeatureGroupType.JAGGED_SEQUENCE if kinds.pop() else FeatureGroupType.DEEP -def _atom_tokens(sid_space: Any) -> List[str]: +def _atom_tokens(sid_space: SidSpace) -> List[str]: """Render the SID atom tokens, one per flat index.""" fmt = sid_space.atom_token_format return [fmt.replace("{i}", str(i)) for i in range(sum(sid_space.codebook))] @@ -132,7 +137,7 @@ def _read_manifest_codebook(path: str) -> Optional[List[int]]: def _build_sid_space( cfg: PromptConfig, tok: Tokenizer, base_vocab: int, has_projection: bool -) -> Optional[SidSpace]: +) -> Optional[ResolvedSidSpace]: """Extend the tokenizer with SID atoms and resolve the token space.""" if not cfg.HasField("sid_space"): return None @@ -180,7 +185,7 @@ def _build_sid_space( lo = [base_vocab + o for o in offsets] hi = [lo[i] + codebook[i] - 1 for i in range(len(codebook))] - return SidSpace( + return ResolvedSidSpace( codebook=tuple(codebook), num_levels=len(codebook), base_vocab=base_vocab, @@ -292,7 +297,7 @@ def compile_prompt( segs[name] = SlotSeg( slot_id=slot_ids[name], name=name, - sources=tuple(slot.feature_names), + feature_names=tuple(slot.feature_names), group_type=types[name], output_key=".sequence" if seq else "", fill=fills[name], @@ -308,7 +313,7 @@ def compile_prompt( for s in body + response if isinstance(s, SlotSeg) and s.fill is FillMode.PROJECTED ) - module_plan = _build_module_plan(projected, slots) + projection_plan = _build_module_plan(projected, slots) plan = PromptPlan( segments=body, @@ -316,7 +321,7 @@ def compile_prompt( max_length=int(cfg.max_length), max_total_length=_max_total_length(body + response), max_holes=_max_holes(projected), - suffix_keep=_suffix_keep(response), + logits_suffix_len=_suffix_keep(response), static_prefix_len=_static_prefix_len(body), length_buckets=tuple(int(b) for b in cfg.length_buckets), projected_slots=projected, @@ -326,10 +331,12 @@ def compile_prompt( return CompiledPrompt( sid_space=sid_space, prompt_plan=plan, - module_plan=module_plan, + projection_plan=projection_plan, tokenizer_dir=tokenizer_dir, vocab_hash=_hash(sid_space, tok.to_str()), - plan_hash=_hash(sid_space, plan, sorted(module_plan.projections), tok.to_str()), + plan_hash=_hash( + sid_space, plan, sorted(projection_plan.projections), tok.to_str() + ), ) @@ -352,7 +359,7 @@ def _weave( def _build_module_plan( projected: Sequence[SlotSeg], slots: Dict[str, PromptSlot] -) -> ModulePlan: +) -> ProjectionPlan: """One module per distinct ``projection_name``, else one per slot.""" projections: Dict[str, PromptProjection] = {} slot_to_module: Dict[int, str] = {} @@ -375,12 +382,12 @@ def _build_module_plan( groups = tuple( FeatureGroupConfig( group_name=seg.name, - feature_names=list(seg.sources), + feature_names=list(seg.feature_names), group_type=seg.group_type, ) for seg in projected ) - return ModulePlan( + return ProjectionPlan( projections=projections, slot_to_module=slot_to_module, feature_groups=groups, @@ -396,8 +403,8 @@ def _max_total_length(segments: Sequence[Segment]) -> Optional[int]: elif seg.width.kind is WidthKind.UNBOUNDED: return None else: - assert seg.width.n is not None - total += seg.width.n + assert seg.width.num_positions is not None + total += seg.width.num_positions return total @@ -410,8 +417,8 @@ def _max_holes(projected: Sequence[SlotSeg]) -> int: f"prompt slot [{seg.name}] is PROJECTED and unbounded, so its " f"hole count is unknowable; give its members a sequence_length." ) - assert seg.width.n is not None - total += seg.width.n + assert seg.width.num_positions is not None + total += seg.width.num_positions return total @@ -434,7 +441,7 @@ def _static_prefix_len(segments: Sequence[Segment]) -> int: def _validate( - cfg: PromptConfig, plan: PromptPlan, sid_space: Optional[SidSpace] + cfg: PromptConfig, plan: PromptPlan, sid_space: Optional[ResolvedSidSpace] ) -> None: """Apply the checks that need the whole plan.""" if plan.max_length and plan.max_total_length is not None: @@ -462,7 +469,7 @@ def _validate( f"static_prefix_len is {plan.static_prefix_len}, which bounds " "what a serving prefix cache may reuse." ) - if plan.response_segments and plan.suffix_keep is None: + if plan.response_segments and plan.logits_suffix_len is None: raise ValueError( "the response has an unbounded slot, so the supervised logits " "window cannot be bounded. A decoder-only model would then " @@ -479,9 +486,10 @@ def _validate( for seg in answer: if ( seg.width.kind is WidthKind.STATIC - and seg.width.n != sid_space.num_levels + and seg.width.num_positions != sid_space.num_levels ): raise ValueError( - f"response slot [{seg.name}] is {seg.width.n} positions but " - f"the codebook has {sid_space.num_levels} levels." + f"response slot [{seg.name}] is " + f"{seg.width.num_positions} positions but the codebook has " + f"{sid_space.num_levels} levels." ) diff --git a/tzrec/prompt/compile_test.py b/tzrec/prompt/compile_test.py index a331f27e8..3594620d2 100644 --- a/tzrec/prompt/compile_test.py +++ b/tzrec/prompt/compile_test.py @@ -114,7 +114,7 @@ def test_scalar_slot_is_one_deep_position(self) -> None: self.assertIs(seg.fill, FillMode.PROJECTED) self.assertEqual(seg.output_key, "") self.assertIs(seg.width.kind, WidthKind.STATIC) - self.assertEqual(seg.width.n, 1) + self.assertEqual(seg.width.num_positions, 1) def test_manifest_mismatch_is_fatal(self) -> None: manifest = os.path.join(self.test_dir, "manifest.json") @@ -210,10 +210,10 @@ def test_answer_width_comes_from_the_codebook(self) -> None: ) # the answer is one SID item, so its width needs no sequence_length self.assertIs(seg.width.kind, WidthKind.STATIC) - self.assertEqual(seg.width.n, 3) + self.assertEqual(seg.width.num_positions, 3) # +1 because HF shifts logits: the window opens one column before the # first supervised label - self.assertEqual(compiled.prompt_plan.suffix_keep, 4) + self.assertEqual(compiled.prompt_plan.logits_suffix_len, 4) def test_unbounded_response_is_rejected(self) -> None: # with no sid_space the response has no codebook-derived width, so the @@ -247,7 +247,7 @@ def test_a_grouped_feature_inherits_the_group_cap(self) -> None: seg = next(s for s in compiled.prompt_plan.segments if isinstance(s, SlotSeg)) self.assertIs(seg.width.kind, WidthKind.BOUNDED) - self.assertEqual(seg.width.n, 16) + self.assertEqual(seg.width.num_positions, 16) if __name__ == "__main__": diff --git a/tzrec/prompt/persist.py b/tzrec/prompt/persist.py index 06184e2a1..f7b52726b 100644 --- a/tzrec/prompt/persist.py +++ b/tzrec/prompt/persist.py @@ -13,7 +13,7 @@ A checkpoint that cannot describe its own vocabulary is a checkpoint serving has to be told about out of band, which is where offline/online skew comes -from. ``ModulePlan`` is deliberately absent: it is model-only and rebuilt from +from. ``ProjectionPlan`` is deliberately absent: it is model-only and rebuilt from config at every ``__init__``. """ diff --git a/tzrec/prompt/plan.py b/tzrec/prompt/plan.py index 8ab49fad8..392120357 100644 --- a/tzrec/prompt/plan.py +++ b/tzrec/prompt/plan.py @@ -11,7 +11,7 @@ """Products of ``compile_prompt``. -This namespace disambiguates ``plan.SidSpace``, the resolved token space, from +This namespace disambiguates ``plan.ResolvedSidSpace``, the resolved token space, from ``prompt_pb2.SidSpace``, the four knobs a user declares. Nothing here stores a physical dimension: the model resolves those at ``__init__``. """ @@ -46,25 +46,25 @@ class Width: Args: kind: STATIC when the count is exact, BOUNDED when only a ceiling is known, UNBOUNDED when neither. - n: the exact count or the ceiling; None when UNBOUNDED. + num_positions: the exact count or the ceiling; None when UNBOUNDED. """ kind: WidthKind - n: Optional[int] = None + num_positions: Optional[int] = None def __post_init__(self) -> None: """Reject a count that contradicts the kind.""" if self.kind is WidthKind.UNBOUNDED: - if self.n is not None: + if self.num_positions is not None: raise ValueError("UNBOUNDED width cannot carry a count.") - elif self.n is None or self.n < 0: + elif self.num_positions is None or self.num_positions < 0: raise ValueError( f"{self.kind.name} width needs a count >= 0, got {self.n}." ) @dataclass(frozen=True) -class SidSpace: +class ResolvedSidSpace: """The resolved SID token space, read by the data layer, model and serving. Three coordinate systems and the constants that convert between them: a @@ -121,7 +121,7 @@ class SlotSeg: Args: slot_id: index into ``PromptPlan.projected_slots`` ordering. name: the placeholder name; also the derived group name. - sources: member feature names. + feature_names: member feature names. group_type: DEEP or JAGGED_SEQUENCE. output_key: "" for DEEP, ".sequence" otherwise. fill: INLINE writes token ids, PROJECTED writes sentinels and a hole. @@ -131,7 +131,7 @@ class SlotSeg: slot_id: int name: str - sources: Tuple[str, ...] + feature_names: Tuple[str, ...] group_type: "FeatureGroupType.ValueType" output_key: str fill: FillMode @@ -152,7 +152,7 @@ class PromptPlan: max_length: validation ceiling; an over-long row is an error. max_total_length: proven ceiling when every slot is bounded, else None. max_holes: per-row projected-position ceiling, not a runtime shape. - suffix_keep: upper bound on the supervised logits window. + logits_suffix_len: upper bound on the supervised logits window. static_prefix_len: leading positions that are request-invariant. length_buckets: sampler and graph-capture buckets. projected_slots: fixes the order hole positions are written in. @@ -163,14 +163,14 @@ class PromptPlan: max_length: int max_total_length: Optional[int] max_holes: int - suffix_keep: Optional[int] + logits_suffix_len: Optional[int] static_prefix_len: int length_buckets: Tuple[int, ...] projected_slots: Tuple[SlotSeg, ...] @dataclass(frozen=True) -class ModulePlan: +class ProjectionPlan: """Projection topology. Model-only, never persisted. Args: @@ -194,15 +194,15 @@ class CompiledPrompt: Args: sid_space: the resolved SID token space. prompt_plan: assembler walk order and ceilings. - module_plan: projection topology. + projection_plan: projection topology. tokenizer_dir: where the extended tokenizer was written. vocab_hash: over sid_space and tokenizer.json; fatal on mismatch. plan_hash: over all four parts; warns on mismatch. """ - sid_space: Optional[SidSpace] + sid_space: Optional[ResolvedSidSpace] prompt_plan: PromptPlan - module_plan: ModulePlan + projection_plan: ProjectionPlan tokenizer_dir: str vocab_hash: str plan_hash: str diff --git a/tzrec/tests/prompt_integration_test.py b/tzrec/tests/prompt_integration_test.py index e2b266125..dcb8e56a9 100644 --- a/tzrec/tests/prompt_integration_test.py +++ b/tzrec/tests/prompt_integration_test.py @@ -154,7 +154,7 @@ def test_every_row_is_supervised_whatever_its_length(self) -> None: int(infos[PROMPT_MAX_SEQLEN]), -100, ) - window = labels[:, -self.prompt.prompt_plan.suffix_keep :] + window = labels[:, -self.prompt.prompt_plan.logits_suffix_len :] supervised = (window != -100).sum(dim=1).tolist() self.assertEqual(supervised[0], supervised[1]) self.assertEqual(supervised[0], self.prompt.sid_space.num_levels) @@ -256,7 +256,7 @@ def setUp(self) -> None: self.prompt = compile_prompt(cfg, self.features, model_dir=self.test_dir) def test_compiler_derives_one_group_per_projected_slot(self) -> None: - groups = self.prompt.module_plan.feature_groups + groups = self.prompt.projection_plan.feature_groups self.assertEqual([g.group_name for g in groups], ["prof"]) self.assertEqual(list(groups[0].feature_names), ["prof"]) # the INLINE slot produces none: its tokens are already in the stream From 11a980fe548aaa7edc35fe0480d64bea3926ffdc Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Wed, 12 Aug 2026 09:54:26 +0000 Subject: [PATCH 84/99] [perf] skip mean resizing when the backbone is still empty The model builds its LM from config, so the rows this resize creates are random and are overwritten in every path: by load_state_dict in init_from_pretrained on a cold start, by the DCP restore otherwise. Only the shape matters, but mean_resizing fits a distribution over the old embedding matrix and samples from it, which costs 1.8x the plain resize at hidden 2048 and grows with hidden size. The cold-start call in init_from_pretrained keeps mean_resizing, since there the distribution is the real pretrained one and it is what seeds the new SID atom rows. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_model.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tzrec/models/prompt_generative_model.py b/tzrec/models/prompt_generative_model.py index b8220a3cb..2dec03f61 100644 --- a/tzrec/models/prompt_generative_model.py +++ b/tzrec/models/prompt_generative_model.py @@ -87,8 +87,10 @@ def __init__( cfg = self._model_config self.lm = self._build_backbone(cfg.hf_model_id, cfg.common.param_dtype) + # only the shape matters here: every path overwrites these rows, from + # init_from_pretrained on a cold start or from the DCP restore otherwise self.lm.resize_token_embeddings( - prompt.sid_space.target_vocab, mean_resizing=True + prompt.sid_space.target_vocab, mean_resizing=False ) self.embedding_group = EmbeddingGroup( self._features, list(self._prompt.projection_plan.feature_groups) From ebac64c54d4441e71b382c14f3ebfa5ec06f93a1 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Thu, 13 Aug 2026 02:18:09 +0000 Subject: [PATCH 85/99] [chore] bump version to 1.3.15 Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tzrec/version.py b/tzrec/version.py index 021551003..73de17c8e 100644 --- a/tzrec/version.py +++ b/tzrec/version.py @@ -9,4 +9,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "1.3.14" +__version__ = "1.3.15" From 6c64e368a5b862d4bd695cdee6feb52e523eecda Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Thu, 13 Aug 2026 05:59:59 +0000 Subject: [PATCH 86/99] [refactor] simplify prompt assembly tests --- tzrec/models/prompt_generative_model_test.py | 45 ++---- tzrec/models/prompt_generative_qwen_test.py | 82 +++-------- tzrec/modules/dynamic_beam_test.py | 39 +----- tzrec/modules/prompt_projection_test.py | 5 - tzrec/optim/lr_scheduler_test.py | 9 -- tzrec/prompt/assembler.py | 138 +++++++++---------- tzrec/prompt/assembler_test.py | 19 +-- tzrec/prompt/compile_test.py | 16 +-- tzrec/prompt/persist_test.py | 23 ++-- tzrec/tests/prompt_integration_test.py | 126 ++--------------- tzrec/tests/prompt_test_util.py | 38 +++++ tzrec/utils/hf_export_util_test.py | 22 --- 12 files changed, 161 insertions(+), 401 deletions(-) create mode 100644 tzrec/tests/prompt_test_util.py diff --git a/tzrec/models/prompt_generative_model_test.py b/tzrec/models/prompt_generative_model_test.py index 38d5fcdf5..a477f980e 100644 --- a/tzrec/models/prompt_generative_model_test.py +++ b/tzrec/models/prompt_generative_model_test.py @@ -25,12 +25,12 @@ from tzrec.prompt.assembler import ( PROMPT_HOLE_POSITIONS, PROMPT_INPUT_IDS, - assemble_into, ) from tzrec.prompt.compile import compile_prompt from tzrec.protos import feature_pb2 from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.prompt_pb2 import PromptConfig +from tzrec.tests.prompt_test_util import assemble_into from tzrec.utils.state_dict_util import init_parameters from tzrec.utils.test_util import create_tiny_causal_lm, make_test_dir @@ -107,41 +107,22 @@ def _batch(self, parsed, prompt=None, sparse=None): ) return batch - def test_tokens_to_local_codes_undoes_both_shifts(self) -> None: + def test_tokens_to_local_codes_undoes_shifts_and_groups_beams(self) -> None: model = self._model() space = self.prompt.sid_space - # one row, one beam: level 0 code 1, level 1 code 2, level 2 code 3 - tokens = torch.tensor( + local_codes = torch.tensor( [ - [ - space.base_vocab + space.level_offsets[0] + 1, - space.base_vocab + space.level_offsets[1] + 2, - space.base_vocab + space.level_offsets[2] + 3, - ] + [0, 1, 3], + [3, 0, 2], + [1, 3, 0], + [2, 2, 1], ] ) - codes = model._tokens_to_local_codes(tokens, batch_size=1) - - self.assertEqual(codes.shape, (1, 1, space.num_levels)) - self.assertEqual(codes[0, 0].tolist(), [1, 2, 3]) - - def test_tokens_to_local_codes_maps_band_edges_to_the_last_code(self) -> None: - model = self._model() - space = self.prompt.sid_space - codes = model._tokens_to_local_codes( - torch.tensor([list(space.band_hi)]), batch_size=1 - ) - - self.assertEqual(codes[0, 0].tolist(), [c - 1 for c in _CODEBOOK]) - - def test_tokens_to_local_codes_groups_beams_under_their_row(self) -> None: - model = self._model() - space = self.prompt.sid_space - base = torch.tensor(space.level_offsets) + space.base_vocab - # four beam rows over a batch of two - codes = model._tokens_to_local_codes(base.repeat(4, 1), batch_size=2) + tokens = local_codes + torch.tensor(space.level_offsets) + space.base_vocab + codes = model._tokens_to_local_codes(tokens, batch_size=2) self.assertEqual(codes.shape, (2, 2, space.num_levels)) + self.assertEqual(codes.tolist(), local_codes.reshape(2, 2, -1).tolist()) def test_rejects_a_model_built_without_a_prompt(self) -> None: with self.assertRaisesRegex(ValueError, "needs a compiled prompt"): @@ -176,7 +157,7 @@ def test_shared_projection_name_requires_matching_widths(self) -> None: with self.assertRaisesRegex(ValueError, "cannot share a module"): self._model(features=features, prompt=prompt) - def test_projected_slot_overwrites_only_its_sentinel_positions(self) -> None: + def test_projected_slot_overwrites_sentinels_and_backpropagates(self) -> None: features = [_feature(_HIST), _feature(_ANSWER), _feature(_projected("prof", 8))] prompt = self._compile( features, @@ -211,6 +192,10 @@ def test_projected_slot_overwrites_only_its_sentinel_positions(self) -> None: changed = ~torch.isclose(embeds, raw).all(dim=-1) self.assertEqual(sorted(changed.nonzero().flatten().tolist()), holes.tolist()) + embeds[holes].sum().backward() + proj = next(iter(model.projections.values())) + self.assertIsNotNone(proj.head.weight.grad) + def test_loss_surfaces_the_ce_computed_in_predict(self) -> None: model = self._model() value = torch.tensor(1.25) diff --git a/tzrec/models/prompt_generative_qwen_test.py b/tzrec/models/prompt_generative_qwen_test.py index d130e464c..69a4e0e4c 100644 --- a/tzrec/models/prompt_generative_qwen_test.py +++ b/tzrec/models/prompt_generative_qwen_test.py @@ -20,76 +20,28 @@ class UnpackTest(unittest.TestCase): """The one adapter where padding lives.""" def test_packs_rows_of_different_lengths(self) -> None: - # rows of 2 and 3 tokens, hidden size 4 - embeds = torch.arange(20, dtype=torch.float32).reshape(5, 4) - cu = torch.tensor([0, 2, 5]) - labels = torch.tensor([10, 11, 20, 21, 22]) - - padded, mask, out = _unpack(embeds, cu, labels, max_seqlen=3, ignore_index=-100) - - self.assertEqual(padded.shape, (2, 3, 4)) - # pads go on the left, so every row ends on a real token - self.assertEqual(mask.tolist(), [[0, 1, 1], [1, 1, 1]]) - torch.testing.assert_close(padded[0, 1:], embeds[:2]) - torch.testing.assert_close(padded[1, :3], embeds[2:]) - # the pad column is zero, and its label is ignored - torch.testing.assert_close(padded[0, 0], torch.zeros(4)) - self.assertEqual(out.tolist(), [[-100, 10, 11], [20, 21, 22]]) - - def test_uses_the_given_width_not_the_observed_max(self) -> None: - # the collator's bucket may exceed the widest row; §7.4 forbids - # deriving the width on device - embeds = torch.ones(3, 2) - cu = torch.tensor([0, 1, 3]) - labels = torch.tensor([1, 2, 3]) - padded, mask, _ = _unpack(embeds, cu, labels, max_seqlen=5, ignore_index=-100) - - self.assertEqual(padded.shape, (2, 5, 2)) - self.assertEqual(mask.sum().item(), 3) - - def test_row_order_survives_the_scatter(self) -> None: - # mask selects row-major, which must match the packing order - embeds = torch.tensor([[1.0], [2.0], [3.0], [4.0]]) - cu = torch.tensor([0, 1, 4]) - labels = torch.zeros(4, dtype=torch.long) - padded, _, _ = _unpack(embeds, cu, labels, max_seqlen=3, ignore_index=-100) - - self.assertEqual(padded[0, :, 0].tolist(), [0.0, 0.0, 1.0]) - self.assertEqual(padded[1, :, 0].tolist(), [2.0, 3.0, 4.0]) - - def test_every_row_ends_on_its_own_last_token(self) -> None: - # decode prefills from [:, -1, :]; a right-padded short row would - # hand it padding instead of the row's final token - embeds = torch.tensor([[1.0], [2.0], [3.0], [11.0], [12.0], [13.0], [14.0]]) - cu = torch.tensor([0, 3, 7]) - labels = torch.zeros(7, dtype=torch.long) - padded, _, _ = _unpack(embeds, cu, labels, max_seqlen=4, ignore_index=-100) - - self.assertEqual(padded[:, -1, 0].tolist(), [3.0, 14.0]) - - def test_a_short_row_keeps_its_answer_in_the_suffix_window(self) -> None: - # the loss scores a fixed-width suffix; every row must contribute the - # same number of supervised positions regardless of its length - embeds = torch.ones(9, 1) + # rows of 4 and 5 tokens, with a padded width larger than either row + embeds = torch.arange(18, dtype=torch.float32).reshape(9, 2) cu = torch.tensor([0, 4, 9]) ignore = -100 - # each row ends on 2 real labels, preceded by unsupervised context labels = torch.tensor([ignore, ignore, 7, 8, ignore, ignore, ignore, 7, 8]) - _, _, out = _unpack(embeds, cu, labels, max_seqlen=5, ignore_index=ignore) - - window = out[:, -2:] - self.assertEqual((window != ignore).sum(dim=1).tolist(), [2, 2]) - self.assertEqual(window.tolist(), [[7, 8], [7, 8]]) - def test_gradient_reaches_the_packed_input(self) -> None: - embeds = torch.ones(3, 2, requires_grad=True) - cu = torch.tensor([0, 1, 3]) - labels = torch.zeros(3, dtype=torch.long) - padded, _, _ = _unpack(embeds, cu, labels, max_seqlen=2, ignore_index=-100) - padded.sum().backward() + padded, mask, out = _unpack( + embeds, cu, labels, max_seqlen=7, ignore_index=ignore + ) - self.assertIsNotNone(embeds.grad) - torch.testing.assert_close(embeds.grad, torch.ones(3, 2)) + self.assertEqual(padded.shape, (2, 7, 2)) + # pads go on the left, so every row ends on a real token + self.assertEqual( + mask.tolist(), + [[0, 0, 0, 1, 1, 1, 1], [0, 0, 1, 1, 1, 1, 1]], + ) + torch.testing.assert_close(padded[0, 3:], embeds[:4]) + torch.testing.assert_close(padded[1, 2:], embeds[4:]) + torch.testing.assert_close(padded[0, :3], torch.zeros(3, 2)) + torch.testing.assert_close(padded[1, :2], torch.zeros(2, 2)) + torch.testing.assert_close(padded[:, -1], torch.stack([embeds[3], embeds[8]])) + self.assertEqual(out.tolist(), [[ignore] * 5 + [7, 8]] * 2) if __name__ == "__main__": diff --git a/tzrec/modules/dynamic_beam_test.py b/tzrec/modules/dynamic_beam_test.py index 7726482f7..b6fcb04eb 100644 --- a/tzrec/modules/dynamic_beam_test.py +++ b/tzrec/modules/dynamic_beam_test.py @@ -9,9 +9,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import ast import itertools -import pathlib import unittest from typing import Any, Dict, List, Tuple @@ -19,7 +17,6 @@ import torch.nn.functional as F from parameterized import parameterized -from tzrec.modules import dynamic_beam from tzrec.modules.dynamic_beam import dynamic_beam_search from tzrec.utils.test_util import create_tiny_causal_lm, parameterized_name_func @@ -80,26 +77,12 @@ def _bruteforce_scores(lm, input_ids, attention_mask, pairs): class DynamicBeamSearchTest(unittest.TestCase): - def test_module_declares_no_tzrec_imports(self) -> None: - # the "torch-only, no tzrec deps" docstring is what makes it liftable. - tree = ast.parse(pathlib.Path(dynamic_beam.__file__).read_text()) - mods = set() - for node in ast.walk(tree): - if isinstance(node, ast.Import): - mods.update(a.name for a in node.names) - elif isinstance(node, ast.ImportFrom) and node.module: - mods.add(node.module) - self.assertEqual( - {m for m in mods if m == "tzrec" or m.startswith("tzrec.")}, set() - ) - @parameterized.expand( [ # every request satisfiable -> the schedule is honoured verbatim [[4, 8, 16], [(20, 27), (28, 34), (35, 40)], [4, 8, 16]], # capped by band x surviving prefixes, not by what was asked [[6, 12], [(20, 21), (22, 24)], [2, 6]], - [[2, 4, 8], [(20, 21), (22, 24), (25, 28)], [2, 4, 8]], # a width-1 band collapses level 0 and bounds everything after it [[4, 8], [(20, 20), (21, 23)], [1, 3]], # a flat (non-doubling) schedule is just as valid to the kernel @@ -126,23 +109,8 @@ def test_rejects_a_schedule_that_does_not_match_the_bands(self) -> None: with self.assertRaisesRegex(ValueError, "must be >= 1"): _decode(lm, ids, pairs, beam_widths=[2, 0, 4]) - def test_tokens_stay_inside_arbitrary_bands(self) -> None: - # bands the PromptGenerativeQwen caller can never produce: descending, disjoint, - # unequal width -- the kernel's contract is per-level (lo, hi), not a - # contiguous codebook layout. - pairs = [(5, 6), (20, 24), (11, 13)] - ids = torch.tensor([[1, 2, 3, 4]]) - out = _decode(create_tiny_causal_lm(vocab_size=30), ids, pairs) - # flat width 8: level 0's 2-wide band caps it, later levels recover - self.assertEqual(tuple(out.shape), (8, 3)) - for level, (lo_j, hi_j) in enumerate(pairs): - col = out[:, level] - self.assertTrue(bool((col >= lo_j).all())) - self.assertTrue(bool((col <= hi_j).all())) - self.assertEqual(len({tuple(r) for r in out.tolist()}), out.shape[0]) - @parameterized.expand( - [[0, 1], [1, 3], [2, 16], [0, 16]], + [[0, 1], [1, 16]], name_func=parameterized_name_func, ) def test_left_padding_matches_unpadded(self, seed, n_pad) -> None: @@ -173,8 +141,9 @@ def test_ragged_batch_rows_match_solo_runs(self) -> None: def test_exhaustive_matches_bruteforce_topk(self) -> None: # widths [2, 6, 12] over 2*3*2 = 12 combinations -> no pruning at any - # level, so the beam must reproduce the exact full-recompute ranking. - pairs = [(20, 21), (22, 24), (25, 26)] + # level. Disjoint, non-monotonic bands also verify that each level uses + # its own bounds rather than assuming one contiguous SID layout. + pairs = [(20, 21), (5, 7), (25, 26)] lm = create_tiny_causal_lm(vocab_size=30) ids = torch.tensor([[5, 6, 7, 8]]) am = torch.ones_like(ids) diff --git a/tzrec/modules/prompt_projection_test.py b/tzrec/modules/prompt_projection_test.py index 6a75e1747..8bbfc87bc 100644 --- a/tzrec/modules/prompt_projection_test.py +++ b/tzrec/modules/prompt_projection_test.py @@ -52,11 +52,6 @@ def test_bias_is_configurable(self) -> None: proj = PromptProjection(config, in_dim=4, hidden_size=6) self.assertIsNone(proj.head.bias) - def test_jagged_rows_project_independently(self) -> None: - proj = PromptProjection(PromptProjectionConfig(), in_dim=4, hidden_size=6) - rows = torch.randn(7, 4) - torch.testing.assert_close(proj(rows)[3], proj(rows[3:4])[0]) - if __name__ == "__main__": unittest.main() diff --git a/tzrec/optim/lr_scheduler_test.py b/tzrec/optim/lr_scheduler_test.py index 93f2c539f..7c1814dbb 100644 --- a/tzrec/optim/lr_scheduler_test.py +++ b/tzrec/optim/lr_scheduler_test.py @@ -83,15 +83,6 @@ def test_manual_step_lr_with_warmup(self) -> None: lr.step() self.assertAlmostEqual(opt.param_groups[0]["lr"], lr_gt) - def test_linear_decay_lr(self) -> None: - params = [torch.tensor([1.0, 2.0])] - opt = torch.optim.Adam(params, lr=0.01) - lr = lr_scheduler.LinearDecayLR(opt, num_training_steps=4) - lr_gts = [0.0075, 0.005, 0.0025, 0.0, 0.0] - for lr_gt in lr_gts: - lr.step() - self.assertAlmostEqual(opt.param_groups[0]["lr"], lr_gt) - def test_linear_decay_lr_with_min_lr(self) -> None: params = [torch.tensor([1.0, 2.0])] opt = torch.optim.Adam(params, lr=0.01) diff --git a/tzrec/prompt/assembler.py b/tzrec/prompt/assembler.py index dbaa90664..1d9e4e5f6 100644 --- a/tzrec/prompt/assembler.py +++ b/tzrec/prompt/assembler.py @@ -22,7 +22,6 @@ import numpy as np from tzrec.prompt.plan import ( - CompiledPrompt, FillMode, PromptPlan, ResolvedSidSpace, @@ -42,10 +41,10 @@ class AssembledPrompt: """One batch of assembled prompts, in packed varlen form. Args: - input_ids: every row's tokens concatenated, ``(total_tokens,)``. - cu_seqlens: row boundaries into ``input_ids``, ``(batch_size + 1,)``. + input_ids: every sample's tokens concatenated, ``(total_tokens,)``. + cu_seqlens: sample boundaries into ``input_ids``, ``(batch_size + 1,)``. hole_positions: absolute indices the projected embeddings overwrite, - in ``PromptPlan.projected_slots`` order within each row. + in ``PromptPlan.projected_slots`` order within each sample. labels: ``ignore_index`` outside the response span. """ @@ -56,7 +55,7 @@ class AssembledPrompt: @property def max_seqlen(self) -> int: - """Widest row, computed on the host so the model never derives it.""" + """Widest sample, computed on the host so the model never derives it.""" if self.cu_seqlens.size < 2: return 0 return int(np.max(np.diff(self.cu_seqlens))) @@ -119,94 +118,105 @@ def _inline_tokens(self, name: str, values: np.ndarray) -> np.ndarray: ) return values.astype(np.int64, copy=False) + self._sid_space.base_vocab - def _emit_row( + def _emit_sample( self, segments: Sequence[object], - row: int, - values: Dict[str, List[np.ndarray]], - slot_lengths: Dict[str, np.ndarray], + sample_index: int, + inline_values: Dict[str, List[np.ndarray]], + projected_lengths: Dict[str, np.ndarray], out: List[int], holes: List[int], - base: int, + sample_start: int, ) -> None: - """Append one row's tokens for one segment list, recording holes.""" + """Append one sample's tokens for one segment list, recording holes.""" for seg in segments: if isinstance(seg, Static): out.extend(seg.token_ids) continue assert isinstance(seg, SlotSeg) if seg.fill is FillMode.INLINE: - out.extend(self._inline_tokens(seg.name, values[seg.name][row])) + out.extend( + self._inline_tokens(seg.name, inline_values[seg.name][sample_index]) + ) else: assert self._sid_space is not None - width = int(slot_lengths[seg.name][row]) - holes.extend(range(base + len(out), base + len(out) + width)) + width = int(projected_lengths[seg.name][sample_index]) + holes.extend( + range( + sample_start + len(out), + sample_start + len(out) + width, + ) + ) out.extend([self._sid_space.sentinel_token_id] * width) def assemble( self, - values: Dict[str, List[np.ndarray]], - slot_lengths: Optional[Dict[str, np.ndarray]] = None, + inline_values: Dict[str, List[np.ndarray]], + projected_lengths: Optional[Dict[str, np.ndarray]] = None, batch_size: Optional[int] = None, ) -> AssembledPrompt: """Assemble one batch. Args: - values: INLINE slot name to its per-row value arrays. - slot_lengths: PROJECTED slot name to its per-row position count. - batch_size: row count; inferred from ``values`` when omitted. + inline_values: INLINE slot name to its per-sample value arrays. + projected_lengths: PROJECTED slot name to its per-sample position count. + batch_size: sample count; inferred from ``inline_values`` when omitted. Returns: The packed streams. """ - slot_lengths = slot_lengths or {} + projected_lengths = projected_lengths or {} if batch_size is None: - if not values: + if not inline_values: raise ValueError("batch_size is required when no INLINE slot exists.") - batch_size = len(next(iter(values.values()))) + first_inline_values = next(iter(inline_values.values())) + batch_size = len(first_inline_values) - ids: List[int] = [] + jagged_token_ids: List[int] = [] labels: List[int] = [] holes: List[int] = [] cu_seqlens = [0] - for row in range(batch_size): - row_ids: List[int] = [] - self._emit_row( + for sample_index in range(batch_size): + sample_token_ids: List[int] = [] + self._emit_sample( self._prompt_plan.segments, - row, - values, - slot_lengths, - row_ids, + sample_index, + inline_values, + projected_lengths, + sample_token_ids, holes, - len(ids), + len(jagged_token_ids), ) - prompt_len = len(row_ids) - self._emit_row( + prompt_len = len(sample_token_ids) + self._emit_sample( self._prompt_plan.response_segments, - row, - values, - slot_lengths, - row_ids, + sample_index, + inline_values, + projected_lengths, + sample_token_ids, holes, - len(ids), + len(jagged_token_ids), ) # supervision covers the response span only; the prompt is context. - row_labels = [self._ignore_index] * prompt_len + row_ids[prompt_len:] + sample_labels = [self._ignore_index] * prompt_len + sample_token_ids[ + prompt_len: + ] if ( self._prompt_plan.max_length - and len(row_ids) > self._prompt_plan.max_length + and len(sample_token_ids) > self._prompt_plan.max_length ): raise ValueError( - f"assembled row {row} is {len(row_ids)} tokens, over " - f"max_length {self._prompt_plan.max_length}. Rows are never " - f"truncated: cap the source features instead." + f"assembled sample {sample_index} is " + f"{len(sample_token_ids)} tokens, over " + f"max_length {self._prompt_plan.max_length}. Samples are " + f"never truncated: cap the source features instead." ) - ids.extend(row_ids) - labels.extend(row_labels) - cu_seqlens.append(len(ids)) + jagged_token_ids.extend(sample_token_ids) + labels.extend(sample_labels) + cu_seqlens.append(len(jagged_token_ids)) return AssembledPrompt( - input_ids=np.asarray(ids, dtype=np.int64), + input_ids=np.asarray(jagged_token_ids, dtype=np.int64), cu_seqlens=np.asarray(cu_seqlens, dtype=np.int64), hole_positions=np.asarray(holes, dtype=np.int64), labels=np.asarray(labels, dtype=np.int64), @@ -222,8 +232,8 @@ def assemble_batch(self, parsed: Dict[str, "np.ndarray"]) -> Dict[str, np.ndarra Returns: The five streams, keyed as ``additional_infos`` expects them. """ - values: Dict[str, List[np.ndarray]] = {} - slot_lengths: Dict[str, np.ndarray] = {} + inline_values: Dict[str, List[np.ndarray]] = {} + projected_lengths: Dict[str, np.ndarray] = {} batch_size = 0 for seg in self._prompt_plan.segments + self._prompt_plan.response_segments: if not isinstance(seg, SlotSeg): @@ -236,13 +246,13 @@ def assemble_batch(self, parsed: Dict[str, "np.ndarray"]) -> Dict[str, np.ndarra # is one code per position, so value_dim is always 1 here flat = np.asarray(parsed[f"{source}.values"]).reshape(-1) bounds = np.concatenate(([0], np.cumsum(lengths))) - values[seg.name] = [ + inline_values[seg.name] = [ flat[bounds[i] : bounds[i + 1]] for i in range(lengths.size) ] else: - slot_lengths[seg.name] = lengths + projected_lengths[seg.name] = lengths - out = self.assemble(values, slot_lengths, batch_size=batch_size) + out = self.assemble(inline_values, projected_lengths, batch_size=batch_size) return { PROMPT_INPUT_IDS: out.input_ids, PROMPT_CU_SEQLENS: out.cu_seqlens, @@ -250,27 +260,3 @@ def assemble_batch(self, parsed: Dict[str, "np.ndarray"]) -> Dict[str, np.ndarra PROMPT_LABELS: out.labels, PROMPT_MAX_SEQLEN: np.asarray(out.max_seqlen, dtype=np.int64), } - - -def assemble_into( - prompt: CompiledPrompt, - parsed: Dict[str, "np.ndarray"], - ignore_index: int = -100, -) -> Dict[str, np.ndarray]: - """Assemble one batch with a throwaway assembler. - - A caller that assembles every batch should hold a ``PromptAssembler`` and - call ``assemble_batch`` instead; this builds one per call. - - Args: - prompt: the compiled prompt. - parsed: ``{feature}.values`` / ``{feature}.lengths`` as the data parser - emits them. - ignore_index: label value outside the supervised span. - - Returns: - The five streams, keyed as ``additional_infos`` expects them. - """ - return PromptAssembler( - prompt.prompt_plan, prompt.sid_space, ignore_index - ).assemble_batch(parsed) diff --git a/tzrec/prompt/assembler_test.py b/tzrec/prompt/assembler_test.py index 449a2a841..c85f1f310 100644 --- a/tzrec/prompt/assembler_test.py +++ b/tzrec/prompt/assembler_test.py @@ -24,6 +24,7 @@ WidthKind, ) from tzrec.protos.model_pb2 import FeatureGroupType +from tzrec.tests.prompt_test_util import assemble_into _BASE = 1000 _SENTINEL = 1099 @@ -109,14 +110,6 @@ def test_projected_emits_sentinels_and_records_holes(self) -> None: # absolute indices into the flat buffer, which is what index_copy needs self.assertEqual(out.hole_positions.tolist(), [1, 2, 4, 5, 6]) - def test_hole_positions_index_the_flat_buffer_exactly(self) -> None: - plan = _plan((_slot("prof", FillMode.PROJECTED, 2),)) - asm = PromptAssembler(plan, _sid_space()) - out = asm.assemble({}, {"prof": np.array([2, 2])}, batch_size=2) - # index_copy requires index.numel() == source.size(0) - self.assertEqual(out.hole_positions.size, 4) - self.assertTrue(np.all(out.input_ids[out.hole_positions] == _SENTINEL)) - def test_labels_cover_the_response_span_only(self) -> None: plan = _plan( (Static((7, 8)),), @@ -164,7 +157,6 @@ def test_inline_without_a_sid_space_is_rejected_at_construction(self) -> None: def test_column_shaped_values_are_flattened(self) -> None: # the data parser emits (total, value_dim) for a dense sequence feature - from tzrec.prompt.assembler import assemble_into from tzrec.prompt.plan import CompiledPrompt, ProjectionPlan plan = _plan((_slot("hist", FillMode.INLINE),)) @@ -184,15 +176,6 @@ def test_column_shaped_values_are_flattened(self) -> None: self.assertEqual(out["prompt_cu_seqlens"].tolist(), [0, 3, 6]) self.assertEqual(out["prompt_input_ids"].tolist()[0], _BASE + 1) - def test_rows_of_different_lengths_pack_without_padding(self) -> None: - plan = _plan((_slot("hist", FillMode.INLINE),)) - asm = PromptAssembler(plan, _sid_space()) - out = asm.assemble( - {"hist": [np.array([1, 6, 11]), np.array([0, 4, 8, 2, 5, 9])]} - ) - self.assertEqual(out.cu_seqlens.tolist(), [0, 3, 9]) - self.assertEqual(out.input_ids.size, 9) - if __name__ == "__main__": unittest.main() diff --git a/tzrec/prompt/compile_test.py b/tzrec/prompt/compile_test.py index 3594620d2..96997242a 100644 --- a/tzrec/prompt/compile_test.py +++ b/tzrec/prompt/compile_test.py @@ -91,6 +91,9 @@ def test_inline_needs_no_group_projected_gets_one(self) -> None: self.assertEqual( [s.name for s in compiled.prompt_plan.projected_slots], ["prof"] ) + groups = compiled.projection_plan.feature_groups + self.assertEqual([g.group_name for g in groups], ["prof"]) + self.assertEqual(list(groups[0].feature_names), ["prof"]) self.assertEqual(compiled.prompt_plan.max_holes, 4) self.assertIsNotNone(compiled.sid_space.sentinel_token_id) @@ -173,19 +176,6 @@ def test_atoms_absent_from_the_base_tokenizer(self) -> None: with self.assertRaisesRegex(ValueError, "already in the base tokenizer"): self._compile(cfg, [_feature(_HIST)]) - def test_vocab_hash_tracks_the_codebook(self) -> None: - def compile_with(sizes): - cfg = self._config(prompt="History : {{hist}}") - cfg.sid_space.codebook.extend(sizes) - return self._compile(cfg, [_feature(_HIST)]) - - self.assertEqual( - compile_with([4, 4]).vocab_hash, compile_with([4, 4]).vocab_hash - ) - self.assertNotEqual( - compile_with([4, 4]).vocab_hash, compile_with([4, 8]).vocab_hash - ) - def test_extended_tokenizer_is_written(self) -> None: cfg = self._config(prompt="History : {{hist}}") cfg.sid_space.codebook.extend([4, 4]) diff --git a/tzrec/prompt/persist_test.py b/tzrec/prompt/persist_test.py index 159249f57..762430220 100644 --- a/tzrec/prompt/persist_test.py +++ b/tzrec/prompt/persist_test.py @@ -86,12 +86,6 @@ def test_sid_space_round_trips_as_plain_json(self) -> None: set(space), {f.name for f in dataclasses.fields(prompt.sid_space)} ) - def test_matching_prompt_passes(self) -> None: - prompt = self._compile() - ckpt = os.path.join(self.test_dir, "model.ckpt-1") - save_prompt_assets(prompt, ckpt) - check_prompt_assets(self._compile(), ckpt) - def test_a_changed_codebook_is_fatal(self) -> None: ckpt = os.path.join(self.test_dir, "model.ckpt-1") save_prompt_assets(self._compile(codebook=(4, 4, 4)), ckpt) @@ -105,16 +99,23 @@ def test_a_changed_template_only_warns(self) -> None: # the vocabulary is untouched, so the weights are still usable self.assertEqual(moved.vocab_hash, read_prompt_hashes(ckpt)["vocab_hash"]) self.assertNotEqual(moved.plan_hash, read_prompt_hashes(ckpt)["plan_hash"]) - check_prompt_assets(moved, ckpt) + with mock.patch("tzrec.prompt.persist.logger.warning") as warning: + check_prompt_assets(moved, ckpt) + warning.assert_called_once() def test_a_checkpoint_without_assets_only_warns(self) -> None: bare = os.path.join(self.test_dir, "model.ckpt-bare") os.makedirs(bare, exist_ok=True) self.assertIsNone(read_prompt_hashes(bare)) - check_prompt_assets(self._compile(), bare) + prompt = self._compile() + with mock.patch("tzrec.prompt.persist.logger.warning") as warning: + check_prompt_assets(prompt, bare) + warning.assert_called_once() def test_no_prompt_config_is_a_no_op(self) -> None: - check_prompt_assets(None, os.path.join(self.test_dir, "nowhere")) + with mock.patch("tzrec.prompt.persist.logger.warning") as warning: + check_prompt_assets(None, os.path.join(self.test_dir, "nowhere")) + warning.assert_not_called() def test_export_carries_the_contract_forward(self) -> None: ckpt = os.path.join(self.test_dir, "model.ckpt-1") @@ -134,7 +135,9 @@ def test_copying_from_a_bare_checkpoint_only_warns(self) -> None: bare = os.path.join(self.test_dir, "bare") os.makedirs(bare, exist_ok=True) export = os.path.join(self.test_dir, "export_bare") - copy_prompt_assets(bare, export) + with mock.patch("tzrec.prompt.persist.logger.warning") as warning: + copy_prompt_assets(bare, export) + warning.assert_called_once() self.assertIsNone(read_prompt_hashes(export)) def test_only_rank_zero_writes(self) -> None: diff --git a/tzrec/tests/prompt_integration_test.py b/tzrec/tests/prompt_integration_test.py index dcb8e56a9..34e43693f 100644 --- a/tzrec/tests/prompt_integration_test.py +++ b/tzrec/tests/prompt_integration_test.py @@ -26,12 +26,12 @@ PROMPT_CU_SEQLENS, PROMPT_LABELS, PROMPT_MAX_SEQLEN, - assemble_into, ) from tzrec.prompt.compile import compile_prompt from tzrec.protos import feature_pb2 from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.prompt_pb2 import PromptConfig +from tzrec.tests.prompt_test_util import assemble_into from tzrec.utils.test_util import create_tiny_causal_lm, make_test_dir _CODEBOOK = [4, 4, 4] @@ -53,17 +53,11 @@ def _tokenizer(path: str) -> str: return path -_PROF = ( - 'sequence_id_feature { feature_name: "prof" expression: "user:prof" ' - "num_buckets: 32 embedding_dim: 8 sequence_length: 4 }" -) - - -def _features(extra=()): +def _features(): text = ( 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }', 'sequence_raw_feature { feature_name: "answer" expression: "item:answer" }', - ) + tuple(extra) + ) out = [] for one in text: fc = feature_pb2.FeatureConfig() @@ -159,14 +153,6 @@ def test_every_row_is_supervised_whatever_its_length(self) -> None: self.assertEqual(supervised[0], supervised[1]) self.assertEqual(supervised[0], self.prompt.sid_space.num_levels) - def test_compiles_a_usable_space(self) -> None: - space = self.prompt.sid_space - self.assertEqual(space.num_levels, 3) - self.assertEqual(space.sid_vocab_size, 12) - # the atoms sit immediately above the base vocabulary - self.assertEqual(space.band_lo[0], space.base_vocab) - self.assertEqual(space.band_hi[-1], space.base_vocab + 11) - def test_model_resizes_to_target_vocab(self) -> None: model = self._model() rows = model.lm.get_input_embeddings().weight.shape[0] @@ -174,38 +160,17 @@ def test_model_resizes_to_target_vocab(self) -> None: # every SID atom has a row self.assertGreater(rows, self.prompt.sid_space.band_hi[-1]) - def test_forward_produces_a_finite_loss(self) -> None: + def test_loss_is_finite_and_backpropagates_into_the_backbone(self) -> None: model = self._model() batch = self._batch([0, 1, 2, 3, 0, 1], [1, 2, 3]) - out = model.predict(batch) - - self.assertIn("loss", out) - self.assertTrue(bool(torch.isfinite(out["loss"]))) - - def test_loss_backpropagates_into_the_backbone(self) -> None: - model = self._model() - batch = self._batch([0, 1, 2, 3, 0, 1], [1, 2, 3]) - model.predict(batch)["loss"].backward() + loss = model.predict(batch)["loss"] + self.assertTrue(bool(torch.isfinite(loss))) + loss.backward() grad = model.lm.get_input_embeddings().weight.grad self.assertIsNotNone(grad) self.assertTrue(bool((grad.abs().sum() > 0))) - def test_assembled_stream_matches_the_template(self) -> None: - batch = self._batch([0, 1, 2], [1, 2, 3]) - ids = batch.additional_infos["prompt_input_ids"] - space = self.prompt.sid_space - # "History :" + 3 history atoms + "." + "Predict :" + 3 answer atoms - self.assertEqual(ids.numel(), 2 + 3 + 1 + 2 + 3) - sid_rows = ids[ids >= space.base_vocab] - self.assertEqual(sid_rows.numel(), 6) - - def test_labels_supervise_only_the_answer(self) -> None: - batch = self._batch([0, 1, 2], [1, 2, 3]) - labels = batch.additional_infos["prompt_labels"] - supervised = labels[labels != -100] - self.assertEqual(supervised.numel(), 3) - def test_training_forward_survives_fx_tracing(self) -> None: # TrainPipelineSparseDist symbolically traces the model whenever a # sharded module exists, and the padded forward reads the collator's @@ -220,82 +185,7 @@ def __init__(self, inner): def forward(self, batch): return self.inner.predict(batch) - graph = torch.fx.symbolic_trace(_Wrapper(model)) - leaves = [ - str(n.target) for n in graph.graph.nodes if "_fx_wrapped" in str(n.target) - ] - self.assertTrue(leaves, "the padded forward must be opaque to FX") - - def test_raw_codes_are_rejected_before_the_model_sees_them(self) -> None: - parsed = { - # not offset: level 1 and 2 fall below their bands - "hist.values": torch.tensor([1, 2, 3]), - "hist.lengths": torch.tensor([3]), - "answer.values": torch.tensor(_offset([1, 2, 3])), - "answer.lengths": torch.tensor([3]), - } - with self.assertRaisesRegex(ValueError, "offset_codebook column"): - assemble_into(self.prompt, parsed) - - -class ProjectedSlotTest(unittest.TestCase): - """A slot whose value reaches the LM through a table and a projection.""" - - def setUp(self) -> None: - self.test_dir = make_test_dir() - self.backbone = _tiny_backbone(os.path.join(self.test_dir, "backbone")) - self.tok = _tokenizer(os.path.join(self.test_dir, "tok.json")) - self.features = _features(extra=(_PROF,)) - - cfg = PromptConfig( - tokenizer=self.tok, - prompt="History : {{hist}} . Predict {{prof}} :", - response="{{answer}}", - ) - cfg.sid_space.codebook.extend(_CODEBOOK) - self.prompt = compile_prompt(cfg, self.features, model_dir=self.test_dir) - - def test_compiler_derives_one_group_per_projected_slot(self) -> None: - groups = self.prompt.projection_plan.feature_groups - self.assertEqual([g.group_name for g in groups], ["prof"]) - self.assertEqual(list(groups[0].feature_names), ["prof"]) - # the INLINE slot produces none: its tokens are already in the stream - self.assertEqual( - [s.name for s in self.prompt.prompt_plan.projected_slots], ["prof"] - ) - - def test_sentinel_is_materialized_and_holes_recorded(self) -> None: - space = self.prompt.sid_space - self.assertIsNotNone(space.sentinel_token_id) - - parsed = { - "hist.values": torch.tensor(_offset([0, 1, 2])).reshape(-1, 1), - "hist.lengths": torch.tensor([3]), - "answer.values": torch.tensor(_offset([1, 2, 3])), - "answer.lengths": torch.tensor([3]), - "prof.values": torch.tensor([5, 9]), - "prof.lengths": torch.tensor([2]), - } - streams = assemble_into(self.prompt, parsed) - # two profile items -> two sentinels -> two holes - self.assertEqual(streams["prompt_hole_positions"].tolist(), [7, 8]) - ids = streams["prompt_input_ids"] - self.assertTrue(all(ids[p] == space.sentinel_token_id for p in [7, 8])) - - def test_projection_receives_gradient(self) -> None: - model_config = ModelConfig() - qwen = model_config.prompt_generative_qwen - qwen.hf_model_id = self.backbone - qwen.common.beam_widths.extend([2, 2, 2]) - qwen.common.num_return_sequences = 2 - model = _create_model( - model_config, self.features, ["answer"], prompt=self.prompt - ) - self.assertEqual(len(model.projections), 1) - - # the scatter is what puts the projection on the autograd path at all - proj = next(iter(model.projections.values())) - self.assertIsNone(proj.head.weight.grad) + torch.fx.symbolic_trace(_Wrapper(model)) if __name__ == "__main__": diff --git a/tzrec/tests/prompt_test_util.py b/tzrec/tests/prompt_test_util.py new file mode 100644 index 000000000..d44070abb --- /dev/null +++ b/tzrec/tests/prompt_test_util.py @@ -0,0 +1,38 @@ +# Copyright (c) 2026, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict + +import numpy as np + +from tzrec.prompt.assembler import PromptAssembler +from tzrec.prompt.plan import CompiledPrompt + + +def assemble_into( + prompt: CompiledPrompt, + parsed: Dict[str, "np.ndarray"], + ignore_index: int = -100, +) -> Dict[str, np.ndarray]: + """Assemble one parsed batch with a temporary assembler. + + Args: + prompt: the compiled prompt. + parsed: ``{feature}.values`` / ``{feature}.lengths`` as the data parser + emits them. + ignore_index: label value outside the supervised span. + + Returns: + The assembled streams keyed for ``additional_infos``. + """ + return PromptAssembler( + prompt.prompt_plan, prompt.sid_space, ignore_index + ).assemble_batch(parsed) diff --git a/tzrec/utils/hf_export_util_test.py b/tzrec/utils/hf_export_util_test.py index bd77a84f4..d41f40d85 100644 --- a/tzrec/utils/hf_export_util_test.py +++ b/tzrec/utils/hf_export_util_test.py @@ -81,15 +81,6 @@ def setUp(self) -> None: def tearDown(self) -> None: shutil.rmtree(self.test_dir, ignore_errors=True) - def test_unwrap_walks_dmp_and_train_wrapper(self) -> None: - inner = _GenRec(_tied_lm()) - self.assertIs(unwrap_to_hf(inner), inner) - self.assertIs(unwrap_to_hf(_TrainWrapper(inner)), inner) - self.assertIs(unwrap_to_hf(_DmpLike(_TrainWrapper(inner))), inner) - - def test_unwrap_returns_none_for_non_hf_model(self) -> None: - self.assertIsNone(unwrap_to_hf(_TrainWrapper(nn.Linear(4, 4)))) - def test_unwrap_terminates_on_a_wrapper_cycle(self) -> None: """A .model/.module cycle must return None, not spin. @@ -151,19 +142,6 @@ def test_dcp_to_hf_round_trip_drops_tied_head(self) -> None: for k, v in lm.state_dict().items(): self.assertTrue(torch.equal(back.state_dict()[k], v), k) - def test_dcp_to_hf_self_heals_a_stale_prefix(self) -> None: - lm = _tied_lm() - ckpt_dir = self._save_ckpt(_TrainWrapper(_GenRec(lm))) - meta_path = os.path.join(ckpt_dir, _HF_EXPORT_META_FILENAME) - with open(meta_path, "w") as f: - json.dump({"backbone_state_dict_prefix": "bogus.wrapper."}, f) - out_dir = os.path.join(self.test_dir, "hf_out_stale") - dcp_to_hf(ckpt_dir, out_dir) # falls back to suffix matching - st = load_file(os.path.join(out_dir, "model.safetensors")) - self.assertTrue( - torch.equal(st["model.embed_tokens.weight"], lm.model.embed_tokens.weight) - ) - def test_dcp_to_hf_refuses_a_mismatched_architecture(self) -> None: ckpt_dir = self._save_ckpt(_TrainWrapper(_GenRec(_tied_lm()))) # widen the recorded architecture so the checkpoint can no longer fill it From 15563a52aec8e47bd52fa02f51c3682d4131327e Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Thu, 13 Aug 2026 07:44:30 +0000 Subject: [PATCH 87/99] [bugfix] align prompt assembly with Qwen execution Projected holes were emitted sample-major while projection outputs are occurrence-major, and prediction still parsed the supervised response. Align assembly with embedding order, keep response out of inference batches, and validate the effective decode contracts. --- docs/source/models/generative.rst | 1 - docs/source/models/prompt_generative_qwen.md | 176 ------------------ tzrec/datasets/dataset.py | 35 +++- tzrec/main.py | 23 ++- tzrec/models/model.py | 19 +- tzrec/models/prompt_generative_model.py | 37 +--- tzrec/models/prompt_generative_model_test.py | 95 +++++----- tzrec/models/prompt_generative_qwen.py | 32 +++- tzrec/models/prompt_generative_qwen_test.py | 2 - tzrec/modules/dynamic_beam.py | 26 ++- tzrec/modules/dynamic_beam_test.py | 18 -- tzrec/modules/prompt_projection.py | 12 +- tzrec/modules/prompt_projection_test.py | 25 ++- tzrec/optim/lr_scheduler_test.py | 3 - tzrec/prompt/assembler.py | 102 +++++++--- tzrec/prompt/assembler_test.py | 117 ++++++++++-- tzrec/prompt/compile.py | 30 +-- tzrec/prompt/compile_test.py | 75 ++++---- tzrec/prompt/persist.py | 4 +- tzrec/prompt/persist_test.py | 30 ++- tzrec/prompt/plan.py | 15 +- tzrec/protos/models/prompt_model.proto | 7 +- tzrec/protos/optimizer.proto | 6 +- tzrec/protos/prompt.proto | 22 +-- .../prompt_generative_qwen_mock.config | 67 +++++++ tzrec/tests/prompt_integration_test.py | 84 +++------ tzrec/tests/prompt_test_util.py | 54 +++++- tzrec/utils/checkpoint_util.py | 8 +- tzrec/utils/hf_export_util.py | 4 +- tzrec/utils/hf_export_util_test.py | 15 +- tzrec/utils/test_util.py | 27 ++- 31 files changed, 609 insertions(+), 562 deletions(-) delete mode 100644 docs/source/models/prompt_generative_qwen.md create mode 100644 tzrec/tests/configs/prompt_generative_qwen_mock.config diff --git a/docs/source/models/generative.rst b/docs/source/models/generative.rst index e92016094..9e0aa39b2 100644 --- a/docs/source/models/generative.rst +++ b/docs/source/models/generative.rst @@ -5,6 +5,5 @@ :maxdepth: 2 dlrm_hstu - prompt_generative_qwen ultra_hstu hstu_match diff --git a/docs/source/models/prompt_generative_qwen.md b/docs/source/models/prompt_generative_qwen.md deleted file mode 100644 index 97ff74b5f..000000000 --- a/docs/source/models/prompt_generative_qwen.md +++ /dev/null @@ -1,176 +0,0 @@ -# Prompt 原生生成式推荐（PromptGenerativeQwen） - -以 Qwen 为骨干，把用户历史的语义 ID（SID）拼进 prompt，让模型直接生成下一个物品的 SID。 - -prompt 的模板、槽位、SID 空间与词表都由新的 `prompt_config` 描述，`model_config` 只保留属于 LM 的部分。 - -## 1. 数据准备 - -SID 必须以 **offset 形式**进入 tzrec，即 SID 生成工具 `resolve_sid_collisions` 输出的 `offset_codebook` 列： - -``` -第 l 层的取值 = level_offsets[l] + code code 属于 [0, codebook[l]) -``` - -以 `codebook: 4 4 4` 为例，`level_offsets` 为 `[0, 4, 8]`，因此一个 item 的三层取值分别落在 `[0,4)`、`[4,8)`、`[8,12)`。 - -```{warning} -只读 `offset_codebook`。`codebook` 与 `origin_codebook` 两列同样格式合法，但前者未加 offset、后者是冲突解析**之前**的 SID；误用不会报格式错，而是训练出静默错误的模型。assembler 的 band 校验能挡住未加 offset 的列，但挡不住手工对 `origin_codebook` 施加 offset 得到的流。 -``` - -一行历史是若干个 item 的三层 code 依次拼平，长度必须是层数的整数倍。 - -## 2. 配置 - -一个最小可运行的配置： - -``` -data_config { - batch_size: 4 - dataset_type: ParquetDataset - fg_mode: FG_NONE - label_fields: "answer" -} - -feature_configs { - sequence_raw_feature { feature_name: "hist" expression: "user:hist" } -} -feature_configs { - sequence_raw_feature { feature_name: "answer" expression: "item:answer" } -} - -prompt_config { - tokenizer: "path/to/tokenizer.json" - prompt: "用户历史行为为：{{hist}}。请预测下一个商品：" - response: "{{answer}}" - sid_space { codebook: 256 codebook: 256 codebook: 256 } - max_length: 4096 -} - -model_config { - prompt_generative_qwen { - hf_model_id: "Qwen/Qwen2.5-0.5B" - common { - beam_widths: 100 - beam_widths: 200 - beam_widths: 400 - num_return_sequences: 50 - } - } -} -``` - -### prompt_config - -| 字段 | 说明 | -| ------------------------- | ------------------------------------------------------------------------------------------------------- | -| `tokenizer` | **基础** tokenizer 的路径或 hub id。注意它与 `hf_model_id` 不同：后者只表示权重，且只在冷启动时读取一次 | -| `prompt` | 模板。`{{name}}` 之间的静态文本自动成为相邻槽位的前后缀，无需逐槽位配置 | -| `response` | 监督目标。定义 loss 覆盖的范围；推理时不生成该段 | -| `sid_space.codebook` | 每层的 SID 词表大小 | -| `sid_space.manifest_path` | 可选。指向 SID manifest，编译期与 `codebook` 逐元素比对，不一致直接报错 | -| `max_length` | 校验上限，**不是**截断开关：超长的行会报错，不会被截断 | - -### 槽位如何被推导 - -`{{name}}` 默认解析为同名特征。槽位的填充方式不需要配置，由成员特征推导： - -| 槽位成员 | 填充方式 | 说明 | -| -------------------------------------------------------- | --------- | --------------------------------------------------- | -| 单个序列特征且不声明 embedding（`sequence_raw_feature`） | INLINE | SID 直接进入 token 流，与答案共享 embedding | -| 其他情形（如 `sequence_id_feature`、标量特征、多成员） | PROJECTED | 走自己的 embedding 表，再经一次投影抵达 LM 隐层维度 | - -PROJECTED 槽位在 token 流中占位为 sentinel，真实取值在前向时写入对应位置。 - -### model_config - -| 字段 | 说明 | -| ---------------------- | ----------------------------------------------------------- | -| `hf_model_id` | 预训练权重的 hub id 或本地目录 | -| `beam_widths` | 每层一个宽度，长度必须等于 `codebook` 的层数 | -| `num_return_sequences` | 不得超过最后一层的宽度 | -| `param_dtype` | 主权重精度，默认 FP32。bf16 会让 Adam 的小更新在 ULP 下丢失 | - -## 3. 训练 - -```bash -torchrun --master_addr=localhost --master_port=32555 --nnodes=1 --nproc-per-node=2 --node_rank=0 \ - -m tzrec.train_eval --pipeline_config_path prompt_qwen.config -``` - -续跑加 `--continue_train`。 - -每个 `model.ckpt-N/` 除权重外还会写出 `prompt/` 目录： - -``` -model.ckpt-N/prompt/ - sid_space.json 解析后的 SID 空间：codebook、level_offsets、band、target_vocab - prompt_plan.json assembler 的遍历顺序与各项上界 - prompt_hashes.json vocab_hash 与 plan_hash - tokenizer/ 扩展后的 tokenizer（含 SID atom） -``` - -即 checkpoint 自带词表契约，服务端无需另行配置。 - -## 4. 预测 - -```bash -torchrun --master_addr=localhost --master_port=32555 --nnodes=1 --nproc-per-node=1 --node_rank=0 \ - -m tzrec.predict --pipeline_config_path experiments/run/pipeline.config \ - --predict_input_path 'data/*.parquet' --predict_output_path out -``` - -输出列 `generated_sids`，形状为 `(num_return_sequences, 层数)`，取值是**局部 0-based** code，可直接与 SID 映射表的 `codebook` 列对齐。 - -## 5. 导出 - -只支持导出为 HuggingFace 目录： - -``` -export_config { export_format: HF } -``` - -```bash -torchrun ... -m tzrec.export --pipeline_config_path experiments/run/pipeline.config \ - --export_dir exported -``` - -产出目录同时包含权重与 prompt 契约，可直接被 `AutoModelForCausalLM.from_pretrained` 加载： - -``` -exported/ - config.json generation_config.json model.safetensors - prompt/ sid_space.json prompt_plan.json prompt_hashes.json tokenizer/ -``` - -```{note} -本模型不支持 TorchScript 导出。它的输入是 dataloader 组装出的 token 流，而导出期的伪造 batch 无法提供。配置为默认的 TORCHSCRIPT 时会直接报错并提示改用 HF。 -``` - -## 6. 常见问题 - -**`SID values must already carry their level offset ... Read the offset_codebook column`** - -读错了列。改用 `offset_codebook`，见第 1 节。 - -**`prompt vocabulary does not match checkpoint`** - -`codebook`、`atom_token_format` 或 tokenizer 变了，与该 checkpoint 训练时的词表不一致。这是硬失败：解码 band 会指向这批权重从未学过的行，继续跑只会产出看似合理的错误结果。要么改回原配置，要么从头训练。 - -若只是模板或槽位变了（`plan_hash` 不同、`vocab_hash` 相同），只会告警，权重仍可用。 - -**`beam_widths has N entries but the codebook has M levels`** - -每层一个宽度，两者长度必须相等。 - -**`assembled row X is N tokens, over max_length`** - -超长的行不会被截断。请在特征上用 `sequence_length` 限制历史长度，而不是调大 `max_length`。 - -**`static_prefix_len is 0`（告警）** - -模板开头就是一个槽位，导致服务端前缀缓存无内容可共享。把静态指令文本放在最前、变长槽位放在最后即可。 - -**`a prompt-native model exports to a HuggingFace directory, not TorchScript`** - -见第 5 节，设置 `export_config.export_format: HF`。 diff --git a/tzrec/datasets/dataset.py b/tzrec/datasets/dataset.py index f85bba1d3..b9a2fa0ac 100644 --- a/tzrec/datasets/dataset.py +++ b/tzrec/datasets/dataset.py @@ -42,7 +42,7 @@ ) from tzrec.features.feature import BaseFeature from tzrec.prompt.assembler import PromptAssembler -from tzrec.prompt.plan import CompiledPrompt +from tzrec.prompt.plan import CompiledPrompt, SlotSeg from tzrec.protos import data_pb2 from tzrec.utils import config_util from tzrec.utils.load_class import get_register_class_meta @@ -100,6 +100,8 @@ class BaseDataset(IterableDataset, metaclass=_dataset_meta_cls): mode (Mode): train or eval or predict. debug_level (int): dataset debug level, when mode=predict and debug_level > 0, will dump fg encoded data to debug_str + prompt (CompiledPrompt, optional): compiled prompt assembly contract. + prompt_ignore_index (int): label value outside the supervised response. """ def __init__( @@ -111,11 +113,16 @@ def __init__( mode: Mode = Mode.EVAL, debug_level: int = 0, prompt: Optional[CompiledPrompt] = None, + prompt_ignore_index: int = -100, ) -> None: super(BaseDataset, self).__init__() - # built once per worker: the plan it walks is fixed for the run self._assembler = ( - PromptAssembler(prompt.prompt_plan, prompt.sid_space) + PromptAssembler( + prompt.prompt_plan, + prompt.sid_space, + ignore_index=prompt_ignore_index, + include_response=mode != Mode.PREDICT, + ) if prompt is not None else None ) @@ -131,8 +138,25 @@ def __init__( else None ) + parser_features = features + if prompt is not None and mode == Mode.PREDICT: + prompt_feature_names = { + feature_name + for segment in prompt.prompt_plan.segments + if isinstance(segment, SlotSeg) + for feature_name in segment.feature_names + } + response_feature_names = { + feature_name + for segment in prompt.prompt_plan.response_segments + if isinstance(segment, SlotSeg) + for feature_name in segment.feature_names + } + response_only = response_feature_names - prompt_feature_names + parser_features = [f for f in features if f.name not in response_only] + self._data_parser = DataParser( - features=features, + features=parser_features, labels=list(data_config.label_fields) if self._mode != Mode.PREDICT else None, @@ -781,6 +805,7 @@ def create_dataloader( debug_level: int = 0, checkpoint_state: Optional[Dict[str, Any]] = None, prompt: Optional[CompiledPrompt] = None, + prompt_ignore_index: int = -100, ) -> DataLoader: """Build dataloader. @@ -797,6 +822,7 @@ def create_dataloader( eager ``iter()`` forks workers so it reaches them. prompt (CompiledPrompt, optional): when set, each batch carries the assembled prompt streams in ``additional_infos``. + prompt_ignore_index (int): label value outside the supervised response. Return: dataloader (dataloader): a DataLoader. @@ -812,6 +838,7 @@ def create_dataloader( mode=mode, debug_level=debug_level, prompt=prompt, + prompt_ignore_index=prompt_ignore_index, ) if checkpoint_state: dataset.load_state_dict(dict(checkpoint_state)) diff --git a/tzrec/main.py b/tzrec/main.py index ecb737a2e..14b3b794d 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -127,11 +127,7 @@ def _create_features( def _compile_prompt( pipeline_config: EasyRecConfig, features: List[BaseFeature] ) -> Optional[CompiledPrompt]: - """Compile prompt_config when the pipeline declares one. - - Runs on every entry point, so the plan the data layer walks and the vocab - the model resizes to are produced by one code path. - """ + """Compile prompt_config for entry points that build prompt-aware objects.""" if not pipeline_config.HasField("prompt_config"): return None return compile_prompt( @@ -139,6 +135,14 @@ def _compile_prompt( ) +def _prompt_ignore_index(pipeline_config: EasyRecConfig) -> int: + """Return the label sentinel configured by a prompt-native model.""" + model_type = pipeline_config.model_config.WhichOneof("model") + if model_type != "prompt_generative_qwen": + return -100 + return int(getattr(pipeline_config.model_config, model_type).common.ignore_index) + + def _get_sampler_type(data_config: DataConfig) -> Optional[str]: try: sampler_type = ( @@ -718,6 +722,7 @@ def train_and_evaluate( # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) prompt = _compile_prompt(pipeline_config, features) + prompt_ignore_index = _prompt_ignore_index(pipeline_config) ckpt_manager = checkpoint_util.CheckpointManager( pipeline_config.model_dir, @@ -771,6 +776,7 @@ def train_and_evaluate( pipeline_config.train_input_path, mode=Mode.TRAIN, prompt=prompt, + prompt_ignore_index=prompt_ignore_index, checkpoint_state=dataloader_state, ) eval_dataloader = None @@ -783,6 +789,7 @@ def train_and_evaluate( pipeline_config.eval_input_path, mode=Mode.EVAL, prompt=prompt, + prompt_ignore_index=prompt_ignore_index, gl_cluster=gl_cluster, ) @@ -999,6 +1006,7 @@ def evaluate( eval_input_path or pipeline_config.eval_input_path, mode=Mode.EVAL, prompt=prompt, + prompt_ignore_index=_prompt_ignore_index(pipeline_config), ) sampler_type = _get_sampler_type(data_config) @@ -1150,8 +1158,7 @@ def export( from tzrec.utils.hf_export_util import dcp_to_hf dcp_to_hf(checkpoint_path, export_dir) - # this branch never builds a model, so save_assets cannot run; the - # checkpoint already carries the contract, so copy it forward + # Carry the prompt contract saved alongside the checkpoint. copy_prompt_assets(checkpoint_path, export_dir) return @@ -1167,7 +1174,6 @@ def export( # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) - prompt = _compile_prompt(pipeline_config, features) # Build model model = _create_model( @@ -1175,7 +1181,6 @@ def export( features, list(data_config.label_fields), sampler_type=None, - prompt=prompt, ) InferWrapper = ScriptWrapper # Flip to inference *before* wrapping so view-dependent state diff --git a/tzrec/models/model.py b/tzrec/models/model.py index 0086966d3..2a88e9c20 100644 --- a/tzrec/models/model.py +++ b/tzrec/models/model.py @@ -95,25 +95,22 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: raise NotImplementedError def save_assets(self, target_dir: str) -> None: - """Write any contract a checkpoint needs to describe itself. + """Write model-specific assets alongside a checkpoint. - Lifecycle hook symmetric with ``init_from_pretrained``, called for every - ``model.ckpt-N/`` and for the export directory. The default is a no-op; - models whose weights are meaningless without a companion artifact -- a - vocabulary, a plan -- override it. + ``CheckpointManager`` calls this hook after saving checkpoint weights. + The default is a no-op; models that require companion artifacts such as + a vocabulary or plan override it. Args: - target_dir: the checkpoint or export directory. + target_dir: the checkpoint directory. """ def init_from_pretrained(self) -> None: """Load pretrained weights at cold start (no checkpoint to restore). - Lifecycle hook the training pipeline calls only on a fresh run - (``ckpt_path is None``), before distributed wrapping. The default is a - no-op; models backed by an external pretrained source (e.g. an HF - backbone) override it. Resume/eval/export never reach it -- they restore - weights from the checkpoint. + The training pipeline calls this hook before distributed wrapping when + no checkpoint was selected. The default is a no-op; models backed by an + external pretrained source, such as an HF backbone, override it. """ def init_loss(self) -> None: diff --git a/tzrec/models/prompt_generative_model.py b/tzrec/models/prompt_generative_model.py index 2dec03f61..20d8f7e77 100644 --- a/tzrec/models/prompt_generative_model.py +++ b/tzrec/models/prompt_generative_model.py @@ -9,17 +9,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Backbone-agnostic half of a prompt-native generative model. +"""Shared causal-LM plumbing for prompt-native generative models. -Everything here is independent of how a family runs its transformer: building -the LM empty, resizing to the compiled vocabulary, wiring slot projections, -converting between the SID coordinate systems, and the checkpoint hooks. - -What a subclass owns is the forward: ``predict``, the teacher-forced loss and -the decode loop. Those differ irreducibly -- a decoder-only model reaches past -``lm(...)`` into body and head so it can score a suffix window, while an -encoder-decoder passes labels and gets a loss back -- so they are abstract here -rather than parameterized. +This layer builds an empty causal LM, resizes its vocabulary, wires slot +projections, converts SID coordinate systems and implements checkpoint hooks. +A family subclass owns its forward and decode path. """ from typing import Any, Dict, List, Optional @@ -87,8 +81,7 @@ def __init__( cfg = self._model_config self.lm = self._build_backbone(cfg.hf_model_id, cfg.common.param_dtype) - # only the shape matters here: every path overwrites these rows, from - # init_from_pretrained on a cold start or from the DCP restore otherwise + # Every run replaces this initialization from pretrained or DCP weights. self.lm.resize_token_embeddings( prompt.sid_space.target_vocab, mean_resizing=False ) @@ -104,14 +97,15 @@ def __init__( ) def _build_backbone(self, hf_model_id: str, param_dtype: int) -> nn.Module: - """Build the LM empty, so HF weights load only on cold start. + """Build the LM from config, so HF weights load only on cold start. Args: - hf_model_id: hub id or local directory naming the weights. + hf_model_id: hub id or local directory naming the architecture and + cold-start weights. param_dtype: master-weight dtype. Returns: - The uninitialized backbone. + A randomly initialized backbone with the requested parameter dtype. """ config = AutoConfig.from_pretrained(hf_model_id) model = AutoModelForCausalLM.from_config(config) @@ -146,7 +140,6 @@ def _build_projections(self) -> None: ) aligned.append(built[module_id]) self.projections = nn.ModuleDict(built) - # zipped with prompt_plan.projected_slots; shared modules appear by reference self._slot_projections = aligned def hf_backbone(self) -> nn.Module: @@ -175,6 +168,7 @@ def _prompt_embeds(self, batch: Batch) -> torch.Tensor: proj(grouped[seg.name + seg.output_key]).reshape(-1, hidden_size) for seg, proj in zip(prompt_plan.projected_slots, self._slot_projections) ] + # The assembler records holes in this projected-occurrence-major order. # out of place: embeds carries grad from the embedding lookup return embeds.index_copy( 0, batch.additional_infos[PROMPT_HOLE_POSITIONS], torch.cat(parts) @@ -196,17 +190,6 @@ def _tokens_to_local_codes( codes = tokens - space.base_vocab - self._level_offsets return codes.view(batch_size, -1, space.num_levels) - def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: - """Run the model over an assembled prompt. - - Args: - batch: carries the packed prompt in ``additional_infos``. - - Returns: - The loss when training, the decoded SIDs otherwise. - """ - raise NotImplementedError - def init_loss(self) -> None: """No-op: an LM computes its own CE inside ``predict``.""" return diff --git a/tzrec/models/prompt_generative_model_test.py b/tzrec/models/prompt_generative_model_test.py index a477f980e..6de99a22e 100644 --- a/tzrec/models/prompt_generative_model_test.py +++ b/tzrec/models/prompt_generative_model_test.py @@ -14,23 +14,24 @@ import numpy as np import torch -from google.protobuf import text_format -from tokenizers import Tokenizer, models, pre_tokenizers from torchrec import KeyedJaggedTensor from transformers import AutoModelForCausalLM from tzrec.datasets.utils import BASE_DATA_GROUP, Batch -from tzrec.features.feature import FgMode, create_features from tzrec.main import _create_model from tzrec.prompt.assembler import ( PROMPT_HOLE_POSITIONS, PROMPT_INPUT_IDS, ) from tzrec.prompt.compile import compile_prompt -from tzrec.protos import feature_pb2 from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.prompt_pb2 import PromptConfig -from tzrec.tests.prompt_test_util import assemble_into +from tzrec.tests.prompt_test_util import ( + assemble_into, + create_prompt_feature, + create_prompt_tokenizer, + offset_sid_codes, +) from tzrec.utils.state_dict_util import init_parameters from tzrec.utils.test_util import create_tiny_causal_lm, make_test_dir @@ -38,27 +39,6 @@ _WORDS = ["History", "Predict", ":", ".", "", "<|im_end|>"] -def _tokenizer(path: str) -> str: - tok = Tokenizer( - models.WordLevel(vocab={w: i for i, w in enumerate(_WORDS)}, unk_token="") - ) - tok.pre_tokenizer = pre_tokenizers.Whitespace() - tok.save(path) - return path - - -def _feature(text: str): - fc = feature_pb2.FeatureConfig() - text_format.Merge(text, fc) - return create_features([fc], fg_mode=FgMode.FG_NONE)[0] - - -def _offset(codes): - """Shift local codes into the flat space, as the SID tool's column does.""" - offsets = np.cumsum([0] + _CODEBOOK[:-1]) - return (np.asarray(codes).reshape(-1, len(_CODEBOOK)) + offsets).reshape(-1) - - _HIST = 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }' _ANSWER = 'sequence_raw_feature { feature_name: "answer" expression: "item:answer" }' @@ -71,14 +51,19 @@ def _projected(name: str, dim: int) -> str: class BasePromptGenerativeModelTest(unittest.TestCase): - """The backbone-agnostic half, reached through its one concrete subclass.""" + """Shared causal-LM behavior, reached through its concrete Qwen subclass.""" def setUp(self) -> None: self.test_dir = make_test_dir() self.backbone = os.path.join(self.test_dir, "backbone") create_tiny_causal_lm(64).save_pretrained(self.backbone) - self.tok = _tokenizer(os.path.join(self.test_dir, "tok.json")) - self.features = [_feature(_HIST), _feature(_ANSWER)] + self.tok = create_prompt_tokenizer( + os.path.join(self.test_dir, "tok.json"), _WORDS + ) + self.features = [ + create_prompt_feature(_HIST), + create_prompt_feature(_ANSWER), + ] self.prompt = self._compile(self.features) def _compile(self, features, template="History : {{hist}} . Predict :", **kwargs): @@ -86,12 +71,18 @@ def _compile(self, features, template="History : {{hist}} . Predict :", **kwargs cfg.sid_space.codebook.extend(_CODEBOOK) return compile_prompt(cfg, features, model_dir=self.test_dir) - def _model(self, features=None, prompt=-1): + def _model( + self, + features=None, + prompt=-1, + beam_widths=(2, 2, 2), + num_return_sequences=2, + ): model_config = ModelConfig() qwen = model_config.prompt_generative_qwen qwen.hf_model_id = self.backbone - qwen.common.beam_widths.extend([2, 2, 2]) - qwen.common.num_return_sequences = 2 + qwen.common.beam_widths.extend(beam_widths) + qwen.common.num_return_sequences = num_return_sequences return _create_model( model_config, self.features if features is None else features, @@ -138,10 +129,10 @@ def test_rejects_a_prompt_that_declares_no_sid_space(self) -> None: def test_shared_projection_name_requires_matching_widths(self) -> None: features = [ - _feature(_HIST), - _feature(_ANSWER), - _feature(_projected("pa", 8)), - _feature(_projected("pb", 16)), + create_prompt_feature(_HIST), + create_prompt_feature(_ANSWER), + create_prompt_feature(_projected("pa", 8)), + create_prompt_feature(_projected("pb", 16)), ] cfg = PromptConfig( tokenizer=self.tok, @@ -158,7 +149,11 @@ def test_shared_projection_name_requires_matching_widths(self) -> None: self._model(features=features, prompt=prompt) def test_projected_slot_overwrites_sentinels_and_backpropagates(self) -> None: - features = [_feature(_HIST), _feature(_ANSWER), _feature(_projected("prof", 8))] + features = [ + create_prompt_feature(_HIST), + create_prompt_feature(_ANSWER), + create_prompt_feature(_projected("prof", 8)), + ] prompt = self._compile( features, template="History : {{hist}} . Predict {{prof}} :", @@ -169,9 +164,11 @@ def test_projected_slot_overwrites_sentinels_and_backpropagates(self) -> None: init_parameters(model, device=torch.device("cpu")) batch = self._batch( { - "hist.values": torch.tensor(_offset([0, 1, 2])).reshape(-1, 1), + "hist.values": torch.tensor( + offset_sid_codes([0, 1, 2], _CODEBOOK) + ).reshape(-1, 1), "hist.lengths": torch.tensor([3]), - "answer.values": torch.tensor(_offset([1, 2, 3])), + "answer.values": torch.tensor(offset_sid_codes([1, 2, 3], _CODEBOOK)), "answer.lengths": torch.tensor([3]), "prof.values": torch.tensor([5, 9]), "prof.lengths": torch.tensor([2]), @@ -196,11 +193,12 @@ def test_projected_slot_overwrites_sentinels_and_backpropagates(self) -> None: proj = next(iter(model.projections.values())) self.assertIsNotNone(proj.head.weight.grad) - def test_loss_surfaces_the_ce_computed_in_predict(self) -> None: - model = self._model() - value = torch.tensor(1.25) - - self.assertEqual(model.loss({"loss": value}, Batch()), {"ce_loss": value}) + def test_beam_config_uses_final_capped_capacity(self) -> None: + with self.assertRaisesRegex(ValueError, "final capped beam width \\(4\\)"): + self._model( + beam_widths=(1, 1, 100), + num_return_sequences=5, + ) def test_metric_averages_the_loss_across_batches(self) -> None: model = self._model() @@ -212,15 +210,6 @@ def test_metric_averages_the_loss_across_batches(self) -> None: model._metric_modules["ce_loss"].compute().item(), 2.0, places=5 ) - def test_save_assets_co_locates_the_prompt_contract(self) -> None: - model = self._model() - target = os.path.join(self.test_dir, "ckpt") - os.makedirs(target, exist_ok=True) - model.save_assets(target) - - self.assertTrue(os.path.isdir(os.path.join(target, "prompt"))) - self.assertTrue(os.listdir(os.path.join(target, "prompt"))) - def test_init_from_pretrained_replaces_the_empty_weights(self) -> None: model = self._model() base_vocab = self.prompt.sid_space.base_vocab diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 2bbe9e5c4..474c1e82c 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -11,12 +11,13 @@ """Decoder-only forward and decode over a Qwen backbone. -Everything backbone-agnostic is in ``BasePromptGenerativeModel``. What is here -is what a decoder-only family does differently: it reaches past ``lm(...)`` into -``lm.model`` and ``lm.lm_head`` so logits are materialized for a suffix window -only, and it decodes by prefilling once and stepping a self-attention cache. +Shared causal-LM plumbing is in ``BasePromptGenerativeModel``. This subclass +reaches past ``lm(...)`` into ``lm.model`` and ``lm.lm_head`` so logits are +materialized for a suffix window only, and it decodes by prefilling once and +stepping a self-attention cache. -The layout is Qwen's, and Llama and Mistral share it exactly. +This implementation targets the Qwen HuggingFace module and cache interfaces; +other causal-LM families need their own compatibility verification. """ from typing import Any, Dict, List, Optional, Tuple @@ -26,7 +27,7 @@ from tzrec.datasets.utils import Batch from tzrec.features.feature import BaseFeature from tzrec.models.prompt_generative_model import BasePromptGenerativeModel -from tzrec.modules.dynamic_beam import dynamic_beam_search +from tzrec.modules.dynamic_beam import _capped_beam_widths, dynamic_beam_search from tzrec.prompt.assembler import ( PROMPT_CU_SEQLENS, PROMPT_LABELS, @@ -85,15 +86,26 @@ def _read_beam_config(self, common: PromptModelConfig) -> None: f"{len(self._beam_widths)} entries but the codebook has " f"{space.num_levels} levels; give one width per level." ) - if self._num_return_sequences > self._beam_widths[-1]: + if any(width < 1 for width in self._beam_widths): + raise ValueError( + f"{type(self).__name__}: beam_widths must be >= 1, got " + f"{self._beam_widths}." + ) + if self._num_return_sequences < 1: + raise ValueError( + f"{type(self).__name__}: num_return_sequences must be >= 1, got " + f"{self._num_return_sequences}." + ) + final_width = _capped_beam_widths(self._beam_widths, space.codebook)[-1] + if self._num_return_sequences > final_width: raise ValueError( f"{type(self).__name__}: num_return_sequences " - f"({self._num_return_sequences}) must not exceed the final beam width " - f"({self._beam_widths[-1]})." + f"({self._num_return_sequences}) must not exceed the final capped " + f"beam width ({final_width})." ) def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: - """Teacher-forced forward over the assembled stream. + """Run teacher-forced loss or inference decode over the assembled stream. Args: batch: carries the packed prompt in ``additional_infos``. diff --git a/tzrec/models/prompt_generative_qwen_test.py b/tzrec/models/prompt_generative_qwen_test.py index 69a4e0e4c..e4d2ab891 100644 --- a/tzrec/models/prompt_generative_qwen_test.py +++ b/tzrec/models/prompt_generative_qwen_test.py @@ -20,7 +20,6 @@ class UnpackTest(unittest.TestCase): """The one adapter where padding lives.""" def test_packs_rows_of_different_lengths(self) -> None: - # rows of 4 and 5 tokens, with a padded width larger than either row embeds = torch.arange(18, dtype=torch.float32).reshape(9, 2) cu = torch.tensor([0, 4, 9]) ignore = -100 @@ -31,7 +30,6 @@ def test_packs_rows_of_different_lengths(self) -> None: ) self.assertEqual(padded.shape, (2, 7, 2)) - # pads go on the left, so every row ends on a real token self.assertEqual( mask.tolist(), [[0, 0, 0, 1, 1, 1, 1], [0, 0, 1, 1, 1, 1, 1]], diff --git a/tzrec/modules/dynamic_beam.py b/tzrec/modules/dynamic_beam.py index 784608b87..1649c8d16 100644 --- a/tzrec/modules/dynamic_beam.py +++ b/tzrec/modules/dynamic_beam.py @@ -21,6 +21,18 @@ from transformers import PreTrainedModel +def _capped_beam_widths( + beam_widths: Sequence[int], band_sizes: Sequence[int] +) -> List[int]: + """Cap each requested width to the candidates reachable at that level.""" + capped_widths: List[int] = [] + previous_width = 1 + for requested, band_size in zip(beam_widths, band_sizes): + capped_widths.append(min(requested, previous_width * band_size)) + previous_width = capped_widths[-1] + return capped_widths + + @torch.no_grad() def dynamic_beam_search( model: PreTrainedModel, @@ -43,7 +55,7 @@ def dynamic_beam_search( Returns: The SID token tail ``(B * W, num_levels)``, score-ordered best-first. - The answer is fixed-length and EOS-free, so no beam bookkeeping. + The fixed-length, EOS-free answer needs no finished-sequence bookkeeping. """ device = prompt_embeds.device batch_size = prompt_embeds.shape[0] @@ -57,17 +69,15 @@ def dynamic_beam_search( raise ValueError( f"dynamic_beam_search: beam_widths must be >= 1, got {list(beam_widths)}." ) - capped_widths: List[int] = [] - prev_width = 1 - for requested, (band_lo, band_hi) in zip(beam_widths, bands): - capped_widths.append(min(requested, prev_width * (band_hi - band_lo + 1))) - prev_width = capped_widths[-1] + capped_widths = _capped_beam_widths( + beam_widths, [band_hi - band_lo + 1 for band_lo, band_hi in bands] + ) def _band_logp(logits: torch.Tensor, level: int) -> torch.Tensor: """Full-vocab log-probs, narrowed to ``level``'s band ``(rows, band)``. - Normalize then slice: same ranking as a full-vocab log_softmax, 21x less - memory at production vocab. + Normalize then slice: same ranking as a full-vocab log_softmax without + materializing a second full-vocabulary log-probability tensor. """ band_lo, band_hi = bands[level] log_z = torch.logsumexp(logits.float(), dim=-1, keepdim=True) diff --git a/tzrec/modules/dynamic_beam_test.py b/tzrec/modules/dynamic_beam_test.py index b6fcb04eb..1a6919f9b 100644 --- a/tzrec/modules/dynamic_beam_test.py +++ b/tzrec/modules/dynamic_beam_test.py @@ -14,7 +14,6 @@ from typing import Any, Dict, List, Tuple import torch -import torch.nn.functional as F from parameterized import parameterized from tzrec.modules.dynamic_beam import dynamic_beam_search @@ -109,23 +108,6 @@ def test_rejects_a_schedule_that_does_not_match_the_bands(self) -> None: with self.assertRaisesRegex(ValueError, "must be >= 1"): _decode(lm, ids, pairs, beam_widths=[2, 0, 4]) - @parameterized.expand( - [[0, 1], [1, 16]], - name_func=parameterized_name_func, - ) - def test_left_padding_matches_unpadded(self, seed, n_pad) -> None: - pairs = [(20, 21), (22, 24), (25, 28)] - lm = create_tiny_causal_lm(vocab_size=30, seed=seed) - short = torch.tensor([[5, 6, 7]]) - plain = _decode(lm, short, pairs) - padded = _decode( - lm, - F.pad(short, (n_pad, 0)), - pairs, - attention_mask=F.pad(torch.ones_like(short), (n_pad, 0)), - ) - self.assertTrue(torch.equal(plain, padded)) - def test_ragged_batch_rows_match_solo_runs(self) -> None: # every row of a ragged batch must decode exactly as if run alone. pairs = [(20, 21), (22, 24), (25, 28)] diff --git a/tzrec/modules/prompt_projection.py b/tzrec/modules/prompt_projection.py index 9ae976026..45be2b1eb 100644 --- a/tzrec/modules/prompt_projection.py +++ b/tzrec/modules/prompt_projection.py @@ -24,10 +24,9 @@ class PromptProjection(nn.Module): """An optional body followed by a bare Linear to the LM hidden size. - The final map never carries an activation: every ``Perceptron`` applies one, - so ending on an MLP would zero half the dimensions feeding the LM input - space. An empty MLP does not degrade to a linear either -- ``output_dim()`` - raises -- which is why the Linear is structural rather than configurable. + The final map has no activation, so it can span the full LM embedding space. + Omit ``mlp`` for a plain linear projection; an explicitly empty ``mlp {}`` + is rejected as an incomplete body configuration. Args: config: the slot's projection config; an empty one is a plain linear. @@ -45,6 +44,11 @@ def __init__( dim = in_dim self.body: Optional[MLP] = None if config.HasField("mlp"): + if not config.mlp.hidden_units: + raise ValueError( + "PromptProjection.mlp.hidden_units must not be empty; omit " + "mlp for a plain linear projection." + ) self.body = MLP(dim, **config_to_kwargs(config.mlp)) dim = self.body.output_dim() self.head = nn.Linear(dim, hidden_size, bias=config.bias) diff --git a/tzrec/modules/prompt_projection_test.py b/tzrec/modules/prompt_projection_test.py index 8bbfc87bc..4e7131857 100644 --- a/tzrec/modules/prompt_projection_test.py +++ b/tzrec/modules/prompt_projection_test.py @@ -20,9 +20,15 @@ class PromptProjectionTest(unittest.TestCase): def test_bodyless_config_is_a_plain_linear(self) -> None: - proj = PromptProjection(PromptProjectionConfig(), in_dim=12, hidden_size=8) + proj = PromptProjection( + PromptProjectionConfig(bias=False), in_dim=12, hidden_size=8 + ) self.assertIsNone(proj.body) self.assertIsInstance(proj.head, nn.Linear) + self.assertIsNone(proj.head.bias) + with torch.no_grad(): + proj.head.weight.fill_(-1.0) + self.assertTrue(bool((proj(torch.ones(1, 12)) < 0).all())) self.assertEqual(proj(torch.randn(5, 12)).shape, (5, 8)) def test_mlp_body_feeds_a_bare_head(self) -> None: @@ -36,21 +42,12 @@ def test_mlp_body_feeds_a_bare_head(self) -> None: self.assertEqual(proj.head.out_features, 8) self.assertEqual(proj(torch.randn(5, 12)).shape, (5, 8)) - def test_head_has_no_activation(self) -> None: + def test_rejects_an_explicitly_empty_mlp(self) -> None: config = PromptProjectionConfig() - config.mlp.hidden_units.extend([16]) - proj = PromptProjection(config, in_dim=4, hidden_size=6) - # a Perceptron would clamp negatives; a bare Linear must not - with torch.no_grad(): - proj.head.weight.fill_(-1.0) - proj.head.bias.fill_(0.0) - out = proj(torch.ones(1, 4)) - self.assertTrue(bool((out < 0).any())) + config.mlp.SetInParent() - def test_bias_is_configurable(self) -> None: - config = PromptProjectionConfig(bias=False) - proj = PromptProjection(config, in_dim=4, hidden_size=6) - self.assertIsNone(proj.head.bias) + with self.assertRaisesRegex(ValueError, "hidden_units must not be empty"): + PromptProjection(config, in_dim=12, hidden_size=8) if __name__ == "__main__": diff --git a/tzrec/optim/lr_scheduler_test.py b/tzrec/optim/lr_scheduler_test.py index 7c1814dbb..a99435ae5 100644 --- a/tzrec/optim/lr_scheduler_test.py +++ b/tzrec/optim/lr_scheduler_test.py @@ -101,13 +101,10 @@ def test_linear_decay_lr_with_warmup(self) -> None: opt, num_training_steps=6, warmup_size=2, warmup_learning_rate=0.002 ) self.assertFalse(lr.by_epoch) - # warmup step 0->1: scale=0.5, lr=0.002+(0.01-0.002)*0.5=0.006 lr.step() self.assertAlmostEqual(opt.param_groups[0]["lr"], 0.006) - # warmup step 1->2: scale=1.0, lr=0.01 lr.step() self.assertAlmostEqual(opt.param_groups[0]["lr"], 0.01) - # decay over remaining 4 steps: 0.0075, 0.005, 0.0025, 0.0 for lr_gt in [0.0075, 0.005, 0.0025, 0.0, 0.0]: lr.step() self.assertAlmostEqual(opt.param_groups[0]["lr"], lr_gt) diff --git a/tzrec/prompt/assembler.py b/tzrec/prompt/assembler.py index 1d9e4e5f6..ea3478fe1 100644 --- a/tzrec/prompt/assembler.py +++ b/tzrec/prompt/assembler.py @@ -28,6 +28,7 @@ SlotSeg, Static, ) +from tzrec.protos.model_pb2 import FeatureGroupType PROMPT_INPUT_IDS = "prompt_input_ids" PROMPT_CU_SEQLENS = "prompt_cu_seqlens" @@ -44,7 +45,8 @@ class AssembledPrompt: input_ids: every sample's tokens concatenated, ``(total_tokens,)``. cu_seqlens: sample boundaries into ``input_ids``, ``(batch_size + 1,)``. hole_positions: absolute indices the projected embeddings overwrite, - in ``PromptPlan.projected_slots`` order within each sample. + grouped by included projected occurrence in + ``PromptPlan.projected_slots`` order, then by sample. labels: ``ignore_index`` outside the response span. """ @@ -68,6 +70,7 @@ class PromptAssembler: prompt_plan: the compiled walk order. sid_space: resolved SID token space; required when a slot renders SIDs. ignore_index: label value outside the supervised span. + include_response: whether to read and emit the supervised response. """ def __init__( @@ -75,10 +78,14 @@ def __init__( prompt_plan: PromptPlan, sid_space: Optional[ResolvedSidSpace] = None, ignore_index: int = -100, + include_response: bool = True, ) -> None: self._prompt_plan = prompt_plan self._sid_space = sid_space self._ignore_index = ignore_index + self._response_segments = ( + prompt_plan.response_segments if include_response else () + ) if sid_space is not None: self._flat_lo = np.asarray(sid_space.level_offsets, dtype=np.int64) self._flat_hi = self._flat_lo + np.asarray( @@ -86,7 +93,7 @@ def __init__( ) inline = [ s - for s in prompt_plan.segments + prompt_plan.response_segments + for s in prompt_plan.segments + self._response_segments if isinstance(s, SlotSeg) and s.fill is FillMode.INLINE ] if inline and sid_space is None: @@ -125,9 +132,10 @@ def _emit_sample( inline_values: Dict[str, List[np.ndarray]], projected_lengths: Dict[str, np.ndarray], out: List[int], - holes: List[int], + holes_by_occurrence: List[List[int]], + projected_occurrence_index: int, sample_start: int, - ) -> None: + ) -> int: """Append one sample's tokens for one segment list, recording holes.""" for seg in segments: if isinstance(seg, Static): @@ -141,13 +149,15 @@ def _emit_sample( else: assert self._sid_space is not None width = int(projected_lengths[seg.name][sample_index]) - holes.extend( + holes_by_occurrence[projected_occurrence_index].extend( range( sample_start + len(out), sample_start + len(out) + width, ) ) + projected_occurrence_index += 1 out.extend([self._sid_space.sentinel_token_id] * width) + return projected_occurrence_index def assemble( self, @@ -174,27 +184,33 @@ def assemble( jagged_token_ids: List[int] = [] labels: List[int] = [] - holes: List[int] = [] + holes_by_occurrence = [ + [] + for seg in self._prompt_plan.segments + self._response_segments + if isinstance(seg, SlotSeg) and seg.fill is FillMode.PROJECTED + ] cu_seqlens = [0] for sample_index in range(batch_size): sample_token_ids: List[int] = [] - self._emit_sample( + projected_occurrence_index = self._emit_sample( self._prompt_plan.segments, sample_index, inline_values, projected_lengths, sample_token_ids, - holes, + holes_by_occurrence, + 0, len(jagged_token_ids), ) prompt_len = len(sample_token_ids) self._emit_sample( - self._prompt_plan.response_segments, + self._response_segments, sample_index, inline_values, projected_lengths, sample_token_ids, - holes, + holes_by_occurrence, + projected_occurrence_index, len(jagged_token_ids), ) # supervision covers the response span only; the prompt is context. @@ -215,6 +231,11 @@ def assemble( labels.extend(sample_labels) cu_seqlens.append(len(jagged_token_ids)) + holes = [ + position + for occurrence_positions in holes_by_occurrence + for position in occurrence_positions + ] return AssembledPrompt( input_ids=np.asarray(jagged_token_ids, dtype=np.int64), cu_seqlens=np.asarray(cu_seqlens, dtype=np.int64), @@ -222,37 +243,76 @@ def assemble( labels=np.asarray(labels, dtype=np.int64), ) - def assemble_batch(self, parsed: Dict[str, "np.ndarray"]) -> Dict[str, np.ndarray]: + def assemble_batch( + self, parsed_features: Dict[str, "np.ndarray"] + ) -> Dict[str, np.ndarray]: """Reshape one parsed batch, assemble it, and key it for the batch. Args: - parsed: ``{feature}.values`` / ``{feature}.lengths`` as the data - parser emits them. + parsed_features: ``{feature}.values`` / ``{feature}.lengths`` as + the data parser emits them. Returns: The five streams, keyed as ``additional_infos`` expects them. """ inline_values: Dict[str, List[np.ndarray]] = {} projected_lengths: Dict[str, np.ndarray] = {} - batch_size = 0 - for seg in self._prompt_plan.segments + self._prompt_plan.response_segments: + batch_size: Optional[int] = None + for seg in self._prompt_plan.segments + self._response_segments: if not isinstance(seg, SlotSeg): continue - source = seg.feature_names[0] - lengths = np.asarray(parsed[f"{source}.lengths"]) - batch_size = max(batch_size, int(lengths.size)) + sources = ( + seg.feature_names + if seg.fill is FillMode.PROJECTED + else seg.feature_names[:1] + ) + member_lengths: List[tuple[str, np.ndarray]] = [] + for source in sources: + lengths_key = f"{source}.lengths" + if seg.group_type == FeatureGroupType.JAGGED_SEQUENCE: + lengths = np.asarray(parsed_features[lengths_key]) + slot_batch_size = int(lengths.size) + member_lengths.append((source, lengths)) + elif lengths_key in parsed_features: + slot_batch_size = int(np.asarray(parsed_features[lengths_key]).size) + else: + values = np.asarray(parsed_features[f"{source}.values"]) + slot_batch_size = int(values.shape[0]) + if batch_size is None: + batch_size = slot_batch_size + elif slot_batch_size != batch_size: + raise ValueError( + f"prompt slot [{seg.name}] has {slot_batch_size} samples, " + f"expected {batch_size}." + ) if seg.fill is FillMode.INLINE: + source, lengths = member_lengths[0] # a dense sequence feature emits (total, value_dim); the stream # is one code per position, so value_dim is always 1 here - flat = np.asarray(parsed[f"{source}.values"]).reshape(-1) + flat = np.asarray(parsed_features[f"{source}.values"]).reshape(-1) bounds = np.concatenate(([0], np.cumsum(lengths))) inline_values[seg.name] = [ flat[bounds[i] : bounds[i + 1]] for i in range(lengths.size) ] - else: + elif seg.group_type == FeatureGroupType.JAGGED_SEQUENCE: + source, lengths = member_lengths[0] + for other_source, other_lengths in member_lengths[1:]: + if not np.array_equal(lengths, other_lengths): + raise ValueError( + f"prompt slot [{seg.name}] PROJECTED features " + f"[{source}] and [{other_source}] have different " + "per-sample lengths." + ) projected_lengths[seg.name] = lengths + else: + assert batch_size is not None + projected_lengths[seg.name] = np.ones(batch_size, dtype=np.int64) - out = self.assemble(inline_values, projected_lengths, batch_size=batch_size) + out = self.assemble( + inline_values, + projected_lengths, + batch_size=batch_size if batch_size is not None else 0, + ) return { PROMPT_INPUT_IDS: out.input_ids, PROMPT_CU_SEQLENS: out.cu_seqlens, diff --git a/tzrec/prompt/assembler_test.py b/tzrec/prompt/assembler_test.py index c85f1f310..93609b60d 100644 --- a/tzrec/prompt/assembler_test.py +++ b/tzrec/prompt/assembler_test.py @@ -49,18 +49,25 @@ def _sid_space(codebook=(4, 4, 4)) -> ResolvedSidSpace: ) -def _slot(name, fill, width_n=None) -> SlotSeg: +def _slot( + name, + fill, + width_n=None, + feature_names=None, + group_type=FeatureGroupType.JAGGED_SEQUENCE, +) -> SlotSeg: return SlotSeg( slot_id=0, name=name, - feature_names=(name,), - group_type=FeatureGroupType.JAGGED_SEQUENCE, - output_key=".sequence", + feature_names=tuple(feature_names) if feature_names is not None else (name,), + group_type=group_type, + output_key=".sequence" + if group_type == FeatureGroupType.JAGGED_SEQUENCE + else "", fill=fill, width=Width(WidthKind.BOUNDED, width_n) if width_n else Width(WidthKind.STATIC, 1), - droppable=False, ) @@ -78,7 +85,6 @@ def _plan(segments, response=(), max_length=0) -> PromptPlan: max_holes=0, logits_suffix_len=None, static_prefix_len=0, - length_buckets=(), projected_slots=projected, ) @@ -101,7 +107,7 @@ def test_projected_emits_sentinels_and_records_holes(self) -> None: asm = PromptAssembler(plan, _sid_space()) out = asm.assemble({}, {"prof": np.array([2, 3])}, batch_size=2) - # row 0: [7, S, S] row 1: [7, S, S, S] + # sample 0: [7, S, S] sample 1: [7, S, S, S] self.assertEqual( out.input_ids.tolist(), [7, _SENTINEL, _SENTINEL, 7, _SENTINEL, _SENTINEL, _SENTINEL], @@ -110,19 +116,40 @@ def test_projected_emits_sentinels_and_records_holes(self) -> None: # absolute indices into the flat buffer, which is what index_copy needs self.assertEqual(out.hole_positions.tolist(), [1, 2, 4, 5, 6]) + def test_holes_are_grouped_by_projected_occurrence_then_sample(self) -> None: + plan = _plan( + ( + _slot("a", FillMode.PROJECTED, 2), + Static((7,)), + _slot("b", FillMode.PROJECTED, 2), + _slot("a", FillMode.PROJECTED, 2), + ) + ) + asm = PromptAssembler(plan, _sid_space()) + out = asm.assemble( + {}, + {"a": np.array([1, 2]), "b": np.array([2, 1])}, + batch_size=2, + ) + + self.assertEqual(out.cu_seqlens.tolist(), [0, 5, 11]) + self.assertEqual(out.hole_positions.tolist(), [0, 5, 6, 2, 3, 8, 4, 9, 10]) + def test_labels_cover_the_response_span_only(self) -> None: plan = _plan( (Static((7, 8)),), response=(Static((9,)), _slot("answer", FillMode.INLINE)), ) - asm = PromptAssembler(plan, _sid_space()) + asm = PromptAssembler(plan, _sid_space(), ignore_index=-7) out = asm.assemble({"answer": [np.array([0, 4, 8])]}) self.assertEqual(out.input_ids.tolist(), [7, 8, 9, _BASE, _BASE + 4, _BASE + 8]) - # the prompt is context; supervision starts at the response - self.assertEqual( - out.labels.tolist(), [-100, -100, 9, _BASE, _BASE + 4, _BASE + 8] - ) + self.assertEqual(out.labels.tolist(), [-7, -7, 9, _BASE, _BASE + 4, _BASE + 8]) + + prompt_only = PromptAssembler( + plan, _sid_space(), ignore_index=-7, include_response=False + ).assemble({}, batch_size=1) + self.assertEqual(prompt_only.input_ids.tolist(), [7, 8]) def test_rejects_raw_codes_that_carry_no_offset(self) -> None: plan = _plan((_slot("hist", FillMode.INLINE),)) @@ -176,6 +203,72 @@ def test_column_shaped_values_are_flattened(self) -> None: self.assertEqual(out["prompt_cu_seqlens"].tolist(), [0, 3, 6]) self.assertEqual(out["prompt_input_ids"].tolist()[0], _BASE + 1) + def test_rejects_inconsistent_slot_batch_sizes(self) -> None: + plan = _plan( + ( + _slot("hist", FillMode.INLINE), + _slot("answer", FillMode.INLINE), + ) + ) + asm = PromptAssembler(plan, _sid_space()) + parsed = { + "hist.values": np.array([1, 6, 11]), + "hist.lengths": np.array([3]), + "answer.values": np.array([0, 4, 8, 1, 6, 11]), + "answer.lengths": np.array([3, 3]), + } + with self.assertRaisesRegex( + ValueError, r"prompt slot \[answer\] has 2 samples, expected 1" + ): + asm.assemble_batch(parsed) + + def test_rejects_mismatched_projected_member_lengths(self) -> None: + plan = _plan( + ( + _slot( + "profile", + FillMode.PROJECTED, + 4, + feature_names=("age", "country"), + ), + ) + ) + asm = PromptAssembler(plan, _sid_space()) + parsed = { + "age.lengths": np.array([2, 1]), + "country.lengths": np.array([2, 2]), + } + + with self.assertRaisesRegex( + ValueError, + r"PROJECTED features \[age\] and \[country\] have different", + ): + asm.assemble_batch(parsed) + + def test_deep_projected_members_emit_one_hole_per_sample(self) -> None: + plan = _plan( + ( + _slot( + "profile", + FillMode.PROJECTED, + feature_names=("dense", "sparse"), + group_type=FeatureGroupType.DEEP, + ), + ) + ) + asm = PromptAssembler(plan, _sid_space()) + out = asm.assemble_batch( + { + "dense.values": np.array([[1.0, 2.0], [3.0, 4.0]]), + "sparse.values": np.array([5, 6, 7]), + "sparse.lengths": np.array([2, 1]), + } + ) + + self.assertEqual(out["prompt_input_ids"].tolist(), [_SENTINEL, _SENTINEL]) + self.assertEqual(out["prompt_cu_seqlens"].tolist(), [0, 1, 2]) + self.assertEqual(out["prompt_hole_positions"].tolist(), [0, 1]) + if __name__ == "__main__": unittest.main() diff --git a/tzrec/prompt/compile.py b/tzrec/prompt/compile.py index ecb2e626c..204a13f62 100644 --- a/tzrec/prompt/compile.py +++ b/tzrec/prompt/compile.py @@ -127,8 +127,8 @@ def _atom_tokens(sid_space: SidSpace) -> List[str]: return [fmt.replace("{i}", str(i)) for i in range(sum(sid_space.codebook))] -def _read_manifest_codebook(path: str) -> Optional[List[int]]: - """Read ``codebook`` from a SID manifest, or None when there is no file.""" +def _read_manifest_codebook(path: str) -> List[int]: + """Read ``codebook`` from a SID manifest.""" if not os.path.exists(path): raise ValueError(f"sid_space.manifest_path [{path}] does not exist.") with open(path, "r") as f: @@ -263,6 +263,12 @@ def compile_prompt( types = {n: _group_type(n, members[n]) for n in slots} fills = {n: _derive_fill(members[n]) for n in slots} + for name in resp_names: + if fills[name] is FillMode.PROJECTED: + raise ValueError( + f"response slot [{name}] is PROJECTED; response slots must be " + "INLINE because the LM generates them as vocabulary tokens." + ) has_projection = any(f is FillMode.PROJECTED for f in fills.values()) for name, slot in slots.items(): @@ -302,7 +308,6 @@ def compile_prompt( output_key=".sequence" if seq else "", fill=fills[name], width=_slot_width(members[name], types[name], levels), - droppable=bool(slot.drop_if_empty), ) body = _weave(body_runs, body_names, segs, tok) @@ -323,7 +328,6 @@ def compile_prompt( max_holes=_max_holes(projected), logits_suffix_len=_suffix_keep(response), static_prefix_len=_static_prefix_len(body), - length_buckets=tuple(int(b) for b in cfg.length_buckets), projected_slots=projected, ) _validate(cfg, plan, sid_space) @@ -455,20 +459,20 @@ def _validate( "prompt has an unbounded slot and max_length is 0; graph-captured " "serving cannot size its buckets." ) - if any(isinstance(s, SlotSeg) for s in plan.segments): - first_slot = next( - i for i, s in enumerate(plan.segments) if isinstance(s, SlotSeg) - ) - later_static = any( - isinstance(s, SlotSeg) and s.width.kind is WidthKind.STATIC - for s in plan.segments[first_slot + 1 :] - ) - if later_static: + variable_slot_seen = False + for seg in plan.segments: + if not isinstance(seg, SlotSeg): + continue + if seg.width.kind is WidthKind.STATIC: + if not variable_slot_seen: + continue logger.warning( "a variable-width prompt slot precedes a fixed-width one; " f"static_prefix_len is {plan.static_prefix_len}, which bounds " "what a serving prefix cache may reuse." ) + break + variable_slot_seen = True if plan.response_segments and plan.logits_suffix_len is None: raise ValueError( "the response has an unbounded slot, so the supervised logits " diff --git a/tzrec/prompt/compile_test.py b/tzrec/prompt/compile_test.py index 96997242a..3bcb86db6 100644 --- a/tzrec/prompt/compile_test.py +++ b/tzrec/prompt/compile_test.py @@ -14,33 +14,22 @@ import unittest from google.protobuf import text_format -from tokenizers import Tokenizer, models, pre_tokenizers +from tokenizers import Tokenizer from tzrec.features.feature import FgMode, create_features from tzrec.prompt.compile import compile_prompt from tzrec.prompt.plan import FillMode, SlotSeg, Static, WidthKind from tzrec.protos import feature_pb2 from tzrec.protos.prompt_pb2 import PromptConfig +from tzrec.tests.prompt_test_util import ( + create_prompt_feature, + create_prompt_tokenizer, +) from tzrec.utils.test_util import make_test_dir _WORDS = ["History", "Profile", "Predict", ":", ".", "", "<|im_end|>"] -def _tokenizer(path: str) -> str: - """Write a minimal word-level tokenizer, so no download is needed.""" - vocab = {w: i for i, w in enumerate(_WORDS)} - tok = Tokenizer(models.WordLevel(vocab=vocab, unk_token="")) - tok.pre_tokenizer = pre_tokenizers.Whitespace() - tok.save(path) - return path - - -def _feature(text: str): - fc = feature_pb2.FeatureConfig() - text_format.Merge(text, fc) - return create_features([fc], fg_mode=FgMode.FG_NONE)[0] - - _HIST = 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }' _PROF = ( 'sequence_id_feature { feature_name: "prof" expression: "user:prof" ' @@ -52,7 +41,9 @@ def _feature(text: str): class CompilePromptTest(unittest.TestCase): def setUp(self) -> None: self.test_dir = make_test_dir() - self.tok_path = _tokenizer(os.path.join(self.test_dir, "tok.json")) + self.tok_path = create_prompt_tokenizer( + os.path.join(self.test_dir, "tok.json"), _WORDS + ) def _config(self, **kwargs) -> PromptConfig: cfg = PromptConfig(tokenizer=self.tok_path, **kwargs) @@ -64,12 +55,12 @@ def _compile(self, cfg, features): def test_sid_space_resolves_offsets_and_bands(self) -> None: cfg = self._config(prompt="History : {{hist}}") cfg.sid_space.codebook.extend([4, 4, 4]) - compiled = self._compile(cfg, [_feature(_HIST)]) + compiled = self._compile(cfg, [create_prompt_feature(_HIST)]) space = compiled.sid_space base = space.base_vocab self.assertEqual(space.num_levels, 3) - self.assertEqual(space.sid_vocab_size, 12) + self.assertEqual(sum(space.codebook), 12) self.assertEqual(space.level_offsets, (0, 4, 8)) self.assertEqual(space.band_lo, (base, base + 4, base + 8)) self.assertEqual(space.band_hi, (base + 3, base + 7, base + 11)) @@ -80,7 +71,9 @@ def test_sid_space_resolves_offsets_and_bands(self) -> None: def test_inline_needs_no_group_projected_gets_one(self) -> None: cfg = self._config(prompt="History : {{hist}} . Profile : {{prof}}") cfg.sid_space.codebook.extend([4, 4, 4]) - compiled = self._compile(cfg, [_feature(_HIST), _feature(_PROF)]) + compiled = self._compile( + cfg, [create_prompt_feature(_HIST), create_prompt_feature(_PROF)] + ) by_name = { s.name: s for s in compiled.prompt_plan.segments if isinstance(s, SlotSeg) @@ -100,7 +93,7 @@ def test_inline_needs_no_group_projected_gets_one(self) -> None: def test_static_runs_are_woven_between_slots(self) -> None: cfg = self._config(prompt="History : {{hist}} . Predict :") cfg.sid_space.codebook.extend([4]) - compiled = self._compile(cfg, [_feature(_HIST)]) + compiled = self._compile(cfg, [create_prompt_feature(_HIST)]) kinds = [ "static" if isinstance(s, Static) else s.name for s in compiled.prompt_plan.segments @@ -112,7 +105,7 @@ def test_static_runs_are_woven_between_slots(self) -> None: def test_scalar_slot_is_one_deep_position(self) -> None: cfg = self._config(prompt="Profile : {{age}}") cfg.sid_space.codebook.extend([4]) - compiled = self._compile(cfg, [_feature(_AGE)]) + compiled = self._compile(cfg, [create_prompt_feature(_AGE)]) seg = next(s for s in compiled.prompt_plan.segments if isinstance(s, SlotSeg)) self.assertIs(seg.fill, FillMode.PROJECTED) self.assertEqual(seg.output_key, "") @@ -127,7 +120,7 @@ def test_manifest_mismatch_is_fatal(self) -> None: cfg.sid_space.codebook.extend([4, 4, 4]) cfg.sid_space.manifest_path = manifest with self.assertRaisesRegex(ValueError, "does not match the manifest"): - self._compile(cfg, [_feature(_HIST)]) + self._compile(cfg, [create_prompt_feature(_HIST)]) def test_manifest_match_compiles(self) -> None: manifest = os.path.join(self.test_dir, "manifest.json") @@ -136,7 +129,10 @@ def test_manifest_match_compiles(self) -> None: cfg = self._config(prompt="History : {{hist}}") cfg.sid_space.codebook.extend([4, 4, 4]) cfg.sid_space.manifest_path = manifest - self.assertEqual(self._compile(cfg, [_feature(_HIST)]).sid_space.num_levels, 3) + self.assertEqual( + self._compile(cfg, [create_prompt_feature(_HIST)]).sid_space.num_levels, + 3, + ) def test_rejects_a_mixed_kind_slot(self) -> None: cfg = self._config(prompt="X : {{both}}") @@ -144,7 +140,9 @@ def test_rejects_a_mixed_kind_slot(self) -> None: slot = cfg.slots.add(name="both") slot.feature_names.extend(["hist", "age"]) with self.assertRaisesRegex(ValueError, "mixes sequence and scalar"): - self._compile(cfg, [_feature(_HIST), _feature(_AGE)]) + self._compile( + cfg, [create_prompt_feature(_HIST), create_prompt_feature(_AGE)] + ) def test_rejects_unknown_feature_and_unreferenced_slot(self) -> None: cfg = self._config(prompt="X : {{hist}}") @@ -152,13 +150,13 @@ def test_rejects_unknown_feature_and_unreferenced_slot(self) -> None: slot = cfg.slots.add(name="hist") slot.feature_names.append("nope") with self.assertRaisesRegex(ValueError, "not in\n?\\s*feature_configs"): - self._compile(cfg, [_feature(_HIST)]) + self._compile(cfg, [create_prompt_feature(_HIST)]) cfg2 = self._config(prompt="X : {{hist}}") cfg2.sid_space.codebook.extend([4]) cfg2.slots.add(name="ghost").feature_names.append("hist") with self.assertRaisesRegex(ValueError, "never referenced"): - self._compile(cfg2, [_feature(_HIST)]) + self._compile(cfg2, [create_prompt_feature(_HIST)]) def test_rejects_a_projection_on_an_inline_slot(self) -> None: cfg = self._config(prompt="X : {{hist}}") @@ -167,19 +165,19 @@ def test_rejects_a_projection_on_an_inline_slot(self) -> None: slot.feature_names.append("hist") slot.projection.bias = True with self.assertRaisesRegex(ValueError, "is INLINE"): - self._compile(cfg, [_feature(_HIST)]) + self._compile(cfg, [create_prompt_feature(_HIST)]) def test_atoms_absent_from_the_base_tokenizer(self) -> None: cfg = self._config(prompt="X : {{hist}}") cfg.sid_space.codebook.extend([4]) cfg.sid_space.atom_token_format = "History" with self.assertRaisesRegex(ValueError, "already in the base tokenizer"): - self._compile(cfg, [_feature(_HIST)]) + self._compile(cfg, [create_prompt_feature(_HIST)]) def test_extended_tokenizer_is_written(self) -> None: cfg = self._config(prompt="History : {{hist}}") cfg.sid_space.codebook.extend([4, 4]) - compiled = self._compile(cfg, [_feature(_HIST)]) + compiled = self._compile(cfg, [create_prompt_feature(_HIST)]) written = os.path.join(compiled.tokenizer_dir, "tokenizer.json") self.assertTrue(os.path.exists(written)) # the atoms round-trip, which is what serving reloads @@ -190,10 +188,10 @@ def test_extended_tokenizer_is_written(self) -> None: def test_answer_width_comes_from_the_codebook(self) -> None: cfg = self._config(prompt="History : {{hist}}", response="{{answer}}") cfg.sid_space.codebook.extend([4, 4, 4]) - answer = _feature( + answer = create_prompt_feature( 'sequence_raw_feature { feature_name: "answer" expression: "item:answer" }' ) - compiled = self._compile(cfg, [_feature(_HIST), answer]) + compiled = self._compile(cfg, [create_prompt_feature(_HIST), answer]) seg = next( s for s in compiled.prompt_plan.response_segments if isinstance(s, SlotSeg) @@ -205,16 +203,25 @@ def test_answer_width_comes_from_the_codebook(self) -> None: # first supervised label self.assertEqual(compiled.prompt_plan.logits_suffix_len, 4) + def test_response_slot_must_be_inline(self) -> None: + cfg = self._config(prompt="History : {{hist}}", response="{{prof}}") + cfg.sid_space.codebook.extend([4, 4, 4]) + + with self.assertRaisesRegex(ValueError, r"response slot \[prof\] is PROJECTED"): + self._compile( + cfg, [create_prompt_feature(_HIST), create_prompt_feature(_PROF)] + ) + def test_unbounded_response_is_rejected(self) -> None: # with no sid_space the response has no codebook-derived width, so the # supervised window is unbounded and the logits would cover every # position cfg = self._config(prompt="History : {{hist}}", response="{{answer}}") - answer = _feature( + answer = create_prompt_feature( 'sequence_raw_feature { feature_name: "answer" expression: "item:answer" }' ) with self.assertRaisesRegex(ValueError, "window cannot be bounded"): - self._compile(cfg, [_feature(_HIST), answer]) + self._compile(cfg, [create_prompt_feature(_HIST), answer]) def test_a_grouped_feature_inherits_the_group_cap(self) -> None: # a SequenceFeature member never sets its own sequence_length; the cap diff --git a/tzrec/prompt/persist.py b/tzrec/prompt/persist.py index f7b52726b..6a63423fc 100644 --- a/tzrec/prompt/persist.py +++ b/tzrec/prompt/persist.py @@ -13,8 +13,8 @@ A checkpoint that cannot describe its own vocabulary is a checkpoint serving has to be told about out of band, which is where offline/online skew comes -from. ``ProjectionPlan`` is deliberately absent: it is model-only and rebuilt from -config at every ``__init__``. +from. ``ProjectionPlan`` is deliberately absent: it is model-only and rebuilt by +``compile_prompt`` from config before model construction. """ import dataclasses diff --git a/tzrec/prompt/persist_test.py b/tzrec/prompt/persist_test.py index 762430220..1d3182978 100644 --- a/tzrec/prompt/persist_test.py +++ b/tzrec/prompt/persist_test.py @@ -15,10 +15,6 @@ import unittest from unittest import mock -from google.protobuf import text_format -from tokenizers import Tokenizer, models, pre_tokenizers - -from tzrec.features.feature import FgMode, create_features from tzrec.prompt.compile import compile_prompt from tzrec.prompt.persist import ( PROMPT_DIR, @@ -27,8 +23,11 @@ read_prompt_hashes, save_prompt_assets, ) -from tzrec.protos import feature_pb2 from tzrec.protos.prompt_pb2 import PromptConfig +from tzrec.tests.prompt_test_util import ( + create_prompt_feature, + create_prompt_tokenizer, +) from tzrec.utils.test_util import make_test_dir _WORDS = ["History", "Predict", ":", "", "<|im_end|>"] @@ -37,21 +36,14 @@ class PromptPersistTest(unittest.TestCase): def setUp(self) -> None: self.test_dir = make_test_dir() - tok_path = os.path.join(self.test_dir, "tok.json") - tok = Tokenizer( - models.WordLevel( - vocab={w: i for i, w in enumerate(_WORDS)}, unk_token="" - ) + self.tok_path = create_prompt_tokenizer( + os.path.join(self.test_dir, "tok.json"), _WORDS ) - tok.pre_tokenizer = pre_tokenizers.Whitespace() - tok.save(tok_path) - self.tok_path = tok_path - - fc = feature_pb2.FeatureConfig() - text_format.Merge( - 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }', fc - ) - self.features = create_features([fc], fg_mode=FgMode.FG_NONE) + self.features = [ + create_prompt_feature( + 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }' + ) + ] def _compile(self, codebook=(4, 4, 4), prompt="History : {{hist}}"): cfg = PromptConfig(tokenizer=self.tok_path, prompt=prompt) diff --git a/tzrec/prompt/plan.py b/tzrec/prompt/plan.py index 392120357..58b620643 100644 --- a/tzrec/prompt/plan.py +++ b/tzrec/prompt/plan.py @@ -59,7 +59,7 @@ def __post_init__(self) -> None: raise ValueError("UNBOUNDED width cannot carry a count.") elif self.num_positions is None or self.num_positions < 0: raise ValueError( - f"{self.kind.name} width needs a count >= 0, got {self.n}." + f"{self.kind.name} width needs a count >= 0, got {self.num_positions}." ) @@ -97,11 +97,6 @@ class ResolvedSidSpace: eos_token_id: int pad_token_id: int - @property - def sid_vocab_size(self) -> int: - """Atoms appended to the backbone vocabulary.""" - return sum(self.codebook) - @dataclass(frozen=True) class Static: @@ -119,14 +114,13 @@ class SlotSeg: """One ``{{name}}`` position in the assembled stream. Args: - slot_id: index into ``PromptPlan.projected_slots`` ordering. + slot_id: stable identifier assigned to the distinct prompt slot. name: the placeholder name; also the derived group name. feature_names: member feature names. group_type: DEEP or JAGGED_SEQUENCE. output_key: "" for DEEP, ".sequence" otherwise. fill: INLINE writes token ids, PROJECTED writes sentinels and a hole. width: position count of this slot. - droppable: whether an empty value removes the slot and its folded text. """ slot_id: int @@ -136,7 +130,6 @@ class SlotSeg: output_key: str fill: FillMode width: Width - droppable: bool Segment = Union[Static, SlotSeg] @@ -154,8 +147,7 @@ class PromptPlan: max_holes: per-row projected-position ceiling, not a runtime shape. logits_suffix_len: upper bound on the supervised logits window. static_prefix_len: leading positions that are request-invariant. - length_buckets: sampler and graph-capture buckets. - projected_slots: fixes the order hole positions are written in. + projected_slots: PROJECTED occurrences in emission order. """ segments: Tuple[Segment, ...] @@ -165,7 +157,6 @@ class PromptPlan: max_holes: int logits_suffix_len: Optional[int] static_prefix_len: int - length_buckets: Tuple[int, ...] projected_slots: Tuple[SlotSeg, ...] diff --git a/tzrec/protos/models/prompt_model.proto b/tzrec/protos/models/prompt_model.proto index c945a973e..48cfba059 100644 --- a/tzrec/protos/models/prompt_model.proto +++ b/tzrec/protos/models/prompt_model.proto @@ -20,7 +20,7 @@ message PromptModelConfig { // REQUIRED. Beam width per SID level, one entry per level; [100, 200, 400] // is the escalating beam. Each entry is capped to what its band supplies. repeated uint32 beam_widths = 2; - // Must not exceed the final beam width. + // Must not exceed the final width after per-level candidate capping. required uint32 num_return_sequences = 3; optional string generated_sids_key = 4 [default = "generated_sids"]; @@ -34,7 +34,8 @@ message PromptModelConfig { message PromptGenerativeQwen { optional PromptModelConfig common = 1; - // HF hub id or local path. Names the WEIGHTS only, and is read solely by - // init_from_pretrained at cold start; the vocabulary is prompt_config's. + // HF hub id or local path. Names the architecture used for every model build + // and the weights loaded by init_from_pretrained at cold start; the vocabulary + // is prompt_config's. optional string hf_model_id = 2 [default = "Qwen/Qwen2.5-0.5B"]; } diff --git a/tzrec/protos/optimizer.proto b/tzrec/protos/optimizer.proto index b2ce86b14..330eb4d30 100644 --- a/tzrec/protos/optimizer.proto +++ b/tzrec/protos/optimizer.proto @@ -245,15 +245,11 @@ message ManualStepLR { message LinearDecayLR { // length of the whole run, measured from step 0 and INCLUDING warmup_size // -- HF Trainer's `num_training_steps`, not the post-warmup horizon that - // decay_size/T_max above measure. Required (must be > 0). + // decay_size/T_max above measure. Must be > 0 when this scheduler is selected. optional uint32 num_training_steps = 1; - // minimum learning rate reached at num_training_steps optional float min_learning_rate = 2 [default = 0.0]; - // warmup start learning rate optional float warmup_learning_rate = 3 [default = 0.0]; - // warmup steps or epochs optional uint32 warmup_size = 4 [default = 0]; - // schedule by epoch or by step. optional bool by_epoch = 5 [default = false]; } diff --git a/tzrec/protos/prompt.proto b/tzrec/protos/prompt.proto index c35b1d268..4528c5716 100644 --- a/tzrec/protos/prompt.proto +++ b/tzrec/protos/prompt.proto @@ -6,11 +6,12 @@ import "tzrec/protos/module.proto"; // Rendering of an LM prompt. Peer of data_config and model_config: extraction // stays in feature_configs, this owns order, literal text and composition. message PromptConfig { - // BASE tokenizer, path or hub id. Distinct from hf_model_id, which names - // the WEIGHTS and is read only at cold start. + reserved 2, 9; + reserved "asset_dir", "length_buckets"; + + // BASE tokenizer JSON file. Distinct from hf_model_id, which supplies the + // model architecture and cold-start weights. required string tokenizer = 1; - // Rewritten by export to the content-addressed asset directory. - optional string asset_dir = 2; // Static text between {{name}} placeholders is the prefix and suffix of // the surrounding slots; there are no per-slot text fields. @@ -30,21 +31,20 @@ message PromptConfig { optional uint32 max_length = 7 [default = 0]; // Reserves a position filled by a projected slot. Materialized only when - // at least one slot has a projection. + // at least one slot is PROJECTED. optional string sentinel_token = 8 [default = "<|pg_hole|>"]; - - // Length buckets for the training sampler and graph capture. - repeated uint32 length_buckets = 9; } message PromptSlot { - // {{name}} in the template; also the name of the derived feature_group. + reserved 3; + reserved "drop_if_empty"; + + // {{name}} in the template; also the derived feature_group name when the + // slot is PROJECTED. required string name = 1; // Features rendered at this position. Several are concatenated per // position; the compiler derives the FeatureGroupConfig from this list. repeated string feature_names = 2; - // Omit the slot and the static text folded into it when the value is empty. - optional bool drop_if_empty = 3 [default = false]; // Reconciles this slot's width with the LM hidden size. Illegal on an // INLINE slot. Carries no dimensions: the model resolves them. optional PromptProjection projection = 4; diff --git a/tzrec/tests/configs/prompt_generative_qwen_mock.config b/tzrec/tests/configs/prompt_generative_qwen_mock.config new file mode 100644 index 000000000..0651228aa --- /dev/null +++ b/tzrec/tests/configs/prompt_generative_qwen_mock.config @@ -0,0 +1,67 @@ +train_input_path: "" +eval_input_path: "" +model_dir: "experiments/prompt_generative_qwen_mock" +train_config { + sparse_optimizer { + adagrad_optimizer { + lr: 0.0 + } + constant_learning_rate { + } + } + dense_optimizer { + adam_optimizer { + lr: 0.0001 + } + constant_learning_rate { + } + } + num_epochs: 1 + save_checkpoints_epochs: 1 +} +eval_config { +} +export_config { + export_format: HF +} +data_config { + batch_size: 4 + dataset_type: ParquetDataset + fg_mode: FG_NONE + label_fields: "answer" + num_workers: 1 +} +feature_configs { + sequence_raw_feature { + feature_name: "hist" + expression: "user:hist" + } +} +feature_configs { + sequence_raw_feature { + feature_name: "answer" + expression: "item:answer" + } +} +prompt_config { + tokenizer: "data/test/tokenizer.json" + prompt: "History : {{hist}} . Predict :" + response: "{{answer}}" + sid_space { + codebook: 4 + codebook: 4 + codebook: 4 + } + max_length: 64 +} +model_config { + prompt_generative_qwen { + hf_model_id: "Qwen/Qwen2.5-0.5B" + common { + beam_widths: 2 + beam_widths: 2 + beam_widths: 2 + num_return_sequences: 2 + } + } +} diff --git a/tzrec/tests/prompt_integration_test.py b/tzrec/tests/prompt_integration_test.py index 34e43693f..614fdab75 100644 --- a/tzrec/tests/prompt_integration_test.py +++ b/tzrec/tests/prompt_integration_test.py @@ -15,11 +15,8 @@ import numpy as np import torch import torch.fx -from google.protobuf import text_format -from tokenizers import Tokenizer, models, pre_tokenizers from tzrec.datasets.utils import Batch -from tzrec.features.feature import FgMode, create_features from tzrec.main import _create_model from tzrec.models.prompt_generative_qwen import _unpack from tzrec.prompt.assembler import ( @@ -28,58 +25,38 @@ PROMPT_MAX_SEQLEN, ) from tzrec.prompt.compile import compile_prompt -from tzrec.protos import feature_pb2 from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.prompt_pb2 import PromptConfig -from tzrec.tests.prompt_test_util import assemble_into +from tzrec.tests.prompt_test_util import ( + assemble_into, + create_prompt_feature, + create_prompt_tokenizer, + offset_sid_codes, +) from tzrec.utils.test_util import create_tiny_causal_lm, make_test_dir _CODEBOOK = [4, 4, 4] _WORDS = ["History", "Predict", ":", ".", "", "<|im_end|>"] -def _tiny_backbone(path: str) -> str: - """The shared tiny Qwen, saved locally so no download is needed.""" - create_tiny_causal_lm(64).save_pretrained(path) - return path - - -def _tokenizer(path: str) -> str: - tok = Tokenizer( - models.WordLevel(vocab={w: i for i, w in enumerate(_WORDS)}, unk_token="") - ) - tok.pre_tokenizer = pre_tokenizers.Whitespace() - tok.save(path) - return path - - -def _features(): - text = ( - 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }', - 'sequence_raw_feature { feature_name: "answer" expression: "item:answer" }', - ) - out = [] - for one in text: - fc = feature_pb2.FeatureConfig() - text_format.Merge(one, fc) - out.append(create_features([fc], fg_mode=FgMode.FG_NONE)[0]) - return out - - -def _offset(codes): - """Shift local codes into the flat space, as the SID tool's column does.""" - offsets = np.cumsum([0] + _CODEBOOK[:-1]) - return (np.asarray(codes).reshape(-1, len(_CODEBOOK)) + offsets).reshape(-1) - - class PromptStackIntegrationTest(unittest.TestCase): """compile -> assemble -> model, on the real code path.""" def setUp(self) -> None: self.test_dir = make_test_dir() - self.backbone = _tiny_backbone(os.path.join(self.test_dir, "backbone")) - self.tok = _tokenizer(os.path.join(self.test_dir, "tok.json")) - self.features = _features() + self.backbone = os.path.join(self.test_dir, "backbone") + create_tiny_causal_lm(64).save_pretrained(self.backbone) + self.tok = create_prompt_tokenizer( + os.path.join(self.test_dir, "tok.json"), _WORDS + ) + self.features = [ + create_prompt_feature(text) + for text in ( + 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }', + 'sequence_raw_feature { feature_name: "answer" ' + 'expression: "item:answer" }', + ) + ] cfg = PromptConfig( tokenizer=self.tok, @@ -100,26 +77,19 @@ def _model(self): ) def _batch(self, hist, answer): - parsed = { - "hist.values": torch.tensor(_offset(hist)), - "hist.lengths": torch.tensor([len(hist)]), - "answer.values": torch.tensor(_offset(answer)), - "answer.lengths": torch.tensor([len(answer)]), - } - streams = assemble_into(self.prompt, parsed) - batch = Batch() - batch.additional_infos.update( - {k: torch.from_numpy(np.asarray(v)) for k, v in streams.items()} - ) - return batch + return self._batch_rows([(hist, answer)]) def _batch_rows(self, rows): hist = [h for h, _ in rows] answer = [a for _, a in rows] parsed = { - "hist.values": torch.tensor(_offset([c for h in hist for c in h])), + "hist.values": torch.tensor( + offset_sid_codes([c for h in hist for c in h], _CODEBOOK) + ), "hist.lengths": torch.tensor([len(h) for h in hist]), - "answer.values": torch.tensor(_offset([c for a in answer for c in a])), + "answer.values": torch.tensor( + offset_sid_codes([c for a in answer for c in a], _CODEBOOK) + ), "answer.lengths": torch.tensor([len(a) for a in answer]), } streams = assemble_into(self.prompt, parsed) @@ -132,7 +102,6 @@ def _batch_rows(self, rows): def test_every_row_is_supervised_whatever_its_length(self) -> None: # a short row must not lose its answer to padding: the loss keeps a # fixed-width suffix, so both rows have to contribute equally - # one history item against two, so the assembled rows differ in width batch = self._batch_rows( [([0, 1, 2], [1, 2, 3]), ([0, 1, 2, 3, 0, 1], [2, 3, 0])] ) @@ -157,7 +126,6 @@ def test_model_resizes_to_target_vocab(self) -> None: model = self._model() rows = model.lm.get_input_embeddings().weight.shape[0] self.assertEqual(rows, self.prompt.sid_space.target_vocab) - # every SID atom has a row self.assertGreater(rows, self.prompt.sid_space.band_hi[-1]) def test_loss_is_finite_and_backpropagates_into_the_backbone(self) -> None: diff --git a/tzrec/tests/prompt_test_util.py b/tzrec/tests/prompt_test_util.py index d44070abb..84f1304fc 100644 --- a/tzrec/tests/prompt_test_util.py +++ b/tzrec/tests/prompt_test_util.py @@ -9,12 +9,64 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict +from typing import Any, Dict, Sequence import numpy as np +from google.protobuf import text_format +from tokenizers import Tokenizer, models, pre_tokenizers +from tzrec.features.feature import BaseFeature, FgMode, create_features from tzrec.prompt.assembler import PromptAssembler from tzrec.prompt.plan import CompiledPrompt +from tzrec.protos import feature_pb2 + + +def create_prompt_tokenizer(path: str, words: Sequence[str]) -> str: + """Write a word-level tokenizer used by prompt tests. + + Args: + path: destination JSON path. + words: vocabulary entries in token-id order. + + Returns: + The destination path. + """ + tokenizer = Tokenizer( + models.WordLevel( + vocab={word: i for i, word in enumerate(words)}, unk_token="" + ) + ) + tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() + tokenizer.save(path) + return path + + +def create_prompt_feature(text: str) -> BaseFeature: + """Create one prompt test feature from text-format protobuf. + + Args: + text: text-format ``FeatureConfig``. + + Returns: + The created feature. + """ + config = feature_pb2.FeatureConfig() + text_format.Merge(text, config) + return create_features([config], fg_mode=FgMode.FG_NONE)[0] + + +def offset_sid_codes(codes: Sequence[Any], codebook: Sequence[int]) -> np.ndarray: + """Shift local SID codes into the flattened per-level space. + + Args: + codes: local codes grouped by SID item. + codebook: vocabulary size for each SID level. + + Returns: + Flat offset codes in item-major order. + """ + offsets = np.cumsum([0, *codebook[:-1]]) + return (np.asarray(codes).reshape(-1, len(codebook)) + offsets).reshape(-1) def assemble_into( diff --git a/tzrec/utils/checkpoint_util.py b/tzrec/utils/checkpoint_util.py index 8866ec8f6..b71da865b 100644 --- a/tzrec/utils/checkpoint_util.py +++ b/tzrec/utils/checkpoint_util.py @@ -420,11 +420,9 @@ def save( ) -> str: """Save a checkpoint at the given step, then request an async prune. - For HF-backed models, co-locates the HF config + tokenizer (no weights) - in this checkpoint dir so each ``model.ckpt-N/`` is self-contained and - convertible to HF. Deliberately not gated on ``export_format``: that is - an export-time knob, and gating it would make a run trained with the - default format permanently unexportable to HF. + For HF-backed models, writes the config, optional tokenizer, and state + dict metadata needed by HF conversion. This is not gated on + ``export_format`` because that setting is selected at export time. """ ckpt_dir = os.path.join(self._model_dir, f"model.ckpt-{step}") save_model(ckpt_dir, model, optimizer, dense_ema) diff --git a/tzrec/utils/hf_export_util.py b/tzrec/utils/hf_export_util.py index 597de2db4..ad24a1fde 100644 --- a/tzrec/utils/hf_export_util.py +++ b/tzrec/utils/hf_export_util.py @@ -42,7 +42,7 @@ def write_hf_assets(wrapped_model: nn.Module, save_dir: str) -> None: - """Co-locate the HF config + tokenizer (NO weights) in a checkpoint dir. + """Write HF config, optional tokenizer, and export metadata beside a checkpoint. The backbone's FQN prefix goes into ``hf_export_meta.json`` so ``dcp_to_hf`` can strip it without hard-coding a wrapper convention. Rank 0 only. @@ -76,7 +76,7 @@ def write_hf_assets(wrapped_model: nn.Module, save_dir: str) -> None: def dcp_to_hf(ckpt_dir: str, out_dir: str) -> None: - """Convert a self-contained checkpoint dir to a ``from_pretrained`` HF dir. + """Convert a checkpoint with co-located HF assets to a ``from_pretrained`` dir. Keys that do not map 1:1 onto the co-located ``config.json`` raise rather than write a partial model. diff --git a/tzrec/utils/hf_export_util_test.py b/tzrec/utils/hf_export_util_test.py index d41f40d85..421d58033 100644 --- a/tzrec/utils/hf_export_util_test.py +++ b/tzrec/utils/hf_export_util_test.py @@ -44,7 +44,7 @@ def save_pretrained(self, save_dir): class _GenRec(nn.Module): - """Stand-in for a prompt-native model: an HF backbone plus unrelated params.""" + """Stand-in for an HF-backed model exposing the optional tokenizer protocol.""" def __init__(self, lm): super().__init__() @@ -93,11 +93,11 @@ def test_unwrap_terminates_on_a_wrapper_cycle(self) -> None: object.__setattr__(a, "model", b) object.__setattr__(b, "model", a) out = [] - t = threading.Thread(target=lambda: out.append(unwrap_to_hf(a))) + t = threading.Thread(target=lambda: out.append(unwrap_to(a, "hf_backbone"))) t.daemon = True t.start() t.join(timeout=5) - self.assertFalse(t.is_alive(), "_unwrap_hf_model did not terminate") + self.assertFalse(t.is_alive(), "unwrap_to did not terminate") self.assertEqual(out, [None]) def test_write_hf_assets_noop_for_non_hf_model(self) -> None: @@ -105,8 +105,8 @@ def test_write_hf_assets_noop_for_non_hf_model(self) -> None: write_hf_assets(_TrainWrapper(nn.Linear(4, 4)), save_dir) self.assertFalse(os.path.exists(save_dir)) - def _save_ckpt(self, wrapped, name="model.ckpt-1"): - ckpt_dir = os.path.join(self.test_dir, name) + def _save_ckpt(self, wrapped): + ckpt_dir = os.path.join(self.test_dir, "model.ckpt-1") save_model(ckpt_dir, wrapped) write_hf_assets(wrapped, ckpt_dir) return ckpt_dir @@ -161,10 +161,5 @@ def test_dcp_to_hf_missing_dcp_dir(self) -> None: dcp_to_hf(empty, os.path.join(self.test_dir, "hf_out_missing")) -def unwrap_to_hf(model): - """The walk write_hf_assets performs, under test.""" - return unwrap_to(model, "hf_backbone") - - if __name__ == "__main__": unittest.main() diff --git a/tzrec/utils/test_util.py b/tzrec/utils/test_util.py index d6b627f69..440f555d1 100644 --- a/tzrec/utils/test_util.py +++ b/tzrec/utils/test_util.py @@ -124,7 +124,6 @@ def create_tiny_causal_lm( vocab_size: int, seed: int = 0, tie_word_embeddings: bool = False, - max_position_embeddings: int = 64, ) -> nn.Module: """A 2-layer Qwen2 causal LM cheap enough to build inside a unit test. @@ -135,25 +134,25 @@ def create_tiny_causal_lm( vocab_size (int): rows in the embedding table. seed (int): torch seed the random init draws from. tie_word_embeddings (bool): tie ``lm_head`` to the input embedding. - max_position_embeddings (int): longest sequence the backbone accepts. Returns: an eval-mode ``Qwen2ForCausalLM``. """ from transformers import Qwen2Config, Qwen2ForCausalLM - torch.manual_seed(seed) - config = Qwen2Config( - vocab_size=vocab_size, - hidden_size=32, - intermediate_size=64, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - max_position_embeddings=max_position_embeddings, - tie_word_embeddings=tie_word_embeddings, - ) - return Qwen2ForCausalLM(config).eval() + with torch.random.fork_rng(devices=[]): + torch.manual_seed(seed) + config = Qwen2Config( + vocab_size=vocab_size, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=64, + tie_word_embeddings=tie_word_embeddings, + ) + return Qwen2ForCausalLM(config).eval() # pyre-ignore [2] From b156f8036224dd19ea85d4576d46ce25b5c2495a Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Thu, 13 Aug 2026 08:58:09 +0000 Subject: [PATCH 88/99] [refactor] clarify prompt compilation flow --- tzrec/prompt/compile.py | 163 +++++++++++++++++++++++----------------- 1 file changed, 94 insertions(+), 69 deletions(-) diff --git a/tzrec/prompt/compile.py b/tzrec/prompt/compile.py index 204a13f62..1b7680de6 100644 --- a/tzrec/prompt/compile.py +++ b/tzrec/prompt/compile.py @@ -67,10 +67,12 @@ def _ceil_to(value: int, multiple: int) -> int: return -(-value // multiple) * multiple -def _resolve_slot(name: str, declared: Dict[str, PromptSlot]) -> PromptSlot: +def _resolve_slot( + name: str, declared_slots_by_name: Dict[str, PromptSlot] +) -> PromptSlot: """Return the declared slot, or an implicit single-feature slot.""" - if name in declared: - return declared[name] + if name in declared_slots_by_name: + return declared_slots_by_name[name] implicit = PromptSlot(name=name) implicit.feature_names.append(name) return implicit @@ -100,25 +102,27 @@ def _slot_width( return Width(WidthKind.BOUNDED, max(caps)) -def _derive_fill(members: Sequence[BaseFeature]) -> FillMode: - """INLINE only for a lone sequence member that declares no embedding.""" - if len(members) == 1 and members[0].is_sequence and not members[0].has_embedding: - return FillMode.INLINE - return FillMode.PROJECTED - - -def _group_type( +def _derive_slot_layout( name: str, members: Sequence[BaseFeature] -) -> "FeatureGroupType.ValueType": - """JAGGED_SEQUENCE for sequence members, DEEP for scalars; never mixed.""" - kinds = {f.is_sequence for f in members} - if len(kinds) != 1: +) -> Tuple["FeatureGroupType.ValueType", FillMode]: + """Derive how a prompt slot is grouped and emitted.""" + sequence_flags = {member.is_sequence for member in members} + if len(sequence_flags) != 1: raise ValueError( f"prompt slot [{name}] mixes sequence and scalar features " - f"{[f.name for f in members]}; a slot must be all one kind, or its " - f"group would carry a '.query' output the prompt cannot place." + f"{[member.name for member in members]}; a slot must be all one kind, " + "or its group would carry a '.query' output the prompt cannot place." ) - return FeatureGroupType.JAGGED_SEQUENCE if kinds.pop() else FeatureGroupType.DEEP + is_sequence = sequence_flags.pop() + group_type = ( + FeatureGroupType.JAGGED_SEQUENCE if is_sequence else FeatureGroupType.DEEP + ) + fill_mode = ( + FillMode.INLINE + if is_sequence and len(members) == 1 and not members[0].has_embedding + else FillMode.PROJECTED + ) + return group_type, fill_mode def _atom_tokens(sid_space: SidSpace) -> List[str]: @@ -237,14 +241,17 @@ def compile_prompt( Returns: The compiled prompt. """ - by_name = {f.name: f for f in features} - declared = {s.name: s for s in cfg.slots} + features_by_name = {feature.name: feature for feature in features} + declared_slots_by_name = {slot.name: slot for slot in cfg.slots} body_runs, body_names = _split_template(cfg.prompt) resp_runs, resp_names = _split_template(cfg.response or "") - slots = {n: _resolve_slot(n, declared) for n in body_names + resp_names} + resolved_slots_by_name = { + name: _resolve_slot(name, declared_slots_by_name) + for name in body_names + resp_names + } - unreferenced = set(declared) - set(slots) + unreferenced = set(declared_slots_by_name) - set(resolved_slots_by_name) if unreferenced: raise ValueError( f"declared prompt slots {sorted(unreferenced)} are never referenced " @@ -252,27 +259,37 @@ def compile_prompt( ) members: Dict[str, List[BaseFeature]] = {} - for name, slot in slots.items(): - missing = [f for f in slot.feature_names if f not in by_name] - if missing: + for name, slot in resolved_slots_by_name.items(): + missing_feature_names = [ + feature_name + for feature_name in slot.feature_names + if feature_name not in features_by_name + ] + if missing_feature_names: raise ValueError( - f"prompt slot [{name}] names features {missing} that are not in " - f"feature_configs." + f"prompt slot [{name}] names features {missing_feature_names} that " + "are not in feature_configs." ) - members[name] = [by_name[f] for f in slot.feature_names] - - types = {n: _group_type(n, members[n]) for n in slots} - fills = {n: _derive_fill(members[n]) for n in slots} - for name in resp_names: - if fills[name] is FillMode.PROJECTED: + members[name] = [ + features_by_name[feature_name] for feature_name in slot.feature_names + ] + + response_slot_names = set(resp_names) + group_types_by_slot_name: Dict[str, "FeatureGroupType.ValueType"] = {} + fill_modes_by_slot_name: Dict[str, FillMode] = {} + for name, slot_members in members.items(): + group_type, fill_mode = _derive_slot_layout(name, slot_members) + if name in response_slot_names and fill_mode is FillMode.PROJECTED: raise ValueError( f"response slot [{name}] is PROJECTED; response slots must be " "INLINE because the LM generates them as vocabulary tokens." ) - has_projection = any(f is FillMode.PROJECTED for f in fills.values()) - - for name, slot in slots.items(): - if fills[name] is FillMode.INLINE and slot.HasField("projection"): + group_types_by_slot_name[name] = group_type + fill_modes_by_slot_name[name] = fill_mode + for name, slot in resolved_slots_by_name.items(): + if fill_modes_by_slot_name[name] is FillMode.INLINE and slot.HasField( + "projection" + ): raise ValueError( f"prompt slot [{name}] is INLINE -- one sequence feature with no " f"embedding -- so it has no group to project; drop its projection." @@ -280,6 +297,10 @@ def compile_prompt( tok = Tokenizer.from_file(cfg.tokenizer) base_vocab = tok.get_vocab_size(with_added_tokens=True) + has_projection = any( + fill_mode is FillMode.PROJECTED + for fill_mode in fill_modes_by_slot_name.values() + ) sid_space = _build_sid_space(cfg, tok, base_vocab, has_projection) tokenizer_dir = "" @@ -288,37 +309,39 @@ def compile_prompt( os.makedirs(tokenizer_dir, exist_ok=True) tok.save(os.path.join(tokenizer_dir, "tokenizer.json")) - slot_ids = {n: i for i, n in enumerate(slots)} - answer_names = set(resp_names) + slot_ids = {name: i for i, name in enumerate(resolved_slots_by_name)} segs: Dict[str, SlotSeg] = {} - for name, slot in slots.items(): - seq = types[name] == FeatureGroupType.JAGGED_SEQUENCE + for name, slot in resolved_slots_by_name.items(): + group_type = group_types_by_slot_name[name] + fill_mode = fill_modes_by_slot_name[name] levels = ( sid_space.num_levels if sid_space is not None - and name in answer_names - and fills[name] is FillMode.INLINE + and name in response_slot_names + and fill_mode is FillMode.INLINE else None ) segs[name] = SlotSeg( slot_id=slot_ids[name], name=name, feature_names=tuple(slot.feature_names), - group_type=types[name], - output_key=".sequence" if seq else "", - fill=fills[name], - width=_slot_width(members[name], types[name], levels), + group_type=group_type, + output_key=( + ".sequence" if group_type == FeatureGroupType.JAGGED_SEQUENCE else "" + ), + fill=fill_mode, + width=_slot_width(members[name], group_type, levels), ) - body = _weave(body_runs, body_names, segs, tok) - response = _weave(resp_runs, resp_names, segs, tok) + body = _build_template_segments(body_runs, body_names, segs, tok) + response = _build_template_segments(resp_runs, resp_names, segs, tok) projected = tuple( s for s in body + response if isinstance(s, SlotSeg) and s.fill is FillMode.PROJECTED ) - projection_plan = _build_module_plan(projected, slots) + projection_plan = _build_projection_plan(projected, resolved_slots_by_name) plan = PromptPlan( segments=body, @@ -344,31 +367,33 @@ def compile_prompt( ) -def _weave( - runs: Sequence[str], - names: Sequence[str], - segs: Dict[str, SlotSeg], - tok: Tokenizer, +def _build_template_segments( + static_text_parts: Sequence[str], + slot_names: Sequence[str], + slot_segments_by_name: Dict[str, SlotSeg], + tokenizer: Tokenizer, ) -> Tuple[Segment, ...]: - """Interleave tokenized static runs with their slots, dropping empty runs.""" - out: List[Segment] = [] - for i, run in enumerate(runs): - if run: - ids = tuple(tok.encode(run, add_special_tokens=False).ids) - out.append(Static(token_ids=ids)) - if i < len(names): - out.append(segs[names[i]]) - return tuple(out) - - -def _build_module_plan( - projected: Sequence[SlotSeg], slots: Dict[str, PromptSlot] + """Interleave tokenized static text parts with their slots.""" + segments: List[Segment] = [] + for index, static_text in enumerate(static_text_parts): + if static_text: + token_ids = tuple( + tokenizer.encode(static_text, add_special_tokens=False).ids + ) + segments.append(Static(token_ids=token_ids)) + if index < len(slot_names): + segments.append(slot_segments_by_name[slot_names[index]]) + return tuple(segments) + + +def _build_projection_plan( + projected: Sequence[SlotSeg], slots_by_name: Dict[str, PromptSlot] ) -> ProjectionPlan: """One module per distinct ``projection_name``, else one per slot.""" projections: Dict[str, PromptProjection] = {} slot_to_module: Dict[int, str] = {} for seg in projected: - slot = slots[seg.name] + slot = slots_by_name[seg.name] module_id = slot.projection_name or seg.name projection = ( slot.projection if slot.HasField("projection") else PromptProjection() From a734fd2411303f713e7b38e528f75d0bd0245a68 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Thu, 13 Aug 2026 12:02:50 +0000 Subject: [PATCH 89/99] [refactor] build prompt labels in model --- tzrec/datasets/dataset.py | 6 --- tzrec/main.py | 12 ------ tzrec/models/prompt_generative_qwen.py | 37 ++++++++++------- tzrec/models/prompt_generative_qwen_test.py | 17 ++++++-- tzrec/prompt/assembler.py | 21 ++++------ tzrec/prompt/assembler_test.py | 9 ++-- tzrec/tests/prompt_integration_test.py | 46 ++------------------- tzrec/tests/prompt_test_util.py | 14 +++---- 8 files changed, 56 insertions(+), 106 deletions(-) diff --git a/tzrec/datasets/dataset.py b/tzrec/datasets/dataset.py index b9a2fa0ac..27dbd889a 100644 --- a/tzrec/datasets/dataset.py +++ b/tzrec/datasets/dataset.py @@ -101,7 +101,6 @@ class BaseDataset(IterableDataset, metaclass=_dataset_meta_cls): debug_level (int): dataset debug level, when mode=predict and debug_level > 0, will dump fg encoded data to debug_str prompt (CompiledPrompt, optional): compiled prompt assembly contract. - prompt_ignore_index (int): label value outside the supervised response. """ def __init__( @@ -113,14 +112,12 @@ def __init__( mode: Mode = Mode.EVAL, debug_level: int = 0, prompt: Optional[CompiledPrompt] = None, - prompt_ignore_index: int = -100, ) -> None: super(BaseDataset, self).__init__() self._assembler = ( PromptAssembler( prompt.prompt_plan, prompt.sid_space, - ignore_index=prompt_ignore_index, include_response=mode != Mode.PREDICT, ) if prompt is not None @@ -805,7 +802,6 @@ def create_dataloader( debug_level: int = 0, checkpoint_state: Optional[Dict[str, Any]] = None, prompt: Optional[CompiledPrompt] = None, - prompt_ignore_index: int = -100, ) -> DataLoader: """Build dataloader. @@ -822,7 +818,6 @@ def create_dataloader( eager ``iter()`` forks workers so it reaches them. prompt (CompiledPrompt, optional): when set, each batch carries the assembled prompt streams in ``additional_infos``. - prompt_ignore_index (int): label value outside the supervised response. Return: dataloader (dataloader): a DataLoader. @@ -838,7 +833,6 @@ def create_dataloader( mode=mode, debug_level=debug_level, prompt=prompt, - prompt_ignore_index=prompt_ignore_index, ) if checkpoint_state: dataset.load_state_dict(dict(checkpoint_state)) diff --git a/tzrec/main.py b/tzrec/main.py index 14b3b794d..be0aded59 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -135,14 +135,6 @@ def _compile_prompt( ) -def _prompt_ignore_index(pipeline_config: EasyRecConfig) -> int: - """Return the label sentinel configured by a prompt-native model.""" - model_type = pipeline_config.model_config.WhichOneof("model") - if model_type != "prompt_generative_qwen": - return -100 - return int(getattr(pipeline_config.model_config, model_type).common.ignore_index) - - def _get_sampler_type(data_config: DataConfig) -> Optional[str]: try: sampler_type = ( @@ -722,7 +714,6 @@ def train_and_evaluate( # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) prompt = _compile_prompt(pipeline_config, features) - prompt_ignore_index = _prompt_ignore_index(pipeline_config) ckpt_manager = checkpoint_util.CheckpointManager( pipeline_config.model_dir, @@ -776,7 +767,6 @@ def train_and_evaluate( pipeline_config.train_input_path, mode=Mode.TRAIN, prompt=prompt, - prompt_ignore_index=prompt_ignore_index, checkpoint_state=dataloader_state, ) eval_dataloader = None @@ -789,7 +779,6 @@ def train_and_evaluate( pipeline_config.eval_input_path, mode=Mode.EVAL, prompt=prompt, - prompt_ignore_index=prompt_ignore_index, gl_cluster=gl_cluster, ) @@ -1006,7 +995,6 @@ def evaluate( eval_input_path or pipeline_config.eval_input_path, mode=Mode.EVAL, prompt=prompt, - prompt_ignore_index=_prompt_ignore_index(pipeline_config), ) sampler_type = _get_sampler_type(data_config) diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 474c1e82c..4c7237829 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -30,8 +30,9 @@ from tzrec.modules.dynamic_beam import _capped_beam_widths, dynamic_beam_search from tzrec.prompt.assembler import ( PROMPT_CU_SEQLENS, - PROMPT_LABELS, + PROMPT_INPUT_IDS, PROMPT_MAX_SEQLEN, + PROMPT_RESPONSE_LENGTHS, ) from tzrec.prompt.plan import CompiledPrompt from tzrec.protos.model_pb2 import ModelConfig @@ -139,9 +140,10 @@ def _forward_loss( padded, mask, labels = _unpack( embeds, infos[PROMPT_CU_SEQLENS], - infos[PROMPT_LABELS], int(infos[PROMPT_MAX_SEQLEN]), - self._ignore_index, + input_ids=infos[PROMPT_INPUT_IDS], + response_lengths=infos[PROMPT_RESPONSE_LENGTHS], + ignore_index=self._ignore_index, ) outputs = self.lm.model(inputs_embeds=padded, attention_mask=mask) @@ -170,9 +172,7 @@ def _generate(self, batch: Batch) -> torch.Tensor: padded, mask, _ = _unpack( embeds, infos[PROMPT_CU_SEQLENS], - None, int(infos[PROMPT_MAX_SEQLEN]), - self._ignore_index, ) space = self._prompt.sid_space tokens = dynamic_beam_search( @@ -189,11 +189,12 @@ def _generate(self, batch: Batch) -> torch.Tensor: def _unpack( embeds: torch.Tensor, cu_seqlens: torch.Tensor, - labels: Optional[torch.Tensor], max_seqlen: int, - ignore_index: int, + input_ids: Optional[torch.Tensor] = None, + response_lengths: Optional[torch.Tensor] = None, + ignore_index: int = -100, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: - """Left-pad a packed varlen batch at the LM boundary. + """Left-pad a packed batch and build response labels for training. Padding lives in this one adapter. ``max_seqlen`` is the collator's, not ``lengths.max()``: deriving it here would sync the device to the host every @@ -207,9 +208,10 @@ def _unpack( Args: embeds: packed embeddings, ``(total_tokens, hidden)``. cu_seqlens: row boundaries, ``(batch_size + 1,)``. - labels: packed labels, or None at inference where nothing is scored. max_seqlen: the collator's padded width. - ignore_index: label value for padding. + input_ids: packed token ids, or None at inference where nothing is scored. + response_lengths: number of supervised response tokens in each row. + ignore_index: label value outside the response span. Returns: Padded embeddings, attention mask and labels. @@ -223,19 +225,22 @@ def _unpack( mask = columns[None, :] >= (max_seqlen - lengths)[:, None] padded = embeds.new_zeros((batch_size, max_seqlen, hidden)) - # mask selects row-major, which is the order embeds and labels are packed in + # mask selects row-major, which is how embeds and input_ids are packed padded[mask] = embeds - if labels is None: + if input_ids is None: return padded, mask.long(), None - out_labels = torch.full( + assert response_lengths is not None + labels = torch.full( (batch_size, max_seqlen), ignore_index, - dtype=labels.dtype, + dtype=input_ids.dtype, device=embeds.device, ) - out_labels[mask] = labels - return padded, mask.long(), out_labels + labels[mask] = input_ids + response_mask = columns[None, :] >= (max_seqlen - response_lengths)[:, None] + labels[~response_mask] = ignore_index + return padded, mask.long(), labels @torch.fx.wrap diff --git a/tzrec/models/prompt_generative_qwen_test.py b/tzrec/models/prompt_generative_qwen_test.py index e4d2ab891..8bb1afc92 100644 --- a/tzrec/models/prompt_generative_qwen_test.py +++ b/tzrec/models/prompt_generative_qwen_test.py @@ -22,11 +22,17 @@ class UnpackTest(unittest.TestCase): def test_packs_rows_of_different_lengths(self) -> None: embeds = torch.arange(18, dtype=torch.float32).reshape(9, 2) cu = torch.tensor([0, 4, 9]) - ignore = -100 - labels = torch.tensor([ignore, ignore, 7, 8, ignore, ignore, ignore, 7, 8]) + input_ids = torch.tensor([1, 2, 7, 8, 3, 4, 5, 7, 8]) + response_lengths = torch.tensor([1, 2]) + ignore = -7 padded, mask, out = _unpack( - embeds, cu, labels, max_seqlen=7, ignore_index=ignore + embeds, + cu, + max_seqlen=7, + input_ids=input_ids, + response_lengths=response_lengths, + ignore_index=ignore, ) self.assertEqual(padded.shape, (2, 7, 2)) @@ -39,7 +45,10 @@ def test_packs_rows_of_different_lengths(self) -> None: torch.testing.assert_close(padded[0, :3], torch.zeros(3, 2)) torch.testing.assert_close(padded[1, :2], torch.zeros(2, 2)) torch.testing.assert_close(padded[:, -1], torch.stack([embeds[3], embeds[8]])) - self.assertEqual(out.tolist(), [[ignore] * 5 + [7, 8]] * 2) + self.assertEqual( + out.tolist(), + [[ignore] * 6 + [8], [ignore] * 5 + [7, 8]], + ) if __name__ == "__main__": diff --git a/tzrec/prompt/assembler.py b/tzrec/prompt/assembler.py index ea3478fe1..2e2ddaf9a 100644 --- a/tzrec/prompt/assembler.py +++ b/tzrec/prompt/assembler.py @@ -33,8 +33,8 @@ PROMPT_INPUT_IDS = "prompt_input_ids" PROMPT_CU_SEQLENS = "prompt_cu_seqlens" PROMPT_HOLE_POSITIONS = "prompt_hole_positions" -PROMPT_LABELS = "prompt_labels" PROMPT_MAX_SEQLEN = "prompt_max_seqlen" +PROMPT_RESPONSE_LENGTHS = "prompt_response_lengths" @dataclass @@ -47,13 +47,13 @@ class AssembledPrompt: hole_positions: absolute indices the projected embeddings overwrite, grouped by included projected occurrence in ``PromptPlan.projected_slots`` order, then by sample. - labels: ``ignore_index`` outside the response span. + response_lengths: number of response tokens in each sample. """ input_ids: np.ndarray cu_seqlens: np.ndarray hole_positions: np.ndarray - labels: np.ndarray + response_lengths: np.ndarray @property def max_seqlen(self) -> int: @@ -69,7 +69,6 @@ class PromptAssembler: Args: prompt_plan: the compiled walk order. sid_space: resolved SID token space; required when a slot renders SIDs. - ignore_index: label value outside the supervised span. include_response: whether to read and emit the supervised response. """ @@ -77,12 +76,10 @@ def __init__( self, prompt_plan: PromptPlan, sid_space: Optional[ResolvedSidSpace] = None, - ignore_index: int = -100, include_response: bool = True, ) -> None: self._prompt_plan = prompt_plan self._sid_space = sid_space - self._ignore_index = ignore_index self._response_segments = ( prompt_plan.response_segments if include_response else () ) @@ -183,7 +180,7 @@ def assemble( batch_size = len(first_inline_values) jagged_token_ids: List[int] = [] - labels: List[int] = [] + response_lengths: List[int] = [] holes_by_occurrence = [ [] for seg in self._prompt_plan.segments + self._response_segments @@ -213,10 +210,7 @@ def assemble( projected_occurrence_index, len(jagged_token_ids), ) - # supervision covers the response span only; the prompt is context. - sample_labels = [self._ignore_index] * prompt_len + sample_token_ids[ - prompt_len: - ] + response_lengths.append(len(sample_token_ids) - prompt_len) if ( self._prompt_plan.max_length and len(sample_token_ids) > self._prompt_plan.max_length @@ -228,7 +222,6 @@ def assemble( f"never truncated: cap the source features instead." ) jagged_token_ids.extend(sample_token_ids) - labels.extend(sample_labels) cu_seqlens.append(len(jagged_token_ids)) holes = [ @@ -240,7 +233,7 @@ def assemble( input_ids=np.asarray(jagged_token_ids, dtype=np.int64), cu_seqlens=np.asarray(cu_seqlens, dtype=np.int64), hole_positions=np.asarray(holes, dtype=np.int64), - labels=np.asarray(labels, dtype=np.int64), + response_lengths=np.asarray(response_lengths, dtype=np.int64), ) def assemble_batch( @@ -317,6 +310,6 @@ def assemble_batch( PROMPT_INPUT_IDS: out.input_ids, PROMPT_CU_SEQLENS: out.cu_seqlens, PROMPT_HOLE_POSITIONS: out.hole_positions, - PROMPT_LABELS: out.labels, PROMPT_MAX_SEQLEN: np.asarray(out.max_seqlen, dtype=np.int64), + PROMPT_RESPONSE_LENGTHS: out.response_lengths, } diff --git a/tzrec/prompt/assembler_test.py b/tzrec/prompt/assembler_test.py index 93609b60d..d56d1641c 100644 --- a/tzrec/prompt/assembler_test.py +++ b/tzrec/prompt/assembler_test.py @@ -135,21 +135,22 @@ def test_holes_are_grouped_by_projected_occurrence_then_sample(self) -> None: self.assertEqual(out.cu_seqlens.tolist(), [0, 5, 11]) self.assertEqual(out.hole_positions.tolist(), [0, 5, 6, 2, 3, 8, 4, 9, 10]) - def test_labels_cover_the_response_span_only(self) -> None: + def test_response_is_optional_and_its_length_is_recorded(self) -> None: plan = _plan( (Static((7, 8)),), response=(Static((9,)), _slot("answer", FillMode.INLINE)), ) - asm = PromptAssembler(plan, _sid_space(), ignore_index=-7) + asm = PromptAssembler(plan, _sid_space()) out = asm.assemble({"answer": [np.array([0, 4, 8])]}) self.assertEqual(out.input_ids.tolist(), [7, 8, 9, _BASE, _BASE + 4, _BASE + 8]) - self.assertEqual(out.labels.tolist(), [-7, -7, 9, _BASE, _BASE + 4, _BASE + 8]) + self.assertEqual(out.response_lengths.tolist(), [4]) prompt_only = PromptAssembler( - plan, _sid_space(), ignore_index=-7, include_response=False + plan, _sid_space(), include_response=False ).assemble({}, batch_size=1) self.assertEqual(prompt_only.input_ids.tolist(), [7, 8]) + self.assertEqual(prompt_only.response_lengths.tolist(), [0]) def test_rejects_raw_codes_that_carry_no_offset(self) -> None: plan = _plan((_slot("hist", FillMode.INLINE),)) diff --git a/tzrec/tests/prompt_integration_test.py b/tzrec/tests/prompt_integration_test.py index 614fdab75..349246eca 100644 --- a/tzrec/tests/prompt_integration_test.py +++ b/tzrec/tests/prompt_integration_test.py @@ -18,12 +18,6 @@ from tzrec.datasets.utils import Batch from tzrec.main import _create_model -from tzrec.models.prompt_generative_qwen import _unpack -from tzrec.prompt.assembler import ( - PROMPT_CU_SEQLENS, - PROMPT_LABELS, - PROMPT_MAX_SEQLEN, -) from tzrec.prompt.compile import compile_prompt from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.prompt_pb2 import PromptConfig @@ -77,20 +71,11 @@ def _model(self): ) def _batch(self, hist, answer): - return self._batch_rows([(hist, answer)]) - - def _batch_rows(self, rows): - hist = [h for h, _ in rows] - answer = [a for _, a in rows] parsed = { - "hist.values": torch.tensor( - offset_sid_codes([c for h in hist for c in h], _CODEBOOK) - ), - "hist.lengths": torch.tensor([len(h) for h in hist]), - "answer.values": torch.tensor( - offset_sid_codes([c for a in answer for c in a], _CODEBOOK) - ), - "answer.lengths": torch.tensor([len(a) for a in answer]), + "hist.values": torch.tensor(offset_sid_codes(hist, _CODEBOOK)), + "hist.lengths": torch.tensor([len(hist)]), + "answer.values": torch.tensor(offset_sid_codes(answer, _CODEBOOK)), + "answer.lengths": torch.tensor([len(answer)]), } streams = assemble_into(self.prompt, parsed) batch = Batch() @@ -99,29 +84,6 @@ def _batch_rows(self, rows): ) return batch - def test_every_row_is_supervised_whatever_its_length(self) -> None: - # a short row must not lose its answer to padding: the loss keeps a - # fixed-width suffix, so both rows have to contribute equally - batch = self._batch_rows( - [([0, 1, 2], [1, 2, 3]), ([0, 1, 2, 3, 0, 1], [2, 3, 0])] - ) - infos = batch.additional_infos - cu = infos[PROMPT_CU_SEQLENS] - lengths = (cu[1:] - cu[:-1]).tolist() - self.assertNotEqual(lengths[0], lengths[1], "rows must differ to be a test") - - _, _, labels = _unpack( - torch.ones(int(cu[-1]), 1), - cu, - infos[PROMPT_LABELS], - int(infos[PROMPT_MAX_SEQLEN]), - -100, - ) - window = labels[:, -self.prompt.prompt_plan.logits_suffix_len :] - supervised = (window != -100).sum(dim=1).tolist() - self.assertEqual(supervised[0], supervised[1]) - self.assertEqual(supervised[0], self.prompt.sid_space.num_levels) - def test_model_resizes_to_target_vocab(self) -> None: model = self._model() rows = model.lm.get_input_embeddings().weight.shape[0] diff --git a/tzrec/tests/prompt_test_util.py b/tzrec/tests/prompt_test_util.py index 84f1304fc..c34d9f147 100644 --- a/tzrec/tests/prompt_test_util.py +++ b/tzrec/tests/prompt_test_util.py @@ -71,20 +71,18 @@ def offset_sid_codes(codes: Sequence[Any], codebook: Sequence[int]) -> np.ndarra def assemble_into( prompt: CompiledPrompt, - parsed: Dict[str, "np.ndarray"], - ignore_index: int = -100, + parsed_features: Dict[str, "np.ndarray"], ) -> Dict[str, np.ndarray]: """Assemble one parsed batch with a temporary assembler. Args: prompt: the compiled prompt. - parsed: ``{feature}.values`` / ``{feature}.lengths`` as the data parser - emits them. - ignore_index: label value outside the supervised span. + parsed_features: ``{feature}.values`` / ``{feature}.lengths`` as the + data parser emits them. Returns: The assembled streams keyed for ``additional_infos``. """ - return PromptAssembler( - prompt.prompt_plan, prompt.sid_space, ignore_index - ).assemble_batch(parsed) + return PromptAssembler(prompt.prompt_plan, prompt.sid_space).assemble_batch( + parsed_features + ) From 5346daa8681962abc30e5b789ca740855bbf7534 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Fri, 14 Aug 2026 06:57:43 +0000 Subject: [PATCH 90/99] [refactor] clarify prompt model naming --- tzrec/datasets/dataset.py | 32 ++++---- tzrec/main.py | 36 ++++----- tzrec/models/prompt_generative_model.py | 75 +++++++++++-------- tzrec/models/prompt_generative_model_test.py | 52 ++++++------- tzrec/models/prompt_generative_qwen.py | 15 ++-- tzrec/prompt/assembler.py | 4 +- tzrec/prompt/assembler_test.py | 37 ++++++--- tzrec/prompt/compile.py | 36 ++++----- tzrec/prompt/compile_test.py | 22 ++++-- tzrec/prompt/persist.py | 36 +++++---- tzrec/prompt/persist_test.py | 33 ++++---- tzrec/prompt/plan.py | 10 +-- tzrec/protos/models/prompt_model.proto | 8 +- tzrec/protos/prompt.proto | 10 +-- .../prompt_generative_qwen_mock.config | 4 +- tzrec/tests/prompt_integration_test.py | 21 ++++-- tzrec/tests/prompt_test_util.py | 10 +-- 17 files changed, 253 insertions(+), 188 deletions(-) diff --git a/tzrec/datasets/dataset.py b/tzrec/datasets/dataset.py index 27dbd889a..5cb56b77a 100644 --- a/tzrec/datasets/dataset.py +++ b/tzrec/datasets/dataset.py @@ -100,7 +100,7 @@ class BaseDataset(IterableDataset, metaclass=_dataset_meta_cls): mode (Mode): train or eval or predict. debug_level (int): dataset debug level, when mode=predict and debug_level > 0, will dump fg encoded data to debug_str - prompt (CompiledPrompt, optional): compiled prompt assembly contract. + compiled_prompt (CompiledPrompt, optional): compiled prompt assembly contract. """ def __init__( @@ -111,16 +111,16 @@ def __init__( reserved_columns: Optional[List[str]] = None, mode: Mode = Mode.EVAL, debug_level: int = 0, - prompt: Optional[CompiledPrompt] = None, + compiled_prompt: Optional[CompiledPrompt] = None, ) -> None: super(BaseDataset, self).__init__() - self._assembler = ( + self._prompt_assembler = ( PromptAssembler( - prompt.prompt_plan, - prompt.sid_space, + compiled_prompt.prompt_plan, + compiled_prompt.sid_space, include_response=mode != Mode.PREDICT, ) - if prompt is not None + if compiled_prompt is not None else None ) self._data_config = data_config @@ -136,16 +136,16 @@ def __init__( ) parser_features = features - if prompt is not None and mode == Mode.PREDICT: + if compiled_prompt is not None and mode == Mode.PREDICT: prompt_feature_names = { feature_name - for segment in prompt.prompt_plan.segments + for segment in compiled_prompt.prompt_plan.segments if isinstance(segment, SlotSeg) for feature_name in segment.feature_names } response_feature_names = { feature_name - for segment in prompt.prompt_plan.response_segments + for segment in compiled_prompt.prompt_plan.response_segments if isinstance(segment, SlotSeg) for feature_name in segment.feature_names } @@ -416,11 +416,13 @@ def _build_batch(self, input_data: Dict[str, pa.Array]) -> Batch: else: batch = self._data_parser.to_batch(output_data) - if self._assembler is not None: + if self._prompt_assembler is not None: batch.additional_infos.update( { k: torch.from_numpy(np.asarray(v)) - for k, v in self._assembler.assemble_batch(output_data).items() + for k, v in self._prompt_assembler.assemble_batch( + output_data + ).items() } ) @@ -801,7 +803,7 @@ def create_dataloader( gl_cluster: Optional[Dict[str, Union[int, str]]] = None, debug_level: int = 0, checkpoint_state: Optional[Dict[str, Any]] = None, - prompt: Optional[CompiledPrompt] = None, + compiled_prompt: Optional[CompiledPrompt] = None, ) -> DataLoader: """Build dataloader. @@ -816,8 +818,8 @@ def create_dataloader( debug_level > 0, will dump fg encoded data to debug_str checkpoint_state (dict, optional): resume state, applied before the eager ``iter()`` forks workers so it reaches them. - prompt (CompiledPrompt, optional): when set, each batch carries the - assembled prompt streams in ``additional_infos``. + compiled_prompt (CompiledPrompt, optional): when set, each batch carries + the assembled prompt streams in ``additional_infos``. Return: dataloader (dataloader): a DataLoader. @@ -832,7 +834,7 @@ def create_dataloader( reserved_columns=reserved_columns, mode=mode, debug_level=debug_level, - prompt=prompt, + compiled_prompt=compiled_prompt, ) if checkpoint_state: dataset.load_state_dict(dict(checkpoint_state)) diff --git a/tzrec/main.py b/tzrec/main.py index be0aded59..3452ba624 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -153,7 +153,7 @@ def _create_model( labels: List[str], sample_weights: Optional[List[str]] = None, sampler_type: Optional[str] = None, - prompt: Optional[CompiledPrompt] = None, + compiled_prompt: Optional[CompiledPrompt] = None, ) -> BaseModel: """Build model. @@ -163,7 +163,7 @@ def _create_model( labels (list): list of label names. sample_weights (list): list of sample weight names. sampler_type (str): negative sampler type - prompt (CompiledPrompt, optional): forwarded to prompt-native models. + compiled_prompt (CompiledPrompt, optional): forwarded to prompt-native models. Return: model: a EasyRec Model. @@ -178,7 +178,7 @@ def _create_model( labels, sample_weights=sample_weights, sampler_type=sampler_type, - prompt=prompt, + compiled_prompt=compiled_prompt, ) kernel = Kernel[KernelProto.Name(model_config.kernel)] @@ -713,7 +713,7 @@ def train_and_evaluate( data_config = pipeline_config.data_config # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) - prompt = _compile_prompt(pipeline_config, features) + compiled_prompt = _compile_prompt(pipeline_config, features) ckpt_manager = checkpoint_util.CheckpointManager( pipeline_config.model_dir, @@ -753,7 +753,7 @@ def train_and_evaluate( # Restore dataloader state before create_dataloader starts its workers dataloader_state: Optional[Dict[str, Any]] = None if ckpt_path: - check_prompt_assets(prompt, ckpt_path) + check_prompt_assets(compiled_prompt, ckpt_path) if ckpt_path and continue_train: dataloader_state = ckpt_manager.restore_dataloader_state(ckpt_path) if dataloader_state and not restore_from_model_dir: @@ -766,7 +766,7 @@ def train_and_evaluate( features, pipeline_config.train_input_path, mode=Mode.TRAIN, - prompt=prompt, + compiled_prompt=compiled_prompt, checkpoint_state=dataloader_state, ) eval_dataloader = None @@ -778,7 +778,7 @@ def train_and_evaluate( features, pipeline_config.eval_input_path, mode=Mode.EVAL, - prompt=prompt, + compiled_prompt=compiled_prompt, gl_cluster=gl_cluster, ) @@ -791,7 +791,7 @@ def train_and_evaluate( list(data_config.label_fields), sample_weights=list(data_config.sample_weight_fields), sampler_type=sampler_type, - prompt=prompt, + compiled_prompt=compiled_prompt, ) # Cold start only; a resumed or fine-tuned run gets its weights from DCP. if ckpt_path is None: @@ -987,14 +987,14 @@ def evaluate( data_config = pipeline_config.data_config # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) - prompt = _compile_prompt(pipeline_config, features) + compiled_prompt = _compile_prompt(pipeline_config, features) eval_dataloader = create_dataloader( data_config, features, eval_input_path or pipeline_config.eval_input_path, mode=Mode.EVAL, - prompt=prompt, + compiled_prompt=compiled_prompt, ) sampler_type = _get_sampler_type(data_config) @@ -1006,7 +1006,7 @@ def evaluate( list(data_config.label_fields), sample_weights=list(data_config.sample_weight_fields), sampler_type=sampler_type, - prompt=prompt, + compiled_prompt=compiled_prompt, ) model = TrainWrapper( model, device=device, mixed_precision=train_config.mixed_precision @@ -1038,7 +1038,7 @@ def evaluate( ) if checkpoint_path: - check_prompt_assets(prompt, checkpoint_path) + check_prompt_assets(compiled_prompt, checkpoint_path) ckpt_manager.restore( checkpoint_path, model, @@ -1353,7 +1353,7 @@ def predict( data_config.drop_remainder = False # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) - prompt = _compile_prompt(pipeline_config, features) + compiled_prompt = _compile_prompt(pipeline_config, features) infer_dataloader = create_dataloader( data_config, @@ -1361,7 +1361,7 @@ def predict( predict_input_path, reserved_columns=reserved_cols, mode=Mode.PREDICT, - prompt=prompt, + compiled_prompt=compiled_prompt, debug_level=debug_level, ) infer_iterator = infer_dataloader.get_iterator() # pyre-ignore[16] @@ -1628,7 +1628,7 @@ def predict_checkpoint( data_config = pipeline_config.data_config # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) - prompt = _compile_prompt(pipeline_config, features) + compiled_prompt = _compile_prompt(pipeline_config, features) # Build dataloader predict_dataloader = create_dataloader( @@ -1637,7 +1637,7 @@ def predict_checkpoint( predict_input_path, reserved_columns=reserved_cols, mode=Mode.PREDICT, - prompt=prompt, + compiled_prompt=compiled_prompt, debug_level=debug_level, ) @@ -1657,7 +1657,7 @@ def predict_checkpoint( pipeline_config.model_config, features, [], - prompt=prompt, + compiled_prompt=compiled_prompt, ) model.set_is_inference(True) model = PredictWrapper( @@ -1693,7 +1693,7 @@ def predict_checkpoint( model.eval() if checkpoint_path: - check_prompt_assets(prompt, checkpoint_path) + check_prompt_assets(compiled_prompt, checkpoint_path) ckpt_manager.restore( checkpoint_path, model, diff --git a/tzrec/models/prompt_generative_model.py b/tzrec/models/prompt_generative_model.py index 20d8f7e77..02673d742 100644 --- a/tzrec/models/prompt_generative_model.py +++ b/tzrec/models/prompt_generative_model.py @@ -53,7 +53,7 @@ class BasePromptGenerativeModel(BaseModel): features: every created feature. labels: data_config label fields. sample_weights: optional sample weight fields. - prompt: the compiled prompt; required. + compiled_prompt: the compiled prompt; required. """ def __init__( @@ -62,54 +62,63 @@ def __init__( features: List[BaseFeature], labels: List[str], sample_weights: Optional[List[str]] = None, - prompt: Optional[CompiledPrompt] = None, + compiled_prompt: Optional[CompiledPrompt] = None, **kwargs: Any, ) -> None: super().__init__(model_config, features, labels, sample_weights, **kwargs) - if prompt is None: + if compiled_prompt is None: raise ValueError( f"{type(self).__name__} needs a compiled prompt; call " f"compile_prompt(pipeline_config.prompt_config, features) and " f"pass it to _create_model." ) - if prompt.sid_space is None: + if compiled_prompt.sid_space is None: raise ValueError( f"{type(self).__name__}: prompt_config declares no sid_space, " f"so there is no SID vocabulary to extend or decode." ) - self._prompt = prompt + self._prompt = compiled_prompt cfg = self._model_config - self.lm = self._build_backbone(cfg.hf_model_id, cfg.common.param_dtype) + self.lm = self._build_backbone( + cfg.hf_model_name_or_path, cfg.common.lm_parameter_dtype + ) # Every run replaces this initialization from pretrained or DCP weights. self.lm.resize_token_embeddings( - prompt.sid_space.target_vocab, mean_resizing=False - ) - self.embedding_group = EmbeddingGroup( - self._features, list(self._prompt.projection_plan.feature_groups) + compiled_prompt.sid_space.target_vocab_size, mean_resizing=False ) - self._build_projections() + self.init_input() + # decode subtracts these every step; a buffer follows the module's device self.register_buffer( "_level_offsets", - torch.tensor(prompt.sid_space.level_offsets), + torch.tensor(compiled_prompt.sid_space.level_offsets), persistent=False, ) - def _build_backbone(self, hf_model_id: str, param_dtype: int) -> nn.Module: + def init_input(self) -> None: + """Build the projected-slot embedding groups and projection modules.""" + self.embedding_group = EmbeddingGroup( + self._features, list(self._prompt.projection_plan.feature_groups) + ) + self._build_projections() + + def _build_backbone( + self, hf_model_name_or_path: str, lm_parameter_dtype: int + ) -> nn.Module: """Build the LM from config, so HF weights load only on cold start. Args: - hf_model_id: hub id or local directory naming the architecture and - cold-start weights. - param_dtype: master-weight dtype. + hf_model_name_or_path: hub id or local directory naming the + architecture and cold-start weights. + lm_parameter_dtype: dtype of the LM parameters. Returns: A randomly initialized backbone with the requested parameter dtype. """ - config = AutoConfig.from_pretrained(hf_model_id) + config = AutoConfig.from_pretrained(hf_model_name_or_path) model = AutoModelForCausalLM.from_config(config) - return model.to(_PARAM_DTYPE[param_dtype]) + return model.to(_PARAM_DTYPE[lm_parameter_dtype]) def _build_projections(self) -> None: """One module per resolved id, aligned with ``prompt_plan.projected_slots``. @@ -121,14 +130,14 @@ def _build_projections(self) -> None: projection_plan = self._prompt.projection_plan hidden_size = int(self.lm.config.hidden_size) - built: Dict[str, PromptProjection] = {} + modules_by_id: Dict[str, PromptProjection] = {} in_dims: Dict[str, int] = {} - aligned: List[PromptProjection] = [] + aligned_modules: List[PromptProjection] = [] for seg in prompt_plan.projected_slots: module_id = projection_plan.slot_to_module[seg.slot_id] in_dim = self.embedding_group.group_total_dim(seg.name + seg.output_key) - if module_id not in built: - built[module_id] = PromptProjection( + if module_id not in modules_by_id: + modules_by_id[module_id] = PromptProjection( projection_plan.projections[module_id], in_dim, hidden_size ) in_dims[module_id] = in_dim @@ -138,16 +147,16 @@ def _build_projections(self) -> None: f"different group dims ({in_dims[module_id]} vs " f"{in_dim}); they cannot share a module." ) - aligned.append(built[module_id]) - self.projections = nn.ModuleDict(built) - self._slot_projections = aligned + aligned_modules.append(modules_by_id[module_id]) + self.projections = nn.ModuleDict(modules_by_id) + self._slot_projections = aligned_modules def hf_backbone(self) -> nn.Module: """The HF module export and checkpointing reach for.""" return self.lm - def _prompt_embeds(self, batch: Batch) -> torch.Tensor: - """Gather the token stream, then overwrite the projected positions. + def build_input(self, batch: Batch) -> torch.Tensor: + """Build packed LM input embeddings and fill projected positions. Args: batch: carries the packed prompt in ``additional_infos``. @@ -164,14 +173,16 @@ def _prompt_embeds(self, batch: Batch) -> torch.Tensor: grouped = self.embedding_group(batch) hidden_size = embeds.shape[-1] - parts = [ + projected_embeddings = [ proj(grouped[seg.name + seg.output_key]).reshape(-1, hidden_size) for seg, proj in zip(prompt_plan.projected_slots, self._slot_projections) ] # The assembler records holes in this projected-occurrence-major order. # out of place: embeds carries grad from the embedding lookup return embeds.index_copy( - 0, batch.additional_infos[PROMPT_HOLE_POSITIONS], torch.cat(parts) + 0, + batch.additional_infos[PROMPT_HOLE_POSITIONS], + torch.cat(projected_embeddings), ) def _tokens_to_local_codes( @@ -187,7 +198,7 @@ def _tokens_to_local_codes( ``(batch_size, beams, num_levels)`` local codes. """ space = self._prompt.sid_space - codes = tokens - space.base_vocab - self._level_offsets + codes = tokens - space.base_vocab_size - self._level_offsets return codes.view(batch_size, -1, space.num_levels) def init_loss(self) -> None: @@ -248,11 +259,11 @@ def save_assets(self, target_dir: str) -> None: def init_from_pretrained(self) -> None: """Load HF weights once, on a cold start only.""" - source = self._model_config.hf_model_id + source = self._model_config.hf_model_name_or_path logger.info(f"loading pretrained weights from [{source}].") pretrained = AutoModelForCausalLM.from_pretrained(source) pretrained.resize_token_embeddings( - self._prompt.sid_space.target_vocab, mean_resizing=True + self._prompt.sid_space.target_vocab_size, mean_resizing=True ) self.lm.load_state_dict(pretrained.state_dict()) del pretrained diff --git a/tzrec/models/prompt_generative_model_test.py b/tzrec/models/prompt_generative_model_test.py index 6de99a22e..ad310cd66 100644 --- a/tzrec/models/prompt_generative_model_test.py +++ b/tzrec/models/prompt_generative_model_test.py @@ -64,34 +64,36 @@ def setUp(self) -> None: create_prompt_feature(_HIST), create_prompt_feature(_ANSWER), ] - self.prompt = self._compile(self.features) + self.compiled_prompt = self._compile(self.features) def _compile(self, features, template="History : {{hist}} . Predict :", **kwargs): - cfg = PromptConfig(tokenizer=self.tok, prompt=template, **kwargs) + cfg = PromptConfig(tokenizer_path=self.tok, prompt=template, **kwargs) cfg.sid_space.codebook.extend(_CODEBOOK) return compile_prompt(cfg, features, model_dir=self.test_dir) def _model( self, features=None, - prompt=-1, + compiled_prompt=-1, beam_widths=(2, 2, 2), num_return_sequences=2, ): model_config = ModelConfig() qwen = model_config.prompt_generative_qwen - qwen.hf_model_id = self.backbone + qwen.hf_model_name_or_path = self.backbone qwen.common.beam_widths.extend(beam_widths) qwen.common.num_return_sequences = num_return_sequences return _create_model( model_config, self.features if features is None else features, ["answer"], - prompt=self.prompt if prompt == -1 else prompt, + compiled_prompt=( + self.compiled_prompt if compiled_prompt == -1 else compiled_prompt + ), ) - def _batch(self, parsed, prompt=None, sparse=None): - streams = assemble_into(prompt or self.prompt, parsed) + def _batch(self, parsed, compiled_prompt=None, sparse=None): + streams = assemble_into(compiled_prompt or self.compiled_prompt, parsed) batch = Batch(sparse_features={BASE_DATA_GROUP: sparse} if sparse else {}) batch.additional_infos.update( {k: torch.from_numpy(np.asarray(v)) for k, v in streams.items()} @@ -100,7 +102,7 @@ def _batch(self, parsed, prompt=None, sparse=None): def test_tokens_to_local_codes_undoes_shifts_and_groups_beams(self) -> None: model = self._model() - space = self.prompt.sid_space + space = self.compiled_prompt.sid_space local_codes = torch.tensor( [ [0, 1, 3], @@ -109,7 +111,7 @@ def test_tokens_to_local_codes_undoes_shifts_and_groups_beams(self) -> None: [2, 2, 1], ] ) - tokens = local_codes + torch.tensor(space.level_offsets) + space.base_vocab + tokens = local_codes + torch.tensor(space.level_offsets) + space.base_vocab_size codes = model._tokens_to_local_codes(tokens, batch_size=2) self.assertEqual(codes.shape, (2, 2, space.num_levels)) @@ -117,15 +119,15 @@ def test_tokens_to_local_codes_undoes_shifts_and_groups_beams(self) -> None: def test_rejects_a_model_built_without_a_prompt(self) -> None: with self.assertRaisesRegex(ValueError, "needs a compiled prompt"): - self._model(prompt=None) + self._model(compiled_prompt=None) def test_rejects_a_prompt_that_declares_no_sid_space(self) -> None: - cfg = PromptConfig(tokenizer=self.tok, prompt="History : {{hist}} .") - prompt = compile_prompt(cfg, self.features, model_dir=self.test_dir) - self.assertIsNone(prompt.sid_space) + cfg = PromptConfig(tokenizer_path=self.tok, prompt="History : {{hist}} .") + compiled_prompt = compile_prompt(cfg, self.features, model_dir=self.test_dir) + self.assertIsNone(compiled_prompt.sid_space) with self.assertRaisesRegex(ValueError, "declares no sid_space"): - self._model(prompt=prompt) + self._model(compiled_prompt=compiled_prompt) def test_shared_projection_name_requires_matching_widths(self) -> None: features = [ @@ -135,7 +137,7 @@ def test_shared_projection_name_requires_matching_widths(self) -> None: create_prompt_feature(_projected("pb", 16)), ] cfg = PromptConfig( - tokenizer=self.tok, + tokenizer_path=self.tok, prompt="History : {{hist}} . {{pa}} {{pb}} Predict :", response="{{answer}}", ) @@ -143,10 +145,10 @@ def test_shared_projection_name_requires_matching_widths(self) -> None: for name in ("pa", "pb"): slot = cfg.slots.add(name=name, projection_name="shared") slot.feature_names.append(name) - prompt = compile_prompt(cfg, features, model_dir=self.test_dir) + compiled_prompt = compile_prompt(cfg, features, model_dir=self.test_dir) with self.assertRaisesRegex(ValueError, "cannot share a module"): - self._model(features=features, prompt=prompt) + self._model(features=features, compiled_prompt=compiled_prompt) def test_projected_slot_overwrites_sentinels_and_backpropagates(self) -> None: features = [ @@ -154,12 +156,12 @@ def test_projected_slot_overwrites_sentinels_and_backpropagates(self) -> None: create_prompt_feature(_ANSWER), create_prompt_feature(_projected("prof", 8)), ] - prompt = self._compile( + compiled_prompt = self._compile( features, template="History : {{hist}} . Predict {{prof}} :", response="{{answer}}", ) - model = self._model(features=features, prompt=prompt) + model = self._model(features=features, compiled_prompt=compiled_prompt) # the embedding table is built on meta until something materializes it init_parameters(model, device=torch.device("cpu")) batch = self._batch( @@ -173,7 +175,7 @@ def test_projected_slot_overwrites_sentinels_and_backpropagates(self) -> None: "prof.values": torch.tensor([5, 9]), "prof.lengths": torch.tensor([2]), }, - prompt=prompt, + compiled_prompt=compiled_prompt, sparse=KeyedJaggedTensor.from_lengths_sync( keys=["prof"], values=torch.tensor([5, 9]), @@ -181,7 +183,7 @@ def test_projected_slot_overwrites_sentinels_and_backpropagates(self) -> None: ), ) - embeds = model._prompt_embeds(batch) + embeds = model.build_input(batch) raw = model.lm.get_input_embeddings()(batch.additional_infos[PROMPT_INPUT_IDS]) holes = batch.additional_infos[PROMPT_HOLE_POSITIONS] self.assertGreater(holes.numel(), 0) @@ -212,14 +214,14 @@ def test_metric_averages_the_loss_across_batches(self) -> None: def test_init_from_pretrained_replaces_the_empty_weights(self) -> None: model = self._model() - base_vocab = self.prompt.sid_space.base_vocab - before = model.lm.get_input_embeddings().weight[:base_vocab].clone() + base_vocab_size = self.compiled_prompt.sid_space.base_vocab_size + before = model.lm.get_input_embeddings().weight[:base_vocab_size].clone() model.init_from_pretrained() - after = model.lm.get_input_embeddings().weight[:base_vocab] + after = model.lm.get_input_embeddings().weight[:base_vocab_size] # the checkpoint rows land verbatim; only the appended SID rows are new reference = AutoModelForCausalLM.from_pretrained(self.backbone) - expected = reference.get_input_embeddings().weight[:base_vocab] + expected = reference.get_input_embeddings().weight[:base_vocab_size] self.assertFalse(torch.allclose(before, expected)) torch.testing.assert_close(after, expected) diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 4c7237829..9c335365e 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -47,7 +47,7 @@ class PromptGenerativeQwen(BasePromptGenerativeModel): features: every created feature. labels: data_config label fields. sample_weights: optional sample weight fields. - prompt: the compiled prompt; required. + compiled_prompt: the compiled prompt; required. """ def __init__( @@ -56,11 +56,16 @@ def __init__( features: List[BaseFeature], labels: List[str], sample_weights: Optional[List[str]] = None, - prompt: Optional[CompiledPrompt] = None, + compiled_prompt: Optional[CompiledPrompt] = None, **kwargs: Any, ) -> None: super().__init__( - model_config, features, labels, sample_weights, prompt, **kwargs + model_config, + features, + labels, + sample_weights, + compiled_prompt=compiled_prompt, + **kwargs, ) common = self._model_config.common self._ignore_index = int(common.ignore_index) @@ -119,7 +124,7 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: # the embedding lookup stays traceable so the train pipeline can still # see the sharded module and prefetch it; only the padding and the LM, # which read the collator's width as a host int, are hidden. - return _fx_wrapped_loss(self, self._prompt_embeds(batch), batch) + return _fx_wrapped_loss(self, self.build_input(batch), batch) def _forward_loss( self, embeds: torch.Tensor, batch: Batch @@ -167,7 +172,7 @@ def _generate(self, batch: Batch) -> torch.Tensor: Returns: ``(B, num_return, num_levels)`` local codes, best first. """ - embeds = self._prompt_embeds(batch) + embeds = self.build_input(batch) infos = batch.additional_infos padded, mask, _ = _unpack( embeds, diff --git a/tzrec/prompt/assembler.py b/tzrec/prompt/assembler.py index 2e2ddaf9a..8cab58acb 100644 --- a/tzrec/prompt/assembler.py +++ b/tzrec/prompt/assembler.py @@ -103,7 +103,7 @@ def _inline_tokens(self, name: str, values: np.ndarray) -> np.ndarray: """Validate offset SID codes against their bands and shift to token ids. The data carries ``level_offsets[l] + code``; the LM vocabulary needs - one further uniform shift by ``base_vocab``. + one further uniform shift by ``base_vocab_size``. """ assert self._sid_space is not None levels = self._sid_space.num_levels @@ -120,7 +120,7 @@ def _inline_tokens(self, name: str, values: np.ndarray) -> np.ndarray: f"[level_offsets[l], level_offsets[l] + codebook[l]). Read the " f"offset_codebook column, not codebook or origin_codebook." ) - return values.astype(np.int64, copy=False) + self._sid_space.base_vocab + return values.astype(np.int64, copy=False) + self._sid_space.base_vocab_size def _emit_sample( self, diff --git a/tzrec/prompt/assembler_test.py b/tzrec/prompt/assembler_test.py index d56d1641c..ecd32f138 100644 --- a/tzrec/prompt/assembler_test.py +++ b/tzrec/prompt/assembler_test.py @@ -26,7 +26,7 @@ from tzrec.protos.model_pb2 import FeatureGroupType from tzrec.tests.prompt_test_util import assemble_into -_BASE = 1000 +_BASE_VOCAB_SIZE = 1000 _SENTINEL = 1099 @@ -38,11 +38,11 @@ def _sid_space(codebook=(4, 4, 4)) -> ResolvedSidSpace: return ResolvedSidSpace( codebook=tuple(codebook), num_levels=len(codebook), - base_vocab=_BASE, + base_vocab_size=_BASE_VOCAB_SIZE, level_offsets=tuple(offsets), - band_lo=tuple(_BASE + o for o in offsets), - band_hi=tuple(_BASE + o + s - 1 for o, s in zip(offsets, codebook)), - target_vocab=1152, + band_lo=tuple(_BASE_VOCAB_SIZE + o for o in offsets), + band_hi=tuple(_BASE_VOCAB_SIZE + o + s - 1 for o, s in zip(offsets, codebook)), + target_vocab_size=1152, sentinel_token_id=_SENTINEL, eos_token_id=2, pad_token_id=3, @@ -97,7 +97,14 @@ def test_inline_sid_gets_the_base_vocab_shift(self) -> None: out = asm.assemble({"hist": [np.array([1, 6, 11])]}) self.assertEqual( - out.input_ids.tolist(), [7, 8, _BASE + 1, _BASE + 6, _BASE + 11] + out.input_ids.tolist(), + [ + 7, + 8, + _BASE_VOCAB_SIZE + 1, + _BASE_VOCAB_SIZE + 6, + _BASE_VOCAB_SIZE + 11, + ], ) self.assertEqual(out.cu_seqlens.tolist(), [0, 5]) self.assertEqual(out.hole_positions.size, 0) @@ -143,7 +150,17 @@ def test_response_is_optional_and_its_length_is_recorded(self) -> None: asm = PromptAssembler(plan, _sid_space()) out = asm.assemble({"answer": [np.array([0, 4, 8])]}) - self.assertEqual(out.input_ids.tolist(), [7, 8, 9, _BASE, _BASE + 4, _BASE + 8]) + self.assertEqual( + out.input_ids.tolist(), + [ + 7, + 8, + 9, + _BASE_VOCAB_SIZE, + _BASE_VOCAB_SIZE + 4, + _BASE_VOCAB_SIZE + 8, + ], + ) self.assertEqual(out.response_lengths.tolist(), [4]) prompt_only = PromptAssembler( @@ -188,7 +205,7 @@ def test_column_shaped_values_are_flattened(self) -> None: from tzrec.prompt.plan import CompiledPrompt, ProjectionPlan plan = _plan((_slot("hist", FillMode.INLINE),)) - prompt = CompiledPrompt( + compiled_prompt = CompiledPrompt( sid_space=_sid_space(), prompt_plan=plan, projection_plan=ProjectionPlan(projections={}, slot_to_module={}), @@ -200,9 +217,9 @@ def test_column_shaped_values_are_flattened(self) -> None: "hist.values": np.array([[1], [6], [11], [0], [4], [8]]), "hist.lengths": np.array([3, 3]), } - out = assemble_into(prompt, parsed) + out = assemble_into(compiled_prompt, parsed) self.assertEqual(out["prompt_cu_seqlens"].tolist(), [0, 3, 6]) - self.assertEqual(out["prompt_input_ids"].tolist()[0], _BASE + 1) + self.assertEqual(out["prompt_input_ids"].tolist()[0], _BASE_VOCAB_SIZE + 1) def test_rejects_inconsistent_slot_batch_sizes(self) -> None: plan = _plan( diff --git a/tzrec/prompt/compile.py b/tzrec/prompt/compile.py index 1b7680de6..e65d2bf48 100644 --- a/tzrec/prompt/compile.py +++ b/tzrec/prompt/compile.py @@ -125,9 +125,9 @@ def _derive_slot_layout( return group_type, fill_mode -def _atom_tokens(sid_space: SidSpace) -> List[str]: - """Render the SID atom tokens, one per flat index.""" - fmt = sid_space.atom_token_format +def _render_sid_tokens(sid_space: SidSpace) -> List[str]: + """Render the SID tokens, one per flat index.""" + fmt = sid_space.token_format return [fmt.replace("{i}", str(i)) for i in range(sum(sid_space.codebook))] @@ -140,9 +140,9 @@ def _read_manifest_codebook(path: str) -> List[int]: def _build_sid_space( - cfg: PromptConfig, tok: Tokenizer, base_vocab: int, has_projection: bool + cfg: PromptConfig, tok: Tokenizer, base_vocab_size: int, has_projection: bool ) -> Optional[ResolvedSidSpace]: - """Extend the tokenizer with SID atoms and resolve the token space.""" + """Extend the tokenizer with SID tokens and resolve the token space.""" if not cfg.HasField("sid_space"): return None space = cfg.sid_space @@ -161,14 +161,16 @@ def _build_sid_space( f"and the decode bands would disagree." ) - atoms = _atom_tokens(space) - present = [a for a in atoms if tok.token_to_id(a) is not None] - if present: + sid_tokens = _render_sid_tokens(space) + existing_sid_tokens = [ + token for token in sid_tokens if tok.token_to_id(token) is not None + ] + if existing_sid_tokens: raise ValueError( - f"SID atoms are already in the base tokenizer, e.g. {present[:3]}; " - f"change sid_space.atom_token_format." + "SID tokens are already in the base tokenizer, e.g. " + f"{existing_sid_tokens[:3]}; change sid_space.token_format." ) - tok.add_special_tokens(atoms) + tok.add_special_tokens(sid_tokens) sentinel_id = None if has_projection: @@ -186,17 +188,17 @@ def _build_sid_space( for size in codebook: offsets.append(running) running += size - lo = [base_vocab + o for o in offsets] + lo = [base_vocab_size + o for o in offsets] hi = [lo[i] + codebook[i] - 1 for i in range(len(codebook))] return ResolvedSidSpace( codebook=tuple(codebook), num_levels=len(codebook), - base_vocab=base_vocab, + base_vocab_size=base_vocab_size, level_offsets=tuple(offsets), band_lo=tuple(lo), band_hi=tuple(hi), - target_vocab=_ceil_to( + target_vocab_size=_ceil_to( tok.get_vocab_size(with_added_tokens=True), space.vocab_pad_to_multiple_of, ), @@ -295,13 +297,13 @@ def compile_prompt( f"embedding -- so it has no group to project; drop its projection." ) - tok = Tokenizer.from_file(cfg.tokenizer) - base_vocab = tok.get_vocab_size(with_added_tokens=True) + tok = Tokenizer.from_file(cfg.tokenizer_path) + base_vocab_size = tok.get_vocab_size(with_added_tokens=True) has_projection = any( fill_mode is FillMode.PROJECTED for fill_mode in fill_modes_by_slot_name.values() ) - sid_space = _build_sid_space(cfg, tok, base_vocab, has_projection) + sid_space = _build_sid_space(cfg, tok, base_vocab_size, has_projection) tokenizer_dir = "" if model_dir: diff --git a/tzrec/prompt/compile_test.py b/tzrec/prompt/compile_test.py index 3bcb86db6..6a9fb18dc 100644 --- a/tzrec/prompt/compile_test.py +++ b/tzrec/prompt/compile_test.py @@ -46,7 +46,7 @@ def setUp(self) -> None: ) def _config(self, **kwargs) -> PromptConfig: - cfg = PromptConfig(tokenizer=self.tok_path, **kwargs) + cfg = PromptConfig(tokenizer_path=self.tok_path, **kwargs) return cfg def _compile(self, cfg, features): @@ -58,15 +58,21 @@ def test_sid_space_resolves_offsets_and_bands(self) -> None: compiled = self._compile(cfg, [create_prompt_feature(_HIST)]) space = compiled.sid_space - base = space.base_vocab + base_vocab_size = space.base_vocab_size self.assertEqual(space.num_levels, 3) self.assertEqual(sum(space.codebook), 12) self.assertEqual(space.level_offsets, (0, 4, 8)) - self.assertEqual(space.band_lo, (base, base + 4, base + 8)) - self.assertEqual(space.band_hi, (base + 3, base + 7, base + 11)) + self.assertEqual( + space.band_lo, + (base_vocab_size, base_vocab_size + 4, base_vocab_size + 8), + ) + self.assertEqual( + space.band_hi, + (base_vocab_size + 3, base_vocab_size + 7, base_vocab_size + 11), + ) # no slot projects, so no sentinel is materialized self.assertIsNone(space.sentinel_token_id) - self.assertEqual(space.target_vocab % 128, 0) + self.assertEqual(space.target_vocab_size % 128, 0) def test_inline_needs_no_group_projected_gets_one(self) -> None: cfg = self._config(prompt="History : {{hist}} . Profile : {{prof}}") @@ -167,10 +173,10 @@ def test_rejects_a_projection_on_an_inline_slot(self) -> None: with self.assertRaisesRegex(ValueError, "is INLINE"): self._compile(cfg, [create_prompt_feature(_HIST)]) - def test_atoms_absent_from_the_base_tokenizer(self) -> None: + def test_sid_tokens_absent_from_the_base_tokenizer(self) -> None: cfg = self._config(prompt="X : {{hist}}") cfg.sid_space.codebook.extend([4]) - cfg.sid_space.atom_token_format = "History" + cfg.sid_space.token_format = "History" with self.assertRaisesRegex(ValueError, "already in the base tokenizer"): self._compile(cfg, [create_prompt_feature(_HIST)]) @@ -180,7 +186,7 @@ def test_extended_tokenizer_is_written(self) -> None: compiled = self._compile(cfg, [create_prompt_feature(_HIST)]) written = os.path.join(compiled.tokenizer_dir, "tokenizer.json") self.assertTrue(os.path.exists(written)) - # the atoms round-trip, which is what serving reloads + # the SID tokens round-trip, which is what serving reloads reloaded = Tokenizer.from_file(written) self.assertIsNotNone(reloaded.token_to_id("<|sid_0|>")) self.assertIsNotNone(reloaded.token_to_id("<|sid_7|>")) diff --git a/tzrec/prompt/persist.py b/tzrec/prompt/persist.py index 6a63423fc..464affd1c 100644 --- a/tzrec/prompt/persist.py +++ b/tzrec/prompt/persist.py @@ -49,14 +49,14 @@ def _plain(value: Any) -> Any: return value -def save_prompt_assets(prompt: CompiledPrompt, target_dir: str) -> None: +def save_prompt_assets(compiled_prompt: CompiledPrompt, target_dir: str) -> None: """Write the prompt contract into a checkpoint or export directory. Rank 0 only: every rank reaches here on save, and concurrent json.dump and copytree to one path can interleave into a truncated file. Args: - prompt: the compiled prompt. + compiled_prompt: the compiled prompt. target_dir: the checkpoint or export directory. """ if int(os.environ.get("RANK", 0)) != 0: @@ -65,19 +65,24 @@ def save_prompt_assets(prompt: CompiledPrompt, target_dir: str) -> None: os.makedirs(out, exist_ok=True) with open(os.path.join(out, _SID_SPACE), "w") as f: - json.dump(_plain(prompt.sid_space), f, indent=2) + json.dump(_plain(compiled_prompt.sid_space), f, indent=2) with open(os.path.join(out, _PROMPT_PLAN), "w") as f: - json.dump(_plain(prompt.prompt_plan), f, indent=2) + json.dump(_plain(compiled_prompt.prompt_plan), f, indent=2) with open(os.path.join(out, _HASHES), "w") as f: json.dump( - {"vocab_hash": prompt.vocab_hash, "plan_hash": prompt.plan_hash}, + { + "vocab_hash": compiled_prompt.vocab_hash, + "plan_hash": compiled_prompt.plan_hash, + }, f, indent=2, ) - if prompt.tokenizer_dir and os.path.isdir(prompt.tokenizer_dir): + if compiled_prompt.tokenizer_dir and os.path.isdir(compiled_prompt.tokenizer_dir): shutil.copytree( - prompt.tokenizer_dir, os.path.join(out, _TOKENIZER), dirs_exist_ok=True + compiled_prompt.tokenizer_dir, + os.path.join(out, _TOKENIZER), + dirs_exist_ok=True, ) @@ -90,7 +95,9 @@ def read_prompt_hashes(source_dir: str) -> Optional[Dict[str, str]]: return json.load(f) -def check_prompt_assets(prompt: Optional[CompiledPrompt], ckpt_dir: str) -> None: +def check_prompt_assets( + compiled_prompt: Optional[CompiledPrompt], ckpt_dir: str +) -> None: """Compare a compiled prompt against what a checkpoint recorded. A ``vocab_hash`` mismatch is fatal: the decode bands would point at token @@ -99,11 +106,11 @@ def check_prompt_assets(prompt: Optional[CompiledPrompt], ckpt_dir: str) -> None warns. Args: - prompt: the freshly compiled prompt, or None when the pipeline declares - no prompt_config. + compiled_prompt: the freshly compiled prompt, or None when the pipeline + declares no prompt_config. ckpt_dir: the checkpoint being restored. """ - if prompt is None: + if compiled_prompt is None: return recorded = read_prompt_hashes(ckpt_dir) if recorded is None: @@ -113,15 +120,16 @@ def check_prompt_assets(prompt: Optional[CompiledPrompt], ckpt_dir: str) -> None ) return - if recorded.get("vocab_hash") != prompt.vocab_hash: + if recorded.get("vocab_hash") != compiled_prompt.vocab_hash: raise ValueError( f"prompt vocabulary does not match checkpoint [{ckpt_dir}]: the " f"checkpoint was trained against {recorded.get('vocab_hash')} but " - f"prompt_config now compiles to {prompt.vocab_hash}. The SID space " + f"prompt_config now compiles to {compiled_prompt.vocab_hash}. The SID " + f"space " f"or the tokenizer changed, so the decode bands no longer address " f"the rows these weights learned." ) - if recorded.get("plan_hash") != prompt.plan_hash: + if recorded.get("plan_hash") != compiled_prompt.plan_hash: logger.warning( f"prompt plan differs from checkpoint [{ckpt_dir}]: the vocabulary " f"matches, so the weights are usable, but the template, slots or " diff --git a/tzrec/prompt/persist_test.py b/tzrec/prompt/persist_test.py index 1d3182978..08dc6c22b 100644 --- a/tzrec/prompt/persist_test.py +++ b/tzrec/prompt/persist_test.py @@ -46,14 +46,14 @@ def setUp(self) -> None: ] def _compile(self, codebook=(4, 4, 4), prompt="History : {{hist}}"): - cfg = PromptConfig(tokenizer=self.tok_path, prompt=prompt) + cfg = PromptConfig(tokenizer_path=self.tok_path, prompt=prompt) cfg.sid_space.codebook.extend(codebook) return compile_prompt(cfg, self.features, model_dir=self.test_dir) def test_writes_a_self_describing_directory(self) -> None: - prompt = self._compile() + compiled_prompt = self._compile() ckpt = os.path.join(self.test_dir, "model.ckpt-1") - save_prompt_assets(prompt, ckpt) + save_prompt_assets(compiled_prompt, ckpt) out = os.path.join(ckpt, PROMPT_DIR) for name in ("sid_space.json", "prompt_plan.json", "prompt_hashes.json"): @@ -64,18 +64,19 @@ def test_writes_a_self_describing_directory(self) -> None: ) def test_sid_space_round_trips_as_plain_json(self) -> None: - prompt = self._compile() + compiled_prompt = self._compile() ckpt = os.path.join(self.test_dir, "model.ckpt-1") - save_prompt_assets(prompt, ckpt) + save_prompt_assets(compiled_prompt, ckpt) with open(os.path.join(ckpt, PROMPT_DIR, "sid_space.json")) as f: space = json.load(f) self.assertEqual(space["codebook"], [4, 4, 4]) self.assertEqual(space["level_offsets"], [0, 4, 8]) - self.assertEqual(space["band_lo"][0], prompt.sid_space.base_vocab) + self.assertEqual(space["band_lo"][0], compiled_prompt.sid_space.base_vocab_size) # every declared field survives, so serving needs no tzrec code self.assertEqual( - set(space), {f.name for f in dataclasses.fields(prompt.sid_space)} + set(space), + {f.name for f in dataclasses.fields(compiled_prompt.sid_space)}, ) def test_a_changed_codebook_is_fatal(self) -> None: @@ -87,21 +88,27 @@ def test_a_changed_codebook_is_fatal(self) -> None: def test_a_changed_template_only_warns(self) -> None: ckpt = os.path.join(self.test_dir, "model.ckpt-1") save_prompt_assets(self._compile(), ckpt) - moved = self._compile(prompt="Predict : {{hist}}") + changed_compiled_prompt = self._compile(prompt="Predict : {{hist}}") # the vocabulary is untouched, so the weights are still usable - self.assertEqual(moved.vocab_hash, read_prompt_hashes(ckpt)["vocab_hash"]) - self.assertNotEqual(moved.plan_hash, read_prompt_hashes(ckpt)["plan_hash"]) + self.assertEqual( + changed_compiled_prompt.vocab_hash, + read_prompt_hashes(ckpt)["vocab_hash"], + ) + self.assertNotEqual( + changed_compiled_prompt.plan_hash, + read_prompt_hashes(ckpt)["plan_hash"], + ) with mock.patch("tzrec.prompt.persist.logger.warning") as warning: - check_prompt_assets(moved, ckpt) + check_prompt_assets(changed_compiled_prompt, ckpt) warning.assert_called_once() def test_a_checkpoint_without_assets_only_warns(self) -> None: bare = os.path.join(self.test_dir, "model.ckpt-bare") os.makedirs(bare, exist_ok=True) self.assertIsNone(read_prompt_hashes(bare)) - prompt = self._compile() + compiled_prompt = self._compile() with mock.patch("tzrec.prompt.persist.logger.warning") as warning: - check_prompt_assets(prompt, bare) + check_prompt_assets(compiled_prompt, bare) warning.assert_called_once() def test_no_prompt_config_is_a_no_op(self) -> None: diff --git a/tzrec/prompt/plan.py b/tzrec/prompt/plan.py index 58b620643..8ae844da3 100644 --- a/tzrec/prompt/plan.py +++ b/tzrec/prompt/plan.py @@ -69,17 +69,17 @@ class ResolvedSidSpace: Three coordinate systems and the constants that convert between them: a local code in ``[0, codebook[l])``, a flat index ``level_offsets[l] + code`` - which is what the data carries, and an LM token id ``base_vocab + flat`` + which is what the data carries, and an LM token id ``base_vocab_size + flat`` which is what ``lm_head`` generates. Args: codebook: per-level vocabulary sizes. num_levels: codes per item; also the answer width. - base_vocab: tokenizer size before the SID atoms were appended. + base_vocab_size: tokenizer size before the SID tokens were appended. level_offsets: ``cumsum(codebook) - codebook``. band_lo: inclusive lower token-id bound of each level. band_hi: inclusive upper token-id bound of each level. - target_vocab: embedding rows after padding, what the LM resizes to. + target_vocab_size: embedding rows after padding, what the LM resizes to. sentinel_token_id: id reserved for projected positions, None when no slot is projected. eos_token_id: end-of-sequence id of the extended tokenizer. @@ -88,11 +88,11 @@ class ResolvedSidSpace: codebook: Tuple[int, ...] num_levels: int - base_vocab: int + base_vocab_size: int level_offsets: Tuple[int, ...] band_lo: Tuple[int, ...] band_hi: Tuple[int, ...] - target_vocab: int + target_vocab_size: int sentinel_token_id: Optional[int] eos_token_id: int pad_token_id: int diff --git a/tzrec/protos/models/prompt_model.proto b/tzrec/protos/models/prompt_model.proto index 48cfba059..a30b98645 100644 --- a/tzrec/protos/models/prompt_model.proto +++ b/tzrec/protos/models/prompt_model.proto @@ -25,9 +25,9 @@ message PromptModelConfig { optional string generated_sids_key = 4 [default = "generated_sids"]; - // MASTER weights. FP32 avoids bf16-ULP underflow of Adam's small updates; - // bf16 COMPUTE comes from mixed_precision, not from this. - optional ParamDtype param_dtype = 5 [default = FP32]; + // LM parameters. FP32 avoids bf16-ULP underflow of Adam's small updates; + // bf16 compute comes from mixed_precision, not from this. + optional ParamDtype lm_parameter_dtype = 5 [default = FP32]; } // Qwen family (Qwen2.5, Qwen3, ...). @@ -37,5 +37,5 @@ message PromptGenerativeQwen { // HF hub id or local path. Names the architecture used for every model build // and the weights loaded by init_from_pretrained at cold start; the vocabulary // is prompt_config's. - optional string hf_model_id = 2 [default = "Qwen/Qwen2.5-0.5B"]; + optional string hf_model_name_or_path = 2 [default = "Qwen/Qwen2.5-0.5B"]; } diff --git a/tzrec/protos/prompt.proto b/tzrec/protos/prompt.proto index 4528c5716..ab255387f 100644 --- a/tzrec/protos/prompt.proto +++ b/tzrec/protos/prompt.proto @@ -9,9 +9,9 @@ message PromptConfig { reserved 2, 9; reserved "asset_dir", "length_buckets"; - // BASE tokenizer JSON file. Distinct from hf_model_id, which supplies the - // model architecture and cold-start weights. - required string tokenizer = 1; + // BASE tokenizer JSON file. Distinct from hf_model_name_or_path, which + // supplies the model architecture and cold-start weights. + required string tokenizer_path = 1; // Static text between {{name}} placeholders is the prefix and suffix of // the surrounding slots; there are no per-slot text fields. @@ -65,8 +65,8 @@ message PromptProjection { message SidSpace { // Per-level sizes; codes are local 0-based in [0, codebook[l]). repeated uint32 codebook = 1; - // {i} is the flat atom index. - optional string atom_token_format = 2 [default = "<|sid_{i}|>"]; + // {i} is the flattened SID code index. + optional string token_format = 2 [default = "<|sid_{i}|>"]; optional uint32 vocab_pad_to_multiple_of = 3 [default = 128]; // SID manifest to cross-check codebook against. A mismatch is fatal: it // means the data and the decode bands disagree. diff --git a/tzrec/tests/configs/prompt_generative_qwen_mock.config b/tzrec/tests/configs/prompt_generative_qwen_mock.config index 0651228aa..ae0457efc 100644 --- a/tzrec/tests/configs/prompt_generative_qwen_mock.config +++ b/tzrec/tests/configs/prompt_generative_qwen_mock.config @@ -44,7 +44,7 @@ feature_configs { } } prompt_config { - tokenizer: "data/test/tokenizer.json" + tokenizer_path: "data/test/tokenizer.json" prompt: "History : {{hist}} . Predict :" response: "{{answer}}" sid_space { @@ -56,7 +56,7 @@ prompt_config { } model_config { prompt_generative_qwen { - hf_model_id: "Qwen/Qwen2.5-0.5B" + hf_model_name_or_path: "Qwen/Qwen2.5-0.5B" common { beam_widths: 2 beam_widths: 2 diff --git a/tzrec/tests/prompt_integration_test.py b/tzrec/tests/prompt_integration_test.py index 349246eca..92d20d67c 100644 --- a/tzrec/tests/prompt_integration_test.py +++ b/tzrec/tests/prompt_integration_test.py @@ -53,21 +53,26 @@ def setUp(self) -> None: ] cfg = PromptConfig( - tokenizer=self.tok, + tokenizer_path=self.tok, prompt="History : {{hist}} . Predict :", response="{{answer}}", ) cfg.sid_space.codebook.extend(_CODEBOOK) - self.prompt = compile_prompt(cfg, self.features, model_dir=self.test_dir) + self.compiled_prompt = compile_prompt( + cfg, self.features, model_dir=self.test_dir + ) def _model(self): model_config = ModelConfig() qwen = model_config.prompt_generative_qwen - qwen.hf_model_id = self.backbone + qwen.hf_model_name_or_path = self.backbone qwen.common.beam_widths.extend([2, 2, 2]) qwen.common.num_return_sequences = 2 return _create_model( - model_config, self.features, ["answer"], prompt=self.prompt + model_config, + self.features, + ["answer"], + compiled_prompt=self.compiled_prompt, ) def _batch(self, hist, answer): @@ -77,18 +82,18 @@ def _batch(self, hist, answer): "answer.values": torch.tensor(offset_sid_codes(answer, _CODEBOOK)), "answer.lengths": torch.tensor([len(answer)]), } - streams = assemble_into(self.prompt, parsed) + streams = assemble_into(self.compiled_prompt, parsed) batch = Batch() batch.additional_infos.update( {k: torch.from_numpy(np.asarray(v)) for k, v in streams.items()} ) return batch - def test_model_resizes_to_target_vocab(self) -> None: + def test_model_resizes_to_target_vocab_size(self) -> None: model = self._model() rows = model.lm.get_input_embeddings().weight.shape[0] - self.assertEqual(rows, self.prompt.sid_space.target_vocab) - self.assertGreater(rows, self.prompt.sid_space.band_hi[-1]) + self.assertEqual(rows, self.compiled_prompt.sid_space.target_vocab_size) + self.assertGreater(rows, self.compiled_prompt.sid_space.band_hi[-1]) def test_loss_is_finite_and_backpropagates_into_the_backbone(self) -> None: model = self._model() diff --git a/tzrec/tests/prompt_test_util.py b/tzrec/tests/prompt_test_util.py index c34d9f147..524ef26d4 100644 --- a/tzrec/tests/prompt_test_util.py +++ b/tzrec/tests/prompt_test_util.py @@ -70,19 +70,19 @@ def offset_sid_codes(codes: Sequence[Any], codebook: Sequence[int]) -> np.ndarra def assemble_into( - prompt: CompiledPrompt, + compiled_prompt: CompiledPrompt, parsed_features: Dict[str, "np.ndarray"], ) -> Dict[str, np.ndarray]: """Assemble one parsed batch with a temporary assembler. Args: - prompt: the compiled prompt. + compiled_prompt: the compiled prompt. parsed_features: ``{feature}.values`` / ``{feature}.lengths`` as the data parser emits them. Returns: The assembled streams keyed for ``additional_infos``. """ - return PromptAssembler(prompt.prompt_plan, prompt.sid_space).assemble_batch( - parsed_features - ) + return PromptAssembler( + compiled_prompt.prompt_plan, compiled_prompt.sid_space + ).assemble_batch(parsed_features) From 029bd954907c7dac9c9b7b1d758d1f3c1e7d67c8 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Fri, 14 Aug 2026 07:47:50 +0000 Subject: [PATCH 91/99] [refactor] encapsulate Qwen input padding --- tzrec/models/prompt_generative_qwen.py | 117 ++++++++------------ tzrec/models/prompt_generative_qwen_test.py | 31 ++++-- 2 files changed, 70 insertions(+), 78 deletions(-) diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 9c335365e..7a5aa75c8 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -141,14 +141,8 @@ def _forward_loss( Returns: The loss. """ - infos = batch.additional_infos - padded, mask, labels = _unpack( - embeds, - infos[PROMPT_CU_SEQLENS], - int(infos[PROMPT_MAX_SEQLEN]), - input_ids=infos[PROMPT_INPUT_IDS], - response_lengths=infos[PROMPT_RESPONSE_LENGTHS], - ignore_index=self._ignore_index, + padded, mask, labels = self._left_pad_packed_inputs( + embeds, batch, build_labels=True ) outputs = self.lm.model(inputs_embeds=padded, attention_mask=mask) @@ -173,12 +167,7 @@ def _generate(self, batch: Batch) -> torch.Tensor: ``(B, num_return, num_levels)`` local codes, best first. """ embeds = self.build_input(batch) - infos = batch.additional_infos - padded, mask, _ = _unpack( - embeds, - infos[PROMPT_CU_SEQLENS], - int(infos[PROMPT_MAX_SEQLEN]), - ) + padded, mask, _ = self._left_pad_packed_inputs(embeds, batch) space = self._prompt.sid_space tokens = dynamic_beam_search( self.lm, @@ -190,62 +179,51 @@ def _generate(self, batch: Batch) -> torch.Tensor: codes = self._tokens_to_local_codes(tokens, padded.shape[0]) return codes[:, : self._num_return_sequences, :] + def _left_pad_packed_inputs( + self, + embeds: torch.Tensor, + batch: Batch, + build_labels: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Left-pad packed prompt embeddings for the causal LM. -def _unpack( - embeds: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - input_ids: Optional[torch.Tensor] = None, - response_lengths: Optional[torch.Tensor] = None, - ignore_index: int = -100, -) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: - """Left-pad a packed batch and build response labels for training. - - Padding lives in this one adapter. ``max_seqlen`` is the collator's, not - ``lengths.max()``: deriving it here would sync the device to the host every - step, which the design forbids. - - Pads go on the left so that every row ends on a real token. Both consumers - index from the right -- the loss keeps a fixed-width suffix and decode - prefills from ``[:, -1, :]`` -- so right-padding would hand a short row its - padding instead of its answer. - - Args: - embeds: packed embeddings, ``(total_tokens, hidden)``. - cu_seqlens: row boundaries, ``(batch_size + 1,)``. - max_seqlen: the collator's padded width. - input_ids: packed token ids, or None at inference where nothing is scored. - response_lengths: number of supervised response tokens in each row. - ignore_index: label value outside the response span. + Args: + embeds: packed embeddings, ``(total_tokens, hidden)``. + batch: carries the packed prompt metadata. + build_labels: whether to build response-only training labels. - Returns: - Padded embeddings, attention mask and labels. - """ - starts = cu_seqlens[:-1] - lengths = cu_seqlens[1:] - starts - batch_size = lengths.numel() - hidden = embeds.shape[-1] - - columns = torch.arange(max_seqlen, device=embeds.device) - mask = columns[None, :] >= (max_seqlen - lengths)[:, None] - - padded = embeds.new_zeros((batch_size, max_seqlen, hidden)) - # mask selects row-major, which is how embeds and input_ids are packed - padded[mask] = embeds - if input_ids is None: - return padded, mask.long(), None - - assert response_lengths is not None - labels = torch.full( - (batch_size, max_seqlen), - ignore_index, - dtype=input_ids.dtype, - device=embeds.device, - ) - labels[mask] = input_ids - response_mask = columns[None, :] >= (max_seqlen - response_lengths)[:, None] - labels[~response_mask] = ignore_index - return padded, mask.long(), labels + Returns: + Padded embeddings, attention mask and optional labels. + """ + infos = batch.additional_infos + cu_seqlens = infos[PROMPT_CU_SEQLENS] + max_seqlen = int(infos[PROMPT_MAX_SEQLEN]) + starts = cu_seqlens[:-1] + lengths = cu_seqlens[1:] - starts + batch_size = lengths.numel() + hidden = embeds.shape[-1] + + columns = torch.arange(max_seqlen, device=embeds.device) + mask = columns[None, :] >= (max_seqlen - lengths)[:, None] + + padded = embeds.new_zeros((batch_size, max_seqlen, hidden)) + # mask selects row-major, which is how embeds and input_ids are packed + padded[mask] = embeds + if not build_labels: + return padded, mask.long(), None + + input_ids = infos[PROMPT_INPUT_IDS] + response_lengths = infos[PROMPT_RESPONSE_LENGTHS] + labels = torch.full( + (batch_size, max_seqlen), + self._ignore_index, + dtype=input_ids.dtype, + device=embeds.device, + ) + labels[mask] = input_ids + response_mask = columns[None, :] >= (max_seqlen - response_lengths)[:, None] + labels[~response_mask] = self._ignore_index + return padded, mask.long(), labels @torch.fx.wrap @@ -255,7 +233,8 @@ def _fx_wrapped_loss( """Hide the padded forward from FX. ``TrainPipelineSparseDist`` symbolically traces the model whenever a - sharded module exists, and ``_unpack`` reads ``max_seqlen`` as a host int. + sharded module exists, and ``_left_pad_packed_inputs`` reads + ``max_seqlen`` as a host int. Args: model: the model whose loss to compute. diff --git a/tzrec/models/prompt_generative_qwen_test.py b/tzrec/models/prompt_generative_qwen_test.py index 8bb1afc92..72fe095fe 100644 --- a/tzrec/models/prompt_generative_qwen_test.py +++ b/tzrec/models/prompt_generative_qwen_test.py @@ -13,10 +13,17 @@ import torch -from tzrec.models.prompt_generative_qwen import _unpack +from tzrec.datasets.utils import Batch +from tzrec.models.prompt_generative_qwen import PromptGenerativeQwen +from tzrec.prompt.assembler import ( + PROMPT_CU_SEQLENS, + PROMPT_INPUT_IDS, + PROMPT_MAX_SEQLEN, + PROMPT_RESPONSE_LENGTHS, +) -class UnpackTest(unittest.TestCase): +class LeftPadPackedInputsTest(unittest.TestCase): """The one adapter where padding lives.""" def test_packs_rows_of_different_lengths(self) -> None: @@ -25,14 +32,20 @@ def test_packs_rows_of_different_lengths(self) -> None: input_ids = torch.tensor([1, 2, 7, 8, 3, 4, 5, 7, 8]) response_lengths = torch.tensor([1, 2]) ignore = -7 + batch = Batch( + additional_infos={ + PROMPT_CU_SEQLENS: cu, + PROMPT_INPUT_IDS: input_ids, + PROMPT_MAX_SEQLEN: torch.tensor(7), + PROMPT_RESPONSE_LENGTHS: response_lengths, + } + ) + model = PromptGenerativeQwen.__new__(PromptGenerativeQwen) + torch.nn.Module.__init__(model) + model._ignore_index = ignore - padded, mask, out = _unpack( - embeds, - cu, - max_seqlen=7, - input_ids=input_ids, - response_lengths=response_lengths, - ignore_index=ignore, + padded, mask, out = model._left_pad_packed_inputs( + embeds, batch, build_labels=True ) self.assertEqual(padded.shape, (2, 7, 2)) From fc9c3c2deb43a180234404e3d367340f01eed77c Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Mon, 17 Aug 2026 07:50:35 +0000 Subject: [PATCH 92/99] [bugfix] give spliced projections the LM dtype _build_backbone casts the LM to lm_parameter_dtype, but _build_projections builds nn.Linear/MLP at the default fp32 and never casts them. index_copy requires self and source to share a dtype, so build_input raised on the first forward for any BF16 or FP16 config with a projected slot -- both dtypes the proto advertises. The spliced values now take embeds.dtype. Casting the projections themselves would put their weights in bf16, which is the Adam underflow the fp32 master-weight default exists to avoid; reading embeds.dtype also adapts under autocast. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_model.py | 2 +- tzrec/models/prompt_generative_model_test.py | 71 ++++++++++++++++++-- 2 files changed, 68 insertions(+), 5 deletions(-) diff --git a/tzrec/models/prompt_generative_model.py b/tzrec/models/prompt_generative_model.py index 02673d742..5d9f6061f 100644 --- a/tzrec/models/prompt_generative_model.py +++ b/tzrec/models/prompt_generative_model.py @@ -182,7 +182,7 @@ def build_input(self, batch: Batch) -> torch.Tensor: return embeds.index_copy( 0, batch.additional_infos[PROMPT_HOLE_POSITIONS], - torch.cat(projected_embeddings), + torch.cat(projected_embeddings).to(embeds.dtype), ) def _tokens_to_local_codes( diff --git a/tzrec/models/prompt_generative_model_test.py b/tzrec/models/prompt_generative_model_test.py index ad310cd66..ee0d95e11 100644 --- a/tzrec/models/prompt_generative_model_test.py +++ b/tzrec/models/prompt_generative_model_test.py @@ -9,22 +9,26 @@ # See the License for the specific language governing permissions and # limitations under the License. +import dataclasses import os import unittest import numpy as np import torch +from parameterized import parameterized from torchrec import KeyedJaggedTensor from transformers import AutoModelForCausalLM from tzrec.datasets.utils import BASE_DATA_GROUP, Batch from tzrec.main import _create_model +from tzrec.models.prompt_generative_model import _PARAM_DTYPE from tzrec.prompt.assembler import ( PROMPT_HOLE_POSITIONS, PROMPT_INPUT_IDS, ) from tzrec.prompt.compile import compile_prompt from tzrec.protos.model_pb2 import ModelConfig +from tzrec.protos.models.prompt_model_pb2 import PromptModelConfig from tzrec.protos.prompt_pb2 import PromptConfig from tzrec.tests.prompt_test_util import ( assemble_into, @@ -33,7 +37,11 @@ offset_sid_codes, ) from tzrec.utils.state_dict_util import init_parameters -from tzrec.utils.test_util import create_tiny_causal_lm, make_test_dir +from tzrec.utils.test_util import ( + create_tiny_causal_lm, + make_test_dir, + parameterized_name_func, +) _CODEBOOK = [4, 4, 4] _WORDS = ["History", "Predict", ":", ".", "", "<|im_end|>"] @@ -67,6 +75,7 @@ def setUp(self) -> None: self.compiled_prompt = self._compile(self.features) def _compile(self, features, template="History : {{hist}} . Predict :", **kwargs): + kwargs.setdefault("response", "{{answer}}") cfg = PromptConfig(tokenizer_path=self.tok, prompt=template, **kwargs) cfg.sid_space.codebook.extend(_CODEBOOK) return compile_prompt(cfg, features, model_dir=self.test_dir) @@ -77,12 +86,15 @@ def _model( compiled_prompt=-1, beam_widths=(2, 2, 2), num_return_sequences=2, + lm_parameter_dtype=None, ): model_config = ModelConfig() qwen = model_config.prompt_generative_qwen qwen.hf_model_name_or_path = self.backbone qwen.common.beam_widths.extend(beam_widths) qwen.common.num_return_sequences = num_return_sequences + if lm_parameter_dtype is not None: + qwen.common.lm_parameter_dtype = lm_parameter_dtype return _create_model( model_config, self.features if features is None else features, @@ -122,9 +134,9 @@ def test_rejects_a_model_built_without_a_prompt(self) -> None: self._model(compiled_prompt=None) def test_rejects_a_prompt_that_declares_no_sid_space(self) -> None: - cfg = PromptConfig(tokenizer_path=self.tok, prompt="History : {{hist}} .") - compiled_prompt = compile_prompt(cfg, self.features, model_dir=self.test_dir) - self.assertIsNone(compiled_prompt.sid_space) + # compile_prompt refuses this config, so the model precondition is + # reachable only by constructing a prompt directly + compiled_prompt = dataclasses.replace(self.compiled_prompt, sid_space=None) with self.assertRaisesRegex(ValueError, "declares no sid_space"): self._model(compiled_prompt=compiled_prompt) @@ -195,6 +207,57 @@ def test_projected_slot_overwrites_sentinels_and_backpropagates(self) -> None: proj = next(iter(model.projections.values())) self.assertIsNotNone(proj.head.weight.grad) + @parameterized.expand( + [[PromptModelConfig.BF16], [PromptModelConfig.FP16]], + name_func=parameterized_name_func, + ) + def test_projected_slot_follows_a_narrow_lm_dtype(self, lm_parameter_dtype) -> None: + features = [ + create_prompt_feature(_HIST), + create_prompt_feature(_ANSWER), + create_prompt_feature(_projected("prof", 8)), + ] + compiled_prompt = self._compile( + features, + template="History : {{hist}} . Predict {{prof}} :", + response="{{answer}}", + ) + model = self._model( + features=features, + compiled_prompt=compiled_prompt, + lm_parameter_dtype=lm_parameter_dtype, + ) + init_parameters(model, device=torch.device("cpu")) + batch = self._batch( + { + "hist.values": torch.tensor( + offset_sid_codes([0, 1, 2], _CODEBOOK) + ).reshape(-1, 1), + "hist.lengths": torch.tensor([3]), + "answer.values": torch.tensor(offset_sid_codes([1, 2, 3], _CODEBOOK)), + "answer.lengths": torch.tensor([3]), + "prof.values": torch.tensor([5, 9]), + "prof.lengths": torch.tensor([2]), + }, + compiled_prompt=compiled_prompt, + sparse=KeyedJaggedTensor.from_lengths_sync( + keys=["prof"], + values=torch.tensor([5, 9]), + lengths=torch.tensor([2]), + ), + ) + + embeds = model.build_input(batch) + self.assertIs(embeds.dtype, _PARAM_DTYPE[lm_parameter_dtype]) + + loss = model.predict(batch)["loss"] + self.assertTrue(bool(torch.isfinite(loss))) + loss.backward() + # the projection keeps fp32 masters, so only the spliced values convert + proj = next(iter(model.projections.values())) + self.assertIs(proj.head.weight.dtype, torch.float32) + self.assertGreater(float(proj.head.weight.grad.abs().sum()), 0.0) + def test_beam_config_uses_final_capped_capacity(self) -> None: with self.assertRaisesRegex(ValueError, "final capped beam width \\(4\\)"): self._model( From 595c226dcda3b47199b9d546edef29111e1916b9 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Mon, 17 Aug 2026 07:51:08 +0000 Subject: [PATCH 93/99] [bugfix] require sid_space and response, and reject a placeholderless token_format Three prompt_configs compiled cleanly and then misbehaved. No response left response_segments empty, so logits_suffix_len collapsed to 0 + 1 and the loss ran over a window holding one ignored position: nan, with no diagnostic. No sid_space produced a CompiledPrompt no model can consume, since every prompt-native model needs one to extend its embedding table and to decode. Both are now required at compile, each reporting its own field rather than a downstream symptom. A token_format with no {i} rendered every SID token as the same string, so add_special_tokens deduplicated them and the tokenizer gained one row where the bands assume sum(codebook) -- every SID above the first then addressed a token id that was never assigned. Requiring sid_space made two checks unreachable: the unbounded-response guard, because an INLINE response slot now always derives STATIC from the codebook, and the response-width check, which compared num_levels against a value just set from num_levels. Both are removed, along with the cfg and sid_space parameters they left unused on _validate. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/prompt/compile.py | 52 +++++++++++++++++------------------- tzrec/prompt/compile_test.py | 52 ++++++++++++++++++++++++++++-------- tzrec/prompt/persist_test.py | 10 +++++-- 3 files changed, 74 insertions(+), 40 deletions(-) diff --git a/tzrec/prompt/compile.py b/tzrec/prompt/compile.py index e65d2bf48..c9bf0189e 100644 --- a/tzrec/prompt/compile.py +++ b/tzrec/prompt/compile.py @@ -128,6 +128,13 @@ def _derive_slot_layout( def _render_sid_tokens(sid_space: SidSpace) -> List[str]: """Render the SID tokens, one per flat index.""" fmt = sid_space.token_format + if "{i}" not in fmt: + raise ValueError( + f"sid_space.token_format [{fmt}] has no '{{i}}' placeholder, so " + f"every SID token would render the same string and the tokenizer " + f"would gain one row where the bands assume " + f"{sum(sid_space.codebook)}." + ) return [fmt.replace("{i}", str(i)) for i in range(sum(sid_space.codebook))] @@ -141,10 +148,14 @@ def _read_manifest_codebook(path: str) -> List[int]: def _build_sid_space( cfg: PromptConfig, tok: Tokenizer, base_vocab_size: int, has_projection: bool -) -> Optional[ResolvedSidSpace]: +) -> ResolvedSidSpace: """Extend the tokenizer with SID tokens and resolve the token space.""" if not cfg.HasField("sid_space"): - return None + raise ValueError( + "prompt_config.sid_space is required: it declares the codebook the " + "SID vocabulary is built from, and every prompt-native model needs " + "one to extend its embedding table and to decode." + ) space = cfg.sid_space codebook = [int(c) for c in space.codebook] if not codebook: @@ -318,9 +329,7 @@ def compile_prompt( fill_mode = fill_modes_by_slot_name[name] levels = ( sid_space.num_levels - if sid_space is not None - and name in response_slot_names - and fill_mode is FillMode.INLINE + if name in response_slot_names and fill_mode is FillMode.INLINE else None ) segs[name] = SlotSeg( @@ -355,7 +364,7 @@ def compile_prompt( static_prefix_len=_static_prefix_len(body), projected_slots=projected, ) - _validate(cfg, plan, sid_space) + _validate(plan) return CompiledPrompt( sid_space=sid_space, @@ -364,7 +373,10 @@ def compile_prompt( tokenizer_dir=tokenizer_dir, vocab_hash=_hash(sid_space, tok.to_str()), plan_hash=_hash( - sid_space, plan, sorted(projection_plan.projections), tok.to_str() + sid_space, + plan, + sorted(projection_plan.projections), + tok.to_str(), ), ) @@ -471,9 +483,7 @@ def _static_prefix_len(segments: Sequence[Segment]) -> int: return total -def _validate( - cfg: PromptConfig, plan: PromptPlan, sid_space: Optional[ResolvedSidSpace] -) -> None: +def _validate(plan: PromptPlan) -> None: """Apply the checks that need the whole plan.""" if plan.max_length and plan.max_total_length is not None: if plan.max_total_length > plan.max_length: @@ -500,27 +510,15 @@ def _validate( ) break variable_slot_seen = True - if plan.response_segments and plan.logits_suffix_len is None: + if not plan.response_segments: raise ValueError( - "the response has an unbounded slot, so the supervised logits " - "window cannot be bounded. A decoder-only model would then " - "materialize logits for every position, which is (batch x length x " - "vocab) and will not fit. Give the response slot a fixed width." + "prompt_config.response is required: it defines the supervised " + "span, and without it the loss window collapses to one ignored " + "position and the loss is nan. Inference drops the response by " + "mode, so a predict-only run keeps it declared." ) if plan.static_prefix_len == 0: logger.warning( "static_prefix_len is 0: no leading run of the prompt is " "request-invariant, so a serving prefix cache can share nothing." ) - if sid_space is not None and cfg.HasField("response"): - answer = [s for s in plan.response_segments if isinstance(s, SlotSeg)] - for seg in answer: - if ( - seg.width.kind is WidthKind.STATIC - and seg.width.num_positions != sid_space.num_levels - ): - raise ValueError( - f"response slot [{seg.name}] is " - f"{seg.width.num_positions} positions but the codebook has " - f"{sid_space.num_levels} levels." - ) diff --git a/tzrec/prompt/compile_test.py b/tzrec/prompt/compile_test.py index 6a9fb18dc..d8895f2cd 100644 --- a/tzrec/prompt/compile_test.py +++ b/tzrec/prompt/compile_test.py @@ -27,10 +27,11 @@ ) from tzrec.utils.test_util import make_test_dir -_WORDS = ["History", "Profile", "Predict", ":", ".", "", "<|im_end|>"] +_WORDS = ["History", "Profile", "Predict", ":", ".", "Histor0", "", "<|im_end|>"] _HIST = 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }' +_ANSWER = 'sequence_raw_feature { feature_name: "answer" expression: "item:answer" }' _PROF = ( 'sequence_id_feature { feature_name: "prof" expression: "user:prof" ' "num_buckets: 768 embedding_dim: 16 sequence_length: 4 }" @@ -46,10 +47,14 @@ def setUp(self) -> None: ) def _config(self, **kwargs) -> PromptConfig: + kwargs.setdefault("response", "{{answer}}") cfg = PromptConfig(tokenizer_path=self.tok_path, **kwargs) return cfg def _compile(self, cfg, features): + named = {f.name for f in features} + if "{{answer}}" in cfg.response and "answer" not in named: + features = list(features) + [create_prompt_feature(_ANSWER)] return compile_prompt(cfg, features, model_dir=self.test_dir) def test_sid_space_resolves_offsets_and_bands(self) -> None: @@ -176,7 +181,8 @@ def test_rejects_a_projection_on_an_inline_slot(self) -> None: def test_sid_tokens_absent_from_the_base_tokenizer(self) -> None: cfg = self._config(prompt="X : {{hist}}") cfg.sid_space.codebook.extend([4]) - cfg.sid_space.token_format = "History" + # renders Histor0..Histor3, and Histor0 is already in the base vocab + cfg.sid_space.token_format = "Histor{i}" with self.assertRaisesRegex(ValueError, "already in the base tokenizer"): self._compile(cfg, [create_prompt_feature(_HIST)]) @@ -218,16 +224,40 @@ def test_response_slot_must_be_inline(self) -> None: cfg, [create_prompt_feature(_HIST), create_prompt_feature(_PROF)] ) - def test_unbounded_response_is_rejected(self) -> None: - # with no sid_space the response has no codebook-derived width, so the - # supervised window is unbounded and the logits would cover every - # position + def test_missing_sid_space_is_rejected(self) -> None: + # the SID vocabulary is what gives the response its codebook-derived + # width; without it no prompt-native model can be built at all cfg = self._config(prompt="History : {{hist}}", response="{{answer}}") - answer = create_prompt_feature( - 'sequence_raw_feature { feature_name: "answer" expression: "item:answer" }' - ) - with self.assertRaisesRegex(ValueError, "window cannot be bounded"): - self._compile(cfg, [create_prompt_feature(_HIST), answer]) + cfg.ClearField("sid_space") + with self.assertRaisesRegex(ValueError, "sid_space is required"): + self._compile(cfg, [create_prompt_feature(_HIST)]) + + def test_token_format_without_a_placeholder_is_rejected(self) -> None: + # without {i} every SID token renders the same string, so the tokenizer + # gains one row while the bands assume sum(codebook) + cfg = self._config(prompt="History : {{hist}}", response="{{answer}}") + cfg.sid_space.codebook.extend([4, 4, 4]) + cfg.sid_space.token_format = "<|sid|>" + with self.assertRaisesRegex(ValueError, "has no '{i}' placeholder"): + self._compile(cfg, [create_prompt_feature(_HIST)]) + + def test_a_custom_token_format_with_a_placeholder_compiles(self) -> None: + cfg = self._config(prompt="History : {{hist}}", response="{{answer}}") + cfg.sid_space.codebook.extend([4, 4, 4]) + cfg.sid_space.token_format = "C{i}" + compiled = self._compile(cfg, [create_prompt_feature(_HIST)]) + + space = compiled.sid_space + self.assertEqual(space.band_hi[-1] - space.band_lo[0] + 1, 12) + + def test_missing_response_is_rejected(self) -> None: + # without a response the supervised window collapses to one ignored + # position and the loss is nan + cfg = self._config(prompt="History : {{hist}}", response="") + cfg.sid_space.codebook.extend([4, 4, 4]) + cfg.ClearField("response") + with self.assertRaisesRegex(ValueError, "response is required"): + self._compile(cfg, [create_prompt_feature(_HIST)]) def test_a_grouped_feature_inherits_the_group_cap(self) -> None: # a SequenceFeature member never sets its own sequence_length; the cap diff --git a/tzrec/prompt/persist_test.py b/tzrec/prompt/persist_test.py index 08dc6c22b..01fa37f7c 100644 --- a/tzrec/prompt/persist_test.py +++ b/tzrec/prompt/persist_test.py @@ -42,11 +42,17 @@ def setUp(self) -> None: self.features = [ create_prompt_feature( 'sequence_raw_feature { feature_name: "hist" expression: "user:hist" }' - ) + ), + create_prompt_feature( + 'sequence_raw_feature { feature_name: "answer" ' + 'expression: "item:answer" }' + ), ] def _compile(self, codebook=(4, 4, 4), prompt="History : {{hist}}"): - cfg = PromptConfig(tokenizer_path=self.tok_path, prompt=prompt) + cfg = PromptConfig( + tokenizer_path=self.tok_path, prompt=prompt, response="{{answer}}" + ) cfg.sid_space.codebook.extend(codebook) return compile_prompt(cfg, self.features, model_dir=self.test_dir) From 47b78443a4ce8de1185f15053b432078f88b7aef Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Mon, 17 Aug 2026 07:51:20 +0000 Subject: [PATCH 94/99] [bugfix] hash projection bodies in plan_hash plan_hash passed sorted(projection_plan.projections) to the digest, and iterating a dict yields its keys, so only the projection module ids were covered. Changing an MLP body while keeping its projection_name produced an identical hash: a slot whose projection went from hidden_units [16] to [256, 128] restored against an old checkpoint with no warning, though the parameter shapes differ. sorted(...) is kept for a stable order across runs. The projection body is the one part of the plan that determines tensor shapes, so it is exactly what the restore guard exists to catch, and the docstring already claimed plan_hash covered it. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/prompt/compile.py | 4 +++- tzrec/prompt/persist_test.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/tzrec/prompt/compile.py b/tzrec/prompt/compile.py index c9bf0189e..d6caf4b63 100644 --- a/tzrec/prompt/compile.py +++ b/tzrec/prompt/compile.py @@ -375,7 +375,9 @@ def compile_prompt( plan_hash=_hash( sid_space, plan, - sorted(projection_plan.projections), + # items(), not the bare dict: iterating a dict yields only its keys, + # which would leave a changed projection body out of the digest + sorted(projection_plan.projections.items()), tok.to_str(), ), ) diff --git a/tzrec/prompt/persist_test.py b/tzrec/prompt/persist_test.py index 01fa37f7c..162c47efe 100644 --- a/tzrec/prompt/persist_test.py +++ b/tzrec/prompt/persist_test.py @@ -91,6 +91,38 @@ def test_a_changed_codebook_is_fatal(self) -> None: with self.assertRaisesRegex(ValueError, "does not match checkpoint"): check_prompt_assets(self._compile(codebook=(8, 8, 8)), ckpt) + def test_a_changed_projection_body_warns(self) -> None: + # iterating a dict yields only its keys, so a body change must not be + # able to hide behind an unchanged projection name + def compile_with(hidden_units): + cfg = PromptConfig( + tokenizer_path=self.tok_path, + prompt="History : {{hist}} {{prof}}", + response="{{answer}}", + ) + cfg.sid_space.codebook.extend((4, 4, 4)) + slot = cfg.slots.add(name="prof") + slot.feature_names.append("prof") + slot.projection.mlp.hidden_units.extend(hidden_units) + features = self.features + [ + create_prompt_feature( + 'sequence_id_feature { feature_name: "prof" ' + 'expression: "user:prof" num_buckets: 16 embedding_dim: 8 ' + "sequence_length: 2 }" + ) + ] + return compile_prompt(cfg, features, model_dir=self.test_dir) + + ckpt = os.path.join(self.test_dir, "model.ckpt-proj") + save_prompt_assets(compile_with([16]), ckpt) + widened = compile_with([256, 128]) + # only the projection changed, so the vocabulary is still usable + self.assertEqual(widened.vocab_hash, read_prompt_hashes(ckpt)["vocab_hash"]) + self.assertNotEqual(widened.plan_hash, read_prompt_hashes(ckpt)["plan_hash"]) + with mock.patch("tzrec.prompt.persist.logger.warning") as warning: + check_prompt_assets(widened, ckpt) + warning.assert_called_once() + def test_a_changed_template_only_warns(self) -> None: ckpt = os.path.join(self.test_dir, "model.ckpt-1") save_prompt_assets(self._compile(), ckpt) From ac05d040cab160ea4c08f23545e29a326716d27c Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Mon, 17 Aug 2026 08:56:29 +0000 Subject: [PATCH 95/99] [bugfix] close two holes in the prompt checkpoint guard plan_hash passed sorted(projection_plan.projections) to the digest, which iterates the dict and so covered only the module ids. Adding items() in the previous commit covered the bodies but left slot_to_module out, so two slots could swap projection_name with matching bodies and still hash the same: keys load, and each slot receives the other's learned weights. Both mappings are now hashed. check_prompt_assets treated a checkpoint with no recorded hashes as a warning and restored anyway. CheckpointManager.save swallows rank-0 asset-write failures, so that combination silently disables the vocabulary guard in exactly the case it exists for -- decode bands addressing rows the weights never learned. Absent assets are now fatal. No legacy bypass: this stack has never been released, so no checkpoint predates the assets. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/prompt/compile.py | 3 +++ tzrec/prompt/persist.py | 14 ++++++++---- tzrec/prompt/persist_test.py | 43 +++++++++++++++++++++++++++++++----- 3 files changed, 51 insertions(+), 9 deletions(-) diff --git a/tzrec/prompt/compile.py b/tzrec/prompt/compile.py index d6caf4b63..ed481097e 100644 --- a/tzrec/prompt/compile.py +++ b/tzrec/prompt/compile.py @@ -378,6 +378,9 @@ def compile_prompt( # items(), not the bare dict: iterating a dict yields only its keys, # which would leave a changed projection body out of the digest sorted(projection_plan.projections.items()), + # routing too: slots can swap projection_name with matching bodies, + # which loads each slot with the other's learned weights + sorted(projection_plan.slot_to_module.items()), tok.to_str(), ), ) diff --git a/tzrec/prompt/persist.py b/tzrec/prompt/persist.py index 464affd1c..2f458b11d 100644 --- a/tzrec/prompt/persist.py +++ b/tzrec/prompt/persist.py @@ -103,22 +103,28 @@ def check_prompt_assets( A ``vocab_hash`` mismatch is fatal: the decode bands would point at token ranges the weights never learned, which produces plausible output rather than an error. A ``plan_hash`` mismatch only reshapes the prompt, so it - warns. + warns. Absent assets are fatal too -- restoring unchecked is the one case + the guard exists to prevent. Args: compiled_prompt: the freshly compiled prompt, or None when the pipeline declares no prompt_config. ckpt_dir: the checkpoint being restored. + + Raises: + ValueError: if the checkpoint records no prompt assets, or its + ``vocab_hash`` disagrees with the compiled prompt. """ if compiled_prompt is None: return recorded = read_prompt_hashes(ckpt_dir) if recorded is None: - logger.warning( + raise ValueError( f"checkpoint [{ckpt_dir}] records no prompt assets, so its " - f"vocabulary cannot be checked against the current prompt_config." + f"vocabulary cannot be checked against the current prompt_config. " + f"Restoring unchecked risks decode bands that address rows these " + f"weights never learned, so this is fatal rather than a warning." ) - return if recorded.get("vocab_hash") != compiled_prompt.vocab_hash: raise ValueError( diff --git a/tzrec/prompt/persist_test.py b/tzrec/prompt/persist_test.py index 162c47efe..5694299a1 100644 --- a/tzrec/prompt/persist_test.py +++ b/tzrec/prompt/persist_test.py @@ -123,6 +123,39 @@ def compile_with(hidden_units): check_prompt_assets(widened, ckpt) warning.assert_called_once() + def test_swapped_projection_routing_warns(self) -> None: + # identical bodies, so only slot_to_module differs -- without it each + # slot would silently load the other's learned weights + def compile_with(pa_module, pb_module): + cfg = PromptConfig( + tokenizer_path=self.tok_path, + prompt="History : {{hist}} {{pa}} {{pb}}", + response="{{answer}}", + ) + cfg.sid_space.codebook.extend((4, 4, 4)) + features = list(self.features) + for name, module_id in (("pa", pa_module), ("pb", pb_module)): + slot = cfg.slots.add(name=name, projection_name=module_id) + slot.feature_names.append(name) + slot.projection.mlp.hidden_units.extend([16]) + features.append( + create_prompt_feature( + f'sequence_id_feature {{ feature_name: "{name}" ' + f'expression: "user:{name}" num_buckets: 16 ' + "embedding_dim: 8 sequence_length: 2 }" + ) + ) + return compile_prompt(cfg, features, model_dir=self.test_dir) + + ckpt = os.path.join(self.test_dir, "model.ckpt-route") + save_prompt_assets(compile_with("X", "Y"), ckpt) + swapped = compile_with("Y", "X") + self.assertEqual(swapped.vocab_hash, read_prompt_hashes(ckpt)["vocab_hash"]) + self.assertNotEqual(swapped.plan_hash, read_prompt_hashes(ckpt)["plan_hash"]) + with mock.patch("tzrec.prompt.persist.logger.warning") as warning: + check_prompt_assets(swapped, ckpt) + warning.assert_called_once() + def test_a_changed_template_only_warns(self) -> None: ckpt = os.path.join(self.test_dir, "model.ckpt-1") save_prompt_assets(self._compile(), ckpt) @@ -140,14 +173,14 @@ def test_a_changed_template_only_warns(self) -> None: check_prompt_assets(changed_compiled_prompt, ckpt) warning.assert_called_once() - def test_a_checkpoint_without_assets_only_warns(self) -> None: + def test_a_checkpoint_without_assets_is_fatal(self) -> None: + # CheckpointManager.save swallows asset-write failures, so a bare + # checkpoint must not restore with the vocabulary guard disabled bare = os.path.join(self.test_dir, "model.ckpt-bare") os.makedirs(bare, exist_ok=True) self.assertIsNone(read_prompt_hashes(bare)) - compiled_prompt = self._compile() - with mock.patch("tzrec.prompt.persist.logger.warning") as warning: - check_prompt_assets(compiled_prompt, bare) - warning.assert_called_once() + with self.assertRaisesRegex(ValueError, "records no prompt assets"): + check_prompt_assets(self._compile(), bare) def test_no_prompt_config_is_a_no_op(self) -> None: with mock.patch("tzrec.prompt.persist.logger.warning") as warning: From 2a8a0855ff9dbc23423886edb1eaee8bf2f20a2a Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Mon, 17 Aug 2026 08:56:29 +0000 Subject: [PATCH 96/99] [bugfix] build prompt embeddings outside both FX leaves predict built the input embeddings outside the leaf on the training path so TrainPipelineSparseDist could see the sharded embedding module and prefetch it, but the inference path passed the whole batch into _fx_wrapped_generate, which hid its build_input call. PredictPipelineSparseDist therefore saw no sharded-module node and the projected-slot lookup ran synchronously inside the leaf. Both paths now build embeddings before the wrapper and pass them in, so only the padding and the LM -- the parts that branch on host ints -- stay hidden. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_qwen.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 7a5aa75c8..02dbe2930 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -119,12 +119,13 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: Returns: The loss when training, the decoded SIDs otherwise. """ + # the embedding lookup stays outside the leaf on both paths, so the + # pipeline still sees the sharded module and can prefetch it; only the + # padding and the LM, which branch on host ints, are hidden + embeds = self.build_input(batch) if self.is_inference: - return {self._generated_sids_key: _fx_wrapped_generate(self, batch)} - # the embedding lookup stays traceable so the train pipeline can still - # see the sharded module and prefetch it; only the padding and the LM, - # which read the collator's width as a host int, are hidden. - return _fx_wrapped_loss(self, self.build_input(batch), batch) + return {self._generated_sids_key: _fx_wrapped_generate(self, embeds, batch)} + return _fx_wrapped_loss(self, embeds, batch) def _forward_loss( self, embeds: torch.Tensor, batch: Batch @@ -157,16 +158,16 @@ def _forward_loss( ) return {"loss": loss} - def _generate(self, batch: Batch) -> torch.Tensor: + def _generate(self, embeds: torch.Tensor, batch: Batch) -> torch.Tensor: """Beam-search the SID answer. Args: - batch: carries the packed prompt and the collator's width. + embeds: the assembled prompt embeddings, packed. + batch: carries the packed prompt and the padded width. Returns: ``(B, num_return, num_levels)`` local codes, best first. """ - embeds = self.build_input(batch) padded, mask, _ = self._left_pad_packed_inputs(embeds, batch) space = self._prompt.sid_space tokens = dynamic_beam_search( @@ -248,7 +249,9 @@ def _fx_wrapped_loss( @torch.fx.wrap -def _fx_wrapped_generate(model: "PromptGenerativeQwen", batch: Batch) -> torch.Tensor: +def _fx_wrapped_generate( + model: "PromptGenerativeQwen", embeds: torch.Tensor, batch: Batch +) -> torch.Tensor: """Hide the decode loop from FX. ``PredictPipelineSparseDist`` FX-traces the model, and beam decode reads @@ -257,9 +260,10 @@ def _fx_wrapped_generate(model: "PromptGenerativeQwen", batch: Batch) -> torch.T Args: model: the model whose decode loop to run. + embeds: the assembled prompt embeddings, packed. batch: the batch to decode. Returns: The decoded local codes. """ - return model._generate(batch) + return model._generate(embeds, batch) From 2b530c0197530a4dd54864ba127ee22bf48653ab Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Mon, 17 Aug 2026 09:03:53 +0000 Subject: [PATCH 97/99] [chore] cut the added inline comments to one line each AGENTS.md asks for one short line of non-obvious "why"; several of the comments added over the review round ran to two or three. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/models/prompt_generative_qwen.py | 4 +--- tzrec/prompt/compile.py | 6 ++---- tzrec/prompt/compile_test.py | 9 +++------ tzrec/prompt/persist_test.py | 9 +++------ 4 files changed, 9 insertions(+), 19 deletions(-) diff --git a/tzrec/models/prompt_generative_qwen.py b/tzrec/models/prompt_generative_qwen.py index 02dbe2930..ae554b8b1 100644 --- a/tzrec/models/prompt_generative_qwen.py +++ b/tzrec/models/prompt_generative_qwen.py @@ -119,9 +119,7 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: Returns: The loss when training, the decoded SIDs otherwise. """ - # the embedding lookup stays outside the leaf on both paths, so the - # pipeline still sees the sharded module and can prefetch it; only the - # padding and the LM, which branch on host ints, are hidden + # outside the leaf on both paths, so the pipeline can prefetch it embeds = self.build_input(batch) if self.is_inference: return {self._generated_sids_key: _fx_wrapped_generate(self, embeds, batch)} diff --git a/tzrec/prompt/compile.py b/tzrec/prompt/compile.py index ed481097e..b933c7b5c 100644 --- a/tzrec/prompt/compile.py +++ b/tzrec/prompt/compile.py @@ -375,11 +375,9 @@ def compile_prompt( plan_hash=_hash( sid_space, plan, - # items(), not the bare dict: iterating a dict yields only its keys, - # which would leave a changed projection body out of the digest + # items(), not the bare dict: iterating one yields only its keys sorted(projection_plan.projections.items()), - # routing too: slots can swap projection_name with matching bodies, - # which loads each slot with the other's learned weights + # routing too: matching bodies can still be wired to other slots sorted(projection_plan.slot_to_module.items()), tok.to_str(), ), diff --git a/tzrec/prompt/compile_test.py b/tzrec/prompt/compile_test.py index d8895f2cd..34988a076 100644 --- a/tzrec/prompt/compile_test.py +++ b/tzrec/prompt/compile_test.py @@ -225,16 +225,14 @@ def test_response_slot_must_be_inline(self) -> None: ) def test_missing_sid_space_is_rejected(self) -> None: - # the SID vocabulary is what gives the response its codebook-derived - # width; without it no prompt-native model can be built at all + # the response width is codebook-derived, so sid_space must exist cfg = self._config(prompt="History : {{hist}}", response="{{answer}}") cfg.ClearField("sid_space") with self.assertRaisesRegex(ValueError, "sid_space is required"): self._compile(cfg, [create_prompt_feature(_HIST)]) def test_token_format_without_a_placeholder_is_rejected(self) -> None: - # without {i} every SID token renders the same string, so the tokenizer - # gains one row while the bands assume sum(codebook) + # without {i} every token renders alike: one row, not sum(codebook) cfg = self._config(prompt="History : {{hist}}", response="{{answer}}") cfg.sid_space.codebook.extend([4, 4, 4]) cfg.sid_space.token_format = "<|sid|>" @@ -251,8 +249,7 @@ def test_a_custom_token_format_with_a_placeholder_compiles(self) -> None: self.assertEqual(space.band_hi[-1] - space.band_lo[0] + 1, 12) def test_missing_response_is_rejected(self) -> None: - # without a response the supervised window collapses to one ignored - # position and the loss is nan + # no response collapses the window to one ignored position: nan loss cfg = self._config(prompt="History : {{hist}}", response="") cfg.sid_space.codebook.extend([4, 4, 4]) cfg.ClearField("response") diff --git a/tzrec/prompt/persist_test.py b/tzrec/prompt/persist_test.py index 5694299a1..df320494a 100644 --- a/tzrec/prompt/persist_test.py +++ b/tzrec/prompt/persist_test.py @@ -92,8 +92,7 @@ def test_a_changed_codebook_is_fatal(self) -> None: check_prompt_assets(self._compile(codebook=(8, 8, 8)), ckpt) def test_a_changed_projection_body_warns(self) -> None: - # iterating a dict yields only its keys, so a body change must not be - # able to hide behind an unchanged projection name + # a body change must not hide behind an unchanged projection name def compile_with(hidden_units): cfg = PromptConfig( tokenizer_path=self.tok_path, @@ -124,8 +123,7 @@ def compile_with(hidden_units): warning.assert_called_once() def test_swapped_projection_routing_warns(self) -> None: - # identical bodies, so only slot_to_module differs -- without it each - # slot would silently load the other's learned weights + # identical bodies, so only slot_to_module differs def compile_with(pa_module, pb_module): cfg = PromptConfig( tokenizer_path=self.tok_path, @@ -174,8 +172,7 @@ def test_a_changed_template_only_warns(self) -> None: warning.assert_called_once() def test_a_checkpoint_without_assets_is_fatal(self) -> None: - # CheckpointManager.save swallows asset-write failures, so a bare - # checkpoint must not restore with the vocabulary guard disabled + # save() swallows asset-write failures, so a bare checkpoint must fail bare = os.path.join(self.test_dir, "model.ckpt-bare") os.makedirs(bare, exist_ok=True) self.assertIsNone(read_prompt_hashes(bare)) From a22d4406bcfa94670df0b3082e8add5ed5defcf5 Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Mon, 17 Aug 2026 11:43:31 +0000 Subject: [PATCH 98/99] [refactor] drop the prompt proto reserved markers and close the tag gaps asset_dir, length_buckets and drop_if_empty were removed earlier and their tags reserved. Nothing has been released, so no serialized PromptConfig exists to protect and the markers guard nothing. Removing them frees the tags, and the remaining fields renumber contiguously: PromptConfig 1-7, PromptSlot 1-4. PromptProjection keeps mlp = 10; that gap belongs to the oneof body block and leaves room for sibling body types. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/protos/prompt.proto | 22 ++++++++-------------- 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/tzrec/protos/prompt.proto b/tzrec/protos/prompt.proto index ab255387f..b3290f0ad 100644 --- a/tzrec/protos/prompt.proto +++ b/tzrec/protos/prompt.proto @@ -6,39 +6,33 @@ import "tzrec/protos/module.proto"; // Rendering of an LM prompt. Peer of data_config and model_config: extraction // stays in feature_configs, this owns order, literal text and composition. message PromptConfig { - reserved 2, 9; - reserved "asset_dir", "length_buckets"; - // BASE tokenizer JSON file. Distinct from hf_model_name_or_path, which // supplies the model architecture and cold-start weights. required string tokenizer_path = 1; // Static text between {{name}} placeholders is the prefix and suffix of // the surrounding slots; there are no per-slot text fields. - required string prompt = 3; + required string prompt = 2; // Supervised answer. Defines the loss span; absent at inference. - optional string response = 4; + optional string response = 3; // A placeholder resolves to a slot with that name, else to an implicit // single-feature slot named after it. - repeated PromptSlot slots = 5; + repeated PromptSlot slots = 4; // Required when any slot renders SIDs. - optional SidSpace sid_space = 6; + optional SidSpace sid_space = 5; // Validation ceiling, not a truncation trigger: an over-long row is an // error, never truncated. - optional uint32 max_length = 7 [default = 0]; + optional uint32 max_length = 6 [default = 0]; // Reserves a position filled by a projected slot. Materialized only when // at least one slot is PROJECTED. - optional string sentinel_token = 8 [default = "<|pg_hole|>"]; + optional string sentinel_token = 7 [default = "<|pg_hole|>"]; } message PromptSlot { - reserved 3; - reserved "drop_if_empty"; - // {{name}} in the template; also the derived feature_group name when the // slot is PROJECTED. required string name = 1; @@ -47,10 +41,10 @@ message PromptSlot { repeated string feature_names = 2; // Reconciles this slot's width with the LM hidden size. Illegal on an // INLINE slot. Carries no dimensions: the model resolves them. - optional PromptProjection projection = 4; + optional PromptProjection projection = 3; // Weight-sharing key across slots. Derived groups dedupe automatically; // modules share only on request. - optional string projection_name = 5; + optional string projection_name = 4; } // An optional body plus a final bare Linear to the LM hidden size. The final From ab99b1cb25e657c71aa33afabbc13f489e4df28f Mon Sep 17 00:00:00 2001 From: root <597191244@qq.com> Date: Tue, 18 Aug 2026 03:18:09 +0000 Subject: [PATCH 99/99] [refactor] mark prompt sid_space required Every prompt-native model in this version emits SIDs, so a config without sid_space is never valid and the declaration now says so. Enforcement still comes from compile_prompt: text_format does not check required fields. Co-Authored-By: Claude Opus 5 (1M context) --- tzrec/protos/prompt.proto | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tzrec/protos/prompt.proto b/tzrec/protos/prompt.proto index b3290f0ad..31b49f4e4 100644 --- a/tzrec/protos/prompt.proto +++ b/tzrec/protos/prompt.proto @@ -20,8 +20,11 @@ message PromptConfig { // single-feature slot named after it. repeated PromptSlot slots = 4; - // Required when any slot renders SIDs. - optional SidSpace sid_space = 5; + // The codebook sizes the answer, extends the vocabulary and fixes the + // decode bands. Required in this version because every prompt-native + // model here emits SIDs; relax it once a non-SID one exists. text_format + // does not enforce required, so compile_prompt raises the actual error. + required SidSpace sid_space = 5; // Validation ceiling, not a truncation trigger: an over-long row is an // error, never truncated.