From 2b84cdb912517339a46b4837029ec5032806e9e9 Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Thu, 1 Oct 2026 10:30:06 +0000 Subject: [PATCH] [6/8] Add a sparse teacher-policy dumper [4/8] added the consumer for data.sparse_data_path but nothing in the repo produces that format, so the path was unusable. This is the counterpart to compute_hidden_states_hf.py: it stores the base's truncated next-token distribution instead of a hidden state per token, which is what the TVD objectives actually read. Position t of the stored policy is the distribution token t was drawn from, conditioned on everything before it -- the model's logits at t - 1. The round-trip test asserts that alignment against SparsePolicyDataset, since dropping the shift still yields a falling loss curve against the wrong target. Signed-off-by: Peter Hicks Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- examples/speculative_decoding/README.md | 31 +- .../collect_hidden_states/common.py | 38 +++ .../compute_hidden_states_hf.py | 38 +-- .../compute_hidden_states_vllm.py | 36 +-- .../compute_sparse_policy_hf.py | 292 ++++++++++++++++++ .../torch/speculative/external/sparse_data.py | 33 +- .../speculative/test_external_draft_sparse.py | 25 ++ .../test_sparse_policy_producer.py | 121 ++++++++ 8 files changed, 537 insertions(+), 77 deletions(-) create mode 100644 examples/speculative_decoding/collect_hidden_states/compute_sparse_policy_hf.py create mode 100644 tests/unit/torch/speculative/test_sparse_policy_producer.py diff --git a/examples/speculative_decoding/README.md b/examples/speculative_decoding/README.md index 4d4076c2bef..786944ca9ef 100644 --- a/examples/speculative_decoding/README.md +++ b/examples/speculative_decoding/README.md @@ -148,16 +148,45 @@ dominate the dump size. Streaming feeds the same hidden states from a live mode. Online training is rejected: it produces no cached base outputs for the draft to be scored against. +The TVD objectives only ever need the base's *truncated* next-token distribution, +not a full hidden state, so they can instead read a much smaller dump of the +base's top-k deployment policy. Produce one with: + +```bash +python collect_hidden_states/compute_sparse_policy_hf.py \ + --model $BASE_MODEL --input-data $CONVERSATIONS --output-dir $POLICY_DIR \ + --top-k 20 --top-p 0.95 --answer-only-loss +``` + +and train against it with `data.sparse_data_path=$POLICY_DIR` in place of +`data.offline_data_path` — the two are alternatives, not a pair: + ```bash torchrun --nproc_per_node 8 main.py \ --config ../../modelopt_recipes/general/speculative_decoding/external_draft.yaml \ model.model_name_or_path=$BASE_MODEL \ draft_model_name_or_path=$DRAFT_MODEL \ - data.offline_data_path=$HIDDEN_STATES_DIR \ + data.sparse_data_path=$POLICY_DIR \ external.external_loss=tvd \ + external.external_top_k=20 external.external_top_p=0.95 \ training.output_dir=$OUTPUT_DIR ``` +Keep `--top-k` / `--top-p` equal to `external.external_top_k` / +`external.external_top_p`: the stored policy *is* the teacher distribution, so a dump +narrower than the objective expects charges the draft for mass on tokens the teacher +actually supported. A top-k mismatch draws a warning under `tvd_deploy`; the dump does not +record its top-p, so a top-p mismatch goes undetected, and only a stored row that does not +sum to 1 draws a warning. This dump +also requires a TVD objective; `soft_ce` needs the full base distribution and is +rejected against `sparse_data_path`. + +Only `tvd_deploy` puts the draft through the same nucleus the dump was written with. +Under `tvd` and `tvd_ce` the draft side is an untruncated softmax, so the +out-of-support term charges it for the mass the dump itself discarded. That is the +same asymmetry dense `tvd` has against a full teacher, and it is intended, but it +means `--top-p` is only symmetric between the two sides under `tvd_deploy`. + `external.external_loss` selects the objective: | value | what it scores | diff --git a/examples/speculative_decoding/collect_hidden_states/common.py b/examples/speculative_decoding/collect_hidden_states/common.py index 742c2d54661..f7733772b30 100644 --- a/examples/speculative_decoding/collect_hidden_states/common.py +++ b/examples/speculative_decoding/collect_hidden_states/common.py @@ -31,6 +31,8 @@ from pathlib import Path import torch +from datasets import load_dataset +from transformers import AutoTokenizer _DFLASH_DEFAULT_NUM_DRAFT_LAYERS = 5 @@ -115,6 +117,42 @@ def add_answer_only_loss_args(parser: argparse.ArgumentParser) -> None: ) +REMOVE_THINK_CHAT_TEMPLATE = ( + "{% if '' in content %}{% set content = content.split('')[-1] %}{% endif %}" +) + + +def load_conversations(input_data: Path, dp_rank: int = 0, dp_world_size: int = 1): + """Conversations from a ``.jsonl`` file or a directory of them, sharded for ``dp_rank``.""" + input_data = Path(input_data) + if input_data.is_file() and input_data.suffix == ".jsonl": + data_files = str(input_data) + elif input_data.is_dir(): + data_files = {"train": f"{input_data}/*.jsonl"} + else: + raise ValueError(f"input_data must be a .jsonl file or directory, got: {input_data}") + dataset = load_dataset("json", data_files=data_files, split="train") + print(f"Loaded {len(dataset)} conversations from {input_data}") + if dp_world_size > 1: + dataset = dataset.shard(num_shards=dp_world_size, index=dp_rank) + return dataset + + +def load_tokenizer(args: argparse.Namespace): + """The base tokenizer, with the same chat-template handling in every dumper.""" + tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=args.trust_remote_code) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + override_template = load_chat_template(args.chat_template) + if override_template is not None: + tokenizer.chat_template = override_template + if tokenizer.chat_template is not None: + tokenizer.chat_template = tokenizer.chat_template.replace(REMOVE_THINK_CHAT_TEMPLATE, "") + if args.answer_only_loss: + verify_generation_tags(tokenizer.chat_template) + return tokenizer + + def load_chat_template(path: Path | None) -> str | None: """Read a Jinja chat template from ``path``, or return ``None`` if not provided.""" if path is None: diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py index 21e1b652527..68a9d86b326 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_hf.py @@ -25,18 +25,13 @@ add_answer_only_loss_args, add_aux_layers_args, add_no_aux_hidden_states_args, - load_chat_template, + load_conversations, + load_tokenizer, resolve_aux_layers, tokenize_with_loss_mask, - verify_generation_tags, ) -from datasets import load_dataset from tqdm import tqdm as tqdm -from transformers import AutoModel, AutoTokenizer - -REMOVE_THINK_CHAT_TEMPLATE = ( - "{% if '' in content %}{% set content = content.split('')[-1] %}{% endif %}" -) +from transformers import AutoModel def parse_args() -> argparse.Namespace: @@ -109,21 +104,7 @@ def parse_args() -> argparse.Namespace: def main(args: argparse.Namespace) -> None: # Load conversations - if args.input_data.is_file() and str(args.input_data).endswith(".jsonl"): - dataset = load_dataset("json", data_files=str(args.input_data), split="train") - elif args.input_data.is_dir(): - dataset = load_dataset( - "json", data_files={"train": f"{args.input_data}/*.jsonl"}, split="train" - ) - else: - raise ValueError( - f"input_data must be a .jsonl file or directory containing .jsonl files, got: {args.input_data}" - ) - print(f"Loaded {len(dataset)} conversations from {args.input_data}") - - # Shard data - if args.dp_world_size > 1: - dataset = dataset.shard(num_shards=args.dp_world_size, index=args.dp_rank) + dataset = load_conversations(args.input_data, args.dp_rank, args.dp_world_size) print( f"Sharded dataset to {len(dataset)} conversations for DP#{args.dp_rank}/{args.dp_world_size}" ) @@ -162,16 +143,7 @@ def keep_conversation(entry): raise ValueError(f"config has no 'num_hidden_layers' attribute: {text_config}") selected_layer_ids = resolve_aux_layers(args, num_hidden_layers) - tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=args.trust_remote_code) - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - override_template = load_chat_template(args.chat_template) - if override_template is not None: - tokenizer.chat_template = override_template - if tokenizer.chat_template is not None: - tokenizer.chat_template = tokenizer.chat_template.replace(REMOVE_THINK_CHAT_TEMPLATE, "") - if args.answer_only_loss: - verify_generation_tags(tokenizer.chat_template) + tokenizer = load_tokenizer(args) output_dir = args.output_dir output_dir.mkdir(parents=True, exist_ok=True) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py index 2579146e12a..7f9ddb25a82 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py @@ -34,17 +34,12 @@ from common import ( add_answer_only_loss_args, add_aux_layers_args, - load_chat_template, + load_conversations, + load_tokenizer, tokenize_with_loss_mask, - verify_generation_tags, ) -from datasets import load_dataset from tqdm import tqdm -from transformers import AutoConfig, AutoTokenizer - -REMOVE_THINK_CHAT_TEMPLATE = ( - "{% if '' in content %}{% set content = content.split('')[-1] %}{% endif %}" -) +from transformers import AutoConfig def _resolve_aux_layers_standalone( @@ -194,19 +189,7 @@ def main(args: argparse.Namespace) -> None: from vllm.inputs import TokensPrompt # Load conversations - if args.input_data.is_file() and str(args.input_data).endswith(".jsonl"): - dataset = load_dataset("json", data_files=str(args.input_data), split="train") - elif args.input_data.is_dir(): - dataset = load_dataset( - "json", data_files={"train": f"{args.input_data}/*.jsonl"}, split="train" - ) - else: - raise ValueError(f"input_data must be a .jsonl file or directory, got: {args.input_data}") - print(f"Loaded {len(dataset)} conversations from {args.input_data}") - - # Shard data - if args.dp_world_size > 1: - dataset = dataset.shard(num_shards=args.dp_world_size, index=args.dp_rank) + dataset = load_conversations(args.input_data, args.dp_rank, args.dp_world_size) print(f"Sharded to {len(dataset)} conversations for DP#{args.dp_rank}/{args.dp_world_size}") # Remove already dumped conversations @@ -243,16 +226,7 @@ def keep_conversation(entry): print(f"Extracting hidden states from layers {extract_layer_ids} (last = final output)") # Tokenize conversations - tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=args.trust_remote_code) - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - override_template = load_chat_template(args.chat_template) - if override_template is not None: - tokenizer.chat_template = override_template - if tokenizer.chat_template is not None: - tokenizer.chat_template = tokenizer.chat_template.replace(REMOVE_THINK_CHAT_TEMPLATE, "") - if args.answer_only_loss: - verify_generation_tags(tokenizer.chat_template) + tokenizer = load_tokenizer(args) # Prepare prompts for vLLM prompts = [] diff --git a/examples/speculative_decoding/collect_hidden_states/compute_sparse_policy_hf.py b/examples/speculative_decoding/collect_hidden_states/compute_sparse_policy_hf.py new file mode 100644 index 00000000000..ed1ab598a7e --- /dev/null +++ b/examples/speculative_decoding/collect_hidden_states/compute_sparse_policy_hf.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Dump a base model's top-k deployment policy for external draft training. + +The counterpart to ``compute_hidden_states_hf.py``: instead of a hidden state per +token it stores the base's truncated next-token distribution, which is what the +TVD objectives actually consume. The dump is far smaller and needs no base +lm_head at training time; feed it to ``main.py`` via ``data.sparse_data_path``. + +Output is gzipped JSONL, one conversation per line, matching the schema +``modelopt.torch.speculative.external.sparse_data`` reads. +""" + +import argparse +import gzip +import json +import os +import sys +from pathlib import Path + +import torch +from common import ( + add_answer_only_loss_args, + load_conversations, + load_tokenizer, + tokenize_with_loss_mask, +) +from tqdm import tqdm +from transformers import AutoModelForCausalLM + +from modelopt.torch.speculative.plugins.hf_external import HFExternalDraftModel + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Dump a base model's top-k deployment policy over conversation responses." + ) + parser.add_argument("--model", type=str, required=True, help="Base (target) model.") + parser.add_argument( + "--input-data", type=Path, required=True, help=".jsonl file or directory of .jsonl files." + ) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument( + "--max-seq-len", + type=int, + default=2048, + help=( + "Skip conversations longer than this. Keep it <= training.training_seq_len: the " + "loader right-truncates to that length, so a record whose context alone exceeds " + "it carries no scored position." + ), + ) + parser.add_argument( + "--top-k", + type=int, + default=20, + help="Width of the stored policy. Must match external.external_top_k at training time.", + ) + parser.add_argument( + "--top-p", + type=float, + default=0.95, + help="Nucleus applied after top-k, matching external.external_top_p.", + ) + parser.add_argument( + "--shard-size", + type=int, + default=2000, + help="Records per output shard. One record per supervised span, so a multi-turn " + "conversation contributes one per assistant turn.", + ) + parser.add_argument("--dp-rank", type=int, default=0) + parser.add_argument("--dp-world-size", type=int, default=1) + # compute_hidden_states_hf.py spells this --trust_remote_code; accept both so the + # dumpers can be driven from one script. + parser.add_argument( + "--trust-remote-code", "--trust_remote_code", action="store_true", dest="trust_remote_code" + ) + parser.add_argument("--debug-max-num-conversations", type=int, default=None) + add_answer_only_loss_args(parser) + return parser.parse_args() + + +def deployment_policy(logits, top_k: int, top_p: float): + """``(ids, log_probs)`` of width ``top_k`` under the filter ``tvd_deploy`` applies to the draft. + + Reuses that filter so the two cannot drift. Slots the nucleus drops hold ``-inf``, + which ``sparse_data`` decodes to zero probability. + """ + probs = HFExternalDraftModel._deployment_probs(logits, top_k, top_p) + ids = logits.float().topk(min(top_k, logits.shape[-1]), dim=-1).indices + return ids, probs.gather(-1, ids).log() + + +def scored_spans(loss_mask): + """Contiguous runs of supervised positions, as half-open ``(start, end)`` pairs. + + A multi-turn conversation masks each assistant turn separately, so the supervised + positions are disjoint; treating everything after the first as the response would + score the later user turns too. + """ + spans, run_start, prev = [], None, None + for i in loss_mask.squeeze(0).nonzero().flatten().tolist(): + if run_start is None: + run_start = i + elif i != prev + 1: + spans.append((run_start, prev + 1)) + run_start = i + prev = i + if run_start is not None: + spans.append((run_start, prev + 1)) + # Position 0 has no preceding context, so it can carry no policy. Without + # --answer-only-loss the mask covers the whole sequence and would start at 0. + return [(max(a, 1), b) for a, b in spans if b > max(a, 1)] + + +def policy_record( + conversation_id, + ids, + first: int, + end: int, + logits, + top_k: int, + top_p: float, + source=None, +): + """Build one shard record for a single supervised span. + + Args: + first: index of the span's first supervised token (>= 1); everything before it is + context. + end: one past the span's last supervised token. + logits: the base's logits over the whole sequence. + """ + # Position t of the stored policy is the distribution token ids[t] was drawn from, + # conditioned on ids[:t] -- that is the model's logits at t - 1. Dropping this shift + # still yields a falling loss curve, against the wrong target. + tok, lp = deployment_policy(logits[first - 1 : end - 1], top_k, top_p) + record = { + "id": str(conversation_id), + "prompt_ids": ids[:first].tolist(), + "gen_ids": ids[first:end].tolist(), + "topk_tok": tok.cpu().tolist(), + "topk_lp": lp.float().cpu().tolist(), + } + if source is not None: + record["source"] = source + return record + + +def main(args: argparse.Namespace) -> None: + if args.top_k < 1: + raise ValueError( + f"--top-k must be at least 1, got {args.top_k}: the dump stores a " + "fixed-width policy and a zero-width row carries no teacher." + ) + dataset = load_conversations(args.input_data, args.dp_rank, args.dp_world_size) + if args.debug_max_num_conversations is not None: + # After DP sharding this rank may hold fewer rows than the debug cap, and + # select() raises on an out-of-range index rather than clamping. + dataset = dataset.select(range(min(args.debug_max_num_conversations, len(dataset)))) + + model = AutoModelForCausalLM.from_pretrained( + args.model, dtype="auto", device_map="auto", trust_remote_code=args.trust_remote_code + ) + model.eval() + tokenizer = load_tokenizer(args) + + args.output_dir.mkdir(parents=True, exist_ok=True) + # A second, shorter export overwrites only the shard names it reaches, so the extra + # shards from the first run stay visible to the trainer's glob and it reads records + # from both. Refuse rather than silently mix two dumps. + existing = sorted(args.output_dir.glob(f"policy-{args.dp_rank:03d}-*.jsonl.gz")) + if existing: + raise FileExistsError( + f"{args.output_dir} already holds {len(existing)} shard(s) for DP rank " + f"{args.dp_rank}, starting with {existing[0].name}. Remove them or pass a " + "fresh --output-dir; a shorter run here would leave the earlier dump's " + "extra shards behind and training would read both." + ) + num_skipped = num_invalid = num_success = 0 + shard_idx = 0 + written_in_shard = 0 + shard_file = None + + def open_shard(idx: int): + # Write-then-rename: a partial shard left by an interrupted run would otherwise + # be picked up by the trainer's glob and decoded as damaged input. + tmp = args.output_dir / f"policy-{args.dp_rank:03d}-{idx:05d}.jsonl.gz.tmp" + # Exclusive create: the existing-shard check above is a point-in-time test, so + # two exports for the same rank can both pass it and then write the same file. + # "xt" makes the second one fail here instead of interleaving records. + try: + return gzip.open(tmp, "xt"), tmp + except FileExistsError: + raise FileExistsError( + f"{tmp} already exists: another export is writing DP rank " + f"{args.dp_rank} into {args.output_dir}, or a previous one was killed. " + "Remove the .tmp files or use a fresh --output-dir." + ) from None + + completed = False + try: + shard_file, shard_tmp = open_shard(shard_idx) + for entry in tqdm(dataset, desc=f"DP#{args.dp_rank} dumping policy"): + conversation_id = entry.get("conversation_id", entry.get("uuid")) + conversations = entry.get("messages") or entry.get("conversations") + if not conversations or not isinstance(conversations, list): + num_invalid += 1 + continue + + input_ids, loss_mask = tokenize_with_loss_mask( + tokenizer, conversations, args.answer_only_loss + ) + if input_ids.shape[1] <= 10 or input_ids.shape[1] > args.max_seq_len: + num_skipped += 1 + continue + + # The supervised span is the response; everything before it is context. + # One record per supervised span: a multi-turn conversation masks each + # assistant turn separately, and each is scored against the context before it. + spans = scored_spans(loss_mask) + if not spans: + num_invalid += 1 + continue + ids = input_ids.squeeze(0) + + with torch.inference_mode(): + logits = model(input_ids=input_ids.to(model.device)).logits.squeeze(0) + for turn, (span_start, span_end) in enumerate(spans): + record = policy_record( + f"{conversation_id}-{turn}" if len(spans) > 1 else conversation_id, + ids, + span_start, + span_end, + logits, + args.top_k, + args.top_p, + entry.get("source"), + ) + shard_file.write(json.dumps(record) + "\n") + written_in_shard += 1 + num_success += 1 + + if written_in_shard >= args.shard_size: + shard_file.close() + os.replace(shard_tmp, shard_tmp.with_suffix("")) + shard_idx += 1 + written_in_shard = 0 + shard_file, shard_tmp = open_shard(shard_idx) + completed = True + finally: + if shard_file is not None: + shard_file.close() + # Publish the last shard only on a clean pass. Renaming it after a failure + # hands the trainer a truncated export that looks complete. + if written_in_shard and completed: + os.replace(shard_tmp, shard_tmp.with_suffix("")) + else: + Path(shard_tmp).unlink(missing_ok=True) + if not completed: + done = sorted(args.output_dir.glob(f"policy-{args.dp_rank:03d}-*.jsonl.gz")) + print( + f"Export failed after {num_success} conversations. {len(done)} completed " + f"shard(s) remain in {args.output_dir} and are a partial dump; remove " + "them before training or re-exporting.", + file=sys.stderr, + ) + + if num_skipped: + print(f"Skipped {num_skipped} conversations due to length constraints.") + if num_invalid: + print(f"Skipped {num_invalid} conversations without a usable response span.") + print(f"Wrote {num_success} conversations to {args.output_dir}") + + +if __name__ == "__main__": + cli_args = parse_args() + main(cli_args) diff --git a/modelopt/torch/speculative/external/sparse_data.py b/modelopt/torch/speculative/external/sparse_data.py index 2f1f7f6e2f2..4bfc4689e47 100644 --- a/modelopt/torch/speculative/external/sparse_data.py +++ b/modelopt/torch/speculative/external/sparse_data.py @@ -59,6 +59,15 @@ def _cache_signature(files: list[str]) -> list[list]: return [[str(f), os.path.getsize(f), int(os.path.getmtime(f))] for f in files] +def _cache_matches(meta: dict, files: list[str], top_k: int | None) -> bool: + """Whether a published cache serves this request; an explicit ``top_k`` must match too.""" + return ( + meta.get("version") == CACHE_VERSION + and meta.get("sources") == _cache_signature(files) + and (top_k is None or meta.get("top_k") == top_k) + ) + + def discover_shards(root) -> list[str]: """Every sparse-policy shard under ``root``, in the order the cache is keyed on. @@ -98,7 +107,9 @@ def _builder_progress(cache: Path) -> tuple[int, int]: return total, newest -def _wait_for_cache(cache: Path, files: list[str], stall_s: float = 1800.0) -> None: +def _wait_for_cache( + cache: Path, files: list[str], top_k: int | None = None, stall_s: float = 1800.0 +) -> None: """Block until rank 0 publishes a cache built from ``files``, without a collective. Waiting on mere existence is not enough. The only way to reach this with a cache @@ -113,13 +124,12 @@ def _wait_for_cache(cache: Path, files: list[str], stall_s: float = 1800.0) -> N temporary directory grows while it works, so a wait that watches it can be patient with a slow build and still fail promptly on a dead builder. """ - want = _cache_signature(files) last_progress = _builder_progress(cache) deadline = time.monotonic() + stall_s while True: try: meta = json.loads((cache / "meta.json").read_text()) - if meta.get("version") == CACHE_VERSION and meta.get("sources") == want: + if _cache_matches(meta, files, top_k): return except (OSError, ValueError): # Absent, or caught mid-publish; either way it is not ours yet. @@ -374,9 +384,10 @@ def __init__( ): """Open (building if needed) the memmapped cache for these shards. - ``top_k`` defaults to the width found in the data. A narrower value - truncates the teacher policy and a wider one pads it with zero-probability - entries; both change the objective silently. + ``top_k`` defaults to the width found in the data. A wider value pads the + stored policy with zero-probability entries, which changes the objective + silently; a narrower one would drop teacher mass and is refused by the + builder instead. ``source_weights`` maps a record's ``source`` field to a sampling weight, oversampling under-represented sources by repeating their index entries. @@ -392,10 +403,8 @@ def __init__( stale = True if (cache / "meta.json").exists(): - meta = json.loads((cache / "meta.json").read_text()) - stale = meta.get("version") != CACHE_VERSION or meta.get("sources") != ( - _cache_signature(files) - ) + # An explicit top_k is part of the request: it truncates or pads the policy. + stale = not _cache_matches(json.loads((cache / "meta.json").read_text()), files, top_k) if stale: # Under DDP every rank sees the same shards; one builds and the rest wait on @@ -413,7 +422,7 @@ def __init__( if rank == 0: _build_cache(cache, files, top_k) else: - _wait_for_cache(cache, files) + _wait_for_cache(cache, files, top_k) meta = json.loads((cache / "meta.json").read_text()) # Re-check rather than trust the build. Two jobs reading different subsets of one @@ -421,7 +430,7 @@ def __init__( # the publish rename returns from _build_cache normally having adopted the # winner's cache -- which it never compared against its own shards. The waiting # ranks are covered by _wait_for_cache; this covers the builder. - if meta.get("version") != CACHE_VERSION or meta.get("sources") != _cache_signature(files): + if not _cache_matches(meta, files, top_k): raise RuntimeError( f"the sparse cache at {cache} was published by a concurrent job from " "different shards; give this run its own cache_dir." diff --git a/tests/unit/torch/speculative/test_external_draft_sparse.py b/tests/unit/torch/speculative/test_external_draft_sparse.py index e87e09a618b..8ce378890ae 100644 --- a/tests/unit/torch/speculative/test_external_draft_sparse.py +++ b/tests/unit/torch/speculative/test_external_draft_sparse.py @@ -210,6 +210,31 @@ def test_sparse_data_path_selects_mode_and_rejects_conflicts(): DataArguments(sparse_data_path="/tmp/x", offline_data_path="/tmp/y") +def test_an_explicit_top_k_rebuilds_a_cache_built_at_another_width(tmp_path): + """A changed top_k must invalidate the cache, not be silently dropped. + + The width is part of the objective: it truncates or pads the stored policy. The + staleness check keys on the shard list, which does not change when only the + requested width does, so without this the run keeps the old width and + ``self.top_k`` reports the cache's rather than the caller's. + """ + a = tmp_path / "a.jsonl.gz" + _write_records(a, n=2) + first = SparsePolicyDataset([str(a)], max_length=32, top_k=TOPK) + assert first.top_k == TOPK + + # Widening is the direction the builder allows: it pads with zero-probability + # entries. Narrowing would drop teacher mass and raises instead, so it cannot + # reach the cache check this test is about. + wider = SparsePolicyDataset([str(a)], max_length=32, top_k=TOPK + 4) + assert wider.top_k == TOPK + 4, "the cache ignored the requested width" + # an unset top_k still accepts whatever the cache holds + assert SparsePolicyDataset([str(a)], max_length=32).top_k == TOPK + 4 + # a non-builder rank must wait for the requested width, not take the one on disk + with pytest.raises(TimeoutError): + sparse_data._wait_for_cache(tmp_path / ".sparse_cache", [str(a)], TOPK, stall_s=0) + + def test_cache_rebuilds_over_an_existing_one(tmp_path): """A stale cache must be replaced, not crash the rebuild with ENOTEMPTY.""" a = tmp_path / "a.jsonl.gz" diff --git a/tests/unit/torch/speculative/test_sparse_policy_producer.py b/tests/unit/torch/speculative/test_sparse_policy_producer.py new file mode 100644 index 00000000000..0cc9fc71a97 --- /dev/null +++ b/tests/unit/torch/speculative/test_sparse_policy_producer.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""The sparse-policy producer must write what SparsePolicyDataset reads.""" + +import gzip +import importlib.util +import json +import sys +from pathlib import Path + +import pytest +import torch + +from modelopt.torch.speculative.external.sparse_data import SparsePolicyDataset + +_PRODUCER = ( + Path(__file__).parents[4] + / "examples" + / "speculative_decoding" + / "collect_hidden_states" + / "compute_sparse_policy_hf.py" +) + +VOCAB = 64 +TOPK = 8 + + +def _load_producer(): + """Import the example script by path; it is not an installed module.""" + sys.path.insert(0, str(_PRODUCER.parent)) + try: + spec = importlib.util.spec_from_file_location("compute_sparse_policy_hf", _PRODUCER) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + finally: + # The script does `from common import ...`, so leaving either the path entry or the + # cached module behind would shadow the several other `common` modules in the tree + # for every test collected after this one in the same process. + sys.path.pop(0) + sys.modules.pop("common", None) + + +pytest.importorskip("datasets") +pytest.importorskip("transformers") + +# Below the guards: the plugins package registers this behind an import_plugin guard, but +# importing the submodule directly does not, so on a torch-only install this module would +# fail to collect rather than skip. +from modelopt.torch.speculative.plugins.hf_external import HFExternalDraftModel + +producer = _load_producer() + + +def test_record_round_trips_with_the_shift_the_loss_expects(tmp_path): + """Position t of the stored policy must be the distribution that produced ids[t].""" + torch.manual_seed(0) + n_prompt, n_gen = 4, 6 + ids = torch.randint(0, VOCAB, (n_prompt + n_gen,)) + logits = torch.randn(n_prompt + n_gen, VOCAB) * 2 + + record = producer.policy_record("c0", ids, n_prompt, ids.shape[0], logits, TOPK, 0.95) + shard = tmp_path / "p.jsonl.gz" + with gzip.open(shard, "wt") as f: + f.write(json.dumps(record) + "\n") + + ds = SparsePolicyDataset([str(shard)], max_length=64, top_k=TOPK) + assert len(ds) == 1 + item = ds[0] + assert torch.equal(item["input_ids"], ids.long()) + assert item["loss_mask"][:n_prompt].sum() == 0, "context must not be scored" + assert item["loss_mask"][n_prompt:].sum() == n_gen + + # The policy stored at sequence index t must be the one derived from logits[t - 1]. + for t in range(n_prompt, n_prompt + n_gen): + expected = logits[t - 1].softmax(-1).topk(TOPK).indices + assert torch.equal(item["teacher_topk_tok"][t], expected.long()), f"misaligned at {t}" + + # The probabilities have to survive the -inf -> exp -> float16 trip through the cache, + # which is the one place the producer's encoding meets the consumer's decoding. The + # tolerance is loose because the memmap stores float16. + scored = item["teacher_topk_prob"][n_prompt:] + assert torch.allclose(scored.sum(-1), torch.ones(n_gen), atol=1e-2) + # and must be the serving filter's distribution, which pins top-p as well as top-k + want = HFExternalDraftModel._deployment_probs(logits[n_prompt - 1 : -1], TOPK, 0.95) + got = torch.zeros_like(want).scatter(-1, item["teacher_topk_tok"][n_prompt:], scored.float()) + assert torch.allclose(got, want, atol=1e-3) + assert (item["teacher_topk_prob"][:n_prompt] == 0).all(), "context carries no policy" + + +def test_multi_turn_spans_do_not_swallow_the_user_turns(): + """Each assistant turn is its own record; the user turns between them are context.""" + # positions 3-5 and 9-11 are assistant turns; 6-8 is a user turn between them + mask = torch.zeros(1, 13) + mask[0, 3:6] = 1.0 + mask[0, 9:12] = 1.0 + assert producer.scored_spans(mask) == [(3, 6), (9, 12)] + # position 0 has no context to condition on, so a mask covering it starts at 1 + assert producer.scored_spans(torch.ones(1, 13)) == [(1, 13)] + + ids = torch.randint(0, VOCAB, (13,)) + logits = torch.randn(13, VOCAB) + first = producer.policy_record("c", ids, 3, 6, logits, TOPK, 0.95) + second = producer.policy_record("c", ids, 9, 12, logits, TOPK, 0.95) + + # the second turn's prompt carries the intervening user turn as context, unscored + assert len(first["gen_ids"]) == 3 and len(second["gen_ids"]) == 3 + assert len(second["prompt_ids"]) == 9