[6/8] Add a sparse teacher-policy dumper

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 |

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 '</think>' in content %}{% set content = content.split('</think>')[-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:

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 '</think>' in content %}{% set content = content.split('</think>')[-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)

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 '</think>' in content %}{% set content = content.split('</think>')[-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 = []