Implementation:Huggingface Open r1 Get Model and Tokenizer
| Knowledge Sources | |
|---|---|
| Domains | NLP, Model_Architecture |
| Last Updated | 2026-02-08 00:00 GMT |
Overview
Concrete tool for loading pretrained causal language models and tokenizers with quantization and attention backend configuration provided by Open-R1.
Description
Two factory functions get_tokenizer and get_model that wrap HuggingFace Transformers' AutoTokenizer and AutoModelForCausalLM with Open-R1-specific defaults. get_tokenizer loads the tokenizer and optionally overrides the chat template. get_model handles dtype resolution, quantization config, attention implementation, and KV cache settings.
get_tokenizerloads anAutoTokenizerfrom the specified model path and revision, respectstrust_remote_code, and optionally overrides the chat template if one is set in the training arguments.get_modelresolves thetorch_dtype(handling the"auto"string or mapping to atorchdtype), obtains a quantization config via TRL'sget_quantization_config, builds a keyword dictionary that includes attention implementation, caching behavior (disabled when gradient checkpointing is on), and device map (set only when quantization is active), then loads the model withAutoModelForCausalLM.from_pretrained.
Usage
Import these functions when setting up a model for SFT or GRPO training in the Open-R1 pipeline. They are the standard model and tokenizer constructors used across all training scripts in the project.
Code Reference
| Property | Value |
|---|---|
| Repository | open-r1 |
| File | src/open_r1/utils/model_utils.py
|
| Lines | L9-42 |
| Import | from open_r1.utils import get_model, get_tokenizer
|
Function Signatures
def get_tokenizer(model_args: ModelConfig, training_args: SFTConfig | GRPOConfig) -> PreTrainedTokenizer:
"""Get the tokenizer for the model."""
def get_model(model_args: ModelConfig, training_args: SFTConfig | GRPOConfig) -> AutoModelForCausalLM:
"""Get the model"""
I/O Contract
Inputs
| Parameter | Type | Required | Description |
|---|---|---|---|
model_args |
ModelConfig |
Yes | Model configuration containing:
|
training_args |
GRPOConfig | Yes | Training configuration containing:
|
Outputs
| Function | Return Type | Description |
|---|---|---|
get_tokenizer |
PreTrainedTokenizer |
Configured tokenizer with optional chat template override applied |
get_model |
AutoModelForCausalLM |
Configured causal language model with dtype, quantization, attention backend, and KV cache settings applied |
Usage Examples
Loading a Model with bfloat16 and Flash Attention
from open_r1.utils import get_model, get_tokenizer
from open_r1.configs import ModelConfig
from trl import SFTConfig
# Configure model arguments
model_args = ModelConfig(
model_name_or_path="Qwen/Qwen2.5-7B",
torch_dtype="bfloat16",
attn_implementation="flash_attention_2",
trust_remote_code=True,
)
# Configure training arguments
training_args = SFTConfig(
output_dir="./output",
gradient_checkpointing=True,
)
# Load tokenizer and model
tokenizer = get_tokenizer(model_args, training_args)
model = get_model(model_args, training_args)