feat: add Domino drafter training backend

ci/domino_gpu_smoke.py

@@ -0,0 +1,147 @@
+"""Hardware smoke test for the Domino drafter training backend.
+
+Drives the real training path on GPU with a real target model (default
+Qwen3-4B): it runs the frozen target forward to collect the DFlash-style
+multi-layer context hidden states, builds the Domino draft via
+``DominoTrainerBackend.build_model``, and runs several optimizer steps through
+``compute_loss`` (which invokes the block-drafter forward with the causal GRU
+correction head and the dual-logit base-anchor curriculum).
+
+The draft is cold-started, so the useful signals are:
+  * ``loss`` / ``final_loss`` (Domino-refined logits CE) trending down,
+  * ``base_loss`` (backbone-only logits CE) trending down,
+  * ``accuracy`` (final) and ``base_accuracy`` rising,
+  * ``lambda_base`` decaying from 1 -> 0 (curriculum handing over to the head).
+
+Run:
+    python ci/domino_gpu_smoke.py --target /path/to/target-model --steps 120
+"""
+
+from __future__ import annotations
+
+import argparse
+
+import torch
+from omegaconf import OmegaConf
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+PROMPTS = [
+    "Explain why the sky appears blue during the day, in a few sentences.",
+    "Write a short Python function that returns the nth Fibonacci number and explain it.",
+    "Summarize the water cycle and its main stages in a short paragraph.",
+    "Describe the main differences between TCP and UDP for a networking student.",
+]
+
+
+def _build_batch(target, tokenizer, target_layer_ids, device):
+    """One packed batch: input_ids, loss_mask, and concatenated context hidden states."""
+    id_chunks, mask_chunks, hidden_chunks = [], [], []
+    for text in PROMPTS:
+        messages = [{"role": "user", "content": text}]
+        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        enc = tokenizer(prompt, return_tensors="pt").to(device)
+        with torch.no_grad():
+            out = target(input_ids=enc["input_ids"], output_hidden_states=True)
+        # hidden_states is a tuple of (num_layers + 1) tensors; pick the context layers.
+        layers = [out.hidden_states[i][0] for i in target_layer_ids]  # each [S, H]
+        id_chunks.append(enc["input_ids"][0])
+        mask_chunks.append(torch.ones(enc["input_ids"].size(1), device=device))
+        hidden_chunks.append(torch.cat(layers, dim=-1))  # [S, num_ctx*H]
+
+    input_ids = torch.cat(id_chunks).unsqueeze(0)
+    loss_mask = torch.cat(mask_chunks).unsqueeze(0)
+    hidden = torch.cat(hidden_chunks).unsqueeze(0).to(torch.bfloat16)
+    return {"input_ids": input_ids, "loss_mask": loss_mask, "hidden_states": hidden, "attention_mask": torch.ones_like(input_ids)}
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--target", required=True, help="path or HF id of the target causal LM")
+    parser.add_argument("--steps", type=int, default=120)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--num-context-layers", type=int, default=5)
+    parser.add_argument("--lambda-decay-steps", type=int, default=60)
+    args = parser.parse_args()
+
+    device = "cuda"
+    torch.manual_seed(0)
+
+    print(f"[smoke] loading target {args.target}")
+    tokenizer = AutoTokenizer.from_pretrained(args.target)
+    target = AutoModelForCausalLM.from_pretrained(args.target, torch_dtype=torch.bfloat16).to(device).eval()
+    target_cfg = AutoConfig.from_pretrained(args.target)
+
+    from verl_speco.models.dflash import build_target_layer_ids
+
+    target_layers = int(getattr(target_cfg, "num_hidden_layers"))
+    target_layer_ids = build_target_layer_ids(args.num_context_layers, target_layers)
+    print(f"[smoke] context layers={target_layer_ids} (of {target_layers})")
+
+    batch = _build_batch(target, tokenizer, target_layer_ids, device)
+    print(f"[smoke] batch seq_len={batch['input_ids'].size(1)} hidden={batch['hidden_states'].size(-1)}")
+
+    cfg = OmegaConf.create(
+        {
+            "rollout": {
+                "drafter": {
+                    "speculative_algorithm": "DOMINO",
+                    "model_path": "/dev/null/does-not-exist",
+                    "training": {
+                        "domino_block_size": 8,
+                        "domino_num_anchors": 128,
+                        "domino_num_target_layers": args.num_context_layers,
+                        "domino_num_hidden_layers": 1,
+                        "domino_lambda_base_decay_steps": args.lambda_decay_steps,
+                        "lr": args.lr,
+                    },
+                }
+            },
+            "model": {"path": args.target},
+        }
+    )
+
+    from verl_speco.backends.domino_trainer_backend import DominoTrainerBackend
+
+    backend = DominoTrainerBackend(cfg, target_cfg)
+    model, drafter_cfg = backend.build_model()
+    model = model.to(device).to(torch.bfloat16).train()
+    backend.target_lm_head = backend.target_lm_head.to(device).to(torch.bfloat16)
+    optimizer = backend.setup_optimizer(model, cfg.rollout.drafter.training)
+    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(
+        f"[smoke] block_size={model.block_size} gru={model.draft_model.gru_hidden_dim} "
+        f"emb_dim={model.draft_model.emb_dim} trainable_params={n_params:,}"
+    )
+
+    first = None
+    for step in range(args.steps):
+        out = backend.compute_loss(model, batch, 0)
+        num_tokens = out["local_num_tokens"].clamp_min(1)
+        loss = out["total_local_ploss"] / num_tokens
+        optimizer.zero_grad()
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+
+        if step % 10 == 0 or step == args.steps - 1:
+            d = out["diagnostics"]
+            fl = float(d["domino_final_loss"])
+            bl = float(d["domino_base_loss"])
+            acc = float(out["accuracy"])
+            bacc = float(d["domino_base_accuracy"])
+            lam = float(d["domino_lambda_base"])
+            if first is None:
+                first = (fl, bl, acc)
+            print(
+                f"[smoke] step {step:3d}  final_loss={fl:.4f}  base_loss={bl:.4f}  "
+                f"final_acc={acc:.4f}  base_acc={bacc:.4f}  lambda_base={lam:.3f}"
+            )
+
+    print(
+        f"[smoke] DONE  final_loss {first[0]:.4f}->{fl:.4f}  base_loss {first[1]:.4f}->{bl:.4f}  "
+        f"final_acc {first[2]:.4f}->{acc:.4f}"
+    )
+
+
+if __name__ == "__main__":
+    main()

tests/integration/test_domino_backend_contract.py

@@ -0,0 +1,173 @@
+"""Contract tests for the Domino drafter backend.
+
+CPU-light: they exercise the Domino projector modules, the lambda-base
+curriculum, the algorithm routing (config, oldlogprob aux layers, vLLM
+guardrail), and the block-drafter classification. The full training forward is
+validated on GPU by ``ci/domino_gpu_smoke.py``.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+
+def _tiny_domino_config():
+    from verl_speco.models.domino import DominoConfig
+
+    return DominoConfig(
+        hidden_size=8,
+        intermediate_size=16,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        num_hidden_layers=1,
+        vocab_size=32,
+        num_target_layers=4,
+        num_context_layers=2,
+        target_hidden_size=8,
+        target_num_hidden_layers=4,
+        target_layer_ids=[1, 3],
+        mask_token_id=31,
+        block_size=4,
+        num_anchors=8,
+        emb_dim=6,
+        gru_hidden_dim=10,
+        pure_draft_prefix_len=1,
+        rms_norm_eps=1e-6,
+        max_position_embeddings=64,
+    )
+
+
+def test_domino_model_builds_projector_head() -> None:
+    pytest.importorskip("torch")
+    pytest.importorskip("transformers")
+    from verl_speco.models.domino import DominoDraftModel
+
+    config = _tiny_domino_config()
+    model = DominoDraftModel(config)
+
+    assert model.projector_type == "domino"
+    # GRU consumes token embeddings (hidden_size) -> gru_hidden_dim.
+    assert model.prefix_gru.input_size == config.hidden_size
+    assert model.prefix_gru.hidden_size == config.gru_hidden_dim
+    # embed_proj: [hidden + gru_hidden] -> emb_dim -> vocab.
+    assert model.embed_proj[0].in_features == config.hidden_size + config.gru_hidden_dim
+    assert model.embed_proj[0].out_features == config.emb_dim
+    assert model.embed_proj[-1].out_features == config.vocab_size
+    # It still carries the DFlash backbone token embedding (used by the GRU).
+    assert model.embed_tokens.num_embeddings == config.vocab_size
+
+
+def test_domino_forward_computes_top5_accuracy() -> None:
+    """top5_correct_count must actually be reduced, not left at its zero init.
+
+    DFlash and DSpark both compute it; a Domino regression here silently reports
+    top5_acc=0 forever (base_trainer derives top5_acc from this counter).
+    """
+    pytest.importorskip("torch")
+    pytest.importorskip("transformers")
+    import torch
+
+    from verl_speco.backends.domino_trainer_backend import DominoTrainingModel
+    from verl_speco.models.domino import DominoDraftModel
+
+    torch.manual_seed(0)
+    config = _tiny_domino_config()
+    model = DominoTrainingModel(
+        draft_model=DominoDraftModel(config),
+        block_size=config.block_size,
+        num_anchors=config.num_anchors,
+        pure_draft_prefix_len=config.pure_draft_prefix_len,
+    )
+
+    bsz, seq_len = 2, 16
+    input_ids = torch.randint(0, config.vocab_size, (bsz, seq_len))
+    hidden_states_list = [torch.randn(bsz, seq_len, config.target_hidden_size) for _ in config.target_layer_ids]
+    loss_mask = torch.ones(bsz, seq_len, dtype=torch.long)
+    lm_head_weight = torch.randn(config.vocab_size, config.hidden_size)
+
+    _, _, _, _, _, diagnostics = model(input_ids, hidden_states_list, loss_mask, lm_head_weight)
+
+    top1 = float(diagnostics["top1_correct_count"])
+    top5 = float(diagnostics["top5_correct_count"])
+    quality = float(diagnostics["quality_token_count"])
+
+    assert quality > 0
+    # top-5 is a superset of top-1, and with vocab_size=32 sampling 5 candidates
+    # over that many tokens must hit at least one target.
+    assert top5 >= top1
+    assert top5 > 0
+
+
+def test_domino_lambda_base_schedule() -> None:
+    # get_lambda_base is pure-python, but its module (domino_trainer_backend)
+    # subclasses the torch-based DFlash backend at import time, so it cannot be
+    # imported without torch. Skip under the torch-free CPU CI like the siblings.
+    pytest.importorskip("torch")
+    pytest.importorskip("transformers")
+    from verl_speco.backends.domino_trainer_backend import get_lambda_base
+
+    assert get_lambda_base(0, decay_steps=100, lambda_start=1.0) == pytest.approx(1.0)
+    assert get_lambda_base(50, decay_steps=100, lambda_start=1.0) == pytest.approx(0.5)
+    assert get_lambda_base(100, decay_steps=100, lambda_start=1.0) == pytest.approx(0.0)
+    assert get_lambda_base(200, decay_steps=100, lambda_start=1.0) == pytest.approx(0.0)
+    assert get_lambda_base(25, decay_steps=100, lambda_start=0.4) == pytest.approx(0.3)
+
+
+def test_domino_backend_is_block_drafter_metadata() -> None:
+    pytest.importorskip("torch")
+    pytest.importorskip("transformers")
+    from omegaconf import OmegaConf
+
+    from verl_speco.backends.domino_trainer_backend import DominoTrainerBackend
+
+    backend = DominoTrainerBackend(
+        OmegaConf.create({"rollout": {"drafter": {"training": {}}}, "model": {"path": "/tmp/none"}}),
+        OmegaConf.create({}),
+    )
+    assert backend.model_type == "domino"
+
+
+def test_domino_config_from_file_routes_to_domino(tmp_path) -> None:
+    pytest.importorskip("transformers")
+    import json
+
+    from verl_speco.models.auto import AutoDraftModelConfig
+    from verl_speco.models.domino import DominoConfig
+
+    config = _tiny_domino_config().to_dict()
+    config["architectures"] = ["DominoDraftModel"]
+    (tmp_path / "config.json").write_text(json.dumps(config), encoding="utf-8")
+
+    loaded = AutoDraftModelConfig.from_file(str(tmp_path / "config.json"))
+    assert isinstance(loaded, DominoConfig)
+    assert loaded.architectures == ["DominoDraftModel"]
+    assert loaded.projector_type == "domino"
+
+
+def test_domino_uses_dflash_aux_layers() -> None:
+    from verl_speco.integration.oldlogprob_layer_ids import resolve_oldlogprob_aux_layer_ids
+
+    layer_ids = resolve_oldlogprob_aux_layer_ids(
+        {"speculative_algorithm": "DOMINO", "training": {"domino_num_target_layers": 5}},
+        target_num_hidden_layers=36,
+    )
+    # Routes down the DFlash multi-context-layer branch (not the EAGLE default triple).
+    assert layer_ids is not None
+    assert len(layer_ids) == 5
+
+
+def test_domino_rejected_by_vllm_config_builder() -> None:
+    from verl_speco.integration.vllm_runtime import _speculative_method_from_drafter
+
+    with pytest.raises(ValueError, match="projector sub-mode"):
+        _speculative_method_from_drafter({"speculative_algorithm": "DOMINO"})
+
+
+def test_domino_rejected_by_sglang_config_builder() -> None:
+    from verl_speco.integration.sglang_runtime import _server_args_overrides_from_drafter
+
+    with pytest.raises(ValueError, match="projector sub-mode"):
+        _server_args_overrides_from_drafter(
+            {"enable": True, "speculative_algorithm": "DOMINO"},
+            supported_fields={"speculative_algorithm"},
+        )