From 8ef06e47a31eabaae67fe9656330f932aab8450d Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Fri, 16 Oct 2026 11:17:52 +0000 Subject: [PATCH 1/3] [4/8] Add the sparse teacher-policy training path Adds data.sparse_data_path: shards holding the base's top-k deployment policy per token instead of its hidden states. The policy is roughly 120 B against 10 KiB per token, is captured during generation, and needs no base lm_head at training time, so the base model costs nothing on the GPU. Only the TVD objectives can consume it; soft_ce needs full base logits and main.py rejects the combination rather than training against a zero-padded distribution. SparsePolicyDataset backs the shards with a memmapped cache and can oversample by the record's source field, since response lengths differ by an order of magnitude across sources. The sparse objective shifts the draft one position left before comparing: the stored policy at t is what the base sampled input_ids[t] from, while the draft's logits at t predict t + 1. 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/eagle_utils.py | 45 +- examples/speculative_decoding/main.py | 59 +- modelopt/torch/speculative/config.py | 4 +- .../torch/speculative/external/__init__.py | 1 + .../torch/speculative/external/sparse_data.py | 549 ++++++++++++++++++ .../torch/speculative/plugins/hf_external.py | 107 +++- .../speculative/plugins/hf_training_args.py | 27 +- .../speculative_decoding/external_draft.yaml | 21 + .../speculative/plugins/test_hf_external.py | 43 ++ .../test_speculative_decoding_main.py | 30 + .../speculative/test_external_draft_sparse.py | 546 +++++++++++++++++ 11 files changed, 1401 insertions(+), 31 deletions(-) create mode 100644 modelopt/torch/speculative/external/sparse_data.py create mode 100644 tests/unit/torch/speculative/test_external_draft_sparse.py 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..8dc65cec657 --- /dev/null +++ b/modelopt/torch/speculative/external/sparse_data.py @@ -0,0 +1,549 @@ +# 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``, with no barrier to use. + + 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. + raise ValueError( + f"{path}: a stored policy row holds {widest} entries " + f"but the cache was sized at {width} from the first " + f"record. Pass top_k= to size it " + f"explicitly." + ) + 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. + # The dataset is often constructed before the process group exists, so fall + # back to the launcher's RANK: without it every worker takes the rank-0 + # branch and they race to delete and rewrite the same cache directory. + if torch.distributed.is_initialized(): + rank, can_barrier = torch.distributed.get_rank(), True + else: + rank, can_barrier = int(os.environ.get("RANK", 0)), False + if rank == 0: + _build_cache(cache, files, top_k) + if can_barrier: + torch.distributed.barrier() + elif rank != 0: + _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 and the barrier. + 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..7da426e72f0 100644 --- a/modelopt_recipes/general/speculative_decoding/external_draft.yaml +++ b/modelopt_recipes/general/speculative_decoding/external_draft.yaml @@ -25,6 +25,21 @@ 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. + 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 +69,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/gpu/torch/speculative/plugins/test_hf_external.py b/tests/gpu/torch/speculative/plugins/test_hf_external.py index 69c79f05892..d6e2815eedb 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_external.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_external.py @@ -30,6 +30,7 @@ import modelopt.torch.speculative as mtsp SEQ_LEN = 16 +TOPK = 8 def _external_draft(**overrides): @@ -73,6 +74,48 @@ def test_base_head_stays_out_of_the_draft_parameters(self): assert not any("_base_lm_head" in k for k in draft.state_dict()) +class TestExternalDraftSparseGPU: + """What the sparse path needs a GPU to prove: that its dedup is order-stable there. + + The forward/backward and the no-teacher-head checks are device-independent and are + covered on CPU in tests/unit; repeating them here buys nothing. The dedup below is + different: it relies on ``sort(..., stable=True)`` keeping the first occurrence of a + repeated id, and unstable tie order is genuinely unspecified on CUDA. + """ + + def test_ragged_row_dedup_is_stable_on_cuda(self): + """A padded row repeating a real token must be charged once, not once per slot. + + The builder pads ragged rows by repeating the trailing id, so the same token can + occupy several slots with the teacher's probability stored only in the first. + Keeping any other occurrence pairs the teacher's mass with the wrong q, which no + loss-magnitude check would catch -- so compare against the deduplicated reference. + """ + draft = _external_draft(external_loss="tvd", external_report_acc=False).cuda() + vocab = draft.config.vocab_size + torch.manual_seed(0) + logits = torch.randn(1, 3, vocab, device="cuda") + # row = [7, 11, 11, 11, ...]: one real entry then pads repeating id 11 + tok = torch.full((1, 3, TOPK), 11, dtype=torch.long, device="cuda") + tok[..., 0] = 7 + prob = torch.zeros(1, 3, TOPK, device="cuda") + prob[..., 0] = 0.6 + prob[..., 1] = 0.4 # the first 11 carries the mass; later slots are pads + + loss, _ = draft.compute_sparse_loss(logits, tok, prob, torch.ones(1, 3, device="cuda")) + + q_all = torch.softmax(logits[:, :-1].float(), dim=-1) + p = prob[:, 1:] / prob[:, 1:].sum(-1, keepdim=True) + q = q_all.gather(-1, tok[:, 1:]) + keep = torch.zeros_like(q, dtype=torch.bool) + keep[..., :2] = True # ids 7 and 11 once each; the rest are duplicates + q = q * keep + 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), ( + f"dedup kept the wrong occurrence on CUDA: {loss} vs {expected.mean()}" + ) + + def test_dense_path_rejects_a_missing_teacher_head(): """Without the head the dense objective cannot be formed; fail loudly, not silently.""" draft = _external_draft().cuda() 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..cddbd552814 --- /dev/null +++ b/tests/unit/torch/speculative/test_external_draft_sparse.py @@ -0,0 +1,546 @@ +# 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, +) +from modelopt.torch.speculative.plugins.hf_external import SPARSE_CAPABLE_LOSSES +from modelopt.torch.speculative.plugins.hf_training_args import DataArguments + +transformers = pytest.importorskip("transformers") + +VOCAB = 64 +HIDDEN = 32 +TOPK = 8 + + +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_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``. + + This pins the property, not the ordering that guarantees it: ``torch.sort`` is + already stable on CPU, so it passes with or without the ``stable=True`` that makes + it hold on CUDA, where the tie order is genuinely unspecified. + """ + 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", "kept"), + [ + # Empty policy: width inference used to raise IndexError and abort the build. + ({"topk_tok": [], "topk_lp": []}, 1), + # topk_lp shorter than topk_tok: the ragged fallback used to index past its end. + ({"topk_lp": [[0.0] * TOPK]}, 2), + ], +) +def test_malformed_policy_records_never_abort_the_build(tmp_path, bad, kept): + """A damaged record is dropped or truncated, never fatal to the whole pass.""" + 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: + f.write(json.dumps({**good, "id": "bad", **bad}) + "\n") + f.write(json.dumps(good) + "\n") + 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_schema_drifted_dump_refuses_rather_than_publishing_a_few_rows(tmp_path): + """Every record failing to parse must fail the build, not publish a near-empty cache. + + The shard-level counter records at most one entry per build for record damage, so on + its own it cannot see "kept 3 of 100000" -- the build warns once, passes, and training + runs on whatever happened to parse with an ordinary-looking loss curve. + """ + # A handful parse and the rest do not: the all-or-nothing case is already caught by + # the empty-index guard, so the gap this closes is the partial one -- a cache that + # looks complete and holds 2 of 42 records. + shard = tmp_path / "drifted.jsonl" + good = { + "id": "ok", + "prompt_ids": [1, 2], + "gen_ids": [3, 4], + "topk_tok": [[3, 5], [4, 6]], + "topk_lp": [[-0.1, -2.0], [-0.1, -2.0]], + } + with shard.open("w") as fh: + for i in range(2): + fh.write(json.dumps({**good, "id": f"ok{i}"}) + "\n") + for i in range(40): + fh.write( + json.dumps( + { + "id": f"r{i}", + "prompt_ids": [1, 2], + "gen_ids": [3, 4], + "topk_tok": [[3, 5], [4, 6]], + "logprobs": [[-0.1, -2.0], [-0.1, -2.0]], # producer renamed topk_lp + } + ) + + "\n" + ) + with pytest.raises(ValueError, match="failed to parse"): + SparsePolicyDataset([str(shard)], max_length=64) + + +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) From e5f9d78199f206593f801ec46d84129d537f5be0 Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Fri, 16 Oct 2026 13:05:06 +0000 Subject: [PATCH 2/3] Skip the sparse test module when transformers is absent The module guarded itself with pytest.importorskip("transformers") but imported plugins.hf_external two lines above the guard. The plugins package registers that module behind import_plugin("transformers"); importing the submodule directly bypasses the guard, so on a torch-only install the module failed to collect instead of skipping, and partial-install (torch) went red. Move the two plugin imports below the guard. Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- .../unit/torch/speculative/test_external_draft_sparse.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/speculative/test_external_draft_sparse.py b/tests/unit/torch/speculative/test_external_draft_sparse.py index cddbd552814..ec8413f279e 100644 --- a/tests/unit/torch/speculative/test_external_draft_sparse.py +++ b/tests/unit/torch/speculative/test_external_draft_sparse.py @@ -31,11 +31,15 @@ SparsePolicyCollator, SparsePolicyDataset, ) -from modelopt.torch.speculative.plugins.hf_external import SPARSE_CAPABLE_LOSSES -from modelopt.torch.speculative.plugins.hf_training_args import DataArguments 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 From 81f4cf605706e8cf7f7e792394318c020f117566 Mon Sep 17 00:00:00 2001 From: Peter Hicks Date: Fri, 16 Oct 2026 13:28:59 +0000 Subject: [PATCH 3/3] Address Keval's second review pass on the sparse path The stall-aware wait was unreachable. HfTrainingArguments builds the process group while parsing arguments (TrainingArguments.__post_init__ touches self.device, which constructs a PartialState with ddp_timeout), so by the time make_speculative_data_module runs is_initialized() is already True and the barrier branch always won. Every non-builder rank was still bounded by the 1800s ddp_timeout, which a first build over a large corpus outlasts. Drop the barrier and give every non-builder rank the file wait, which needs no process group and bounds time without progress instead. Also: - Drop TestExternalDraftSparseGPU. It existed to pin sort(..., stable=True) on CUDA; the dedup now masks on topk_prob > 0 and _tvd_sparse does not sort, so nothing in it is device-dependent. Its CPU twin covers the computation. The twin's docstring described the same removed sort and is corrected. - Fold the schema-drift case into test_malformed_policy_records_never_abort_ the_build as one more (bad, n_bad, kept) row with a raises expectation, rather than a third test of damaged input. - Point the top-k width refusal at the single-process pre-build, which accepts top_k, since no recipe field threads one through. Document that in the recipe. Co-Authored-By: Claude Opus 5 Signed-off-by: Peter Hicks --- .../torch/speculative/external/sparse_data.py | 33 +++--- .../speculative_decoding/external_draft.yaml | 4 +- .../speculative/plugins/test_hf_external.py | 43 ------- .../speculative/test_external_draft_sparse.py | 106 ++++++++++-------- 4 files changed, 82 insertions(+), 104 deletions(-) diff --git a/modelopt/torch/speculative/external/sparse_data.py b/modelopt/torch/speculative/external/sparse_data.py index 8dc65cec657..2f1f7f6e2f2 100644 --- a/modelopt/torch/speculative/external/sparse_data.py +++ b/modelopt/torch/speculative/external/sparse_data.py @@ -99,7 +99,7 @@ def _builder_progress(cache: Path) -> tuple[int, int]: 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``, with no barrier to use. + """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 @@ -228,11 +228,16 @@ def _build_cache(cache: Path, files: list[str], top_k: int | None) -> None: # 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. Pass top_k= to size it " - f"explicitly." + 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) @@ -393,19 +398,21 @@ def __init__( ) if stale: - # Under DDP every rank sees the same shards; one builds and the rest wait. - # The dataset is often constructed before the process group exists, so fall - # back to the launcher's RANK: without it every worker takes the rank-0 - # branch and they race to delete and rewrite the same cache directory. + # 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, can_barrier = torch.distributed.get_rank(), True + rank = torch.distributed.get_rank() else: - rank, can_barrier = int(os.environ.get("RANK", 0)), False + rank = int(os.environ.get("RANK", 0)) if rank == 0: _build_cache(cache, files, top_k) - if can_barrier: - torch.distributed.barrier() - elif rank != 0: + else: _wait_for_cache(cache, files) meta = json.loads((cache / "meta.json").read_text()) @@ -413,7 +420,7 @@ def __init__( # 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 and the barrier. + # 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 " diff --git a/modelopt_recipes/general/speculative_decoding/external_draft.yaml b/modelopt_recipes/general/speculative_decoding/external_draft.yaml index 7da426e72f0..c7415a9faf2 100644 --- a/modelopt_recipes/general/speculative_decoding/external_draft.yaml +++ b/modelopt_recipes/general/speculative_decoding/external_draft.yaml @@ -35,7 +35,9 @@ data: # 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. + # 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. diff --git a/tests/gpu/torch/speculative/plugins/test_hf_external.py b/tests/gpu/torch/speculative/plugins/test_hf_external.py index d6e2815eedb..69c79f05892 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_external.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_external.py @@ -30,7 +30,6 @@ import modelopt.torch.speculative as mtsp SEQ_LEN = 16 -TOPK = 8 def _external_draft(**overrides): @@ -74,48 +73,6 @@ def test_base_head_stays_out_of_the_draft_parameters(self): assert not any("_base_lm_head" in k for k in draft.state_dict()) -class TestExternalDraftSparseGPU: - """What the sparse path needs a GPU to prove: that its dedup is order-stable there. - - The forward/backward and the no-teacher-head checks are device-independent and are - covered on CPU in tests/unit; repeating them here buys nothing. The dedup below is - different: it relies on ``sort(..., stable=True)`` keeping the first occurrence of a - repeated id, and unstable tie order is genuinely unspecified on CUDA. - """ - - def test_ragged_row_dedup_is_stable_on_cuda(self): - """A padded row repeating a real token must be charged once, not once per slot. - - The builder pads ragged rows by repeating the trailing id, so the same token can - occupy several slots with the teacher's probability stored only in the first. - Keeping any other occurrence pairs the teacher's mass with the wrong q, which no - loss-magnitude check would catch -- so compare against the deduplicated reference. - """ - draft = _external_draft(external_loss="tvd", external_report_acc=False).cuda() - vocab = draft.config.vocab_size - torch.manual_seed(0) - logits = torch.randn(1, 3, vocab, device="cuda") - # row = [7, 11, 11, 11, ...]: one real entry then pads repeating id 11 - tok = torch.full((1, 3, TOPK), 11, dtype=torch.long, device="cuda") - tok[..., 0] = 7 - prob = torch.zeros(1, 3, TOPK, device="cuda") - prob[..., 0] = 0.6 - prob[..., 1] = 0.4 # the first 11 carries the mass; later slots are pads - - loss, _ = draft.compute_sparse_loss(logits, tok, prob, torch.ones(1, 3, device="cuda")) - - q_all = torch.softmax(logits[:, :-1].float(), dim=-1) - p = prob[:, 1:] / prob[:, 1:].sum(-1, keepdim=True) - q = q_all.gather(-1, tok[:, 1:]) - keep = torch.zeros_like(q, dtype=torch.bool) - keep[..., :2] = True # ids 7 and 11 once each; the rest are duplicates - q = q * keep - 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), ( - f"dedup kept the wrong occurrence on CUDA: {loss} vs {expected.mean()}" - ) - - def test_dense_path_rejects_a_missing_teacher_head(): """Without the head the dense objective cannot be formed; fail loudly, not silently.""" draft = _external_draft().cuda() diff --git a/tests/unit/torch/speculative/test_external_draft_sparse.py b/tests/unit/torch/speculative/test_external_draft_sparse.py index ec8413f279e..c91d9619cc6 100644 --- a/tests/unit/torch/speculative/test_external_draft_sparse.py +++ b/tests/unit/torch/speculative/test_external_draft_sparse.py @@ -43,6 +43,8 @@ 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): @@ -266,6 +268,41 @@ def test_losing_a_publish_race_to_different_shards_is_refused(tmp_path): 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. @@ -273,9 +310,8 @@ def test_ragged_row_whose_support_includes_id_zero(): 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``. - This pins the property, not the ordering that guarantees it: ``torch.sort`` is - already stable on CPU, so it passes with or without the ``stable=True`` that makes - it hold on CUDA, where the tie order is genuinely unspecified. + 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]]]) @@ -419,16 +455,22 @@ def test_source_weights_reject_an_unknown_source(tmp_path): @pytest.mark.parametrize( - ("bad", "kept"), + ("bad", "n_bad", "kept"), [ # Empty policy: width inference used to raise IndexError and abort the build. - ({"topk_tok": [], "topk_lp": []}, 1), + ({"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]}, 2), + ({"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, kept): - """A damaged record is dropped or truncated, never fatal to the whole pass.""" +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 = { @@ -439,8 +481,16 @@ def test_malformed_policy_records_never_abort_the_build(tmp_path, bad, kept): "topk_lp": torch.randn(4, TOPK).log_softmax(-1).tolist(), } with gzip.open(p, "wt") as f: - f.write(json.dumps({**good, "id": "bad", **bad}) + "\n") + 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 @@ -481,44 +531,6 @@ def test_sample_size_draws_across_sources_not_a_prefix(tmp_path): assert kept == {"alpha", "beta"}, f"sample came from {kept}, not both shards" -def test_a_schema_drifted_dump_refuses_rather_than_publishing_a_few_rows(tmp_path): - """Every record failing to parse must fail the build, not publish a near-empty cache. - - The shard-level counter records at most one entry per build for record damage, so on - its own it cannot see "kept 3 of 100000" -- the build warns once, passes, and training - runs on whatever happened to parse with an ordinary-looking loss curve. - """ - # A handful parse and the rest do not: the all-or-nothing case is already caught by - # the empty-index guard, so the gap this closes is the partial one -- a cache that - # looks complete and holds 2 of 42 records. - shard = tmp_path / "drifted.jsonl" - good = { - "id": "ok", - "prompt_ids": [1, 2], - "gen_ids": [3, 4], - "topk_tok": [[3, 5], [4, 6]], - "topk_lp": [[-0.1, -2.0], [-0.1, -2.0]], - } - with shard.open("w") as fh: - for i in range(2): - fh.write(json.dumps({**good, "id": f"ok{i}"}) + "\n") - for i in range(40): - fh.write( - json.dumps( - { - "id": f"r{i}", - "prompt_ids": [1, 2], - "gen_ids": [3, 4], - "topk_tok": [[3, 5], [4, 6]], - "logprobs": [[-0.1, -2.0], [-0.1, -2.0]], # producer renamed topk_lp - } - ) - + "\n" - ) - with pytest.raises(ValueError, match="failed to parse"): - SparsePolicyDataset([str(shard)], max_length=64) - - 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.