diff --git a/examples/speculative_decoding/doc/dflash.md b/examples/speculative_decoding/doc/dflash.md index c468c7652be..640207b065c 100644 --- a/examples/speculative_decoding/doc/dflash.md +++ b/examples/speculative_decoding/doc/dflash.md @@ -281,6 +281,30 @@ DFlash supports checkpoint resume transparently. Rotary embeddings are lazily initialized on first forward (matching EAGLE3's `_maybe_init_rope` pattern), avoiding meta-tensor issues during `from_pretrained` model construction. +### Warm Start / Fine-Tuning (`dflash_init_checkpoint`) + +To fine-tune an already-trained drafter instead of training from scratch, set +`dflash.dflash_init_checkpoint` to a local directory containing an exported +DFlash checkpoint (the z-lab-compatible `model.safetensors` layout produced by +`export_hf_checkpoint.py`; for a public drafter, download it first, e.g. +`hf download z-lab/Qwen3-8B-DFlash-b16`): + +```yaml +dflash: + dflash_init_checkpoint: /path/to/Qwen3-8B-DFlash-b16 +``` + +The weights are loaded into `model.dflash_module` right after conversion. +`dflash_architecture_config` must describe the same draft architecture as the +checkpoint — note that draft head/MLP dims default to `Qwen3Config` values, not +the base model's, so set `num_attention_heads` / `num_key_value_heads` / +`head_dim` / `intermediate_size` explicitly. `dflash_block_size` may differ from +the checkpoint's (weights are block-size agnostic); a `mask_token_id` or +`target_layer_ids` mismatch loads but logs a degradation warning. Restore of a +saved training checkpoint ignores the field. See +`tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash_finetune.yaml` for a +full fine-tuning pipeline. + ### Export ```bash diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index 4df57d3035f..4b9c43e0986 100644 --- a/modelopt/torch/speculative/config.py +++ b/modelopt/torch/speculative/config.py @@ -144,6 +144,19 @@ class DFlashConfig(ModeloptBaseConfig): default={}, description="Config for the DFlash draft module architecture." ) + dflash_init_checkpoint: str | None = ModeloptField( + default=None, + description=( + "Warm-start the draft module from an exported DFlash checkpoint before " + "fine-tuning. Must be a local directory containing ``model.safetensors`` in the " + "DFlashExporter / z-lab layout (e.g. a downloaded copy of " + "``z-lab/Qwen3-8B-DFlash-b16``). ``dflash_architecture_config`` must describe " + "the same draft architecture as the checkpoint (weights are block-size " + "agnostic, so ``dflash_block_size`` may differ). Ignored on restore — restored " + "models load their own trained weights." + ), + ) + dflash_use_torch_compile: bool = ModeloptField( default=True, description="Whether to use torch.compile on DFlash forward/loss methods.", diff --git a/modelopt/torch/speculative/dflash/conversion.py b/modelopt/torch/speculative/dflash/conversion.py index 0516ab7181d..7ec118ef04d 100644 --- a/modelopt/torch/speculative/dflash/conversion.py +++ b/modelopt/torch/speculative/dflash/conversion.py @@ -80,4 +80,7 @@ def restore_dflash_model( ) -> nn.Module: """Function for restoring a previously converted model to a DFlash model.""" assert not metadata, "No metadata expected!" + # Never warm-start on restore: the restored model loads its own trained weights, + # and the init checkpoint may no longer exist where the model was trained. + config.dflash_init_checkpoint = None return convert_to_dflash_model(model, config)[0] diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 7679ff0020b..ab6d6a219e0 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -71,7 +71,9 @@ lazy rope pattern needed for MLA models. """ +import json import logging +import os import torch import torch.nn.functional as F @@ -306,6 +308,12 @@ def modify(self, config): if base_device.type != "meta": self.dflash_module.to(self._base_model.dtype).to(base_device) + # Warm-start the draft module from an exported DFlash checkpoint (fine-tuning). + # Restore paths never reach here with a checkpoint set: restore_dflash_model + # clears the field since restored weights come from the restored model itself. + if config.dflash_init_checkpoint: + self._load_draft_init_weights(config.dflash_init_checkpoint) + # Delete base model layers for offline training (save memory) if self.dflash_offline: self._base_model._modules.pop("layers") @@ -317,6 +325,91 @@ def _build_draft_module(self, dflash_config): """Build the draft module. Subclasses override to use an augmented module.""" return DFlashModule(dflash_config) + def _load_draft_init_weights(self, checkpoint: str): + """Warm-start ``self.dflash_module`` from an exported DFlash draft checkpoint. + + ``checkpoint`` is a local directory holding ``model.safetensors`` (DFlashExporter / + z-lab layout). Keys map 1:1 onto ``dflash_module``. Missing keys are tolerated + only for submodules entirely absent from the checkpoint (e.g. the Domino/DSpark + heads when warm-starting from a plain DFlash backbone); partial overlap within a + submodule means an architecture mismatch and raises. + """ + from safetensors.torch import load_file + + weights_path = os.path.join(checkpoint, "model.safetensors") + if not os.path.isfile(weights_path): + raise ValueError( + f"dflash_init_checkpoint must be a local directory containing " + f"model.safetensors; not found at {weights_path}." + ) + + self._check_init_checkpoint_config(checkpoint, weights_path) + + param = next(self.dflash_module.parameters()) + state_dict = { + key: value.to(param.dtype) + for key, value in load_file(weights_path, device=str(param.device)).items() + } + arch_hint = ( + "Set dflash_architecture_config to the checkpoint's architecture " + "(num_hidden_layers, heads, intermediate_size, ...)." + ) + try: + # strict=False tolerates missing/unexpected keys but still raises on shape + # mismatches (e.g. a different draft depth changing the fc fusion width). + missing, unexpected = self.dflash_module.load_state_dict(state_dict, strict=False) + except RuntimeError as e: + raise ValueError( + f"DFlash init checkpoint {checkpoint} does not match the draft architecture " + f"built from dflash_architecture_config: {e} {arch_hint}" + ) from e + ckpt_submodules = {key.split(".", 1)[0] for key in state_dict} + hard_missing = [key for key in missing if key.split(".", 1)[0] in ckpt_submodules] + if unexpected or hard_missing: + raise ValueError( + f"DFlash init checkpoint {checkpoint} does not match the draft architecture " + f"built from dflash_architecture_config (unexpected keys: {unexpected}, " + f"missing keys: {hard_missing}). {arch_hint}" + ) + if missing: + logger.warning( + "DFlash: submodules not in init checkpoint %s train from scratch: %s", + checkpoint, + sorted({key.split(".", 1)[0] for key in missing}), + ) + logger.info( + "DFlash: warm-started draft module from %s (%d tensors).", + checkpoint, + len(state_dict), + ) + + def _check_init_checkpoint_config(self, checkpoint: str, weights_path: str): + """Warn when the init checkpoint's config disagrees with this conversion. + + A ``mask_token_id`` or ``target_layer_ids`` mismatch loads cleanly (shapes match) + but silently degrades the warm start: the mask embedding and the fc fusion columns + were trained for the checkpoint's values. + """ + config_path = os.path.join(os.path.dirname(weights_path), "config.json") + if not os.path.isfile(config_path): + return + with open(config_path) as f: + ckpt_dflash_config = json.load(f).get("dflash_config", {}) + for name, ours in ( + ("mask_token_id", self.mask_token_id), + ("target_layer_ids", list(self.target_layer_ids)), + ): + ckpt_value = ckpt_dflash_config.get(name) + if ckpt_value is not None and ckpt_value != ours: + logger.warning( + "DFlash: init checkpoint %s was trained with %s=%s but this conversion " + "uses %s — the warm start will be degraded.", + checkpoint, + name, + ckpt_value, + ours, + ) + def get_exporter(self): """Get the exporter for the DFlash draft model.""" from modelopt.torch.export.plugins.hf_spec_export import DFlashExporter diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index ef2eec2ad07..a8ff21f0d52 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -279,6 +279,69 @@ def test_save_and_restore(self, tmp_path): tf_modelopt_state_and_output_tester(model_ref, model_test) +class TestDFlashWarmStart: + """Test warm-starting the draft module from an exported DFlash checkpoint.""" + + def _export_drafter(self, tmp_path, **config_overrides): + """Convert a tiny model and export its drafter in z-lab layout.""" + model = get_tiny_llama(num_hidden_layers=4) + config = _get_dflash_config() + config.update(config_overrides) + mtsp.convert(model, [("dflash", config)]) + export_dir = tmp_path / "exported_drafter" + model.get_exporter().export(export_dir) + return model, export_dir + + def test_warm_start_loads_exported_weights(self, tmp_path): + """A fresh conversion with dflash_init_checkpoint matches the exported drafter. + + Fine-tunes at a different block size: weights are block-size agnostic. + """ + model_ref, export_dir = self._export_drafter(tmp_path) + + model = get_tiny_llama(num_hidden_layers=4) + config = _get_dflash_config(block_size=2 * BLOCK_SIZE) + config["dflash_init_checkpoint"] = str(export_dir) + mtsp.convert(model, [("dflash", config)]) + + ref_sd = model_ref.dflash_module.state_dict() + for key, value in model.dflash_module.state_dict().items(): + assert torch.equal(value, ref_sd[key]), f"Mismatch after warm start: {key}" + + def test_warm_start_bad_checkpoint_raises(self, tmp_path): + """A missing directory or a different draft architecture fails loudly.""" + _, export_dir = self._export_drafter(tmp_path) + + config = _get_dflash_config() + config["dflash_init_checkpoint"] = str(tmp_path / "no_such_dir") + with pytest.raises(ValueError, match="must be a local directory"): + mtsp.convert(get_tiny_llama(num_hidden_layers=4), [("dflash", config)]) + + config = _get_dflash_config(num_layers=NUM_DRAFT_LAYERS + 1) + config["dflash_init_checkpoint"] = str(export_dir) + with pytest.raises(ValueError, match="does not match the draft architecture"): + mtsp.convert(get_tiny_llama(num_hidden_layers=4), [("dflash", config)]) + + def test_restore_ignores_init_checkpoint(self, tmp_path): + """Save/restore of a warm-started model must not re-read the init checkpoint.""" + import shutil + + _, export_dir = self._export_drafter(tmp_path) + + mto.enable_huggingface_checkpointing() + model_ref = get_tiny_llama(num_hidden_layers=4) + config = _get_dflash_config() + config["dflash_init_checkpoint"] = str(export_dir) + mtsp.convert(model_ref, [("dflash", config)]) + model_ref.save_pretrained(tmp_path / "warm_started_model") + + # The init checkpoint is gone; restore must still work. + shutil.rmtree(export_dir) + model_test = AutoModelForCausalLM.from_pretrained(tmp_path / "warm_started_model") + assert isinstance(model_test, HFDFlashModel) + tf_modelopt_state_and_output_tester(model_ref, model_test) + + class TestDFlashLazyRotaryEmb: """Test lazy rotary embedding initialization (matching EAGLE3 pattern). diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash_finetune.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash_finetune.yaml new file mode 100644 index 00000000000..5d2f23032fb --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash_finetune.yaml @@ -0,0 +1,107 @@ +# DFlash online FINE-TUNING for Qwen3-8B, warm-started from an exported DFlash checkpoint. +# +# Instead of training the drafter from scratch, dflash.dflash_init_checkpoint loads an +# already-trained draft module (DFlashExporter / z-lab layout: model.safetensors + +# config.json) into model.dflash_module before training. It must be a LOCAL directory — +# stage the public drafter first (hf download z-lab/Qwen3-8B-DFlash-b16) or point +# init_checkpoint at your own export directory. +# +# NOTE: dflash_architecture_config must describe the SAME draft architecture as the +# checkpoint. modify() takes num_attention_heads/num_key_value_heads/head_dim/ +# intermediate_size from Qwen3Config defaults (NOT the base model), so set them +# explicitly below; hidden_size/vocab/rope_theta are forced to the base model and need +# not be set. A mismatch fails at conversion with a shape error. dflash_block_size may +# differ from the checkpoint's (weights are block-size agnostic), but mask_token_id and +# target_layer_ids should match — the loader warns if they don't. +# +# 3-step pipeline: +# task_0: Online DFlash fine-tuning (warm start) +# task_1: vLLM smoke test with DFlash speculative decoding +# task_2: MT-Bench per-category HF AR evaluation (1 GPU) +# +# The warm start is visible in the very first training logs: step 100 accuracy should +# already be near the converged from-scratch level (~0.2, see hf_online_dflash.yaml) +# instead of ~0.03. +# +# Regression criteria (set via environment): +# MAX_FINAL_LOSS: final loss must be below this (default: 4.5) +# MIN_FINAL_ACC: final accuracy must be above this (default: 0.18) +# +# Usage: +# uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_online_dflash_finetune.yaml --yes +# uv run slurm.py --yaml modules/Model-Optimizer/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash_finetune.yaml --yes + +job_name: Qwen3-8B_DFlash_online_finetune +pipeline: + global_vars: + hf_model: /hf-local/Qwen/Qwen3-8B + # Local dir with model.safetensors (a staged z-lab/Qwen3-8B-DFlash-b16 or your own export). + init_checkpoint: /hf-local/z-lab/Qwen3-8B-DFlash-b16 + + # Step 1: Online DFlash fine-tuning from the warm-start checkpoint + task_0: + script: common/specdec/dflash_online_training.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dflash.yaml + - model.model_name_or_path=<> + - data.data_path=/hf-local/modelopt/Speculative-Decoding-Dataset-v1-Qwen3-8B/sample-100K-openai.jsonl + - data.chat_template=examples/Qwen/Qwen3-8B/chat_template_train.jinja + - training.output_dir=/scratchspace/dflash_bs16_finetune + - training.per_device_train_batch_size=1 + - training.num_train_epochs=1 + - training.training_seq_len=4096 + - training.save_steps=5000 + - training.logging_steps=100 + - training.disable_tqdm=true + - training.answer_only_loss=true + # Fine-tuning: lower LR than the 6e-4 from-scratch default. + - training.learning_rate=1.0e-4 + - dflash.dflash_init_checkpoint=<> + # Match z-lab/Qwen3-8B-DFlash-b16: block_size 16, mask token 151669, + # 5 draft layers with Qwen3-8B dims. + - dflash.dflash_block_size=16 + - dflash.dflash_num_anchors=512 + - dflash.dflash_loss_decay_factor=7 + - dflash.dflash_mask_token_id=151669 + - dflash.dflash_architecture_config.num_hidden_layers=5 + - dflash.dflash_architecture_config.num_attention_heads=32 + - dflash.dflash_architecture_config.num_key_value_heads=8 + - dflash.dflash_architecture_config.head_dim=128 + - dflash.dflash_architecture_config.intermediate_size=12288 + environment: + - MAX_FINAL_LOSS: "4.5" + - MIN_FINAL_ACC: "0.18" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + + # Step 2: vLLM smoke test (uses exported checkpoint from training) + task_1: + script: common/specdec/vllm_smoke_test.sh + environment: + - HF_MODEL_CKPT: <> + - DRAFT_CKPT_DIR: /scratchspace/dflash_bs16_finetune + - SPEC_METHOD: "dflash" + - NUM_SPEC_TOKENS: "7" + - MIN_ACCEPTANCE_LENGTH: "1.4" + slurm_config: + _factory_: "slurm_factory" + container: "vllm/vllm-openai:nightly" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + + # Step 3: HF AR evaluation + task_2: + script: common/specdec/ar_eval_mtbench.sh + args: + - --ckpt_dir /scratchspace/dflash_bs16_finetune + - --osl 512 + - --steps 15 + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1