Model Summary Granite-SWASH-2B is a 2B parameter decoder-only dense language model with per-layer Sliding Window Attention (SWA) and learnable attention sinks (LSE-based). Trained on a mix of open source and proprietary data, it is a broad, general-purpose, English-language base model that serves as an early exploration and small-scale preview for upcoming Granite series model releases.
Granite-SWASH-2B is based on a decoder-only dense transformer architecture. Core components of this architecture are: GQA, RoPE, MLP with SwiGLU, RMSNorm, shared input/output embeddings, Sliding Window Attention, and Attention Sinks.
Parameter
Value
Hidden size
2560
Layers
24
Attention heads
20 (4 KV heads, GQA)
Intermediate size
8192 (SwiGLU)
Max position embeddings
8192
Sliding window size
128
Vocabulary
100,352
Parameters
2,144,463,840
Attention Pattern: 7 full-attention layers (indices 0, 3, 7, 11, 15, 19, 23) and 17 sliding-window layers (window=128). All layers have learnable per-head attention sinks.
Where lse is the log-sum-exp of attention scores, and sinks is a learnable per-head parameter. Flash Attention 3 and 4 implementations natively incorporate and compute the sink token under the hood. Note that SDPA does not express sink tokens or sink_scale, and so is not a supported backend.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained(
"ibm-granite/granite-swash-2b",
dtype=torch.bfloat16,
device_map="auto",
attn_implementation="flash_attention_3", # or "flash_attention_4" or "eager", NOT "sdpa"
)
tokenizer = AutoTokenizer.from_pretrained("ibm-granite/granite-swash-2b")
inputs = tokenizer("The capital of France is", return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=50, do_sample=False)
print(tokenizer.decode(out[0], skip_special_tokens=True))
The use of Large Language Models involves risks and ethical considerations people must be aware of, including but not limited to: bias and fairness, misinformation, and autonomous decision-making. Granite-SWASH-2B is not an exception in this regard. Even though this model is suited for multiple generative AI tasks, it has not undergone any safety alignment and it may produce problematic outputs. Additionally, it remains uncertain whether smaller models might exhibit increased susceptibility to hallucination in generation scenarios by copying text verbatim from the training dataset due to their reduced sizes and memorization capacities. This aspect is currently an active area of research, and we anticipate more rigorous exploration, comprehension, and mitigations in this domain. Regarding ethics, a latent risk associated with all Large Language Models is their malicious utilization. We urge the community to use Granite-SWASH-2B model with ethical intentions and in a responsible way. To enhance safety in enterprise deployments, we recommend using Granite Language models alongside Granite Guardian, a model designed to detect and flag risks in inputs and outputs across key dimensions outlined in the IBM AI Risk Atlas.
import os
os.environ["HF_ENDPOINT"] = "https://freeweights.allternit.com"
from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("ibm-granite/granite-swash-2b", revision="af1e3227100b61088eead48389ab5409b5d0e39c")
model = AutoModelForCausalLM.from_pretrained("ibm-granite/granite-swash-2b", revision="af1e3227100b61088eead48389ab5409b5d0e39c", device_map="auto")
