Skip to content

Update granite.md #37791

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 10 commits into
base: main
Choose a base branch
from
75 changes: 61 additions & 14 deletions docs/source/en/model_doc/granite.md
Original file line number Diff line number Diff line change
Expand Up @@ -9,33 +9,52 @@ Unless required by applicable law or agreed to in writing, software distributed
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
specific language governing permissions and limitations under the License.

⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
⚠️ Note that this file is in Markdown but contains specific syntax for our doc-builder (similar to MDX) that may not be
rendered properly in your Markdown viewer.

-->

# Granite

<div class="flex flex-wrap space-x-1">
<img alt="PyTorch" src="https://img.shields.io/badge/PyTorch-DE3412?style=flat&logo=pytorch&logoColor=white">
<img alt="FlashAttention" src="https://img.shields.io/badge/%E2%9A%A1%EF%B8%8E%20FlashAttention-eae0c8?style=flat">
<img alt="SDPA" src="https://img.shields.io/badge/SDPA-DE3412?style=flat&logo=pytorch&logoColor=white">
</div>

## Overview
# Granite

[Granite](https://huggingface.co/ibm-granite) model was proposed in [Power Scheduler: A Batch Size and Token Number Agnostic Learning Rate Scheduler](https://arxiv.org/abs/2408.13359) by Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adriana Meza Soria, David D. Cox, and Rameswar Panda.
Granite model is a series of AI language models built for businesses, not just for general use. It is a super-focused tool for enterprise tasks like summarizing reports, answering questions, or generating code. It is trained on curated datasets (think finance, legal, code, and academic stuff) and filters out junk like duplicates or harmful content with their HAP detector. It comes in different sizes, from super lightweight (sub-billion parameters) up to 34 billion, so you can pick what fits your needs.

You can find all the original Granite checkpoints under the [Granite](https://huggingface.co/ibm-granite) collection.

> [!TIP]
> Click on the Granite models in the right sidebar for more examples of how to apply Granite to different language tasks.

The Granite model was proposed in [Power Scheduler: A Batch Size and Token Number Agnostic Learning Rate Scheduler](https://arxiv.org/abs/2408.13359) by Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adriana Meza Soria, David D. Cox and Rameswar Panda.
The example below demonstrates how to generate text based on an image with [`Pipeline`] or the [`AutoModel`] class.

PowerLM-3B is a 3B state-of-the-art small language model trained with the Power learning rate scheduler. It is trained on a wide range of open-source and synthetic datasets with permissive licenses. PowerLM-3B has shown promising results compared to other models in the size categories across various benchmarks, including natural language multi-choices, code generation, and math reasoning.
<hfoptions id="usage">
<hfoption id="Pipeline">

The abstract from the paper is the following:
```py
from transformers import pipeline

*Finding the optimal learning rate for language model pretraining is a challenging task.
This is not only because there is a complicated correlation between learning rate, batch size, number of training tokens, model size, and other hyperparameters but also because it is prohibitively expensive to perform a hyperparameter search for large language models with Billions or Trillions of parameters. Recent studies propose using small proxy models and small corpus to perform hyperparameter searches and transposing the optimal parameters to large models and large corpus. While the zero-shot transferability is theoretically and empirically proven for model size related hyperparameters, like depth and width, the zero-shot transfer from small corpus to large corpus is underexplored.
In this paper, we study the correlation between optimal learning rate, batch size, and number of training tokens for the recently proposed WSD scheduler. After thousands of small experiments, we found a power-law relationship between variables and demonstrated its transferability across model sizes. Based on the observation, we propose a new learning rate scheduler, Power scheduler, that is agnostic about the number of training tokens and batch size. The experiment shows that combining the Power scheduler with Maximum Update Parameterization (\mup) can consistently achieve impressive performance with one set of hyperparameters regardless of the number of training tokens, batch size, model size, and even model architecture. Our 3B dense and MoE models trained with the Power scheduler achieve comparable performance as state-of-the-art small language models.
We [open source](https://huggingface.co/collections/ibm/power-lm-66be64ae647ddf11b9808000) these pretrained models.*
pipe = pipeline(
"text-generation",
model="ibm-granite/granite-3.3-8b-instruct",
tokenizer=tokenizer,
device_map = "auto",
torch_dtype = torch.bfloat16,

Tips:
)

prompt = "Explain quantum computing in simple terms:"
output = pipe(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)
print(output[0]["generated_text"])
```

</hfoption>
<hfoption id="AutoModel">

```python
import torch
Expand All @@ -48,7 +67,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
model.eval()

# change input text as desired
# Change input text as desired
prompt = "Write a code to find the maximum value in a list of numbers."

# tokenize the text
Expand All @@ -57,14 +76,42 @@ input_tokens = tokenizer(prompt, return_tensors="pt")
output = model.generate(**input_tokens, max_new_tokens=100)
# decode output tokens into text
output = tokenizer.batch_decode(output)
# loop over the batch to print, in this example the batch size is 1
# loop over the batch to print, in this example, the batch size is 1
for i in output:
print(i)
```

This model was contributed by [mayank-mishra](https://huggingface.co/mayank-mishra).
</hfoption>
</hfoptions>

Quantization reduces the memory burden of large models by representing the weights in a lower precision. Refer to the [Quantization](../quantization/overview) overview for more available quantization backends.

The example below uses [bitsandbytes](../quantization/bitsandbytes) to only quantize the weights to int4.

```python
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_4bit=True)

tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-27b-it")
model = AutoModelForCausalLM.from_pretrained(
"google/gemma-2-27b-it",
quantization_config=quantization_config,
)

input_text = "Explain artificial intelligence to a 10 year old"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=32)
print(tokenizer.decode(outputs[0]))
```

## Notes

- Granite models support context lengths up to 128K tokens (e.g., in Granite-3.3), thanks to Rotary Position Embeddings (RoPE). For long-context tasks like document analysis, ensure your input sequences are padded appropriately to avoid attention issues.
- Don’t use the torch_dtype parameter in from_pretrained() if you’re using FlashAttention-2, as it only supports fp16 or bf16. Granite models are optimized for mixed precision, so use Automatic Mixed Precision; set fp16 or bf16 to True if
using Trainer, or use torch.autocast.

## GraniteConfig

[[autodoc]] GraniteConfig
Expand Down