From afeeee762ad3c49a863a82e4a3563e6a6fc8bf85 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Fri, 30 Oct 2026 10:41:47 +0000 Subject: [PATCH 1/2] feat(speculative): Muse-Spark-Open DFlash2 recipes that reproduce the released drafter Add the two-stage recipe behind nvidia/Muse-Spark-Open-816B-A42B-DFlash2 under modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/, with a streaming launcher example per stage and a README on the data, the compute and the results: - dflash2.yaml: from scratch, 8 epochs at 4K tokens; stage 2 continues from step 390,000. - dflash2_long_context.yaml: the stage-1 export continued at 12K tokens on the corpus's long samples; the release is step 25,000. Training against this target needed three things main did not have: - The fake base reproduces what Muse-Spark does outside embed_tokens and lm_head: a weightless RMSNorm on the token embeddings, and cap * tanh(multiplier * x / cap) on the logits, for the draft and the teacher alike. It reads the raw config.json when transformers does not know the model_type, as with muse_spark. The final norm is a plain RMSNorm. - DFlash exports write the target's output_multiplier and final_logit_softcapping into dflash_config, the only place vLLM's DFlash2 reads them. Without them the served draft's logits are unscaled and uncapped, which costs acceptance under sampling. - dflash_fused_loss accepts the soft-capped head and applies its transform. The streaming launcher gains SERVE_ENABLE_EP for expert-parallel serve replicas. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- CHANGELOG.rst | 1 + .../torch/export/plugins/hf_spec_export.py | 10 ++ .../torch/speculative/plugins/hf_dflash.py | 15 ++- .../speculative/plugins/modeling_fakebase.py | 126 +++++++++++++++++- .../plugins/modeling_final_norm.py | 3 + .../speculative_decoding/README.md | 62 +++++++++ .../speculative_decoding/dflash2.yaml | 95 +++++++++++++ .../dflash2_long_context.yaml | 87 ++++++++++++ .../speculative/plugins/test_hf_dflash.py | 26 +++- .../torch/export/test_hf_spec_rope_export.py | 1 + .../speculative/plugins/test_fakebase.py | 49 +++++++ .../speculative/plugins/test_hf_dflash.py | 29 ++-- .../speculative/plugins/test_hf_dflash2.py | 25 +++- .../common/eagle3/train_eagle_streaming.sh | 3 + .../hf_streaming_dflash2.yaml | 91 +++++++++++++ .../hf_streaming_dflash2_long_context.yaml | 61 +++++++++ 16 files changed, 658 insertions(+), 26 deletions(-) create mode 100644 modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/README.md create mode 100644 modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2.yaml create mode 100644 modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2_long_context.yaml create mode 100644 tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2.yaml create mode 100644 tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2_long_context.yaml diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1274a86f5e3..cb8ef1e6d89 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -31,6 +31,7 @@ Changelog - With ``training.dataloader_num_workers`` > 0, ``examples/speculative_decoding/main.py`` now yields each training batch as soon as a worker finishes it, so one slow sample, such as a slow hidden-state fetch in streaming training, no longer stalls every DDP rank. Sample order within the prefetch window is then no longer deterministic; set ``training.dataloader_in_order=true`` to restore it (context-parallel runs always keep it). - The ``general/speculative_decoding/dflash2`` recipe now enables ``dflash_use_torch_compile`` and ``dflash_use_flex_attention``, and turns ``tf32`` off, which compilation needs on PyTorch 2.9 and 2.10. To train eagerly as before, start the run with ``dflash.dflash_use_torch_compile=false dflash.dflash_use_flex_attention=false``; a run resumed from a checkpoint keeps the checkpoint's DFlash settings. - Add ``dflash_fused_loss`` to the DFlash config. DFlash, DFlash2 and LiLiCorr training then computes its loss in row chunks fused with the frozen ``lm_head`` projection instead of materializing the draft's full-vocabulary logits, which saves time and memory on large vocabularies. +- Add DFlash2 training recipes for ``meta-models/Muse-Spark-Open-816B-A42B`` under ``modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/``, with streaming launcher examples, that reproduce the released ``nvidia/Muse-Spark-Open-816B-A42B-DFlash2`` drafter. Streaming and offline training against a Muse-Spark target now apply its embedding norm and logit soft cap, and DFlash exports record the target's ``output_multiplier`` and ``final_logit_softcapping`` for vLLM. *Megatron Framework (M-LM / M-Bridge)* diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index f832d8a2961..8f97ddea858 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -453,6 +453,16 @@ def _export_config(self): config["dflash_config"]["attention_sink_bias"] = True config["attention_sink_bias"] = True + # The target's logit transform, which vLLM's DFlash2 reads only from dflash_config + # (default: none). Read off the head the draft logits were trained through. + head = model._base_model_lm_head + for key, attr in ( + ("output_multiplier", "logits_multiplier"), + ("final_logit_softcapping", "logits_soft_cap"), + ): + if (value := getattr(head, attr, None)) is not None: + config["dflash_config"][key] = float(value) + # Inject the export-time YaRN rope_scaling from the dflash_export_rope_scaling # config field (empty dict disables). Mirrors eagle's eagle_export_rope_scaling. export_rope_scaling = getattr(self.model, "dflash_export_rope_scaling", None) diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 6c468d96bec..7cf69968f2b 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -96,6 +96,7 @@ _BASE_MODEL_PATHS, _EMBED_TOKENS_PATHS, _LM_HEAD_PATHS, + _SoftCappedLMHead, find_final_norm_path, ) from .modeling_final_norm import _apply_base_final_norm @@ -610,9 +611,10 @@ def _check_fused_loss(self): ) head = self._base_model_lm_head # The fused pass reads lm_head.weight directly: whatever a bias, a quantizer or a - # hook adds in lm_head's forward would be silently skipped. + # hook adds in lm_head's forward would be silently skipped. The fake base's soft cap + # is the one transform it reproduces (see _fused_head_transform). plain_forward = ( - type(head).forward is torch.nn.Linear.forward + type(head).forward in (torch.nn.Linear.forward, _SoftCappedLMHead.forward) and "forward" not in vars(head) and not (head._forward_hooks or head._forward_pre_hooks) ) @@ -1122,6 +1124,14 @@ def _fused_head_extras(self): """``(statistics with gradient, top_k)`` a variant's own terms read from the fused pass.""" return set(), 0 + def _fused_head_transform(self): + """The logit transform the base ``lm_head`` applies after its projection, if any.""" + head = self._base_model_lm_head + return { + "logit_scale": getattr(head, "logits_multiplier", None), + "soft_cap": getattr(head, "logits_soft_cap", None), + } + def _fused_head_stats(self, draft_hidden, base_outputs, label_positions, flat_targets): """Every statistic of the draft's logits the loss reads, from one fused pass.""" grads, top_k = self._fused_head_extras() @@ -1145,6 +1155,7 @@ def _fused_head_stats(self, draft_hidden, base_outputs, label_positions, flat_ta grads=grads, top_k=top_k, **teacher, + **self._fused_head_transform(), ) def forward( diff --git a/modelopt/torch/speculative/plugins/modeling_fakebase.py b/modelopt/torch/speculative/plugins/modeling_fakebase.py index 85a3fe3a9f4..84a87f3bbc4 100644 --- a/modelopt/torch/speculative/plugins/modeling_fakebase.py +++ b/modelopt/torch/speculative/plugins/modeling_fakebase.py @@ -17,6 +17,8 @@ import json import os +import types +import warnings import torch import torch.nn as nn @@ -96,6 +98,92 @@ def find_final_norm_path(model) -> str | None: # Single-file safetensors names to try, in order. Mistral native checkpoints # use ``consolidated.safetensors`` instead of the HF-standard ``model.safetensors``. _SAFETENSORS_SINGLE_FILENAMES = ["model.safetensors", "consolidated.safetensors"] +# Muse-Spark applies a weightless RMSNorm to its token embeddings and a multiplier and tanh +# soft cap to its logits. Both live in the model's code, not its config, hence the table. +_MUSE_SPARK_MODEL_TYPES = frozenset({"muse_spark", "muse_spark_text"}) + + +class _NormedEmbedding(nn.Embedding): + """``nn.Embedding`` followed by an RMSNorm with no gain, computed in fp32. + + The draft embeds its block tokens through ``embed_tokens``, so the norm has to live in + the module every caller resolves rather than at the call sites. + """ + + def __init__(self, *args, eps: float, **kwargs): + super().__init__(*args, **kwargs) + self.eps = eps + + def forward(self, input_ids): + h = super().forward(input_ids) + h32 = h.float() + return (h32 * torch.rsqrt(h32.pow(2).mean(-1, keepdim=True) + self.eps)).to(h.dtype) + + +class _SoftCappedLMHead(nn.Linear): + """``lm_head`` followed by ``cap * tanh(multiplier * logits / cap)``, computed in fp32. + + Draft and teacher logits both come through this head, so the draft trains against the + distribution the target actually samples from. + """ + + def __init__(self, *args, logits_multiplier=None, logits_soft_cap=None, **kwargs): + super().__init__(*args, **kwargs) + self.logits_multiplier = logits_multiplier + self.logits_soft_cap = logits_soft_cap + + def forward(self, x): + logits = super().forward(x) + out = logits.float() + if self.logits_multiplier is not None: + out = out * self.logits_multiplier + if self.logits_soft_cap is not None: + out = self.logits_soft_cap * torch.tanh(out / self.logits_soft_cap) + return out.to(logits.dtype) + + +def _select_base_transforms(base_cfg) -> dict: + """FakeBaseConfig fields for what the base does outside ``embed_tokens`` and ``lm_head``.""" + if getattr(base_cfg, "model_type", None) not in _MUSE_SPARK_MODEL_TYPES: + return {} + return { + "embed_norm": True, + "logits_multiplier": getattr(base_cfg, "output_multiplier", None), + "logits_soft_cap": getattr(base_cfg, "final_logit_softcapping", None), + } + + +def _load_base_config(source: str, trust_remote_code: bool): + """The checkpoint's config, or its raw ``config.json`` if transformers lacks the model_type. + + FakeBaseModel reads only scalars, so the JSON is enough for them. Anything the model does + outside embed_tokens and lm_head is not in the JSON, which is why _select_base_transforms + and the final-norm table are keyed by model_type. + """ + try: + return transformers.AutoConfig.from_pretrained(source, trust_remote_code=trust_remote_code) + except ValueError as exc: + if "does not recognize this architecture" not in str(exc): + raise + if os.path.isdir(source): + path = os.path.join(source, "config.json") + else: + path = hf_hub_download(repo_id=source, filename="config.json") + with open(path) as f: + raw = json.load(f) + warnings.warn( + f"transformers does not know model_type {raw.get('model_type')!r}; reading {path} " + "directly for the dimensions FakeBaseModel needs.", + stacklevel=2, + ) + return types.SimpleNamespace( + **{ + key: types.SimpleNamespace(**value) + if key in _VLM_CONFIG_ATTRS and isinstance(value, dict) + else value + for key, value in raw.items() + } + ) class FakeBaseConfig(PretrainedConfig): @@ -118,6 +206,9 @@ def __init__( rms_norm_eps=1e-6, rope_theta=None, final_norm_type=None, + embed_norm=False, + logits_multiplier=None, + logits_soft_cap=None, **kwargs, ): """Initialize FakeBaseConfig with minimal model configuration parameters.""" @@ -127,6 +218,11 @@ def __init__( # (model whose final-norm type we don't know). See _FINAL_NORM_CLASSES / # _FINAL_NORM_TYPE_BY_MODEL_TYPE. Persisted so a reloaded config rebuilds the same norm. self.final_norm_type = final_norm_type + # The base's embedding and logit transforms (see _select_base_transforms), persisted + # so a reloaded config rebuilds the same modules. + self.embed_norm = embed_norm + self.logits_multiplier = logits_multiplier + self.logits_soft_cap = logits_soft_cap self.num_hidden_layers = num_hidden_layers # Mirror the original base layer count. The non-fake offline path loads with # num_hidden_layers=0 and stashes the real count here (see utils.load_vlm_or_llm); @@ -184,10 +280,27 @@ def __init__(self, config: FakeBaseConfig, **kwargs): self.model = nn.Module() self.model.layers = nn.ModuleList() self.model.dtype = config.dtype - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, dtype=config.dtype) - self.lm_head = nn.Linear( - config.hidden_size, config.vocab_size, bias=False, dtype=config.dtype - ) + if config.embed_norm: + self.embed_tokens = _NormedEmbedding( + config.vocab_size, config.hidden_size, dtype=config.dtype, eps=config.rms_norm_eps + ) + else: + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, dtype=config.dtype + ) + if config.logits_multiplier is not None or config.logits_soft_cap is not None: + self.lm_head = _SoftCappedLMHead( + config.hidden_size, + config.vocab_size, + bias=False, + dtype=config.dtype, + logits_multiplier=config.logits_multiplier, + logits_soft_cap=config.logits_soft_cap, + ) + else: + self.lm_head = nn.Linear( + config.hidden_size, config.vocab_size, bias=False, dtype=config.dtype + ) # Final pre-lm_head norm, applied before lm_head when reconstructing base logits for # self-logit-distillation (vLLM-captured final hidden states are un-normed). Built ONLY # when the base model's final-norm type is known (config.final_norm_type set); otherwise @@ -212,9 +325,7 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM """ from modelopt.torch.export.plugins.hf_spec_export import _get_rope_theta - orig_config = transformers.AutoConfig.from_pretrained( - source, trust_remote_code=trust_remote_code - ) + orig_config = _load_base_config(source, trust_remote_code) # For vlms, detect language model config based on _VLM_CONFIG_ATTRS base_cfg = next( ( @@ -242,6 +353,7 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM final_norm_type=_select_final_norm_type( getattr(base_cfg, "model_type", None), base_cfg ), + **_select_base_transforms(base_cfg), ) model = cls(config) # Load lm_head, embed_tokens, and (for known models) the final norm into the model. diff --git a/modelopt/torch/speculative/plugins/modeling_final_norm.py b/modelopt/torch/speculative/plugins/modeling_final_norm.py index 28b1ff5a0db..8da7ce11025 100644 --- a/modelopt/torch/speculative/plugins/modeling_final_norm.py +++ b/modelopt/torch/speculative/plugins/modeling_final_norm.py @@ -104,6 +104,9 @@ def __init__(self, hidden_size, eps=1e-6, dtype=torch.bfloat16): # use_gemma_norm flag still gets the correct flavor instead of silently dropping the +1. "minimax_m3_vl_text": "gemma_rmsnorm", "gpt_oss": "gpt_oss_rmsnorm", + # Muse-Spark's final norm uses its weight as stored: no (1 + weight) as in Gemma. + "muse_spark": "rmsnorm", + "muse_spark_text": "rmsnorm", } diff --git a/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/README.md b/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/README.md new file mode 100644 index 00000000000..43537f39ec9 --- /dev/null +++ b/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/README.md @@ -0,0 +1,62 @@ +# Muse-Spark-Open-816B-A42B DFlash2 + +These two recipes train the DFlash2 drafter released as +[nvidia/Muse-Spark-Open-816B-A42B-DFlash2](https://huggingface.co/nvidia/Muse-Spark-Open-816B-A42B-DFlash2) +for [meta-models/Muse-Spark-Open-816B-A42B](https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B). + +| Stage | Recipe | Launcher example | Steps | Sequence | Learning rate | +| --- | --- | --- | --- | --- | --- | +| 1, from scratch | `dflash2.yaml` | `hf_streaming_dflash2.yaml` | 390,000 of 434,744 (8 epochs) | 4,096 | 1e-3, 4% warmup, linear | +| 2, long context | `dflash2_long_context.yaml` | `hf_streaming_dflash2_long_context.yaml` | 25,000 of 59,502 (6 epochs) | 12,288 | 3e-5, 1,500 warmup steps, linear | + +The launcher examples are under `tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/`. +Stage 2 starts from the export of stage 1's checkpoint-390000, and the release is stage 2's +checkpoint-25000. Both stages stream hidden states from the NVFP4 target, +[meta-models/Muse-Spark-Open-816B-A42B-NVFP4](https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B-NVFP4), +served by vLLM. + +## Training data + +The drafter learns from the target's own outputs. The prompts come from +[nvidia/Nemotron-Post-Training-Dataset-v2](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2); +its responses are not used. + +- The target answered each prompt at temperature 1, with up to 8,192 output tokens, at one of + its five reasoning-effort levels: xhigh for 78% of the samples, and minimal, low, medium and + high for about 200,000 each. +- Each response is split into its reasoning and its answer. Responses that reached the token + limit without an answer, and a few malformed ones, were dropped, leaving 3.48 million samples. +- Each row is a `conversations` list: a system turn that sets the reasoning effort the way + the model's chat template writes it, the user turn, and an assistant turn whose `content` is + the answer and whose `reasoning` is the reasoning. + +Stage 2 uses the samples whose conversation, rendered with the chat template, is longer than +4,096 tokens: 634,723 samples, 18% of the corpus. + +Synthesizing the corpus took about 525 node-hours on 4-GPU GB300 nodes, one TP4 replica of the +target per node. + +## Compute + +Both stages ran on 22 nodes of 4 GB300 GPUs: 14 vLLM serve replicas of the target and 8 +trainer nodes. + +| Stage | Run time | GPU-hours | +| --- | --- | --- | +| 1 (390,000 steps) | 37.7 h | ~3,320 | +| 2 (25,000 steps) | 6.0 h | ~530 | + +## Results + +Acceptance length with vLLM, draft length 7, reasoning effort xhigh, 4,096 output tokens, +concurrency 32, on one GB300 node (TP4, EP4). At temperature 1 the draft samples with +`draft_sample_method: probabilistic`. + +| Benchmark | Greedy | Temperature 1 | +| --- | --- | --- | +| MT-Bench | 4.01 | 3.66 | +| SPEED-Bench, qualitative | 4.08 | 3.68 | +| SPEED-Bench, throughput 32K (first 100 per category) | 3.59 | 3.26 | + +Against stage 1's checkpoint-390000, stage 2 raised SPEED-Bench acceptance by 0.5-1.9% and +moved MT-Bench by -0.9% (greedy) and +0.4% (temperature 1). diff --git a/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2.yaml b/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2.yaml new file mode 100644 index 00000000000..52556a0c29d --- /dev/null +++ b/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2.yaml @@ -0,0 +1,95 @@ +# DFlash2 drafter for meta-models/Muse-Spark-Open-816B-A42B, as released in +# nvidia/Muse-Spark-Open-816B-A42B-DFlash2. Stage 1 of 2: trained from scratch for 8 +# epochs at 4K tokens; dflash2_long_context.yaml continues its export at 12K tokens. +# See README.md in this directory for the data, the compute and the results. +# +# Streaming (data.mode=streaming): vLLM serves the NVFP4 target and the trainer pulls its +# hidden states over RDMA, so the trainer holds only the target's embedding, final norm and +# lm_head. The launcher example +# tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2.yaml +# runs it as 14 serve + 8 trainer nodes of 4 GB300 GPUs each. + +# modelopt-schema: modelopt.recipe.config.ModelOptDFlashRecipe +metadata: + description: Muse-Spark-Open DFlash2 drafter, stage 1 (from scratch, 4K tokens, streaming). + +model: + model_name_or_path: + trust_remote_code: false + # The trainer needs only embed_tokens, the final norm and lm_head; the fake base also + # reproduces Muse-Spark's embedding norm and logit soft cap. + use_fake_base_for_offline: true + +data: + mode: streaming + # The target's own responses at temperature 1, five reasoning-effort levels, split into + # reasoning and answer; see README.md. + data_path: + offline_data_path: + # The checkpoint's own template already has {% generation %} tags, and every row + # carries its reasoning effort in its system turn. + chat_template: + +training: + output_dir: + num_train_epochs: 8 + # Global batch 64 = 2 per GPU x 32 trainer GPUs. + per_device_train_batch_size: 2 + gradient_accumulation_steps: 1 + learning_rate: 1.0e-3 + warmup_ratio: 0.04 + training_seq_len: 4096 + logging_steps: 20 + save_steps: 2500 + save_total_limit: 3 + cp_size: 1 + dp_shard_size: 1 + disable_tqdm: true + estimate_ar: false + ar_validate_steps: 0 + answer_only_loss: true + do_eval: false + lr_scheduler_type: linear + save_strategy: steps + weight_decay: 0.0 + max_grad_norm: 1.0 + dataloader_drop_last: true + # A resubmitted job skips the batches it already trained on instead of replaying the + # head of the epoch; with a fixed data_seed that makes a resume exact. + ignore_data_skip: false + data_seed: 20261025 + bf16: true + tf32: true + remove_unused_columns: false + ddp_find_unused_parameters: true + ddp_timeout: 1800 + report_to: tensorboard + +dflash: + dflash_block_size: 8 + dflash_num_anchors: 256 + dflash_swa_window_size: 2048 + dflash_use_torch_compile: true + dflash_use_flex_attention: true + dflash_fused_loss: true + # Distill from the target's logits, D-PACE position weighting. + dflash_self_logit_distillation: true + dflash_lk_loss_type: ce + dflash_loss_objective: dpace + # Muse-Spark has no mask token; 201817 is an unused id. + dflash_mask_token_id: 201817 + dflash_selector_loss_alpha: 1.0 + dflash_fp32_master_weights: true + dflash_architecture_config: + num_hidden_layers: 5 + num_attention_heads: 128 + num_key_value_heads: 16 + head_dim: 64 + intermediate_size: 24576 + projector_type: dflash2 + # The 62-layer target's hidden states after these layers (zero-based) feed the draft. + target_layer_ids: [3, 19, 31, 47, 59] + conv_kernel_size: 2 + conv_group_size: 16 + selector_rank: 256 + selector_top_k: 16 diff --git a/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2_long_context.yaml b/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2_long_context.yaml new file mode 100644 index 00000000000..ad7ed5ae69d --- /dev/null +++ b/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2_long_context.yaml @@ -0,0 +1,87 @@ +# DFlash2 drafter for meta-models/Muse-Spark-Open-816B-A42B, as released in +# nvidia/Muse-Spark-Open-816B-A42B-DFlash2. Stage 2 of 2: continues the stage-1 export +# (dflash2.yaml, step 390,000) at 12K tokens on the corpus's long samples. The release is +# step 25,000 of this schedule; later steps did not improve acceptance. See README.md. +# +# Launcher example: +# tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2_long_context.yaml + +# modelopt-schema: modelopt.recipe.config.ModelOptDFlashRecipe +metadata: + description: Muse-Spark-Open DFlash2 drafter, stage 2 (long-context continuation, 12K tokens). + +model: + model_name_or_path: + trust_remote_code: false + use_fake_base_for_offline: true + +data: + mode: streaming + # The stage-1 corpus, keeping only the samples whose rendered conversation is longer than + # 4,096 tokens (about 18% of it). + data_path: + offline_data_path: + chat_template: + +training: + output_dir: + num_train_epochs: 6 + # Global batch 64 = 1 per GPU x 2 accumulation steps x 32 trainer GPUs. + per_device_train_batch_size: 1 + gradient_accumulation_steps: 2 + # A continuation: low peak, short warmup. + learning_rate: 3.0e-5 + warmup_steps: 1500 + training_seq_len: 12288 + logging_steps: 20 + save_steps: 2500 + save_total_limit: 3 + cp_size: 1 + dp_shard_size: 1 + disable_tqdm: true + estimate_ar: false + ar_validate_steps: 0 + answer_only_loss: true + do_eval: false + lr_scheduler_type: linear + save_strategy: steps + weight_decay: 0.0 + max_grad_norm: 1.0 + dataloader_drop_last: true + ignore_data_skip: false + data_seed: 153313048 + bf16: true + tf32: true + remove_unused_columns: false + ddp_find_unused_parameters: true + ddp_timeout: 1800 + report_to: tensorboard + +dflash: + # The stage-1 drafter as exported (a deployment checkpoint, not a training one). + dflash_init_checkpoint: + dflash_block_size: 8 + # More anchors per sample for the 3x longer sequences. + dflash_num_anchors: 512 + dflash_swa_window_size: 2048 + dflash_use_torch_compile: true + dflash_use_flex_attention: true + dflash_fused_loss: true + dflash_self_logit_distillation: true + dflash_lk_loss_type: ce + dflash_loss_objective: dpace + dflash_mask_token_id: 201817 + dflash_selector_loss_alpha: 1.0 + dflash_fp32_master_weights: true + dflash_architecture_config: + num_hidden_layers: 5 + num_attention_heads: 128 + num_key_value_heads: 16 + head_dim: 64 + intermediate_size: 24576 + projector_type: dflash2 + target_layer_ids: [3, 19, 31, 47, 59] + conv_kernel_size: 2 + conv_group_size: 16 + selector_rank: 256 + selector_top_k: 16 diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py index 9b9569d1205..e2bee8b8edf 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py @@ -25,6 +25,7 @@ import modelopt.torch.speculative as mtsp from modelopt.torch.speculative.plugins.dflash_flex_attention import is_block_mask +from modelopt.torch.speculative.plugins.modeling_fakebase import _SoftCappedLMHead BLOCK_SIZE = 4 NUM_DRAFT_LAYERS = 2 @@ -442,8 +443,17 @@ class TestDFlashFusedLossGPU: } @staticmethod - def _step(fused, overrides, arch): + def _step(fused, overrides, arch, soft_capped_head=False): model = get_tiny_llama(num_hidden_layers=4) + if soft_capped_head: # as the fake base builds it for Muse-Spark + head = _SoftCappedLMHead( + *model.lm_head.weight.shape[::-1], + bias=False, + logits_multiplier=0.5, + logits_soft_cap=3.0, + ) + head.weight.data.copy_(model.lm_head.weight.data) + model.lm_head = head config = get_dflash_config() config.update(dflash_fused_loss=fused, **overrides) config["dflash_architecture_config"].update(arch) @@ -458,19 +468,21 @@ def _step(fused, overrides, arch): return out, {name: p.grad for name, p in params if p.grad is not None} @pytest.mark.parametrize( - ("overrides", "arch"), + ("overrides", "arch", "soft_capped_head"), [ - ({}, {}), # KD on the base's logits, D-PACE weighting + ({}, {}, False), # KD on the base's logits, D-PACE weighting ( {"dflash_self_logit_distillation": False, "dflash_lk_loss_type": "lambda"}, DFLASH2_ARCH, + False, ), + ({}, DFLASH2_ARCH, True), # the Muse-Spark recipe's loss ], - ids=["dflash-kd", "dflash2-selector"], + ids=["dflash-kd", "dflash2-selector", "dflash2-kd-soft-capped-head"], ) - def test_matches_the_materialized_logits(self, overrides, arch): - eager, eager_grads = self._step(False, overrides, arch) - fused, fused_grads = self._step(True, overrides, arch) + def test_matches_the_materialized_logits(self, overrides, arch, soft_capped_head): + eager, eager_grads = self._step(False, overrides, arch, soft_capped_head) + fused, fused_grads = self._step(True, overrides, arch, soft_capped_head) torch.testing.assert_close(fused.loss, eager.loss) assert fused.train_acc == eager.train_acc assert fused_grads.keys() == eager_grads.keys() diff --git a/tests/unit/torch/export/test_hf_spec_rope_export.py b/tests/unit/torch/export/test_hf_spec_rope_export.py index 737aaa46d82..c7b853754d5 100644 --- a/tests/unit/torch/export/test_hf_spec_rope_export.py +++ b/tests/unit/torch/export/test_hf_spec_rope_export.py @@ -122,6 +122,7 @@ def _make_dflash_exporter(dflash_export_rope_scaling=None, base_rope_theta=50000 mask_token_id=999, target_layer_ids=[1, 3, 5, 7], dflash_export_rope_scaling=dflash_export_rope_scaling, + _base_model_lm_head=torch.nn.Linear(128, 1000, bias=False), # no logit transform ) exporter = DFlashExporter.__new__(DFlashExporter) exporter.model = model diff --git a/tests/unit/torch/speculative/plugins/test_fakebase.py b/tests/unit/torch/speculative/plugins/test_fakebase.py index 44cc19ed90c..5521d724116 100644 --- a/tests/unit/torch/speculative/plugins/test_fakebase.py +++ b/tests/unit/torch/speculative/plugins/test_fakebase.py @@ -278,3 +278,52 @@ def __init__(self): self.model = torch.nn.Linear(4, 4) assert find_final_norm_path(_NoNorm()) is None + + +def test_muse_spark_base_reproduces_its_embedding_and_logit_transforms(tmp_path): + """The draft's embeddings and both logit streams must match what the target computes. + + transformers has no ``muse_spark`` config class, so the fake base reads the raw JSON. + The embedding norm and the logit transform are in the model's code, not its config: + dropping either still trains, against the wrong inputs or the wrong teacher. + """ + multiplier, cap, eps = 0.25, 5.0, 1e-5 + text_config = { + "model_type": "muse_spark_text", + "hidden_size": _HIDDEN_SIZE, + "vocab_size": _VOCAB_SIZE, + "num_hidden_layers": 2, + "rms_norm_eps": eps, + "rope_parameters": {"rope_theta": 500000.0, "rope_type": "default"}, + "output_multiplier": multiplier, + "final_logit_softcapping": cap, + "tie_word_embeddings": False, + } + config = {"model_type": "muse_spark", "dtype": "bfloat16", "text_config": text_config} + (tmp_path / "config.json").write_text(json.dumps(config)) + embed, head = torch.randn(_VOCAB_SIZE, _HIDDEN_SIZE), torch.randn(_VOCAB_SIZE, _HIDDEN_SIZE) + safetensors.torch.save_file( + { + "model.language_model.embed_tokens.weight": embed, + "lm_head.weight": head, + "model.language_model.norm.weight": torch.ones(_HIDDEN_SIZE), + }, + tmp_path / "model.safetensors", + ) + + with pytest.warns(UserWarning, match="muse_spark"): + model = FakeBaseModel.from_source(str(tmp_path)) + + ids = torch.tensor([3, 7]) + rows = embed[ids] + expected = rows * torch.rsqrt(rows.pow(2).mean(-1, keepdim=True) + eps) + tol = {"rtol": 2e-2, "atol": 2e-2} # the fake base keeps the checkpoint's bf16 + torch.testing.assert_close(model.embed_tokens(ids).float(), expected, **tol) + x = torch.randn(2, _HIDDEN_SIZE) + logits = model.lm_head(x.to(model.lm_head.weight.dtype)).float() + torch.testing.assert_close(logits, cap * torch.tanh(multiplier * (x @ head.T) / cap), **tol) + assert model.config.final_norm_type == "rmsnorm" + assert model.config.rope_theta == 500000.0 + rebuilt = FakeBaseModel(FakeBaseConfig(**model.config.to_dict())) + assert type(rebuilt.embed_tokens) is type(model.embed_tokens) + assert type(rebuilt.lm_head) is type(model.lm_head) diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index 0aae907c171..e67f3cb42a4 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -56,6 +56,7 @@ MasterWeightAdamW, VerifyMasterWeightsCallback, ) +from modelopt.torch.speculative.plugins.modeling_fakebase import _SoftCappedLMHead from modelopt.torch.speculative.utils import AcceptanceRateValidation from modelopt.torch.utils.plugins.transformers_dataset import LanguageDataCollator @@ -1733,10 +1734,19 @@ class TestDFlashFusedLoss: """``dflash_fused_loss`` trains exactly what the loss on materialized logits trains.""" @staticmethod - def _step(fused, offline, **overrides): + def _step(fused, offline, soft_capped_head=False, **overrides): """One training step's output and draft gradients, from fixed weights, data and anchors.""" torch.manual_seed(0) model = get_tiny_llama(num_hidden_layers=4) + if soft_capped_head: # as the fake base builds it for Muse-Spark + head = _SoftCappedLMHead( + *model.lm_head.weight.shape[::-1], + bias=False, + logits_multiplier=0.5, + logits_soft_cap=3.0, + ) + head.weight.data.copy_(model.lm_head.weight.data) + model.lm_head = head if offline: model.config.num_orig_hidden_layers = 4 config = get_dflash_config(offline=offline) @@ -1757,17 +1767,18 @@ def _step(fused, offline, **overrides): return out, {name: p.grad for name, p in params if p.grad is not None} @pytest.mark.parametrize( - ("offline", "overrides"), + ("offline", "soft_capped_head", "overrides"), [ - (False, {}), # KD on the base's logits, D-PACE weighting - (True, {"dflash_loss_objective": "decay", "dflash_loss_decay_factor": 4.0}), - (False, {"dflash_self_logit_distillation": False}), # cross-entropy, D-PACE + (False, False, {}), # KD on the base's logits, D-PACE weighting + (True, False, {"dflash_loss_objective": "decay", "dflash_loss_decay_factor": 4.0}), + (False, False, {"dflash_self_logit_distillation": False}), # cross-entropy, D-PACE + (True, True, {}), # KD on the base's hidden through a multiplier and soft cap ], - ids=["kd-online", "kd-offline-decay", "ce"], + ids=["kd-online", "kd-offline-decay", "ce", "kd-offline-soft-capped-head"], ) - def test_matches_the_materialized_logits(self, offline, overrides): - eager, eager_grads = self._step(False, offline, **overrides) - fused, fused_grads = self._step(True, offline, **overrides) + def test_matches_the_materialized_logits(self, offline, soft_capped_head, overrides): + eager, eager_grads = self._step(False, offline, soft_capped_head, **overrides) + fused, fused_grads = self._step(True, offline, soft_capped_head, **overrides) torch.testing.assert_close(fused.loss, eager.loss) assert fused.train_acc == eager.train_acc assert fused_grads.keys() == eager_grads.keys() diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py index 67219179bbe..416e852988c 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py @@ -45,6 +45,7 @@ DFlash2Module, DFlashGroupedConv, ) +from modelopt.torch.speculative.plugins.modeling_fakebase import _SoftCappedLMHead BLOCK_SIZE = 4 NUM_DRAFT_LAYERS = 2 @@ -608,8 +609,11 @@ def test_distillation_conflict_is_rejected(self, loss_type): class TestDFlash2Export: """Test the DFlash2 export format (weights + config).""" - def _export(self, tmp_path): + def _export(self, tmp_path, lm_head=None): model = get_tiny_llama(num_hidden_layers=4) + if lm_head is not None: + lm_head.weight.data.copy_(model.lm_head.weight.data) + model.lm_head = lm_head mtsp.convert(model, [("dflash", _get_dflash2_config())]) export_dir = tmp_path / "exported" model.get_exporter().export(export_dir) @@ -654,3 +658,22 @@ def test_export_config_declares_dflash2_architecture(self, tmp_path): assert dflash_config["selector_top_k"] == SELECTOR_TOP_K assert "mask_token_id" in dflash_config assert "target_layer_ids" in dflash_config + assert "output_multiplier" not in dflash_config + assert "final_logit_softcapping" not in dflash_config + + def test_export_config_carries_the_target_logit_transform(self, tmp_path): + """vLLM's DFlash2 applies the target's logit transform only if dflash_config says so. + + Without it the served draft's logits come out unscaled and uncapped while training + saw them transformed: the draft distribution is too sharp, which costs acceptance + under sampling. + """ + model = get_tiny_llama(num_hidden_layers=4) + hidden, vocab = model.config.hidden_size, model.config.vocab_size + head = _SoftCappedLMHead( + hidden, vocab, bias=False, logits_multiplier=0.25, logits_soft_cap=20.0 + ) + with open(self._export(tmp_path, lm_head=head) / "config.json") as f: + dflash_config = json.load(f)["dflash_config"] + assert dflash_config["output_multiplier"] == 0.25 + assert dflash_config["final_logit_softcapping"] == 20.0 diff --git a/tools/launcher/common/eagle3/train_eagle_streaming.sh b/tools/launcher/common/eagle3/train_eagle_streaming.sh index 9edf07c24dc..8675fffb61b 100755 --- a/tools/launcher/common/eagle3/train_eagle_streaming.sh +++ b/tools/launcher/common/eagle3/train_eagle_streaming.sh @@ -50,6 +50,8 @@ # SERVE_BLOCK_SIZE KV-cache block size (e.g. 128 for MiniMax-M3 MSA sparse attention). # Needs its own knob: multi-token SERVE_EXTRA_ARGS values are mangled # by nemo_run's unquoted env export, so "--block-size 128" cannot ride it. +# SERVE_ENABLE_EP 1 -> `vllm serve --enable-expert-parallel` for MoE targets; its own +# knob for the same reason as SERVE_BLOCK_SIZE. # SERVE_HOST single-node: bind/connect host. default 127.0.0.1 # SERVE_GPU single-node: CUDA_VISIBLE_DEVICES for vllm. default "0" # SERVE_TP tensor-parallel size. default 1 single-node / all serve-node GPUs @@ -182,6 +184,7 @@ launch_vllm() { [ -n "${SERVE_MAX_MODEL_LEN:-}" ] && opt_args+=(--max-model-len "$SERVE_MAX_MODEL_LEN") [ -n "${SERVE_MAX_NUM_SEQS:-}" ] && opt_args+=(--max-num-seqs "$SERVE_MAX_NUM_SEQS") [ -n "${SERVE_BLOCK_SIZE:-}" ] && opt_args+=(--block-size "$SERVE_BLOCK_SIZE") + [ "${SERVE_ENABLE_EP:-0}" = 1 ] && opt_args+=(--enable-expert-parallel) # --no-enable-chunked-prefill / --no-enable-prefix-caching: connector captures hidden states during prefill; both skip recomputing cached/partial prefixes, yielding short/empty hidden_states. Required. # --no-enable-flashinfer-autotune: on NVFP4 MoE the autotuner re-tunes on the first serving step and stalls a worker past vLLM's execute-model timeout, killing EngineCore. # Hidden states move serve -> trainer over NIXL RDMA (no disk round-trip): one diff --git a/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2.yaml b/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2.yaml new file mode 100644 index 00000000000..2a9f18de8e3 --- /dev/null +++ b/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2.yaml @@ -0,0 +1,91 @@ +# DFlash2 drafter for meta-models/Muse-Spark-Open-816B-A42B, stage 1 of 2 (streaming): +# reproduces the training of nvidia/Muse-Spark-Open-816B-A42B-DFlash2 up to the checkpoint +# stage 2 starts from. Stage 2 is hf_streaming_dflash2_long_context.yaml next to this file. +# +# The drafter, objective and schedule live in the recipe this points at, +# modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/ +# dflash2.yaml +# (its README.md covers the training data and the results); this file holds the cluster and +# serve settings. +# +# Shape: 14 serve + 8 trainer nodes of 4 GB300 GPUs, ~0.35 s/step (390,000 steps in ~38 h). +# Each serve node runs one replica of the NVFP4 target with TP4 and expert parallelism, and +# one replica delivers ~12 samples/s of 4K tokens, so fewer serve nodes slow the trainers down. +# +# Muse-Spark specifics: +# * EAGLE_CAPTURE_IDS = recipe target_layer_ids [3,19,31,47,59] + 1 (post-layer, 1-based), +# plus 62 = the final layer, whose hidden state is the distillation teacher. +# * SERVE_BLOCK_SIZE=256: vLLM pads the hidden-state page up to the attention page, which +# needs block size >= 32 x the 6 captured layers. +# * --language-model-only: the training data is text, so the serve skips the vision and +# audio towers. +# * HS_MAX_TOKENS, the hidden-state slot size, must cover training.training_seq_len. +# * The script installs examples/speculative_decoding/requirements.txt on every node, which +# pins transformers<5.13, and that version rejects the `window_attention` layer type in the +# checkpoint's vision_config: vllm serve then fails to parse config.json. Serve a copy whose +# config.json says `full_attention` there instead (the weights can be symlinks). The vision +# tower is never built under --language-model-only, so nothing else changes. +# +# Run ON the cluster login node: +# export SLURM_HOST=localhost SLURM_ACCOUNT= \ +# SLURM_PARTITION= \ +# SLURM_HF_LOCAL= \ +# SLURM_JOB_DIR= \ +# NEMORUN_HOME=$PWD +# uv run launch.py --yaml examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2.yaml \ +# identity=$HOME/.ssh/id_ecdsa detach=True --yes +# +# A resubmitted job resumes from the latest checkpoint in training.output_dir. When it ends, +# the final checkpoint is exported to /scratchspace/export; stage 2 starts from the export of +# checkpoint-390000 (scripts/export_hf_checkpoint.py --model_path /checkpoint-390000). + +job_name: Muse-Spark-Open-816B-A42B_DFlash2_streaming +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + # The NVFP4 checkpoint, one replica per 4-GPU GB300 node; its vision_config patched (see header). + hf_model: /hf-local/meta-models/Muse-Spark-Open-816B-A42B-NVFP4 + + task_0: + script: common/eagle3/train_eagle_streaming.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2.yaml + - model.model_name_or_path=<> + - data.data_path= + - training.output_dir=/scratchspace/dflash2 + # The vLLM serve container has no tensorboard. + - training.report_to=none + environment: + - HF_MODEL_CKPT: <> + - EAGLE_CAPTURE_IDS: "[4,20,32,48,60,62]" + - SERVE_NODES: "14" + - SERVE_TP: "4" + - SERVE_ENABLE_EP: "1" + - SERVE_EXTRA_ARGS: "--language-model-only" + - SERVE_BLOCK_SIZE: "256" + - SERVE_MAX_MODEL_LEN: "4608" + - SERVE_MAX_NUM_SEQS: "32" + - SERVE_GPU_MEM_UTIL: "0.9" + # Loading 500 GB of weights on every replica can take tens of minutes. + - SERVE_READY_TIMEOUT: "3600" + - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200" + - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200" + - HS_MAX_TOKENS: "4608" + - HS_POOL_SLOTS: "64" + - STREAMING_NUM_WORKERS: "8" + # RDMA transport is UCX (InfiniBand) by default. On AWS EFA, uncomment: + # - NIXL_BACKENDS: "LIBFABRIC" + # - FI_PROVIDER: "efa" + # - NCCL_IB_DISABLE: "1" + slurm_config: + _factory_: "slurm_factory" + nodes: 22 + ntasks_per_node: 1 + gpus_per_node: 4 + # The launcher asks for 4 h unless told otherwise. + time: "16:00:00" + # A vLLM build that serves Muse-Spark-Open with the extract_hidden_states method. + container: diff --git a/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2_long_context.yaml b/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2_long_context.yaml new file mode 100644 index 00000000000..0e2a9f66d9a --- /dev/null +++ b/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2_long_context.yaml @@ -0,0 +1,61 @@ +# DFlash2 drafter for meta-models/Muse-Spark-Open-816B-A42B, stage 2 of 2 (streaming): +# continues the stage-1 export (hf_streaming_dflash2.yaml, checkpoint-390000) at 12K tokens +# on the corpus's long samples. The released nvidia/Muse-Spark-Open-816B-A42B-DFlash2 is +# checkpoint-25000 of this run (~6 h on this shape); later checkpoints did not improve +# acceptance. Export it with +# scripts/export_hf_checkpoint.py --model_path /checkpoint-25000 +# +# Recipe: +# modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/ +# dflash2_long_context.yaml +# Serve settings as in stage 1, except that each request is 3x longer: HS_MAX_TOKENS and +# SERVE_MAX_MODEL_LEN cover the 12,288-token sequences and SERVE_MAX_NUM_SEQS is lower. +# +# Run ON the cluster login node, with the environment from stage 1: +# uv run launch.py --yaml examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2_long_context.yaml \ +# identity=$HOME/.ssh/id_ecdsa detach=True --yes + +job_name: Muse-Spark-Open-816B-A42B_DFlash2_long_context_streaming +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + # The same patched copy as stage 1 (see that file's header). + hf_model: /hf-local/meta-models/Muse-Spark-Open-816B-A42B-NVFP4 + + task_0: + script: common/eagle3/train_eagle_streaming.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2_long_context.yaml + - model.model_name_or_path=<> + # The stage-1 export (a deployment checkpoint, not a training one). + - dflash.dflash_init_checkpoint= + - data.data_path= + - training.output_dir=/scratchspace/dflash2_long_context + - training.report_to=none + environment: + - HF_MODEL_CKPT: <> + - EAGLE_CAPTURE_IDS: "[4,20,32,48,60,62]" + - SERVE_NODES: "14" + - SERVE_TP: "4" + - SERVE_ENABLE_EP: "1" + - SERVE_EXTRA_ARGS: "--language-model-only" + - SERVE_BLOCK_SIZE: "256" + - SERVE_MAX_MODEL_LEN: "12800" + - SERVE_MAX_NUM_SEQS: "8" + - SERVE_GPU_MEM_UTIL: "0.9" + - SERVE_READY_TIMEOUT: "3600" + - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200" + - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200" + - HS_MAX_TOKENS: "12288" + - HS_POOL_SLOTS: "64" + - STREAMING_NUM_WORKERS: "8" + slurm_config: + _factory_: "slurm_factory" + nodes: 22 + ntasks_per_node: 1 + gpus_per_node: 4 + time: "16:00:00" + container: From 23d2a6c41991638f103931acb9925bb78ae9ce2c Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Fri, 30 Oct 2026 16:18:14 +0000 Subject: [PATCH 2/2] fix(speculative): pin the Muse-Spark-Open stage-2 horizon; say when to stop each stage Stage 2 now pins max_steps at the release's 59,502. Left to num_train_epochs, transformers rounds each epoch of the 634,723 long samples up to 9,918 optimizer steps, a 59,508-step schedule. Stage 2 starts from stage 1's checkpoint-390000 and the release is stage 2's checkpoint-25000, but neither is the last checkpoint of its run and both recipes keep only the last three, so a run left to finish deletes it. The launcher examples and the README now say to stop each run once its checkpoint is written, or to copy the checkpoint out. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../speculative_decoding/README.md | 4 +++- .../speculative_decoding/dflash2_long_context.yaml | 3 +++ .../Muse-Spark-Open-816B-A42B/hf_streaming_dflash2.yaml | 8 +++++--- .../hf_streaming_dflash2_long_context.yaml | 3 ++- 4 files changed, 13 insertions(+), 5 deletions(-) diff --git a/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/README.md b/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/README.md index 43537f39ec9..2052d452359 100644 --- a/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/README.md +++ b/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/README.md @@ -11,7 +11,9 @@ for [meta-models/Muse-Spark-Open-816B-A42B](https://huggingface.co/meta-models/M The launcher examples are under `tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/`. Stage 2 starts from the export of stage 1's checkpoint-390000, and the release is stage 2's -checkpoint-25000. Both stages stream hidden states from the NVFP4 target, +checkpoint-25000. Neither is the last checkpoint of its run, and the recipes keep only the last +three, so stop each run once its checkpoint is written, or copy the checkpoint out. Both stages +stream hidden states from the NVFP4 target, [meta-models/Muse-Spark-Open-816B-A42B-NVFP4](https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B-NVFP4), served by vLLM. diff --git a/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2_long_context.yaml b/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2_long_context.yaml index ad7ed5ae69d..5c63a36958b 100644 --- a/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2_long_context.yaml +++ b/modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/dflash2_long_context.yaml @@ -26,6 +26,9 @@ data: training: output_dir: num_train_epochs: 6 + # The release's horizon, 6 epochs at global batch 64: from num_train_epochs alone, + # transformers rounds each epoch up to 9,918 steps (59,508 in all). + max_steps: 59502 # Global batch 64 = 1 per GPU x 2 accumulation steps x 32 trainer GPUs. per_device_train_batch_size: 1 gradient_accumulation_steps: 2 diff --git a/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2.yaml b/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2.yaml index 2a9f18de8e3..4a05a7f828b 100644 --- a/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2.yaml +++ b/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2.yaml @@ -35,9 +35,11 @@ # uv run launch.py --yaml examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2.yaml \ # identity=$HOME/.ssh/id_ecdsa detach=True --yes # -# A resubmitted job resumes from the latest checkpoint in training.output_dir. When it ends, -# the final checkpoint is exported to /scratchspace/export; stage 2 starts from the export of -# checkpoint-390000 (scripts/export_hf_checkpoint.py --model_path /checkpoint-390000). +# A resubmitted job resumes from the latest checkpoint in training.output_dir. Stage 2 starts +# from the export of checkpoint-390000, not of the last checkpoint, and the recipe keeps only +# the last three: stop the job once checkpoint-390000 is written (the release did not use the +# remaining 44,744 steps), or copy it out before checkpoint-397500 is written. Export it with +# scripts/export_hf_checkpoint.py --model_path /checkpoint-390000 job_name: Muse-Spark-Open-816B-A42B_DFlash2_streaming pipeline: diff --git a/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2_long_context.yaml b/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2_long_context.yaml index 0e2a9f66d9a..9d0889908c6 100644 --- a/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2_long_context.yaml +++ b/tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2_long_context.yaml @@ -2,7 +2,8 @@ # continues the stage-1 export (hf_streaming_dflash2.yaml, checkpoint-390000) at 12K tokens # on the corpus's long samples. The released nvidia/Muse-Spark-Open-816B-A42B-DFlash2 is # checkpoint-25000 of this run (~6 h on this shape); later checkpoints did not improve -# acceptance. Export it with +# acceptance. The recipe keeps only the last three checkpoints, so stop the job once +# checkpoint-25000 is written, or copy it out before checkpoint-32500 is written. Export it with # scripts/export_hf_checkpoint.py --model_path /checkpoint-25000 # # Recipe: