diff --git a/scripts/train.py b/scripts/train.py index e4a103713..c9cb89d34 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -961,6 +961,20 @@ 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", @@ -1129,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/core.py b/src/speculators/models/dflash/core.py index af04c5222..e6b02433e 100644 --- a/src/speculators/models/dflash/core.py +++ b/src/speculators/models/dflash/core.py @@ -207,10 +207,16 @@ 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) @@ -393,6 +399,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( @@ -412,6 +420,8 @@ def forward( self.block_size, gamma=gamma, loss_config=loss_config, + per_position_loss_weight=per_position_loss_weight, + dpace_alpha=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..561a6039f 100644 --- a/src/speculators/models/dflash/metrics.py +++ b/src/speculators/models/dflash/metrics.py @@ -10,6 +10,7 @@ compound_loss, compute_accuracy_multi_step, dflash_loss_decay, + dpace_loss_decay, kl_div_loss, ) @@ -23,6 +24,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 +36,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 +51,23 @@ def compute_metrics( pos_idx = torch.arange(seq_len, device=logits.device) % block_size pos_idx = pos_idx.unsqueeze(0) # shape: [1, T] + 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=partial(dflash_loss_decay, gamma=gamma), + decay_fn=decay_fn, ) 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..34c4e0765 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__ = [ @@ -31,13 +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)) + 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 +54,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 70f0f7c4a..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): +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,6 +278,57 @@ def exp_loss_decay(pos_idx: torch.Tensor, gamma: float): return gamma**pos_idx +def dpace_loss_decay( + pos_idx: torch.Tensor, # noqa: ARG001 + loss_mask: torch.Tensor, + block_size: int, + dpace_alpha: float, + elementwise_loss: torch.Tensor, + **_kwargs, +): + """ + Per-position block-drafting loss weight based on D-PACE + + Args: + 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(): + # convert CE to per-position confidence + q = torch.exp(-elementwise_loss).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) + + # smoothed confidence for numerical stability + smooth = (1.0 - dpace_alpha) * q + dpace_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, @@ -360,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. @@ -377,7 +428,12 @@ 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, ) if multi: term_losses[f"{name}_loss"] = term.detach() @@ -391,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. @@ -412,7 +468,9 @@ def loss_function( 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