diff --git a/configs/sglang_qwen3_8b_domino_2gpu.yaml b/configs/sglang_qwen3_8b_domino_2gpu.yaml new file mode 100644 index 00000000..a0471f57 --- /dev/null +++ b/configs/sglang_qwen3_8b_domino_2gpu.yaml @@ -0,0 +1,90 @@ +# Domino training config for Qwen3-8B on a 2-GPU RunPod. +# +# GPU allocation: +# - 1 GPU for SGLang target inference +# - 1 GPU for Domino training +# +# First-run calibration: +# CUDA_VISIBLE_DEVICES=0,1 python -m torchspec.train_entry \ +# --config configs/sglang_qwen3_8b_domino_2gpu.yaml \ +# dataset.train_data_path=/path/to/perfectblend.jsonl \ +# training.num_train_steps=100 \ +# output_dir=./outputs/qwen3-8b-domino-2gpu-100 +# +# Review experiment: +# Run matching DFlash and Domino jobs with the same dataset, seed, and step count. + +model: + target_model_path: Qwen/Qwen3-8B + trust_remote_code: true + draft_model_config: torchspec/config/domino_draft_config.json + +dataset: + train_data_path: ../examples/data/sample_conversations.jsonl + eval_data_path: null + eval_interval: 100 + chat_template: qwen + prompt_key: conversations + min_loss_tokens: 32 + +training: + attention_backend: flex_attention + micro_batch_size: 1 + draft_accumulation_steps: 16 + learning_rate: 6e-4 + min_lr: 0.0 + weight_decay: 0.0 + max_concurrent_batches: 1 + max_grad_norm: 1.0 + max_seq_length: 3072 + num_epochs: 3 + num_train_steps: 100 + seed: 42 + training_num_gpus_per_node: 1 + training_num_nodes: 1 + ttt_length: 7 + fsdp_strategy: REPLICATE + fsdp_reduce_dtype: bfloat16 + prefetch_depth: 4 + save_interval: 500 + save_per_epoch: false + max_checkpoints: 2 + warmup_ratio: 0.04 + + dflash_block_size: 16 + dflash_num_anchors: 256 + dflash_loss_decay_gamma: 7.0 + dflash_num_target_layers: 5 + + # None means lambda decays across lr_total_steps / num_train_steps. + domino_curriculum_steps: null + +inference: + inference_engine_type: sgl + store_last_hidden_states: false + inference_num_gpus: 1 + inference_num_gpus_per_engine: 1 + inference_num_gpus_per_node: 2 + max_sample_pool_size: 64 + inference_buffer_threshold: 32 + inference_batch_size: 8 + sglang: + tp_size: 1 + mem_fraction_static: 0.7 + +mooncake: + master_server_address: null + metadata_server: null + protocol: tcp + global_segment_size: 16GB + local_buffer_size: 4GB + enable_hard_pin: true + +output_dir: ./outputs/qwen3-8b-domino-2gpu +cache_dir: ./cache/qwen3-8b-domino-2gpu +model_download_dir: null + +debug: + save_debug_train_data: null + debug_train_only: false + debug_inference_only: false diff --git a/examples/qwen3-8b-domino-8h100/run.sh b/examples/qwen3-8b-domino-8h100/run.sh new file mode 100755 index 00000000..2fdadc89 --- /dev/null +++ b/examples/qwen3-8b-domino-8h100/run.sh @@ -0,0 +1,116 @@ +#!/bin/bash +# Qwen3-8B Domino 8-GPU verification run. +# +# Default GPU allocation: +# - 4 GPUs for SGLang inference, one full model copy per engine +# - 4 GPUs for Domino training with FSDP FULL_SHARD +# +# Usage: +# TRAIN_DATA_PATH=/path/to/perfectblend_10k.jsonl \ +# OUTPUT_ROOT=/path/to/durable/output \ +# ./examples/qwen3-8b-domino-8h100/run.sh +# +# Optional: +# ./examples/qwen3-8b-domino-8h100/run.sh configs/sglang_qwen3_8b_domino_2gpu.yaml \ +# training.num_train_steps=20 + +set -euo pipefail +set -x + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" +export TORCHSPEC_LOG_LEVEL="${TORCHSPEC_LOG_LEVEL:-INFO}" +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" +export PYTORCH_ALLOC_CONF="${PYTORCH_ALLOC_CONF:-expandable_segments:True}" +export TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS="${TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS:-ATEN,TRITON}" +export MC_STORE_MEMCPY="${MC_STORE_MEMCPY:-0}" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +ROOT_DIR="$(dirname "$(dirname "$SCRIPT_DIR")")" +cd "$ROOT_DIR" + +CONFIG_FILE="${1:-$ROOT_DIR/configs/sglang_qwen3_8b_domino_2gpu.yaml}" +if [[ -f "$CONFIG_FILE" ]]; then + shift 1 || true +elif [[ -f "$ROOT_DIR/$CONFIG_FILE" ]]; then + CONFIG_FILE="$ROOT_DIR/$CONFIG_FILE" + shift 1 || true +else + CONFIG_FILE="$ROOT_DIR/configs/sglang_qwen3_8b_domino_2gpu.yaml" +fi + +IFS=',' read -ra GPU_ARRAY <<< "$CUDA_VISIBLE_DEVICES" +TOTAL_GPUS="${#GPU_ARRAY[@]}" +if [[ "$TOTAL_GPUS" -lt 8 ]]; then + echo "Expected at least 8 visible GPUs, got ${TOTAL_GPUS}: ${CUDA_VISIBLE_DEVICES}" >&2 + exit 1 +fi + +RUN_NAME="${RUN_NAME:-qwen3-8b-domino-8h100-100}" +OUTPUT_ROOT="${OUTPUT_ROOT:-$ROOT_DIR/outputs}" +OUTPUT_DIR="${OUTPUT_DIR:-$OUTPUT_ROOT/$RUN_NAME}" +TRAIN_DATA_PATH="${TRAIN_DATA_PATH:-$ROOT_DIR/data/perfectblend_10k.jsonl}" +NUM_TRAIN_STEPS="${NUM_TRAIN_STEPS:-100}" +SAVE_INTERVAL="${SAVE_INTERVAL:-100}" +MAX_CHECKPOINTS="${MAX_CHECKPOINTS:-1}" + +TRAIN_GPUS="${TRAIN_GPUS:-4}" +INFERENCE_GPUS="${INFERENCE_GPUS:-4}" +TP_SIZE="${TP_SIZE:-1}" +DRAFT_ACCUMULATION_STEPS="${DRAFT_ACCUMULATION_STEPS:-4}" +PREFETCH_DEPTH="${PREFETCH_DEPTH:-8}" + +export HF_HOME="${HF_HOME:-$ROOT_DIR/hf-cache}" +export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-$ROOT_DIR/cache/compiled_kernels}" +export TORCHSPEC_LOG_DIR="${TORCHSPEC_LOG_DIR:-$OUTPUT_DIR/rank-logs}" + +if [[ ! -f "$TRAIN_DATA_PATH" ]]; then + echo "Training data not found: ${TRAIN_DATA_PATH}" >&2 + echo "Set TRAIN_DATA_PATH to the JSONL dataset used for the verification run." >&2 + exit 1 +fi + +mkdir -p "$OUTPUT_DIR" "$TORCHSPEC_LOG_DIR" "$ROOT_DIR/cache" + +LOG_FILE="$OUTPUT_DIR/launcher.log" +exec > >(tee -a "$LOG_FILE") 2>&1 + +echo "==============================================" +echo "Qwen3-8B Domino 8-GPU verification" +echo "==============================================" +echo "Config: $CONFIG_FILE" +echo "Run name: $RUN_NAME" +echo "Output dir: $OUTPUT_DIR" +echo "Training data: $TRAIN_DATA_PATH" +echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" +echo "Training GPUs: $TRAIN_GPUS" +echo "Inference GPUs: $INFERENCE_GPUS" +echo "Inference TP size: $TP_SIZE" +echo "Steps: $NUM_TRAIN_STEPS" +echo "Save interval: $SAVE_INTERVAL" +echo "Max checkpoints: $MAX_CHECKPOINTS" +echo "Extra args: $*" +echo "==============================================" + +python3 -m torchspec.train_entry \ + --config "$CONFIG_FILE" \ + dataset.train_data_path="$TRAIN_DATA_PATH" \ + training.num_train_steps="$NUM_TRAIN_STEPS" \ + training.save_interval="$SAVE_INTERVAL" \ + training.max_checkpoints="$MAX_CHECKPOINTS" \ + training.training_num_nodes=1 \ + training.training_num_gpus_per_node="$TRAIN_GPUS" \ + training.fsdp_strategy=FULL_SHARD \ + training.draft_accumulation_steps="$DRAFT_ACCUMULATION_STEPS" \ + training.prefetch_depth="$PREFETCH_DEPTH" \ + inference.inference_num_gpus="$INFERENCE_GPUS" \ + inference.inference_num_gpus_per_engine="$TP_SIZE" \ + inference.inference_num_gpus_per_node="$TOTAL_GPUS" \ + inference.sglang.tp_size="$TP_SIZE" \ + debug.enable_perf_metrics=true \ + output_dir="$OUTPUT_DIR" \ + cache_dir="$ROOT_DIR/cache/$RUN_NAME" \ + "$@" + +echo "==============================================" +echo "Training completed. Checkpoints: $OUTPUT_DIR/checkpoints" +echo "==============================================" diff --git a/tests/test_domino.py b/tests/test_domino.py new file mode 100644 index 00000000..7e80c13c --- /dev/null +++ b/tests/test_domino.py @@ -0,0 +1,207 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Licensed under the MIT License (see repository LICENSE / file headers). + +"""Correctness tests for the Domino draft model (CPU/float32). + +Verifies the new logic - the Domino causal-correction head and the base-anchored +curriculum loss - in isolation from flex attention/CUDA. The reused DFlash +backbone/anchor/mask code is covered by test_dflash.py. +""" + +import pytest + +torch = pytest.importorskip("torch") + +from torchspec.models.domino import DominoModel # noqa: E402 +from torchspec.models.draft.auto import AutoDraftModelConfig # noqa: E402 +from torchspec.models.draft.domino import DominoConfig, DominoDraftModel # noqa: E402 +from torchspec.training.domino_trainer import DominoTrainer # noqa: E402 + +DEV = torch.device("cpu") +DT = torch.float32 +B, SEQ, BLOCK, NANCH = 2, 24, 4, 4 + + +def _cfg(): + return DominoConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=32, + num_target_layers=2, + target_hidden_size=32, + target_num_hidden_layers=4, + mask_token_id=31, + gru_hidden_size=16, + correction_rank=8, + ) + + +def _build(cfg, *, loss_objective="decay", loss_decay_gamma=7.0): + draft = DominoDraftModel(cfg).to(DEV, DT) + return DominoModel( + draft_model=draft, + block_size=BLOCK, + num_anchors=NANCH, + loss_objective=loss_objective, + dpace_alpha=0.5, + loss_decay_gamma=loss_decay_gamma, + ).to(DEV, DT) + + +def _batch(cfg): + torch.manual_seed(123) + return ( + torch.randint(0, cfg.vocab_size, (B, SEQ), device=DEV), + [ + torch.randn(B, SEQ, cfg.target_hidden_size, device=DEV, dtype=DT) + for _ in range(cfg.num_target_layers) + ], + torch.ones(B, SEQ, device=DEV, dtype=DT), + torch.randn(cfg.vocab_size, cfg.hidden_size, device=DEV, dtype=DT) * 0.02, + ) + + +def _fwd(model, batch, lam): + model.curriculum_lambda = lam + return model(batch[0], batch[1], batch[2], batch[3]) + + +def test_domino_json_resolves_to_domino_config(): + cfg = AutoDraftModelConfig.from_file("torchspec/config/domino_draft_config.json") + assert isinstance(cfg, DominoConfig) + + +def test_head_present_and_trainable(): + model = _build(_cfg()) + for name in ("causal_gru", "correction_w1", "correction_w2"): + module = getattr(model.draft_model, name) + params = list(module.parameters()) + assert params and all(p.requires_grad for p in params) + + +def test_forward_output_contract(): + cfg = _cfg() + model = _build(cfg) + loss, acc, loss_pp, acc_pp, count_pp, aux_metrics = _fwd(model, _batch(cfg), 0.0) + assert loss.ndim == 0 and acc.ndim == 0 + assert loss_pp.shape == (BLOCK,) and acc_pp.shape == (BLOCK,) + assert count_pp.shape == (BLOCK,) + assert torch.isfinite(loss) + assert set(aux_metrics) == { + "base_loss", + "final_loss", + "correction_norm", + "correction_abs_mean", + } + assert all(torch.isfinite(v) for v in aux_metrics.values()) + + +def test_trainer_forward_preserves_loss_components_slot(monkeypatch): + cfg = _cfg() + model = _build(cfg) + trainer = object.__new__(DominoTrainer) + trainer.model = model + trainer.target_lm_head_weight = torch.randn( + cfg.vocab_size, cfg.hidden_size, device=DEV, dtype=DT + ) + trainer.num_target_layers = cfg.num_target_layers + + monkeypatch.setattr("torchspec.training.domino_trainer.torch.device", lambda _: DEV) + + input_ids, hidden_states_list, loss_mask, _ = _batch(cfg) + batch = { + "input_ids": input_ids, + "hidden_states": torch.cat(hidden_states_list, dim=-1), + "loss_mask": loss_mask, + } + output = DominoTrainer._forward(trainer, batch) + + assert len(output) == 6 + assert set(output[-1]) == { + "base_loss", + "final_loss", + "correction_norm", + "correction_abs_mean", + } + + +def test_dpace_objective_changes_domino_loss_weights(): + cfg = _cfg() + batch = _batch(cfg) + torch.manual_seed(7) + decay_model = _build(cfg, loss_objective="decay", loss_decay_gamma=None) + torch.manual_seed(7) + dpace_model = _build(cfg, loss_objective="dpace") + dpace_model.load_state_dict(decay_model.state_dict()) + + torch.manual_seed(11) + decay_loss = _fwd(decay_model, batch, 0.0)[0] + torch.manual_seed(11) + dpace_loss = _fwd(dpace_model, batch, 0.0)[0] + + assert torch.isfinite(dpace_loss) + assert not torch.allclose(decay_loss, dpace_loss) + + +def test_curriculum_lambda_is_noop_when_correction_zeroed(): + cfg = _cfg() + model = _build(cfg) + batch = _batch(cfg) + with torch.no_grad(): + model.draft_model.correction_w2.weight.zero_() + torch.manual_seed(7) + loss_base = _fwd(model, batch, 1.0)[0].item() + torch.manual_seed(7) + loss_final = _fwd(model, batch, 0.0)[0].item() + assert abs(loss_base - loss_final) < 1e-5 + + +def test_curriculum_selects_base_vs_final_when_correction_active(): + cfg = _cfg() + model = _build(cfg) + batch = _batch(cfg) + torch.manual_seed(7) + base = _fwd(model, batch, 1.0)[0].item() + torch.manual_seed(7) + final = _fwd(model, batch, 0.0)[0].item() + assert abs(base - final) > 1e-4 + + +def test_gradients_flow_to_domino_head(): + cfg = _cfg() + model = _build(cfg) + model.zero_grad(set_to_none=True) + _fwd(model, _batch(cfg), 0.0)[0].backward() + for name in ("causal_gru", "correction_w1", "correction_w2"): + module = getattr(model.draft_model, name) + grad_norm = sum(p.grad.norm().item() for p in module.parameters() if p.grad is not None) + assert grad_norm > 0, name + + +def test_model_learns_under_curriculum(): + torch.manual_seed(0) + cfg = _cfg() + model = _build(cfg) + batch = _batch(cfg) + optimizer = torch.optim.AdamW( + [p for p in model.parameters() if p.requires_grad], + lr=3e-3, + ) + steps = 200 + first_loss = last_loss = first_acc = last_acc = None + for step in range(steps): + lam = max(0.0, 1.0 - step / (steps * 0.5)) + optimizer.zero_grad(set_to_none=True) + loss, acc, *_ = _fwd(model, batch, lam) + loss.backward() + optimizer.step() + if step == 0: + first_loss, first_acc = loss.item(), acc.item() + last_loss, last_acc = loss.item(), acc.item() + + assert last_loss < first_loss * 0.5 + assert last_acc > first_acc + 0.2 diff --git a/torchspec/config/domino_draft_config.json b/torchspec/config/domino_draft_config.json new file mode 100644 index 00000000..19e3cbd8 --- /dev/null +++ b/torchspec/config/domino_draft_config.json @@ -0,0 +1,21 @@ +{ + "architectures": ["DominoDraftModel"], + "model_type": "domino", + "hidden_size": 4096, + "intermediate_size": 12288, + "num_hidden_layers": 5, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "vocab_size": 151936, + "rms_norm_eps": 1e-6, + "max_position_embeddings": 40960, + "rope_theta": 1000000.0, + "num_target_layers": 5, + "target_hidden_size": 4096, + "target_num_hidden_layers": 36, + "target_layer_ids": [1, 9, 17, 25, 33], + "mask_token_id": 151669, + "gru_hidden_size": 1024, + "correction_rank": 256, + "tie_word_embeddings": false +} diff --git a/torchspec/config/train_config.py b/torchspec/config/train_config.py index ba50bfbc..73758ddb 100644 --- a/torchspec/config/train_config.py +++ b/torchspec/config/train_config.py @@ -164,6 +164,12 @@ class TrainingConfig: dspark_l1_loss_alpha: float = 0.9 dspark_confidence_head_alpha: float = 1.0 + # Domino-specific parameters (causal-correction head + base-anchored curriculum) + domino_gru_hidden_size: int = 1024 + domino_correction_rank: int = 256 + # Steps over which the curriculum lambda anneals 1 -> 0 (None = lr_total_steps) + domino_curriculum_steps: Optional[int] = None + @dataclass class DecodeConfig: diff --git a/torchspec/models/domino.py b/torchspec/models/domino.py new file mode 100644 index 00000000..4ce7f780 --- /dev/null +++ b/torchspec/models/domino.py @@ -0,0 +1,231 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Domino training wrapper. + +Extends the DFlash training wrapper with the Domino causal-correction head and the +base-anchored training curriculum (arXiv:2605.29707, Sec. 4): + + L = (1 - lambda) * L_final + lambda * L_base + +where L_base is the cross entropy on the parallel backbone's base logits and +L_final is the cross entropy on (base + causal correction). Lambda is annealed from +1 -> 0 over training so the backbone is strengthened first and the Domino head +gradually takes over the residual correction. The causal encoder is teacher +forced on ground-truth token embeddings (Sec. 4.1). Both losses reuse DFlash's +configurable position weighting. + +Reuses DFlash's anchor sampling, block-causal mask, and label/weight construction +verbatim, so Domino's data contract is identical to DFlash. +""" + +from typing import Dict, List, Tuple + +import torch +import torch.nn.functional as F + +from torchspec.models.dflash import DFlashModel, _create_dflash_mask_mod, _dpace_position_weights +from torchspec.models.ops.flex_attention import compile_friendly_create_block_mask + + +class DominoModel(DFlashModel): + """DFlash training wrapper + Domino causal correction + curriculum loss. + + ``curriculum_lambda`` is set by the trainer before each step (1 -> 0 schedule). + At lambda=1 the loss is pure base (backbone); at lambda=0 pure final (with correction). + """ + + def __init__(self, *args, curriculum_lambda: float = 0.0, **kwargs): + super().__init__(*args, **kwargs) + # Mutable attribute: the trainer updates it each step (avoids threading kwargs + # through the FSDP-wrapped forward). + self.curriculum_lambda = curriculum_lambda + + def _weighted_ce( + self, + logits: torch.Tensor, + target_ids: torch.Tensor, + objective_weights: torch.Tensor, + ) -> torch.Tensor: + """Objective-weighted mean cross entropy (fp32 for stability).""" + flat_logits = logits.reshape(-1, logits.size(-1)).float() + flat_targets = target_ids.reshape(-1) + loss_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none") + w = objective_weights.reshape(-1) + return (loss_per_token * w).sum() / w.sum().clamp(min=1e-6) + + def forward( + self, + input_ids: torch.Tensor, + hidden_states_list: List[torch.Tensor], + loss_mask: torch.Tensor, + lm_head_weight: torch.Tensor, + ) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + Dict[str, torch.Tensor], + ]: + bsz, seq_len = input_ids.shape + device = input_ids.device + + # Steps 1-6: identical to DFlash (reuse parent helpers). + context_feature = self.draft_model.extract_context_feature(hidden_states_list) + anchor_positions, block_keep_mask = self._sample_anchor_positions( + seq_len, loss_mask, device + ) + n_blocks = anchor_positions.shape[1] + noise_embedding = self._create_noise_embed(input_ids, anchor_positions, block_keep_mask) + context_position_ids, draft_position_ids = self._create_position_ids( + anchor_positions, seq_len + ) + + draft_len = n_blocks * self.block_size + kv_len = seq_len + draft_len + block_mask = None + if device.type == "cuda": + mask_mod = _create_dflash_mask_mod( + anchor_positions=anchor_positions, + block_keep_mask=block_keep_mask, + ctx_len=seq_len, + block_size=self.block_size, + ) + block_mask = compile_friendly_create_block_mask( + mask_mod=mask_mod, + B=bsz, + H=None, + Q_LEN=draft_len, + KV_LEN=kv_len, + device=device, + ) + + draft_hidden = self.draft_model( + draft_input_ids=None, + context_feature=context_feature, + draft_position_ids=draft_position_ids, + context_position_ids=context_position_ids, + block_mask=block_mask, + noise_embedding=noise_embedding, + ) + + # Base logits via the frozen target LM head (the DFlash backbone). + base_logits = F.linear(draft_hidden, lm_head_weight) + + # Labels and binary weight mask are identical to DFlash. + label_offsets = torch.arange(0, self.block_size, device=device).view(1, 1, -1) + label_indices = anchor_positions.unsqueeze(-1) + label_offsets + valid_label_mask = label_indices < seq_len + safe_label_indices = label_indices.clamp(max=seq_len - 1) + target_ids = torch.gather( + input_ids.unsqueeze(1).expand(-1, n_blocks, -1), + 2, + safe_label_indices, + ) + + weight_mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, self.block_size).float() + weight_mask = weight_mask * valid_label_mask.float() + pos_in_block = torch.arange(self.block_size, device=device).view(1, 1, -1) + weight_mask = weight_mask * (pos_in_block > 0).float() + original_loss_mask_gathered = torch.gather( + loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), + 2, + safe_label_indices, + ) + weight_mask = weight_mask * original_loss_mask_gathered + binary_eval_mask = weight_mask.view(bsz, n_blocks, self.block_size) + + # Domino causal correction (teacher-forced on ground-truth tokens). + hidden_dim = draft_hidden.shape[-1] + h_blocked = draft_hidden.reshape(bsz * n_blocks, self.block_size, hidden_dim) + gt_token_embeds = self.draft_model.embed_tokens(target_ids).reshape( + bsz * n_blocks, self.block_size, hidden_dim + ) + delta = self.draft_model.domino_correction(h_blocked, gt_token_embeds) + delta = delta.reshape(bsz, n_blocks, self.block_size, -1).reshape(bsz, draft_len, -1) + final_logits = base_logits.float() + delta.float() + + # Objective weighting mirrors DFlash. For D-PACE, final logits drive the + # detached dynamic weights because they are the deployed Domino outputs. + objective_weights = weight_mask + if ( + self.loss_objective == "decay" + and self.loss_decay_gamma is not None + and self.loss_decay_gamma > 0 + ): + k = torch.arange(self.block_size, device=device).view(1, 1, -1) + decay_weights = torch.exp(-(k - 1).clamp(min=0).float() / self.loss_decay_gamma) + objective_weights = weight_mask * decay_weights + elif self.loss_objective == "dpace": + dpace_weights = torch.ones_like(weight_mask) + if self.block_size > 1: + with torch.no_grad(): + final_loss_per_position = F.cross_entropy( + final_logits.reshape(-1, final_logits.size(-1)).float(), + target_ids.reshape(-1), + reduction="none", + ).view(bsz, n_blocks, self.block_size) + target_confidences = torch.exp(-final_loss_per_position[..., 1:].float()) + dpace_pred_weights = _dpace_position_weights( + target_confidences, + self.dpace_alpha, + ).to(dtype=weight_mask.dtype) + dpace_weights[..., 1:] = dpace_pred_weights + objective_weights = weight_mask * dpace_weights + + # Base-anchored curriculum loss. + loss_base = self._weighted_ce(base_logits, target_ids, objective_weights) + loss_final = self._weighted_ce(final_logits, target_ids, objective_weights) + lam = float(self.curriculum_lambda) + loss = (1.0 - lam) * loss_final + lam * loss_base + + with torch.no_grad(): + correction_norm = delta.float().pow(2).mean().sqrt() + correction_abs_mean = delta.float().abs().mean() + + # Metrics from the final logits (what gets deployed). + with torch.no_grad(): + flat_final = final_logits.reshape(-1, final_logits.size(-1)) + flat_targets = target_ids.reshape(-1) + loss_per_token = F.cross_entropy(flat_final.float(), flat_targets, reduction="none") + bem = binary_eval_mask.reshape(-1) + pred_ids = torch.argmax(flat_final, dim=-1) + correct = (pred_ids == flat_targets) & (bem > 0.5) + accuracy = correct.sum().float() / bem.sum().clamp(min=1e-6) + + bw = bem.view(bsz, n_blocks, self.block_size) + count_per_position = bw.sum(dim=(0, 1)) + count_pp = count_per_position.clamp(min=1.0) + loss_per_position = (loss_per_token.view(bsz, n_blocks, self.block_size) * bw).sum( + dim=(0, 1) + ) / count_pp + acc_per_position = ( + correct.view(bsz, n_blocks, self.block_size).float().sum(dim=(0, 1)) / count_pp + ) + + aux_metrics = { + "base_loss": loss_base.detach(), + "final_loss": loss_final.detach(), + "correction_norm": correction_norm.detach(), + "correction_abs_mean": correction_abs_mean.detach(), + } + + return loss, accuracy, loss_per_position, acc_per_position, count_per_position, aux_metrics diff --git a/torchspec/models/draft/auto.py b/torchspec/models/draft/auto.py index 550b6839..8aef06ca 100644 --- a/torchspec/models/draft/auto.py +++ b/torchspec/models/draft/auto.py @@ -28,6 +28,7 @@ from torchspec.models.draft.deepseek_eagle import Eagle3DeepseekV2ForCausalLM from torchspec.models.draft.dflash import DFlashConfig, DFlashDraftModel +from torchspec.models.draft.domino import DominoConfig, DominoDraftModel from torchspec.models.draft.dspark import DSparkConfig, DSparkDraftModel from torchspec.models.draft.llama3_eagle import LlamaForCausalLMEagle3 from torchspec.utils.logging import logger @@ -38,6 +39,7 @@ class AutoEagle3DraftModel(AutoModelForCausalLMBase): LlamaConfig: LlamaForCausalLMEagle3, DeepseekV3Config: Eagle3DeepseekV2ForCausalLM, DFlashConfig: DFlashDraftModel, + DominoConfig: DominoDraftModel, DSparkConfig: DSparkDraftModel, } @@ -79,6 +81,7 @@ class AutoDraftModelConfig: "LlamaForCausalLMEagle3": LlamaConfig, "Eagle3DeepseekV2ForCausalLM": DeepseekV3Config, "DFlashDraftModel": DFlashConfig, + "DominoDraftModel": DominoConfig, "Qwen3DSparkModel": DSparkConfig, } diff --git a/torchspec/models/draft/domino.py b/torchspec/models/draft/domino.py new file mode 100644 index 00000000..60350783 --- /dev/null +++ b/torchspec/models/draft/domino.py @@ -0,0 +1,119 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Domino draft model: DFlash parallel backbone + a lightweight causal-correction head. + +Domino (arXiv:2605.29707) decouples autoregressive modelling from parallel +draft execution. It keeps the cheap block-parallel DFlash backbone (which produces +base logits via the frozen target LM head) and adds a small "Domino head" that +injects token-dependence: + + logits_i = base_logits_i + Delta_i, Delta_i = g(z_i, s_{i-1}) + +where s_{i-1} = GRU(E(t_{ r -> vocab (logit-space residual). + self.correction_w1 = nn.Linear( + config.hidden_size + self.gru_hidden_size, + self.correction_rank, + bias=False, + ) + self.correction_w2 = nn.Linear( + self.correction_rank, + config.vocab_size, + bias=False, + ) + + def domino_correction( + self, + draft_hidden_blocked: torch.Tensor, + gt_token_embeds: torch.Tensor, + ) -> torch.Tensor: + """Compute the Domino logit-space residual correction for one or more blocks. + + Args: + draft_hidden_blocked: [N, block, D] - backbone hidden states + gt_token_embeds: [N, block, D] - embeddings of the ground-truth tokens + at block positions (teacher forcing) + + Returns: + delta_logits: [N, block, vocab]. Delta at position 0 uses s_{-1}=0. + """ + # CUDA RNNs are happiest in fp32; run the (small) GRU there, then cast back. + gru_dtype = self.causal_gru.weight_ih_l0.dtype + states, _ = self.causal_gru(gt_token_embeds.to(gru_dtype)) + + # Delta_i = g(z_i, s_{i-1}); position 0 sees no tokens before it. + s_prev = torch.cat([torch.zeros_like(states[:, :1]), states[:, :-1]], dim=1) + s_prev = s_prev.to(draft_hidden_blocked.dtype) + x = torch.cat([draft_hidden_blocked, s_prev], dim=-1) + return self.correction_w2(F.silu(self.correction_w1(x))) diff --git a/torchspec/training/dflash_trainer.py b/torchspec/training/dflash_trainer.py index 2fcebe09..d1e289ad 100644 --- a/torchspec/training/dflash_trainer.py +++ b/torchspec/training/dflash_trainer.py @@ -149,6 +149,7 @@ def init_model( f"{frozen_count:,} frozen (embedding) parameters" ) + dflash_model = self._build_training_wrapper(draft_model) dflash_model = self._build_training_wrapper(draft_model) full_state = dflash_model.state_dict() if dist.get_rank() == 0 else {} diff --git a/torchspec/training/domino_trainer.py b/torchspec/training/domino_trainer.py new file mode 100644 index 00000000..2507a860 --- /dev/null +++ b/torchspec/training/domino_trainer.py @@ -0,0 +1,135 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Domino trainer - DFlash trainer + Domino head + base-anchored curriculum. + +Reuses DFlashTrainer's hook-based init/model path and adds the curriculum: +lambda is linearly annealed 1 -> 0 and pushed onto the DominoModel before each step. +""" + +from argparse import Namespace + +import torch +import torch.distributed as dist + +from torchspec.models.domino import DominoModel +from torchspec.models.draft.domino import DominoConfig, DominoDraftModel +from torchspec.training.dflash_trainer import DFlashTrainer + + +class DominoTrainer(DFlashTrainer): + """Domino-specific trainer (extends DFlashTrainer with the curriculum).""" + + _draft_config_class = DominoConfig + + def __init__(self, args: Namespace): + super().__init__(args) + # Steps over which lambda: 1 -> 0. Defaults to the LR schedule length. + self.curriculum_steps = getattr(args, "domino_curriculum_steps", None) or getattr( + args, "lr_total_steps", None + ) + + def _build_draft_model(self, config): + return DominoDraftModel(config) + + def _build_training_wrapper(self, draft_model): + return DominoModel( + draft_model=draft_model, + block_size=self.block_size, + num_anchors=self.num_anchors, + loss_objective=self.loss_objective, + dpace_alpha=self.dpace_alpha, + loss_decay_gamma=self.loss_decay_gamma, + ce_loss_alpha=self.ce_loss_alpha, + l1_loss_alpha=self.l1_loss_alpha, + ) + + def _compute_curriculum_lambda(self, step: int) -> float: + """Linear anneal lambda from 1 (pure base) to 0 (pure final) over curriculum_steps.""" + total = self.curriculum_steps + if not total or total <= 0: + return 0.0 + return max(0.0, 1.0 - float(step) / float(total)) + + def _forward(self, batch: dict): + device = torch.device("cuda") + input_ids = batch["input_ids"].to(device, non_blocking=True) + hidden_states = batch["hidden_states"].to(device, non_blocking=True) + + loss_mask = batch["loss_mask"] + if loss_mask.dim() == 3: + loss_mask = loss_mask.squeeze(-1) + loss_mask = loss_mask.to(device, non_blocking=True) + + hidden_states_list = self._split_hidden_states(hidden_states) + del hidden_states + + output = self.model( + input_ids=input_ids, + hidden_states_list=hidden_states_list, + loss_mask=loss_mask, + lm_head_weight=self.target_lm_head_weight, + ) + loss, accuracy, loss_pp, acc_pp, count_pp, aux_metrics = output + self._last_domino_aux_metrics = aux_metrics + return loss, accuracy, loss_pp, acc_pp, count_pp, aux_metrics + + def _train_step( + self, + batch: dict, + accumulation_steps: int, + step: int, + batch_idx: int, + num_batches: int, + ) -> dict: + # Push the current curriculum weight onto the unwrapped DominoModel so + # forward() picks it up - avoids threading kwargs through the FSDP wrapper. + lam = self._compute_curriculum_lambda(step) + self.dflash.curriculum_lambda = lam + metrics = super()._train_step(batch, accumulation_steps, step, batch_idx, num_batches) + metrics["train/curriculum_lambda"] = lam + for name, value in getattr(self, "_last_domino_aux_metrics", {}).items(): + metrics[f"domino/{name}"] = value.detach() + return metrics + + def _aggregate_metrics( + self, all_step_metrics: list[dict], step: int, *, grad_norm: torch.Tensor = None + ) -> dict: + metrics = super()._aggregate_metrics(all_step_metrics, step, grad_norm=grad_norm) + + aux_keys = [ + "domino/base_loss", + "domino/final_loss", + "domino/correction_norm", + "domino/correction_abs_mean", + ] + for key in aux_keys: + values = [m[key].float() for m in all_step_metrics if key in m] + if not values: + continue + value = torch.stack(values).mean() + dist.all_reduce(value, op=dist.ReduceOp.SUM) + value = value / dist.get_world_size() + metrics[f"train/{key.removeprefix('domino/')}"] = value.item() + + if all_step_metrics and "train/curriculum_lambda" in all_step_metrics[-1]: + metrics["train/curriculum_lambda"] = all_step_metrics[-1]["train/curriculum_lambda"] + + return metrics diff --git a/torchspec/training/trainer_actor.py b/torchspec/training/trainer_actor.py index a9ce6457..d5485442 100644 --- a/torchspec/training/trainer_actor.py +++ b/torchspec/training/trainer_actor.py @@ -26,6 +26,7 @@ from torchspec import AutoDraftModelConfig from torchspec.models.draft.dflash import DFlashConfig +from torchspec.models.draft.domino import DominoConfig from torchspec.models.draft.dspark import DSparkConfig from torchspec.ray.ray_actor import RayActor from torchspec.training.eagle3_trainer import Eagle3Trainer @@ -77,14 +78,19 @@ def init(self, args: Namespace, role: str, mooncake_config=None, with_ref: bool draft_model_config = AutoDraftModelConfig.from_file(args.draft_model_config) # Config-based trainer dispatch. - # DSparkConfig subclasses DFlashConfig, so it must be checked first. - # - DSparkConfig → DSparkTrainer - # - DFlashConfig → DFlashTrainer + # DSparkConfig and DominoConfig subclass DFlashConfig, so they must be checked first. + # - DSparkConfig -> DSparkTrainer + # - DominoConfig -> DominoTrainer + # - DFlashConfig -> DFlashTrainer # - else Eagle3. if isinstance(draft_model_config, DSparkConfig): from torchspec.training.dspark_trainer import DSparkTrainer self._trainer = DSparkTrainer(args) + elif isinstance(draft_model_config, DominoConfig): + from torchspec.training.domino_trainer import DominoTrainer + + self._trainer = DominoTrainer(args) elif isinstance(draft_model_config, DFlashConfig): from torchspec.training.dflash_trainer import DFlashTrainer