diff --git a/requirements/runtime.txt b/requirements/runtime.txt index 8036e8678..df6a58d8a 100644 --- a/requirements/runtime.txt +++ b/requirements/runtime.txt @@ -23,3 +23,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 diff --git a/tzrec/datasets/dataset.py b/tzrec/datasets/dataset.py index 07024c05e..5cb56b77a 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 PromptAssembler +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 @@ -97,6 +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 + compiled_prompt (CompiledPrompt, optional): compiled prompt assembly contract. """ def __init__( @@ -107,8 +111,18 @@ def __init__( reserved_columns: Optional[List[str]] = None, mode: Mode = Mode.EVAL, debug_level: int = 0, + compiled_prompt: Optional[CompiledPrompt] = None, ) -> None: super(BaseDataset, self).__init__() + self._prompt_assembler = ( + PromptAssembler( + compiled_prompt.prompt_plan, + compiled_prompt.sid_space, + include_response=mode != Mode.PREDICT, + ) + if compiled_prompt is not None + else None + ) self._data_config = data_config self._features = features self._input_path = input_path @@ -121,8 +135,25 @@ def __init__( else None ) + parser_features = features + if compiled_prompt is not None and mode == Mode.PREDICT: + prompt_feature_names = { + feature_name + 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 compiled_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, @@ -385,6 +416,16 @@ def _build_batch(self, input_data: Dict[str, pa.Array]) -> Batch: else: batch = self._data_parser.to_batch(output_data) + if self._prompt_assembler is not None: + batch.additional_infos.update( + { + k: torch.from_numpy(np.asarray(v)) + for k, v in self._prompt_assembler.assemble_batch( + output_data + ).items() + } + ) + # Set checkpoint info on batch batch.checkpoint_info = checkpoint_info batch.data_timestamp = data_timestamp @@ -762,6 +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, + compiled_prompt: Optional[CompiledPrompt] = None, ) -> DataLoader: """Build dataloader. @@ -776,6 +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. + compiled_prompt (CompiledPrompt, optional): when set, each batch carries + the assembled prompt streams in ``additional_infos``. Return: dataloader (dataloader): a DataLoader. @@ -790,6 +834,7 @@ def create_dataloader( reserved_columns=reserved_columns, mode=mode, debug_level=debug_level, + 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 6da197bfa..02ce22aa8 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -72,12 +72,17 @@ 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.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 from tzrec.protos.eval_pb2 import EvalConfig from tzrec.protos.export_pb2 import ExportConfig 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, predict_util from tzrec.utils.delta_embedding_dump import DeltaEmbeddingDumper @@ -119,6 +124,17 @@ def _create_features( return features +def _compile_prompt( + pipeline_config: EasyRecConfig, features: List[BaseFeature] +) -> Optional[CompiledPrompt]: + """Compile prompt_config for entry points that build prompt-aware objects.""" + 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 = ( @@ -137,6 +153,7 @@ def _create_model( labels: List[str], sample_weights: Optional[List[str]] = None, sampler_type: Optional[str] = None, + compiled_prompt: Optional[CompiledPrompt] = None, ) -> BaseModel: """Build model. @@ -146,6 +163,8 @@ def _create_model( labels (list): list of label names. sample_weights (list): list of sample weight names. sampler_type (str): negative sampler type + compiled_prompt (CompiledPrompt, optional): forwarded to prompt-native models. + Return: model: a EasyRec Model. """ @@ -159,6 +178,7 @@ def _create_model( labels, sample_weights=sample_weights, sampler_type=sampler_type, + compiled_prompt=compiled_prompt, ) kernel = Kernel[KernelProto.Name(model_config.kernel)] @@ -693,6 +713,7 @@ def train_and_evaluate( data_config = pipeline_config.data_config # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) + compiled_prompt = _compile_prompt(pipeline_config, features) ckpt_manager = checkpoint_util.CheckpointManager( pipeline_config.model_dir, @@ -731,6 +752,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(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: @@ -743,6 +766,7 @@ def train_and_evaluate( features, pipeline_config.train_input_path, mode=Mode.TRAIN, + compiled_prompt=compiled_prompt, checkpoint_state=dataloader_state, ) eval_dataloader = None @@ -754,6 +778,7 @@ def train_and_evaluate( features, pipeline_config.eval_input_path, mode=Mode.EVAL, + compiled_prompt=compiled_prompt, gl_cluster=gl_cluster, ) @@ -766,7 +791,11 @@ def train_and_evaluate( list(data_config.label_fields), sample_weights=list(data_config.sample_weight_fields), sampler_type=sampler_type, + compiled_prompt=compiled_prompt, ) + # Cold start only; a resumed or fine-tuned run gets its weights from DCP. + if ckpt_path is None: + model.init_from_pretrained() model = TrainWrapper( model, device=device, mixed_precision=train_config.mixed_precision ) @@ -957,12 +986,14 @@ def evaluate( data_config = pipeline_config.data_config # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) + 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, + compiled_prompt=compiled_prompt, ) sampler_type = _get_sampler_type(data_config) @@ -974,6 +1005,7 @@ def evaluate( list(data_config.label_fields), sample_weights=list(data_config.sample_weight_fields), sampler_type=sampler_type, + compiled_prompt=compiled_prompt, ) model = TrainWrapper( model, device=device, mixed_precision=train_config.mixed_precision @@ -1005,6 +1037,7 @@ def evaluate( ) if checkpoint_path: + check_prompt_assets(compiled_prompt, checkpoint_path) ckpt_manager.restore( checkpoint_path, model, @@ -1077,6 +1110,53 @@ 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() + + # 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")): + 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.hf_export_util import dcp_to_hf + + dcp_to_hf(checkpoint_path, export_dir) + # Carry the prompt contract saved alongside the checkpoint. + 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 @@ -1096,18 +1176,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() - if isinstance(model.model, MatchModel): for name, module in model.model.named_children(): if isinstance(module, MatchTower) or isinstance(module, MatchTowerWoEG): @@ -1284,6 +1352,7 @@ def predict( data_config.drop_remainder = False # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) + compiled_prompt = _compile_prompt(pipeline_config, features) infer_dataloader = create_dataloader( data_config, @@ -1291,6 +1360,7 @@ def predict( predict_input_path, reserved_columns=reserved_cols, mode=Mode.PREDICT, + compiled_prompt=compiled_prompt, debug_level=debug_level, ) infer_iterator = infer_dataloader.get_iterator() # pyre-ignore[16] @@ -1557,6 +1627,7 @@ def predict_checkpoint( data_config = pipeline_config.data_config # Build feature features = _create_features(list(pipeline_config.feature_configs), data_config) + compiled_prompt = _compile_prompt(pipeline_config, features) # Build dataloader predict_dataloader = create_dataloader( @@ -1565,6 +1636,7 @@ def predict_checkpoint( predict_input_path, reserved_columns=reserved_cols, mode=Mode.PREDICT, + compiled_prompt=compiled_prompt, debug_level=debug_level, ) @@ -1584,6 +1656,7 @@ def predict_checkpoint( pipeline_config.model_config, features, [], + compiled_prompt=compiled_prompt, ) model.set_is_inference(True) model = PredictWrapper( @@ -1619,6 +1692,7 @@ def predict_checkpoint( model.eval() if checkpoint_path: + check_prompt_assets(compiled_prompt, checkpoint_path) ckpt_manager.restore( checkpoint_path, model, diff --git a/tzrec/main_test.py b/tzrec/main_test.py index 2afbbb133..4c1f96ad3 100644 --- a/tzrec/main_test.py +++ b/tzrec/main_test.py @@ -20,14 +20,15 @@ import pyarrow as pa import torch +from google.protobuf import text_format from parameterized import parameterized from tzrec.datasets.utils import RecordBatchTensor -from tzrec.main import _train_and_evaluate, predict, predict_checkpoint +from tzrec.main import _train_and_evaluate, export, predict, predict_checkpoint from tzrec.optim.ema import DenseEMA from tzrec.protos.data_pb2 import DataConfig 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 from tzrec.utils import predict_util @@ -162,6 +163,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 PredictionLifecycleTest(unittest.TestCase): """Tests for prediction lifecycle wiring.""" diff --git a/tzrec/models/model.py b/tzrec/models/model.py index 26ec63dbc..2a88e9c20 100644 --- a/tzrec/models/model.py +++ b/tzrec/models/model.py @@ -94,6 +94,25 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: """ raise NotImplementedError + def save_assets(self, target_dir: str) -> None: + """Write model-specific assets alongside a checkpoint. + + ``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 directory. + """ + + def init_from_pretrained(self) -> None: + """Load pretrained weights at cold start (no checkpoint to restore). + + 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: """Initialize loss modules.""" raise NotImplementedError diff --git a/tzrec/models/prompt_generative_model.py b/tzrec/models/prompt_generative_model.py new file mode 100644 index 000000000..5d9f6061f --- /dev/null +++ b/tzrec/models/prompt_generative_model.py @@ -0,0 +1,269 @@ +# 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. + +"""Shared causal-LM plumbing for prompt-native generative models. + +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 + +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 +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. + compiled_prompt: the compiled prompt; required. + """ + + def __init__( + self, + model_config: ModelConfig, + features: List[BaseFeature], + labels: List[str], + sample_weights: Optional[List[str]] = None, + compiled_prompt: Optional[CompiledPrompt] = None, + **kwargs: Any, + ) -> None: + super().__init__(model_config, features, labels, sample_weights, **kwargs) + 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 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 = compiled_prompt + cfg = self._model_config + + 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( + compiled_prompt.sid_space.target_vocab_size, mean_resizing=False + ) + self.init_input() + + # decode subtracts these every step; a buffer follows the module's device + self.register_buffer( + "_level_offsets", + torch.tensor(compiled_prompt.sid_space.level_offsets), + persistent=False, + ) + + 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_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_name_or_path) + model = AutoModelForCausalLM.from_config(config) + return model.to(_PARAM_DTYPE[lm_parameter_dtype]) + + def _build_projections(self) -> None: + """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``. + """ + prompt_plan = self._prompt.prompt_plan + projection_plan = self._prompt.projection_plan + hidden_size = int(self.lm.config.hidden_size) + + modules_by_id: Dict[str, PromptProjection] = {} + in_dims: Dict[str, int] = {} + 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 modules_by_id: + modules_by_id[module_id] = PromptProjection( + projection_plan.projections[module_id], in_dim, hidden_size + ) + 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 dims ({in_dims[module_id]} vs " + f"{in_dim}); they cannot share a module." + ) + 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 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``. + + Returns: + ``(total_tokens, hidden_size)``. + """ + ids = batch.additional_infos[PROMPT_INPUT_IDS] + embeds = self.lm.get_input_embeddings()(ids) + + prompt_plan = self._prompt.prompt_plan + if not prompt_plan.projected_slots: + return embeds + + grouped = self.embedding_group(batch) + hidden_size = embeds.shape[-1] + 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(projected_embeddings).to(embeds.dtype), + ) + + 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: + 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 + codes = tokens - space.base_vocab_size - self._level_offsets + return codes.view(batch_size, -1, space.num_levels) + + 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_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_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 new file mode 100644 index 000000000..ee0d95e11 --- /dev/null +++ b/tzrec/models/prompt_generative_model_test.py @@ -0,0 +1,293 @@ +# 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 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, + 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, + parameterized_name_func, +) + +_CODEBOOK = [4, 4, 4] +_WORDS = ["History", "Predict", ":", ".", "", "<|im_end|>"] + + +_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): + """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 = create_prompt_tokenizer( + os.path.join(self.test_dir, "tok.json"), _WORDS + ) + self.features = [ + create_prompt_feature(_HIST), + create_prompt_feature(_ANSWER), + ] + 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) + + def _model( + self, + features=None, + 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, + ["answer"], + compiled_prompt=( + self.compiled_prompt if compiled_prompt == -1 else compiled_prompt + ), + ) + + 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()} + ) + return batch + + def test_tokens_to_local_codes_undoes_shifts_and_groups_beams(self) -> None: + model = self._model() + space = self.compiled_prompt.sid_space + local_codes = torch.tensor( + [ + [0, 1, 3], + [3, 0, 2], + [1, 3, 0], + [2, 2, 1], + ] + ) + 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)) + 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"): + self._model(compiled_prompt=None) + + def test_rejects_a_prompt_that_declares_no_sid_space(self) -> None: + # 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) + + def test_shared_projection_name_requires_matching_widths(self) -> None: + features = [ + create_prompt_feature(_HIST), + create_prompt_feature(_ANSWER), + create_prompt_feature(_projected("pa", 8)), + create_prompt_feature(_projected("pb", 16)), + ] + cfg = PromptConfig( + tokenizer_path=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) + compiled_prompt = compile_prompt(cfg, features, model_dir=self.test_dir) + + with self.assertRaisesRegex(ValueError, "cannot share a module"): + self._model(features=features, compiled_prompt=compiled_prompt) + + def test_projected_slot_overwrites_sentinels_and_backpropagates(self) -> 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) + # 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_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) + 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()) + + embeds[holes].sum().backward() + 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( + beam_widths=(1, 1, 100), + num_return_sequences=5, + ) + + 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_init_from_pretrained_replaces_the_empty_weights(self) -> None: + model = self._model() + 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_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_size] + 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.py b/tzrec/models/prompt_generative_qwen.py new file mode 100644 index 000000000..ae554b8b1 --- /dev/null +++ b/tzrec/models/prompt_generative_qwen.py @@ -0,0 +1,267 @@ +# 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. + +"""Decoder-only forward and decode over a Qwen backbone. + +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. + +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 + +import torch + +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 _capped_beam_widths, dynamic_beam_search +from tzrec.prompt.assembler import ( + PROMPT_CU_SEQLENS, + PROMPT_INPUT_IDS, + PROMPT_MAX_SEQLEN, + PROMPT_RESPONSE_LENGTHS, +) +from tzrec.prompt.plan import CompiledPrompt +from tzrec.protos.model_pb2 import ModelConfig +from tzrec.protos.models.prompt_model_pb2 import PromptModelConfig + + +class PromptGenerativeQwen(BasePromptGenerativeModel): + """Qwen family (Qwen2.5, Qwen3, ...) 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. + compiled_prompt: the compiled prompt; required. + """ + + def __init__( + self, + model_config: ModelConfig, + features: List[BaseFeature], + labels: List[str], + sample_weights: Optional[List[str]] = None, + compiled_prompt: Optional[CompiledPrompt] = None, + **kwargs: Any, + ) -> None: + super().__init__( + model_config, + features, + labels, + sample_weights, + compiled_prompt=compiled_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) + + def _read_beam_config(self, common: PromptModelConfig) -> None: + """Parse the decode knobs; the schedule must match the codebook. + + Args: + common: the shared model config. + """ + space = self._prompt.sid_space + 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( + 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 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 capped " + f"beam width ({final_width})." + ) + + def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: + """Run teacher-forced loss or inference decode over the assembled stream. + + Args: + batch: carries the packed prompt in ``additional_infos``. + + Returns: + The loss when training, the decoded SIDs otherwise. + """ + # 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)} + return _fx_wrapped_loss(self, embeds, batch) + + def _forward_loss( + self, embeds: torch.Tensor, batch: Batch + ) -> Dict[str, torch.Tensor]: + """Run the LM over the assembled embeddings and score the response. + + Body and head are called separately so logits cover the supervised + window only: a full (batch, length, vocab) upcast does not fit. + + Args: + embeds: the assembled prompt embeddings, packed. + batch: carries the row boundaries and the collator's width. + + Returns: + The loss. + """ + padded, mask, labels = self._left_pad_packed_inputs( + embeds, batch, build_labels=True + ) + outputs = self.lm.model(inputs_embeds=padded, attention_mask=mask) + + 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( + logits=logits, + labels=labels[:, window], + vocab_size=self.lm.config.vocab_size, + ignore_index=self._ignore_index, + ) + return {"loss": loss} + + def _generate(self, embeds: torch.Tensor, batch: Batch) -> torch.Tensor: + """Beam-search the SID answer. + + Args: + 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. + """ + padded, mask, _ = self._left_pad_packed_inputs(embeds, batch) + space = self._prompt.sid_space + tokens = dynamic_beam_search( + self.lm, + padded, + mask, + self._beam_widths, + list(zip(space.band_lo, space.band_hi)), + ) + 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. + + 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 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 +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 ``_left_pad_packed_inputs`` 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", embeds: torch.Tensor, 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. + + 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(embeds, batch) diff --git a/tzrec/models/prompt_generative_qwen_test.py b/tzrec/models/prompt_generative_qwen_test.py new file mode 100644 index 000000000..72fe095fe --- /dev/null +++ b/tzrec/models/prompt_generative_qwen_test.py @@ -0,0 +1,68 @@ +# 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.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 LeftPadPackedInputsTest(unittest.TestCase): + """The one adapter where padding lives.""" + + 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]) + 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 = model._left_pad_packed_inputs( + embeds, batch, build_labels=True + ) + + self.assertEqual(padded.shape, (2, 7, 2)) + 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] * 6 + [8], [ignore] * 5 + [7, 8]], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/modules/dynamic_beam.py b/tzrec/modules/dynamic_beam.py new file mode 100644 index 000000000..1649c8d16 --- /dev/null +++ b/tzrec/modules/dynamic_beam.py @@ -0,0 +1,134 @@ +# 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. + +"""Band-restricted beam SID decode (no tzrec deps). + +The caller owns the width schedule; this module only enforces what each level +can supply. +""" + +from typing import List, Sequence, Tuple + +import torch +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, + prompt_embeds: torch.Tensor, + attention_mask: torch.Tensor, + beam_widths: List[int], + bands: Sequence[Tuple[int, int]], +) -> torch.Tensor: + """Decode SID answers with a caller-supplied per-level beam width. + + Args: + model: an HF causal LM exposing ``.model`` / ``.lm_head`` (Qwen layout). + 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. + 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. + The fixed-length, EOS-free answer needs no finished-sequence bookkeeping. + """ + device = prompt_embeds.device + batch_size = prompt_embeds.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 " + 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)}." + ) + 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 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) + return logits[:, band_lo : band_hi + 1].float() - log_z + + position_ids = (attention_mask.long().cumsum(-1) - 1).clamp(min=0) + outputs = model.model( + inputs_embeds=prompt_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + use_cache=True, + ) + cache = outputs.past_key_values + scores = _band_logp(model.lm_head(outputs.last_hidden_state[:, -1, :]), 0) + 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])) + beam_mask = attention_mask.repeat_interleave(capped_widths[0], dim=0) + width = capped_widths[0] + + 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_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=beam_mask, + position_ids=step_position, + past_key_values=cache, + use_cache=True, + 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 + ) + 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) + width = capped_widths[level] + if level + 1 < num_levels: + # the last level never reads the cache; skip the largest reorder copy. + cache.reorder_cache(parent) + beam_mask = beam_mask[parent] + return seq diff --git a/tzrec/modules/dynamic_beam_test.py b/tzrec/modules/dynamic_beam_test.py new file mode 100644 index 000000000..1a6919f9b --- /dev/null +++ b/tzrec/modules/dynamic_beam_test.py @@ -0,0 +1,145 @@ +# 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 itertools +import unittest +from typing import Any, Dict, List, Tuple + +import torch +from parameterized import parameterized + +from tzrec.modules.dynamic_beam import dynamic_beam_search +from tzrec.utils.test_util import create_tiny_causal_lm, parameterized_name_func + + +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, + lm.get_input_embeddings()(ids), + torch.ones_like(ids) if attention_mask is None else attention_mask, + beam_widths=beam_widths, + bands=pairs, + ) + + +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 get_input_embeddings(self): + return self._lm.get_input_embeddings() + + def model(self, **kwargs: Any) -> Any: + first = kwargs.get("input_ids") + if first is None: + first = kwargs["inputs_embeds"] + self.rows.append(first.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 DynamicBeamSearchTest(unittest.TestCase): + @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]], + # 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, 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, 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_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)] + 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 = _decode(lm, ids, pairs, attention_mask=am) + width = out.shape[0] // 2 + 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. 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) + 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)) + 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/modules/prompt_projection.py b/tzrec/modules/prompt_projection.py new file mode 100644 index 000000000..45be2b1eb --- /dev/null +++ b/tzrec/modules/prompt_projection.py @@ -0,0 +1,67 @@ +# 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 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. + 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__() + 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) + + 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..4e7131857 --- /dev/null +++ b/tzrec/modules/prompt_projection_test.py @@ -0,0 +1,54 @@ +# 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(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: + 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_rejects_an_explicitly_empty_mlp(self) -> None: + config = PromptProjectionConfig() + config.mlp.SetInParent() + + with self.assertRaisesRegex(ValueError, "hidden_units must not be empty"): + PromptProjection(config, in_dim=12, hidden_size=8) + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/optim/lr_scheduler.py b/tzrec/optim/lr_scheduler.py index 38680f60c..9df083a95 100644 --- a/tzrec/optim/lr_scheduler.py +++ b/tzrec/optim/lr_scheduler.py @@ -159,6 +159,67 @@ 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, + with optional linear warmup. Mirrors HuggingFace Trainer's + ``lr_scheduler_type: linear``. + + Args: + optimizer (Optimizer): an instance of Optimizer. + 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. + by_epoch (bool): schedule by epoch or by step. + """ + + def __init__( + self, + optimizer: Optimizer, + 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 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"num_training_steps ({num_training_steps})" + ) + 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 + 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 + ] + 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 + ] + + 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..a99435ae5 100644 --- a/tzrec/optim/lr_scheduler_test.py +++ b/tzrec/optim/lr_scheduler_test.py @@ -83,6 +83,32 @@ 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_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, 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() + 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, num_training_steps=6, warmup_size=2, warmup_learning_rate=0.002 + ) + self.assertFalse(lr.by_epoch) + lr.step() + self.assertAlmostEqual(opt.param_groups[0]["lr"], 0.006) + lr.step() + self.assertAlmostEqual(opt.param_groups[0]["lr"], 0.01) + 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/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/assembler.py b/tzrec/prompt/assembler.py new file mode 100644 index 000000000..8cab58acb --- /dev/null +++ b/tzrec/prompt/assembler.py @@ -0,0 +1,315 @@ +# 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, + ResolvedSidSpace, + SlotSeg, + Static, +) +from tzrec.protos.model_pb2 import FeatureGroupType + +PROMPT_INPUT_IDS = "prompt_input_ids" +PROMPT_CU_SEQLENS = "prompt_cu_seqlens" +PROMPT_HOLE_POSITIONS = "prompt_hole_positions" +PROMPT_MAX_SEQLEN = "prompt_max_seqlen" +PROMPT_RESPONSE_LENGTHS = "prompt_response_lengths" + + +@dataclass +class AssembledPrompt: + """One batch of assembled prompts, in packed varlen form. + + Args: + 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, + grouped by included projected occurrence in + ``PromptPlan.projected_slots`` order, then by sample. + response_lengths: number of response tokens in each sample. + """ + + input_ids: np.ndarray + cu_seqlens: np.ndarray + hole_positions: np.ndarray + response_lengths: np.ndarray + + @property + def max_seqlen(self) -> int: + """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))) + + +class PromptAssembler: + """Walks a ``PromptPlan`` to build token streams. + + Args: + prompt_plan: the compiled walk order. + sid_space: resolved SID token space; required when a slot renders SIDs. + include_response: whether to read and emit the supervised response. + """ + + def __init__( + self, + prompt_plan: PromptPlan, + sid_space: Optional[ResolvedSidSpace] = None, + include_response: bool = True, + ) -> None: + self._prompt_plan = prompt_plan + self._sid_space = sid_space + 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( + sid_space.codebook, dtype=np.int64 + ) + inline = [ + s + 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: + 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_size``. + """ + 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._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_space.base_vocab_size + + def _emit_sample( + self, + segments: Sequence[object], + sample_index: int, + inline_values: Dict[str, List[np.ndarray]], + projected_lengths: Dict[str, np.ndarray], + out: List[int], + holes_by_occurrence: List[List[int]], + projected_occurrence_index: int, + sample_start: int, + ) -> int: + """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, inline_values[seg.name][sample_index]) + ) + else: + assert self._sid_space is not None + width = int(projected_lengths[seg.name][sample_index]) + 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, + 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: + 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. + """ + projected_lengths = projected_lengths or {} + if batch_size is None: + if not inline_values: + raise ValueError("batch_size is required when no INLINE slot exists.") + first_inline_values = next(iter(inline_values.values())) + batch_size = len(first_inline_values) + + jagged_token_ids: List[int] = [] + response_lengths: 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] = [] + projected_occurrence_index = self._emit_sample( + self._prompt_plan.segments, + sample_index, + inline_values, + projected_lengths, + sample_token_ids, + holes_by_occurrence, + 0, + len(jagged_token_ids), + ) + prompt_len = len(sample_token_ids) + self._emit_sample( + self._response_segments, + sample_index, + inline_values, + projected_lengths, + sample_token_ids, + holes_by_occurrence, + projected_occurrence_index, + len(jagged_token_ids), + ) + 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 + ): + raise ValueError( + 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." + ) + jagged_token_ids.extend(sample_token_ids) + 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), + hole_positions=np.asarray(holes, dtype=np.int64), + response_lengths=np.asarray(response_lengths, dtype=np.int64), + ) + + 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_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: Optional[int] = None + for seg in self._prompt_plan.segments + self._response_segments: + if not isinstance(seg, SlotSeg): + continue + 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_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) + ] + 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 if batch_size is not None else 0, + ) + return { + PROMPT_INPUT_IDS: out.input_ids, + PROMPT_CU_SEQLENS: out.cu_seqlens, + PROMPT_HOLE_POSITIONS: out.hole_positions, + 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 new file mode 100644 index 000000000..ecd32f138 --- /dev/null +++ b/tzrec/prompt/assembler_test.py @@ -0,0 +1,292 @@ +# 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, + ResolvedSidSpace, + SlotSeg, + Static, + Width, + WidthKind, +) +from tzrec.protos.model_pb2 import FeatureGroupType +from tzrec.tests.prompt_test_util import assemble_into + +_BASE_VOCAB_SIZE = 1000 +_SENTINEL = 1099 + + +def _sid_space(codebook=(4, 4, 4)) -> ResolvedSidSpace: + offsets, running = [], 0 + for size in codebook: + offsets.append(running) + running += size + return ResolvedSidSpace( + codebook=tuple(codebook), + num_levels=len(codebook), + base_vocab_size=_BASE_VOCAB_SIZE, + level_offsets=tuple(offsets), + 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, + ) + + +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=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), + ) + + +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, + logits_suffix_len=None, + static_prefix_len=0, + projected_slots=projected, + ) + + +class PromptAssemblerTest(unittest.TestCase): + def test_inline_sid_gets_the_base_vocab_shift(self) -> None: + 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])]}) + + self.assertEqual( + 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) + + def test_projected_emits_sentinels_and_records_holes(self) -> None: + 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) + + # sample 0: [7, S, S] sample 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_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_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()) + out = asm.assemble({"answer": [np.array([0, 4, 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( + 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),)) + 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)), _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_column_shaped_values_are_flattened(self) -> None: + # the data parser emits (total, value_dim) for a dense sequence feature + from tzrec.prompt.plan import CompiledPrompt, ProjectionPlan + + plan = _plan((_slot("hist", FillMode.INLINE),)) + compiled_prompt = CompiledPrompt( + sid_space=_sid_space(), + prompt_plan=plan, + projection_plan=ProjectionPlan(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(compiled_prompt, parsed) + self.assertEqual(out["prompt_cu_seqlens"].tolist(), [0, 3, 6]) + self.assertEqual(out["prompt_input_ids"].tolist()[0], _BASE_VOCAB_SIZE + 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 new file mode 100644 index 000000000..b933c7b5c --- /dev/null +++ b/tzrec/prompt/compile.py @@ -0,0 +1,527 @@ +# 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, + ProjectionPlan, + PromptPlan, + ResolvedSidSpace, + Segment, + SlotSeg, + Static, + Width, + WidthKind, +) +from tzrec.protos.model_pb2 import FeatureGroupConfig, FeatureGroupType +from tzrec.protos.prompt_pb2 import ( + PromptConfig, + PromptProjection, + PromptSlot, + SidSpace, +) +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_slots_by_name: Dict[str, PromptSlot] +) -> PromptSlot: + """Return the declared slot, or an implicit single-feature slot.""" + if name in declared_slots_by_name: + return declared_slots_by_name[name] + implicit = PromptSlot(name=name) + implicit.feature_names.append(name) + return implicit + + +def _slot_width( + 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. 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) + # 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)) + + +def _derive_slot_layout( + name: str, members: Sequence[BaseFeature] +) -> 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"{[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." + ) + 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 _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))] + + +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: + return [int(c) for c in json.load(f)["codebook"]] + + +def _build_sid_space( + cfg: PromptConfig, tok: Tokenizer, base_vocab_size: int, has_projection: bool +) -> ResolvedSidSpace: + """Extend the tokenizer with SID tokens and resolve the token space.""" + if not cfg.HasField("sid_space"): + 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: + 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." + ) + + 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( + "SID tokens are already in the base tokenizer, e.g. " + f"{existing_sid_tokens[:3]}; change sid_space.token_format." + ) + tok.add_special_tokens(sid_tokens) + + 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_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_size=base_vocab_size, + level_offsets=tuple(offsets), + band_lo=tuple(lo), + band_hi=tuple(hi), + target_vocab_size=_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. + """ + 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 "") + resolved_slots_by_name = { + name: _resolve_slot(name, declared_slots_by_name) + for name in body_names + resp_names + } + + unreferenced = set(declared_slots_by_name) - set(resolved_slots_by_name) + 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 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_feature_names} that " + "are not in feature_configs." + ) + 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." + ) + 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." + ) + + 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_size, 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 = {name: i for i, name in enumerate(resolved_slots_by_name)} + segs: Dict[str, SlotSeg] = {} + 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 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=group_type, + output_key=( + ".sequence" if group_type == FeatureGroupType.JAGGED_SEQUENCE else "" + ), + fill=fill_mode, + width=_slot_width(members[name], group_type, levels), + ) + + 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_projection_plan(projected, resolved_slots_by_name) + + 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), + logits_suffix_len=_suffix_keep(response), + static_prefix_len=_static_prefix_len(body), + projected_slots=projected, + ) + _validate(plan) + + return CompiledPrompt( + sid_space=sid_space, + prompt_plan=plan, + projection_plan=projection_plan, + tokenizer_dir=tokenizer_dir, + vocab_hash=_hash(sid_space, tok.to_str()), + plan_hash=_hash( + sid_space, + plan, + # items(), not the bare dict: iterating one yields only its keys + sorted(projection_plan.projections.items()), + # routing too: matching bodies can still be wired to other slots + sorted(projection_plan.slot_to_module.items()), + tok.to_str(), + ), + ) + + +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 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_by_name[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 + + groups = tuple( + FeatureGroupConfig( + group_name=seg.name, + feature_names=list(seg.feature_names), + group_type=seg.group_type, + ) + for seg in projected + ) + return ProjectionPlan( + projections=projections, + slot_to_module=slot_to_module, + feature_groups=groups, + ) + + +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.num_positions is not None + total += seg.width.num_positions + 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.num_positions is not None + total += seg.width.num_positions + 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(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: + 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." + ) + 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 not plan.response_segments: + raise ValueError( + "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." + ) diff --git a/tzrec/prompt/compile_test.py b/tzrec/prompt/compile_test.py new file mode 100644 index 000000000..34988a076 --- /dev/null +++ b/tzrec/prompt/compile_test.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. + +import json +import os +import unittest + +from google.protobuf import text_format +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", ":", ".", "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 }" +) +_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 = create_prompt_tokenizer( + os.path.join(self.test_dir, "tok.json"), _WORDS + ) + + 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: + cfg = self._config(prompt="History : {{hist}}") + cfg.sid_space.codebook.extend([4, 4, 4]) + compiled = self._compile(cfg, [create_prompt_feature(_HIST)]) + space = compiled.sid_space + + 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_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_size % 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, [create_prompt_feature(_HIST), create_prompt_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"] + ) + 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) + + 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, [create_prompt_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, [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, "") + self.assertIs(seg.width.kind, WidthKind.STATIC) + self.assertEqual(seg.width.num_positions, 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, [create_prompt_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, [create_prompt_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, [create_prompt_feature(_HIST), create_prompt_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, [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, [create_prompt_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, [create_prompt_feature(_HIST)]) + + def test_sid_tokens_absent_from_the_base_tokenizer(self) -> None: + cfg = self._config(prompt="X : {{hist}}") + cfg.sid_space.codebook.extend([4]) + # 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)]) + + 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, [create_prompt_feature(_HIST)]) + written = os.path.join(compiled.tokenizer_dir, "tokenizer.json") + self.assertTrue(os.path.exists(written)) + # 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|>")) + + 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 = create_prompt_feature( + 'sequence_raw_feature { feature_name: "answer" expression: "item: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) + ) + # the answer is one SID item, so its width needs no sequence_length + self.assertIs(seg.width.kind, WidthKind.STATIC) + 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.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_missing_sid_space_is_rejected(self) -> None: + # 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 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|>" + 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: + # 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") + 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 + # 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.num_positions, 16) + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/prompt/persist.py b/tzrec/prompt/persist.py new file mode 100644 index 000000000..2f458b11d --- /dev/null +++ b/tzrec/prompt/persist.py @@ -0,0 +1,160 @@ +# 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. ``ProjectionPlan`` is deliberately absent: it is model-only and rebuilt by +``compile_prompt`` from config before model construction. +""" + +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(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: + compiled_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) + + with open(os.path.join(out, _SID_SPACE), "w") as f: + json.dump(_plain(compiled_prompt.sid_space), f, indent=2) + with open(os.path.join(out, _PROMPT_PLAN), "w") as f: + json.dump(_plain(compiled_prompt.prompt_plan), f, indent=2) + with open(os.path.join(out, _HASHES), "w") as f: + json.dump( + { + "vocab_hash": compiled_prompt.vocab_hash, + "plan_hash": compiled_prompt.plan_hash, + }, + f, + indent=2, + ) + + if compiled_prompt.tokenizer_dir and os.path.isdir(compiled_prompt.tokenizer_dir): + shutil.copytree( + compiled_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( + 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 + ranges the weights never learned, which produces plausible output rather + than an error. A ``plan_hash`` mismatch only reshapes the prompt, so it + 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: + raise ValueError( + f"checkpoint [{ckpt_dir}] records no prompt assets, so its " + 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." + ) + + 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 {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") != 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 " + 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 new file mode 100644 index 000000000..df320494a --- /dev/null +++ b/tzrec/prompt/persist_test.py @@ -0,0 +1,223 @@ +# 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 unittest import mock + +from tzrec.prompt.compile import compile_prompt +from tzrec.prompt.persist import ( + PROMPT_DIR, + check_prompt_assets, + copy_prompt_assets, + read_prompt_hashes, + save_prompt_assets, +) +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|>"] + + +class PromptPersistTest(unittest.TestCase): + def setUp(self) -> None: + self.test_dir = make_test_dir() + self.tok_path = create_prompt_tokenizer( + os.path.join(self.test_dir, "tok.json"), _WORDS + ) + 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, response="{{answer}}" + ) + 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: + compiled_prompt = self._compile() + ckpt = os.path.join(self.test_dir, "model.ckpt-1") + 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"): + 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: + compiled_prompt = self._compile() + ckpt = os.path.join(self.test_dir, "model.ckpt-1") + 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], 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(compiled_prompt.sid_space)}, + ) + + 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_projection_body_warns(self) -> None: + # a body change must not 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_swapped_projection_routing_warns(self) -> None: + # identical bodies, so only slot_to_module differs + 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) + changed_compiled_prompt = self._compile(prompt="Predict : {{hist}}") + # the vocabulary is untouched, so the weights are still usable + 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(changed_compiled_prompt, ckpt) + warning.assert_called_once() + + def test_a_checkpoint_without_assets_is_fatal(self) -> None: + # 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)) + 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: + 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") + 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") + 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: + 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() diff --git a/tzrec/prompt/plan.py b/tzrec/prompt/plan.py new file mode 100644 index 000000000..8ae844da3 --- /dev/null +++ b/tzrec/prompt/plan.py @@ -0,0 +1,199 @@ +# 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.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__``. +""" + +from dataclasses import dataclass +from enum import Enum +from typing import Mapping, Optional, Tuple, Union + +from tzrec.protos.model_pb2 import FeatureGroupConfig, 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. + num_positions: the exact count or the ceiling; None when UNBOUNDED. + """ + + kind: WidthKind + 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.num_positions is not 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.num_positions}." + ) + + +@dataclass(frozen=True) +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 + 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_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_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_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. + pad_token_id: padding id of the extended tokenizer. + """ + + codebook: Tuple[int, ...] + num_levels: int + base_vocab_size: int + level_offsets: Tuple[int, ...] + band_lo: Tuple[int, ...] + band_hi: Tuple[int, ...] + target_vocab_size: int + sentinel_token_id: Optional[int] + eos_token_id: int + pad_token_id: int + + +@dataclass(frozen=True) +class Static: + """A run of literal template tokens. + + Args: + token_ids: the tokenized run. + """ + + token_ids: Tuple[int, ...] + + +@dataclass(frozen=True) +class SlotSeg: + """One ``{{name}}`` position in the assembled stream. + + Args: + 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. + """ + + slot_id: int + name: str + feature_names: Tuple[str, ...] + group_type: "FeatureGroupType.ValueType" + output_key: str + fill: FillMode + width: Width + + +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. + logits_suffix_len: upper bound on the supervised logits window. + static_prefix_len: leading positions that are request-invariant. + projected_slots: PROJECTED occurrences in emission order. + """ + + segments: Tuple[Segment, ...] + response_segments: Tuple[Segment, ...] + max_length: int + max_total_length: Optional[int] + max_holes: int + logits_suffix_len: Optional[int] + static_prefix_len: int + projected_slots: Tuple[SlotSeg, ...] + + +@dataclass(frozen=True) +class ProjectionPlan: + """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. + 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) +class CompiledPrompt: + """Everything ``compile_prompt`` produces. + + Args: + sid_space: the resolved SID token space. + prompt_plan: assembler walk order and ceilings. + 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[ResolvedSidSpace] + prompt_plan: PromptPlan + projection_plan: ProjectionPlan + tokenizer_dir: str + vocab_hash: str + plan_hash: str diff --git a/tzrec/protos/export.proto b/tzrec/protos/export.proto index 68dc0d5c0..5df8a4f78 100644 --- a/tzrec/protos/export.proto +++ b/tzrec/protos/export.proto @@ -1,6 +1,11 @@ syntax = "proto2"; package tzrec.protos; +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 @@ -23,4 +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; + optional ExportFormat export_format = 8 [default = TORCHSCRIPT]; } diff --git a/tzrec/protos/model.proto b/tzrec/protos/model.proto index d2c34ae0f..d72b1f206 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/prompt_model.proto"; import "tzrec/protos/models/sid_model.proto"; import "tzrec/protos/loss.proto"; import "tzrec/protos/metric.proto"; @@ -81,6 +82,9 @@ message ModelConfig { // SID generation models SidRqvae sid_rqvae = 600; SidRqkmeans sid_rqkmeans = 601; + + // Generative (causal-LM) models; the 700-block keeps clear of the SID 600s. + 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..a30b98645 --- /dev/null +++ b/tzrec/protos/models/prompt_model.proto @@ -0,0 +1,41 @@ +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 width after per-level candidate capping. + required uint32 num_return_sequences = 3; + + optional string generated_sids_key = 4 [default = "generated_sids"]; + + // 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, ...). +message PromptGenerativeQwen { + optional PromptModelConfig common = 1; + + // 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_name_or_path = 2 [default = "Qwen/Qwen2.5-0.5B"]; +} diff --git a/tzrec/protos/optimizer.proto b/tzrec/protos/optimizer.proto index 8e5305807..330eb4d30 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; } } @@ -43,6 +44,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; optional EMAConfig ema = 202; @@ -64,6 +66,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; } } @@ -239,6 +242,17 @@ message ManualStepLR { optional bool by_epoch = 4 [default = false]; } +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. Must be > 0 when this scheduler is selected. + optional uint32 num_training_steps = 1; + optional float min_learning_rate = 2 [default = 0.0]; + optional float warmup_learning_rate = 3 [default = 0.0]; + optional uint32 warmup_size = 4 [default = 0]; + 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/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..31b49f4e4 --- /dev/null +++ b/tzrec/protos/prompt.proto @@ -0,0 +1,71 @@ +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 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 = 2; + // Supervised answer. Defines the loss span; absent at inference. + 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 = 4; + + // 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. + 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 = 7 [default = "<|pg_hole|>"]; +} + +message PromptSlot { + // {{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; + // 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 = 3; + // Weight-sharing key across slots. Derived groups dedupe automatically; + // modules share only on request. + optional string projection_name = 4; +} + +// 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 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. + optional string manifest_path = 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..ae0457efc --- /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_path: "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_name_or_path: "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 new file mode 100644 index 000000000..92d20d67c --- /dev/null +++ b/tzrec/tests/prompt_integration_test.py @@ -0,0 +1,127 @@ +# 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 +import torch.fx + +from tzrec.datasets.utils import Batch +from tzrec.main import _create_model +from tzrec.prompt.compile import compile_prompt +from tzrec.protos.model_pb2 import ModelConfig +from tzrec.protos.prompt_pb2 import PromptConfig +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|>"] + + +class PromptStackIntegrationTest(unittest.TestCase): + """compile -> assemble -> model, on the real code path.""" + + 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 = 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_path=self.tok, + prompt="History : {{hist}} . Predict :", + response="{{answer}}", + ) + cfg.sid_space.codebook.extend(_CODEBOOK) + 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_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"], + compiled_prompt=self.compiled_prompt, + ) + + def _batch(self, hist, answer): + parsed = { + "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.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_size(self) -> None: + model = self._model() + rows = model.lm.get_input_embeddings().weight.shape[0] + 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() + batch = self._batch([0, 1, 2, 3, 0, 1], [1, 2, 3]) + 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_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) + + torch.fx.symbolic_trace(_Wrapper(model)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/tests/prompt_test_util.py b/tzrec/tests/prompt_test_util.py new file mode 100644 index 000000000..524ef26d4 --- /dev/null +++ b/tzrec/tests/prompt_test_util.py @@ -0,0 +1,88 @@ +# 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, 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( + compiled_prompt: CompiledPrompt, + parsed_features: Dict[str, "np.ndarray"], +) -> Dict[str, np.ndarray]: + """Assemble one parsed batch with a temporary assembler. + + Args: + 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( + compiled_prompt.prompt_plan, compiled_prompt.sid_space + ).assemble_batch(parsed_features) diff --git a/tzrec/utils/checkpoint_util.py b/tzrec/utils/checkpoint_util.py index 2beb37f38..b71da865b 100644 --- a/tzrec/utils/checkpoint_util.py +++ b/tzrec/utils/checkpoint_util.py @@ -332,6 +332,34 @@ def best_checkpoint( return latest_checkpoint(model_dir) +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, attr): + if id(inner) in seen: + return None + seen.add(id(inner)) + if hasattr(inner, "module"): + inner = inner.module + elif hasattr(inner, "model"): + inner = inner.model + else: + return None + return inner + + class CheckpointManager: """Saves training checkpoints and prunes old ones asynchronously. @@ -390,9 +418,34 @@ def save( dataloader_state: Optional[Dict[str, Any]] = None, dense_ema: Optional[DenseEMA] = 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. + + 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) + # 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." + ) + try: + 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 " + 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 diff --git a/tzrec/utils/hf_export_util.py b/tzrec/utils/hf_export_util.py new file mode 100644 index 000000000..ad24a1fde --- /dev/null +++ b/tzrec/utils/hf_export_util.py @@ -0,0 +1,161 @@ +# 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. + +Kept out of ``export_util`` so ``checkpoint_util`` can call it 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 + +_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 write_hf_assets(wrapped_model: nn.Module, save_dir: str) -> None: + """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. + """ + if int(os.environ.get("RANK", 0)) != 0: + return + inner = checkpoint_util.unwrap_to(wrapped_model, "hf_backbone") + 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) + # 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( + (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 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. + """ + 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 + + 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: + raise RuntimeError( + "dcp_to_hf: cannot map the DCP state dict onto the backbone " + 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." + ) + + # 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..421d58033 --- /dev/null +++ b/tzrec/utils/hf_export_util_test.py @@ -0,0 +1,165 @@ +# 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 threading +import unittest +from unittest import mock + +import torch +from safetensors.torch import load_file +from torch import nn + +from tzrec.utils.checkpoint_util import save_model, unwrap_to +from tzrec.utils.hf_export_util import ( + _HF_EXPORT_META_FILENAME, + dcp_to_hf, + write_hf_assets, +) +from tzrec.utils.test_util import create_tiny_causal_lm, 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.""" + + 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 an HF-backed model exposing the optional tokenizer protocol.""" + + 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 + + +class HfExportUtilTest(unittest.TestCase): + def setUp(self) -> None: + self.test_dir = make_test_dir() + # 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) + + 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_to(a, "hf_backbone"))) + t.daemon = True + t.start() + t.join(timeout=5) + 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: + 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): + 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 + + def test_write_hf_assets_records_state_dict_prefix(self) -> None: + lm = _tied_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 = _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) + + 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_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 + 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() diff --git a/tzrec/utils/test_util.py b/tzrec/utils/test_util.py index 4f398684f..440f555d1 100644 --- a/tzrec/utils/test_util.py +++ b/tzrec/utils/test_util.py @@ -120,6 +120,41 @@ 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, +) -> 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. + + Returns: + an eval-mode ``Qwen2ForCausalLM``. + """ + from transformers import Qwen2Config, Qwen2ForCausalLM + + 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] def parameterized_name_func(func, num, p) -> str: """Name func for parameterized."""