Hugging Face's logo

OpenModel-Labs
/
Comet-1-Preview

Text Generation
Transformers
Safetensors
COMET
English
openmodel-1
causal-lm
base-model
hybrid-architecture
custom_code
Model card Files
xet
 Community
1
Comet-1-Preview / modeling_comet.py
Bc-AI's picture
Bc-AI
Add Transformers support (modeling + config code), tokenizer/generation configs and model card
9df61c9 verified
Raw History Blame Contribute Delete
25.2 kB
"""
Comet-1 (OpenModel-1) for Hugging Face Transformers.
Inference/fine-tuning port of the ``OpenModel1`` module defined in the training
pipeline (``train.py``). Parameter names and shapes are identical to the training
checkpoint, so ``model.safetensors`` loads without any key remapping.
Architecture
------------
A pre-norm decoder with tied input/output embeddings. Blocks follow a repeating
schedule (``config.layer_pattern``, by default ``[global, global, global, recall]``):
* ``global`` blocks use causal grouped-query attention with per-head RMSNorm on
queries and keys (QK-norm) followed by rotary position embeddings.
* ``recall`` blocks replace attention with a gated, depthwise, dilated *causal*
1-D convolution over a narrower lane (``config.memory_width``).
Every block then applies a SwiGLU feed-forward. Both residual branches are scaled
by a learned per-channel vector (LayerScale).
Compared with ``train.py`` this file adds what generation needs and the training
loop did not: an incremental cache (attention keys/values *and* convolution
history), padding-aware attention masks, and the standard Transformers outputs.
"""
import math
from typing import Optional, Tuple, Union
import torch
import torch.nn.functional as F
from torch import nn
from transformers import DynamicCache, GenerationMixin, PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from .configuration_comet import CometConfig
def _cache_length(cache) -> int:
"""Number of tokens already stored in ``cache`` (0 for ``None`` or an empty cache)."""
if cache is None:
return 0
try:
return int(cache.get_seq_length())
except Exception:
return 0
class CometRMSNorm(nn.Module):
def __init__(self, size: int, eps: float = 1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(size))
self.eps = eps
def forward(self, x: torch.Tensor) -> torch.Tensor:
weight = self.weight.to(x.dtype)
if hasattr(F, "rms_norm"):
return F.rms_norm(x, (x.shape[-1],), weight, self.eps)
variance = x.float().pow(2).mean(dim=-1, keepdim=True)
return weight * (x * torch.rsqrt(variance + self.eps).to(x.dtype))
def extra_repr(self) -> str:
return f"{self.weight.shape[0]}, eps={self.eps}"
def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
"""Rotate interleaved (even, odd) channel pairs. ``x``: [B, H, T, D]; ``cos``/``sin``: [B|1, 1, T, D/2]."""
even = x[..., ::2]
odd = x[..., 1::2]
return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
class CometAttention(nn.Module):
"""Causal grouped-query attention with QK-norm and RoPE."""
def __init__(self, config: CometConfig):
super().__init__()
self.heads = config.num_attention_heads
self.kv_heads = config.num_key_value_heads
self.head_dim = config.hidden_size // config.num_attention_heads
self.q_dim = self.heads * self.head_dim
self.kv_dim = self.kv_heads * self.head_dim
self.fused_projections = config.fused_projections
if self.fused_projections:
self.qkv_proj = nn.Linear(config.hidden_size, self.q_dim + 2 * self.kv_dim, bias=False)
else:
self.q_proj = nn.Linear(config.hidden_size, self.q_dim, bias=False)
self.k_proj = nn.Linear(config.hidden_size, self.kv_dim, bias=False)
self.v_proj = nn.Linear(config.hidden_size, self.kv_dim, bias=False)
self.o_proj = nn.Linear(self.q_dim, config.hidden_size, bias=False)
self.o_proj._comet_residual_out = True
if config.qk_norm:
self.q_norm = CometRMSNorm(self.head_dim, eps=config.rms_norm_eps)
self.k_norm = CometRMSNorm(self.head_dim, eps=config.rms_norm_eps)
else:
self.q_norm = nn.Identity()
self.k_norm = nn.Identity()
def forward(
self,
x: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
attn_mask: Optional[torch.Tensor] = None,
is_causal: bool = False,
cache=None,
layer_idx: int = 0,
) -> torch.Tensor:
batch, sequence, _ = x.shape
if self.fused_projections:
q, k, v = self.qkv_proj(x).split((self.q_dim, self.kv_dim, self.kv_dim), dim=-1)
else:
q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x)
q = q.view(batch, sequence, self.heads, self.head_dim).transpose(1, 2)
k = k.view(batch, sequence, self.kv_heads, self.head_dim).transpose(1, 2)
v = v.view(batch, sequence, self.kv_heads, self.head_dim).transpose(1, 2)
cos = cos.to(device=q.device, dtype=q.dtype)
sin = sin.to(device=q.device, dtype=q.dtype)
q = _apply_rope(self.q_norm(q), cos, sin)
k = _apply_rope(self.k_norm(k), cos, sin)
if cache is not None:
# Keys are stored already normalised and rotated, at key/value-head granularity.
k, v = cache.update(k, v, layer_idx)
groups = self.heads // self.kv_heads
if groups > 1:
# Query head h reads key/value head h // groups (same grouping as `enable_gqa`).
length = k.shape[-2]
k = k[:, :, None].expand(batch, self.kv_heads, groups, length, self.head_dim)
k = k.reshape(batch, self.heads, length, self.head_dim)
v = v[:, :, None].expand(batch, self.kv_heads, groups, length, self.head_dim)
v = v.reshape(batch, self.heads, length, self.head_dim)
if attn_mask is not None:
y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask.to(q.device))
else:
y = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
return self.o_proj(y.transpose(1, 2).reshape(batch, sequence, self.q_dim))
class CometRecallMixer(nn.Module):
"""Gated local recall branch: depthwise, dilated, causal convolution over a narrow lane."""
def __init__(self, config: CometConfig, dilation: int):
super().__init__()
width = config.memory_width
self.width = width
self.kernel_size = config.memory_kernel_size
self.dilation = dilation
# Number of past frames a single output position can see.
self.history = (self.kernel_size - 1) * self.dilation
self.in_proj = nn.Linear(config.hidden_size, 2 * width, bias=False)
self.depthwise = nn.Conv1d(
width,
width,
kernel_size=self.kernel_size,
dilation=dilation,
groups=width,
bias=False,
)
self.out_proj = nn.Linear(width, config.hidden_size, bias=False)
self.out_proj._comet_residual_out = True
def forward(
self,
x: torch.Tensor,
token_mask: Optional[torch.Tensor] = None,
cache=None,
layer_idx: int = 0,
) -> torch.Tensor:
batch, sequence, _ = x.shape
gate, value = self.in_proj(x).chunk(2, dim=-1)
if token_mask is not None:
# Padding tokens contribute nothing to the convolution, exactly like the
# zero padding that precedes the first real token during training.
value = value * token_mask.to(device=value.device, dtype=value.dtype).unsqueeze(-1)
if cache is not None:
# The convolution input is cached in the "key" slot as [B, 1, T, width] so
# that every generic cache operation (beam reordering, cropping, batch
# expansion) applies to it unchanged. The "value" slot is a placeholder.
frames, _ = cache.update(value.unsqueeze(1), value.new_zeros(batch, 1, sequence, 1), layer_idx)
context = frames[:, 0, -(self.history + sequence) :]
else:
context = value
left = self.history + sequence - context.shape[1]
context = F.pad(context.transpose(1, 2), (left, 0))
value = self.depthwise(context).transpose(1, 2)
return self.out_proj(torch.sigmoid(gate) * value)
class CometSwiGLU(nn.Module):
def __init__(self, config: CometConfig):
super().__init__()
self.fused_projections = config.fused_projections
if self.fused_projections:
self.gate_up = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
else:
self.gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
self.up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
self.down._comet_residual_out = True
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.fused_projections:
gate, up = self.gate_up(x).chunk(2, dim=-1)
else:
gate, up = self.gate(x), self.up(x)
return self.down(F.silu(gate) * up)
class CometBlock(nn.Module):
def __init__(self, config: CometConfig, index: int):
super().__init__()
self.index = index
self.kind = config.block_kind(index)
self.norm1 = CometRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.norm2 = CometRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
if self.kind == "global":
self.mixer = CometAttention(config)
else:
self.mixer = CometRecallMixer(config, config.block_dilation(index))
self.mlp = CometSwiGLU(config)
self.mixer_scale = nn.Parameter(torch.full((config.hidden_size,), float(config.layer_scale_init)))
self.mlp_scale = nn.Parameter(torch.full((config.hidden_size,), float(config.layer_scale_init)))
self.dropout = config.dropout
def forward(
self,
x: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
attn_mask: Optional[torch.Tensor] = None,
token_mask: Optional[torch.Tensor] = None,
is_causal: bool = False,
cache=None,
) -> torch.Tensor:
h = self.norm1(x)
if self.kind == "global":
h = self.mixer(h, cos, sin, attn_mask=attn_mask, is_causal=is_causal, cache=cache, layer_idx=self.index)
else:
h = self.mixer(h, token_mask=token_mask, cache=cache, layer_idx=self.index)
if self.dropout:
h = F.dropout(h, p=self.dropout, training=self.training)
x = x + h * self.mixer_scale.to(h.dtype)
h = self.mlp(self.norm2(x))
if self.dropout:
h = F.dropout(h, p=self.dropout, training=self.training)
return x + h * self.mlp_scale.to(h.dtype)
class CometPreTrainedModel(PreTrainedModel):
config_class = CometConfig
base_model_prefix = ""
supports_gradient_checkpointing = True
_no_split_modules = ["CometBlock"]
_skip_keys_device_placement = ["past_key_values"]
_supports_cache_class = True
def _init_weights(self, module: nn.Module) -> None:
std = self.config.initializer_range
if isinstance(module, nn.Linear):
if getattr(module, "_comet_residual_out", False):
std = std / math.sqrt(2.0 * self.config.num_hidden_layers)
nn.init.normal_(module.weight, mean=0.0, std=std)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, (nn.Embedding, nn.Conv1d)):
nn.init.normal_(module.weight, mean=0.0, std=std)
elif isinstance(module, CometRMSNorm):
nn.init.ones_(module.weight)
elif isinstance(module, CometBlock):
nn.init.constant_(module.mixer_scale, float(self.config.layer_scale_init))
nn.init.constant_(module.mlp_scale, float(self.config.layer_scale_init))
class CometModel(CometPreTrainedModel):
"""Bare Comet trunk: token embeddings, blocks and the final norm. Returns hidden states."""
def __init__(self, config: CometConfig):
super().__init__(config)
self._build_trunk(config)
self.post_init()
def _build_trunk(self, config: CometConfig) -> None:
self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
self.blocks = nn.ModuleList(CometBlock(config, index) for index in range(config.num_hidden_layers))
self.norm = CometRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.gradient_checkpointing = False
# RoPE tables are built lazily on first use and kept as plain attributes rather than
# buffers: they are derived constants, and this keeps them out of every loading path
# (meta-device instantiation, device maps, dtype casts).
self._rope_tables = None
def _rope(self, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
tables = self._rope_tables
if tables is None or tables[0].device != device:
config = self.config
head_dim = config.hidden_size // config.num_attention_heads
exponent = torch.arange(0, head_dim, 2, dtype=torch.float32, device="cpu") / head_dim
inv_freq = 1.0 / (config.rope_theta**exponent)
positions = torch.arange(config.max_position_embeddings, dtype=torch.float32, device="cpu")
frequencies = torch.outer(positions, inv_freq)
tables = (frequencies.cos().to(device), frequencies.sin().to(device))
self._rope_tables = tables
return tables
def get_input_embeddings(self) -> nn.Embedding:
return self.embed
def set_input_embeddings(self, value: nn.Embedding) -> None:
self.embed = value
def _run_trunk(
self,
input_ids: Optional[torch.LongTensor] = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values=None,
inputs_embeds: Optional[torch.Tensor] = None,
use_cache: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
):
if (input_ids is None) == (inputs_embeds is None):
raise ValueError("Pass exactly one of `input_ids` or `inputs_embeds`.")
if inputs_embeds is None:
inputs_embeds = self.embed(input_ids)
batch, sequence = inputs_embeds.shape[:2]
device = inputs_embeds.device
use_cache = self.config.use_cache if use_cache is None else use_cache
if self.training and past_key_values is None:
use_cache = False
cache = past_key_values
if cache is not None and not hasattr(cache, "update"):
raise TypeError(
"`past_key_values` must be a `transformers.Cache` instance (for example `DynamicCache()`); "
"legacy tuple caches are not supported."
)
if cache is None and use_cache:
cache = DynamicCache()
past = _cache_length(cache)
total = past + sequence
if total > self.config.max_position_embeddings:
raise ValueError(
f"Sequence length {total} exceeds max_position_embeddings={self.config.max_position_embeddings}."
)
# ---- masks -------------------------------------------------------------
attn_mask = None # [B|1, 1, T, past + T] boolean, True = may attend
token_mask = None # [B, T] boolean, True = real token (current chunk only)
is_causal = False
keep = None
if attention_mask is not None:
if attention_mask.dim() != 2:
raise ValueError("`attention_mask` must have shape [batch, sequence].")
keep = attention_mask.to(device=device, dtype=torch.bool)
if keep.shape[1] == sequence and past > 0:
# Only the new chunk was described; everything cached is treated as real.
keep = torch.cat([keep.new_ones(batch, past), keep], dim=1)
if keep.shape[1] != total:
raise ValueError(
f"`attention_mask` has length {attention_mask.shape[1]} but the model sees {total} positions "
f"({past} cached + {sequence} new)."
)
if bool(keep.all()):
keep = None
if keep is not None or (past > 0 and sequence > 1):
query_positions = past + torch.arange(sequence, device=device)
key_positions = torch.arange(total, device=device)
attn_mask = (key_positions[None, :] <= query_positions[:, None])[None, None]
if keep is not None:
# Every position may always attend to itself so that padded rows never
# become fully masked (which would produce NaNs).
diagonal = (key_positions[None, :] == query_positions[:, None])[None, None]
attn_mask = (attn_mask & keep[:, None, None, :]) | diagonal
token_mask = keep[:, past:]
elif past == 0 and sequence > 1:
is_causal = True
# ---- positions ---------------------------------------------------------
if position_ids is None:
if keep is not None:
# Real tokens are numbered 0, 1, 2, ... regardless of how much padding precedes them.
position_ids = (keep.long().cumsum(-1) - 1).clamp_min(0)[:, past:]
else:
position_ids = torch.arange(past, total, device=device)[None, :]
else:
position_ids = position_ids.to(device)
if position_ids.dim() == 1:
position_ids = position_ids[None, :]
if position_ids.shape[-1] != sequence:
position_ids = position_ids[..., -sequence:]
if int(position_ids.max()) >= self.config.max_position_embeddings:
raise ValueError("`position_ids` exceed max_position_embeddings.")
rope_cos, rope_sin = self._rope(device)
cos = rope_cos[position_ids].unsqueeze(1)
sin = rope_sin[position_ids].unsqueeze(1)
# ---- blocks ------------------------------------------------------------
x = inputs_embeds
all_hidden_states = () if output_hidden_states else None
for block in self.blocks:
if output_hidden_states:
all_hidden_states += (x,)
if self.gradient_checkpointing and self.training and cache is None:
x = torch.utils.checkpoint.checkpoint(
block, x, cos, sin, attn_mask, token_mask, is_causal, use_reentrant=False
)
else:
x = block(x, cos, sin, attn_mask=attn_mask, token_mask=token_mask, is_causal=is_causal, cache=cache)
x = self.norm(x)
if output_hidden_states:
all_hidden_states += (x,)
return x, (cache if use_cache else None), all_hidden_states
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values=None,
inputs_embeds: Optional[torch.Tensor] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
**kwargs,
) -> Union[Tuple, BaseModelOutputWithPast]:
if output_hidden_states is None:
output_hidden_states = bool(getattr(self.config, "output_hidden_states", False))
if return_dict is None:
return_dict = bool(getattr(self.config, "return_dict", True))
hidden, cache, all_hidden_states = self._run_trunk(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_hidden_states=output_hidden_states,
)
if not return_dict:
return tuple(item for item in (hidden, cache, all_hidden_states) if item is not None)
return BaseModelOutputWithPast(
last_hidden_state=hidden,
past_key_values=cache,
hidden_states=all_hidden_states,
)
class CometForCausalLM(CometModel, GenerationMixin):
"""
Comet trunk with the language-modelling head.
The state-dict layout is flat and identical to the training checkpoint
(``embed.weight``, ``blocks.N.*``, ``norm.weight``). With
``config.tie_word_embeddings=True`` there is no separate head parameter: logits are
computed directly against ``embed.weight``.
"""
def __init__(self, config: CometConfig):
CometPreTrainedModel.__init__(self, config)
self._build_trunk(config)
if config.tie_word_embeddings:
self.lm_head = None
else:
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
self.post_init()
def get_output_embeddings(self):
# `None` when tied: there is nothing for Transformers to tie or resize separately.
return self.lm_head
def set_output_embeddings(self, value) -> None:
self.lm_head = value
def _logits(self, hidden: torch.Tensor) -> torch.Tensor:
weight = self.embed.weight if self.lm_head is None else self.lm_head.weight
return F.linear(hidden.to(weight.device), weight)
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values=None,
inputs_embeds: Optional[torch.Tensor] = None,
labels: Optional[torch.LongTensor] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
**kwargs,
) -> Union[Tuple, CausalLMOutputWithPast]:
r"""
labels (`torch.LongTensor` of shape `(batch, sequence)`, *optional*):
Next-token targets. They are shifted inside the model, so pass ``labels=input_ids``;
positions set to ``-100`` are ignored.
"""
if output_hidden_states is None:
output_hidden_states = bool(getattr(self.config, "output_hidden_states", False))
if return_dict is None:
return_dict = bool(getattr(self.config, "return_dict", True))
hidden, cache, all_hidden_states = self._run_trunk(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_hidden_states=output_hidden_states,
)
logits = self._logits(hidden)
loss = None
if labels is not None:
shift_logits = logits[:, :-1].float()
shift_labels = labels[:, 1:].to(shift_logits.device)
total = F.cross_entropy(
shift_logits.reshape(-1, shift_logits.shape[-1]),
shift_labels.reshape(-1),
ignore_index=-100,
reduction="sum",
)
loss = total / shift_labels.ne(-100).sum().clamp_min(1)
if not return_dict:
return tuple(item for item in (loss, logits, cache, all_hidden_states) if item is not None)
return CausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=cache,
hidden_states=all_hidden_states,
)
def prepare_inputs_for_generation(
self,
input_ids: torch.LongTensor,
past_key_values=None,
attention_mask: Optional[torch.Tensor] = None,
inputs_embeds: Optional[torch.Tensor] = None,
**kwargs,
) -> dict:
"""Feed only the tokens the cache has not seen yet; keep the full-length attention mask."""
past = _cache_length(past_key_values)
model_inputs = {}
if inputs_embeds is not None and past == 0:
model_inputs["inputs_embeds"] = inputs_embeds
else:
seen = past
if inputs_embeds is not None:
# Generation started from embeddings: `input_ids` holds only the generated tokens.
seen = past - inputs_embeds.shape[1]
unseen = input_ids.shape[1] - seen
if past > 0:
input_ids = input_ids[:, -unseen:] if unseen > 0 else input_ids[:, -1:]
model_inputs["input_ids"] = input_ids.contiguous()
model_inputs["past_key_values"] = past_key_values
model_inputs["attention_mask"] = attention_mask
model_inputs["use_cache"] = kwargs.get("use_cache", True)
if kwargs.get("output_hidden_states") is not None:
model_inputs["output_hidden_states"] = kwargs["output_hidden_states"]
return model_inputs
__all__ = ["CometPreTrainedModel", "CometModel", "CometForCausalLM"]