berna-r8 / code /modeling_berna_v2.py

Upload folder using huggingface_hub

d997548 verified 8 days ago

14.6 kB

"""
Berna R5-integrated Model Architecture.
Adds CellAwareMLP with DNA-based gates.
"""
import math
from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

from transformers import PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import CausalLMOutputWithPast
from cell_memory import CellMemory


class BernaConfig(PretrainedConfig):
    model_type = "berna"
    def __init__(
        self,
        vocab_size=32000,
        hidden_size=1536,
        intermediate_size=6144,
        num_hidden_layers=28,
        num_attention_heads=16,
        num_key_value_heads=4,
        hidden_act="silu",
        max_position_embeddings=4096,
        initializer_range=0.02,
        rms_norm_eps=1e-6,
        rope_theta=10000.0,
        attention_bias=False,
        mlp_bias=False,
        tie_word_embeddings=False,
        pad_token_id=0,
        bos_token_id=1,
        eos_token_id=2,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.hidden_act = hidden_act
        self.max_position_embeddings = max_position_embeddings
        self.initializer_range = initializer_range
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.attention_bias = attention_bias
        self.mlp_bias = mlp_bias
        self.pad_token_id = pad_token_id
        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id
        super().__init__(
            tie_word_embeddings=tie_word_embeddings,
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            **kwargs,
        )


class RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)


class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=4096, base=10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.max_seq_len_cached = max_position_embeddings
        self._set_cos_sin_cache(max_position_embeddings)

    def _set_cos_sin_cache(self, seq_len):
        self.max_seq_len_cached = seq_len
        t = torch.arange(seq_len, dtype=self.inv_freq.dtype, device=self.inv_freq.device)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos(), persistent=False)
        self.register_buffer("sin_cached", emb.sin(), persistent=False)

    def forward(self, x, seq_len=None):
        if seq_len is None:
            seq_len = x.shape[1]
        if seq_len > self.max_seq_len_cached:
            self._set_cos_sin_cache(seq_len)
        return (
            self.cos_cached[:seq_len].to(dtype=x.dtype),
            self.sin_cached[:seq_len].to(dtype=x.dtype),
        )


def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2:]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin):
    cos = cos.unsqueeze(0).unsqueeze(0)
    sin = sin.unsqueeze(0).unsqueeze(0)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


class Attention(nn.Module):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads

        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)

        self.rotary_emb = RotaryEmbedding(
            self.head_dim,
            max_position_embeddings=config.max_position_embeddings,
            base=config.rope_theta,
        )

    def forward(self, hidden_states, attention_mask=None):
        bsz, q_len, _ = hidden_states.size()
        query = self.q_proj(hidden_states)
        key = self.k_proj(hidden_states)
        value = self.v_proj(hidden_states)

        query = query.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        key = key.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
        value = value.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

        cos, sin = self.rotary_emb(value, seq_len=q_len)
        query, key = apply_rotary_pos_emb(query, key, cos, sin)

        if self.num_key_value_groups > 1:
            key = key.repeat_interleave(self.num_key_value_groups, dim=1)
            value = value.repeat_interleave(self.num_key_value_groups, dim=1)

        attn_output = F.scaled_dot_product_attention(
            query, key, value,
            attn_mask=attention_mask,
            is_causal=(attention_mask is None),
        )
        attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, self.hidden_size)
        return self.o_proj(attn_output)


class CellAwareMLP(nn.Module):
    """MLP with per-cell DNA-based gates on intermediate neurons.
    Init: DNA has a-gene = 1.0 → gate = 1.0 → identity (checkpoint-compatible).
    On split: DNA mutates → gate drifts → cell's neuron range scaled.
    """
    def __init__(self, config, layer_idx, cells, dnas):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias)
        self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_bias)
        self.experience = nn.Parameter(torch.zeros(config.intermediate_size))
        self.layer_idx = layer_idx
        self.cells = cells        # list[Cell]
        self.dnas = dnas          # dict[str, DNA]
        self._gate_cache = None
        self._slot_by_cell = {}   # cell_id -> slot_idx
        self._cache_version = -1
        self._last_inter_detached = None   # for cell measurement
        self._measure_now = False

    def _gate_version(self):
        return tuple(sorted((c.cell_id, c.state) for c in self.cells))

    def _rebuild_gate_mask(self, device, dtype, interm_size):
        mask = torch.ones(1, 1, interm_size, device=device, dtype=dtype)
        from cells.dna_kernel import signal_route
        for c in self.cells:
            if c.state != "active":
                continue
            dna = self.dnas.get(c.cell_id)
            g = signal_route(dna) if dna is not None else 1.0
            mask[..., c.neuron_start:c.neuron_end] = g
        return mask

    def assign_slots(self, mapping):
        """mapping: {cell_id: slot_idx}. Set externally by CellSystem."""
        self._slot_by_cell = dict(mapping)

    def get_cell_slots(self):
        """Return (slot_ids [N], slot_mask [N]) for active cells."""
        ids = []
        for c in self.cells:
            if c.state != "active":
                continue
            sid = self._slot_by_cell.get(c.cell_id, None)
            if sid is None:
                continue
            ids.append(sid)
        if not ids:
            import torch as _t
            return (_t.zeros(0, dtype=_t.long),
                    _t.zeros(0, dtype=_t.float))
        import torch as _t
        return (_t.tensor(ids, dtype=_t.long),
                _t.ones(len(ids), dtype=_t.float))

    def forward(self, x):
        inter_raw = F.silu(self.gate_proj(x)) * self.up_proj(x)
        # store detached version ONLY when measuring (saves 5.6 GB)
        if self._measure_now:
            self._last_inter_detached = inter_raw.detach()

        version = self._gate_version()
        if self._gate_cache is None or self._cache_version != version:
            self._gate_cache = self._rebuild_gate_mask(
                inter_raw.device, inter_raw.dtype, inter_raw.shape[-1]
            )
            self._cache_version = version

        inter_gated = inter_raw * self._gate_cache
        # R8: apply native experience
        if hasattr(self, "experience"):
            inter_gated = inter_gated * (1.0 + self.experience.to(inter_gated.dtype))
        return self.down_proj(inter_gated)


class DecoderLayer(nn.Module):
    def __init__(self, config, layer_idx, cells, dnas):
        super().__init__()
        self.self_attn = Attention(config, layer_idx)
        self.mlp = CellAwareMLP(config, layer_idx, cells, dnas)
        self.cell_memory = CellMemory(config.hidden_size, max_cells=256, n_heads=config.num_attention_heads)
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_cell_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(self, hidden_states, attention_mask=None, cell_slot_ids=None, cell_slot_mask=None):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(hidden_states, attention_mask=attention_mask)
        hidden_states = residual + hidden_states

        # Cell memory (residual, identity at start)
        if cell_slot_ids is not None and cell_slot_ids.numel() > 0:
            residual = hidden_states
            mem_in = self.post_cell_norm(hidden_states)
            mem_out = self.cell_memory(mem_in, cell_slot_ids, cell_slot_mask)
            hidden_states = residual + mem_out

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
        return hidden_states


class BernaModel(nn.Module):
    def __init__(self, config, cells_by_layer=None, dnas=None):
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        cells_by_layer = cells_by_layer or {}
        dnas = dnas or {}
        self.layers = nn.ModuleList([
            DecoderLayer(config, i, cells_by_layer.get(i, []), dnas)
            for i in range(config.num_hidden_layers)
        ])
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.gradient_checkpointing = False

    def forward(self, input_ids):
        hidden_states = self.embed_tokens(input_ids)
        for layer in self.layers:
            slot_ids, slot_mask = layer.mlp.get_cell_slots()
            if self.training and self.gradient_checkpointing:
                hidden_states = checkpoint(
                    layer, hidden_states,
                    cell_slot_ids=slot_ids, cell_slot_mask=slot_mask,
                    use_reentrant=False,
                )
            else:
                hidden_states = layer(
                    hidden_states,
                    cell_slot_ids=slot_ids, cell_slot_mask=slot_mask,
                )
        hidden_states = self.norm(hidden_states)
        return hidden_states


class BernaForCausalLM(PreTrainedModel):
    supports_gradient_checkpointing = True
    config_class = BernaConfig

    def _set_gradient_checkpointing(self, module, value=False):
        if hasattr(module, "gradient_checkpointing"):
            module.gradient_checkpointing = value

    def __init__(self, config, cells_by_layer=None, dnas=None):
        super().__init__(config)
        self.model = BernaModel(config, cells_by_layer, dnas)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def forward(self, input_ids, labels=None, **kwargs):
        hidden_states = self.model(input_ids)
        logits = self.lm_head(hidden_states)
        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )
        return CausalLMOutputWithPast(loss=loss, logits=logits)

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=50, temperature=0.8, top_k=50, **kwargs):
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = input_ids[:, -self.config.max_position_embeddings:]
            outputs = self.forward(idx_cond)
            logits = outputs.logits[:, -1, :]
            logits = torch.nan_to_num(logits, nan=0.0, posinf=0.0, neginf=0.0)
            if temperature > 0:
                logits = logits / temperature
            if top_k > 0:
                k = min(top_k, logits.size(-1))
                top_vals, top_idx = torch.topk(logits, k, dim=-1)
                new_logits = torch.full_like(logits, float("-inf"))
                new_logits.scatter_(-1, top_idx, top_vals)
                logits = new_logits
            probs = F.softmax(logits, dim=-1)
            if torch.isnan(probs).any():
                probs = torch.nan_to_num(probs, nan=0.0)
                probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-8)
            next_token = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat([input_ids, next_token], dim=-1)
        return input_ids

	"""
	Berna R5-integrated Model Architecture.
	Adds CellAwareMLP with DNA-based gates.
	"""
	import math
	from typing import Optional
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from torch.utils.checkpoint import checkpoint

	from transformers import PreTrainedModel, PretrainedConfig
	from transformers.modeling_outputs import CausalLMOutputWithPast
	from cell_memory import CellMemory


	class BernaConfig(PretrainedConfig):
	model_type = "berna"
	def __init__(
	self,
	vocab_size=32000,
	hidden_size=1536,
	intermediate_size=6144,
	num_hidden_layers=28,
	num_attention_heads=16,
	num_key_value_heads=4,
	hidden_act="silu",
	max_position_embeddings=4096,
	initializer_range=0.02,
	rms_norm_eps=1e-6,
	rope_theta=10000.0,
	attention_bias=False,
	mlp_bias=False,
	tie_word_embeddings=False,
	pad_token_id=0,
	bos_token_id=1,
	eos_token_id=2,
	**kwargs,
	):
	self.vocab_size = vocab_size
	self.hidden_size = hidden_size
	self.intermediate_size = intermediate_size
	self.num_hidden_layers = num_hidden_layers
	self.num_attention_heads = num_attention_heads
	self.num_key_value_heads = num_key_value_heads
	self.hidden_act = hidden_act
	self.max_position_embeddings = max_position_embeddings
	self.initializer_range = initializer_range
	self.rms_norm_eps = rms_norm_eps
	self.rope_theta = rope_theta
	self.attention_bias = attention_bias
	self.mlp_bias = mlp_bias
	self.pad_token_id = pad_token_id
	self.bos_token_id = bos_token_id
	self.eos_token_id = eos_token_id
	super().__init__(
	tie_word_embeddings=tie_word_embeddings,
	pad_token_id=pad_token_id,
	bos_token_id=bos_token_id,
	eos_token_id=eos_token_id,
	**kwargs,
	)


	class RMSNorm(nn.Module):
	def __init__(self, hidden_size, eps=1e-6):
	super().__init__()
	self.weight = nn.Parameter(torch.ones(hidden_size))
	self.variance_epsilon = eps

	def forward(self, hidden_states):
	input_dtype = hidden_states.dtype
	hidden_states = hidden_states.to(torch.float32)
	variance = hidden_states.pow(2).mean(-1, keepdim=True)
	hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
	return self.weight * hidden_states.to(input_dtype)


	class RotaryEmbedding(nn.Module):
	def __init__(self, dim, max_position_embeddings=4096, base=10000.0):
	super().__init__()
	inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
	self.register_buffer("inv_freq", inv_freq, persistent=False)
	self.max_seq_len_cached = max_position_embeddings
	self._set_cos_sin_cache(max_position_embeddings)

	def _set_cos_sin_cache(self, seq_len):
	self.max_seq_len_cached = seq_len
	t = torch.arange(seq_len, dtype=self.inv_freq.dtype, device=self.inv_freq.device)
	freqs = torch.outer(t, self.inv_freq)
	emb = torch.cat((freqs, freqs), dim=-1)
	self.register_buffer("cos_cached", emb.cos(), persistent=False)
	self.register_buffer("sin_cached", emb.sin(), persistent=False)

	def forward(self, x, seq_len=None):
	if seq_len is None:
	seq_len = x.shape[1]
	if seq_len > self.max_seq_len_cached:
	self._set_cos_sin_cache(seq_len)
	return (
	self.cos_cached[:seq_len].to(dtype=x.dtype),
	self.sin_cached[:seq_len].to(dtype=x.dtype),
	)


	def rotate_half(x):
	x1 = x[..., : x.shape[-1] // 2]
	x2 = x[..., x.shape[-1] // 2:]
	return torch.cat((-x2, x1), dim=-1)


	def apply_rotary_pos_emb(q, k, cos, sin):
	cos = cos.unsqueeze(0).unsqueeze(0)
	sin = sin.unsqueeze(0).unsqueeze(0)
	q_embed = (q * cos) + (rotate_half(q) * sin)
	k_embed = (k * cos) + (rotate_half(k) * sin)
	return q_embed, k_embed


	class Attention(nn.Module):
	def __init__(self, config, layer_idx):
	super().__init__()
	self.config = config
	self.layer_idx = layer_idx
	self.hidden_size = config.hidden_size
	self.num_heads = config.num_attention_heads
	self.head_dim = self.hidden_size // self.num_heads
	self.num_key_value_heads = config.num_key_value_heads
	self.num_key_value_groups = self.num_heads // self.num_key_value_heads

	self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
	self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
	self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
	self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)

	self.rotary_emb = RotaryEmbedding(
	self.head_dim,
	max_position_embeddings=config.max_position_embeddings,
	base=config.rope_theta,
	)

	def forward(self, hidden_states, attention_mask=None):
	bsz, q_len, _ = hidden_states.size()
	query = self.q_proj(hidden_states)
	key = self.k_proj(hidden_states)
	value = self.v_proj(hidden_states)

	query = query.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
	key = key.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
	value = value.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

	cos, sin = self.rotary_emb(value, seq_len=q_len)
	query, key = apply_rotary_pos_emb(query, key, cos, sin)

	if self.num_key_value_groups > 1:
	key = key.repeat_interleave(self.num_key_value_groups, dim=1)
	value = value.repeat_interleave(self.num_key_value_groups, dim=1)

	attn_output = F.scaled_dot_product_attention(
	query, key, value,
	attn_mask=attention_mask,
	is_causal=(attention_mask is None),
	)
	attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, self.hidden_size)
	return self.o_proj(attn_output)


	class CellAwareMLP(nn.Module):
	"""MLP with per-cell DNA-based gates on intermediate neurons.

	Init: DNA has a-gene = 1.0 → gate = 1.0 → identity (checkpoint-compatible).
	On split: DNA mutates → gate drifts → cell's neuron range scaled.
	"""
	def __init__(self, config, layer_idx, cells, dnas):
	super().__init__()
	self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias)
	self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias)
	self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_bias)
	self.experience = nn.Parameter(torch.zeros(config.intermediate_size))
	self.layer_idx = layer_idx
	self.cells = cells # list[Cell]
	self.dnas = dnas # dict[str, DNA]
	self._gate_cache = None
	self._slot_by_cell = {} # cell_id -> slot_idx
	self._cache_version = -1
	self._last_inter_detached = None # for cell measurement
	self._measure_now = False

	def _gate_version(self):
	return tuple(sorted((c.cell_id, c.state) for c in self.cells))

	def _rebuild_gate_mask(self, device, dtype, interm_size):
	mask = torch.ones(1, 1, interm_size, device=device, dtype=dtype)
	from cells.dna_kernel import signal_route
	for c in self.cells:
	if c.state != "active":
	continue
	dna = self.dnas.get(c.cell_id)
	g = signal_route(dna) if dna is not None else 1.0
	mask[..., c.neuron_start:c.neuron_end] = g
	return mask

	def assign_slots(self, mapping):
	"""mapping: {cell_id: slot_idx}. Set externally by CellSystem."""
	self._slot_by_cell = dict(mapping)

	def get_cell_slots(self):
	"""Return (slot_ids [N], slot_mask [N]) for active cells."""
	ids = []
	for c in self.cells:
	if c.state != "active":
	continue
	sid = self._slot_by_cell.get(c.cell_id, None)
	if sid is None:
	continue
	ids.append(sid)
	if not ids:
	import torch as _t
	return (_t.zeros(0, dtype=_t.long),
	_t.zeros(0, dtype=_t.float))
	import torch as _t
	return (_t.tensor(ids, dtype=_t.long),
	_t.ones(len(ids), dtype=_t.float))

	def forward(self, x):
	inter_raw = F.silu(self.gate_proj(x)) * self.up_proj(x)
	# store detached version ONLY when measuring (saves 5.6 GB)
	if self._measure_now:
	self._last_inter_detached = inter_raw.detach()

	version = self._gate_version()
	if self._gate_cache is None or self._cache_version != version:
	self._gate_cache = self._rebuild_gate_mask(
	inter_raw.device, inter_raw.dtype, inter_raw.shape[-1]
	)
	self._cache_version = version

	inter_gated = inter_raw * self._gate_cache
	# R8: apply native experience
	if hasattr(self, "experience"):
	inter_gated = inter_gated * (1.0 + self.experience.to(inter_gated.dtype))
	return self.down_proj(inter_gated)


	class DecoderLayer(nn.Module):
	def __init__(self, config, layer_idx, cells, dnas):
	super().__init__()
	self.self_attn = Attention(config, layer_idx)
	self.mlp = CellAwareMLP(config, layer_idx, cells, dnas)
	self.cell_memory = CellMemory(config.hidden_size, max_cells=256, n_heads=config.num_attention_heads)
	self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
	self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
	self.post_cell_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

	def forward(self, hidden_states, attention_mask=None, cell_slot_ids=None, cell_slot_mask=None):
	residual = hidden_states
	hidden_states = self.input_layernorm(hidden_states)
	hidden_states = self.self_attn(hidden_states, attention_mask=attention_mask)
	hidden_states = residual + hidden_states

	# Cell memory (residual, identity at start)
	if cell_slot_ids is not None and cell_slot_ids.numel() > 0:
	residual = hidden_states
	mem_in = self.post_cell_norm(hidden_states)
	mem_out = self.cell_memory(mem_in, cell_slot_ids, cell_slot_mask)
	hidden_states = residual + mem_out

	residual = hidden_states
	hidden_states = self.post_attention_layernorm(hidden_states)
	hidden_states = self.mlp(hidden_states)
	hidden_states = residual + hidden_states
	return hidden_states


	class BernaModel(nn.Module):
	def __init__(self, config, cells_by_layer=None, dnas=None):
	super().__init__()
	self.config = config
	self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
	cells_by_layer = cells_by_layer or {}
	dnas = dnas or {}
	self.layers = nn.ModuleList([
	DecoderLayer(config, i, cells_by_layer.get(i, []), dnas)
	for i in range(config.num_hidden_layers)
	])
	self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
	self.gradient_checkpointing = False

	def forward(self, input_ids):
	hidden_states = self.embed_tokens(input_ids)
	for layer in self.layers:
	slot_ids, slot_mask = layer.mlp.get_cell_slots()
	if self.training and self.gradient_checkpointing:
	hidden_states = checkpoint(
	layer, hidden_states,
	cell_slot_ids=slot_ids, cell_slot_mask=slot_mask,
	use_reentrant=False,
	)
	else:
	hidden_states = layer(
	hidden_states,
	cell_slot_ids=slot_ids, cell_slot_mask=slot_mask,
	)
	hidden_states = self.norm(hidden_states)
	return hidden_states


	class BernaForCausalLM(PreTrainedModel):
	supports_gradient_checkpointing = True
	config_class = BernaConfig

	def _set_gradient_checkpointing(self, module, value=False):
	if hasattr(module, "gradient_checkpointing"):
	module.gradient_checkpointing = value

	def __init__(self, config, cells_by_layer=None, dnas=None):
	super().__init__(config)
	self.model = BernaModel(config, cells_by_layer, dnas)
	self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
	self.post_init()

	def forward(self, input_ids, labels=None, **kwargs):
	hidden_states = self.model(input_ids)
	logits = self.lm_head(hidden_states)
	loss = None
	if labels is not None:
	shift_logits = logits[..., :-1, :].contiguous()
	shift_labels = labels[..., 1:].contiguous()
	loss = F.cross_entropy(
	shift_logits.view(-1, shift_logits.size(-1)),
	shift_labels.view(-1),
	ignore_index=-100,
	)
	return CausalLMOutputWithPast(loss=loss, logits=logits)

	@torch.no_grad()
	def generate(self, input_ids, max_new_tokens=50, temperature=0.8, top_k=50, **kwargs):
	self.eval()
	for _ in range(max_new_tokens):
	idx_cond = input_ids[:, -self.config.max_position_embeddings:]
	outputs = self.forward(idx_cond)
	logits = outputs.logits[:, -1, :]
	logits = torch.nan_to_num(logits, nan=0.0, posinf=0.0, neginf=0.0)
	if temperature > 0:
	logits = logits / temperature
	if top_k > 0:
	k = min(top_k, logits.size(-1))
	top_vals, top_idx = torch.topk(logits, k, dim=-1)
	new_logits = torch.full_like(logits, float("-inf"))
	new_logits.scatter_(-1, top_idx, top_vals)
	logits = new_logits
	probs = F.softmax(logits, dim=-1)
	if torch.isnan(probs).any():
	probs = torch.nan_to_num(probs, nan=0.0)
	probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-8)
	next_token = torch.multinomial(probs, num_samples=1)
	input_ids = torch.cat([input_ids, next_token], dim=-1)
	return input_ids