caduceus_base / modeling_caduceus.py

Update modeling_caduceus.py

c699cf9 verified 2 months ago

No virus

26.3 kB

	"""Caduceus model for Hugging Face.

	"""

	import math
	from functools import partial
	from typing import Optional, Tuple, Union

	import torch
	from mamba_ssm.modules.mamba_simple import Mamba, Block
	from torch import nn
	from torch.nn import functional as F
	from transformers import PreTrainedModel
	from transformers.modeling_outputs import BaseModelOutputWithNoAttention, MaskedLMOutput, SequenceClassifierOutput

	try:
	from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn
	except ImportError:
	RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None

	from .configuration_caduceus import CaduceusConfig
	from .modeling_rcps import RCPSAddNormWrapper, RCPSEmbedding, RCPSLMHead, RCPSMambaBlock


	def create_block(
	d_model,
	ssm_cfg=None,
	norm_epsilon=1e-5,
	rms_norm=False,
	residual_in_fp32=False,
	fused_add_norm=False,
	layer_idx=None,
	bidirectional=True,
	bidirectional_strategy="add",
	bidirectional_weight_tie=True,
	rcps=False,
	device=None,
	dtype=None,
	):
	"""Create Caduceus block.

	Adapted from: https://github.com/state-spaces/mamba/blob/main/mamba_ssm/models/mixer_seq_simple.py
	"""
	if ssm_cfg is None:
	ssm_cfg = {}
	factory_kwargs = {"device": device, "dtype": dtype}
	bidirectional_kwargs = {
	"bidirectional": bidirectional,
	"bidirectional_strategy": bidirectional_strategy,
	"bidirectional_weight_tie": bidirectional_weight_tie,
	}
	mixer_cls = partial(BiMambaWrapper, layer_idx=layer_idx, ssm_cfg, bidirectional_kwargs, **factory_kwargs)
	norm_cls = partial(
	nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs
	)
	block_cls = RCPSMambaBlock if rcps else Block
	block = block_cls(
	d_model,
	mixer_cls,
	norm_cls=norm_cls,
	fused_add_norm=fused_add_norm,
	residual_in_fp32=residual_in_fp32,
	)
	block.layer_idx = layer_idx
	return block


	class BiMambaWrapper(nn.Module):
	"""Thin wrapper around Mamba to support bi-directionality."""

	def __init__(
	self,
	d_model: int,
	bidirectional: bool = True,
	bidirectional_strategy: Optional[str] = "add",
	bidirectional_weight_tie: bool = True,
	**mamba_kwargs,
	):
	super().__init__()
	if bidirectional and bidirectional_strategy is None:
	bidirectional_strategy = "add" # Default strategy: `add`
	if bidirectional and bidirectional_strategy not in ["add", "ew_multiply"]:
	raise NotImplementedError(f"`{bidirectional_strategy}` strategy for bi-directionality is not implemented!")
	self.bidirectional = bidirectional
	self.bidirectional_strategy = bidirectional_strategy
	self.mamba_fwd = Mamba(
	d_model=d_model,
	**mamba_kwargs
	)
	if bidirectional:
	self.mamba_rev = Mamba(
	d_model=d_model,
	**mamba_kwargs
	)
	if bidirectional_weight_tie: # Tie in and out projections (where most of param count lies)
	self.mamba_rev.in_proj.weight = self.mamba_fwd.in_proj.weight
	self.mamba_rev.in_proj.bias = self.mamba_fwd.in_proj.bias
	self.mamba_rev.out_proj.weight = self.mamba_fwd.out_proj.weight
	self.mamba_rev.out_proj.bias = self.mamba_fwd.out_proj.bias
	else:
	self.mamba_rev = None

	def forward(self, hidden_states, inference_params=None):
	"""Bidirectional-enabled forward pass

	hidden_states: (B, L, D)
	Returns: same shape as hidden_states
	"""
	out = self.mamba_fwd(hidden_states, inference_params=inference_params)
	if self.bidirectional:
	out_rev = self.mamba_rev(
	hidden_states.flip(dims=(1,)), # Flip along the sequence length dimension
	inference_params=inference_params
	).flip(dims=(1,)) # Flip back for combining with forward hidden states
	if self.bidirectional_strategy == "add":
	out = out + out_rev
	elif self.bidirectional_strategy == "ew_multiply":
	out = out * out_rev
	else:
	raise NotImplementedError(f"`{self.bidirectional_strategy}` for bi-directionality not implemented!")
	return out


	class CaduceusEmbeddings(nn.Module):
	def __init__(
	self,
	config: CaduceusConfig,
	device=None,
	dtype=None,
	):
	super().__init__()
	factory_kwargs = {"device": device, "dtype": dtype}
	if config.rcps:
	self.word_embeddings = RCPSEmbedding(
	config.vocab_size, config.d_model, config.complement_map, **factory_kwargs
	)
	else:
	self.word_embeddings = nn.Embedding(config.vocab_size, config.d_model, **factory_kwargs)

	def forward(self, input_ids):
	"""
	input_ids: (batch, seqlen)
	"""
	return self.word_embeddings(input_ids)


	class CaduceusMixerModel(nn.Module):
	def __init__(
	self,
	config: CaduceusConfig,
	device=None,
	dtype=None,
	) -> None:
	super().__init__()
	factory_kwargs = {"device": device, "dtype": dtype}

	self.fused_add_norm = config.fused_add_norm
	self.rcps = config.rcps
	self.residual_in_fp32 = config.residual_in_fp32

	self.embeddings = CaduceusEmbeddings(config, **factory_kwargs)

	# Mamba changes the order of residual and layer norm:
	# Instead of LN -> Attn / MLP -> Add, we do:
	# Add -> LN -> Attn / MLP / Mixer, returning both the residual branch (output of Add) and
	# the main branch (output of MLP / Mixer). The model definition is unchanged.
	# This is for performance reason: we can fuse add + layer_norm.
	if config.fused_add_norm:
	if layer_norm_fn is None or rms_norm_fn is None:
	raise ImportError("Failed to import Triton LayerNorm / RMSNorm kernels")

	self.layers = nn.ModuleList(
	[
	create_block(
	config.d_model,
	ssm_cfg=config.ssm_cfg,
	norm_epsilon=config.norm_epsilon,
	rms_norm=config.rms_norm,
	residual_in_fp32=config.residual_in_fp32,
	fused_add_norm=config.fused_add_norm,
	layer_idx=i,
	bidirectional=config.bidirectional,
	bidirectional_strategy=config.bidirectional_strategy,
	bidirectional_weight_tie=config.bidirectional_weight_tie,
	rcps=config.rcps,
	**factory_kwargs,
	)
	for i in range(config.n_layer)
	]
	)

	norm_f = (nn.LayerNorm if not config.rms_norm else RMSNorm)(
	config.d_model, eps=config.norm_epsilon, **factory_kwargs
	)
	self.norm_f = norm_f if (config.fused_add_norm or not config.rcps) else RCPSAddNormWrapper(norm_f)

	def forward(self, input_ids, inputs_embeds=None, output_hidden_states=False):
	"""Mixer forward."""
	all_hidden_states = []
	if inputs_embeds is not None:
	hidden_states = inputs_embeds
	else:
	hidden_states = self.embeddings(input_ids)

	residual = None
	for layer in self.layers:
	if output_hidden_states:
	all_hidden_states.append(hidden_states)
	# TODO: Add support for gradient checkpointing
	hidden_states, residual = layer(
	hidden_states, residual, inference_params=None
	)

	if not self.fused_add_norm:
	if self.rcps:
	# Set prenorm=False here since we don't need the residual
	hidden_states = self.norm_f(hidden_states, residual=residual, prenorm=False)
	else:
	residual = (hidden_states + residual) if residual is not None else hidden_states
	hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))
	else:
	fused_add_norm_fn = rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn
	if self.rcps:
	# Set prenorm=False here since we don't need the residual
	hidden_states_fwd = fused_add_norm_fn(
	hidden_states[..., :hidden_states.shape[-1] // 2],
	self.norm_f.weight,
	self.norm_f.bias,
	eps=self.norm_f.eps,
	residual=residual[..., :hidden_states.shape[-1] // 2],
	prenorm=False,
	residual_in_fp32=self.residual_in_fp32,
	)
	hidden_states_rc = fused_add_norm_fn(
	hidden_states[..., hidden_states.shape[-1] // 2:].flip(dims=[-2, -1]),
	self.norm_f.weight,
	self.norm_f.bias,
	eps=self.norm_f.eps,
	residual=residual[..., hidden_states.shape[-1] // 2:].flip(dims=[-2, -1]),
	prenorm=False,
	residual_in_fp32=self.residual_in_fp32,
	)
	hidden_states = torch.cat([hidden_states_fwd, hidden_states_rc.flip(dims=[-2, -1])], dim=-1)
	else:
	# Set prenorm=False here since we don't need the residual
	hidden_states = fused_add_norm_fn(
	hidden_states,
	self.norm_f.weight,
	self.norm_f.bias,
	eps=self.norm_f.eps,
	residual=residual,
	prenorm=False,
	residual_in_fp32=self.residual_in_fp32,
	)
	if output_hidden_states:
	all_hidden_states.append(hidden_states)
	return hidden_states, all_hidden_states


	def cross_entropy(logits, y, ignore_index=-100):
	"""Cross entropy loss."""
	logits = logits.view(-1, logits.shape[-1])
	y = y.view(-1)
	return F.cross_entropy(logits, y, ignore_index=ignore_index)


	def weighted_cross_entropy(logits, y, loss_weight, ignore_index=-100):
	"""Weighted cross entropy loss (discounts certain tokens, e.g., repeated base pairs in genome)."""
	logits = logits.view(-1, logits.shape[-1])
	y = y.view(-1)
	ce = F.cross_entropy(logits, y, ignore_index=ignore_index, reduction="none")
	loss_weight = loss_weight.view(-1)
	loss_weight[y == ignore_index] = 0.0
	# TODO: Follows GPN implementation, but should we remove weight normalization?
	return (ce * (loss_weight / loss_weight.sum())).sum()


	class CaduceusPreTrainedModel(PreTrainedModel):
	"""PreTrainedModel wrapper for Caduceus backbone."""
	config_class = CaduceusConfig
	base_model_prefix = "caduceus"
	supports_gradient_checkpointing = False
	_no_split_modules = ["BiMambaWrapper"]

	def _init_weights(
	self,
	module,
	initializer_range=0.02, # Now only used for embedding layer.
	**kwargs,
	):
	"""Adapted from: https://github.com/state-spaces/mamba/blob/main/mamba_ssm/models/mixer_seq_simple.py"""

	n_layer = self.config.n_layer
	initialized_cfg = self.config.initializer_cfg if self.config.initializer_cfg is not None else {}
	rescale_prenorm_residual = initialized_cfg.get("rescale_prenorm_residual", True)
	initializer_range = initialized_cfg.get("initializer_range", initializer_range)
	n_residuals_per_layer = initialized_cfg.get("n_residuals_per_layer", 1)

	if isinstance(module, nn.Linear):
	if module.bias is not None:
	if not getattr(module.bias, "_no_reinit", False):
	nn.init.zeros_(module.bias)
	elif isinstance(module, nn.Embedding):
	nn.init.normal_(module.weight, std=initializer_range)

	if rescale_prenorm_residual:
	# Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
	# > A modified initialization which accounts for the accumulation on the residual path with model depth.
	# > Scale the weights of residual layers at initialization by a factor of 1/√N where N is the # of
	# residual layers.
	# > -- GPT-2 :: https://openai.com/blog/better-language-models/
	#
	# Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
	for name, p in module.named_parameters():
	if name in ["out_proj.weight", "fc2.weight"]:
	# Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
	# Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
	# We need to reinit p since this code could be called multiple times
	# Having just p *= scale would repeatedly scale it down
	nn.init.kaiming_uniform_(p, a=math.sqrt(5))
	with torch.no_grad():
	p /= math.sqrt(n_residuals_per_layer * n_layer)


	class Caduceus(CaduceusPreTrainedModel):
	"""Caduceus model that can be instantiated using HF patterns."""
	def __init__(self, config: CaduceusConfig, device=None, dtype=None, **kwargs):
	super().__init__(config)

	if config.rcps:
	assert config.complement_map is not None, "Complement map must be provided for RCPS."

	# Adjust vocab size and complement maps if vocab padding is set.
	if config.vocab_size % config.pad_vocab_size_multiple != 0:
	config.vocab_size += config.pad_vocab_size_multiple - (config.vocab_size % config.pad_vocab_size_multiple)
	if config.complement_map is not None and config.vocab_size > len(config.complement_map):
	for i in range(len(config.complement_map), config.vocab_size):
	config.complement_map[i] = i

	self.config = config
	factory_kwargs = {"device": device, "dtype": dtype}
	self.backbone = CaduceusMixerModel(config, factory_kwargs, kwargs)

	def forward(
	self,
	input_ids: torch.LongTensor = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[torch.Tensor, Tuple, BaseModelOutputWithNoAttention]:
	"""HF-compatible forward method."""
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	hidden_states, all_hidden_states = self.backbone(
	input_ids,
	inputs_embeds=inputs_embeds,
	output_hidden_states=output_hidden_states
	)
	if return_dict:
	return BaseModelOutputWithNoAttention(
	last_hidden_state=hidden_states,
	hidden_states=all_hidden_states if output_hidden_states else None
	)
	elif output_hidden_states:
	return hidden_states, all_hidden_states
	else:
	return hidden_states


	class CaduceusForMaskedLM(CaduceusPreTrainedModel):
	"""HF-compatible Caduceus model for masked language modeling."""

	def __init__(self, config: CaduceusConfig, device=None, dtype=None, **kwargs):
	super().__init__(config, **kwargs)
	factory_kwargs = {"device": device, "dtype": dtype}
	self.caduceus = Caduceus(config, factory_kwargs, kwargs)
	if config.rcps:
	self.lm_head = RCPSLMHead(
	complement_map=self.config.complement_map, # Use caduceus config as it might have been updated
	vocab_size=self.config.vocab_size, # Use caduceus config as it might have been updated
	true_dim=config.d_model,
	dtype=dtype
	)
	else:
	self.lm_head = nn.Linear(
	config.d_model,
	self.config.vocab_size, # Use caduceus config as it might have been updated
	bias=False,
	**factory_kwargs
	)

	# Initialize weights and apply final processing
	self.post_init()

	def get_input_embeddings(self):
	return self.caduceus.backbone.embeddings.word_embeddings

	def set_input_embeddings(self, value):
	if self.config.rcps:
	raise NotImplementedError("Setting input embeddings for RCPS LM is not supported.")
	self.caduceus.backbone.embeddings.word_embeddings = value

	def get_output_embeddings(self):
	return self.lm_head

	def set_output_embeddings(self, new_embeddings):
	"""Overrides output embeddings."""
	if self.config.rcps:
	raise NotImplementedError("Setting output embeddings for RCPS LM is not supported.")
	self.lm_head = new_embeddings

	def tie_weights(self):
	"""Tie weights, accounting for RCPS."""
	if self.config.rcps:
	self.lm_head.set_weight(self.get_input_embeddings().weight)
	else:
	super().tie_weights()

	def get_decoder(self):
	"""Get decoder (backbone) for the model."""
	return self.caduceus

	def set_decoder(self, decoder):
	"""Set decoder (backbone) for the model."""
	self.caduceus = decoder

	def forward(
	self,
	input_ids: torch.LongTensor = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	labels: Optional[torch.LongTensor] = None,
	loss_weights: Optional[torch.FloatTensor] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, MaskedLMOutput]:
	"""HF-compatible forward method."""

	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
	outputs = self.caduceus(
	input_ids=input_ids,
	inputs_embeds=inputs_embeds,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	hidden_states = outputs[0]
	logits = self.lm_head(hidden_states)
	logits = logits.float()

	loss = None
	if labels is not None:
	if loss_weights is not None:
	loss = weighted_cross_entropy(logits, labels, loss_weights, ignore_index=self.config.pad_token_id)
	else:
	loss = cross_entropy(logits, labels, ignore_index=self.config.pad_token_id)

	if not return_dict:
	output = (logits,) + outputs[1:]
	return (loss,) + output if loss is not None else output

	return MaskedLMOutput(
	loss=loss,
	logits=logits,
	hidden_states=outputs.hidden_states,
	)


	class CaduceusForSequenceClassification(CaduceusPreTrainedModel):
	def __init__(
	self,
	config: CaduceusConfig,
	pooling_strategy: str = "mean",
	conjoin_train: bool = False,
	conjoin_eval: bool = False,
	device=None,
	dtype=None,
	**kwargs):
	super().__init__(config, **kwargs)
	if pooling_strategy not in ["mean", "max", "first", "last"]:
	raise NotImplementedError(f"Pooling strategy `{pooling_strategy}` not implemented.")
	self.pooling_strategy = pooling_strategy
	factory_kwargs = {"device": device, "dtype": dtype}
	self.num_labels = kwargs.get("num_labels", config.num_labels)
	self.caduceus = Caduceus(config, factory_kwargs, kwargs)
	self.score = nn.Linear(config.d_model, self.num_labels, bias=False)

	self.conjoin_train = conjoin_train
	self.conjoin_eval = conjoin_eval

	# Initialize weights and apply final processing
	self.post_init()

	def get_input_embeddings(self):
	return self.caduceus.backbone.embeddings.word_embeddings

	def set_input_embeddings(self, value):
	if self.config.rcps:
	raise NotImplementedError("Setting input embeddings for RCPS LM is not supported.")
	self.caduceus.backbone.embeddings.word_embeddings = value

	def pool_hidden_states(self, hidden_states, sequence_length_dim=1):
	"""Pools hidden states along sequence length dimension."""
	if self.pooling_strategy == "mean": # Mean pooling along sequence length dimension
	return hidden_states.mean(dim=sequence_length_dim)
	if self.pooling_strategy == "max": # Max pooling along sequence length dimension
	return hidden_states.max(dim=sequence_length_dim).values
	if self.pooling_strategy == "last": # Use embedding of last token in the sequence
	return hidden_states.moveaxis(hidden_states, sequence_length_dim, 0)[-1, ...]
	if self.pooling_strategy == "first": # Use embedding of first token in the sequence
	return hidden_states.moveaxis(hidden_states, sequence_length_dim, 0)[0, ...]

	def forward(
	self,
	input_ids: torch.LongTensor = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	labels: Optional[torch.LongTensor] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, SequenceClassifierOutput]:
	r"""
	labels (`torch.LongTensor` of shape `(batch_size,)`, optional):
	Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
	config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
	`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
	"""
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	# Get hidden representations from the backbone
	if self.config.rcps: # Hidden states have 2 * d_model channels for RCPS
	transformer_outputs = self.caduceus(
	input_ids,
	inputs_embeds=inputs_embeds,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)
	hidden_states = torch.stack(
	[
	transformer_outputs[0][..., :self.config.d_model // 2],
	torch.flip(transformer_outputs[0][..., self.config.d_model // 2:], dims=[1, 2])
	],
	dim=-1
	)
	elif self.conjoin_train or (self.conjoin_eval and not self.training): # For conjoining / post-hoc conjoining
	assert input_ids is not None, "`input_ids` must be provided for conjoining."
	assert input_ids.ndim == 3, "`input_ids` must be 3D tensor: channels corresponds to forward and rc strands."
	transformer_outputs = self.caduceus(
	input_ids[..., 0],
	inputs_embeds=None,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)
	transformer_outputs_rc = self.caduceus(
	input_ids[..., 1],
	inputs_embeds=None,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)
	# Stack along channel dimension (dim=-1)
	hidden_states = torch.stack([transformer_outputs[0], transformer_outputs_rc[0]], dim=-1)
	else:
	transformer_outputs = self.caduceus(
	input_ids,
	inputs_embeds=None,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)
	hidden_states = transformer_outputs[0]

	# Pool and get logits
	pooled_hidden_states = self.pool_hidden_states(hidden_states)
	# Potentially run `score` twice (with parameters shared) for conjoining
	if hidden_states.ndim == 4: # bsz, seq_len, hidden_dim, 2 where last channel has the stacked fwd and rc reps
	logits_fwd = self.score(pooled_hidden_states[..., 0])
	logits_rc = self.score(pooled_hidden_states[..., 1])
	logits = (logits_fwd + logits_rc) / 2
	else:
	logits = self.score(pooled_hidden_states)

	loss = None
	if labels is not None:
	labels = labels.to(logits.device)
	if self.config.problem_type is None:
	if self.num_labels == 1:
	self.config.problem_type = "regression"
	elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
	self.config.problem_type = "single_label_classification"
	else:
	self.config.problem_type = "multi_label_classification"

	if self.config.problem_type == "regression":
	if self.num_labels == 1:
	loss = F.mse_loss(logits.squeeze(), labels.squeeze())
	else:
	loss = F.mse_loss(logits, labels)
	elif self.config.problem_type == "single_label_classification":
	loss = F.cross_entropy(logits.view(-1, self.num_labels), labels.view(-1))
	elif self.config.problem_type == "multi_label_classification":
	loss = F.binary_cross_entropy_with_logits(logits, labels)
	if not return_dict:
	output = (logits,) + transformer_outputs[1:]
	return ((loss,) + output) if loss is not None else output

	return SequenceClassifierOutput(
	loss=loss,
	logits=logits,
	hidden_states=transformer_outputs.hidden_states,
	)