Instructions to use BananaMind/normal-model-ablation with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use BananaMind/normal-model-ablation with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="BananaMind/normal-model-ablation", trust_remote_code=True)

# pip install -U transformers accelerate
# Load model directly
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("BananaMind/normal-model-ablation", trust_remote_code=True, device_map="auto")

Notebooks
Google Colab
Kaggle
Local Apps Settings

vLLM

How to use BananaMind/normal-model-ablation with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "BananaMind/normal-model-ablation"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "BananaMind/normal-model-ablation",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker

docker model run hf.co/BananaMind/normal-model-ablation

SGLang

How to use BananaMind/normal-model-ablation with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "BananaMind/normal-model-ablation" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "BananaMind/normal-model-ablation",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "BananaMind/normal-model-ablation" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "BananaMind/normal-model-ablation",
		"prompt": "Once upon a time,",
		"max_tokens": 512,
		"temperature": 0.5
	}'

Docker Model Runner
How to use BananaMind/normal-model-ablation with Docker Model Runner:
```
docker model run hf.co/BananaMind/normal-model-ablation
```

normal-model-ablation / modeling_ablation.py

BananaMindBot

Upload model + benchmarks

88b4577 verified 26 days ago

Raw History Blame Contribute Delete

7.94 kB

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_ablation import AblationConfig


class RMSNorm(nn.Module):
    def __init__(self, dim, eps):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        dt = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return (x * self.weight.float()).to(dt)


def rope_cache(T, dim, base, device):
    inv = 1.0 / (base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))
    f = torch.outer(torch.arange(T, device=device, dtype=torch.float32), inv)
    return f.cos(), f.sin()  # (T, dim/2)


def apply_rope(x, cos, sin):
    # x: (B, T, H, D)
    x1, x2 = x.float().chunk(2, dim=-1)
    c, s = cos[None, :, None, :], sin[None, :, None, :]
    return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1).to(x.dtype)


class DSAIndexer(nn.Module):
    """DeepSeek Sparse Attention "lightning indexer":
    I[t, s] = sum_j w[t, j] * ReLU(q[t, j] . k[s]).  Used to pick the top-k keys for every query."""

    def __init__(self, cfg):
        super().__init__()
        self.h, self.d, self.base = cfg.dsa_index_heads, cfg.dsa_index_dim, cfg.rope_base
        self.wq = nn.Linear(cfg.hidden_size, self.h * self.d, bias=False)
        self.wk = nn.Linear(cfg.hidden_size, self.d, bias=False)
        self.ww = nn.Linear(cfg.hidden_size, self.h, bias=False)

    def forward(self, x):
        B, T, _ = x.shape
        cos, sin = rope_cache(T, self.d, self.base, x.device)
        q = apply_rope(self.wq(x).view(B, T, self.h, self.d), cos, sin)
        k = apply_rope(self.wk(x).view(B, T, 1, self.d), cos, sin)[:, :, 0]
        w = self.ww(x) * (self.d ** -0.5)
        s = torch.einsum("bthd,bsd->bhts", q, k).relu()
        return torch.einsum("bhts,bth->bts", s, w).float()  # (B, T, T)


class Attention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.h = cfg.num_attention_heads
        self.d = cfg.hidden_size // self.h
        self.topk = cfg.dsa_topk
        self.qkv = nn.Linear(cfg.hidden_size, 3 * cfg.hidden_size, bias=False)
        self.o = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False)
        self.o.is_resid = True

    def forward(self, x, cos, sin, indexer=None, dense_warmup=False):
        B, T, C = x.shape
        q, k, v = self.qkv(x).view(B, T, 3, self.h, self.d).unbind(2)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        q, k, v = (t.transpose(1, 2) for t in (q, k, v))  # (B, H, T, D)
        aux = None
        if indexer is None or (not self.training and T <= self.topk):
            # standard causal attention (DSA with T <= topk selects every causal token, i.e. is identical)
            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        else:
            idx_scores = indexer(x.detach())  # indexer is trained only through its own KL loss
            causal = torch.ones(T, T, dtype=torch.bool, device=x.device).tril()
            dense = dense_warmup or T <= self.topk
            if dense:
                sel = causal.expand(B, T, T)
                y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
            else:
                top = idx_scores.masked_fill(~causal, float("-inf")).topk(self.topk, dim=-1).indices
                sel = torch.zeros(B, T, T, dtype=torch.bool, device=x.device).scatter_(-1, top, True) & causal
                y = F.scaled_dot_product_attention(q, k, v, attn_mask=sel[:, None])
            if self.training:
                # indexer distillation: KL(main attention (head-avg, detached) || indexer) over selected tokens
                with torch.no_grad():
                    s = (q @ k.transpose(-1, -2)).float() * (self.d ** -0.5)
                    p = s.masked_fill(~sel[:, None], float("-inf")).softmax(-1).mean(1)
                logq = idx_scores.masked_fill(~sel, float("-inf")).log_softmax(-1).masked_fill(~sel, 0.0)
                aux = (p * (p.clamp_min(1e-9).log() - logq)).sum(-1).mean()
        return self.o(y.transpose(1, 2).reshape(B, T, C)), aux


class MLP(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.gate = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
        self.up = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
        self.down = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)
        self.down.is_resid = True

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))


class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm1 = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.attn = Attention(cfg)
        self.norm2 = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.mlp = MLP(cfg)

    def forward(self, x, cos, sin, indexer=None, dense_warmup=False):
        a, aux = self.attn(self.norm1(x), cos, sin, indexer, dense_warmup)
        x = x + a
        return x + self.mlp(self.norm2(x)), aux


class AblationPreTrainedModel(PreTrainedModel):
    config_class = AblationConfig
    base_model_prefix = "ablation"

    def _init_weights(self, m):
        std = self.config.initializer_range
        if isinstance(m, nn.Linear):
            if getattr(m, "is_resid", False):
                std = std / math.sqrt(2 * self.config.num_hidden_layers)
            nn.init.normal_(m.weight, 0.0, std)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, 0.0, std)
        elif isinstance(m, RMSNorm):
            nn.init.ones_(m.weight)


class AblationForCausalLM(AblationPreTrainedModel, GenerationMixin):
    def __init__(self, config):
        super().__init__(config)
        self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([Block(config) for _ in range(config.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        # indexers are created last so the shared weights get identical init for both attention types
        use_dsa = config.attention_type == "dsa"
        self.indexers = nn.ModuleList([DSAIndexer(config) for _ in range(config.num_hidden_layers)] if use_dsa else [])
        self.dsa_warmup = False  # set by the trainer: dense attention while the indexer warms up
        self.aux_loss = None
        self.post_init()

    def get_input_embeddings(self):
        return self.embed

    def set_input_embeddings(self, value):
        self.embed = value

    def get_output_embeddings(self):
        return None

    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
        x = self.embed(input_ids)
        hd = self.config.hidden_size // self.config.num_attention_heads
        cos, sin = rope_cache(x.shape[1], hd, self.config.rope_base, x.device)
        aux_total = None
        for i, layer in enumerate(self.layers):
            idx = self.indexers[i] if len(self.indexers) else None
            x, aux = layer(x, cos, sin, idx, self.dsa_warmup)
            if aux is not None:
                aux_total = aux if aux_total is None else aux_total + aux
        self.aux_loss = aux_total
        logits = F.linear(self.norm(x), self.embed.weight)  # tied embeddings
        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits[:, :-1].float().reshape(-1, logits.size(-1)), labels[:, 1:].reshape(-1), ignore_index=-100
            )
        return CausalLMOutputWithPast(loss=loss, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        # no KV cache: recompute the full sequence every step
        return {"input_ids": input_ids, "use_cache": False}

	import math

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from transformers import PreTrainedModel, GenerationMixin
	from transformers.modeling_outputs import CausalLMOutputWithPast

	from .configuration_ablation import AblationConfig


	class RMSNorm(nn.Module):
	def __init__(self, dim, eps):
	super().__init__()
	self.eps = eps
	self.weight = nn.Parameter(torch.ones(dim))

	def forward(self, x):
	dt = x.dtype
	x = x.float()
	x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
	return (x * self.weight.float()).to(dt)


	def rope_cache(T, dim, base, device):
	inv = 1.0 / (base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))
	f = torch.outer(torch.arange(T, device=device, dtype=torch.float32), inv)
	return f.cos(), f.sin() # (T, dim/2)


	def apply_rope(x, cos, sin):
	# x: (B, T, H, D)
	x1, x2 = x.float().chunk(2, dim=-1)
	c, s = cos[None, :, None, :], sin[None, :, None, :]
	return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1).to(x.dtype)


	class DSAIndexer(nn.Module):
	"""DeepSeek Sparse Attention "lightning indexer":
	I[t, s] = sum_j w[t, j] * ReLU(q[t, j] . k[s]). Used to pick the top-k keys for every query."""

	def __init__(self, cfg):
	super().__init__()
	self.h, self.d, self.base = cfg.dsa_index_heads, cfg.dsa_index_dim, cfg.rope_base
	self.wq = nn.Linear(cfg.hidden_size, self.h * self.d, bias=False)
	self.wk = nn.Linear(cfg.hidden_size, self.d, bias=False)
	self.ww = nn.Linear(cfg.hidden_size, self.h, bias=False)

	def forward(self, x):
	B, T, _ = x.shape
	cos, sin = rope_cache(T, self.d, self.base, x.device)
	q = apply_rope(self.wq(x).view(B, T, self.h, self.d), cos, sin)
	k = apply_rope(self.wk(x).view(B, T, 1, self.d), cos, sin)[:, :, 0]
	w = self.ww(x) * (self.d ** -0.5)
	s = torch.einsum("bthd,bsd->bhts", q, k).relu()
	return torch.einsum("bhts,bth->bts", s, w).float() # (B, T, T)


	class Attention(nn.Module):
	def __init__(self, cfg):
	super().__init__()
	self.h = cfg.num_attention_heads
	self.d = cfg.hidden_size // self.h
	self.topk = cfg.dsa_topk
	self.qkv = nn.Linear(cfg.hidden_size, 3 * cfg.hidden_size, bias=False)
	self.o = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False)
	self.o.is_resid = True

	def forward(self, x, cos, sin, indexer=None, dense_warmup=False):
	B, T, C = x.shape
	q, k, v = self.qkv(x).view(B, T, 3, self.h, self.d).unbind(2)
	q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
	q, k, v = (t.transpose(1, 2) for t in (q, k, v)) # (B, H, T, D)
	aux = None
	if indexer is None or (not self.training and T <= self.topk):
	# standard causal attention (DSA with T <= topk selects every causal token, i.e. is identical)
	y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
	else:
	idx_scores = indexer(x.detach()) # indexer is trained only through its own KL loss
	causal = torch.ones(T, T, dtype=torch.bool, device=x.device).tril()
	dense = dense_warmup or T <= self.topk
	if dense:
	sel = causal.expand(B, T, T)
	y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
	else:
	top = idx_scores.masked_fill(~causal, float("-inf")).topk(self.topk, dim=-1).indices
	sel = torch.zeros(B, T, T, dtype=torch.bool, device=x.device).scatter_(-1, top, True) & causal
	y = F.scaled_dot_product_attention(q, k, v, attn_mask=sel[:, None])
	if self.training:
	# indexer distillation: KL(main attention (head-avg, detached) \|\| indexer) over selected tokens
	with torch.no_grad():
	s = (q @ k.transpose(-1, -2)).float() * (self.d ** -0.5)
	p = s.masked_fill(~sel[:, None], float("-inf")).softmax(-1).mean(1)
	logq = idx_scores.masked_fill(~sel, float("-inf")).log_softmax(-1).masked_fill(~sel, 0.0)
	aux = (p * (p.clamp_min(1e-9).log() - logq)).sum(-1).mean()
	return self.o(y.transpose(1, 2).reshape(B, T, C)), aux


	class MLP(nn.Module):
	def __init__(self, cfg):
	super().__init__()
	self.gate = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
	self.up = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
	self.down = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)
	self.down.is_resid = True

	def forward(self, x):
	return self.down(F.silu(self.gate(x)) * self.up(x))


	class Block(nn.Module):
	def __init__(self, cfg):
	super().__init__()
	self.norm1 = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
	self.attn = Attention(cfg)
	self.norm2 = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
	self.mlp = MLP(cfg)

	def forward(self, x, cos, sin, indexer=None, dense_warmup=False):
	a, aux = self.attn(self.norm1(x), cos, sin, indexer, dense_warmup)
	x = x + a
	return x + self.mlp(self.norm2(x)), aux


	class AblationPreTrainedModel(PreTrainedModel):
	config_class = AblationConfig
	base_model_prefix = "ablation"

	def _init_weights(self, m):
	std = self.config.initializer_range
	if isinstance(m, nn.Linear):
	if getattr(m, "is_resid", False):
	std = std / math.sqrt(2 * self.config.num_hidden_layers)
	nn.init.normal_(m.weight, 0.0, std)
	elif isinstance(m, nn.Embedding):
	nn.init.normal_(m.weight, 0.0, std)
	elif isinstance(m, RMSNorm):
	nn.init.ones_(m.weight)


	class AblationForCausalLM(AblationPreTrainedModel, GenerationMixin):
	def __init__(self, config):
	super().__init__(config)
	self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
	self.layers = nn.ModuleList([Block(config) for _ in range(config.num_hidden_layers)])
	self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
	# indexers are created last so the shared weights get identical init for both attention types
	use_dsa = config.attention_type == "dsa"
	self.indexers = nn.ModuleList([DSAIndexer(config) for _ in range(config.num_hidden_layers)] if use_dsa else [])
	self.dsa_warmup = False # set by the trainer: dense attention while the indexer warms up
	self.aux_loss = None
	self.post_init()

	def get_input_embeddings(self):
	return self.embed

	def set_input_embeddings(self, value):
	self.embed = value

	def get_output_embeddings(self):
	return None

	def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
	x = self.embed(input_ids)
	hd = self.config.hidden_size // self.config.num_attention_heads
	cos, sin = rope_cache(x.shape[1], hd, self.config.rope_base, x.device)
	aux_total = None
	for i, layer in enumerate(self.layers):
	idx = self.indexers[i] if len(self.indexers) else None
	x, aux = layer(x, cos, sin, idx, self.dsa_warmup)
	if aux is not None:
	aux_total = aux if aux_total is None else aux_total + aux
	self.aux_loss = aux_total
	logits = F.linear(self.norm(x), self.embed.weight) # tied embeddings
	loss = None
	if labels is not None:
	loss = F.cross_entropy(
	logits[:, :-1].float().reshape(-1, logits.size(-1)), labels[:, 1:].reshape(-1), ignore_index=-100
	)
	return CausalLMOutputWithPast(loss=loss, logits=logits)

	def prepare_inputs_for_generation(self, input_ids, **kwargs):
	# no KV cache: recompute the full sequence every step
	return {"input_ids": input_ids, "use_cache": False}