# coding=utf-8 # Copyright 2024 The WeDLM team and the HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """PyTorch WeDLM model.""" from typing import Optional, Tuple, Union, Dict, List, Callable import math import torch from torch import nn import torch.nn.functional as F from transformers import PreTrainedModel, GenerationMixin from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.processing_utils import Unpack from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple from transformers.utils.generic import check_model_inputs from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask from transformers.modeling_layers import GradientCheckpointingLayer from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS # Import attention-related utilities try: from transformers.modeling_flash_attention_utils import FlashAttentionKwargs except ImportError: FlashAttentionKwargs = dict from .configuration_wedlm import WeDLMConfig import logging logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) # ============================================================================ # Flow Matching / Rectified Flow helpers # ============================================================================ class WeDLMFlowTimeEmbedding(nn.Module): """Sinusoidal timestep embedding + MLP, used to condition Flow Matching / Rectified Flow. The module is intentionally lightweight and conditions the velocity field on continuous timesteps. Timesteps are assumed to be normalized to [0, 1] (float). Internally, a configurable scale is applied before sinusoidal features are computed. """ def __init__(self, config: WeDLMConfig): super().__init__() self.hidden_size = config.hidden_size self.time_embed_dim = int(getattr(config, "flow_time_embedding_dim", 256)) self.max_period = int(getattr(config, "flow_time_embedding_max_period", 10000)) self.time_scale = float(getattr(config, "flow_time_scale", 1000.0)) self.linear_1 = nn.Linear(self.time_embed_dim, self.hidden_size) self.act = nn.SiLU() self.linear_2 = nn.Linear(self.hidden_size, self.hidden_size) @staticmethod def _sinusoidal_embedding(timesteps: torch.Tensor, dim: int, max_period: int) -> torch.Tensor: """Create sinusoidal timestep embeddings. timesteps: (batch,) float tensor. Returns: (batch, dim) float tensor. """ if timesteps.ndim != 1: timesteps = timesteps.view(-1) half = dim // 2 device = timesteps.device dtype = torch.float32 freqs = torch.exp( -math.log(max_period) * torch.arange(0, half, device=device, dtype=dtype) / max(half, 1) ) args = timesteps.to(dtype)[:, None] * freqs[None] emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) if dim % 2 == 1: emb = torch.cat([emb, torch.zeros((emb.shape[0], 1), device=device, dtype=dtype)], dim=-1) return emb def forward(self, timesteps: torch.Tensor) -> torch.Tensor: # timesteps expected in [0,1]; scale to a more typical diffusion timestep range. t = timesteps.to(dtype=torch.float32) if t.ndim == 0: t = t[None] if t.ndim != 1: t = t.view(-1) t = t.clamp(0.0, 1.0) * self.time_scale emb = self._sinusoidal_embedding(t, self.time_embed_dim, self.max_period) # NOTE: In bf16/fp16 training (and especially under PEFT/LoRA), the wrapped Linear's base weights # can be low-precision (e.g. bfloat16) while the sinusoidal features are float32. # Torch's F.linear requires the input and weight dtypes to match, so we explicitly cast here # to the *base* layer's weight dtype. base_linear_1 = getattr(self.linear_1, "base_layer", self.linear_1) w1 = getattr(base_linear_1, "weight", None) if w1 is not None: emb = emb.to(dtype=w1.dtype) emb = self.linear_1(emb) emb = self.act(emb) base_linear_2 = getattr(self.linear_2, "base_layer", self.linear_2) w2 = getattr(base_linear_2, "weight", None) if w2 is not None: emb = emb.to(dtype=w2.dtype) emb = self.linear_2(emb) return emb # ============================================================================ # ============================================================================ # Core Components (self-contained, no Qwen2 dependency) # ============================================================================ class WeDLMMLP(nn.Module): """WeDLM MLP module with SwiGLU activation.""" def __init__(self, config: WeDLMConfig): super().__init__() self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x: torch.Tensor) -> torch.Tensor: down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) return down_proj class WeDLMRMSNorm(nn.Module): """WeDLM RMSNorm, equivalent to T5LayerNorm.""" def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) def extra_repr(self) -> str: return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" class WeDLMRotaryEmbedding(nn.Module): """WeDLM Rotary Position Embedding.""" def __init__(self, config: WeDLMConfig, device=None): super().__init__() # Determine rope_type from config if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict): self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type", "default")) else: self.rope_type = "default" self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config # Get initialization function if self.rope_type == "default": inv_freq, self.attention_scaling = self._compute_default_rope_parameters(config, device) else: rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = rope_init_fn(config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.original_inv_freq = self.inv_freq @staticmethod def _compute_default_rope_parameters( config: WeDLMConfig, device: Optional[torch.device] = None, ) -> Tuple[torch.Tensor, float]: """ Computes the inverse frequencies for default RoPE. Args: config: Model configuration device: Device to place the tensors on Returns: Tuple of (inv_freq tensor, attention_scaling factor) """ base = config.rope_theta dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads # Compute the inverse frequencies inv_freq = 1.0 / ( base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) ) attention_factor = 1.0 return inv_freq, attention_factor @torch.no_grad() def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """ Compute rotary position embeddings. Args: x: Input tensor, used for dtype and device position_ids: Position indices Returns: Tuple of (cos, sin) tensors """ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) position_ids_expanded = position_ids[:, None, :].float() device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" # Force float32 computation for numerical stability with torch.amp.autocast(device_type=device_type, enabled=False): freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * self.attention_scaling sin = emb.sin() * self.attention_scaling return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) # ============================================================================ # Attention Utilities # ============================================================================ def rotate_half(x: torch.Tensor) -> torch.Tensor: """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, position_ids: Optional[torch.Tensor] = None, unsqueeze_dim: int = 1 ) -> Tuple[torch.Tensor, torch.Tensor]: """Applies Rotary Position Embedding to the query and key tensors.""" cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ Repeats key/value heads to match the number of query heads (for GQA). Equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). """ batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: Optional[torch.Tensor], scaling: float, dropout: float = 0.0, **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: """Eager (standard) attention implementation.""" key_states = repeat_kv(key, module.num_key_value_groups) value_states = repeat_kv(value, module.num_key_value_groups) attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling if attention_mask is not None: causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] attn_weights = attn_weights + causal_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights # ============================================================================ # Attention Layer # ============================================================================ class WeDLMAttention(nn.Module): """ WeDLM Attention module. Supports both: - Qwen2.5 style: with QKV bias, no QK Norm - Qwen3 style: configurable QKV bias, with QK Norm """ def __init__(self, config: WeDLMConfig, layer_idx: int): super().__init__() self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None self.config = config self.layer_idx = layer_idx self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads self.scaling = self.head_dim ** -0.5 self.attention_dropout = config.attention_dropout self.is_causal = True # Support configurable attention_bias (Qwen2.5: True, Qwen3: False by default) attention_bias = getattr(config, "attention_bias", True) self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=attention_bias) self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=attention_bias) self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=attention_bias) self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) # Support optional QK Norm (Qwen3 feature) self.qk_norm = getattr(config, "qk_norm", False) if self.qk_norm: self.q_norm = WeDLMRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = WeDLMRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None def forward( self, hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_values: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) if self.qk_norm: # Qwen3 style: apply norm after projection, before transpose query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) else: # Qwen2 style: no norm query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_values is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) # Select attention implementation attention_interface: Callable = eager_attention_forward if self.config._attn_implementation != "eager" and self.config._attn_implementation in ALL_ATTENTION_FUNCTIONS: attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, sliding_window=self.sliding_window, **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights # ============================================================================ # Decoder Layer # ============================================================================ class WeDLMDecoderLayer(GradientCheckpointingLayer): """WeDLM Decoder Layer with pre-norm architecture.""" def __init__(self, config: WeDLMConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.self_attn = WeDLMAttention(config=config, layer_idx=layer_idx) self.mlp = WeDLMMLP(config) self.input_layernorm = WeDLMRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = WeDLMRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.attention_type = config.layer_types[layer_idx] def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs: Unpack[TransformersKwargs], ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Args: hidden_states: Input tensor of shape `(batch, seq_len, embed_dim)` attention_mask: Attention mask of size `(batch, sequence_length)` position_ids: Position indices past_key_values: Cached past key and value projection states output_attentions: Whether to return attention weights use_cache: Whether to use KV cache cache_position: Position in the cache position_embeddings: Tuple of (cos, sin) for rotary embeddings """ residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states, self_attn_weights = self.self_attn( hidden_states=hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, cache_position=cache_position, **kwargs, ) hidden_states = residual + hidden_states # Feed Forward residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) return outputs # ============================================================================ # Model Classes # ============================================================================ @auto_docstring class WeDLMPreTrainedModel(PreTrainedModel): """Base class for WeDLM models.""" config_class = WeDLMConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["WeDLMDecoderLayer"] _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn = True _supports_sdpa = True _supports_flex_attn = True _can_compile_fullgraph = True _supports_attention_backend = True _can_record_outputs = { "hidden_states": WeDLMDecoderLayer, "attentions": WeDLMAttention, } @auto_docstring class WeDLMModel(WeDLMPreTrainedModel): """ WeDLM base model outputting raw hidden states. """ def __init__(self, config: WeDLMConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList( [WeDLMDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = WeDLMRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = WeDLMRotaryEmbedding(config=config) self.gradient_checkpointing = False self.has_sliding_layers = "sliding_attention" in self.config.layer_types # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value @check_model_inputs @auto_docstring def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs: Unpack[TransformersKwargs], ) -> Union[Tuple, BaseModelOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device ) if position_ids is None: position_ids = cache_position.unsqueeze(0) # Prepare attention masks if not isinstance(causal_mask_mapping := attention_mask, dict): mask_kwargs = { "config": self.config, "input_embeds": inputs_embeds, "attention_mask": attention_mask, "cache_position": cache_position, "past_key_values": past_key_values, "position_ids": position_ids, } causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), } if self.has_sliding_layers: causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs) hidden_states = inputs_embeds # Create position embeddings to be shared across the decoder layers position_embeddings = self.rotary_emb(hidden_states, position_ids) # Decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None for decoder_layer in self.layers[: self.config.num_hidden_layers]: if output_hidden_states: all_hidden_states += (hidden_states,) layer_outputs = decoder_layer( hidden_states, attention_mask=causal_mask_mapping[decoder_layer.attention_type], position_ids=position_ids, past_key_values=past_key_values, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, **kwargs, ) hidden_states = layer_outputs[0] if output_attentions: all_self_attns += (layer_outputs[1],) hidden_states = self.norm(hidden_states) if output_hidden_states: all_hidden_states += (hidden_states,) if not return_dict: return tuple(v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None, hidden_states=all_hidden_states, attentions=all_self_attns, ) @auto_docstring class WeDLMForCausalLM(WeDLMPreTrainedModel, GenerationMixin): """ WeDLM Model for Flow-Matching language modeling (Rectified Flow in token-embedding space). - Training (`labels` provided): optimizes a Flow Matching objective on a selected subset of token positions. Large-vocabulary projection (`lm_head`) is skipped by default during training for lower cost. - Inference (no `labels`): behaves like a standard causal LM (returns logits). - Fast decoding: use `generate_wedlm` (Flow-Matching block decoding). """ _tied_weights_keys = ["lm_head.weight"] def __init__(self, config: WeDLMConfig): super().__init__(config) self.model = WeDLMModel(config) self.vocab_size = config.vocab_size # Token discretization head (used for inference / evaluation / discretization at the end of each flow block) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) # Flow Matching modules self.flow_time_embed = WeDLMFlowTimeEmbedding(config) self.flow_head = nn.Linear(config.hidden_size, config.hidden_size, bias=False) # Initialize weights and apply final processing self.post_init() # ---------------------------- # Embedding plumbing # ---------------------------- def get_input_embeddings(self): return self.model.embed_tokens def set_input_embeddings(self, value): self.model.embed_tokens = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def set_decoder(self, decoder): self.model = decoder def get_decoder(self): return self.model # ---------------------------- # Flow Matching utilities # ---------------------------- def _select_flow_targets( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor], labels: Optional[torch.LongTensor], flow_target_mask: Optional[torch.BoolTensor], ) -> torch.BoolTensor: """ Determine which token positions participate in Flow Matching loss. Priority: 1) `flow_target_mask` argument 2) `config.flow_train_strategy` """ bsz, seq_len = input_ids.shape device = input_ids.device if attention_mask is not None: valid = attention_mask.to(dtype=torch.bool, device=device) else: # Best-effort fallback: treat non-pad as valid. pad_id = getattr(self.config, "pad_token_id", None) if pad_id is None: valid = torch.ones((bsz, seq_len), dtype=torch.bool, device=device) else: valid = input_ids.ne(pad_id) if labels is not None: valid = valid & labels.ne(-100) if flow_target_mask is not None: target = flow_target_mask.to(dtype=torch.bool, device=device) return target & valid strategy = str(getattr(self.config, "flow_train_strategy", "suffix_block")).lower() min_targets = int(getattr(self.config, "flow_train_min_target_tokens", 1)) if strategy == "random": ratio = float(getattr(self.config, "flow_train_mask_ratio", 0.15)) # Sample only from valid positions; guarantee at least `min_targets` if possible. rand = torch.rand((bsz, seq_len), device=device) target = (rand < ratio) & valid if min_targets > 0: for b in range(bsz): if valid[b].any() and target[b].sum().item() < min_targets: valid_idx = valid[b].nonzero(as_tuple=True)[0] # Select the last positions (deterministic tie-break) to fill up. need = min(min_targets - target[b].sum().item(), valid_idx.numel()) if need > 0: target[b, valid_idx[-need:]] = True return target # Default: suffix_block block_size = int(getattr(self.config, "flow_train_block_size", 64)) target = torch.zeros((bsz, seq_len), dtype=torch.bool, device=device) for b in range(bsz): valid_idx = valid[b].nonzero(as_tuple=True)[0] if valid_idx.numel() == 0: continue # Do not target the very first valid token by default (no context); if needed, user can pass flow_target_mask. if valid_idx.numel() == 1: continue # Suffix contiguous block among valid positions. k = min(block_size, valid_idx.numel() - 1) k = max(k, min_targets) k = min(k, valid_idx.numel() - 1) # ensure at least one context token remains if k <= 0: continue target_pos = valid_idx[-k:] target[b, target_pos] = True return target def _normalize_timesteps( self, timesteps: Optional[torch.FloatTensor], target_mask: torch.BoolTensor, ) -> torch.FloatTensor: """ Returns per-token normalized timesteps t in [0, 1], shape (bsz, seq_len). """ device = target_mask.device bsz, seq_len = target_mask.shape if timesteps is None: t = torch.rand((bsz, seq_len), device=device, dtype=torch.float32) # Non-target positions: t=1.0 (data endpoint) so that any accidental use is benign. t = torch.where(target_mask, t, torch.ones_like(t)) return t t_in = timesteps.to(device=device, dtype=torch.float32) if t_in.ndim == 0: t = t_in.view(1, 1).expand(bsz, seq_len) elif t_in.ndim == 1: if t_in.shape[0] == 1 and bsz > 1: t = t_in.view(1, 1).expand(bsz, seq_len) elif t_in.shape[0] == bsz: t = t_in.view(bsz, 1).expand(bsz, seq_len) else: raise ValueError( f"flow_timesteps must be scalar, shape (bsz,), or shape (bsz, seq_len); got {tuple(t_in.shape)}" ) elif t_in.ndim == 2: if t_in.shape != (bsz, seq_len): raise ValueError( f"flow_timesteps must have shape (bsz, seq_len) == {(bsz, seq_len)}; got {tuple(t_in.shape)}" ) t = t_in else: raise ValueError( f"flow_timesteps must be scalar, 1D, or 2D; got ndim={t_in.ndim} with shape {tuple(t_in.shape)}" ) # Clamp into [0,1] for numerical safety. t = torch.clamp(t, 0.0, 1.0) t = torch.where(target_mask, t, torch.ones_like(t)) return t def _build_flow_inputs( self, input_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor], labels: Optional[torch.LongTensor], flow_target_mask: Optional[torch.BoolTensor], flow_timesteps: Optional[torch.FloatTensor], flow_noise: Optional[torch.FloatTensor], ) -> Tuple[torch.FloatTensor, torch.BoolTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: """ Prepare (inputs_embeds, target_mask, t, noise, clean_embeds) for Flow Matching training. """ clean_embeds = self.model.embed_tokens(input_ids) bsz, seq_len, hidden = clean_embeds.shape device = clean_embeds.device target_mask = self._select_flow_targets( input_ids=input_ids, attention_mask=attention_mask, labels=labels, flow_target_mask=flow_target_mask, ) t = self._normalize_timesteps(flow_timesteps, target_mask=target_mask) sigma = float(getattr(self.config, "flow_init_sigma", 1.0)) if flow_noise is None: noise = torch.randn_like(clean_embeds, dtype=torch.float32) * sigma else: noise = flow_noise.to(device=device, dtype=torch.float32) if noise.shape != clean_embeds.shape: raise ValueError( f"flow_noise must have the same shape as token embeddings {tuple(clean_embeds.shape)}; got {tuple(noise.shape)}" ) # Rectified Flow path: X_t = (1 - t) * X_0 + t * X_1 # Here X_0 is noise, X_1 is data (token embeddings). t_exp = t.unsqueeze(-1) x_t = (1.0 - t_exp) * noise + t_exp * clean_embeds.to(dtype=torch.float32) # Time conditioning is provided by adding a learned timestep embedding to the *input embeddings* # for flow-target positions (so the Transformer can use t). inputs_embeds = clean_embeds.to(dtype=torch.float32) if target_mask.any(): # Compute time embedding only for targets to reduce overhead. t_flat = t[target_mask].reshape(-1) time_cond = self.flow_time_embed(t_flat).to(dtype=inputs_embeds.dtype) inputs_embeds = inputs_embeds.clone() inputs_embeds[target_mask] = x_t[target_mask] + time_cond else: inputs_embeds = inputs_embeds.clone() return inputs_embeds.to(dtype=clean_embeds.dtype), target_mask, t, noise, clean_embeds # ---------------------------- # Decoding: Flow Matching block decoding # ---------------------------- def _top_k_top_p_filtering( self, logits: torch.Tensor, top_k: int = 0, top_p: float = 1.0, filter_value: float = -float("inf"), ) -> torch.Tensor: """Apply top-k and/or nucleus (top-p) filtering to logits.""" if top_k is not None and top_k > 0: top_k = min(top_k, logits.size(-1)) indices_to_remove = logits < torch.topk(logits, top_k, dim=-1).values[..., -1, None] logits = logits.masked_fill(indices_to_remove, filter_value) if top_p is not None and 0.0 < top_p < 1.0: sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) sorted_indices_to_remove = cumulative_probs > top_p sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] = False indices_to_remove = sorted_indices_to_remove.scatter(-1, sorted_indices, sorted_indices_to_remove) logits = logits.masked_fill(indices_to_remove, filter_value) return logits def _sample_from_logits( self, logits: torch.Tensor, temperature: float = 1.0, top_p: float = 1.0, top_k: int = 0, ) -> torch.Tensor: """Sample token IDs from logits with temperature + (top-k, top-p) filtering.""" if temperature is None or temperature <= 0: temperature = 1.0 logits = logits / float(temperature) logits = self._top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p) probs = F.softmax(logits, dim=-1) return torch.multinomial(probs, num_samples=1).squeeze(-1) @torch.no_grad() def generate_wedlm( self, input_ids: torch.LongTensor, max_new_tokens: int, block_size: Optional[int] = None, num_steps: Optional[int] = None, flow_init_sigma: Optional[float] = None, discretization: Optional[str] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, top_k: Optional[int] = None, pad_token_id: Optional[int] = None, eos_token_id: Optional[int] = None, return_stats: bool = True, **kwargs, ) -> Union[torch.LongTensor, Dict]: """ Flow-Matching block decoding. Generates `block_size` tokens per block using `num_steps` Euler steps, and runs the vocabulary projection (`lm_head`) only once per block (final discretization). """ device = input_ids.device if pad_token_id is None: pad_token_id = self.config.pad_token_id if eos_token_id is None: eos_token_id = getattr(self.config, "eos_token_id", None) if block_size is None: block_size = int(getattr(self.config, "flow_block_size", 64)) if num_steps is None: num_steps = int(getattr(self.config, "flow_inference_steps", 8)) if flow_init_sigma is None: flow_init_sigma = float(getattr(self.config, "flow_init_sigma", 1.0)) if discretization is None: discretization = str(getattr(self.config, "flow_discretization", "argmax")).lower() if temperature is None: temperature = float(getattr(self.config, "flow_temperature", 1.0)) if top_p is None: top_p = float(getattr(self.config, "flow_top_p", 1.0)) if top_k is None: top_k = int(getattr(self.config, "flow_top_k", 0)) batch_size = input_ids.shape[0] all_generated: List[torch.Tensor] = [] all_sample_stats: List[Dict] = [] num_blocks = (max_new_tokens + block_size - 1) // block_size for batch_idx in range(batch_size): sample_ids = input_ids[batch_idx] if pad_token_id is not None: pad_mask = sample_ids.ne(pad_token_id) if pad_mask.any(): valid_length = int(pad_mask.sum().item()) prefix_ids = sample_ids[:valid_length] else: prefix_ids = sample_ids else: prefix_ids = sample_ids prefix_ids = prefix_ids.clone() prefix_length = prefix_ids.shape[0] sample_stats = { "input_length": prefix_length, "num_blocks": num_blocks, "block_size": block_size, "num_steps": num_steps, "sigma": float(flow_init_sigma), "generated_tokens": 0, "blocks": [], } current_ids = prefix_ids for block_idx in range(num_blocks): remaining = max_new_tokens - block_idx * block_size cur_block = min(block_size, remaining) if cur_block <= 0: break # State variable: current embedding estimates for the block (initialized from Gaussian noise) x = torch.randn((cur_block, self.config.hidden_size), device=device, dtype=torch.float32) * float(flow_init_sigma) dt = 1.0 / float(num_steps) # Euler integration from t=0 (noise) to t=1 (data) for step in range(num_steps): t = float(step) / float(num_steps) # Create per-token timesteps matching training behavior (all same t for this step) t_tensor = torch.full((cur_block,), t, device=device, dtype=torch.float32) # Build embeddings in *model dtype* to avoid dtype mismatch in Linear layers when AMP is off. context_embeds = self.model.embed_tokens(current_ids.unsqueeze(0)) ctx_dtype = context_embeds.dtype # Per-token time conditioning (cur_block,) -> (cur_block, H) t_cond = self.flow_time_embed(t_tensor).to(dtype=ctx_dtype) # (cur_block, H) # Context tokens are discrete; block tokens are continuous (x) + time conditioning. block_embeds = x.to(dtype=ctx_dtype) + t_cond # (cur_block, H) block_embeds = block_embeds.view(1, cur_block, -1) # (1, cur_block, H) inputs_embeds = torch.cat([context_embeds, block_embeds], dim=1) seq_len = inputs_embeds.shape[1] attention_mask = torch.ones((1, seq_len), dtype=torch.long, device=device) position_ids = torch.arange(seq_len, dtype=torch.long, device=device).unsqueeze(0) outputs = self.model( input_ids=None, inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, use_cache=False, return_dict=True, ) h_new = outputs.last_hidden_state[:, -cur_block:, :] # (1, cur_block, H) # Predict velocity in a dtype-compatible way, then accumulate in fp32. h_in = h_new base_flow = getattr(self.flow_head, "base_layer", self.flow_head) w_flow = getattr(base_flow, "weight", None) if w_flow is not None: h_in = h_in.to(dtype=w_flow.dtype) v = self.flow_head(h_in).to(dtype=torch.float32).squeeze(0) # (cur_block, H) x = x + dt * v # Discretize final embeddings into token IDs base_lm = getattr(self.lm_head, "base_layer", self.lm_head) w_lm = getattr(base_lm, "weight", None) x_lm = x if w_lm is not None: x_lm = x_lm.to(dtype=w_lm.dtype) logits = self.lm_head(x_lm).to(dtype=torch.float32) # (cur_block, vocab) if discretization == "sample": next_ids = self._sample_from_logits(logits, temperature=temperature, top_p=top_p, top_k=top_k) else: next_ids = torch.argmax(logits, dim=-1) # Optional early stop on EOS within the block if eos_token_id is not None: eos_positions = (next_ids == eos_token_id).nonzero(as_tuple=True)[0] if eos_positions.numel() > 0: cut = int(eos_positions[0].item()) + 1 next_ids = next_ids[:cut] current_ids = torch.cat([current_ids, next_ids.to(dtype=torch.long)], dim=0) sample_stats["generated_tokens"] += int(next_ids.numel()) sample_stats["blocks"].append( { "block_idx": block_idx, "target_block_size": int(cur_block), "actual_block_tokens": int(next_ids.numel()), } ) if eos_token_id is not None and next_ids.numel() > 0 and next_ids[-1].item() == int(eos_token_id): break sample_stats["output_length"] = int(current_ids.numel()) all_generated.append(current_ids) all_sample_stats.append(sample_stats) # Pad to max length max_len = max(seq.numel() for seq in all_generated) if all_generated else 0 padded = [] for seq in all_generated: if seq.numel() < max_len: pad = torch.full( (max_len - seq.numel(),), int(pad_token_id) if pad_token_id is not None else 0, dtype=torch.long, device=device, ) seq = torch.cat([seq, pad], dim=0) padded.append(seq) sequences = torch.stack(padded, dim=0) if padded else torch.empty((0, 0), dtype=torch.long, device=device) if not return_stats: return sequences total_steps = int(num_steps) * int(num_blocks) * int(batch_size) return { "sequences": sequences, "stats": { "batch_size": int(batch_size), "max_new_tokens": int(max_new_tokens), "block_size": int(block_size), "num_steps": int(num_steps), "discretization": discretization, "temperature": float(temperature), "top_p": float(top_p), "top_k": int(top_k), "total_flow_evals": total_steps, "per_sample_stats": all_sample_stats, }, } # ---------------------------- # Generate (override to use Flow Matching by default) # ---------------------------- @torch.no_grad() def generate( self, input_ids: Optional[torch.LongTensor] = None, generation_config=None, max_new_tokens: Optional[int] = None, max_length: Optional[int] = None, do_sample: Optional[bool] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, top_k: Optional[int] = None, pad_token_id: Optional[int] = None, eos_token_id: Optional[int] = None, use_flow_matching: bool = True, block_size: Optional[int] = None, num_steps: Optional[int] = None, flow_init_sigma: Optional[float] = None, streamer=None, **kwargs, ): """ Override generate() to use Flow Matching decoding by default. Set `use_flow_matching=False` to fall back to standard AR generation. """ # Fall back to standard AR generation if requested if not use_flow_matching: return super().generate( input_ids=input_ids, generation_config=generation_config, max_new_tokens=max_new_tokens, max_length=max_length, do_sample=do_sample, temperature=temperature, top_p=top_p, top_k=top_k, pad_token_id=pad_token_id, eos_token_id=eos_token_id, streamer=streamer, **kwargs, ) # Extract parameters from generation_config if provided if generation_config is not None: if max_new_tokens is None: max_new_tokens = getattr(generation_config, "max_new_tokens", None) if max_length is None: max_length = getattr(generation_config, "max_length", None) if do_sample is None: do_sample = getattr(generation_config, "do_sample", None) if temperature is None: temperature = getattr(generation_config, "temperature", None) if top_p is None: top_p = getattr(generation_config, "top_p", None) if top_k is None: top_k = getattr(generation_config, "top_k", None) if pad_token_id is None: pad_token_id = getattr(generation_config, "pad_token_id", None) if eos_token_id is None: eos_token_id = getattr(generation_config, "eos_token_id", None) # Determine max_new_tokens if max_new_tokens is None: if max_length is not None and input_ids is not None: max_new_tokens = max_length - input_ids.shape[1] else: max_new_tokens = 256 # Default max_new_tokens = max(1, max_new_tokens) # Map do_sample to discretization discretization = "sample" if do_sample else "argmax" # Use config defaults for flow parameters if block_size is None: block_size = getattr(self.config, "flow_block_size", 64) if num_steps is None: num_steps = getattr(self.config, "flow_inference_steps", 8) if flow_init_sigma is None: flow_init_sigma = getattr(self.config, "flow_init_sigma", 1.0) if temperature is None: temperature = getattr(self.config, "flow_temperature", 1.0) if top_p is None: top_p = getattr(self.config, "flow_top_p", 1.0) if top_k is None: top_k = getattr(self.config, "flow_top_k", 0) if pad_token_id is None: pad_token_id = self.config.pad_token_id if eos_token_id is None: eos_token_id = self.config.eos_token_id # Call Flow Matching generation result = self.generate_wedlm( input_ids=input_ids, max_new_tokens=max_new_tokens, block_size=block_size, num_steps=num_steps, flow_init_sigma=flow_init_sigma, discretization=discretization, temperature=temperature, top_p=top_p, top_k=top_k, pad_token_id=pad_token_id, eos_token_id=eos_token_id, return_stats=False, ) # Handle streamer if provided (basic support) if streamer is not None: for token_id in result[0, input_ids.shape[1]:]: streamer.put(token_id.unsqueeze(0).unsqueeze(0)) streamer.end() return result # ---------------------------- # Forward # ---------------------------- @can_return_tuple @auto_docstring def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, # Flow Matching controls (training-time) flow_target_mask: Optional[torch.BoolTensor] = None, flow_timesteps: Optional[torch.FloatTensor] = None, flow_noise: Optional[torch.FloatTensor] = None, return_logits: bool = False, **kwargs: Unpack[TransformersKwargs], ) -> Union[Tuple, CausalLMOutputWithPast]: """ When `labels` is provided: computes Flow Matching loss. Otherwise: returns logits like a standard causal LM. Args: flow_target_mask (`torch.BoolTensor`, *optional*): Boolean mask indicating which positions are flow targets. flow_timesteps (`torch.FloatTensor`, *optional*): Timesteps for flow matching, in range [0, 1]. flow_noise (`torch.FloatTensor`, *optional*): Noise tensor for flow matching interpolation. return_logits (`bool`, defaults to `False`): If True, also compute vocabulary logits during Flow-Matching training. """ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict # ------------------------------------------------------------ # Flow Matching training path (primary) # ------------------------------------------------------------ if labels is not None: if input_ids is None: raise ValueError("Flow-Matching training requires input_ids (token IDs) when labels is provided.") if inputs_embeds is not None: raise ValueError("Do not pass inputs_embeds when training with labels; Flow-Matching builds embeds internally.") inputs_embeds, target_mask, t, noise, clean_embeds = self._build_flow_inputs( input_ids=input_ids, attention_mask=attention_mask, labels=labels, flow_target_mask=flow_target_mask, flow_timesteps=flow_timesteps, flow_noise=flow_noise, ) outputs = self.model( input_ids=None, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_position, **kwargs, ) hidden_states = outputs.last_hidden_state loss = None logits = None if target_mask.any(): # Target velocity: d/dt X_t = X_1 - X_0 (data - noise) on straight-line path. v_target = (clean_embeds.to(dtype=torch.float32) - noise.to(dtype=torch.float32))[target_mask] # PEFT/LoRA can wrap `flow_head` and keep base weights in bf16/fp16. # Ensure dtype alignment for the Linear matmul. hs = hidden_states[target_mask] base_flow = getattr(self.flow_head, "base_layer", self.flow_head) w_flow = getattr(base_flow, "weight", None) if w_flow is not None: hs = hs.to(dtype=w_flow.dtype) v_pred = self.flow_head(hs).to(dtype=torch.float32) flow_loss = F.mse_loss(v_pred, v_target, reduction="mean") w = float(getattr(self.config, "flow_loss_weight", 1.0)) loss = w * flow_loss else: # No valid targets -> zero loss (avoid NaNs). loss = hidden_states.new_tensor(0.0) if return_logits: slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]) if not return_dict: output = (logits,) + (outputs.past_key_values, outputs.hidden_states, outputs.attentions) return (loss,) + output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) # ------------------------------------------------------------ # Standard causal LM path (no labels): logits for evaluation / AR generation # ------------------------------------------------------------ outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_position, **kwargs, ) hidden_states = outputs.last_hidden_state slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]) if not return_dict: output = (logits,) + (outputs.past_key_values, outputs.hidden_states, outputs.attentions) return output return CausalLMOutputWithPast( loss=None, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, position_ids=None, use_cache=True, **kwargs, ): if past_key_values is not None: if inputs_embeds is not None: input_ids = input_ids[:, -cache_position.shape[0] :] elif input_ids.shape[1] != cache_position.shape[0]: input_ids = input_ids[:, cache_position] if attention_mask is not None and position_ids is None: position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) if past_key_values: position_ids = position_ids[:, -input_ids.shape[1] :] if inputs_embeds is not None and cache_position is not None and cache_position[0] == 0: model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None} else: model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format), "inputs_embeds": None} if isinstance(past_key_values, DynamicCache) and attention_mask is not None and attention_mask.ndim == 2: model_inputs["cache_position"] = cache_position model_inputs["past_key_values"] = past_key_values model_inputs["use_cache"] = use_cache model_inputs["position_ids"] = position_ids model_inputs["attention_mask"] = attention_mask return model_inputs model_inputs.update( { "position_ids": position_ids, "cache_position": cache_position, "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, } ) return model_inputs __all__ = [ "WeDLMConfig", "WeDLMPreTrainedModel", "WeDLMModel", "WeDLMForCausalLM", ]