DFlash speculator (polishing, bf16) from DFlash_Stage2/25000

README.md

@@ -0,0 +1,13 @@
+---
+license: other
+tags: [dflash, speculative-decoding, laguna]
+---
+
+# Laguna-S-2.1-DFlash-polishing
+
+DFlash speculator (drafter) for Laguna-S-2.1, bf16.
+
+- Architecture: `DFlashLagunaForCausalLM` (6 sliding-attention layers, block_size 16).
+- Shares token embedding + lm_head with the target; `draft_vocab_size == vocab_size` (no d2t/t2d).
+- Source checkpoint: `s3://poolside.titan.checkpoints.us-east-2/adam/checkpoints/ft_sft_e0630_rhiemann_baseline_titan_sft_training/0008400/DFlash_Stage2/25000`
+- Loads under vLLM (native `laguna_dflash`) and TRT-LLM (pytorch DFlash backend) as the draft model in a speculative config.

config.json

@@ -0,0 +1,61 @@
+{
+  "attention_bias": false,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 3072,
+  "intermediate_size": 12288,
+  "max_position_embeddings": 262144,
+  "model_type": "laguna",
+  "num_attention_heads": 72,
+  "num_hidden_layers": 6,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 500000.0,
+  "sliding_window": 512,
+  "vocab_size": 100352,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention"
+  ],
+  "gating": "per-head",
+  "architectures": [
+    "DFlashLagunaForCausalLM"
+  ],
+  "num_experts": 0,
+  "sliding_windows": [
+    512,
+    512,
+    512,
+    512,
+    512,
+    512
+  ],
+  "draft_vocab_size": 100352,
+  "torch_dtype": "bfloat16",
+  "eagle_aux_hidden_state_layer_ids": [
+    2,
+    11,
+    20,
+    30,
+    39,
+    48
+  ],
+  "dflash_config": {
+    "block_size": 16,
+    "mask_token_id": 12,
+    "num_target_layers": 48,
+    "target_layer_ids": [
+      1,
+      10,
+      19,
+      29,
+      38,
+      47
+    ],
+    "causal": true
+  }
+}

config.py

@@ -0,0 +1,188 @@
+from typing import Any, Literal
+
+from pydantic import Field, field_serializer, field_validator
+from transformers import AutoConfig, PretrainedConfig
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3Config,
+)
+
+from speculators import SpeculatorModelConfig
+
+__all__ = [
+    "DFlashSpeculatorConfig",
+]
+
+
+@SpeculatorModelConfig.register("dflash")
+class DFlashSpeculatorConfig(SpeculatorModelConfig):
+    """
+    Configuration for DFlash speculator with vocabulary mapping.
+
+    DFlash features vocabulary mapping between draft (64K) and target (128K)
+    vocabularies, enabling cross-tokenizer speculation.
+
+    :param transformer_layer_config: Configuration for the transformer decoder layer
+    :param draft_vocab_size: Size of draft model vocabulary for speculation
+    """
+
+    speculators_model_type: Literal["dflash"] = "dflash"
+    architectures: list[str] = Field(
+        default_factory=lambda: ["DFlashSpeculator"],
+        description="Model architectures that can load these weights",
+    )
+
+    transformer_layer_config: PretrainedConfig = Field(
+        default_factory=Qwen3Config,
+        description="Configuration for the transformer decoder layer",
+    )
+
+    draft_vocab_size: int = Field(
+        default=32000,
+        description="Size of draft model vocabulary for speculation",
+    )
+
+    block_size: int = Field(
+        default=8,
+        description=(
+            "Default size of the draft block predicted with a forward pass of the model"
+        ),
+    )
+
+    max_anchors: int = Field(
+        default=256,
+        description=(
+            "Maximum number of anchor positions to sample during training "
+            "(controls memory usage and training efficiency)"
+        ),
+    )
+
+    target_hidden_size: int | None = Field(
+        default=None,
+        description="Hidden size of the target model (if different from draft model)",
+    )
+
+    aux_hidden_state_layer_ids: list[int] | None = Field(
+        default=None,
+        description="Layer IDs of the DFlash auxiliary hidden state layers",
+    )
+
+    decoder_layer_type: Literal["qwen3", "laguna_xs"] = Field(
+        default="qwen3",
+        description="Decoder layer implementation used by the DFlash drafter.",
+    )
+
+    mask_token_id: int | None = Field(
+        default=None,
+        description="Token ID used for masking",
+    )
+
+    sliding_window_non_causal: bool = Field(
+        default=False,
+        description="Use non-causal synthetic block attention for sliding-window layers.",
+    )
+
+    sliding_window_base: Literal["fixed_anchor", "moving_query"] = Field(
+        default="moving_query",
+        description=(
+            "Base-token sliding-window lower-bound policy. 'moving_query' matches "
+            "FlashAttention-style SWA during inference; 'fixed_anchor' preserves "
+            "the legacy DFlash training mask."
+        ),
+    )
+
+    loss_type: Literal["distill", "dflash", "lk", "tv"] = Field(
+        default="distill",
+        description="DFlash objective. 'lk' uses hard-label LK loss.",
+    )
+
+    ce_weight: float | None = Field(
+        default=None,
+        description="Additive weight for hard-label DFlash CE.",
+    )
+
+    tv_weight: float | None = Field(
+        default=None,
+        description="Additive weight for full-distribution TV loss.",
+    )
+
+    kl_weight: float | None = Field(
+        default=None,
+        description="Additive weight for full-distribution KL distillation.",
+    )
+
+    lk_lambda: float = Field(
+        default=0.5,
+        description="Blend coefficient for hard-label LK loss.",
+    )
+
+    tv_temperature: float = Field(
+        default=1.0,
+        description="Teacher softmax temperature for TV/KL terms.",
+    )
+
+    cumacc_weight: bool = Field(
+        default=False,
+        description="Weight hard-label DFlash CE by draft cumulative acceptance.",
+    )
+
+    veri_cum_acc: bool = Field(
+        default=False,
+        description="Weight DFlash loss by verifier cumulative acceptance.",
+    )
+
+    veri_acc_temperature: float = Field(
+        default=1.0,
+        description="Temperature for verifier cumulative acceptance weighting.",
+    )
+
+    static_decay_weight: bool = Field(
+        default=True,
+        description="Apply DFlash position decay to hard-label CE.",
+    )
+
+    kl_distill_weight: float = Field(
+        default=0.0,
+        description="Back-compatible alias for kl_weight when kl_weight is unset.",
+    )
+
+    compile_decoder_layers: bool = Field(
+        default=True,
+        description=(
+            "If True, torch.compile each decoder layer forward during training. "
+            "The DFlash loss remains eager."
+        ),
+    )
+
+    @field_serializer("transformer_layer_config")
+    def serialize_transformer_config(self, value: PretrainedConfig) -> dict:
+        """Serialize transformer config to dict."""
+        return value.to_diff_dict()
+
+    @field_validator("transformer_layer_config", mode="before")
+    @classmethod
+    def validate_transformer_config(cls, value: Any) -> PretrainedConfig:
+        """Validate and convert transformer config."""
+        if isinstance(value, dict):
+            config_class: type[PretrainedConfig] = Qwen3Config
+            if "model_type" in value:
+                config_class = AutoConfig.for_model(
+                    model_type=value["model_type"]
+                ).__class__
+            return config_class(**value)
+        return value
+
+    @property
+    def target_vocab_size(self) -> int:
+        """Get target vocabulary size from transformer config."""
+        return self.transformer_layer_config.vocab_size
+
+    def resolve_loss_weights(self) -> tuple[float, float, float]:
+        if self.loss_type == "tv":
+            ce_default, tv_default = 0.0, 1.0
+        else:
+            ce_default, tv_default = 1.0, 0.0
+
+        ce = ce_default if self.ce_weight is None else self.ce_weight
+        tv = tv_default if self.tv_weight is None else self.tv_weight
+        kl = self.kl_distill_weight if self.kl_weight is None else self.kl_weight
+        return float(ce), float(tv), float(kl)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85c6f3f097358226136f366f4a91c90bd662d544745f1f4944e09668db8fbbe7
+size 2229962896