diff --git a/examples/speculative_decoding/eagle_utils.py b/examples/speculative_decoding/eagle_utils.py index a22438716a6..a0f76a02ef4 100644 --- a/examples/speculative_decoding/eagle_utils.py +++ b/examples/speculative_decoding/eagle_utils.py @@ -74,6 +74,14 @@ def make_speculative_data_module( aux_hidden_states_optional: accept dumps written with ``--no-aux-hidden-states``. Only modes that never read the aux planes may set this. """ + + # Validated once for every mode: each branch consumes sample_size differently + # (row select, shard slice, dataset limit) but the contract is the same. + # getattr: data_args is duck-typed across callers and some supply only the fields + # their own mode reads, so a hoisted check must not demand one they never set. + sample_size = getattr(data_args, "sample_size", -1) + if sample_size == 0 or sample_size < -1: + raise ValueError("sample_size must be -1 (use all samples) or a positive integer") # Load chat template from file if provided chat_template = None if getattr(data_args, "chat_template", None): @@ -155,6 +163,41 @@ def make_speculative_data_module( chat_template=chat_template, ) + elif mode == "sparse": + # Sparse teacher policy: the dump stores the base's top-k deployment + # distribution per response token instead of its hidden states, so no base + # forward (and no base lm_head) is needed at train time. Only the TVD + # objectives can consume it -- soft_ce wants full logits. + # Imported here, like the streaming dataset above, so the other modes do not + # pay for a module only this branch uses. + from modelopt.torch.speculative.external.sparse_data import ( + SparsePolicyCollator, + SparsePolicyDataset, + discover_shards, + ) + + print_rank_0("Loading sparse teacher-policy data for external draft training...") + assert not data_args.vlm_processor, "Sparse data is not supported for VLM." + # Shared with the pre-build instructions: a hand-built cache is only reused if it + # was keyed on the same shard list, so discovery lives in one place. + files = discover_shards(data_args.sparse_data_path) + if not files: + raise ValueError(f"No .jsonl.gz or .jsonl files found in {data_args.sparse_data_path}") + train_dataset = SparsePolicyDataset( + files, + max_length=train_len, + limit=data_args.sample_size if data_args.sample_size > 0 else None, + source_weights=data_args.sparse_source_weights, + cache_dir=data_args.sparse_cache_dir, + ) + print_rank_0( + f"Loaded {len(train_dataset)} records from {len(files)} shard(s), " + f"teacher top-k={train_dataset.top_k}" + ) + if data_args.sparse_source_weights: + print_rank_0(f"[sparse] source weights applied: {data_args.sparse_source_weights}") + data_collator = SparsePolicyCollator(train_len=train_len) + else: print_rank_0("Loading pre-processed data for offline training...") assert not data_args.vlm_processor, "Offline data is not supported for VLM." @@ -165,8 +208,6 @@ def make_speculative_data_module( raise ValueError(f"No .pt files found in {data_args.offline_data_path}") # sample_size=-1 means use all samples; positive integer selects that many - if data_args.sample_size == 0 or data_args.sample_size < -1: - raise ValueError("sample_size must be -1 (use all samples) or a positive integer") if data_args.sample_size > 0: dumped_files = dumped_files[: data_args.sample_size] train_dataset = OfflineSupervisedDataset( diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index 0d01cd81c48..f03ec0efa23 100644 --- a/examples/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -61,6 +61,7 @@ from modelopt.torch.speculative.config import ExternalDraftConfig from modelopt.torch.speculative.plugins.hf_dflash import HFDFlashModel from modelopt.torch.speculative.plugins.hf_domino import DominoLambdaCallback +from modelopt.torch.speculative.plugins.hf_external import SPARSE_CAPABLE_LOSSES from modelopt.torch.speculative.plugins.hf_training_args import ( TrainingArguments as SpecTrainingArgs, ) @@ -69,6 +70,7 @@ from modelopt.torch.speculative.utils import load_vlm_or_llm, patch_transformers5_params_loading from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.distributed import is_master, local_rank +from modelopt.torch.utils.logging import warn_rank_0 torch.manual_seed(0) mto.enable_huggingface_checkpointing() @@ -236,6 +238,14 @@ def setup_external_draft( model_max_length=training_args.training_seq_len, trust_remote_code=recipe.model.trust_remote_code, ) + # Sparse data carries only the teacher's truncated policy. Objectives that need full + # teacher logits would silently train against a zero-padded distribution, so gate on + # the list the plugin exports. + if recipe.data.mode == "sparse" and external_cfg["external_loss"] not in SPARSE_CAPABLE_LOSSES: + raise ValueError( + f"data.sparse_data_path requires external.external_loss in " + f"{list(SPARSE_CAPABLE_LOSSES)}, got {external_cfg['external_loss']!r}." + ) restored_modes = {m for m, _ in modelopt_state(draft)["modelopt_state_dict"]} if "external" in restored_modes: # The restore replayed the config stored in the checkpoint, so modify() has already @@ -269,16 +279,21 @@ def setup_external_draft( # The norm is not always at base_model.model.norm -- FakeBaseModel keeps it at .norm # -- so search the same paths the other offline modes use. It stays optional; # _teacher_logits raises only if the dump declares a pre-norm hidden. - norm_path = find_final_norm_path(base_model) - base_final_norm = base_model.get_submodule(norm_path) if norm_path else None - # Attach on the device training will run on: the draft is loaded with - # device_map="cpu", so otherwise the first forward moves the copies mid-step on - # every rank. training_args.device is the device accelerate assigned this process, - # which is correct both when every rank sees all GPUs and when CUDA_VISIBLE_DEVICES - # pins one GPU per rank -- local_rank arithmetic gets the second case wrong. - draft.attach_base_lm_head( - base_model.get_output_embeddings(), base_final_norm, device=training_args.device - ) + # The sparse path carries the teacher's policy in the data, so the base lm_head is + # dead weight on the GPU; the base is still loaded for the vocabulary check, which is + # cheap under use_fake_base_for_offline. + if recipe.data.mode != "sparse": + norm_path = find_final_norm_path(base_model) + base_final_norm = base_model.get_submodule(norm_path) if norm_path else None + # Attach on the device training will run on: the draft is loaded with + # device_map="cpu", so otherwise the first forward moves the copies mid-step on + # every rank. training_args.device is the device accelerate assigned this + # process, which is correct both when every rank sees all GPUs and when + # CUDA_VISIBLE_DEVICES pins one GPU per rank -- local_rank arithmetic gets the + # second case wrong. + draft.attach_base_lm_head( + base_model.get_output_embeddings(), base_final_norm, device=training_args.device + ) return draft, tokenizer @@ -453,6 +468,14 @@ def train(): print_rank_0("Loading dataset...") is_dflash = isinstance(recipe, ModelOptDFlashRecipe) + if recipe.data.mode == "sparse" and not isinstance(recipe, ModelOptExternalDraftRecipe): + # Only the external draft reads a sparse policy. Any other mode would build the + # dataset, hand the collator's teacher_topk_* to a forward that ignores them, and + # train on an all-zero loss -- so refuse instead of running an empty objective. + raise ValueError( + f"data.sparse_data_path is only supported by the external draft recipe; " + f"{type(recipe).__name__} cannot consume a sparse teacher policy." + ) data_module = make_speculative_data_module( tokenizer, recipe.data, @@ -465,6 +488,22 @@ def train(): # loudly on a dump that omits them. aux_hidden_states_optional=isinstance(recipe, ModelOptExternalDraftRecipe), ) + if recipe.data.mode == "sparse" and recipe.external.external_loss == "tvd_deploy": + # tvd_deploy filters the draft side by external_top_k, while the teacher side is + # whatever the producer baked into the dump. If the two disagree the objective + # compares two differently-truncated policies: the loss stays well-formed and + # still falls, so nothing surfaces except a draft that accepts worse than it + # should. Checked here because this is the first point where both the dump's + # stored top_k and the configured one are known. A wider draft filter can be a + # deliberate choice, so warn rather than refuse. + dump_top_k = getattr(data_module.get("train_dataset"), "top_k", None) + if dump_top_k and dump_top_k != recipe.external.external_top_k: + warn_rank_0( + f"the dump stores the teacher's top-{dump_top_k} policy but " + f"external.external_top_k={recipe.external.external_top_k} filters the draft; " + f"tvd_deploy then compares two differently-truncated policies. Set " + f"external.external_top_k={dump_top_k} to match the dump." + ) callbacks = [EagleTrainingPlot(training_args.ar_validate_steps, training_args.estimate_ar)] if ( diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index 9e792ea8ce9..eaaf37b62cb 100644 --- a/modelopt/torch/speculative/config.py +++ b/modelopt/torch/speculative/config.py @@ -484,7 +484,9 @@ class ExternalDraftConfig(ModeloptBaseConfig): "``tvd`` is total-variation distance between the two untruncated " "distributions. ``tvd_deploy`` puts both through the serving filter first " "(``external_top_k`` then ``external_top_p``), making it the acceptance the " - "deployed pair will see: ``alpha = 1 - TVD``." + "deployed pair will see: ``alpha = 1 - TVD``. Sparse data requires one of the " + "TVD objectives, since the dump is already a truncated policy and " + "``soft_ce`` needs full base logits." ), ) diff --git a/modelopt/torch/speculative/external/__init__.py b/modelopt/torch/speculative/external/__init__.py index 7af203db446..d7379c45d76 100644 --- a/modelopt/torch/speculative/external/__init__.py +++ b/modelopt/torch/speculative/external/__init__.py @@ -18,3 +18,4 @@ from .conversion import * from .external_model import * +from .sparse_data import * diff --git a/modelopt/torch/speculative/external/sparse_data.py b/modelopt/torch/speculative/external/sparse_data.py new file mode 100644 index 00000000000..2f1f7f6e2f2 --- /dev/null +++ b/modelopt/torch/speculative/external/sparse_data.py @@ -0,0 +1,556 @@ +# 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. + +"""Sparse teacher-policy dataset for external draft training. + +An alternative to caching base hidden states: the TVD objectives need only the +teacher's deployment policy (top-k then top-p at the serving temperature), which +is far smaller than a hidden state per token and is captured during generation +rather than by a second forward pass. + +Records are gzipped JSONL, one conversation per line:: + + { + "id": str, + "prompt_ids": [int], # context, not scored + "gen_ids": [int], # response tokens, scored + "topk_tok": [[int]], # per response position, teacher support + "topk_lp": [[float]], + } # matching log-probabilities + +They are decoded once into a memmapped cache beside the shards, stored as int32 +ids and float16 probabilities; holding them as Python objects does not fit in +memory for a corpus of any size. +""" + +import gzip +import json +import os +import shutil +import tempfile +import time +import warnings +from pathlib import Path + +import numpy as np +import torch + +from modelopt.torch.utils.logging import warn_rank_0 + +__all__ = ["SparsePolicyCollator", "SparsePolicyDataset", "discover_shards"] + +CACHE_VERSION = 2 + + +def _cache_signature(files: list[str]) -> list[list]: + """Identify the source shards so a stale cache is rebuilt rather than reused.""" + return [[str(f), os.path.getsize(f), int(os.path.getmtime(f))] for f in files] + + +def discover_shards(root) -> list[str]: + """Every sparse-policy shard under ``root``, in the order the cache is keyed on. + + Shared with the training entrypoint, and named in the pre-build instructions, so a + cache built by hand is keyed identically to the one training would build. A command + that globbed differently -- one pattern, or non-recursively -- would produce a + different ``_cache_signature``, and training would rebuild from scratch under + torchrun having apparently pre-built nothing. + + Both suffixes are combined rather than preferred: an ``or`` would silently drop every + ``.jsonl`` shard from a directory that also holds a ``.jsonl.gz``. + """ + root = Path(root) + return sorted( + str(p) + for pattern in ("*.jsonl.gz", "*.jsonl") + for p in root.rglob(pattern) + if not p.name.startswith(".") + ) + + +def _builder_progress(cache: Path) -> tuple[int, int]: + """Bytes written and newest mtime across rank 0's in-progress cache directories. + + The builder writes into a sibling ``.tmp.*`` directory and renames it into + place at the end, so its growth is the only signal that it is still alive. + """ + total = newest = 0 + for tmp in cache.parent.glob(f"{cache.name}.tmp.*"): + for f in tmp.glob("*"): + try: + st = f.stat() + except OSError: + continue + total += st.st_size + newest = max(newest, int(st.st_mtime)) + return total, newest + + +def _wait_for_cache(cache: Path, files: list[str], 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 + already on disk is that it was found stale, so ``meta.json`` is there from the + previous dump; returning on it would leave every rank but the builder memmapped to + the old shards -- and a mapping survives the rename, so they would train a whole run + on the wrong corpus against a perfectly ordinary loss curve. + + ``stall_s`` bounds time *without progress*, not the whole build. A first build over a + large corpus runs for hours, and a total deadline would fail every multi-GPU run that + has not pre-built its cache -- the one case this function exists to serve. Rank 0's + 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: + return + except (OSError, ValueError): + # Absent, or caught mid-publish; either way it is not ours yet. + pass + progress = _builder_progress(cache) + if progress != last_progress: + last_progress = progress + deadline = time.monotonic() + stall_s + elif time.monotonic() > deadline: + raise TimeoutError( + f"the sparse cache at {cache} has not grown for {stall_s:.0f}s and no cache " + "matching these shards was published. Rank 0 is not building: pre-build it " + "in a single process before torchrun, using the same discovery training " + "does, or the cache is keyed differently and rebuilt anyway:\n" + ' python -c "from modelopt.torch.speculative.external.sparse_data import ' + "SparsePolicyDataset as D, discover_shards as F; " + 'D(F(DUMP), max_length=SEQ_LEN, cache_dir=CACHE_DIR)"' + ) + time.sleep(2.0) + + +def _build_cache(cache: Path, files: list[str], top_k: int | None) -> None: + """Decode the JSONL shards into flat memmaps. One sequential pass.""" + # A private scratch directory per builder. Training several drafts against one + # dumped dataset is the normal workflow, so two jobs can reach this at the same + # time; a shared scratch name had each one delete the files the other still held + # open, leaving it writing to unlinked inodes and publishing an index whose + # offsets no longer matched its data -- a cache that looks complete and is not. + cache.parent.mkdir(parents=True, exist_ok=True) + tmp = Path(tempfile.mkdtemp(prefix=cache.name + ".tmp.", dir=cache.parent)) + try: + index: list[tuple[int, int, int, int]] = [] + srcs: list[str] = [] + n_ids = n_pos = dropped = 0 + mass_sum = mass_n = 0.0 + width = top_k + damaged: list[tuple[str, str]] = [] + with ( + open(tmp / "ids.i32", "wb") as f_ids, + open(tmp / "tok.i32", "wb") as f_tok, + open(tmp / "prob.f16", "wb") as f_prob, + ): + for path in files: + opener = gzip.open if str(path).endswith(".gz") else open + # A shard truncated by a killed writer must not destroy the whole pass: + # this loop is an hour of work on a large corpus, records are independent, + # and a retry would hit the same bad byte. Keep what parsed, report the rest. + try: + fh = opener(path, "rt") + except OSError as exc: + damaged.append((str(path), repr(exc))) + continue + with fh: + while True: + try: + line = fh.readline() + except (OSError, EOFError) as exc: + damaged.append((str(path), repr(exc))) + break + if not line: + break + try: + r = json.loads(line) + prompt, gen = r["prompt_ids"], r["gen_ids"] + # Both inside the guard: a record with an empty or absent + # policy must be dropped like any other damaged one, not + # abort a build the rest of this loop works to survive. + n_gen = min(len(gen), len(r["topk_tok"]), len(r["topk_lp"])) + if width is None: + # Widest row of the record, not row 0. With top_p < 1 + # the stored nucleus varies per position, and row 0 is + # the first generated token -- conditioned on the whole + # prompt, so typically the lowest-entropy and narrowest + # row there is. Inferring from it clips every wider row + # in the corpus. + width = max(len(x) for x in r["topk_tok"]) + except (ValueError, KeyError, IndexError) as exc: + dropped += 1 + if dropped == 1: + damaged.append((str(path), f"unparseable record: {exc!r}")) + continue + if n_gen == 0: + # Nothing supervised: an index entry for it would contribute an + # empty span to every epoch. + dropped += 1 + if dropped == 1: + damaged.append((str(path), "record has no scored positions")) + continue + + ids = np.asarray(prompt + gen[:n_gen], dtype=np.int32) + # Ragged rows would silently misalign the flat memmap, so pad or + # truncate every row to the one width the cache is indexed by. + # Dumps are normally rectangular, so try the whole block in one + # numpy call first: the per-position loop below costs 500M+ + # iterations on a 100k-conversation corpus. + raw_t, raw_p = r["topk_tok"][:n_gen], r["topk_lp"][:n_gen] + try: + tok = np.array(raw_t, dtype=np.int32) + # exp(-inf) = 0 for tokens the teacher's nucleus excluded. + prob = np.exp(np.array(raw_p, dtype=np.float32)).astype(np.float16) + if tok.ndim != 2 or tok.shape[1] != width: + raise ValueError("width mismatch") + except (ValueError, TypeError): + widest = max((len(x) for x in raw_t), default=0) + if widest > width: + # Truncating here would drop real teacher mass and bake + # the narrow width into meta["top_k"], and the stored-mass + # warning cannot see it. Refuse with the remedy instead. + # The recipe has no top_k field, so the remedy is the + # single-process pre-build, which takes one. + raise ValueError( + f"{path}: a stored policy row holds {widest} entries " + f"but the cache was sized at {width} from the first " + f"record. Build the cache once with an explicit " + f"width, then point the run at it: " + f"SparsePolicyDataset(discover_shards(dump), " + f"max_length=..., top_k=, cache_dir=)." + ) + tok = np.zeros((n_gen, width), dtype=np.int32) + prob = np.zeros((n_gen, width), dtype=np.float16) + for j in range(n_gen): + t, lp = raw_t[j][:width], raw_p[j][:width] + tok[j, : len(t)] = t + prob[j, : len(lp)] = np.exp(np.asarray(lp, dtype=np.float32)) + + # tvd treats the stored policy as the teacher's whole + # distribution, so it must not be a truncation of a wider + # nucleus. Sampled with the same top_k the dump requested, + # this sums to 1; built with a smaller k it does not, and the + # objective then charges the draft for mass on tokens the + # teacher actually supported. + if n_gen: + # Whole record, not prob[0]: the first response token is + # conditioned on the prompt alone and is typically the + # lowest-entropy position in the generation, so its top-k mass + # is near 1 even when a too-narrow top_k is badly clipping the + # nucleus at later, higher-entropy positions -- the condition + # this check exists to catch. + mass_sum += float(prob.sum(axis=-1).mean()) + mass_n += 1 + + f_ids.write(ids.tobytes()) + f_tok.write(tok.tobytes()) + f_prob.write(prob.tobytes()) + index.append((n_ids, len(prompt), n_gen, n_pos)) + srcs.append(r.get("source") or "") + n_ids += ids.shape[0] + n_pos += n_gen + + if not index: + raise ValueError(f"No sparse-policy records found in {len(files)} file(s)") + + if damaged: + for path, why in damaged[:10]: + warnings.warn(f"sparse shard {path}: {why}") + # Tolerate the odd truncated shard, but a widespread failure means the dump + # is wrong and training on it would quietly use a fraction of the corpus. + # `damaged` holds at most one entry per *build* for record-level damage + # (the dropped == 1 guards above), so it cannot express "most records + # failed to parse". A dump whose schema drifted -- a renamed topk_lp, say + # -- raises KeyError on every record, warns once, passes the shard check + # below, and publishes a cache holding a handful of rows. Gate on the + # record drop rate too, which is what the shard threshold is reasoning + # about anyway. + kept = len(index) + if dropped > max(1, (kept + dropped) // 20): + raise ValueError( + f"{dropped} of {kept + dropped} sparse-policy records failed to parse " + f"(kept {kept}); the dump's schema does not match what this reader " + "expects." + ) + if len(damaged) > max(1, len(files) // 20): + raise ValueError( + f"{len(damaged)} of {len(files)} sparse shards are unreadable; " + "the dump looks broken rather than merely truncated." + ) + print( + f"[sparse] skipped {len(damaged)} damaged shard(s) and {dropped} record(s); " + f"kept {len(index)}", + flush=True, + ) + + if mass_n and mass_sum / mass_n < 0.999: + warnings.warn( + f"stored teacher policy sums to {mass_sum / mass_n:.4f}, not ~1: the dump's " + "top-k is narrower than the nucleus it was sampled from. The TVD objectives " + "renormalise over the stored support, so the draft is charged for mass on " + "tokens the teacher supported but the dump dropped; even a draft matching the " + "teacher scores a non-zero loss. Re-dump with a larger top-k.", + stacklevel=2, + ) + + np.save(tmp / "index.npy", np.asarray(index, dtype=np.int64)) + # Per-record source, kept as codes into source_names so the dataset can weight + # sources without re-reading the shards. + source_names = sorted(set(srcs)) + codes = {name: i for i, name in enumerate(source_names)} + np.save(tmp / "sources.npy", np.asarray([codes[x] for x in srcs], dtype=np.int16)) + (tmp / "meta.json").write_text( + json.dumps( + { + "version": CACHE_VERSION, + "top_k": width, + "n_records": len(index), + "n_ids": n_ids, + "n_pos": n_pos, + "sources": _cache_signature(files), + "source_names": source_names, + } + ) + ) + # Publish by rename: a half-written cache that looks complete is worse than + # none. os.rename onto an existing non-empty directory raises ENOTEMPTY, so + # move any stale cache aside first and only delete it once the new one is in + # place -- a rebuild must never leave the shards with no usable cache. + # The hand-off name is per-builder too, for the same reason as the scratch one. + stale = cache.with_name(f"{cache.name}.stale.{tmp.name.rsplit('.', 1)[-1]}") + shutil.rmtree(stale, ignore_errors=True) + try: + if cache.exists(): + cache.rename(stale) + tmp.rename(cache) + except OSError: + # Another builder published between the check above and this rename, so + # the name is taken by a cache built from the same shards. Drop ours + # rather than fight for it -- but never leave the shards with no cache. + if not cache.exists() and stale.exists(): + stale.rename(cache) + else: + shutil.rmtree(tmp, ignore_errors=True) + finally: + shutil.rmtree(stale, ignore_errors=True) + except BaseException: + # Leave no half-decoded scratch behind; the name is unique per builder, + # so this can only ever remove our own. + shutil.rmtree(tmp, ignore_errors=True) + raise + + +class SparsePolicyDataset(torch.utils.data.Dataset): + """Conversations with the teacher's sparse deployment policy per response token.""" + + def __init__( + self, + files: list[str], + max_length: int = 4096, + top_k: int | None = None, + cache_dir: str | None = None, + limit: int | None = None, + source_weights: dict[str, float] | None = None, + ): + """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. + + ``source_weights`` maps a record's ``source`` field to a sampling weight, + oversampling under-represented sources by repeating their index entries. + Weight by *tokens*, not rows: sources differ by an order of magnitude in + response length, so balanced row counts can still leave one source holding + most of the trained positions. + """ + super().__init__() + if not files: + raise ValueError("No sparse-policy shards given") + self.max_length = max_length + cache = Path(cache_dir) if cache_dir else Path(files[0]).parent / ".sparse_cache" + + 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) + ) + + if stale: + # Under DDP every rank sees the same shards; one builds and the rest wait on + # the files rather than on a barrier. A collective here would be bounded by + # ddp_timeout (1800s by default, and HfTrainingArguments has already created + # the process group with it by the time a dataset is built), which a first + # build over a large corpus outlasts. _wait_for_cache needs no process group + # and bounds time without progress instead, so it serves both cases. + # Fall back to the launcher's RANK when there is no group: without it every + # worker takes the rank-0 branch and they race to rewrite the same cache. + if torch.distributed.is_initialized(): + rank = torch.distributed.get_rank() + else: + rank = int(os.environ.get("RANK", 0)) + if rank == 0: + _build_cache(cache, files, top_k) + else: + _wait_for_cache(cache, files) + + meta = json.loads((cache / "meta.json").read_text()) + # Re-check rather than trust the build. Two jobs reading different subsets of one + # dumped directory resolve to the same default cache path, and the one that loses + # 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): + raise RuntimeError( + f"the sparse cache at {cache} was published by a concurrent job from " + "different shards; give this run its own cache_dir." + ) + self.top_k = meta["top_k"] + self.index = np.load(cache / "index.npy") + self.source_names = meta.get("source_names", []) + self.sources = None + if (cache / "sources.npy").exists(): + self.sources = np.load(cache / "sources.npy")[: self.index.shape[0]] + # __getitem__ right-truncates to max_length, so a record whose context alone + # fills it keeps no scored position: an all-zero loss_mask that still counts + # towards the record total and the loss average. Report it rather than let the + # run look larger than it is. Counted before any oversampling, so the number + # describes the corpus the advice below asks the user to re-dump, not the + # repeated index entries. Rank 0 only: every rank builds this object. + n_dead = int((self.index[:, 1] >= self.max_length).sum()) + if n_dead: + warn_rank_0( + f"{n_dead} of {self.index.shape[0]} sparse-policy records have a context of " + f"{self.max_length} tokens or more and so contribute no scored positions at " + f"max_length={self.max_length}. Raise training_seq_len, or re-dump with a " + f"smaller --max-seq-len." + ) + if source_weights: + self.index = self._apply_source_weights(source_weights) + if limit is not None and limit < self.index.shape[0]: + # A seeded sample, not index[:limit]. index.npy is written in shard-then-line + # order, so a prefix is a prefix of the shard list: with one source per shard + # -- the natural layout -- it can contain a single source, and weighting that + # gives a mixture nothing like the one asked for. Taken after weighting so the + # sample is drawn from the mixture the user configured. + take = np.sort( + np.random.default_rng(0).choice(self.index.shape[0], limit, replace=False) + ) + self.index = self.index[take] + if self.sources is not None: + self.sources = self.sources[take] + self._ids = np.memmap(cache / "ids.i32", dtype=np.int32, mode="r", shape=(meta["n_ids"],)) + self._tok = np.memmap( + cache / "tok.i32", dtype=np.int32, mode="r", shape=(meta["n_pos"], self.top_k) + ) + self._prob = np.memmap( + cache / "prob.f16", dtype=np.float16, mode="r", shape=(meta["n_pos"], self.top_k) + ) + + def _apply_source_weights(self, source_weights: dict[str, float]) -> np.ndarray: + """Repeat index entries so each source's share matches its weight.""" + if self.sources is None: + raise ValueError( + "source_weights needs a cache built with CACHE_VERSION >= 2; " + "delete the .sparse_cache directory to rebuild it" + ) + unknown = set(source_weights) - set(self.source_names) + if unknown: + raise ValueError( + f"source_weights names {sorted(unknown)} which are not in the data " + f"{self.source_names}" + ) + w = np.ones(self.index.shape[0], dtype=np.float64) + for name, weight in source_weights.items(): + if weight < 0: + raise ValueError(f"source weight for {name!r} is negative") + w[self.sources == self.source_names.index(name)] = weight + # Split the weight into guaranteed repeats plus a seeded draw for the + # remainder, so a weight of 2.5 gives every record 2 copies and half of + # them a third -- deterministic across ranks and restarts. + counts = np.floor(w).astype(np.int64) + frac = w - counts + counts += (np.random.default_rng(0).random(w.shape[0]) < frac).astype(np.int64) + # sources indexes index row-for-row and is read back by callers (and by the + # sampling below), so it has to be repeated identically or the two desynchronise. + self.sources = np.repeat(self.sources, counts, axis=0) + return np.repeat(self.index, counts, axis=0) + + def __len__(self): + return self.index.shape[0] + + def __getitem__(self, i): + ids_start, n_prompt, n_gen, pos_start = (int(x) for x in self.index[i]) + n = min(n_prompt + n_gen, self.max_length) + input_ids = torch.from_numpy( + np.asarray(self._ids[ids_start : ids_start + n], dtype=np.int64) + ) + + # Only response positions carry a teacher policy, so only they are scored. + loss_mask = torch.zeros(n, dtype=torch.float) + loss_mask[n_prompt:n] = 1.0 + + tok = torch.zeros(n, self.top_k, dtype=torch.long) + prob = torch.zeros(n, self.top_k, dtype=torch.float) + keep = max(0, n - n_prompt) + if keep: + sl = slice(pos_start, pos_start + keep) + tok[n_prompt:n] = torch.from_numpy(np.asarray(self._tok[sl], dtype=np.int64)) + prob[n_prompt:n] = torch.from_numpy(np.asarray(self._prob[sl], dtype=np.float32)) + + return { + "input_ids": input_ids, + "attention_mask": torch.ones_like(input_ids), + "loss_mask": loss_mask, + "teacher_topk_tok": tok, + "teacher_topk_prob": prob, + } + + +class SparsePolicyCollator: + """Pad or truncate a batch of sparse-policy records to ``train_len``.""" + + def __init__(self, train_len: int): + """Set the fixed sequence length every batch is padded or truncated to.""" + self.train_len = train_len + + def _fit(self, x): + n = x.shape[0] + if n >= self.train_len: + return x[: self.train_len] + pad = torch.zeros((self.train_len - n, *x.shape[1:]), dtype=x.dtype) + return torch.cat([x, pad], dim=0) + + def __call__(self, features): + """Collate records into a padded batch.""" + return { + k: torch.stack([self._fit(f[k]) for f in features]) + for k in ( + "input_ids", + "attention_mask", + "loss_mask", + "teacher_topk_tok", + "teacher_topk_prob", + ) + } diff --git a/modelopt/torch/speculative/plugins/hf_external.py b/modelopt/torch/speculative/plugins/hf_external.py index 97dbdf46c08..e9095a640f8 100644 --- a/modelopt/torch/speculative/plugins/hf_external.py +++ b/modelopt/torch/speculative/plugins/hf_external.py @@ -37,7 +37,16 @@ from ..external.external_model import ExternalDraftModel from .modeling_final_norm import _maybe_apply_base_final_norm -__all__ = ["HFExternalDraftModel"] +# Objectives with a sparse implementation: they read only the teacher's stored +# top-k policy. soft_ce needs full teacher logits and so cannot run on that data. +# main.py gates data.sparse_data_path on this, so it must be defined beside the +# implementations rather than duplicated there. +# On the sparse path the base's truncation is baked into the dump, so both TVD +# variants are computable there; they differ only in whether the *draft* is put +# through the serving filter too. soft_ce needs full base logits and is dense-only. +SPARSE_CAPABLE_LOSSES = ("tvd", "tvd_deploy") + +__all__ = ["SPARSE_CAPABLE_LOSSES", "HFExternalDraftModel"] def _train_acc(acc): @@ -121,6 +130,8 @@ def forward( attention_mask=None, loss_mask=None, base_model_outputs=None, + teacher_topk_tok=None, + teacher_topk_prob=None, **kwargs, ): """Score the draft against cached base hidden states. @@ -128,24 +139,32 @@ def forward( Only ``base_model_hidden_states`` is consumed; ``aux_hidden_states`` exists in the dump for the head-based modes and is not used here. """ - if base_model_outputs is None: + sparse = teacher_topk_tok is not None + if not sparse and base_model_outputs is None: raise ValueError( "External draft training expects offline base_model_outputs in the batch; " "online training against a live base model is not supported." ) draft_out = super().forward(input_ids=input_ids, attention_mask=attention_mask, **kwargs) - base_logits = self._teacher_logits( - base_model_outputs["base_model_hidden_states"], base_model_outputs - ) if loss_mask is None: # Derive from attention_mask so padding is not trained on, matching the # other speculative plugins; all-ones only when there is no mask at all. + # Both paths default identically: the sparse collator always supplies a mask, + # but when it does not, training on padding is as wrong here as it is there. loss_mask = ( attention_mask.clone().detach() if attention_mask is not None else torch.ones_like(input_ids) ).to(draft_out.logits.dtype) - loss, acc = self.compute_loss(draft_out.logits, base_logits, loss_mask) + if sparse: + loss, acc = self.compute_sparse_loss( + draft_out.logits, teacher_topk_tok, teacher_topk_prob, loss_mask + ) + else: + base_logits = self._teacher_logits( + base_model_outputs["base_model_hidden_states"], base_model_outputs + ) + loss, acc = self.compute_loss(draft_out.logits, base_logits, loss_mask) draft_out.loss = loss if acc is not None: draft_out.train_acc = _train_acc(acc) @@ -270,6 +289,70 @@ def _deployment_probs(logits, top_k, top_p): p = p * torch.zeros_like(keep).scatter(-1, idx, keep) return p / p.sum(-1, keepdim=True).clamp_min(1e-9) + def _draft_probs(self, draft_logits): + """The draft's distribution under the configured objective. + + ``tvd_deploy`` puts the draft through the serving filter so the loss is the + acceptance the deployed pair will see; the other objectives use it unfiltered. + Shared so the two loss paths cannot drift apart on which filter they applied. + """ + if self.external_loss == "tvd_deploy": + return self._deployment_probs(draft_logits, self.external_top_k, self.external_top_p) + return torch.softmax(draft_logits.float(), dim=-1) + + @staticmethod + def _top1_accuracy(teacher_top1, draft_logits, mask): + """Fraction of scored positions where the draft's argmax matches the teacher's.""" + with torch.no_grad(): + valid = mask.bool() + correct = (teacher_top1 == draft_logits.detach().argmax(dim=-1)) & valid + return correct.sum().float() / valid.sum().clamp_min(1).float() + + def _tvd_sparse(self, draft_logits, topk_tok, topk_prob): + """Plain TVD against the stored teacher policy, over the whole vocabulary. + + The draft is not renormalised over the teacher's support, so probability + it places outside that support counts as error. The stored policy is the + teacher's deployment distribution and is exactly zero off its own support, + which makes the out-of-support term simply ``1 - sum(q_topk)``. + """ + p = topk_prob / topk_prob.sum(-1, keepdim=True).clamp_min(1e-9) + q = self._draft_probs(draft_logits).gather(-1, topk_tok) + # A ragged dump is padded to a rectangular width, and the pad slots must not + # charge the draft again for a token the row already holds. Pads identify + # themselves: both builder paths zero-initialise, so a padded slot is exactly + # one with zero stored probability. Keying on that rather than on id equality + # needs no assumption about builder ordering, and leaves a genuinely repeated + # id alone -- an id-equality collapse would drop its second occurrence from q + # while p kept both entries. + q = q * (topk_prob > 0) + # q sums to 1 over its own support, so the mass it placed off the base's + # support is whatever is left over. + outside = (1.0 - q.sum(-1)).clamp_min(0.0) + return 0.5 * ((p - q).abs().sum(-1) + outside) + + def compute_sparse_loss(self, draft_logits, topk_tok, topk_prob, loss_mask): + """``tvd`` against a sparse teacher policy; returns ``(loss, accuracy)``. + + The sparse policy is indexed by the token it produced: position ``t`` holds the + distribution the base sampled ``input_ids[t]`` from, conditioned on ``< t``. The + draft's logits at ``t`` are conditioned on ``<= t`` and predict ``t + 1``, so the + draft is shifted one position left before the two are compared. Dropping the shift + still yields a falling loss curve -- it just trains against the wrong target. + """ + seq = min(draft_logits.shape[1], topk_tok.shape[1]) + draft = draft_logits[:, : seq - 1] + tok, prob = topk_tok[:, 1:seq], topk_prob[:, 1:seq] + mask = loss_mask[:, 1:seq].to(draft_logits.dtype) + denom = mask.sum() + 1e-5 + tvd = self._tvd_sparse(draft, tok, prob) + loss = (tvd * mask).sum() / denom + accuracy = None + if self.external_report_acc: + teacher_top1 = tok.gather(-1, prob.argmax(-1, keepdim=True)).squeeze(-1) + accuracy = self._top1_accuracy(teacher_top1, draft, mask) + return loss, accuracy + def compute_loss(self, draft_logits, base_logits, loss_mask, base_predict_tok=None): """Compute the configured objective against the base model's distribution. @@ -294,10 +377,9 @@ def compute_loss(self, draft_logits, base_logits, loss_mask, base_predict_tok=No # Both sides through the serving filter, so the loss is the # acceptance the deployed pair will actually see. p = self._deployment_probs(base_logits, self.external_top_k, self.external_top_p) - q = self._deployment_probs(draft_logits, self.external_top_k, self.external_top_p) else: p = torch.softmax(base_logits.float(), dim=-1) - q = torch.softmax(draft_logits.float(), dim=-1) + q = self._draft_probs(draft_logits) loss = (0.5 * (p - q).abs().sum(-1) * mask).sum() / denom else: # Identical in form to hf_eagle._eagle_loss so switching modes does not @@ -308,11 +390,8 @@ def compute_loss(self, draft_logits, base_logits, loss_mask, base_predict_tok=No accuracy = None if self.external_report_acc: - with torch.no_grad(): - if base_predict_tok is None: - base_predict_tok = base_logits.argmax(dim=-1) - valid = mask.bool() - correct = (base_predict_tok == draft_logits.detach().argmax(dim=-1)) & valid - accuracy = correct.sum().float() / valid.sum().clamp_min(1).float() + if base_predict_tok is None: + base_predict_tok = base_logits.argmax(dim=-1) + accuracy = self._top1_accuracy(base_predict_tok, draft_logits, mask) return loss, accuracy diff --git a/modelopt/torch/speculative/plugins/hf_training_args.py b/modelopt/torch/speculative/plugins/hf_training_args.py index b121d9f260d..917f300e215 100644 --- a/modelopt/torch/speculative/plugins/hf_training_args.py +++ b/modelopt/torch/speculative/plugins/hf_training_args.py @@ -58,9 +58,18 @@ class DataArguments(BaseModel): # Derived in ``_check_mode_requirements`` from the data-source fields; accepted as input # only for backward compatibility (existing ``mode:`` keys / overrides), then overwritten. - mode: Literal["online", "offline", "streaming"] | None = None + mode: Literal["online", "offline", "streaming", "sparse"] | None = None data_path: str | None = None offline_data_path: str | None = None + # Dumped top-k teacher deployment policy (external draft mode); see sparse_data.py. + sparse_data_path: str | None = None + # Oversample sources by their ``source`` field, e.g. {"cnn": 5.0}. Weight by tokens, + # not rows: response lengths differ by an order of magnitude across sources. + sparse_source_weights: dict[str, float] | None = None + # Where to build the memmapped cache. Defaults to a .sparse_cache beside the shards, + # which fails when the dump lives on a read-only mount -- the normal case for a shared + # corpus -- and leaves concurrent runs over the same shards fighting over one cache. + sparse_cache_dir: str | None = None lazy_preprocess: bool = True draft_vocab_cache: str | None = None chat_template: str | None = None @@ -89,11 +98,21 @@ def _check_mode_requirements(self) -> DataArguments: # override round-trip and silently select the wrong training path. has_offline = self.offline_data_path is not None has_streaming = self.streaming_server_url is not None - if has_offline and has_streaming: + has_sparse = self.sparse_data_path is not None + if sum([has_offline, has_streaming, has_sparse]) > 1: raise ValueError( - "ambiguous: set only one of data.offline_data_path / data.streaming_server_url" + "ambiguous: set only one of data.offline_data_path / " + "data.streaming_server_url / data.sparse_data_path" ) - self.mode = "offline" if has_offline else "streaming" if has_streaming else "online" + self.mode = ( + "sparse" + if has_sparse + else "offline" + if has_offline + else "streaming" + if has_streaming + else "online" + ) if self.mode == "streaming" and not self.streaming_model_name: raise ValueError( "data.mode='streaming' requires data.streaming_server_url and " diff --git a/modelopt_recipes/general/speculative_decoding/external_draft.yaml b/modelopt_recipes/general/speculative_decoding/external_draft.yaml index 106d9072b79..c7415a9faf2 100644 --- a/modelopt_recipes/general/speculative_decoding/external_draft.yaml +++ b/modelopt_recipes/general/speculative_decoding/external_draft.yaml @@ -25,6 +25,23 @@ data: mode: data_path: offline_data_path: + # Dumped top-k teacher deployment policy. Selects data.mode=sparse, which needs no + # base lm_head at train time and requires a TVD objective (tvd or tvd_deploy). + # + # The first run over a dump builds a memmapped cache beside the shards, which takes + # hours on a large corpus. Under torchrun only rank 0 builds it and the others wait, + # so build it once in a single process first: + # python -c "from modelopt.torch.speculative.external.sparse_data import \ + # SparsePolicyDataset as D, discover_shards as F; \ + # D(F('DUMP'), max_length=12288, cache_dir=None)" + # Use discover_shards and the same cache_dir training will use: a cache keyed on a + # different shard list is ignored, and torchrun rebuilds it from scratch. A ragged + # third-party dump also needs top_k= here, since the cache width is + # otherwise inferred from the first record and there is no recipe field for it. + sparse_data_path: + # Where that cache lives. Defaults to a .sparse_cache beside the shards, which fails + # on a read-only dump and makes concurrent runs over the same shards contend. + sparse_cache_dir: draft_vocab_cache: vlm_img_dir: vlm_processor: @@ -54,5 +71,11 @@ external: # soft_ce matches the EAGLE offline objective. tvd is plain total-variation # distance; tvd_deploy filters both sides to the serving policy first and is the # acceptance the deployed pair sees. + # + # How much of the base's distribution the loss can see depends on the data: + # offline_data_path hidden states -> base lm_head -> the FULL distribution, + # so tvd is untruncated total-variation distance + # sparse_data_path only the dumped top-k policy, so the base side is + # truncated by the dump no matter which objective is used external_loss: tvd external_report_acc: true diff --git a/tests/unit/examples/test_speculative_decoding_main.py b/tests/unit/examples/test_speculative_decoding_main.py index 13cd3fc3597..c73ead5659f 100644 --- a/tests/unit/examples/test_speculative_decoding_main.py +++ b/tests/unit/examples/test_speculative_decoding_main.py @@ -264,3 +264,33 @@ def test_resuming_a_different_speculative_mode_is_refused(main_mod, stub_tokeniz base_model=_tiny_causal_lm(hidden=2 * HIDDEN), draft=restored, ) + + +def test_sparse_requires_a_tvd_objective(main_mod, stub_tokenizers): + """soft_ce needs full teacher logits; a sparse dump only stores the truncated policy. + + Without the gate the run trains against a zero-padded distribution and the loss curve + looks ordinary, so the refusal has to happen before training starts. + """ + recipe = _recipe(sparse_data_path="/tmp/policy") + recipe.external.external_loss = "soft_ce" + with pytest.raises(ValueError, match="external_loss"): + main_mod.setup_external_draft( + recipe, + _training_args(), + base_model=_tiny_causal_lm(hidden=2 * HIDDEN), + draft=_tiny_causal_lm(), + ) + + +def test_sparse_path_skips_the_base_head_attach(main_mod, stub_tokenizers): + """The sparse dump carries the teacher's policy, so the base lm_head is never needed.""" + draft, _ = main_mod.setup_external_draft( + _recipe(sparse_data_path="/tmp/policy"), + _training_args(), + base_model=_tiny_causal_lm(hidden=2 * HIDDEN), + draft=_tiny_causal_lm(), + ) + assert getattr(draft, "_base_lm_head", None) is None, ( + "the sparse path attached the base head it does not use" + ) diff --git a/tests/unit/torch/speculative/test_external_draft_sparse.py b/tests/unit/torch/speculative/test_external_draft_sparse.py new file mode 100644 index 00000000000..c91d9619cc6 --- /dev/null +++ b/tests/unit/torch/speculative/test_external_draft_sparse.py @@ -0,0 +1,562 @@ +# 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. + +"""Sparse teacher-policy path for the ``external`` speculative decoding mode.""" + +import contextlib +import gzip +import json +import math +import threading +from unittest import mock + +import pytest +import torch + +import modelopt.torch.speculative as mtsp +from modelopt.torch.speculative.external import sparse_data +from modelopt.torch.speculative.external.sparse_data import ( + SparsePolicyCollator, + SparsePolicyDataset, +) + +transformers = pytest.importorskip("transformers") + +# Below the guard: the plugins package registers these behind an import_plugin +# guard, but importing the submodules directly does not, so on an install +# without transformers this module would fail to collect rather than skip. +from modelopt.torch.speculative.plugins.hf_external import SPARSE_CAPABLE_LOSSES +from modelopt.torch.speculative.plugins.hf_training_args import DataArguments + +VOCAB = 64 +HIDDEN = 32 +TOPK = 8 +# Marks a key the parametrization removes from a record rather than overwrites. +_DROP = object() + + +def _causal_lm(hidden=HIDDEN, vocab_size=VOCAB): + cfg = transformers.LlamaConfig( + vocab_size=vocab_size, + hidden_size=hidden, + intermediate_size=2 * hidden, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + max_position_embeddings=128, + ) + return transformers.LlamaForCausalLM(cfg) + + +def _draft(**overrides): + m = _causal_lm() + mtsp.convert(m, [("external", {"external_loss": "tvd_deploy", **overrides})]) + return m + + +def _write_records(path, n=4, n_prompt=3, n_gen=5): + torch.manual_seed(0) + with gzip.open(path, "wt") as f: + for i in range(n): + lp = torch.randn(n_gen, TOPK).log_softmax(-1) + f.write( + json.dumps( + { + "id": f"r{i}", + "prompt_ids": torch.randint(0, VOCAB, (n_prompt,)).tolist(), + "gen_ids": torch.randint(0, VOCAB, (n_gen,)).tolist(), + "topk_tok": torch.randint(0, VOCAB, (n_gen, TOPK)).tolist(), + "topk_lp": lp.tolist(), + } + ) + + "\n" + ) + + +def test_dataset_masks_response_positions_and_pads_to_train_len(tmp_path): + """Only response positions carry a teacher policy, and top_k is read from the data. + + A config top_k narrower than the records would silently truncate the policy. + """ + p = tmp_path / "s.jsonl.gz" + _write_records(p, n=3, n_prompt=3, n_gen=5) + ds = SparsePolicyDataset([str(p)], max_length=32) + assert ds.top_k == TOPK + + item = ds[0] + assert item["input_ids"].shape[0] == 8 + assert item["loss_mask"][:3].sum() == 0 + assert item["loss_mask"][3:8].sum() == 5 + + batch = SparsePolicyCollator(train_len=16)([ds[i] for i in range(3)]) + assert batch["input_ids"].shape == (3, 16) + assert batch["teacher_topk_tok"].shape == (3, 16, TOPK) + assert batch["teacher_topk_prob"].shape == (3, 16, TOPK) + + +# The pure-TVD objectives, named rather than taken from SPARSE_CAPABLE_LOSSES: the +# reference below is the TVD formula, and a later objective that adds a term to it +# would silently join this parametrization and fail against the wrong reference. +@pytest.mark.parametrize("loss_fn", ["tvd", "tvd_deploy"]) +def test_sparse_tvd_matches_one_minus_acceptance(loss_fn): + """The sparse path must compute the same acceptance objective as the dense one. + + The reference below is the whole contract: it applies the one-position shift, the + per-objective draft filter, and the out-of-support term. Agreeing with it on random + inputs therefore also pins next-token alignment, the zero at draft == stored policy, + and the penalty on mass the draft places off the teacher's support. + """ + model = _draft(external_loss=loss_fn, external_report_acc=False) + torch.manual_seed(0) + draft_logits = torch.randn(1, 4, VOCAB) + # Distinct ids per position. randint would collide most rows out of 8-of-64, and + # _tvd_sparse deduplicates a repeated id -- which the reference below deliberately + # does not model, because the dedup has its own test. (tvd_deploy hides the + # difference: its filter zeroes the duplicate slots anyway.) + tok = torch.stack([torch.randperm(VOCAB)[:TOPK] for _ in range(4)]).unsqueeze(0) + prob = torch.randn(1, 4, TOPK).softmax(-1) + + loss, _ = model.compute_sparse_loss(draft_logits, tok, prob, torch.ones(1, 4)) + + # Teacher position t is scored against draft position t-1; see compute_sparse_loss. + # p is zero off its support, so TVD over the full vocabulary is the on-support + # difference plus everything the draft placed elsewhere. + p = prob[:, 1:] / prob[:, 1:].sum(-1, keepdim=True) + q_full = ( + model._deployment_probs(draft_logits[:, :-1], model.external_top_k, model.external_top_p) + if loss_fn == "tvd_deploy" + else torch.softmax(draft_logits[:, :-1].float(), dim=-1) + ) + q = q_full.gather(-1, tok[:, 1:]) + expected = 0.5 * ((p - q).abs().sum(-1) + (1.0 - q.sum(-1)).clamp_min(0.0)) + assert torch.allclose(loss, expected.mean(), atol=1e-5) + + +def test_every_sparse_capable_loss_runs_and_backprops(tmp_path): + """The guard in main.py and the implementations here must not drift apart, and the + sparse path must train without a teacher lm_head.""" + p = tmp_path / "s.jsonl.gz" + _write_records(p, n=2, n_prompt=2, n_gen=4) + ds = SparsePolicyDataset([str(p)], max_length=32, top_k=TOPK) + batch = SparsePolicyCollator(train_len=8)([ds[0], ds[1]]) + + for loss_fn in SPARSE_CAPABLE_LOSSES: + model = _draft(external_loss=loss_fn, external_report_acc=False) + assert getattr(model, "_base_lm_head", None) is None, "sparse path needs no teacher head" + out = model(**batch) + assert torch.isfinite(out.loss), f"{loss_fn} produced {out.loss}" + out.loss.backward() + assert any(q.grad is not None and q.grad.abs().sum() > 0 for q in model.parameters()) + + +@pytest.mark.parametrize( + ("kind", "expect"), + [("truncated", "warn"), ("unparseable", "warn"), ("widespread", "raise"), ("empty", "raise")], +) +def test_damaged_shards(tmp_path, kind, expect): + """One bad shard must not discard an hour of parsing, but losing most of the corpus + is a broken dump rather than a truncation to tolerate.""" + files = [] + if kind == "empty": + p = tmp_path / "empty.jsonl.gz" + gzip.open(p, "wt").close() + files = [str(p)] + elif kind == "unparseable": + p = tmp_path / "s.jsonl.gz" + _write_records(p, n=2) + with gzip.open(p, "at") as f: + f.write("{not json}\n") + files = [str(p)] + else: + good = tmp_path / "good.jsonl.gz" + _write_records(good, n=3 if kind == "truncated" else 1) + files = [str(good)] + for i in range(1 if kind == "truncated" else 6): + bad = tmp_path / f"bad{i}.jsonl.gz" + bad.write_bytes(b"\x9f?not a gzip stream at all") + files.append(str(bad)) + + if expect == "raise": + match = "No sparse-policy records" if kind == "empty" else "unreadable" + with pytest.raises(ValueError, match=match): + SparsePolicyDataset(files, max_length=32) + else: + with pytest.warns(UserWarning): + ds = SparsePolicyDataset(files, max_length=32) + assert len(ds) == (3 if kind == "truncated" else 2) + + +def test_sparse_data_path_selects_mode_and_rejects_conflicts(): + assert DataArguments(sparse_data_path="/tmp/x").mode == "sparse" + assert DataArguments(offline_data_path="/tmp/x").mode == "offline" + assert DataArguments().mode == "online" + with pytest.raises(ValueError, match="ambiguous"): + DataArguments(sparse_data_path="/tmp/x", offline_data_path="/tmp/y") + + +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" + _write_records(a, n=3) + first = SparsePolicyDataset([str(a)], max_length=32) + assert len(first) == 3 + + # a second shard changes the signature, so the cache is stale and rebuilds + b = tmp_path / "b.jsonl.gz" + _write_records(b, n=2) + second = SparsePolicyDataset([str(a), str(b)], max_length=32) + assert len(second) == 5 + assert not list(tmp_path.glob(".sparse_cache.stale*")) + + +def test_wait_for_cache_ignores_a_stale_one(tmp_path): + """A non-builder rank must wait for *these* shards, not for any meta.json. + + The only way to reach the wait with a cache on disk is that it was found stale, so + returning on mere existence hands the waiter the previous dump. + """ + a = tmp_path / "a.jsonl.gz" + _write_records(a, n=3) + SparsePolicyDataset([str(a)], max_length=32) # publishes a cache for [a] + cache = tmp_path / ".sparse_cache" + assert (cache / "meta.json").exists() + + b = tmp_path / "b.jsonl.gz" + _write_records(b, n=2) + # [a, b] has a different signature, so the cache on disk is stale for it: the wait + # must time out rather than accept the one already published. + with pytest.raises(TimeoutError): + sparse_data._wait_for_cache(cache, [str(a), str(b)], stall_s=0.1) + # and it returns promptly for the shards it was actually built from + sparse_data._wait_for_cache(cache, [str(a)], stall_s=5.0) + + +def test_losing_a_publish_race_to_different_shards_is_refused(tmp_path): + """Two runs over different subsets of one dump share the default cache path. + + The loser of the publish rename returns from the build normally, having adopted the + winner's cache without ever comparing it against its own shards. It must refuse that + rather than train a whole run on the wrong corpus. + """ + a, b = tmp_path / "a.jsonl.gz", tmp_path / "b.jsonl.gz" + _write_records(a, n=3) + _write_records(b, n=2) + # publish a cache for {a, b}, then ask for one over {a} alone + SparsePolicyDataset([str(a), str(b)], max_length=32) + cache = tmp_path / ".sparse_cache" + meta = json.loads((cache / "meta.json").read_text()) + + # simulate losing the race: the build is skipped and the other job's cache is in place + with ( + mock.patch.object(sparse_data, "_build_cache", return_value=None), + pytest.raises(RuntimeError, match="concurrent job"), + ): + SparsePolicyDataset([str(a)], max_length=32) + # the published cache is untouched by the refusal + assert json.loads((cache / "meta.json").read_text()) == meta + + +def test_a_non_builder_rank_waits_on_the_files_not_a_barrier(tmp_path): + """With a process group already up, rank != 0 must still take the file wait. + + HfTrainingArguments builds the process group -- with ddp_timeout, 1800s by default -- + while parsing arguments, so by the time any dataset is constructed a collective here + is both reachable and bounded. A first build over a large corpus outlasts that bound + and every non-builder rank dies at the barrier. _wait_for_cache needs no group and + bounds time without progress instead, so it is the branch both cases must take. + """ + a, b = tmp_path / "a.jsonl.gz", tmp_path / "b.jsonl.gz" + _write_records(a, n=3) + _write_records(b, n=2) + SparsePolicyDataset([str(a)], max_length=32) # a cache for [a], stale for [a, b] + + waited = [] + real_build = sparse_data._build_cache + + def spy(cache, files, *args, **kwargs): + waited.append(list(files)) + real_build(cache, files, None) # stand in for the rank-0 process + + barrier = mock.Mock() + with ( + mock.patch.object(sparse_data.torch.distributed, "is_initialized", return_value=True), + mock.patch.object(sparse_data.torch.distributed, "get_rank", return_value=1), + mock.patch.object(sparse_data.torch.distributed, "barrier", barrier), + mock.patch.object(sparse_data, "_wait_for_cache", spy), + ): + ds = SparsePolicyDataset([str(a), str(b)], max_length=32) + + assert waited == [[str(a), str(b)]], "rank 1 did not wait on the files" + barrier.assert_not_called() + assert len(ds) == 5 + + +def test_ragged_row_whose_support_includes_id_zero(): + """A padded row whose support holds id 0 still scores ~0 for a matching draft. + + The builder pads a ragged row by repeating id 0, so a row legitimately containing + id 0 has it twice, and ``p`` is not deduplicated -- collapsing to the pad instead of + the real occurrence would pair the teacher's probability with the pad's ``q``. + + The collapse keys on stored probability rather than on token id, so the pad slots + are identified directly and the real occurrence is never the one dropped. + """ + tok = torch.tensor([[[0, 7, 0, 0]]]) + prob = torch.tensor([[[0.6, 0.4, 0.0, 0.0]]]) + # a draft that matches the teacher exactly on this row + logits = torch.full((1, 1, VOCAB), -30.0) + logits[0, 0, 0] = math.log(0.6) + 30.0 + logits[0, 0, 7] = math.log(0.4) + 30.0 + m = _draft(external_loss="tvd", external_report_acc=False) + # one scored position needs two columns: the policy at t is scored against draft t-1 + tok2 = torch.cat([tok, tok], dim=1) + prob2 = torch.cat([prob, prob], dim=1) + logits2 = torch.cat([logits, logits], dim=1) + loss, _ = m.compute_sparse_loss(logits2, tok2, prob2, torch.ones(1, 2)) + assert loss.item() < 0.05, f"matching the teacher should score ~0, got {loss.item()}" + + +def test_concurrent_builders_do_not_corrupt_the_cache(tmp_path): + """Several drafts train from one dumped dataset, so two jobs can build at once. + + They must not share a scratch directory: deleting files the other still holds open + leaves it writing to unlinked inodes and publishing an index whose offsets no longer + match its data, which reads back as a complete cache full of nonsense. + """ + shard = tmp_path / "a.jsonl.gz" + _write_records(shard, n=6) + files = [str(shard)] + cache = tmp_path / ".sparse_cache" + + # Hold both builders at the first shard read, so each has created its scratch + # directory before either starts writing -- the interleaving that loses data. + barrier = threading.Barrier(2) + real_open = gzip.open + + def synced_open(*args, **kwargs): + with contextlib.suppress(threading.BrokenBarrierError): + barrier.wait(timeout=30) + return real_open(*args, **kwargs) + + errors = [] + + def build(): + try: + sparse_data._build_cache(cache, files, None) + except BaseException as exc: + errors.append(exc) + + with mock.patch.object(sparse_data.gzip, "open", synced_open): + threads = [threading.Thread(target=build) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=120) + assert not errors, errors + + ds = SparsePolicyDataset(files, max_length=32) + assert len(ds) == 6 + for i in range(len(ds)): + item = ds[i] + scored = item["loss_mask"].bool() + # Garbage offsets show up here: a row read from the wrong place does not + # sum to 1, and an unwritten one is all zeros. + assert torch.allclose( + item["teacher_topk_prob"][scored].sum(-1), + torch.ones(int(scored.sum())), + atol=1e-2, + ) + assert not list(tmp_path.glob(".sparse_cache.tmp*")) + assert not list(tmp_path.glob(".sparse_cache.stale*")) + + +def test_warns_when_stored_policy_is_a_truncated_nucleus(tmp_path): + """tvd treats the stored policy as complete; a partial one must not pass silently.""" + p = tmp_path / "partial.jsonl.gz" + torch.manual_seed(0) + with gzip.open(p, "wt") as f: + # deliberately store only half the mass, as a too-small top-k would + lp = (torch.rand(5, TOPK) * 0.01).log() + math.log(0.5 / TOPK) + f.write( + json.dumps( + { + "id": "r0", + "prompt_ids": [1, 2], + "gen_ids": [3, 4, 5, 6, 7], + "topk_tok": torch.randint(0, VOCAB, (5, TOPK)).tolist(), + "topk_lp": lp.tolist(), + } + ) + + "\n" + ) + with pytest.warns(UserWarning, match="stored teacher policy sums to"): + SparsePolicyDataset([str(p)], max_length=32) + + +def _write_sourced(path, sources): + """One record per entry in ``sources``, tagged with that source.""" + torch.manual_seed(0) + with gzip.open(path, "wt") as f: + for i, src in enumerate(sources): + lp = torch.randn(4, TOPK).log_softmax(-1) + f.write( + json.dumps( + { + "id": f"r{i}", + "source": src, + "prompt_ids": torch.randint(0, VOCAB, (3,)).tolist(), + "gen_ids": torch.randint(0, VOCAB, (4,)).tolist(), + "topk_tok": torch.randint(0, VOCAB, (4, TOPK)).tolist(), + "topk_lp": lp.tolist(), + } + ) + + "\n" + ) + + +@pytest.mark.parametrize( + ("weights", "expected"), + [ + (None, 4), # unweighted: one entry per record + ({"a": 3.0}, 3 * 2 + 2), # 2 'a' records tripled, 2 'b' untouched + ({"a": 2.0, "b": 0.0}, 4), # a zero weight drops a source entirely + ], +) +def test_source_weights_oversample_by_source(tmp_path, weights, expected): + p = tmp_path / "s.jsonl.gz" + _write_sourced(p, ["a", "a", "b", "b"]) + plain = SparsePolicyDataset([str(p)], max_length=32, top_k=TOPK) + ds = SparsePolicyDataset([str(p)], max_length=32, top_k=TOPK, source_weights=weights) + assert len(ds) == expected + # oversampling repeats index entries; it must not corrupt what they point at + base = {tuple(plain[i]["input_ids"].tolist()) for i in range(len(plain))} + for i in range(len(ds)): + assert tuple(ds[i]["input_ids"].tolist()) in base + + +def test_source_weights_reject_an_unknown_source(tmp_path): + p = tmp_path / "s.jsonl.gz" + _write_sourced(p, ["a", "b"]) + # silently ignoring a typo would leave the corpus unweighted and the run looking fine + with pytest.raises(ValueError, match="not in the data"): + SparsePolicyDataset([str(p)], max_length=32, top_k=TOPK, source_weights={"typo": 2.0}) + + +@pytest.mark.parametrize( + ("bad", "n_bad", "kept"), + [ + # Empty policy: width inference used to raise IndexError and abort the build. + ({"topk_tok": [], "topk_lp": []}, 1, 1), + # topk_lp shorter than topk_tok: the ragged fallback used to index past its end. + ({"topk_lp": [[0.0] * TOPK]}, 1, 2), + # A renamed field drops every record it touches. One dropped record is a + # truncation to tolerate, but a dump-wide rename is schema drift: publishing + # whatever happened to parse would train a whole run on a handful of rows + # against an ordinary-looking loss curve, so the drop rate has to fail it. + ({"topk_lp": _DROP, "logprobs": [[0.0] * TOPK]}, 40, "raises"), + ], +) +def test_malformed_policy_records_never_abort_the_build(tmp_path, bad, n_bad, kept): + """A damaged record is dropped or truncated, never fatal to the whole pass -- until + enough of them are damaged that the dump, not the record, is what is wrong.""" + p = tmp_path / "s.jsonl.gz" + torch.manual_seed(0) + good = { + "id": "ok", + "prompt_ids": torch.randint(0, VOCAB, (3,)).tolist(), + "gen_ids": torch.randint(0, VOCAB, (4,)).tolist(), + "topk_tok": torch.randint(0, VOCAB, (4, TOPK)).tolist(), + "topk_lp": torch.randn(4, TOPK).log_softmax(-1).tolist(), + } + with gzip.open(p, "wt") as f: + for i in range(n_bad): + record = {**good, "id": f"bad{i}", **bad} + f.write(json.dumps({k: v for k, v in record.items() if v is not _DROP}) + "\n") + f.write(json.dumps(good) + "\n") + + if kept == "raises": + with pytest.raises(ValueError, match="failed to parse"): + SparsePolicyDataset([str(p)], max_length=32, top_k=TOPK) + return + + ds = SparsePolicyDataset([str(p)], max_length=32, top_k=TOPK) + assert len(ds) == kept + # the intact record must still be intact + assert ds[len(ds) - 1]["input_ids"].shape[0] == len(good["prompt_ids"]) + len(good["gen_ids"]) + + +def test_sample_size_draws_across_sources_not_a_prefix(tmp_path): + """`limit` must sample the weighted index, not take index.npy's prefix. + + index.npy is written in shard-then-line order, so a prefix is a prefix of the shard + list: with one source per shard -- the natural layout -- it can hold a single source, + and the configured mixture is then applied to whatever survived. + """ + shards = [] + for name in ("alpha", "beta"): + path = tmp_path / f"{name}.jsonl" + with path.open("w") as fh: + for i in range(40): + fh.write( + json.dumps( + { + "id": f"{name}-{i}", + "source": name, + "prompt_ids": [1, 2], + "gen_ids": [3, 4], + "topk_tok": [[3, 5], [4, 6]], + "topk_lp": [[-0.1, -2.0], [-0.1, -2.0]], + } + ) + + "\n" + ) + shards.append(str(path)) + + ds = SparsePolicyDataset(shards, max_length=64, limit=20) + assert ds.sources is not None + assert len(ds.sources) == len(ds), "sources desynchronised from index" + kept = {ds.source_names[s] for s in ds.sources} + assert kept == {"alpha", "beta"}, f"sample came from {kept}, not both shards" + + +def test_a_wider_row_refuses_rather_than_losing_teacher_mass(tmp_path): + """Width comes from the first record; a later wider row must not be silently clipped. + + With top_p < 1 the stored nucleus varies per position, so a corpus whose later rows + are wider than the first record's is ordinary -- and truncating them drops real + teacher probability that no stored-mass check can see afterwards. + """ + shard = tmp_path / "ragged.jsonl" + rows = [ + { + "id": "narrow", + "prompt_ids": [1, 2], + "gen_ids": [3, 4], + "topk_tok": [[3, 5], [4, 6]], + "topk_lp": [[-0.1, -2.0], [-0.1, -2.0]], + }, + { + "id": "wide", + "prompt_ids": [1, 2], + "gen_ids": [3, 4], + "topk_tok": [[3, 5, 7, 9], [4, 6, 8, 10]], + "topk_lp": [[-0.1, -2.0, -3.0, -4.0], [-0.1, -2.0, -3.0, -4.0]], + }, + ] + with shard.open("w") as fh: + for r in rows: + fh.write(json.dumps(r) + "\n") + with pytest.raises(ValueError, match="sized at"): + SparsePolicyDataset([str(shard)], max_length=64)