[trainer, model, data] feat: integrate DeepSpec draft-model training - #882
[trainer, model, data] feat: integrate DeepSpec draft-model training#882zhangxin81 wants to merge 2 commits into
Conversation
|
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an integration between DeepSpec speculative-decoding draft models and the VeOmni distributed training engine. It includes a trainer implementation, data adapters for DeepSpec's target cache, and configuration utilities to support training DSpark and Eagle3 models using FSDP2. The review identified several critical issues, including a missing return value in the training step, incorrect parameter freezing logic, potential type mismatches in dataset length calculations, and issues with attribute access on FSDP-wrapped models. Additionally, a parameter naming error in the initialization script was corrected.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| total_loss_dict.update(self._flush_deepspec_metrics()) | ||
|
|
||
| self.on_step_end(loss=total_loss, loss_dict=total_loss_dict, grad_norm=grad_norm) |
There was a problem hiding this comment.
The train_step method is annotated to return Dict[str, float], but it currently lacks a return statement, effectively returning None. This will cause runtime errors or static type checking failures in the training loop. Additionally, total_loss_dict does not contain the standard "loss" key because forward_backward_step returns an empty dictionary for its metrics, which prevents the standard loss metric from being logged to WandB/TensorBoard. Please update total_loss_dict with the "loss" key and return it at the end of the method.
| total_loss_dict.update(self._flush_deepspec_metrics()) | |
| self.on_step_end(loss=total_loss, loss_dict=total_loss_dict, grad_norm=grad_norm) | |
| total_loss_dict.update(self._flush_deepspec_metrics()) | |
| total_loss_dict["loss"] = total_loss | |
| self.on_step_end(loss=total_loss, loss_dict=total_loss_dict, grad_norm=grad_norm) | |
| return total_loss_dict |
| for name, param in model.named_parameters(): | ||
| if name.startswith("embed_tokens") or name.startswith("lm_head"): | ||
| param.requires_grad_(False) |
There was a problem hiding this comment.
In HuggingFace models, parameters are typically nested under a sub-module (e.g., model.embed_tokens.weight or model.lm_head.weight). Using name.startswith(...) will fail to match these nested names because they start with "model." instead of the layer name directly. This means the fallback path will silently fail to freeze the embeddings and language modeling head, leading to incorrect training behavior and potential memory issues. Using "embed_tokens" in name or "lm_head" in name is much more robust.
| for name, param in model.named_parameters(): | |
| if name.startswith("embed_tokens") or name.startswith("lm_head"): | |
| param.requires_grad_(False) | |
| for name, param in model.named_parameters(): | |
| if "embed_tokens" in name or "lm_head" in name: | |
| param.requires_grad_(False) |
| self.train_dataset = build_target_cache_dataset(cache_dir) | ||
| self._validate_cache_against_model(self.train_dataset) | ||
|
|
||
| dataset_length = len(self.train_dataset) / get_parallel_state().dp_size |
There was a problem hiding this comment.
The dataset length represents the number of samples per rank, which must be an integer. Using float division / can result in a float value, which might cause type mismatches or unexpected behavior in downstream step calculations. Please use integer floor division // instead.
| dataset_length = len(self.train_dataset) / get_parallel_state().dp_size | |
| dataset_length = len(self.train_dataset) // get_parallel_state().dp_size |
| def _validate_cache_against_model(self, dataset): | ||
| """Mirror DeepSpec's ``validate_train_cache`` on the built model.""" | ||
| model = self.model | ||
| expected_layer_ids = [int(x) for x in getattr(model, "target_layer_ids", [])] |
There was a problem hiding this comment.
When the model is wrapped by FSDP, custom attributes like target_layer_ids are not forwarded from the underlying module to the FSDP wrapper. This means getattr(model, "target_layer_ids", []) will return [] on wrapped models, silently bypassing the target layer validation check. To ensure this safety check always runs correctly, please access target_layer_ids from self.model_config instead.
| expected_layer_ids = [int(x) for x in getattr(model, "target_layer_ids", [])] | |
| expected_layer_ids = [int(x) for x in getattr(self.model_config, "target_layer_ids", [])] |
| ttt_length=int(getattr(self.model, "ttt_length", self.model_config.ttt_length)), | ||
| step_loss_decay=float(getattr(self.model, "step_loss_decay", self.model_config.step_loss_decay)), |
There was a problem hiding this comment.
When the model is wrapped by FSDP, custom attributes like ttt_length and step_loss_decay are not forwarded from the underlying module to the FSDP wrapper. This means getattr(self.model, ...) will fail to find them on the wrapper and fall back to the config values. To be safe and consistent, please access these attributes directly from self.model_config.
| ttt_length=int(getattr(self.model, "ttt_length", self.model_config.ttt_length)), | |
| step_loss_decay=float(getattr(self.model, "step_loss_decay", self.model_config.step_loss_decay)), | |
| ttt_length=int(getattr(self.model_config, "ttt_length")), | |
| step_loss_decay=float(getattr(self.model_config, "step_loss_decay")), |
|
|
||
| print("[prepare_draft_init] Copying frozen target embeddings + lm_head...") | ||
| target_model = ( | ||
| AutoModelForCausalLM.from_pretrained(args.target_model_name_or_path, dtype=torch.float32) |
There was a problem hiding this comment.
In HuggingFace Transformers, the parameter to specify the model's data type is torch_dtype, not dtype. Passing dtype can raise a TypeError or be ignored depending on the Transformers version. Please use torch_dtype instead.
| AutoModelForCausalLM.from_pretrained(args.target_model_name_or_path, dtype=torch.float32) | |
| AutoModelForCausalLM.from_pretrained(args.target_model_name_or_path, torch_dtype=torch.float32) |
| asserts a consistent metric schema across ranks, so every rank must call | ||
| it exactly once per step. | ||
| """ | ||
| try: |
There was a problem hiding this comment.
Maybe it's better to pin a commit for Deepspec and add it to pyproject.toml
| model.set_embedding_head_trainable(False) | ||
| else: # defensive: freeze by name if the helper is absent | ||
| for name, param in model.named_parameters(): | ||
| if name.startswith("embed_tokens") or name.startswith("lm_head"): |
There was a problem hiding this comment.
sub string match like "embed_tokens" in name or "lm_head" in name maybe better
| micro_batch = self.preforward(micro_batch) | ||
|
|
||
| with self.model_fwd_context: | ||
| loss = self._compute_deepspec_loss(micro_batch) |
There was a problem hiding this comment.
We need get the loss reduce on global tokens, it seems _compute_deepspec_loss does not return the global loss.
Install DeepSpec as a pinned git dependency instead of a sys.path-only bootstrap, and address reviewer comments on the draft trainer. - Add an opt-in `deepspec` extra: git source pinned by commit in [tool.uv.sources] (the commit adds DeepSpec's pyproject.toml, proposed upstream in deepseek-ai/DeepSpec#54; repoint to upstream once merged) and an empty [[tool.uv.dependency-metadata]] so DeepSpec's own torch/ transformers pins do not conflict with VeOmni's stack. DEEPSPEC_PATH / sibling checkout stays as a dev fallback. - draft_trainer: report the true global-token loss (mean all-reduce over the fsdp group inverts DeepSpec's world-size backward scaling) instead of DeepSpec's per-rank mean-of-means metric. - draft_trainer: freeze embed_tokens / lm_head by substring match so nested HF param names (model.embed_tokens.weight) are matched in the fallback path. - Update integration docs, uv.md, and integration docstrings to describe the deepspec extra as the primary install path.
|
@FoolPlayer thanks for the review — addressed in the latest commit. Summary of the three comments: 1. Freeze 2. Loss should reduce over global tokens — done. 3. Pin DeepSpec and add it to
|
Re-review after the "address review feedback" update (
|
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Review summary
Baseline: main@7be22df074b4, ahead by 2, behind by 55.
1. Does it solve a real problem? Yes. Training speculative-decoding draft models on VeOmni's engine instead of DeepSpec's FSDP.NO_SHARD setup is a genuine capability, and the reasoning is unusually well argued: the module docstring explains exactly why DeepSpec's ×world_size loss scaling cancels FSDP2's gradient averaging only when the shard group spans the world, and _validate_parallelism then refuses every other parallel dim rather than silently mis-scaling gradients. The registry integration is also done correctly — registering under a dedicated deepspec_draft model type avoids shadowing the real Qwen3, the factory signature matches loader.py:143's MODELING_REGISTRY[model_type](arch_name) contract, and deepspec is imported lazily inside _import_draft_class, so import veomni.models.transformers still works without the optional extra. That last point is easy to get wrong and worth calling out as done right.
2. Are comments, docs and tests complete? Docs, yes — a 281-line design document correctly registered in the docs/index.md toctree, a 137-line config README, two example YAMLs, and genuinely explanatory docstrings throughout. Tests, no: roughly 1,600 lines of new source (a BaseTrainer subclass that overrides train_step and forward_backward_step, a data adapter, a config class, a prep script, a patchgen target) with no test file at all. At minimum the pure-Python pieces are testable without DeepSpec installed: _infer_algorithm, _dspark_loss_weights precedence, _validate_cache_against_model's two error paths, and _validate_parallelism rejecting each parallel dim. Note also that CI will fail as-is — see the inline comment on the patchgen config.
3. Is it clean, and does it conflict with main? The code is tidy and well organized, but there are three structural problems. First, the new dependency resolves from a personal fork (github.com/zhangxin81/DeepSpec), which the comment itself flags as temporary pending upstream deepseek-ai/DeepSpec#54. Second, [[tool.uv.dependency-metadata]] requires-dist = [] suppresses DeepSpec's own pins, which include transformers==5.10.2 against VeOmni's 5.9.0 — and the generated draft modeling is produced from that code. Third, both trainer overrides silently drop behaviour the base class provides. Textually the PR also conflicts: git apply --check fails on docs/index.md, pyproject.toml and uv.lock (mergeable: false, rebaseable: false); the lockfile in particular has to be regenerated with uv lock, not merged by hand.
Overall: needs changes. The design and documentation are the strongest part of this PR and the loss-scaling analysis is convincing. The blockers are the fork-hosted dependency, the patchgen CI failure, and the silently dropped base-trainer features; tests and a rebase are needed on top.
Sent by Cursor Automation: VeOmni-Review-Agent(Outer)
| # DeepSpec's pyproject.toml (upstream PR deepseek-ai/DeepSpec#54); repoint to the | ||
| # upstream commit once that merges. Only resolved when the `deepspec` extra is | ||
| # selected. | ||
| deepspec = { git = "https://github.com/zhangxin81/DeepSpec", rev = "57a01adb496229f3d93afd10679adf5ced7094a2" } |
There was a problem hiding this comment.
This makes VeOmni's lockfile depend on a personal fork of a third-party project. Even pinned by commit that is a durability problem for a first-party repo: the fork can be renamed, made private, or have its history rewritten, and every uv sync --extra deepspec (plus any CI job that resolves it) then breaks with no recourse from this repo. The comment already acknowledges the intent to repoint once deepseek-ai/DeepSpec#54 merges.
Options, roughly in order of preference: wait for the upstream merge and pin deepseek-ai/DeepSpec; or point at the upstream commit now and cover the missing pyproject.toml with the DEEPSPEC_PATH / checkout fallback that deepspec_path.py already implements; or keep this PR as the trainer/data/registry integration only, with DeepSpec left as a documented manual install and no entry in pyproject.toml or uv.lock.
| [[tool.uv.dependency-metadata]] | ||
| name = "deepspec" | ||
| version = "0.1.0" | ||
| requires-dist = [] |
There was a problem hiding this comment.
requires-dist = [] is an established technique in this file (flash-qla above does the same), but the justification here is riskier than for a compiled kernel package. DeepSpec pins torch==2.9.1 and transformers==5.10.2; VeOmni ships torch 2.11.0+cu130 and transformers 5.9.0, so this is not just "drop a redundant pin", it silently resolves a version conflict in the opposite direction — DeepSpec's modeling code was written against a newer transformers than VeOmni pins, and the generated draft modeling in this PR is derived from it.
Please state in the comment which transformers APIs the imported deepspec.modeling.* / deepspec.data / deepspec.utils.* modules actually touch and that they exist in 5.9.0, and enumerate DeepSpec's non-pinned runtime imports so a reader can confirm nothing is genuinely missing from the environment.
|
|
||
|
|
||
| config = PatchConfig( | ||
| source_module="deepspec.modeling.dspark.qwen3.modeling", |
There was a problem hiding this comment.
This will break check_patchgen CI. That workflow runs uv run --extra gpu --dev patchgen --check (.github/workflows/check_patchgen.yml:45) — it does not select the deepspec extra, and this PR does not change the workflow. patchgen's resolver (patchgen-pkg/patchgen/codegen.py:83-99) first looks for deepspec/modeling/dspark/qwen3/modeling.py under sys.path, then falls back to importlib.import_module and raises CodegenError: Cannot import module 'deepspec.modeling.dspark.qwen3.modeling' when that fails. With DeepSpec absent, neither path can succeed.
So the check needs --extra deepspec added, which in turn makes every patchgen CI run fetch and execute code from the git dependency discussed above. That coupling is worth designing deliberately rather than inheriting: consider whether this patchgen target belongs in the repo at all, or whether the OpSlot wrapping is better expressed as a runtime subclass in veomni/models/transformers/deepspec_draft/ that needs no codegen against an external source module.
| self.LOG_SAMPLE = False | ||
| return micro_batch | ||
|
|
||
| def forward_backward_step(self, micro_batch: Dict[str, torch.Tensor]): |
There was a problem hiding this comment.
This override silently drops several things BaseTrainer.forward_backward_step does (veomni/trainer/base.py:718-742):
use_parallel_state("base")chunk_mbs_context(self._chunk_mbs_ranges)set_batch_invariant_mode(self.args.train.enable_batch_invariant_mode)- the
channel_loss_callbackmicro-step context,strip_model_inputs, and forward context
The practical consequence is that a user who sets --train.enable_batch_invariant_mode or configures channel loss for a draft run gets no error and no effect. Since _init_callbacks still constructs channel_loss_callback when configured, that one is a silent no-op rather than a clean refusal. Please either keep the contexts that still apply, or refuse the flags explicitly in _validate_parallelism with a warning_rank0 so the omission is visible.
| self._num_micro_steps = max(1, len(micro_batches)) | ||
|
|
||
| for micro_step, micro_batch in enumerate(micro_batches): | ||
| self.model_reshard(micro_step, self._num_micro_steps) |
There was a problem hiding this comment.
Two omissions relative to the base loop (veomni/trainer/base.py:806-822): mark_compile_step_begin(getattr(self.model, "_veomni_compile_uses_cuda_graphs", False)) and self._configure_hsdp_allreduce(micro_step, num_micro_steps). The HSDP one is defensible — _validate_parallelism already rejects dp_replicate_size > 1, so it would be a no-op — but saying so in a comment would keep the next person from "fixing" it. mark_compile_step_begin is not covered by any guard, so a compiled draft model would step incorrectly. The base also wraps grad clipping in with use_parallel_state("base"), which this override calls bare.
Related: self._num_micro_steps is set here and read in forward_backward_step, so calling forward_backward_step from anywhere else (evaluation, a test) raises AttributeError. A class-level default of 1 would make it safe.
| return {} | ||
| try: | ||
| summary = deepspec_metrics.flush() | ||
| except Exception as exc: # never let metric logging break training |
There was a problem hiding this comment.
This wraps a collective in try/except, which is the one place that pattern is dangerous. The docstring correctly says flush() all-reduces and therefore every rank must call it exactly once per step — but if it raises on some ranks and not others, those ranks fall through to reset() and continue while the rest are still inside the collective, so the group desynchronizes and the next step hangs or aborts. And because the handler uses warning_rank0, a failure on any rank other than 0 leaves no trace at all.
Suggestions: log with logger.warning plus exc_info=True (or logger.exception) so the failing rank is identifiable, and consider letting a flush() failure propagate — an inconsistent metric schema across ranks is a real misconfiguration, not a cosmetic logging problem. The broad except Exception on the import just above is also redundant, since ensure_deepspec_importable() already ran in __init__.
| ) | ||
|
|
||
| if root not in sys.path: | ||
| sys.path.insert(0, root) |
There was a problem hiding this comment.
Prepending to sys.path[0] from library code means a DeepSpec checkout's top-level directory names take precedence over everything, including the standard library, for the rest of the process. Combined with the automatic sibling-directory scan above (../DeepSpec, ./DeepSpec), behaviour depends on where the repo happens to sit on disk, which is hard to debug when it goes wrong.
Consider narrowing this to the explicit DEEPSPEC_PATH case only (dropping the layout guessing now that the deepspec extra is the supported install path) and using sys.path.append instead. The lru_cache is a good touch — worth noting it also means a DEEPSPEC_PATH change after first call is ignored, which is fine but worth a line in the docstring.
| } | ||
|
|
||
|
|
||
| def _import_draft_class(architecture: str) -> Type: |
There was a problem hiding this comment.
Worth affirming: keeping _ARCHITECTURE_IMPORTS as strings and importing inside this factory is the right call. veomni/models/transformers/__init__.py imports deepspec_draft unconditionally, so an eager import deepspec here would break import veomni.models.transformers for every user who has not installed the optional extra. As written, the dependency is only touched at model-build time.
| super()._setup() | ||
| self._validate_parallelism() | ||
|
|
||
| def _validate_parallelism(self): |
There was a problem hiding this comment.
Good guard, and the docstring explains the gradient-scaling identity clearly. One gap: _setup calls this, but fsdp_size/dp_shard is not checked against the world size directly — the argument relies on "every other dim is 1 implies the shard group spans the world". If fsdp_size can ever be set smaller than the world size independently of dp_replicate_size, the identity breaks without tripping any of these checks. Worth asserting get_parallel_state().fsdp_size == world_size after _setup so the invariant is verified rather than inferred.
3 participants
What does this PR do?
Integrate DeepSpec speculative-decoding draft-model training into VeOmni, so draft models (DSpark / DFlash / Eagle3) train through VeOmni's distributed engine instead of DeepSpec's original single-node
mp.spawn+FSDP(NO_SHARD)loop.The training loop remains decoupled from the target model: DeepSpec builds the target hidden-state cache offline, and VeOmni trains only the draft
PreTrainedModelover that tensor cache. This PR reuses DeepSpec's modeling, loss, and dataset code, and adapts only the interfaces VeOmni needs.This update also routes the Qwen3 DSpark draft model through a patchgen-generated VeOmni modeling file so DSpark's imported Qwen3
RMSNormandMLPcan dispatch throughOpSlot("rms_norm", "standard")andOpSlot("swiglu_mlp", "standard")for Liger-backed runs. RoPE and DeepSpec's custom DSpark loss stay eager because their shape/loss contracts differ from VeOmni's existing generic slots.Checklist Before Starting
[{modules}] {type}: {description}:[trainer, model, data] feat: integrate DeepSpec draft-model training.Test
Validation completed:
compileallfor the DeepSpec integration, model registry, trainer, task, and prep script paths.DEEPSPEC_PATH; missing DeepSpec checkout reports an actionable error.init_device=meta, DCP checkpoint writes, and HF checkpoint export.global_step_2and savedglobal_step_3.rms_norm/swiglu_mlpLiger dispatch.CPU pytest regressions were not clean on the CPU-only master environment for environment reasons, not PR-specific failures: missing
einops, dataloaderpin_memoryrequiring an accelerator, and checkpoint tests requiring CUDA/distributed support.Performance summary from the synthetic production-shape benchmark:
NO_SHARDnvidia-smisamplingInterpretation:
NO_SHARD.6.00 / 5.45 ~= 1.10xover the eager VeOmni path, and about6.00 / 5.00 ~= 1.20xover the native DeepSpec estimate in this synthetic benchmark.API and Usage Example
New task entry, prep script, launcher, configs, and docs. No existing public VeOmni API is changed.
Design & Code Changes
Full integration write-up:
docs/design/deepspec_draft_integration.md.Correctness anchor: DeepSpec scales its loss by
world_sizeto cancel DDP/FSDP gradient averaging. Under pure FSDP data parallelism, where the FSDP shard group spans the full world and all other parallel dimensions are size 1, FSDP2's gradient averaging cancels DeepSpec'sx world_sizescaling exactly.DraftModelTrainer._validate_parallelismenforces this by rejecting SP/TP/EP/PP/CP/HSDP settings for now.Change list:
veomni/integrations/deepspec/: locate and lazily import a DeepSpec checkout viaDEEPSPEC_PATHor a siblingDeepSpec/checkout.veomni/models/transformers/deepspec_draft/: registermodel_type=deepspec_draft,DeepSpecDraftConfig, and DeepSpec draft architectures.veomni/models/transformers/deepspec_draft/dspark_qwen3_gpu_patch_gen_config.py: patchgen config for Qwen3 DSpark draft RMSNorm/SwiGLU OpSlot routing.veomni/models/transformers/deepspec_draft/generated/: generated Qwen3 DSpark GPU modeling file and generated diff.veomni/data/deepspec/: adapt DeepSpecCacheDatasetandCacheCollatorto VeOmni's mapping dataloader contract.veomni/trainer/deepspec/draft_trainer.py:DraftModelTrainer(BaseTrainer)for model build, frozen target embeds/lm_head, dataset/collator setup, DeepSpec loss, grad accumulation, metrics, and checkpoint-compatible training.tasks/train_deepspec_draft.py,scripts/deepspec/*,configs/deepspec/*: task entry, init-checkpoint prep, launcher, and sample configs.Current limits:
Checklist Before Submitting
make quality.docs/design/deepspec_draft_integration.md,configs/deepspec/README.md).tasks/scripts were moved/renamed; a new task entry was added.