From 89e6f99bdd13705601254f8b2faa22dc6b708b97 Mon Sep 17 00:00:00 2001 From: Weifan Jiang Date: Mon, 6 Jul 2026 11:21:29 -0400 Subject: [PATCH 1/5] add DPACE loss support for DFlash training Signed-off-by: Weifan Jiang --- scripts/train.py | 13 ++++++ src/speculators/models/dflash/config.py | 14 +++++++ src/speculators/models/dflash/core.py | 10 +++++ src/speculators/models/dflash/metrics.py | 19 ++++++++- src/speculators/models/metrics.py | 50 +++++++++++++++++++++++- 5 files changed, 103 insertions(+), 3 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index e4a103713..154ca1daf 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -961,6 +961,19 @@ def parse_args(): default=4.0, help="Decay gamma for DFlash/DSpark loss weighting (default: 4.0)", ) + # D-Pace specific arguments (loss weight option + smoothing) + parser.add_argument( + "--per-position-loss-weight", + choices=["fixed-exp-decay", "dpace"], + default="fixed-exp-decay", + help="Per-position loss weight option for D-PACE support (default: fixed-exp-decay)", + ) + parser.add_argument( + "--dpace-alpha", + type=float, + default=0.5, + help="Smoothing constant for D-PACE loss (default: 0.5)" + ) # DSpark-specific arguments (sequential Markov head + confidence head). parser.add_argument( "--markov-rank", diff --git a/src/speculators/models/dflash/config.py b/src/speculators/models/dflash/config.py index 0e504e3ab..147fb4669 100644 --- a/src/speculators/models/dflash/config.py +++ b/src/speculators/models/dflash/config.py @@ -70,6 +70,20 @@ class DFlashSpeculatorConfig(SpeculatorModelConfig): "bidirectional.", ) + per_position_loss_weight: Literal["fixed-exp-decay", "dpace"] = Field( + default="fixed-exp-decay", + description="Per-position weighting of the block-drafting loss: " + "fixed-exp-decay uses original D-Flash loss. dpace uses confidence" + "to change per-position weights dynamically." + ) + + dpace_alpha: float = Field( + default=0.5, + ge=0.0, + le=1.0, + description="Smoothing constant for D-PACE loss" + ) + @field_serializer("transformer_layer_config") def serialize_transformer_config(self, value: PretrainedConfig) -> dict: """Serialize transformer config to dict.""" diff --git a/src/speculators/models/dflash/core.py b/src/speculators/models/dflash/core.py index af04c5222..7f59d0275 100644 --- a/src/speculators/models/dflash/core.py +++ b/src/speculators/models/dflash/core.py @@ -105,6 +105,8 @@ def __init__( ) self.verifier_norm.weight.requires_grad = False self.block_size = config.block_size + self.per_position_loss_weight = config.per_position_loss_weight + self.dpace_alpha = config.dpace_alpha self.post_init() @property @@ -181,6 +183,8 @@ def _build_base_config_kwargs( "aux_hidden_state_layer_ids": target_layer_ids, "mask_token_id": kwargs.get("mask_token_id"), "sliding_window_non_causal": kwargs.get("sliding_window_non_causal", False), + "per_position_loss_weight": kwargs.get("per_position_loss_weight", "fixed-exp-decay"), + "dpace_alpha": kwargs.get("dpace_alpha", 0.5), "speculators_config": SpeculatorsConfig( algorithm=algorithm, proposal_methods=[ @@ -207,10 +211,14 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]: loss_config = resolve_loss_config(kwargs["loss_fn"]) gamma = kwargs.get("dflash_decay_gamma", 4.0) max_anchors = kwargs.get("max_anchors", 3072) + per_position_loss_weight = kwargs.get("per_position_loss_weight", "fixed-exp-decay") + dpace_alpha = kwargs.get("dpace_alpha", 0.5) shared = { "loss_config": loss_config, "gamma": gamma, "max_anchors": max_anchors, + "per_position_loss_weight": per_position_loss_weight, + "dpace_alpha": dpace_alpha, } return dict(shared), dict(shared) @@ -412,6 +420,8 @@ def forward( self.block_size, gamma=gamma, loss_config=loss_config, + per_position_loss_weight=self.per_position_loss_weight, + dpace_alpha=self.dpace_alpha, ) draft_tokens = torch.argmax(logits, dim=-1) diff --git a/src/speculators/models/dflash/metrics.py b/src/speculators/models/dflash/metrics.py index d8867a92a..6e3dc8e1c 100644 --- a/src/speculators/models/dflash/metrics.py +++ b/src/speculators/models/dflash/metrics.py @@ -11,6 +11,8 @@ compute_accuracy_multi_step, dflash_loss_decay, kl_div_loss, + ce_loss, + dpace_loss_weight ) _DEFAULT_LOSS_CONFIG: LossConfig = {"kl_div": (kl_div_loss, 1.0)} @@ -23,6 +25,8 @@ def compute_metrics( block_size: int = 1, gamma: float = 4.0, loss_config: LossConfig | None = None, + per_position_loss_weight: str = "fixed-exp-decay", + dpace_alpha: float = 0.5, ) -> tuple[torch.Tensor, dict]: """Compute loss and accuracy metrics for draft model predictions. @@ -33,6 +37,8 @@ def compute_metrics( block_size: Block size for per-position metrics gamma: Temperature for exponential decay in loss weighting loss_config: Mapping of ``{name: (loss_fn, weight)}`` + per_position_loss_weight: Weighting option for per-position block-drafting loss + dpace_alpha: Smoothing constant for D-Pace loss weighting Returns: Tuple of (loss, metrics_dict) where metrics_dict contains: @@ -46,13 +52,24 @@ def compute_metrics( pos_idx = torch.arange(seq_len, device=logits.device) % block_size pos_idx = pos_idx.unsqueeze(0) # shape: [1, T] + elementwise_ce = None + if per_position_loss_weight == "dpace": + elementwise_ce = ce_loss(logits, targets) + decay_mult = dpace_loss_weight( + elementwise_ce, loss_mask, block_size, alpha=dpace_alpha + ) + decay_fn = lambda pos_idx: decay_mult + else: + decay_fn = partial(dflash_loss_decay, gamma=gamma) + loss, term_losses = compound_loss( logits, targets, loss_mask, pos_idx, loss_config=loss_config, - decay_fn=partial(dflash_loss_decay, gamma=gamma), + decay_fn=decay_fn, + dpace_precomputed_ce=elementwise_ce, ) pred_ids = torch.argmax(logits, dim=-1) diff --git a/src/speculators/models/metrics.py b/src/speculators/models/metrics.py index 70f0f7c4a..d4e8e5b42 100644 --- a/src/speculators/models/metrics.py +++ b/src/speculators/models/metrics.py @@ -278,6 +278,46 @@ def exp_loss_decay(pos_idx: torch.Tensor, gamma: float): return gamma**pos_idx +def dpace_loss_weight( + neg_log_q: torch.Tensor, # shape: [1, seq_len] + loss_mask: torch.Tensor, # shape: [1, seq_len] + block_size: int, + alpha: float = 0.5, +): + """ + Per-position block-drafting loss weight based on D-PACE + + Args: + neg_log_q: per-draft-position confidence (negative log-likelihood) + alpha: confidence smoothing constant + """ + with torch.no_grad(): + # convert CE to per-position confidence + q = torch.exp(-neg_log_q).float() + + # reshape loss to [num_anchors, block_size] + # for intra-block cumulative multiplication + num_anchors = q.shape[1] // block_size + q = q.reshape(num_anchors, block_size) + mask = loss_mask.reshape(num_anchors, block_size).to(q.dtype) + + # smoothed confidence for numerical stability + smooth = (1.0 - alpha) * q + alpha + smooth = torch.where(mask > 0, smooth, torch.ones_like(smooth)) + + # prefix cumulative production + prefix = torch.cumprod(smooth, dim=-1) + + # suffix summation: flip -> cumsum -> flip + weight = torch.flip( + torch.cumsum(torch.flip(prefix * mask, dims=[-1]), dim=-1), dims=[-1] + ) + weight = weight * mask + + # reshape weight + return weight.reshape(1, -1) + + _LOSS_FN_MAP: dict[str, Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = { "kl_div": kl_div_loss, "rkl": reverse_kl_div_loss, @@ -361,6 +401,7 @@ def compound_loss( pos_idx: torch.Tensor, loss_config: LossConfig, decay_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, + dpace_precomputed_ce: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """Compute a weighted sum of loss terms. @@ -377,7 +418,8 @@ def compound_loss( multi = len(loss_config) > 1 for name, (fn, weight) in loss_config.items(): term = loss_function( - logits, targets, loss_mask, pos_idx, loss_fn=fn, decay_fn=decay_fn + logits, targets, loss_mask, pos_idx, loss_fn=fn, decay_fn=decay_fn, + dpace_precomputed_ce=dpace_precomputed_ce ) if multi: term_losses[f"{name}_loss"] = term.detach() @@ -392,6 +434,7 @@ def loss_function( pos_idx: torch.Tensor, # shape: [1, seq_len] loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = kl_div_loss, decay_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, + dpace_precomputed_ce: torch.Tensor | None = None, ): """Compute masked, optionally position-decayed training loss. @@ -406,7 +449,10 @@ def loss_function( Returns: Scalar mean loss across the batch. """ - elementwise_loss = loss_fn(logits, targets) # shape: [1, seq_len] + if dpace_precomputed_ce is not None: + elementwise_loss = dpace_precomputed_ce + else: + elementwise_loss = loss_fn(logits, targets) # shape: [1, seq_len] loss_mask = loss_mask.to(elementwise_loss.dtype) elementwise_loss = elementwise_loss * loss_mask From ccf3465a8b780882e95995f077f534c6a9668dfe Mon Sep 17 00:00:00 2001 From: Weifan Jiang Date: Mon, 6 Jul 2026 12:11:03 -0400 Subject: [PATCH 2/5] remove redundant argument path Signed-off-by: Weifan Jiang --- src/speculators/models/dflash/config.py | 14 -------------- src/speculators/models/dflash/core.py | 10 ++++------ 2 files changed, 4 insertions(+), 20 deletions(-) diff --git a/src/speculators/models/dflash/config.py b/src/speculators/models/dflash/config.py index 147fb4669..0e504e3ab 100644 --- a/src/speculators/models/dflash/config.py +++ b/src/speculators/models/dflash/config.py @@ -70,20 +70,6 @@ class DFlashSpeculatorConfig(SpeculatorModelConfig): "bidirectional.", ) - per_position_loss_weight: Literal["fixed-exp-decay", "dpace"] = Field( - default="fixed-exp-decay", - description="Per-position weighting of the block-drafting loss: " - "fixed-exp-decay uses original D-Flash loss. dpace uses confidence" - "to change per-position weights dynamically." - ) - - dpace_alpha: float = Field( - default=0.5, - ge=0.0, - le=1.0, - description="Smoothing constant for D-PACE loss" - ) - @field_serializer("transformer_layer_config") def serialize_transformer_config(self, value: PretrainedConfig) -> dict: """Serialize transformer config to dict.""" diff --git a/src/speculators/models/dflash/core.py b/src/speculators/models/dflash/core.py index 7f59d0275..1fe12c6b8 100644 --- a/src/speculators/models/dflash/core.py +++ b/src/speculators/models/dflash/core.py @@ -105,8 +105,6 @@ def __init__( ) self.verifier_norm.weight.requires_grad = False self.block_size = config.block_size - self.per_position_loss_weight = config.per_position_loss_weight - self.dpace_alpha = config.dpace_alpha self.post_init() @property @@ -183,8 +181,6 @@ def _build_base_config_kwargs( "aux_hidden_state_layer_ids": target_layer_ids, "mask_token_id": kwargs.get("mask_token_id"), "sliding_window_non_causal": kwargs.get("sliding_window_non_causal", False), - "per_position_loss_weight": kwargs.get("per_position_loss_weight", "fixed-exp-decay"), - "dpace_alpha": kwargs.get("dpace_alpha", 0.5), "speculators_config": SpeculatorsConfig( algorithm=algorithm, proposal_methods=[ @@ -401,6 +397,8 @@ def forward( loss_config: LossConfig | None = None, gamma: float = 4.0, max_anchors: int = 3072, + per_position_loss_weight: str = "fixed-exp-decay", + dpace_alpha: float = 0.5, **kwargs, ): _, logits, targets, aligned_loss_mask, _ = self._backbone_forward( @@ -420,8 +418,8 @@ def forward( self.block_size, gamma=gamma, loss_config=loss_config, - per_position_loss_weight=self.per_position_loss_weight, - dpace_alpha=self.dpace_alpha, + per_position_loss_weight=per_position_loss_weight, + dpace_alpha=dpace_alpha, ) draft_tokens = torch.argmax(logits, dim=-1) From b76c3590ea67e94a2ed79afc556c69cdbebe0fc9 Mon Sep 17 00:00:00 2001 From: Weifan Jiang Date: Tue, 7 Jul 2026 12:50:18 -0400 Subject: [PATCH 3/5] fix make style / make quality issues Signed-off-by: Weifan Jiang --- scripts/train.py | 5 +++-- src/speculators/models/dflash/core.py | 4 +++- src/speculators/models/dflash/metrics.py | 10 +++++++--- src/speculators/models/metrics.py | 16 ++++++++++++++-- 4 files changed, 27 insertions(+), 8 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index 154ca1daf..803139f78 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -966,13 +966,14 @@ def parse_args(): "--per-position-loss-weight", choices=["fixed-exp-decay", "dpace"], default="fixed-exp-decay", - help="Per-position loss weight option for D-PACE support (default: fixed-exp-decay)", + help="Per-position loss weight option for D-PACE support" + "default: fixed-exp-decay", ) parser.add_argument( "--dpace-alpha", type=float, default=0.5, - help="Smoothing constant for D-PACE loss (default: 0.5)" + help="Smoothing constant for D-PACE loss (default: 0.5)", ) # DSpark-specific arguments (sequential Markov head + confidence head). parser.add_argument( diff --git a/src/speculators/models/dflash/core.py b/src/speculators/models/dflash/core.py index 1fe12c6b8..e6b02433e 100644 --- a/src/speculators/models/dflash/core.py +++ b/src/speculators/models/dflash/core.py @@ -207,7 +207,9 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]: loss_config = resolve_loss_config(kwargs["loss_fn"]) gamma = kwargs.get("dflash_decay_gamma", 4.0) max_anchors = kwargs.get("max_anchors", 3072) - per_position_loss_weight = kwargs.get("per_position_loss_weight", "fixed-exp-decay") + per_position_loss_weight = kwargs.get( + "per_position_loss_weight", "fixed-exp-decay" + ) dpace_alpha = kwargs.get("dpace_alpha", 0.5) shared = { "loss_config": loss_config, diff --git a/src/speculators/models/dflash/metrics.py b/src/speculators/models/dflash/metrics.py index 6e3dc8e1c..b3a5fe053 100644 --- a/src/speculators/models/dflash/metrics.py +++ b/src/speculators/models/dflash/metrics.py @@ -7,12 +7,12 @@ from speculators.models.metrics import ( LossConfig, + ce_loss, compound_loss, compute_accuracy_multi_step, dflash_loss_decay, + dpace_loss_weight, kl_div_loss, - ce_loss, - dpace_loss_weight ) _DEFAULT_LOSS_CONFIG: LossConfig = {"kl_div": (kl_div_loss, 1.0)} @@ -58,7 +58,11 @@ def compute_metrics( decay_mult = dpace_loss_weight( elementwise_ce, loss_mask, block_size, alpha=dpace_alpha ) - decay_fn = lambda pos_idx: decay_mult + + def decay_fn( + _pos_idx: torch.Tensor, _weight: torch.Tensor = decay_mult + ) -> torch.Tensor: + return _weight else: decay_fn = partial(dflash_loss_decay, gamma=gamma) diff --git a/src/speculators/models/metrics.py b/src/speculators/models/metrics.py index d4e8e5b42..20fb0d43a 100644 --- a/src/speculators/models/metrics.py +++ b/src/speculators/models/metrics.py @@ -292,11 +292,18 @@ def dpace_loss_weight( alpha: confidence smoothing constant """ with torch.no_grad(): + if not 0.0 < alpha <= 1.0: + raise ValueError(f"alpha must be in (0, 1], got {alpha}") # convert CE to per-position confidence q = torch.exp(-neg_log_q).float() # reshape loss to [num_anchors, block_size] # for intra-block cumulative multiplication + if q.shape[1] % block_size != 0: + raise ValueError( + f"q.shape[1] ({q.shape[1]}) must be divisible by " + f"block_size ({block_size})" + ) num_anchors = q.shape[1] // block_size q = q.reshape(num_anchors, block_size) mask = loss_mask.reshape(num_anchors, block_size).to(q.dtype) @@ -418,8 +425,13 @@ def compound_loss( multi = len(loss_config) > 1 for name, (fn, weight) in loss_config.items(): term = loss_function( - logits, targets, loss_mask, pos_idx, loss_fn=fn, decay_fn=decay_fn, - dpace_precomputed_ce=dpace_precomputed_ce + logits, + targets, + loss_mask, + pos_idx, + loss_fn=fn, + decay_fn=decay_fn, + dpace_precomputed_ce=dpace_precomputed_ce, ) if multi: term_losses[f"{name}_loss"] = term.detach() From ce287cd1d30a04bfdc208f88d647a4bd5fd86ea2 Mon Sep 17 00:00:00 2001 From: Weifan Jiang Date: Tue, 7 Jul 2026 17:54:24 -0400 Subject: [PATCH 4/5] refactor loss computation; enable d-pace for d-spark Signed-off-by: Weifan Jiang --- scripts/train.py | 7 +++++ src/speculators/models/dflash/metrics.py | 18 ++++------- src/speculators/models/dspark/core.py | 10 ++++++ src/speculators/models/dspark/metrics.py | 16 ++++++++-- src/speculators/models/metrics.py | 40 ++++++++++++------------ 5 files changed, 57 insertions(+), 34 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index 803139f78..c9cb89d34 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -1143,6 +1143,13 @@ def parse_args(): provided = explicitly_provided_dests(parser, DECODER_SHAPING_FLAGS) validate_draft_init_args(parser, args, provided) resolve_loss_config(args.loss_fn) + + if args.per_position_loss_weight == "dpace": + if args.loss_fn != "ce": + parser.error("--per-position-loss-weight=dpace requires --loss-fn=ce") + if not 0.0 < args.dpace_alpha <= 1.0: + raise ValueError(f"alpha must be in (0, 1], got {args.dpace_alpha}") + return args diff --git a/src/speculators/models/dflash/metrics.py b/src/speculators/models/dflash/metrics.py index b3a5fe053..561a6039f 100644 --- a/src/speculators/models/dflash/metrics.py +++ b/src/speculators/models/dflash/metrics.py @@ -7,11 +7,10 @@ from speculators.models.metrics import ( LossConfig, - ce_loss, compound_loss, compute_accuracy_multi_step, dflash_loss_decay, - dpace_loss_weight, + dpace_loss_decay, kl_div_loss, ) @@ -52,17 +51,13 @@ def compute_metrics( pos_idx = torch.arange(seq_len, device=logits.device) % block_size pos_idx = pos_idx.unsqueeze(0) # shape: [1, T] - elementwise_ce = None if per_position_loss_weight == "dpace": - elementwise_ce = ce_loss(logits, targets) - decay_mult = dpace_loss_weight( - elementwise_ce, loss_mask, block_size, alpha=dpace_alpha + decay_fn = partial( + dpace_loss_decay, + loss_mask=loss_mask, + block_size=block_size, + dpace_alpha=dpace_alpha, ) - - def decay_fn( - _pos_idx: torch.Tensor, _weight: torch.Tensor = decay_mult - ) -> torch.Tensor: - return _weight else: decay_fn = partial(dflash_loss_decay, gamma=gamma) @@ -73,7 +68,6 @@ def decay_fn( pos_idx, loss_config=loss_config, decay_fn=decay_fn, - dpace_precomputed_ce=elementwise_ce, ) pred_ids = torch.argmax(logits, dim=-1) diff --git a/src/speculators/models/dspark/core.py b/src/speculators/models/dspark/core.py index 6c4726a49..61276b306 100644 --- a/src/speculators/models/dspark/core.py +++ b/src/speculators/models/dspark/core.py @@ -84,11 +84,17 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]: gamma = kwargs.get("dflash_decay_gamma", 4.0) max_anchors = kwargs.get("max_anchors", 3072) confidence_head_alpha = kwargs.get("confidence_head_alpha", 1.0) + per_position_loss_weight = kwargs.get( + "per_position_loss_weight", "fixed-exp-decay" + ) + dpace_alpha = kwargs.get("dpace_alpha", 0.5) shared = { "loss_config": loss_config, "gamma": gamma, "max_anchors": max_anchors, "confidence_head_alpha": confidence_head_alpha, + "per_position_loss_weight": per_position_loss_weight, + "dpace_alpha": dpace_alpha, } return dict(shared), dict(shared) @@ -105,6 +111,8 @@ def forward( gamma: float = 4.0, max_anchors: int = 3072, confidence_head_alpha: float = 1.0, + per_position_loss_weight: str = "fixed-exp-decay", + dpace_alpha: float = 0.5, **kwargs, ): hidden, logits, targets, aligned_loss_mask, anchored_block_indices = ( @@ -167,6 +175,8 @@ def forward( loss_config=loss_config or _DEFAULT_LOSS_CONFIG, gamma=gamma, confidence_head_alpha=confidence_head_alpha, + per_position_loss_weight=per_position_loss_weight, + dpace_alpha=dpace_alpha, ) draft_tokens = torch.argmax(logits, dim=-1) return draft_tokens, loss, metrics diff --git a/src/speculators/models/dspark/metrics.py b/src/speculators/models/dspark/metrics.py index a2187cfa3..630e24267 100644 --- a/src/speculators/models/dspark/metrics.py +++ b/src/speculators/models/dspark/metrics.py @@ -18,6 +18,7 @@ compound_loss, compute_accuracy_multi_step, dflash_loss_decay, + dpace_loss_decay, ) __all__ = [ @@ -37,7 +38,8 @@ def _masked_decayed_mean( loss_mask = loss_mask.to(elementwise.dtype) weighted = elementwise * loss_mask if decay_fn is not None: - weighted = weighted * decay_fn(pos_idx.to(weighted.dtype)) + weighted = weighted * decay_fn(pos_idx.to(weighted.dtype), + elementwise_loss=elementwise) denominator = loss_mask.sum(dim=1) + _EPS return (weighted.sum(dim=1) / denominator).mean() @@ -51,13 +53,23 @@ def compute_metrics( loss_config: LossConfig, gamma: float = 4.0, confidence_head_alpha: float = 1.0, + per_position_loss_weight: str = "fixed-exp-decay", + dpace_alpha: float = 0.5, ) -> tuple[torch.Tensor, dict]: """Compute the DSpark loss and a metrics dict (``*_sum``/``*_total`` pairs).""" device = logits.device seq_len = logits.shape[1] pos_idx = (torch.arange(seq_len, device=device) % block_size).unsqueeze(0) - decay_fn = partial(dflash_loss_decay, gamma=gamma) + if per_position_loss_weight == "dpace": + decay_fn = partial( + dpace_loss_decay, + loss_mask=loss_mask, + block_size=block_size, + dpace_alpha=dpace_alpha, + ) + else: + decay_fn = partial(dflash_loss_decay, gamma=gamma) loss, term_losses = compound_loss( logits, targets, loss_mask, pos_idx, loss_config=loss_config, decay_fn=decay_fn diff --git a/src/speculators/models/metrics.py b/src/speculators/models/metrics.py index 20fb0d43a..0a790f0a8 100644 --- a/src/speculators/models/metrics.py +++ b/src/speculators/models/metrics.py @@ -244,7 +244,7 @@ def lk_hybrid_loss( return elementwise_loss # noqa: RET504 -def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float): +def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float, **kwargs): """Compute DFlash-style exponential decay weights per position. Position 0 gets weight 0, position 1 gets weight 1, and subsequent positions @@ -265,7 +265,7 @@ def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float): return decay_mult # noqa: RET504 -def exp_loss_decay(pos_idx: torch.Tensor, gamma: float): +def exp_loss_decay(pos_idx: torch.Tensor, gamma: float, **kwargs): """Compute simple exponential decay weights as gamma^pos_idx. Args: @@ -278,24 +278,28 @@ def exp_loss_decay(pos_idx: torch.Tensor, gamma: float): return gamma**pos_idx -def dpace_loss_weight( - neg_log_q: torch.Tensor, # shape: [1, seq_len] - loss_mask: torch.Tensor, # shape: [1, seq_len] +def dpace_loss_decay( + pos_idx: torch.Tensor, + loss_mask: torch.Tensor, block_size: int, - alpha: float = 0.5, + dpace_alpha: float, + elementwise_loss: torch.Tensor, + **kwargs, ): """ Per-position block-drafting loss weight based on D-PACE Args: - neg_log_q: per-draft-position confidence (negative log-likelihood) - alpha: confidence smoothing constant + elementwise_loss: requires to be cross-entropy loss, negative log-likelihood + of per-position confidence + dpace_alpha: confidence smoothing constant + + Returns: + Decay multiplier tensor with same shape as pos_idx. """ with torch.no_grad(): - if not 0.0 < alpha <= 1.0: - raise ValueError(f"alpha must be in (0, 1], got {alpha}") # convert CE to per-position confidence - q = torch.exp(-neg_log_q).float() + q = torch.exp(-elementwise_loss).float() # reshape loss to [num_anchors, block_size] # for intra-block cumulative multiplication @@ -309,7 +313,7 @@ def dpace_loss_weight( mask = loss_mask.reshape(num_anchors, block_size).to(q.dtype) # smoothed confidence for numerical stability - smooth = (1.0 - alpha) * q + alpha + smooth = (1.0 - dpace_alpha) * q + dpace_alpha smooth = torch.where(mask > 0, smooth, torch.ones_like(smooth)) # prefix cumulative production @@ -408,7 +412,6 @@ def compound_loss( pos_idx: torch.Tensor, loss_config: LossConfig, decay_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, - dpace_precomputed_ce: torch.Tensor | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """Compute a weighted sum of loss terms. @@ -431,7 +434,6 @@ def compound_loss( pos_idx, loss_fn=fn, decay_fn=decay_fn, - dpace_precomputed_ce=dpace_precomputed_ce, ) if multi: term_losses[f"{name}_loss"] = term.detach() @@ -446,7 +448,6 @@ def loss_function( pos_idx: torch.Tensor, # shape: [1, seq_len] loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = kl_div_loss, decay_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, - dpace_precomputed_ce: torch.Tensor | None = None, ): """Compute masked, optionally position-decayed training loss. @@ -461,16 +462,15 @@ def loss_function( Returns: Scalar mean loss across the batch. """ - if dpace_precomputed_ce is not None: - elementwise_loss = dpace_precomputed_ce - else: - elementwise_loss = loss_fn(logits, targets) # shape: [1, seq_len] + elementwise_loss = loss_fn(logits, targets) # shape: [1, seq_len] loss_mask = loss_mask.to(elementwise_loss.dtype) elementwise_loss = elementwise_loss * loss_mask if decay_fn is not None: - decay_mult = decay_fn(pos_idx.to(elementwise_loss.dtype)) + decay_mult = decay_fn( + pos_idx.to(elementwise_loss.dtype), elementwise_loss=elementwise_loss + ) elementwise_loss = elementwise_loss * decay_mult denominator = loss_mask.sum(dim=1) + _EPS From fb2e9593d5a7bb0b0d060b2ae0f5c5b3ef90a32a Mon Sep 17 00:00:00 2001 From: Weifan Jiang Date: Wed, 8 Jul 2026 10:14:10 -0400 Subject: [PATCH 5/5] fix make style/quality issues Signed-off-by: Weifan Jiang --- src/speculators/models/dspark/metrics.py | 7 ++++--- src/speculators/models/metrics.py | 12 ++++++------ 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/src/speculators/models/dspark/metrics.py b/src/speculators/models/dspark/metrics.py index 630e24267..34c4e0765 100644 --- a/src/speculators/models/dspark/metrics.py +++ b/src/speculators/models/dspark/metrics.py @@ -32,14 +32,15 @@ def _masked_decayed_mean( elementwise: torch.Tensor, # [1, T] loss_mask: torch.Tensor, # [1, T] pos_idx: torch.Tensor, # [1, T] - decay_fn: Callable[[torch.Tensor], torch.Tensor] | None, + decay_fn: Callable[..., torch.Tensor] | None, ) -> torch.Tensor: """Masked, optionally position-decayed mean of a precomputed per-position term.""" loss_mask = loss_mask.to(elementwise.dtype) weighted = elementwise * loss_mask if decay_fn is not None: - weighted = weighted * decay_fn(pos_idx.to(weighted.dtype), - elementwise_loss=elementwise) + weighted = weighted * decay_fn( + pos_idx.to(weighted.dtype), elementwise_loss=elementwise + ) denominator = loss_mask.sum(dim=1) + _EPS return (weighted.sum(dim=1) / denominator).mean() diff --git a/src/speculators/models/metrics.py b/src/speculators/models/metrics.py index 0a790f0a8..0af8978f7 100644 --- a/src/speculators/models/metrics.py +++ b/src/speculators/models/metrics.py @@ -244,7 +244,7 @@ def lk_hybrid_loss( return elementwise_loss # noqa: RET504 -def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float, **kwargs): +def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float, **_kwargs): """Compute DFlash-style exponential decay weights per position. Position 0 gets weight 0, position 1 gets weight 1, and subsequent positions @@ -265,7 +265,7 @@ def dflash_loss_decay(pos_idx: torch.Tensor, gamma: float, **kwargs): return decay_mult # noqa: RET504 -def exp_loss_decay(pos_idx: torch.Tensor, gamma: float, **kwargs): +def exp_loss_decay(pos_idx: torch.Tensor, gamma: float, **_kwargs): """Compute simple exponential decay weights as gamma^pos_idx. Args: @@ -279,12 +279,12 @@ def exp_loss_decay(pos_idx: torch.Tensor, gamma: float, **kwargs): def dpace_loss_decay( - pos_idx: torch.Tensor, + pos_idx: torch.Tensor, # noqa: ARG001 loss_mask: torch.Tensor, block_size: int, dpace_alpha: float, elementwise_loss: torch.Tensor, - **kwargs, + **_kwargs, ): """ Per-position block-drafting loss weight based on D-PACE @@ -411,7 +411,7 @@ def compound_loss( loss_mask: torch.Tensor, pos_idx: torch.Tensor, loss_config: LossConfig, - decay_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, + decay_fn: Callable[..., torch.Tensor] | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """Compute a weighted sum of loss terms. @@ -447,7 +447,7 @@ def loss_function( loss_mask: torch.Tensor, # shape: [1, seq_len] pos_idx: torch.Tensor, # shape: [1, seq_len] loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = kl_div_loss, - decay_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, + decay_fn: Callable[..., torch.Tensor] | None = None, ): """Compute masked, optionally position-decayed training loss.