diff --git a/src/proxyz/generate.py b/src/proxyz/generate.py index b0d8ce0..fda6de2 100644 --- a/src/proxyz/generate.py +++ b/src/proxyz/generate.py @@ -182,6 +182,7 @@ def resolve_model_path(model_dir: str) -> str: with torch.no_grad(): out = model.generate( **input_ids, + processor=processor, generation_config=gen_config, logits_processor=logits_processor, ) diff --git a/src/proxyz/models/configuration_xyz.py b/src/proxyz/models/configuration_xyz.py index d30c029..739928d 100644 --- a/src/proxyz/models/configuration_xyz.py +++ b/src/proxyz/models/configuration_xyz.py @@ -4,7 +4,10 @@ # the file from the modular. If any change should be done, please apply the change to the # modular_xyz.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +import contextlib + from huggingface_hub.dataclasses import strict +from proxyz.utils import attr from transformers.configuration_utils import PreTrainedConfig from transformers.modeling_rope_utils import RopeParameters @@ -15,36 +18,101 @@ @auto_docstring(checkpoint="bigict/ProXYZ") @strict class XYZConfig(PreTrainedConfig): - r""" - ```python - >>> from transformers import XYZModel, XYZConfig + """ + Configuration for the XYZ U-Net style protein language model. + + A three-stage architecture: + + 1. **Char encoder** — transformer at character (residue) granularity + 2. **Token trunk** — transformer at BPE-token granularity + 3. **Char decoder** — transformer at character granularity with U-Net skip + connections from the encoder + + The char encoder/decoder share the same ``head_dim`` as the trunk so that + a single ``RotaryEmbedding`` instance can be used across all three stacks. + + This configuration extends [`LlamaConfig`]. All token-level fields + (``hidden_size``, ``intermediate_size``, ``num_hidden_layers``, etc.) are + inherited and control the *trunk* transformer. The char-level fields + below control the encoder and decoder. + + Constraint: ``char_hidden_size == char_num_attention_heads × head_dim``. + + char_hidden_size (`int`, *optional*, defaults to 2048): + Hidden size of the char encoder / decoder transformer. + char_intermediate_size (`int`, *optional*, defaults to 5504): + FFN intermediate size for char-level layers. + char_num_hidden_layers (`int`, *optional*, defaults to 4): + Number of transformer layers in the char encoder and char decoder. + char_num_attention_heads (`int`, *optional*, defaults to 16): + Number of attention heads for char-level self-attention. + char_num_key_value_heads (`int`, *optional*): + Number of KV heads for GQA at char level. Defaults to + ``char_num_attention_heads`` (i.e. multi-head attention) when + not specified. + + use_char_position_ids (`bool`, *optional*, defaults to `False`): + If ``True``, token-level ``position_ids`` are gathered from + ``char_position_ids`` via ``repr_char_idx`` instead of being + computed independently. This keeps RoPE aligned between the + char and token stacks. + + has_char_lm_head (`bool`, *optional*, defaults to `False`): + If ``True``, create a next-character prediction head + (``char_lm_head``) on top of the char decoder output. + + has_cle_lm_head (`bool`, *optional*, defaults to `False`): + If ``True``, create a CLE (Cα–Local–Environment) classification + head on top of the char decoder output. + + has_distogram_lm_head (`bool`, *optional*, defaults to `False`): + If ``True``, create a distogram prediction head on top of the + char decoder output. Predicts pairwise residue–residue distance + bins via an outer-sum MLP. + + distogram_bins_num (`int`, *optional*, defaults to 64): + Number of distance bins for the distogram head. - >>> # Initializing a XYZ x_y_z-7b style configuration - >>> configuration = XYZConfig() + distogram_intermediate_size (`int`, *optional*, defaults to 32): + Inner dimension for the distogram pair representation. Each + residue is projected to this size by separate left/right linear + layers; the outer product yields a ``distogram_intermediate_size²`` + feature per residue pair before the final classification layer. - >>> # Initializing a model from the x_y_z-7b style configuration - >>> model = XYZModel(configuration) + distogram_chunk_size (`int`, *optional*, defaults to 0): + If > 0, compute the distogram in chunks of this size along the + first sequence dimension to reduce peak memory. ``0`` disables + chunking (compute the full ``L×L`` matrix at once). - >>> # Accessing the model configuration - >>> configuration = model.config - ```""" + Note: + ``char_head_dim`` is intentionally omitted — the char stacks reuse the + token-level ``head_dim`` so that RoPE can be shared. - model_type = "x_y_z" - keys_to_ignore_at_inference = ["past_key_values"] - # Default tensor parallel plan for base model `XYZModel` + Example: + ```python + >>> from proxyz.models import XYZConfig, XYZForCausalLM + >>> config = XYZConfig() + >>> model = XYZForCausalLM(config) + ``` + """ + + model_type = "xyz" + keys_to_ignore_at_inference = ["past_key_values", "char_past_key_values", "offset_mapping"] + + # ---- Tensor-parallel / pipeline-parallel plans (inherited from Llama) ---- base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", + ".*.layers.*.self_attn.q_proj": "colwise", + ".*.layers.*.self_attn.k_proj": "colwise", + ".*.layers.*.self_attn.v_proj": "colwise", + ".*.layers.*.self_attn.o_proj": "rowwise", + ".*.layers.*.mlp.gate_proj": "colwise", + ".*.layers.*.mlp.up_proj": "colwise", + ".*.layers.*.mlp.down_proj": "rowwise", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), + ".*.layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + ".*.norm": (["hidden_states"], ["hidden_states"]), } vocab_size: int = 32000 @@ -69,7 +137,26 @@ class XYZConfig(PreTrainedConfig): mlp_bias: bool = False head_dim: int | None = None + char_hidden_size: int = 2048 + char_intermediate_size: int = 5504 + char_num_hidden_layers: int = 4 + char_num_attention_heads: int = 16 + char_num_key_value_heads: int | None = None + + use_char_position_ids: bool = False + + has_char_lm_head: bool = False + has_cle_lm_head: bool = False + has_distogram_lm_head: bool = False + + distogram_bins_num: int = 64 + distogram_intermediate_size: int = 32 + distogram_chunk_size: int = 0 # 0 = no chunking (compute full L×L at once) + def __post_init__(self, **kwargs): + # Default char KV heads to char query heads (MHA) when unspecified. + if self.char_num_key_value_heads is None: + self.char_num_key_value_heads = self.char_num_attention_heads if self.head_dim is None: self.head_dim = self.hidden_size // self.num_attention_heads if self.num_key_value_heads is None: @@ -79,11 +166,49 @@ def __post_init__(self, **kwargs): def validate_architecture(self): """Part of `@strict`-powered validation. Validates the architecture of the config.""" + # Ensure char_hidden_size is compatible with (char_num_attention_heads, head_dim). + if self.has_characterization: + if self.char_hidden_size != self.char_num_attention_heads * self.head_dim: + raise ValueError( + f"The char hidden size ({self.char_hidden_size}) must equal " + f"char_num_attention_heads ({self.char_num_attention_heads}) × head_dim ({self.head_dim})." + ) if self.hidden_size % self.num_attention_heads != 0: raise ValueError( f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention " f"heads ({self.num_attention_heads})." ) + @contextlib.contextmanager + def tokenization(self): + """Context manager that yields *self* with token-level config active. + + This is a no-op passthrough — the inherited Llama fields already hold + the token-level values. Provided for symmetry with ``characterization``. + """ + yield self + + @contextlib.contextmanager + def characterization(self): + """Context manager that temporarily swaps token-level fields with + char-level equivalents so that ``XYZDecoderLayer`` / ``XYZDecoderLayers`` + can be constructed or invoked with char-granularity dimensions. + + On exit, all fields are restored to their original (token-level) values. + """ + with attr( + self, + hidden_size=self.char_hidden_size, + intermediate_size=self.char_intermediate_size, + num_hidden_layers=self.char_num_hidden_layers, + num_attention_heads=self.char_num_attention_heads, + num_key_value_heads=self.char_num_key_value_heads, + ): + yield self + + @property + def has_characterization(self): + return any([self.has_char_lm_head, self.has_cle_lm_head, self.has_distogram_lm_head]) + __all__ = ["XYZConfig"] diff --git a/src/proxyz/models/modeling_xyz.py b/src/proxyz/models/modeling_xyz.py index 07f1d74..b245c24 100644 --- a/src/proxyz/models/modeling_xyz.py +++ b/src/proxyz/models/modeling_xyz.py @@ -4,8 +4,9 @@ # the file from the modular. If any change should be done, please apply the change to the # modular_xyz.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 - +import contextlib from collections.abc import Callable +from dataclasses import dataclass from typing import Optional import torch @@ -21,7 +22,7 @@ GenericForTokenClassification, GradientCheckpointingLayer, ) -from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, TokenClassifierOutput from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from transformers.processing_utils import Unpack @@ -30,6 +31,7 @@ from transformers.utils.output_capturing import capture_outputs from .configuration_xyz import XYZConfig +from .processing_xyz import XYZProcessor @use_kernel_forward_from_hub("RMSNorm") @@ -134,6 +136,39 @@ def forward(self, x): return down_proj +class XYZDistogram(nn.Module): + def __init__(self, config: XYZConfig): + super().__init__() + self.config = config + + self.left_proj = nn.Linear(config.char_hidden_size, config.distogram_intermediate_size) + self.right_proj = nn.Linear(config.char_hidden_size, config.distogram_intermediate_size) + self.out_proj = nn.Linear(config.distogram_intermediate_size**2, config.distogram_bins_num) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + # x = torch.einsum( + # "... i c, ... j d -> ... i j c d", self.left_proj(x), self.right_proj(x) + # ) + # x = x.view(*x.shape[:-2], -1) + # return self.out_proj(self.act_fn((x + x.transpose(-2, -3)) / 2)) # symmetrize + + # Decompose the bilinear form to avoid materializing (B, L, L, D²). + # Reshape out_proj weight: (bins, D²) -> (bins, D, D) + d = self.config.distogram_intermediate_size + x = torch.einsum( + "k c d, ... i c, ... j d -> ... i j k", + self.out_proj.weight.view(self.config.distogram_bins_num, d, d), + self.act_fn(self.left_proj(x)), + self.act_fn(self.right_proj(x)), + ) + # Symmetrize on (B, L, L, bins) and add bias + x = (x + x.transpose(-2, -3)) / 2 + if self.out_proj.bias is not None: + x = x + self.out_proj.bias + return x + + def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -315,13 +350,81 @@ def forward( return hidden_states +class XYZDecoderLayers(nn.Module): + def __init__(self, config: XYZConfig): + super().__init__() + self.config = config # FIX: AttributeError: 'XYZDecoderLayers' object has no attribute 'config' + + self.layers = nn.ModuleList( + [XYZDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = XYZRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + inputs_embeds: torch.FloatTensor, + causal_mask: torch.Tensor, + position_embeddings: torch.Tensor, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + """ + Runs a stack of [`XYZDecoderLayer`] modules followed by RMS normalization. + + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch, seq_len, hidden_size)`): + Input embeddings (output of a projection or embedding layer). + causal_mask (`torch.Tensor`): + Pre-computed causal attention mask. + position_embeddings (`torch.Tensor`): + RoPE cos/sin embeddings for the current positions. + position_ids (`torch.LongTensor`, *optional*): + Position indices. Inferred from *past_key_values* when *None*. + past_key_values (`Cache`, *optional*): + Key-value cache for incremental decoding. + use_cache (`bool`, *optional*): + Whether to populate and return the KV cache. + + Returns: + [`BaseModelOutputWithPast`] with the normalized hidden states and + (optionally) the updated KV cache. + """ + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + hidden_states = inputs_embeds + for decoder_layer in self.layers: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + @auto_docstring class XYZPreTrainedModel(PreTrainedModel): config: XYZConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["XYZDecoderLayer"] - _skip_keys_device_placement = ["past_key_values"] + _skip_keys_device_placement = ["past_key_values", "char_past_key_values"] _supports_flash_attn = True _supports_sdpa = True _supports_flex_attn = True @@ -335,22 +438,65 @@ class XYZPreTrainedModel(PreTrainedModel): @auto_docstring +@dataclass +class XYZModelOutputWithPast(BaseModelOutputWithPast): + """ + Output of [`XYZModel`]. + + Extends [`BaseModelOutputWithPast`] with character-granularity outputs from + the U-Net decoder branch. + + char_last_hidden_state (`torch.Tensor` of shape `(batch, char_seq_len, char_hidden_size)`, *optional*): + Character-level hidden states from the char decoder. + char_past_key_values (`tuple[Cache, Cache]`, *optional*): + Pair of KV caches — one for the char encoder, one for the char + decoder — used during incremental generation. + char_hidden_states (`tuple[torch.FloatTensor]`, *optional*): + Char-decoder hidden states at every layer (returned when + ``output_hidden_states=True``). Each element has shape + `(batch, char_seq_len, char_hidden_size)`. + char_attentions (`tuple[torch.FloatTensor]`, *optional*): + Char-decoder self-attention weights at every layer (returned when + ``output_attentions=True``). Each element has shape + `(batch, num_heads, char_seq_len, char_seq_len)`. + """ + + char_last_hidden_state: torch.FloatTensor | None = None + char_past_key_values: tuple[Cache, Cache] | None = None + char_hidden_states: tuple[torch.FloatTensor, ...] | None = None + char_attentions: tuple[torch.FloatTensor, ...] | None = None + + class XYZModel(XYZPreTrainedModel): def __init__(self, config: XYZConfig): 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( - [XYZDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] - ) - self.norm = XYZRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = XYZRotaryEmbedding(config=config) - self.gradient_checkpointing = False - # Initialize weights and apply final processing - self.post_init() + # Encoder (character granularity) + if config.has_characterization: + self.to_char_encoder = nn.Sequential(nn.Linear(config.hidden_size, config.char_hidden_size, bias=False)) + with self.config.characterization() as config: + self.char_encoder = XYZDecoderLayers(config) + self.from_char_encoder = nn.Sequential( + nn.Linear(config.char_hidden_size, config.hidden_size, bias=False), + ) + + # Trunk + with self.config.tokenization() as config: + self.trunk = XYZDecoderLayers(config) + + # Decoder (character granularity) + if config.has_characterization: + self.to_char_decoder = nn.Sequential(nn.Linear(config.hidden_size, config.char_hidden_size, bias=False)) + with self.config.characterization() as config: + self.char_decoder = XYZDecoderLayers(config) + + self.gradient_checkpointing = False @merge_with_config_defaults @capture_outputs @@ -362,9 +508,95 @@ def forward( position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, + char_input_ids: torch.LongTensor | None = None, + char_attention_mask: torch.Tensor | None = None, + char_position_ids: torch.LongTensor | None = None, + char_past_key_values: tuple[Cache, Cache] | None = None, + char_inputs_embeds: torch.FloatTensor | None = None, + repr_char_idx: torch.LongTensor | None = None, use_cache: bool | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPast: + ) -> XYZModelOutputWithPast: + """ + U-Net forward pass: char encoder → token trunk → char decoder. + + **Stage 1 — Char encoder.** Token embeddings are projected to + *char_hidden_size* and processed by the char encoder stack. + + **Stage 2 — Token trunk.** Char hidden states are gathered at + representative positions (``repr_char_idx``), projected back to + *hidden_size*, and added to the token embeddings. The result is + processed by the trunk transformer. + + **Stage 3 — Char decoder.** Trunk output is projected to + *char_hidden_size* and scattered back to char positions via + ``scatter_add`` (skip connection from the trunk). The char decoder + stack refines the representation with U-Net skip connections from the + encoder. + + Args: + input_ids (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + BPE token indices. Mutually exclusive with *inputs_embeds*. + attention_mask (`torch.Tensor` of shape `(batch, token_seq_len)`, *optional*): + Token-level attention mask (1 = attend, 0 = mask). + position_ids (`torch.LongTensor`, *optional*): + Token position indices. Inferred from *past_key_values* when *None*. + past_key_values (`Cache`, *optional*): + Trunk KV cache for incremental decoding. + inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed token embeddings. Mutually exclusive with *input_ids*. + char_input_ids (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Character-level token indices. Mutually exclusive with + *char_inputs_embeds*. + char_attention_mask (`torch.Tensor` of shape `(batch, char_seq_len)`, *optional*): + Character-level attention mask. + char_position_ids (`torch.LongTensor`, *optional*): + Character position indices. Inferred from *char_past_key_values* + when *None*. + char_past_key_values (`tuple[Cache, Cache]`, *optional*): + Pair of KV caches for the char encoder and char decoder. + char_inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed char embeddings. Mutually exclusive with *char_input_ids*. + repr_char_idx (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + Maps each BPE token to its representative character position in + the char sequence. Used to gather encoder output into the trunk + and scatter trunk output back to the decoder. + use_cache (`bool`, *optional*): + Whether to populate and return KV caches. + + Returns: + [`XYZModelOutputWithPast`]: + - **last_hidden_state** — token-level output from the trunk. + - **char_last_hidden_state** — char-level output from the char + decoder. + - **past_key_values** / **char_past_key_values** — updated caches. + """ + # character level. + if self.config.has_characterization: + if (char_input_ids is None) ^ (char_inputs_embeds is not None): + raise ValueError("You must specify exactly one of char_input_ids or char_inputs_embeds") + + if char_inputs_embeds is None: + char_inputs_embeds: torch.Tensor = self.to_char_encoder(self.embed_tokens(char_input_ids)) + + # if use_cache and char_past_key_values is None: + # with self.config.characterization() as config: + # char_past_key_values = ( + # DynamicCache(config=config), DynamicCache(config=config) + # ) + + if char_position_ids is None: + char_past_seen_tokens = ( + char_past_key_values[0].get_seq_length() if char_past_key_values is not None else 0 + ) + char_position_ids = ( + torch.arange(char_inputs_embeds.shape[1], device=char_inputs_embeds.device) + char_past_seen_tokens + ) + char_position_ids = char_position_ids.unsqueeze(0) + + char_position_embeddings = self.rotary_emb(char_inputs_embeds, position_ids=char_position_ids) + + # token level if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -372,28 +604,60 @@ def forward( inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) if use_cache and past_key_values is None: - past_key_values = DynamicCache(config=self.config) - - if position_ids is None: + with self.config.tokenization() as config: + past_key_values = DynamicCache(config=self.config) + + if self.config.has_characterization and self.config.use_char_position_ids: # gather from `char_position_ids` + position_ids = char_position_ids + if char_position_ids.shape[0] != repr_char_idx.shape[0]: + assert char_position_ids.shape[0] == 1 + position_ids = position_ids.expand(repr_char_idx.shape[0], char_position_ids.shape[1]) + position_ids = position_ids.gather(1, repr_char_idx) + elif position_ids is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.unsqueeze(0) + position_embeddings = self.rotary_emb(inputs_embeds, position_ids=position_ids) + + # encode + if self.config.has_characterization: + with self.config.characterization() as config: + char_causal_mask = create_causal_mask( + config=config, + inputs_embeds=char_inputs_embeds, + attention_mask=char_attention_mask, + past_key_values=char_past_key_values[0] if char_past_key_values else None, + position_ids=char_position_ids, + ) + encoder_outputs: BaseModelOutputWithPast = self.char_encoder( + inputs_embeds=char_inputs_embeds, + causal_mask=char_causal_mask, + position_embeddings=char_position_embeddings, + position_ids=char_position_ids, + past_key_values=char_past_key_values[0] if char_past_key_values else None, + use_cache=use_cache, + **kwargs, + ) + inputs_embeds = inputs_embeds + self.from_char_encoder( + encoder_outputs.last_hidden_state.gather( + 1, + repr_char_idx[..., None].expand(*repr_char_idx.shape, encoder_outputs.last_hidden_state.shape[2]), + ) + ) - causal_mask = create_causal_mask( - config=self.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - past_key_values=past_key_values, - position_ids=position_ids, - ) - - hidden_states = inputs_embeds - position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + # trunk + with self.config.tokenization() as config: + causal_mask = create_causal_mask( + config=config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) - for decoder_layer in self.layers[: self.config.num_hidden_layers]: - hidden_states = decoder_layer( - hidden_states, - attention_mask=causal_mask, + trunk_outputs: BaseModelOutputWithPast = self.trunk( + inputs_embeds=inputs_embeds, + causal_mask=causal_mask, position_embeddings=position_embeddings, position_ids=position_ids, past_key_values=past_key_values, @@ -401,25 +665,100 @@ def forward( **kwargs, ) - hidden_states = self.norm(hidden_states) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=past_key_values, + # decode + if self.config.has_characterization: + trunk_last_hidden_state = self.to_char_decoder(trunk_outputs.last_hidden_state) + char_inputs_embeds = encoder_outputs.last_hidden_state.scatter_add( # skip connection + 1, repr_char_idx[..., None].expand_as(trunk_last_hidden_state), trunk_last_hidden_state + ) + with self.config.characterization() as config: + decoder_outputs: BaseModelOutputWithPast = self.char_decoder( + inputs_embeds=char_inputs_embeds, + causal_mask=char_causal_mask, + position_embeddings=char_position_embeddings, + position_ids=char_position_ids, + past_key_values=char_past_key_values[1] if char_past_key_values else None, + use_cache=use_cache, + **kwargs, + ) + + return XYZModelOutputWithPast( + last_hidden_state=trunk_outputs.last_hidden_state, + past_key_values=trunk_outputs.past_key_values, + char_last_hidden_state=decoder_outputs.last_hidden_state, + char_past_key_values=(encoder_outputs.past_key_values, decoder_outputs.past_key_values), + ) + + return XYZModelOutputWithPast( + last_hidden_state=trunk_outputs.last_hidden_state, + past_key_values=trunk_outputs.past_key_values, ) +@auto_docstring +@dataclass +class XYZCausalLMOutputWithPast(CausalLMOutputWithPast): + """ + Output of [`XYZCausalLMOutputWithPast`]. + + Extends [`CausalLMOutputWithPast`] with character-granularity logits and + hidden states for the auxiliary next-character prediction head. + + offset_mapping (`torch.LongTensor` of shape `(batch, token_seq_len, 2)`, *optional*): + A list of (start_char, end_char) tuples for each token, mapping tokens back to + the original text characters. + char_logits (`torch.FloatTensor` of shape `(batch, char_seq_len, vocab_size)`, *optional*): + Next-character prediction logits from the char decoder LM head. + char_past_key_values (`tuple[Cache, Cache]`, *optional*): + Pair of KV caches — one for the char encoder, one for the char + decoder — used during incremental generation. + char_hidden_states (`tuple[torch.FloatTensor]`, *optional*): + Char-decoder hidden states at every layer (returned when + ``output_hidden_states=True``). + char_attentions (`tuple[torch.FloatTensor]`, *optional*): + Char-decoder self-attention weights at every layer (returned when + ``output_attentions=True``). + cle_logits (`torch.FloatTensor` of shape `(batch, char_seq_len, 26)`, *optional*): + CLE (Cα–Local–Environment) classification logits from the + auxiliary CLE head. + distogram_logits (`torch.FloatTensor` of shape `(batch, char_seq_len, char_seq_len, distogram_bins_num)`, *optional*): + Pairwise residue–residue distance-bin logits from the distogram + head. + """ + + offset_mapping: torch.LongTensor | None = None + char_logits: torch.FloatTensor | None = None + char_past_key_values: tuple[Cache, Cache] | None = None + char_hidden_states: tuple[torch.FloatTensor, ...] | None = None + char_attentions: tuple[torch.FloatTensor, ...] | None = None + cle_logits: torch.FloatTensor | None = None + distogram_logits: torch.FloatTensor | None = None + + @auto_docstring class XYZForCausalLM(XYZPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_gather_output"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} - def __init__(self, config): + def __init__(self, config: XYZConfig): super().__init__(config) self.model = XYZModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + if config.has_char_lm_head: + self.char_lm_head = nn.Linear(config.char_hidden_size, config.vocab_size, bias=False) + if config.has_cle_lm_head: + self.cle_lm_head = nn.Linear(config.char_hidden_size, 26, bias=False) + if config.has_distogram_lm_head: + # self.distogram_head = nn.Sequential( + # nn.Linear(config.char_hidden_size, config.char_hidden_size, bias=False), + # nn.GELU(), + # nn.Linear(config.char_hidden_size, config.distogram_bins_num, bias=False), + # ) + self.distogram_head = XYZDistogram(config) + # Initialize weights and apply final processing self.post_init() @@ -433,37 +772,98 @@ def forward( past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, - use_cache: bool | None = None, logits_to_keep: int | torch.Tensor = 0, + char_input_ids: torch.LongTensor | None = None, + char_attention_mask: torch.Tensor | None = None, + char_position_ids: torch.LongTensor | None = None, + char_past_key_values: tuple[Cache, Cache] | None = None, + char_inputs_embeds: torch.FloatTensor | None = None, + char_labels: torch.LongTensor | None = None, + char_logits_to_keep: int | torch.Tensor = 0, + cle_labels: torch.LongTensor | None = None, + distogram_labels: torch.LongTensor | None = None, + repr_char_idx: torch.LongTensor | None = None, + use_cache: bool | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> CausalLMOutputWithPast: - r""" - Example: + ) -> XYZCausalLMOutputWithPast: + """ + Causal language-modeling forward pass with dual prediction heads. - ```python - >>> from transformers import AutoTokenizer, XYZForCausalLM + Runs the U-Net backbone ([`XYZModel`]) and applies two LM heads: - >>> model = XYZForCausalLM.from_pretrained("meta-x_y_z/XYZ-2-7b-hf") - >>> tokenizer = AutoTokenizer.from_pretrained("meta-x_y_z/XYZ-2-7b-hf") + - **Next-token** — ``lm_head`` on trunk output → token logits + - **Next-character** — ``char_lm_head`` on char-decoder output → char logits - >>> prompt = "Hey, are you conscious? Can you talk to me?" - >>> inputs = tokenizer(prompt, return_tensors="pt") + Losses are summed when both *labels* and *char_labels* are provided. + + Args: + input_ids (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + BPE token indices. + attention_mask (`torch.Tensor`, *optional*): + Token-level attention mask. + position_ids (`torch.LongTensor`, *optional*): + Token position indices. + past_key_values (`Cache`, *optional*): + Trunk KV cache. + inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed token embeddings. + labels (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + Ground-truth token ids for next-token loss. Shifted internally. + logits_to_keep (`int` or `torch.Tensor`, *optional*, defaults to 0): + If an int, only compute logits for the last ``logits_to_keep`` + token positions (saves memory during training). If a tensor, + used as an explicit slice / boolean mask. + char_input_ids (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Character-level token indices. + char_attention_mask (`torch.Tensor`, *optional*): + Character-level attention mask. + char_position_ids (`torch.LongTensor`, *optional*): + Character position indices. + char_past_key_values (`tuple[Cache, Cache]`, *optional*): + Char encoder / decoder KV caches. + char_inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed char embeddings. + char_labels (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Ground-truth char ids for next-character loss. Shifted internally. + char_logits_to_keep (`int` or `torch.Tensor`, *optional*, defaults to 0): + Same as ``logits_to_keep`` but for the char-level heads + (``char_lm_head``, ``cle_lm_head``). + cle_labels (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Ground-truth CLE class ids for the CLE classification loss. + distogram_labels (`torch.LongTensor` of shape `(batch, char_seq_len, char_seq_len)`, *optional*): + Ground-truth distance-bin indices for the distogram head. + Each value is a bin index in ``[0, distogram_bins_num)``. + repr_char_idx (`torch.LongTensor`, *optional*): + Token → representative-char mapping. + use_cache (`bool`, *optional*): + Whether to populate and return KV caches. - >>> # Generate - >>> generate_ids = model.generate(inputs.input_ids, max_length=30) - >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." - ```""" - outputs: BaseModelOutputWithPast = self.model( + Returns: + [`XYZForCausalLMOutput`]: + - **loss** — combined next-token + next-character + aux CE loss. + - **logits** — next-token logits (used by `generate()`). + - **char_logits** — next-character logits. + - **cle_logits** — CLE classification logits. + - **distogram_logits** — pairwise distance-bin logits + of shape `(batch, char_seq_len, char_seq_len, distogram_bins_num)`. + """ + outputs: XYZModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, + char_input_ids=char_input_ids, + char_attention_mask=char_attention_mask, + char_position_ids=char_position_ids, + char_past_key_values=char_past_key_values, + char_inputs_embeds=char_inputs_embeds, + repr_char_idx=repr_char_idx, use_cache=use_cache, **kwargs, ) + # trunk hidden_states = outputs.last_hidden_state # Only compute necessary logits, and do not upcast them to float if we are not computing the loss slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep @@ -471,29 +871,231 @@ def forward( loss = None if labels is not None: - loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + loss = self.loss_function( + logits=logits, + labels=labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) + + # decoder + aux_loss, aux_logits = [], {} - return CausalLMOutputWithPast( + if self.config.has_characterization: + char_hidden_states = outputs.char_last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = ( + slice(-char_logits_to_keep, None) if isinstance(char_logits_to_keep, int) else char_logits_to_keep + ) + + if self.config.has_char_lm_head: + char_logits = self.char_lm_head(char_hidden_states[:, slice_indices, :]) + if char_labels is not None: + with self.num_items_in_batch( + kwargs, outputs.char_last_hidden_state.size(-2) / outputs.last_hidden_state.size(-2) + ): + aux_loss.append( + self.loss_function( + logits=char_logits, + labels=char_labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) + ) + aux_logits["char_logits"] = char_logits + if self.config.has_cle_lm_head: + cle_logits = self.cle_lm_head(char_hidden_states[:, slice_indices, :]) + if cle_labels is not None: + with self.num_items_in_batch( + kwargs, outputs.char_last_hidden_state.size(-2) / outputs.last_hidden_state.size(-2) + ): + aux_loss.append( + self.loss_function( + logits=cle_logits, + labels=cle_labels, + vocab_size=26, + shift_labels=cle_labels, # FIX: Fake the `ForCausalLMLoss` + **kwargs, + ) + ) + aux_logits["cle_logits"] = cle_logits + + if self.config.has_distogram_lm_head: + # pairwise outer sum: (B, L, H) -> (B, L, L, H) -> (B, L, L, num_bins) + # h = char_hidden_states + # pair_repr = h.unsqueeze(1) + h.unsqueeze(2) + distogram_logits = self.distogram_head(char_hidden_states[:, slice_indices, :]) + if distogram_labels is not None: + ignore_index = kwargs.get("ignore_index", -100) + # shift_distogram_labels = nn.functional.pad( + # distogram_labels, (0, 1, 0, 1), value=ignore_index + # ) + # shift_distogram_labels = shift_distogram_labels[..., 1:, 1:].contiguous() + with self.num_items_in_batch( + kwargs, outputs.char_last_hidden_state.size(-2) / outputs.last_hidden_state.size(-2) + ): + num_items_in_batch = kwargs.get("num_items_in_batch") + if num_items_in_batch is None: + aux_loss.append( + self.loss_function( + logits=distogram_logits, + labels=distogram_labels, + vocab_size=self.config.distogram_bins_num, + shift_labels=distogram_labels, # FIX: Fake the `ForCausalLMLoss` + **kwargs, + ) + ) + else: + eps = 1e-8 + distogram_loss = nn.functional.cross_entropy( + distogram_logits.reshape(-1, self.config.distogram_bins_num), + distogram_labels.reshape(-1), + ignore_index=ignore_index, + reduction="none", + ).view_as(distogram_labels) + distogram_loss = ( + distogram_loss.sum(-1) / ((distogram_labels != ignore_index).sum(-1) + eps) + ).sum() + # just in case users pass an int for num_items_in_batch, which could be the case for custom trainer + if torch.is_tensor(num_items_in_batch): + num_items_in_batch = num_items_in_batch.to(distogram_loss.device) + aux_loss.append(distogram_loss / num_items_in_batch) + # aux_loss.append( + # self.loss_function( + # logits=distogram_logits, + # labels=distogram_labels, + # vocab_size=self.config.distogram_bins_num, + # labels=distogram_labels, + # **kwargs, + # ) + # ) + aux_logits["distogram_logits"] = distogram_logits + + if loss is not None and aux_loss: + loss = loss + sum(aux_loss) + elif aux_loss: + loss = sum(aux_loss) + + return XYZCausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, + offset_mapping=kwargs.get("offset_mapping"), + char_past_key_values=outputs.char_past_key_values, + char_hidden_states=outputs.char_hidden_states, + char_attentions=outputs.char_attentions, + **aux_logits, ) + @contextlib.contextmanager + def num_items_in_batch(self, kwargs, amplify: float | None): + if "num_items_in_batch" in kwargs: + num_items_in_batch = kwargs["num_items_in_batch"] + if amplify is not None: + kwargs["num_items_in_batch"] = num_items_in_batch * amplify + else: + kwargs["num_items_in_batch"] = None + yield kwargs + if "num_items_in_batch" in kwargs: + kwargs["num_items_in_batch"] = num_items_in_batch + + def prepare_inputs_for_generation( + self, + input_ids: torch.LongTensor, + next_sequence_length: int | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + is_first_iteration: bool | None = False, + offset_mapping: torch.LongTensor | None = None, + processor: XYZProcessor | None = None, + **kwargs, + ): + model_inputs = super().prepare_inputs_for_generation( + input_ids=input_ids, + next_sequence_length=next_sequence_length, + past_key_values=past_key_values, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + is_first_iteration=is_first_iteration, + **kwargs, + ) + if not self.config.has_characterization: + return model_inputs + + # append offset_mapping + if not is_first_iteration: # not prefill stage + tokenized = processor.to_tokenization(processor.to_example(model_inputs["input_ids"]), generate=True) + tokenized["offset_mapping"] = tokenized["offset_mapping"].to(device=offset_mapping.device) + if kwargs.get("use_cache"): + tokenized["offset_mapping"] = tokenized["offset_mapping"] + offset_mapping[:, -1:, -1:] + offset_mapping = torch.cat((offset_mapping, tokenized["offset_mapping"]), dim=1) + else: + offset_mapping = tokenized["offset_mapping"] + # make char-level inputs + model_inputs.update(offset_mapping=offset_mapping) + model_inputs.update( + processor.to_characterization(processor.to_example(input_ids), model_inputs, generate=True) + ) + return model_inputs + + def _update_model_kwargs_for_generation( + self, + outputs: XYZCausalLMOutputWithPast, + model_kwargs: dict, + is_encoder_decoder: bool = False, + num_new_tokens: int = 1, + ) -> dict: + model_kwargs = super()._update_model_kwargs_for_generation( + outputs, model_kwargs, is_encoder_decoder=is_encoder_decoder, num_new_tokens=num_new_tokens + ) + if outputs.offset_mapping is not None: + model_kwargs["offset_mapping"] = outputs.offset_mapping + return model_kwargs + class XYZForSequenceClassification(GenericForSequenceClassification, XYZPreTrainedModel): pass +@auto_docstring +@dataclass +class XYZTokenClassifierOutput(TokenClassifierOutput): + """ + Output of [`XYZForTokenClassification`]. + + Extends [`TokenClassifierOutput`] with character-granularity logits and + hidden states for the auxiliary next-character prediction head. + + lm_output (`XYZCausalLMOutputWithPast`, *optional*): + Full causal-LM output from the U-Net backbone, including next-token + logits (``logits``) and next-character logits (``char_logits``) + cle_logits (`torch.FloatTensor` of shape `(batch, char_seq_len, 26)`, *optional*): + CLE (Cα–Local–Environment) classification logits from the + auxiliary CLE head. + distogram_logits (`torch.FloatTensor` of shape `(batch, char_seq_len, char_seq_len, distogram_bins_num)`, *optional*): + Pairwise residue–residue distance-bin logits from the distogram + head. + """ + + lm_output: XYZCausalLMOutputWithPast | None = None + cle_logits: torch.FloatTensor | None = None + distogram_logits: torch.FloatTensor | None = None + + class XYZForTokenClassification(GenericForTokenClassification, XYZPreTrainedModel): pass __all__ = [ "XYZPreTrainedModel", + "XYZModelOutputWithPast", "XYZModel", + "XYZCausalLMOutputWithPast", "XYZForCausalLM", "XYZForSequenceClassification", + "XYZTokenClassifierOutput", "XYZForTokenClassification", ] diff --git a/src/proxyz/models/modular_xyz.py b/src/proxyz/models/modular_xyz.py index 55cf125..b6418b4 100644 --- a/src/proxyz/models/modular_xyz.py +++ b/src/proxyz/models/modular_xyz.py @@ -1,6 +1,18 @@ +"""U-Net style language model. + +Three-stage architecture: + -> char encoder (character granularity) + -> token transformmer (token granularity) + -> char decoder (character granularity) + => predict next token +""" +import contextlib +from dataclasses import dataclass +import functools import random import torch +from torch import nn from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import ( @@ -49,15 +61,946 @@ @auto_docstring(checkpoint="bigict/ProXYZ") @strict class XYZConfig(LlamaConfig): + """ + Configuration for the XYZ U-Net style protein language model. + + A three-stage architecture: + + 1. **Char encoder** — transformer at character (residue) granularity + 2. **Token trunk** — transformer at BPE-token granularity + 3. **Char decoder** — transformer at character granularity with U-Net skip + connections from the encoder + + The char encoder/decoder share the same ``head_dim`` as the trunk so that + a single ``RotaryEmbedding`` instance can be used across all three stacks. + + This configuration extends [`LlamaConfig`]. All token-level fields + (``hidden_size``, ``intermediate_size``, ``num_hidden_layers``, etc.) are + inherited and control the *trunk* transformer. The char-level fields + below control the encoder and decoder. + + Constraint: ``char_hidden_size == char_num_attention_heads × head_dim``. + + char_hidden_size (`int`, *optional*, defaults to 2048): + Hidden size of the char encoder / decoder transformer. + char_intermediate_size (`int`, *optional*, defaults to 5504): + FFN intermediate size for char-level layers. + char_num_hidden_layers (`int`, *optional*, defaults to 4): + Number of transformer layers in the char encoder and char decoder. + char_num_attention_heads (`int`, *optional*, defaults to 16): + Number of attention heads for char-level self-attention. + char_num_key_value_heads (`int`, *optional*): + Number of KV heads for GQA at char level. Defaults to + ``char_num_attention_heads`` (i.e. multi-head attention) when + not specified. + + use_char_position_ids (`bool`, *optional*, defaults to `False`): + If ``True``, token-level ``position_ids`` are gathered from + ``char_position_ids`` via ``repr_char_idx`` instead of being + computed independently. This keeps RoPE aligned between the + char and token stacks. + + has_char_lm_head (`bool`, *optional*, defaults to `False`): + If ``True``, create a next-character prediction head + (``char_lm_head``) on top of the char decoder output. + + has_cle_lm_head (`bool`, *optional*, defaults to `False`): + If ``True``, create a CLE (Cα–Local–Environment) classification + head on top of the char decoder output. + + has_distogram_lm_head (`bool`, *optional*, defaults to `False`): + If ``True``, create a distogram prediction head on top of the + char decoder output. Predicts pairwise residue–residue distance + bins via an outer-sum MLP. + + distogram_bins_num (`int`, *optional*, defaults to 64): + Number of distance bins for the distogram head. + + distogram_intermediate_size (`int`, *optional*, defaults to 32): + Inner dimension for the distogram pair representation. Each + residue is projected to this size by separate left/right linear + layers; the outer product yields a ``distogram_intermediate_size²`` + feature per residue pair before the final classification layer. + + distogram_chunk_size (`int`, *optional*, defaults to 0): + If > 0, compute the distogram in chunks of this size along the + first sequence dimension to reduce peak memory. ``0`` disables + chunking (compute the full ``L×L`` matrix at once). + + Note: + ``char_head_dim`` is intentionally omitted — the char stacks reuse the + token-level ``head_dim`` so that RoPE can be shared. + + Example: + ```python + >>> from proxyz.models import XYZConfig, XYZForCausalLM + >>> config = XYZConfig() + >>> model = XYZForCausalLM(config) + ``` + """ + + model_type = "xyz" + keys_to_ignore_at_inference = [ + "past_key_values", "char_past_key_values", "offset_mapping" + ] + + # ---- Tensor-parallel / pipeline-parallel plans (inherited from Llama) ---- + base_model_tp_plan = { + ".*.layers.*.self_attn.q_proj": "colwise", + ".*.layers.*.self_attn.k_proj": "colwise", + ".*.layers.*.self_attn.v_proj": "colwise", + ".*.layers.*.self_attn.o_proj": "rowwise", + ".*.layers.*.mlp.gate_proj": "colwise", + ".*.layers.*.mlp.up_proj": "colwise", + ".*.layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + ".*.layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + ".*.norm": (["hidden_states"], ["hidden_states"]), + } + + char_hidden_size: int = 2048 + char_intermediate_size: int = 5504 + char_num_hidden_layers: int = 4 + char_num_attention_heads: int = 16 + char_num_key_value_heads: int | None = None + + use_char_position_ids: bool = False + + has_char_lm_head: bool = False + has_cle_lm_head: bool = False + has_distogram_lm_head: bool = False + + distogram_bins_num: int = 64 + distogram_intermediate_size: int = 32 + distogram_chunk_size: int = 0 # 0 = no chunking (compute full L×L at once) + + def __post_init__(self, **kwargs): + # Default char KV heads to char query heads (MHA) when unspecified. + if self.char_num_key_value_heads is None: + self.char_num_key_value_heads = self.char_num_attention_heads + super().__post_init__(**kwargs) + + def validate_architecture(self): + # Ensure char_hidden_size is compatible with (char_num_attention_heads, head_dim). + if self.has_characterization: + if self.char_hidden_size != self.char_num_attention_heads * self.head_dim: + raise ValueError( + f"The char hidden size ({self.char_hidden_size}) must equal " + f"char_num_attention_heads ({self.char_num_attention_heads}) × head_dim ({self.head_dim})." + ) + super().validate_architecture() + + @contextlib.contextmanager + def tokenization(self): + """Context manager that yields *self* with token-level config active. + + This is a no-op passthrough — the inherited Llama fields already hold + the token-level values. Provided for symmetry with ``characterization``. + """ + yield self + + @contextlib.contextmanager + def characterization(self): + """Context manager that temporarily swaps token-level fields with + char-level equivalents so that ``XYZDecoderLayer`` / ``XYZDecoderLayers`` + can be constructed or invoked with char-granularity dimensions. + + On exit, all fields are restored to their original (token-level) values. + """ + with attr( + self, + hidden_size=self.char_hidden_size, + intermediate_size=self.char_intermediate_size, + num_hidden_layers=self.char_num_hidden_layers, + num_attention_heads=self.char_num_attention_heads, + num_key_value_heads=self.char_num_key_value_heads, + ): + yield self + + @property + def has_characterization(self): + return any( + [self.has_char_lm_head, self.has_cle_lm_head, self.has_distogram_lm_head] + ) + + +class XYZRMSNorm(LlamaRMSNorm): pass -class XYZForCausalLM(LlamaForCausalLM): + + +class XYZRotaryEmbedding(LlamaRotaryEmbedding): pass +class XYZMLP(LlamaMLP): + pass + + +class XYZDistogram(nn.Module): + def __init__(self, config: XYZConfig): + super().__init__() + self.config = config + + self.left_proj = nn.Linear( + config.char_hidden_size, config.distogram_intermediate_size + ) + self.right_proj = nn.Linear( + config.char_hidden_size, config.distogram_intermediate_size + ) + self.out_proj = nn.Linear( + config.distogram_intermediate_size**2, config.distogram_bins_num + ) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + # x = torch.einsum( + # "... i c, ... j d -> ... i j c d", self.left_proj(x), self.right_proj(x) + # ) + # x = x.view(*x.shape[:-2], -1) + # return self.out_proj(self.act_fn((x + x.transpose(-2, -3)) / 2)) # symmetrize + + # Decompose the bilinear form to avoid materializing (B, L, L, D²). + # Reshape out_proj weight: (bins, D²) -> (bins, D, D) + d = self.config.distogram_intermediate_size + x = torch.einsum( + "k c d, ... i c, ... j d -> ... i j k", + self.out_proj.weight.view(self.config.distogram_bins_num, d, d), + self.act_fn(self.left_proj(x)), + self.act_fn(self.right_proj(x)), + ) + # Symmetrize on (B, L, L, bins) and add bias + x = (x + x.transpose(-2, -3)) / 2 + if self.out_proj.bias is not None: + x = x + self.out_proj.bias + return x + + +class XYZDecoderLayer(LlamaDecoderLayer): + pass + + +class XYZDecoderLayers(nn.Module): + def __init__(self, config: XYZConfig): + super().__init__() + self.config = config # FIX: AttributeError: 'XYZDecoderLayers' object has no attribute 'config' + + self.layers = nn.ModuleList( + [XYZDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = XYZRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + inputs_embeds: torch.FloatTensor, + causal_mask: torch.Tensor, + position_embeddings: torch.Tensor, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + """ + Runs a stack of [`XYZDecoderLayer`] modules followed by RMS normalization. + + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch, seq_len, hidden_size)`): + Input embeddings (output of a projection or embedding layer). + causal_mask (`torch.Tensor`): + Pre-computed causal attention mask. + position_embeddings (`torch.Tensor`): + RoPE cos/sin embeddings for the current positions. + position_ids (`torch.LongTensor`, *optional*): + Position indices. Inferred from *past_key_values* when *None*. + past_key_values (`Cache`, *optional*): + Key-value cache for incremental decoding. + use_cache (`bool`, *optional*): + Whether to populate and return the KV cache. + + Returns: + [`BaseModelOutputWithPast`] with the normalized hidden states and + (optionally) the updated KV cache. + """ + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + hidden_states = inputs_embeds + for decoder_layer in self.layers: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +class XYZPreTrainedModel(LlamaPreTrainedModel): + _skip_keys_device_placement = ["past_key_values", "char_past_key_values"] + + +@auto_docstring +@dataclass +class XYZModelOutputWithPast(BaseModelOutputWithPast): + """ + Output of [`XYZModel`]. + + Extends [`BaseModelOutputWithPast`] with character-granularity outputs from + the U-Net decoder branch. + + char_last_hidden_state (`torch.Tensor` of shape `(batch, char_seq_len, char_hidden_size)`, *optional*): + Character-level hidden states from the char decoder. + char_past_key_values (`tuple[Cache, Cache]`, *optional*): + Pair of KV caches — one for the char encoder, one for the char + decoder — used during incremental generation. + char_hidden_states (`tuple[torch.FloatTensor]`, *optional*): + Char-decoder hidden states at every layer (returned when + ``output_hidden_states=True``). Each element has shape + `(batch, char_seq_len, char_hidden_size)`. + char_attentions (`tuple[torch.FloatTensor]`, *optional*): + Char-decoder self-attention weights at every layer (returned when + ``output_attentions=True``). Each element has shape + `(batch, num_heads, char_seq_len, char_seq_len)`. + """ + + char_last_hidden_state: torch.FloatTensor | None = None + char_past_key_values: tuple[Cache, Cache] | None = None + char_hidden_states: tuple[torch.FloatTensor, ...] | None = None + char_attentions: tuple[torch.FloatTensor, ...] | None = None + + +class XYZModel(XYZPreTrainedModel): + def __init__(self, config: XYZConfig): + 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.rotary_emb = XYZRotaryEmbedding(config=config) + + # Encoder (character granularity) + if config.has_characterization: + self.to_char_encoder = nn.Sequential( + nn.Linear(config.hidden_size, config.char_hidden_size, bias=False) + ) + with self.config.characterization() as config: + self.char_encoder = XYZDecoderLayers(config) + self.from_char_encoder = nn.Sequential( + nn.Linear(config.char_hidden_size, config.hidden_size, bias=False), + ) + + # Trunk + with self.config.tokenization() as config: + self.trunk = XYZDecoderLayers(config) + + # Decoder (character granularity) + if config.has_characterization: + self.to_char_decoder = nn.Sequential( + nn.Linear(config.hidden_size, config.char_hidden_size, bias=False) + ) + with self.config.characterization() as config: + self.char_decoder = XYZDecoderLayers(config) + + self.gradient_checkpointing = False + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + char_input_ids: torch.LongTensor | None = None, + char_attention_mask: torch.Tensor | None = None, + char_position_ids: torch.LongTensor | None = None, + char_past_key_values: tuple[Cache, Cache] | None = None, + char_inputs_embeds: torch.FloatTensor | None = None, + repr_char_idx: torch.LongTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> XYZModelOutputWithPast: + """ + U-Net forward pass: char encoder → token trunk → char decoder. + + **Stage 1 — Char encoder.** Token embeddings are projected to + *char_hidden_size* and processed by the char encoder stack. + + **Stage 2 — Token trunk.** Char hidden states are gathered at + representative positions (``repr_char_idx``), projected back to + *hidden_size*, and added to the token embeddings. The result is + processed by the trunk transformer. + + **Stage 3 — Char decoder.** Trunk output is projected to + *char_hidden_size* and scattered back to char positions via + ``scatter_add`` (skip connection from the trunk). The char decoder + stack refines the representation with U-Net skip connections from the + encoder. + + Args: + input_ids (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + BPE token indices. Mutually exclusive with *inputs_embeds*. + attention_mask (`torch.Tensor` of shape `(batch, token_seq_len)`, *optional*): + Token-level attention mask (1 = attend, 0 = mask). + position_ids (`torch.LongTensor`, *optional*): + Token position indices. Inferred from *past_key_values* when *None*. + past_key_values (`Cache`, *optional*): + Trunk KV cache for incremental decoding. + inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed token embeddings. Mutually exclusive with *input_ids*. + char_input_ids (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Character-level token indices. Mutually exclusive with + *char_inputs_embeds*. + char_attention_mask (`torch.Tensor` of shape `(batch, char_seq_len)`, *optional*): + Character-level attention mask. + char_position_ids (`torch.LongTensor`, *optional*): + Character position indices. Inferred from *char_past_key_values* + when *None*. + char_past_key_values (`tuple[Cache, Cache]`, *optional*): + Pair of KV caches for the char encoder and char decoder. + char_inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed char embeddings. Mutually exclusive with *char_input_ids*. + repr_char_idx (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + Maps each BPE token to its representative character position in + the char sequence. Used to gather encoder output into the trunk + and scatter trunk output back to the decoder. + use_cache (`bool`, *optional*): + Whether to populate and return KV caches. + + Returns: + [`XYZModelOutputWithPast`]: + - **last_hidden_state** — token-level output from the trunk. + - **char_last_hidden_state** — char-level output from the char + decoder. + - **past_key_values** / **char_past_key_values** — updated caches. + """ + # character level. + if self.config.has_characterization: + if (char_input_ids is None) ^ (char_inputs_embeds is not None): + raise ValueError("You must specify exactly one of char_input_ids or char_inputs_embeds") + + if char_inputs_embeds is None: + char_inputs_embeds: torch.Tensor = self.to_char_encoder(self.embed_tokens(char_input_ids)) + + # if use_cache and char_past_key_values is None: + # with self.config.characterization() as config: + # char_past_key_values = ( + # DynamicCache(config=config), DynamicCache(config=config) + # ) + + if char_position_ids is None: + char_past_seen_tokens = char_past_key_values[0].get_seq_length() if char_past_key_values is not None else 0 + char_position_ids = torch.arange(char_inputs_embeds.shape[1], device=char_inputs_embeds.device) + char_past_seen_tokens + char_position_ids = char_position_ids.unsqueeze(0) + + char_position_embeddings = self.rotary_emb(char_inputs_embeds, position_ids=char_position_ids) + + # token level + 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: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + with self.config.tokenization() as config: + past_key_values = DynamicCache(config=self.config) + + if self.config.has_characterization and self.config.use_char_position_ids: # gather from `char_position_ids` + position_ids = char_position_ids + if char_position_ids.shape[0] != repr_char_idx.shape[0]: + assert char_position_ids.shape[0] == 1 + position_ids = position_ids.expand( + repr_char_idx.shape[0], char_position_ids.shape[1] + ) + position_ids = position_ids.gather(1, repr_char_idx) + elif position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + position_embeddings = self.rotary_emb(inputs_embeds, position_ids=position_ids) + + # encode + if self.config.has_characterization: + with self.config.characterization() as config: + char_causal_mask = create_causal_mask( + config=config, + inputs_embeds=char_inputs_embeds, + attention_mask=char_attention_mask, + past_key_values=char_past_key_values[0] if char_past_key_values else None, + position_ids=char_position_ids, + ) + encoder_outputs: BaseModelOutputWithPast = self.char_encoder( + inputs_embeds=char_inputs_embeds, + causal_mask=char_causal_mask, + position_embeddings=char_position_embeddings, + position_ids=char_position_ids, + past_key_values=char_past_key_values[0] if char_past_key_values else None, + use_cache=use_cache, + **kwargs, + ) + inputs_embeds = inputs_embeds + self.from_char_encoder( + encoder_outputs.last_hidden_state.gather( + 1, repr_char_idx[..., None].expand( + *repr_char_idx.shape, encoder_outputs.last_hidden_state.shape[2] + ) + ) + ) + + # trunk + with self.config.tokenization() as config: + causal_mask = create_causal_mask( + config=config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + trunk_outputs: BaseModelOutputWithPast = self.trunk( + inputs_embeds=inputs_embeds, + causal_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + # decode + if self.config.has_characterization: + trunk_last_hidden_state = self.to_char_decoder(trunk_outputs.last_hidden_state) + char_inputs_embeds = encoder_outputs.last_hidden_state.scatter_add( # skip connection + 1, repr_char_idx[..., None].expand_as(trunk_last_hidden_state), trunk_last_hidden_state + ) + with self.config.characterization() as config: + decoder_outputs: BaseModelOutputWithPast = self.char_decoder( + inputs_embeds=char_inputs_embeds, + causal_mask=char_causal_mask, + position_embeddings=char_position_embeddings, + position_ids=char_position_ids, + past_key_values=char_past_key_values[1] if char_past_key_values else None, + use_cache=use_cache, + **kwargs, + ) + + return XYZModelOutputWithPast( + last_hidden_state=trunk_outputs.last_hidden_state, + past_key_values=trunk_outputs.past_key_values, + char_last_hidden_state=decoder_outputs.last_hidden_state, + char_past_key_values=( + encoder_outputs.past_key_values, decoder_outputs.past_key_values + ) + ) + + return XYZModelOutputWithPast( + last_hidden_state=trunk_outputs.last_hidden_state, + past_key_values=trunk_outputs.past_key_values, + ) + + +@auto_docstring +@dataclass +class XYZCausalLMOutputWithPast(CausalLMOutputWithPast): + """ + Output of [`XYZCausalLMOutputWithPast`]. + + Extends [`CausalLMOutputWithPast`] with character-granularity logits and + hidden states for the auxiliary next-character prediction head. + + offset_mapping (`torch.LongTensor` of shape `(batch, token_seq_len, 2)`, *optional*): + A list of (start_char, end_char) tuples for each token, mapping tokens back to + the original text characters. + char_logits (`torch.FloatTensor` of shape `(batch, char_seq_len, vocab_size)`, *optional*): + Next-character prediction logits from the char decoder LM head. + char_past_key_values (`tuple[Cache, Cache]`, *optional*): + Pair of KV caches — one for the char encoder, one for the char + decoder — used during incremental generation. + char_hidden_states (`tuple[torch.FloatTensor]`, *optional*): + Char-decoder hidden states at every layer (returned when + ``output_hidden_states=True``). + char_attentions (`tuple[torch.FloatTensor]`, *optional*): + Char-decoder self-attention weights at every layer (returned when + ``output_attentions=True``). + cle_logits (`torch.FloatTensor` of shape `(batch, char_seq_len, 26)`, *optional*): + CLE (Cα–Local–Environment) classification logits from the + auxiliary CLE head. + distogram_logits (`torch.FloatTensor` of shape `(batch, char_seq_len, char_seq_len, distogram_bins_num)`, *optional*): + Pairwise residue–residue distance-bin logits from the distogram + head. + """ + + offset_mapping: torch.LongTensor | None = None + char_logits: torch.FloatTensor | None = None + char_past_key_values: tuple[Cache, Cache] | None = None + char_hidden_states: tuple[torch.FloatTensor, ...] | None = None + char_attentions: tuple[torch.FloatTensor, ...] | None = None + cle_logits: torch.FloatTensor | None = None + distogram_logits: torch.FloatTensor | None = None + + +class XYZForCausalLM(LlamaForCausalLM): + def __init__(self, config: XYZConfig): + super().__init__(config) + + if config.has_char_lm_head: + self.char_lm_head = nn.Linear( + config.char_hidden_size, config.vocab_size, bias=False + ) + if config.has_cle_lm_head: + self.cle_lm_head = nn.Linear(config.char_hidden_size, 26, bias=False) + if config.has_distogram_lm_head: + # self.distogram_head = nn.Sequential( + # nn.Linear(config.char_hidden_size, config.char_hidden_size, bias=False), + # nn.GELU(), + # nn.Linear(config.char_hidden_size, config.distogram_bins_num, bias=False), + # ) + self.distogram_head = XYZDistogram(config) + + @contextlib.contextmanager + def num_items_in_batch(self, kwargs, amplify: float | None): + if "num_items_in_batch" in kwargs: + num_items_in_batch = kwargs["num_items_in_batch"] + if amplify is not None: + kwargs["num_items_in_batch"] = num_items_in_batch * amplify + else: + kwargs["num_items_in_batch"] = None + yield kwargs + if "num_items_in_batch" in kwargs: + kwargs["num_items_in_batch"] = num_items_in_batch + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + char_input_ids: torch.LongTensor | None = None, + char_attention_mask: torch.Tensor | None = None, + char_position_ids: torch.LongTensor | None = None, + char_past_key_values: tuple[Cache, Cache] | None = None, + char_inputs_embeds: torch.FloatTensor | None = None, + char_labels: torch.LongTensor | None = None, + char_logits_to_keep: int | torch.Tensor = 0, + cle_labels: torch.LongTensor | None = None, + distogram_labels: torch.LongTensor | None = None, + repr_char_idx: torch.LongTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> XYZCausalLMOutputWithPast: + """ + Causal language-modeling forward pass with dual prediction heads. + + Runs the U-Net backbone ([`XYZModel`]) and applies two LM heads: + + - **Next-token** — ``lm_head`` on trunk output → token logits + - **Next-character** — ``char_lm_head`` on char-decoder output → char logits + + Losses are summed when both *labels* and *char_labels* are provided. + + Args: + input_ids (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + BPE token indices. + attention_mask (`torch.Tensor`, *optional*): + Token-level attention mask. + position_ids (`torch.LongTensor`, *optional*): + Token position indices. + past_key_values (`Cache`, *optional*): + Trunk KV cache. + inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed token embeddings. + labels (`torch.LongTensor` of shape `(batch, token_seq_len)`, *optional*): + Ground-truth token ids for next-token loss. Shifted internally. + logits_to_keep (`int` or `torch.Tensor`, *optional*, defaults to 0): + If an int, only compute logits for the last ``logits_to_keep`` + token positions (saves memory during training). If a tensor, + used as an explicit slice / boolean mask. + char_input_ids (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Character-level token indices. + char_attention_mask (`torch.Tensor`, *optional*): + Character-level attention mask. + char_position_ids (`torch.LongTensor`, *optional*): + Character position indices. + char_past_key_values (`tuple[Cache, Cache]`, *optional*): + Char encoder / decoder KV caches. + char_inputs_embeds (`torch.FloatTensor`, *optional*): + Pre-computed char embeddings. + char_labels (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Ground-truth char ids for next-character loss. Shifted internally. + char_logits_to_keep (`int` or `torch.Tensor`, *optional*, defaults to 0): + Same as ``logits_to_keep`` but for the char-level heads + (``char_lm_head``, ``cle_lm_head``). + cle_labels (`torch.LongTensor` of shape `(batch, char_seq_len)`, *optional*): + Ground-truth CLE class ids for the CLE classification loss. + distogram_labels (`torch.LongTensor` of shape `(batch, char_seq_len, char_seq_len)`, *optional*): + Ground-truth distance-bin indices for the distogram head. + Each value is a bin index in ``[0, distogram_bins_num)``. + repr_char_idx (`torch.LongTensor`, *optional*): + Token → representative-char mapping. + use_cache (`bool`, *optional*): + Whether to populate and return KV caches. + + Returns: + [`XYZForCausalLMOutput`]: + - **loss** — combined next-token + next-character + aux CE loss. + - **logits** — next-token logits (used by `generate()`). + - **char_logits** — next-character logits. + - **cle_logits** — CLE classification logits. + - **distogram_logits** — pairwise distance-bin logits + of shape `(batch, char_seq_len, char_seq_len, distogram_bins_num)`. + """ + outputs: XYZModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + char_input_ids=char_input_ids, + char_attention_mask=char_attention_mask, + char_position_ids=char_position_ids, + char_past_key_values=char_past_key_values, + char_inputs_embeds=char_inputs_embeds, + repr_char_idx=repr_char_idx, + use_cache=use_cache, + **kwargs, + ) + + # trunk + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + 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, :]) + + loss = None + if labels is not None: + loss = self.loss_function( + logits=logits, + labels=labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) + + # decoder + aux_loss, aux_logits = [], {} + + if self.config.has_characterization: + char_hidden_states = outputs.char_last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-char_logits_to_keep, None) if isinstance( + char_logits_to_keep, int + ) else char_logits_to_keep + + if self.config.has_char_lm_head: + char_logits = self.char_lm_head(char_hidden_states[:, slice_indices, :]) + if char_labels is not None: + with self.num_items_in_batch( + kwargs, outputs.char_last_hidden_state.size(-2) / outputs.last_hidden_state.size(-2) + ): + aux_loss.append( + self.loss_function( + logits=char_logits, + labels=char_labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) + ) + aux_logits["char_logits"] = char_logits + if self.config.has_cle_lm_head: + cle_logits = self.cle_lm_head(char_hidden_states[:, slice_indices, :]) + if cle_labels is not None: + with self.num_items_in_batch( + kwargs, outputs.char_last_hidden_state.size(-2) / outputs.last_hidden_state.size(-2) + ): + aux_loss.append( + self.loss_function( + logits=cle_logits, + labels=cle_labels, + vocab_size=26, + shift_labels=cle_labels, # FIX: Fake the `ForCausalLMLoss` + **kwargs, + ) + ) + aux_logits["cle_logits"] = cle_logits + + if self.config.has_distogram_lm_head: + # pairwise outer sum: (B, L, H) -> (B, L, L, H) -> (B, L, L, num_bins) + # h = char_hidden_states + # pair_repr = h.unsqueeze(1) + h.unsqueeze(2) + distogram_logits = self.distogram_head(char_hidden_states[:, slice_indices, :]) + if distogram_labels is not None: + ignore_index = kwargs.get("ignore_index", -100) + # shift_distogram_labels = nn.functional.pad( + # distogram_labels, (0, 1, 0, 1), value=ignore_index + # ) + # shift_distogram_labels = shift_distogram_labels[..., 1:, 1:].contiguous() + with self.num_items_in_batch( + kwargs, outputs.char_last_hidden_state.size(-2) / outputs.last_hidden_state.size(-2) + ): + num_items_in_batch = kwargs.get("num_items_in_batch") + if num_items_in_batch is None: + aux_loss.append( + self.loss_function( + logits=distogram_logits, + labels=distogram_labels, + vocab_size=self.config.distogram_bins_num, + shift_labels=distogram_labels, # FIX: Fake the `ForCausalLMLoss` + **kwargs, + ) + ) + else: + eps = 1e-8 + distogram_loss = nn.functional.cross_entropy( + distogram_logits.reshape(-1, self.config.distogram_bins_num), + distogram_labels.reshape(-1), + ignore_index=ignore_index, + reduction="none", + ).view_as(distogram_labels) + distogram_loss = ( + distogram_loss.sum(-1) / ((distogram_labels != ignore_index).sum(-1) + eps) + ).sum() + # just in case users pass an int for num_items_in_batch, which could be the case for custom trainer + if torch.is_tensor(num_items_in_batch): + num_items_in_batch = num_items_in_batch.to(distogram_loss.device) + aux_loss.append(distogram_loss / num_items_in_batch) + # aux_loss.append( + # self.loss_function( + # logits=distogram_logits, + # labels=distogram_labels, + # vocab_size=self.config.distogram_bins_num, + # labels=distogram_labels, + # **kwargs, + # ) + # ) + aux_logits["distogram_logits"] = distogram_logits + + if loss is not None and aux_loss: + loss = loss + sum(aux_loss) + elif aux_loss: + loss = sum(aux_loss) + + return XYZCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + offset_mapping=kwargs.get("offset_mapping"), + char_past_key_values=outputs.char_past_key_values, + char_hidden_states=outputs.char_hidden_states, + char_attentions=outputs.char_attentions, + **aux_logits, + ) + + def prepare_inputs_for_generation( + self, + input_ids: torch.LongTensor, + next_sequence_length: int | None = None, + past_key_values: Cache | None = None, + attention_mask: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + is_first_iteration: bool | None = False, + offset_mapping: torch.LongTensor | None = None, + processor: XYZProcessor | None = None, + **kwargs, + ): + model_inputs = super().prepare_inputs_for_generation( + input_ids=input_ids, + next_sequence_length=next_sequence_length, + past_key_values=past_key_values, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + is_first_iteration=is_first_iteration, + **kwargs, + ) + if not self.config.has_characterization: + return model_inputs + + # append offset_mapping + if not is_first_iteration: # not prefill stage + tokenized = processor.to_tokenization( + processor.to_example(model_inputs["input_ids"]), generate=True + ) + tokenized["offset_mapping"] = tokenized["offset_mapping"].to(device=offset_mapping.device) + if kwargs.get("use_cache"): + tokenized["offset_mapping"] = tokenized["offset_mapping"] + offset_mapping[:, -1:, -1:] + offset_mapping = torch.cat((offset_mapping, tokenized["offset_mapping"]), dim=1) + else: + offset_mapping = tokenized["offset_mapping"] + # make char-level inputs + model_inputs.update(offset_mapping=offset_mapping) + model_inputs.update( + processor.to_characterization( + processor.to_example(input_ids), model_inputs, generate=True + ) + ) + return model_inputs + + def _update_model_kwargs_for_generation( + self, + outputs: XYZCausalLMOutputWithPast, + model_kwargs: dict, + is_encoder_decoder: bool = False, + num_new_tokens: int = 1, + ) -> dict: + model_kwargs = super()._update_model_kwargs_for_generation( + outputs, model_kwargs, is_encoder_decoder=is_encoder_decoder, num_new_tokens=num_new_tokens + ) + if outputs.offset_mapping is not None: + model_kwargs["offset_mapping"] = outputs.offset_mapping + return model_kwargs + + class XYZForSequenceClassification(LlamaForSequenceClassification): pass +@auto_docstring +@dataclass +class XYZTokenClassifierOutput(TokenClassifierOutput): + """ + Output of [`XYZForTokenClassification`]. + + Extends [`TokenClassifierOutput`] with character-granularity logits and + hidden states for the auxiliary next-character prediction head. + + lm_output (`XYZCausalLMOutputWithPast`, *optional*): + Full causal-LM output from the U-Net backbone, including next-token + logits (``logits``) and next-character logits (``char_logits``) + cle_logits (`torch.FloatTensor` of shape `(batch, char_seq_len, 26)`, *optional*): + CLE (Cα–Local–Environment) classification logits from the + auxiliary CLE head. + distogram_logits (`torch.FloatTensor` of shape `(batch, char_seq_len, char_seq_len, distogram_bins_num)`, *optional*): + Pairwise residue–residue distance-bin logits from the distogram + head. + """ + + lm_output: XYZCausalLMOutputWithPast | None = None + cle_logits: torch.FloatTensor | None = None + distogram_logits: torch.FloatTensor | None = None + + class XYZForTokenClassification(LlamaForTokenClassification): pass @@ -82,7 +1025,9 @@ class XYZProcessor(ProcessorMixin): FIM_MIDDLE = "" FIM_TOKENS = [FIM_PREFIX, FIM_SUFFIX, FIM_MIDDLE] - def __init__(self, tokenizer, text_column: str = "text", **kwargs) -> None: + def __init__( + self, tokenizer, text_column: str = "text", features: list[str] = None, **kwargs + ) -> None: tokenizer.add_special_tokens({"additional_special_tokens": self.FIM_TOKENS}) logger.info( f"Added FIM tokens: {self.FIM_TOKENS} to tokenizer. " @@ -97,13 +1042,21 @@ def __init__(self, tokenizer, text_column: str = "text", **kwargs) -> None: super().__init__(tokenizer, **kwargs) self.text_column = text_column + self.features = features self.ignore_index = kwargs.get("ignore_index", -100) + self.distogram_bins = torch.linspace( + kwargs.get("distogram_bins_min", 2.3125), + kwargs.get("distogram_bins_max", 21.6875), + steps=kwargs.get("distogram_bins_num", 64) - 1, + ) + def __call__( self, examples: dict, *, bpe_dropout: float | None = None, + char_apply: bool = False, fim_apply: bool = False, fim_spm_rate: float = 0.5, fim_sft_style: bool = False, @@ -123,6 +1076,43 @@ def __call__( fim_sft_style=fim_sft_style, generate=generate, ) + if char_apply: + tokenized.update( + self.to_characterization( + examples, tokenized, + fim_apply=fim_apply, + fim_sft_style=fim_sft_style, + generate=generate, + ) + ) + # del tokenized["offset_mapping"] + + # copy features + if self.features is not None: + + def feat_collate(feat, is_label: bool = False): + if isinstance(feat, tuple): + return tuple(feat_collate(feat[t]) for t in range(len(feat))) + max_len = max(feat[k].shape[0] for k in range(len(feat))) + for k in range(len(feat)): + pad = feat[k].new_full( + (max_len - feat[k].shape[0], *feat[k].shape[1:]), + self.ignore_index if is_label else 0, + ) + feat[k] = torch.cat((feat[k], pad)) + return torch.stack(feat) + + for column in self.features: + if column in examples: + tokenized[column] = feat_collate( + examples[column], "labels" in column + ) + + if "distogram_labels" in self.features: + if "distogram_labels" in tokenized: + tokenized["distogram_labels"] = self.to_distogram( + *tokenized["distogram_labels"] + ) return tokenized @@ -148,6 +1138,66 @@ def to_tokenization( ) return tokenized + def to_characterization( + self, + examples: dict, + tokenized: dict, + bpe_dropout: float = 1.0, + fim_apply: bool = False, + fim_sft_style: bool = False, + generate: bool = False, + ) -> dict: + assert "offset_mapping" in tokenized and "attention_mask" in tokenized + + characterized = {} + for k, v in self._tokenize_with_dropout( + examples, + bpe_dropout=bpe_dropout, + prefix="char_", + fim_apply=fim_apply, + fim_sft_style=fim_sft_style, + generate=generate, + ).items(): + if torch.is_tensor(v): + v = v.to(device=tokenized["offset_mapping"].device) + characterized[k] = v + + # Align characters with tokenized *text* + # NOTE: special token is treat as one char + # NOTE: pad offset_mapping with the maxinum offset + char_offset_mapping = characterized["char_offset_mapping"].where( + characterized["char_attention_mask"][..., None] > 0, + characterized["char_offset_mapping"].max() + 1 + ) + offset_mapping = tokenized["offset_mapping"].where( + tokenized["attention_mask"][..., None] > 0, + tokenized["offset_mapping"].max() + 1 + ) + token_to_char_map = torch.searchsorted( + char_offset_mapping.transpose(-1, -2).contiguous(), + offset_mapping.transpose(-1, -2).contiguous(), + right=True, + ) + token_to_char_cnt = token_to_char_map[:, 1, :] - token_to_char_map[:, 0, :] + 1 + characterized["repr_char_idx"] = token_to_char_cnt.cumsum(1) - 1 + characterized["repr_char_idx"] = characterized["repr_char_idx"].where( + tokenized["attention_mask"] > 0, 0 + ) + # characterized["char_to_token_mask"] = torch.arange(self.max_chars_within_token).view( + # 1, 1, self.max_chars_within_token + # ) + # characterized["char_to_token_mask"] = ( + # characterized["char_to_token_mask"] < token_to_char_cnt[:, :, None] + # ).where(tokenized["attention_mask"][:, :, None] > 0, 0) + # characterized["char_to_token_idx"] = torch.searchsorted( + # tokenized["offset_mapping"][:, :, 0].contiguous(), + # characterized["char_offset_mapping"][:, :, 0].contiguous(), + # right=True, + # ) - 1 + + del characterized["char_offset_mapping"] + return characterized + def to_example(self, input_ids: torch.LongTensor) -> dict: text = self.tokenizer.decode(input_ids.tolist(), skip_special_tokens=False) return {self.text_column: [t.replace(" ", "") for t in text]} @@ -206,6 +1256,38 @@ def apply_fim(self, examples: dict, spm_rate: float = 0.5) -> dict: examples[self.text_column][idx] = ( first_tag + first + second_tag + second + self.FIM_MIDDLE + middle ) + if self.features is not None: + def feat_fim(feat, is_label: bool = False): + pad = feat.new_full( + (1, *feat.shape[1:]), self.ignore_index if is_label else 0 + ) + if is_spm: + return torch.cat( + ( + pad, feat[cut2:, ...], + pad, feat[:cut1, ...], + pad, feat[cut1:cut2, ...], + ) + ) + return torch.cat( + ( + pad, feat[:cut1, ...], + pad, feat[cut2:, ...], + pad, feat[cut1:cut2, ...], + ) + ) + + for column in self.features: + if column in examples: + if isinstance(examples[column], tuple): + for t in range(len(examples[column])): + examples[column][t][idx] = feat_fim( + examples[column][t][idx] + ) + else: + examples[column][idx] = feat_fim( + examples[column][idx], "labels" in column + ) return examples def apply_crop(self, examples: dict, max_length: int | None = None) -> dict: @@ -215,6 +1297,19 @@ def apply_crop(self, examples: dict, max_length: int | None = None) -> dict: cut = random.randint(0, n - max_length) text = text[cut:cut + max_length] examples[self.text_column][idx] = text + if self.features is not None: + def feat_crop(feat): + return feat[cut:cut + max_length, ...] + + for column in self.features: + if column in examples: + if isinstance(examples[column], tuple): + for t in range(len(examples[column])): + examples[column][t][idx] = feat_crop( + examples[column][t][idx] + ) + else: + examples[column][idx] = feat_crop(examples[column][idx]) return examples def apply_wrap(self, examples: dict, add_eos_token: bool = True) -> dict: @@ -223,14 +1318,45 @@ def apply_wrap(self, examples: dict, add_eos_token: bool = True) -> dict: if add_eos_token: text = f"{text}{self.tokenizer.eos_token}" examples[self.text_column][idx] = text + if self.features is not None: + def feat_wrap(feat, is_label: bool = False): + pad = feat.new_full( + (1, *feat.shape[1:]), self.ignore_index if is_label else 0 + ) + if add_eos_token: + return torch.cat((pad, feat, pad)) + return torch.cat((pad, feat)) + + for column in self.features: + if column in examples: + if isinstance(examples[column], tuple): + for t in range(len(examples[column])): + examples[column][t][idx] = feat_wrap( + examples[column][t][idx] + ) + else: + examples[column][idx] = feat_wrap( + examples[column][idx], "labels" in column + ) return examples + def to_distogram(self, pseudo_beta, pseudo_beta_mask): + distogram_labels = torch.cdist(pseudo_beta, pseudo_beta) + distogram_labels = (distogram_labels[..., None] > self.distogram_bins).sum(-1) + return distogram_labels.where( + pseudo_beta_mask[..., :, None] * pseudo_beta_mask[..., None, :], + self.ignore_index + ) + __all__ = [ "XYZPreTrainedModel", + "XYZModelOutputWithPast", "XYZModel", + "XYZCausalLMOutputWithPast", "XYZForCausalLM", "XYZForSequenceClassification", + "XYZTokenClassifierOutput", "XYZForTokenClassification", "XYZConfig", "XYZProcessor", diff --git a/src/proxyz/models/processing_xyz.py b/src/proxyz/models/processing_xyz.py index 04364c5..e8f3dd9 100644 --- a/src/proxyz/models/processing_xyz.py +++ b/src/proxyz/models/processing_xyz.py @@ -36,7 +36,7 @@ class XYZProcessor(ProcessorMixin): FIM_MIDDLE = "" FIM_TOKENS = [FIM_PREFIX, FIM_SUFFIX, FIM_MIDDLE] - def __init__(self, tokenizer, text_column: str = "text", **kwargs) -> None: + def __init__(self, tokenizer, text_column: str = "text", features: list[str] = None, **kwargs) -> None: tokenizer.add_special_tokens({"additional_special_tokens": self.FIM_TOKENS}) logger.info(f"Added FIM tokens: {self.FIM_TOKENS} to tokenizer. (vocabu size: {len(tokenizer)})") @@ -46,13 +46,21 @@ def __init__(self, tokenizer, text_column: str = "text", **kwargs) -> None: super().__init__(tokenizer, **kwargs) self.text_column = text_column + self.features = features self.ignore_index = kwargs.get("ignore_index", -100) + self.distogram_bins = torch.linspace( + kwargs.get("distogram_bins_min", 2.3125), + kwargs.get("distogram_bins_max", 21.6875), + steps=kwargs.get("distogram_bins_num", 64) - 1, + ) + def __call__( self, examples: dict, *, bpe_dropout: float | None = None, + char_apply: bool = False, fim_apply: bool = False, fim_spm_rate: float = 0.5, fim_sft_style: bool = False, @@ -72,6 +80,40 @@ def __call__( fim_sft_style=fim_sft_style, generate=generate, ) + if char_apply: + tokenized.update( + self.to_characterization( + examples, + tokenized, + fim_apply=fim_apply, + fim_sft_style=fim_sft_style, + generate=generate, + ) + ) + # del tokenized["offset_mapping"] + + # copy features + if self.features is not None: + + def feat_collate(feat, is_label: bool = False): + if isinstance(feat, tuple): + return tuple(feat_collate(feat[t]) for t in range(len(feat))) + max_len = max(feat[k].shape[0] for k in range(len(feat))) + for k in range(len(feat)): + pad = feat[k].new_full( + (max_len - feat[k].shape[0], *feat[k].shape[1:]), + self.ignore_index if is_label else 0, + ) + feat[k] = torch.cat((feat[k], pad)) + return torch.stack(feat) + + for column in self.features: + if column in examples: + tokenized[column] = feat_collate(examples[column], "labels" in column) + + if "distogram_labels" in self.features: + if "distogram_labels" in tokenized: + tokenized["distogram_labels"] = self.to_distogram(*tokenized["distogram_labels"]) return tokenized @@ -97,6 +139,62 @@ def to_tokenization( ) return tokenized + def to_characterization( + self, + examples: dict, + tokenized: dict, + bpe_dropout: float = 1.0, + fim_apply: bool = False, + fim_sft_style: bool = False, + generate: bool = False, + ) -> dict: + assert "offset_mapping" in tokenized and "attention_mask" in tokenized + + characterized = {} + for k, v in self._tokenize_with_dropout( + examples, + bpe_dropout=bpe_dropout, + prefix="char_", + fim_apply=fim_apply, + fim_sft_style=fim_sft_style, + generate=generate, + ).items(): + if torch.is_tensor(v): + v = v.to(device=tokenized["offset_mapping"].device) + characterized[k] = v + + # Align characters with tokenized *text* + # NOTE: special token is treat as one char + # NOTE: pad offset_mapping with the maxinum offset + char_offset_mapping = characterized["char_offset_mapping"].where( + characterized["char_attention_mask"][..., None] > 0, characterized["char_offset_mapping"].max() + 1 + ) + offset_mapping = tokenized["offset_mapping"].where( + tokenized["attention_mask"][..., None] > 0, tokenized["offset_mapping"].max() + 1 + ) + token_to_char_map = torch.searchsorted( + char_offset_mapping.transpose(-1, -2).contiguous(), + offset_mapping.transpose(-1, -2).contiguous(), + right=True, + ) + token_to_char_cnt = token_to_char_map[:, 1, :] - token_to_char_map[:, 0, :] + 1 + characterized["repr_char_idx"] = token_to_char_cnt.cumsum(1) - 1 + characterized["repr_char_idx"] = characterized["repr_char_idx"].where(tokenized["attention_mask"] > 0, 0) + # characterized["char_to_token_mask"] = torch.arange(self.max_chars_within_token).view( + # 1, 1, self.max_chars_within_token + # ) + # characterized["char_to_token_mask"] = ( + # characterized["char_to_token_mask"] < token_to_char_cnt[:, :, None] + # ).where(tokenized["attention_mask"][:, :, None] > 0, 0) + # characterized["char_to_token_idx"] = torch.searchsorted( + # tokenized["offset_mapping"][:, :, 0].contiguous(), + # characterized["char_offset_mapping"][:, :, 0].contiguous(), + # right=True, + # ) - 1 + + del characterized["char_offset_mapping"] + return characterized + def to_example(self, input_ids: torch.LongTensor) -> dict: text = self.tokenizer.decode(input_ids.tolist(), skip_special_tokens=False) return {self.text_column: [t.replace(" ", "") for t in text]} @@ -151,6 +249,39 @@ def apply_fim(self, examples: dict, spm_rate: float = 0.5) -> dict: first, second = prefix, suffix examples[self.text_column][idx] = first_tag + first + second_tag + second + self.FIM_MIDDLE + middle + if self.features is not None: + + def feat_fim(feat, is_label: bool = False): + pad = feat.new_full((1, *feat.shape[1:]), self.ignore_index if is_label else 0) + if is_spm: + return torch.cat( + ( + pad, + feat[cut2:, ...], + pad, + feat[:cut1, ...], + pad, + feat[cut1:cut2, ...], + ) + ) + return torch.cat( + ( + pad, + feat[:cut1, ...], + pad, + feat[cut2:, ...], + pad, + feat[cut1:cut2, ...], + ) + ) + + for column in self.features: + if column in examples: + if isinstance(examples[column], tuple): + for t in range(len(examples[column])): + examples[column][t][idx] = feat_fim(examples[column][t][idx]) + else: + examples[column][idx] = feat_fim(examples[column][idx], "labels" in column) return examples def apply_crop(self, examples: dict, max_length: int | None = None) -> dict: @@ -160,6 +291,18 @@ def apply_crop(self, examples: dict, max_length: int | None = None) -> dict: cut = random.randint(0, n - max_length) text = text[cut : cut + max_length] examples[self.text_column][idx] = text + if self.features is not None: + + def feat_crop(feat): + return feat[cut : cut + max_length, ...] + + for column in self.features: + if column in examples: + if isinstance(examples[column], tuple): + for t in range(len(examples[column])): + examples[column][t][idx] = feat_crop(examples[column][t][idx]) + else: + examples[column][idx] = feat_crop(examples[column][idx]) return examples def apply_wrap(self, examples: dict, add_eos_token: bool = True) -> dict: @@ -168,7 +311,29 @@ def apply_wrap(self, examples: dict, add_eos_token: bool = True) -> dict: if add_eos_token: text = f"{text}{self.tokenizer.eos_token}" examples[self.text_column][idx] = text + if self.features is not None: + + def feat_wrap(feat, is_label: bool = False): + pad = feat.new_full((1, *feat.shape[1:]), self.ignore_index if is_label else 0) + if add_eos_token: + return torch.cat((pad, feat, pad)) + return torch.cat((pad, feat)) + + for column in self.features: + if column in examples: + if isinstance(examples[column], tuple): + for t in range(len(examples[column])): + examples[column][t][idx] = feat_wrap(examples[column][t][idx]) + else: + examples[column][idx] = feat_wrap(examples[column][idx], "labels" in column) return examples + def to_distogram(self, pseudo_beta, pseudo_beta_mask): + distogram_labels = torch.cdist(pseudo_beta, pseudo_beta) + distogram_labels = (distogram_labels[..., None] > self.distogram_bins).sum(-1) + return distogram_labels.where( + pseudo_beta_mask[..., :, None] * pseudo_beta_mask[..., None, :], self.ignore_index + ) + __all__ = ["XYZProcessor"] diff --git a/src/proxyz/train.py b/src/proxyz/train.py index b66d1fb..ce45272 100644 --- a/src/proxyz/train.py +++ b/src/proxyz/train.py @@ -90,6 +90,50 @@ default=4, help="Model Grouped-Query Attention (GQA) for speed.", ) +@click.option( + "--model_char_hidden_size", + type=int, + default=768, + help="Character: Model width.", +) +@click.option( + "--model_char_intermediate_size", + type=int, + default=2064, + help="Character: Model SwiGLU hidden dimension (usually ~8/3 of hidden_size).", +) +@click.option( + "--model_char_num_hidden_layers", + type=int, + default=2, + help="Character: Model depth.", +) +@click.option( + "--model_char_num_attention_heads", + type=int, + default=6, + help="Character: Model attention heads.", +) +@click.option( + "--model_use_char_position_ids", + is_flag=True, + help="Character: Model gather position_ids from char_position_ids", +) +@click.option( + "--model_has_char_lm_head", + is_flag=True, + help="Character: Model has char_lm_head", +) +@click.option( + "--model_has_cle_lm_head", + is_flag=True, + help="Character: Model has cle_lm_head", +) +@click.option( + "--model_has_distogram_lm_head", + is_flag=True, + help="Character: Model has distogram_lm_head", +) @click.option( "--max_position_embeddings", type=int, default=4096, help="Context window length." ) @@ -212,6 +256,26 @@ def main(**args): random.seed(args.random_seed) + features = [] + label_names = [("labels", 1)] + keys_to_ignore_at_inference = [ + "past_key_values", "char_past_key_values", "offset_mapping" + ] + if args.model_has_char_lm_head: + label_names += [("char_labels", 1)] + else: + keys_to_ignore_at_inference += ["char_logits"] + if args.model_has_cle_lm_head: + label_names += [("cle_labels", 0)] + features += ["cle_labels"] + else: + keys_to_ignore_at_inference += ["cle_logits"] + if args.model_has_distogram_lm_head: + label_names += [("distogram_labels", 0)] + features += ["distogram_labels"] + else: + keys_to_ignore_at_inference += ["distogram_logits"] + # ========================================== # 0. CHECK DATA SOURCE IS PROVIDED # ========================================== @@ -237,7 +301,7 @@ def main(**args): eos_token="[EOS]", ) processor = XYZProcessor( - tokenizer=tokenizer, text_column=args.text_column + tokenizer=tokenizer, text_column=args.text_column, features=features ) # Ensure the embedding layer matches this size exactly @@ -266,6 +330,15 @@ def main(**args): attn_implementation=args.attn_implementation, torch_dtype=torch.bfloat16, tie_word_embeddings=False, + keys_to_ignore_at_inference=keys_to_ignore_at_inference, + char_hidden_size=args.model_char_hidden_size, + char_intermediate_size=args.model_char_intermediate_size, + char_num_hidden_layers=args.model_char_num_hidden_layers, + char_num_attention_heads=args.model_char_num_attention_heads, + use_char_position_ids=args.model_use_char_position_ids, + has_char_lm_head=args.model_has_char_lm_head, + has_cle_lm_head=args.model_has_cle_lm_head, + has_distogram_lm_head=args.model_has_distogram_lm_head, ) model = XYZForCausalLM(config) @@ -278,6 +351,8 @@ def main(**args): total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"--- Dense DeepSeek-Style Model ---") + if config.has_characterization: + print("Use U-net style XYZForCausalLM instead of standard Llama attention.") print(f"Attention backend: {args.attn_implementation}") print(f"Total Parameters: {total_params:,}") print(f"Trainable Parameters: {trainable_params:,}") @@ -299,6 +374,7 @@ def tokenize_function(examples): examples = processor( examples, bpe_dropout=args.tokenizer_bpe_dropout, + char_apply=config.has_characterization, fim_apply=fim_apply, fim_spm_rate=args.fim_spm_rate, fim_sft_style=args.fim_sft_style, @@ -385,7 +461,9 @@ def compute_loss( ): # Always cache data for FIM loss tracking (training and eval) # Cache data BEFORE calling super (which may modify inputs) - labels = inputs["labels"].clone() + labels = tuple(inputs[k].clone() for k in self.args.label_names) + if len(labels) == 1: + labels = labels[0] # Call parent compute_loss (handles label smoothing, loss scaling, etc.) loss, outputs = super().compute_loss( @@ -397,7 +475,24 @@ def compute_loss( # ONLY track training metrics if the model is actively training if model.training and self.compute_metrics is not None: - logits = outputs.logits + ignore_keys = [] + + module = model + # FIX: DistributedDataParallel + if not hasattr(module, "config") and hasattr(module, "module"): + module = module.module + if hasattr(module, "config"): + ignore_keys = getattr( + module.config, + "keys_to_ignore_at_inference", + ["past_key_values"] + ) + + logits = tuple( + v for k, v in outputs.items() if k not in ignore_keys + ["loss"] + ) + if len(logits) == 1: + logits = logits[0] if self.preprocess_logits_for_metrics is not None: logits = self.preprocess_logits_for_metrics(logits, labels) for key, val in self.compute_metrics( @@ -410,6 +505,32 @@ def compute_loss( return (loss, outputs) if return_outputs else loss + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys): + loss, logits, labels = super().prediction_step( + model, inputs, prediction_loss_only, ignore_keys + ) + + def pad_across_processes(tensors, dims=None): + if isinstance(tensors, tuple): + return tuple( + map(functools.partial(pad_across_processes, dims=dims), tensors) + ) + if dims is None: + dims = tensors.dim() + elif dims < 0: + dims = dims + tensors.dim() + for dim in range(2, dims): + tensors = self.accelerator.pad_across_processes( + tensors, dim=dim, pad_index=-100 + ) + return tensors + + if logits is not None: + logits = pad_across_processes(logits, dims=-1) + if labels is not None: + labels = pad_across_processes(labels) + return loss, logits, labels + def log(self, logs, start_time=None): if self._logs: for key, val in self._logs.items(): @@ -421,13 +542,26 @@ def log(self, logs, start_time=None): super().log(logs, start_time=start_time) @classmethod - def aux_preprocess_logits_for_metrics(cls, logits, labels): + def aux_preprocess_logits_for_metrics(cls, logits, labels, label_names=None): loss_fct = torch.nn.CrossEntropyLoss(reduction="none") + if isinstance(logits, tuple) and isinstance(labels, tuple): + return tuple( + cls.aux_preprocess_logits_for_metrics( + p, l, label_names=label_names[i] if label_names else None + ) for i, (p, l) in enumerate(zip(logits, labels)) + ) assert not isinstance(logits, tuple), len(logits) assert not isinstance(labels, tuple), len(labels) - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() + n_shift = 0 + if label_names is None or label_names[1]: + n_shift = labels.dim() - 1 + logits_slices = [slice(0, -1)] * n_shift + labels_slices = [slice(1, None)] * n_shift + + shift_logits = logits[..., *logits_slices, :].contiguous() + shift_labels = labels[..., *labels_slices].contiguous() + loss_per_token = loss_fct( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1) ) @@ -435,7 +569,7 @@ def aux_preprocess_logits_for_metrics(cls, logits, labels): @classmethod def aux_metric_calculator( - cls, metrics, preds, compute_result=False, update_metrics=True, prefix="eval_" + cls, metrics, preds, compute_result=False, update_metrics=True, prefix="eval_", label_names=None ): """Computes metrics: n_fim / n_std, loss_fim / loss_std etc""" loss_per_token, labels = preds @@ -443,7 +577,34 @@ def aux_metric_calculator( logs = {} if isinstance(loss_per_token, tuple) and isinstance(labels, tuple): - assert False + assert len(label_names) == len(loss_per_token) + for i, (p, l) in enumerate(zip(loss_per_token, labels)): + label, _ = label_names[i] + if label.endswith("labels"): + label = label[:-len("labels")] + logs.update( + cls.aux_metric_calculator( + metrics, (p, l), + update_metrics=False, + prefix=f"{prefix}{label}", + label_names=label_names[i], + ) + ) + elif isinstance(loss_per_token, list) and isinstance(labels, list): + assert len(loss_per_token) % len(label_names) == 0, ([l.shape for l in loss_per_token], [l.shape for l in labels], label_names) + gather_logs = defaultdict(list) + + for g in range(len(loss_per_token) // len(label_names)): + i, j = g * len(label_names), (g + 1) * len(label_names) + for key, value in cls.aux_metric_calculator( + metrics, (tuple(loss_per_token[i:j]), tuple(labels[i:j])), + update_metrics=False, + prefix=prefix, + label_names=label_names, + ).items(): + gather_logs[key].append(value) + + logs.update({k: sum(v)/len(v) for k, v in gather_logs.items()}) else: with torch.no_grad(): # Detect FIM examples: labels start with -100 @@ -456,13 +617,18 @@ def aux_metric_calculator( ) ) is_fim = is_fim.any(tuple(range(1, is_fim.dim()))) + + n_shift = 0 + if label_names is None or label_names[1]: + n_shift = labels.dim() - 1 + labels_slices = [slice(1, None)] * n_shift for tag, mask in [("fim", is_fim), ("std", ~is_fim)]: # Add FIM/standard counts to logs logs[f"{prefix}n_{tag}"] = mask.sum().item() # Add FIM/standard loss to logs if mask.any(): - shift_labels = labels[mask][..., 1:] + shift_labels = labels[mask][..., *labels_slices] valid = (shift_labels != -100).reshape(-1) if valid.any(): loss = loss_per_token[mask].reshape(-1) @@ -501,14 +667,17 @@ def aux_metric_calculator( eval_strategy=args.eval_strategy if (args.eval_files or args.dataset_eval_split) else "no", eval_steps=args.eval_steps if args.eval_strategy == "steps" else None, per_device_eval_batch_size=args.per_device_train_batch_size, + eval_use_gather_object=len(label_names) > 1, eval_accumulation_steps=args.gradient_accumulation_steps, eval_on_start=True if args.eval_files else False, batch_eval_metrics=True, + label_names=[label for label, _ in label_names], bf16=use_cuda, # bf16 is preferred over fp16 on modern GPUs ddp_find_unused_parameters=False, # disabled warning num_train_epochs=args.num_train_epochs, max_steps=args.max_steps, dataloader_num_workers=args.dataloader_num_workers, + dataloader_drop_last=True, remove_unused_columns=False, report_to=report_to, # SwanLab + TensorBoard run_name=args.run_name, @@ -521,10 +690,14 @@ def aux_metric_calculator( eval_dataset=eval_dataset, processing_class=processor, # transformers >=5 renamed `tokenizer` train_sampler=train_sampler, - preprocess_logits_for_metrics=FIMTrainer.aux_preprocess_logits_for_metrics, + preprocess_logits_for_metrics=functools.partial( + FIMTrainer.aux_preprocess_logits_for_metrics, + label_names=label_names if len(label_names) > 1 else label_names[0], + ), compute_metrics=functools.partial( FIMTrainer.aux_metric_calculator, defaultdict(list), + label_names=label_names if len(label_names) > 1 else label_names[0], ), )