Add SmolLM3 (#38755)
* init smollm3 * integration tests * config quirks * docs stub * rests round 2 * tests round 3 * tests round 4 * bring SWA back * config checker pls * final checkpoint * style and copies * Update src/transformers/models/smollm3/modular_smollm3.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * Update src/transformers/models/smollm3/modular_smollm3.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> --------- Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
This commit is contained in:
@@ -1053,6 +1053,8 @@
|
|||||||
title: SigLIP
|
title: SigLIP
|
||||||
- local: model_doc/siglip2
|
- local: model_doc/siglip2
|
||||||
title: SigLIP2
|
title: SigLIP2
|
||||||
|
- local: model_doc/smollm3
|
||||||
|
title: SmolLM3
|
||||||
- local: model_doc/smolvlm
|
- local: model_doc/smolvlm
|
||||||
title: SmolVLM
|
title: SmolVLM
|
||||||
- local: model_doc/speech-encoder-decoder
|
- local: model_doc/speech-encoder-decoder
|
||||||
|
|||||||
173
docs/source/en/model_doc/smollm3.md
Normal file
173
docs/source/en/model_doc/smollm3.md
Normal file
@@ -0,0 +1,173 @@
|
|||||||
|
<!--Copyright 2025 The HuggingFace Team. All rights reserved.
|
||||||
|
|
||||||
|
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
|
||||||
|
the License. You may obtain a copy of the License at
|
||||||
|
|
||||||
|
http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
|
||||||
|
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
|
||||||
|
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
|
||||||
|
rendered properly in your Markdown viewer.
|
||||||
|
|
||||||
|
-->
|
||||||
|
|
||||||
|
<div style="float: right;">
|
||||||
|
<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>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
# SmolLM3
|
||||||
|
|
||||||
|
SmolLM3 is a fully open, compact language model designed for efficient deployment while maintaining strong performance. It uses a Transformer decoder architecture with Grouped Query Attention (GQA) to reduce the kv cache, and no RoPE, enabling improved performance on long-context tasks. It is trained using a multi-stage training approach on high-quality public datasets across web, code, and math domains. The model is multilingual and supports very large context lengths. The instruct variant is optimized for reasoning and tool use.
|
||||||
|
|
||||||
|
> [!TIP]
|
||||||
|
> Click on the SmolLM3 models in the right sidebar for more examples of how to apply SmolLM3 to different language tasks.
|
||||||
|
|
||||||
|
The example below demonstrates how to generate text with [`Pipeline`], [`AutoModel`], and from the command line using the instruction-tuned models.
|
||||||
|
|
||||||
|
<hfoptions id="usage">
|
||||||
|
<hfoption id="Pipeline">
|
||||||
|
|
||||||
|
```python
|
||||||
|
import torch
|
||||||
|
from transformers import pipeline
|
||||||
|
|
||||||
|
pipe = pipeline(
|
||||||
|
task="text-generation",
|
||||||
|
model="HuggingFaceTB/SmolLM3-3B",
|
||||||
|
torch_dtype=torch.bfloat16,
|
||||||
|
device_map=0
|
||||||
|
)
|
||||||
|
|
||||||
|
messages = [
|
||||||
|
{"role": "system", "content": "You are a helpful assistant."},
|
||||||
|
{"role": "user", "content": "Tell me about yourself."},
|
||||||
|
]
|
||||||
|
outputs = pipe(messages, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
|
||||||
|
print(outputs[0]["generated_text"][-1]['content'])
|
||||||
|
```
|
||||||
|
|
||||||
|
</hfoption>
|
||||||
|
<hfoption id="AutoModel">
|
||||||
|
|
||||||
|
```python
|
||||||
|
import torch
|
||||||
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||||
|
|
||||||
|
model = AutoModelForCausalLM.from_pretrained(
|
||||||
|
"HuggingFaceTB/SmolLM3-3B",
|
||||||
|
torch_dtype=torch.bfloat16,
|
||||||
|
device_map="auto",
|
||||||
|
attn_implementation="sdpa"
|
||||||
|
)
|
||||||
|
tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM3-3B")
|
||||||
|
|
||||||
|
prompt = "Give me a short introduction to large language models."
|
||||||
|
messages = [
|
||||||
|
{"role": "system", "content": "You are a helpful assistant."},
|
||||||
|
{"role": "user", "content": prompt}
|
||||||
|
]
|
||||||
|
text = tokenizer.apply_chat_template(
|
||||||
|
messages,
|
||||||
|
tokenize=False,
|
||||||
|
add_generation_prompt=True
|
||||||
|
)
|
||||||
|
model_inputs = tokenizer([text], return_tensors="pt").to("cuda")
|
||||||
|
|
||||||
|
generated_ids = model.generate(
|
||||||
|
model_inputs.input_ids,
|
||||||
|
cache_implementation="static",
|
||||||
|
max_new_tokens=512,
|
||||||
|
do_sample=True,
|
||||||
|
temperature=0.7,
|
||||||
|
top_k=50,
|
||||||
|
top_p=0.95
|
||||||
|
)
|
||||||
|
generated_ids = [
|
||||||
|
output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
|
||||||
|
]
|
||||||
|
|
||||||
|
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
|
||||||
|
print(response)
|
||||||
|
```
|
||||||
|
|
||||||
|
</hfoption>
|
||||||
|
<hfoption id="transformers CLI">
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# pip install -U flash-attn --no-build-isolation
|
||||||
|
transformers chat HuggingFaceTB/SmolLM3-3B --torch_dtype auto --attn_implementation flash_attention_2 --device 0
|
||||||
|
```
|
||||||
|
|
||||||
|
</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 quantize the weights to 4-bits.
|
||||||
|
|
||||||
|
```python
|
||||||
|
# pip install -U flash-attn --no-build-isolation
|
||||||
|
import torch
|
||||||
|
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
|
||||||
|
|
||||||
|
quantization_config = BitsAndBytesConfig(
|
||||||
|
load_in_4bit=True,
|
||||||
|
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||||
|
bnb_4bit_quant_type="nf4",
|
||||||
|
bnb_4bit_use_double_quant=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM3-3B")
|
||||||
|
model = AutoModelForCausalLM.from_pretrained(
|
||||||
|
"HuggingFaceTB/SmolLM3-3B",
|
||||||
|
torch_dtype=torch.bfloat16,
|
||||||
|
device_map="auto",
|
||||||
|
quantization_config=quantization_config,
|
||||||
|
attn_implementation="flash_attention_2"
|
||||||
|
)
|
||||||
|
|
||||||
|
inputs = tokenizer("Gravity is the force", return_tensors="pt").to("cuda")
|
||||||
|
outputs = model.generate(**inputs, max_new_tokens=100)
|
||||||
|
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
|
||||||
|
```
|
||||||
|
|
||||||
|
|
||||||
|
## Notes
|
||||||
|
|
||||||
|
- Ensure your Transformers library version is up-to-date. SmolLM3 requires Transformers>=4.53.0 for full support.
|
||||||
|
|
||||||
|
## SmolLM3Config
|
||||||
|
|
||||||
|
[[autodoc]] SmolLM3Config
|
||||||
|
|
||||||
|
## SmolLM3Model
|
||||||
|
|
||||||
|
[[autodoc]] SmolLM3Model
|
||||||
|
- forward
|
||||||
|
|
||||||
|
## SmolLM3ForCausalLM
|
||||||
|
|
||||||
|
[[autodoc]] SmolLM3ForCausalLM
|
||||||
|
- forward
|
||||||
|
|
||||||
|
## SmolLM3ForSequenceClassification
|
||||||
|
|
||||||
|
[[autodoc]] SmolLM3ForSequenceClassification
|
||||||
|
- forward
|
||||||
|
|
||||||
|
## SmolLM3ForTokenClassification
|
||||||
|
|
||||||
|
[[autodoc]] SmolLM3ForTokenClassification
|
||||||
|
- forward
|
||||||
|
|
||||||
|
## SmolLM3ForQuestionAnswering
|
||||||
|
|
||||||
|
[[autodoc]] SmolLM3ForQuestionAnswering
|
||||||
|
- forward
|
||||||
@@ -315,6 +315,7 @@ CONFIG_MAPPING_NAMES = OrderedDict[str, str](
|
|||||||
("siglip", "SiglipConfig"),
|
("siglip", "SiglipConfig"),
|
||||||
("siglip2", "Siglip2Config"),
|
("siglip2", "Siglip2Config"),
|
||||||
("siglip_vision_model", "SiglipVisionConfig"),
|
("siglip_vision_model", "SiglipVisionConfig"),
|
||||||
|
("smollm3", "SmolLM3Config"),
|
||||||
("smolvlm", "SmolVLMConfig"),
|
("smolvlm", "SmolVLMConfig"),
|
||||||
("smolvlm_vision", "SmolVLMVisionConfig"),
|
("smolvlm_vision", "SmolVLMVisionConfig"),
|
||||||
("speech-encoder-decoder", "SpeechEncoderDecoderConfig"),
|
("speech-encoder-decoder", "SpeechEncoderDecoderConfig"),
|
||||||
@@ -705,6 +706,7 @@ MODEL_NAMES_MAPPING = OrderedDict[str, str](
|
|||||||
("siglip2", "SigLIP2"),
|
("siglip2", "SigLIP2"),
|
||||||
("siglip2_vision_model", "Siglip2VisionModel"),
|
("siglip2_vision_model", "Siglip2VisionModel"),
|
||||||
("siglip_vision_model", "SiglipVisionModel"),
|
("siglip_vision_model", "SiglipVisionModel"),
|
||||||
|
("smollm3", "SmolLM3"),
|
||||||
("smolvlm", "SmolVLM"),
|
("smolvlm", "SmolVLM"),
|
||||||
("smolvlm_vision", "SmolVLMVisionTransformer"),
|
("smolvlm_vision", "SmolVLMVisionTransformer"),
|
||||||
("speech-encoder-decoder", "Speech Encoder decoder"),
|
("speech-encoder-decoder", "Speech Encoder decoder"),
|
||||||
|
|||||||
@@ -295,6 +295,7 @@ MODEL_MAPPING_NAMES = OrderedDict(
|
|||||||
("siglip", "SiglipModel"),
|
("siglip", "SiglipModel"),
|
||||||
("siglip2", "Siglip2Model"),
|
("siglip2", "Siglip2Model"),
|
||||||
("siglip_vision_model", "SiglipVisionModel"),
|
("siglip_vision_model", "SiglipVisionModel"),
|
||||||
|
("smollm3", "SmolLM3Model"),
|
||||||
("smolvlm", "SmolVLMModel"),
|
("smolvlm", "SmolVLMModel"),
|
||||||
("smolvlm_vision", "SmolVLMVisionTransformer"),
|
("smolvlm_vision", "SmolVLMVisionTransformer"),
|
||||||
("speech_to_text", "Speech2TextModel"),
|
("speech_to_text", "Speech2TextModel"),
|
||||||
@@ -644,6 +645,7 @@ MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
|
|||||||
("roc_bert", "RoCBertForCausalLM"),
|
("roc_bert", "RoCBertForCausalLM"),
|
||||||
("roformer", "RoFormerForCausalLM"),
|
("roformer", "RoFormerForCausalLM"),
|
||||||
("rwkv", "RwkvForCausalLM"),
|
("rwkv", "RwkvForCausalLM"),
|
||||||
|
("smollm3", "SmolLM3ForCausalLM"),
|
||||||
("speech_to_text_2", "Speech2Text2ForCausalLM"),
|
("speech_to_text_2", "Speech2Text2ForCausalLM"),
|
||||||
("stablelm", "StableLmForCausalLM"),
|
("stablelm", "StableLmForCausalLM"),
|
||||||
("starcoder2", "Starcoder2ForCausalLM"),
|
("starcoder2", "Starcoder2ForCausalLM"),
|
||||||
@@ -1158,6 +1160,7 @@ MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
|
|||||||
("roberta-prelayernorm", "RobertaPreLayerNormForSequenceClassification"),
|
("roberta-prelayernorm", "RobertaPreLayerNormForSequenceClassification"),
|
||||||
("roc_bert", "RoCBertForSequenceClassification"),
|
("roc_bert", "RoCBertForSequenceClassification"),
|
||||||
("roformer", "RoFormerForSequenceClassification"),
|
("roformer", "RoFormerForSequenceClassification"),
|
||||||
|
("smollm3", "SmolLM3ForSequenceClassification"),
|
||||||
("squeezebert", "SqueezeBertForSequenceClassification"),
|
("squeezebert", "SqueezeBertForSequenceClassification"),
|
||||||
("stablelm", "StableLmForSequenceClassification"),
|
("stablelm", "StableLmForSequenceClassification"),
|
||||||
("starcoder2", "Starcoder2ForSequenceClassification"),
|
("starcoder2", "Starcoder2ForSequenceClassification"),
|
||||||
@@ -1244,6 +1247,7 @@ MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
|
|||||||
("roberta-prelayernorm", "RobertaPreLayerNormForQuestionAnswering"),
|
("roberta-prelayernorm", "RobertaPreLayerNormForQuestionAnswering"),
|
||||||
("roc_bert", "RoCBertForQuestionAnswering"),
|
("roc_bert", "RoCBertForQuestionAnswering"),
|
||||||
("roformer", "RoFormerForQuestionAnswering"),
|
("roformer", "RoFormerForQuestionAnswering"),
|
||||||
|
("smollm3", "SmolLM3ForQuestionAnswering"),
|
||||||
("splinter", "SplinterForQuestionAnswering"),
|
("splinter", "SplinterForQuestionAnswering"),
|
||||||
("squeezebert", "SqueezeBertForQuestionAnswering"),
|
("squeezebert", "SqueezeBertForQuestionAnswering"),
|
||||||
("t5", "T5ForQuestionAnswering"),
|
("t5", "T5ForQuestionAnswering"),
|
||||||
@@ -1352,6 +1356,7 @@ MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
|
|||||||
("roberta-prelayernorm", "RobertaPreLayerNormForTokenClassification"),
|
("roberta-prelayernorm", "RobertaPreLayerNormForTokenClassification"),
|
||||||
("roc_bert", "RoCBertForTokenClassification"),
|
("roc_bert", "RoCBertForTokenClassification"),
|
||||||
("roformer", "RoFormerForTokenClassification"),
|
("roformer", "RoFormerForTokenClassification"),
|
||||||
|
("smollm3", "SmolLM3ForTokenClassification"),
|
||||||
("squeezebert", "SqueezeBertForTokenClassification"),
|
("squeezebert", "SqueezeBertForTokenClassification"),
|
||||||
("stablelm", "StableLmForTokenClassification"),
|
("stablelm", "StableLmForTokenClassification"),
|
||||||
("starcoder2", "Starcoder2ForTokenClassification"),
|
("starcoder2", "Starcoder2ForTokenClassification"),
|
||||||
|
|||||||
@@ -172,6 +172,8 @@ class MimiEncoderOutput(ModelOutput):
|
|||||||
|
|
||||||
If `past_key_values` are used, the user can optionally input only the last `audio_values` or `audio_codes (those that don't
|
If `past_key_values` are used, the user can optionally input only the last `audio_values` or `audio_codes (those that don't
|
||||||
have their past key value states given to this model).
|
have their past key value states given to this model).
|
||||||
|
padding_cache (<fill_type>):
|
||||||
|
<fill_docstring>
|
||||||
"""
|
"""
|
||||||
|
|
||||||
audio_codes: Optional[torch.LongTensor] = None
|
audio_codes: Optional[torch.LongTensor] = None
|
||||||
|
|||||||
27
src/transformers/models/smollm3/__init__.py
Normal file
27
src/transformers/models/smollm3/__init__.py
Normal file
@@ -0,0 +1,27 @@
|
|||||||
|
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on 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.
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from ...utils import _LazyModule
|
||||||
|
from ...utils.import_utils import define_import_structure
|
||||||
|
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from .configuration_smollm3 import *
|
||||||
|
from .modeling_smollm3 import *
|
||||||
|
else:
|
||||||
|
import sys
|
||||||
|
|
||||||
|
_file = globals()["__file__"]
|
||||||
|
sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
|
||||||
245
src/transformers/models/smollm3/configuration_smollm3.py
Normal file
245
src/transformers/models/smollm3/configuration_smollm3.py
Normal file
@@ -0,0 +1,245 @@
|
|||||||
|
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
||||||
|
# This file was automatically generated from src/transformers/models/smollm3/modular_smollm3.py.
|
||||||
|
# Do NOT edit this file manually as any edits will be overwritten by the generation of
|
||||||
|
# the file from the modular. If any change should be done, please apply the change to the
|
||||||
|
# modular_smollm3.py file directly. One of our CI enforces this.
|
||||||
|
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
||||||
|
# coding=utf-8
|
||||||
|
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on 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.
|
||||||
|
|
||||||
|
from ...configuration_utils import PretrainedConfig, layer_type_validation
|
||||||
|
from ...modeling_rope_utils import rope_config_validation
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3Config(PretrainedConfig):
|
||||||
|
r"""
|
||||||
|
This is the configuration class to store the configuration of a [`SmolLM3Model`]. It is used to instantiate a
|
||||||
|
SmolLM3 model according to the specified arguments, defining the model architecture. Instantiating a configuration
|
||||||
|
with the defaults will yield a similar configuration to that of the SmolLM3 3B.
|
||||||
|
e.g. [HuggingFaceTB/SmolLM3-3B](https://huggingface.co/HuggingFaceTB/SmolLM3-3B)
|
||||||
|
|
||||||
|
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||||
|
documentation from [`PretrainedConfig`] for more information.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
vocab_size (`int`, *optional*, defaults to 128256):
|
||||||
|
Vocabulary size of the SmolLM3 model. Defines the number of different tokens that can be represented by the
|
||||||
|
`inputs_ids` passed when calling [`SmolLM3Model`]
|
||||||
|
hidden_size (`int`, *optional*, defaults to 2048):
|
||||||
|
Dimension of the hidden representations.
|
||||||
|
intermediate_size (`int`, *optional*, defaults to 11008):
|
||||||
|
Dimension of the MLP representations.
|
||||||
|
num_hidden_layers (`int`, *optional*, defaults to 36):
|
||||||
|
Number of hidden layers in the Transformer encoder.
|
||||||
|
num_attention_heads (`int`, *optional*, defaults to 16):
|
||||||
|
Number of attention heads for each attention layer in the Transformer encoder.
|
||||||
|
num_key_value_heads (`int`, *optional*, defaults to 4):
|
||||||
|
This is the number of key_value heads that should be used to implement Grouped Query Attention. If
|
||||||
|
`num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
|
||||||
|
`num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
|
||||||
|
converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
|
||||||
|
by meanpooling all the original heads within that group. For more details checkout [this
|
||||||
|
paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `16`.
|
||||||
|
hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
|
||||||
|
The non-linear activation function (function or string) in the decoder.
|
||||||
|
max_position_embeddings (`int`, *optional*, defaults to 32768):
|
||||||
|
The maximum sequence length that this model might ever be used with.
|
||||||
|
initializer_range (`float`, *optional*, defaults to 0.02):
|
||||||
|
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
||||||
|
rms_norm_eps (`float`, *optional*, defaults to 1e-06):
|
||||||
|
The epsilon used by the rms normalization layers.
|
||||||
|
use_cache (`bool`, *optional*, defaults to `True`):
|
||||||
|
Whether or not the model should return the last key/values attentions (not used by all models). Only
|
||||||
|
relevant if `config.is_decoder=True`.
|
||||||
|
pad_token_id (`int`, *optional*, defaults to 128004):
|
||||||
|
The id of the padding token.
|
||||||
|
bos_token_id (`int`, *optional*, defaults to 128000):
|
||||||
|
The id of the beginning of sentence token.
|
||||||
|
eos_token_id (`int`, *optional*, defaults to 128001):
|
||||||
|
The id of the end of sentence token.
|
||||||
|
rope_theta (`float`, *optional*, defaults to 2000000.0):
|
||||||
|
The base period of the RoPE embeddings.
|
||||||
|
rope_scaling (`Dict`, *optional*):
|
||||||
|
Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
|
||||||
|
and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
|
||||||
|
accordingly.
|
||||||
|
Expected contents:
|
||||||
|
`rope_type` (`str`):
|
||||||
|
The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
|
||||||
|
'llama3'], with 'default' being the original RoPE implementation.
|
||||||
|
`factor` (`float`, *optional*):
|
||||||
|
Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
|
||||||
|
most scaling types, a `factor` of x will enable the model to handle sequences of length x *
|
||||||
|
original maximum pre-trained length.
|
||||||
|
`original_max_position_embeddings` (`int`, *optional*):
|
||||||
|
Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
|
||||||
|
pretraining.
|
||||||
|
`attention_factor` (`float`, *optional*):
|
||||||
|
Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
|
||||||
|
computation. If unspecified, it defaults to value recommended by the implementation, using the
|
||||||
|
`factor` field to infer the suggested value.
|
||||||
|
`beta_fast` (`float`, *optional*):
|
||||||
|
Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
|
||||||
|
ramp function. If unspecified, it defaults to 32.
|
||||||
|
`beta_slow` (`float`, *optional*):
|
||||||
|
Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
|
||||||
|
ramp function. If unspecified, it defaults to 1.
|
||||||
|
`short_factor` (`List[float]`, *optional*):
|
||||||
|
Only used with 'longrope'. The scaling factor to be applied to short contexts (<
|
||||||
|
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
||||||
|
size divided by the number of attention heads divided by 2
|
||||||
|
`long_factor` (`List[float]`, *optional*):
|
||||||
|
Only used with 'longrope'. The scaling factor to be applied to long contexts (<
|
||||||
|
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
||||||
|
size divided by the number of attention heads divided by 2
|
||||||
|
`low_freq_factor` (`float`, *optional*):
|
||||||
|
Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
|
||||||
|
`high_freq_factor` (`float`, *optional*):
|
||||||
|
Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
|
||||||
|
use_sliding_window (`bool`, *optional*, defaults to `False`):
|
||||||
|
Whether to use sliding window attention.
|
||||||
|
sliding_window (`int`, *optional*):
|
||||||
|
Sliding window attention (SWA) window size. If not specified, will default to `None`.
|
||||||
|
no_rope_layers (`List[int]`, *optional*):
|
||||||
|
List with at least the same length as the number of layers in the model.
|
||||||
|
A `1` at an index position indicates that the corresponding layer will use RoPE,
|
||||||
|
while a `0` indicates that it's a NoPE layer.
|
||||||
|
no_rope_layer_interval (`int`, *optional*, defaults to 4):
|
||||||
|
If `no_rope_layers` is `None`, it will be created using a NoPE layer every
|
||||||
|
`no_rope_layer_interval` layers.
|
||||||
|
layer_types (`list`, *optional*):
|
||||||
|
Attention pattern for each layer. Automatically computed based on sliding window and NoPE settings.
|
||||||
|
attention_bias (`bool`, *optional*, defaults to `False`):
|
||||||
|
Whether to use a bias in the query, key, value and output projection layers during self-attention.
|
||||||
|
attention_dropout (`float`, *optional*, defaults to 0.0):
|
||||||
|
The dropout ratio for the attention probabilities.
|
||||||
|
|
||||||
|
```python
|
||||||
|
>>> from transformers import SmolLM3Model, SmolLM3Config
|
||||||
|
|
||||||
|
>>> # Initializing a SmolLM3 style configuration
|
||||||
|
>>> configuration = SmolLM3Config()
|
||||||
|
|
||||||
|
>>> # Initializing a model from the SmolLM3 style configuration
|
||||||
|
>>> model = SmolLM3Model(configuration)
|
||||||
|
|
||||||
|
>>> # Accessing the model configuration
|
||||||
|
>>> configuration = model.config
|
||||||
|
```"""
|
||||||
|
|
||||||
|
model_type = "smollm3"
|
||||||
|
keys_to_ignore_at_inference = ["past_key_values"]
|
||||||
|
|
||||||
|
base_model_tp_plan = {
|
||||||
|
"layers.*.self_attn.q_proj": "colwise",
|
||||||
|
"layers.*.self_attn.k_proj": "colwise",
|
||||||
|
"layers.*.self_attn.v_proj": "colwise",
|
||||||
|
"layers.*.self_attn.o_proj": "rowwise",
|
||||||
|
"layers.*.mlp.gate_proj": "colwise",
|
||||||
|
"layers.*.mlp.up_proj": "colwise",
|
||||||
|
"layers.*.mlp.down_proj": "rowwise",
|
||||||
|
}
|
||||||
|
base_model_pp_plan = {
|
||||||
|
"embed_tokens": (["input_ids"], ["inputs_embeds"]),
|
||||||
|
"layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
|
||||||
|
"norm": (["hidden_states"], ["hidden_states"]),
|
||||||
|
}
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
vocab_size=128256,
|
||||||
|
hidden_size=2048,
|
||||||
|
intermediate_size=11008,
|
||||||
|
num_hidden_layers=36,
|
||||||
|
num_attention_heads=16,
|
||||||
|
num_key_value_heads=4,
|
||||||
|
hidden_act="silu",
|
||||||
|
max_position_embeddings=32768,
|
||||||
|
initializer_range=0.02,
|
||||||
|
rms_norm_eps=1e-6,
|
||||||
|
use_cache=True,
|
||||||
|
pad_token_id=128004,
|
||||||
|
bos_token_id=128000,
|
||||||
|
eos_token_id=128001,
|
||||||
|
rope_theta=2000000.0,
|
||||||
|
rope_scaling=None,
|
||||||
|
use_sliding_window=False,
|
||||||
|
sliding_window=None,
|
||||||
|
no_rope_layers=None,
|
||||||
|
no_rope_layer_interval=4,
|
||||||
|
layer_types=None,
|
||||||
|
attention_bias=False,
|
||||||
|
attention_dropout=0.0,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
super().__init__(
|
||||||
|
pad_token_id=pad_token_id,
|
||||||
|
bos_token_id=bos_token_id,
|
||||||
|
eos_token_id=eos_token_id,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
self.vocab_size = vocab_size
|
||||||
|
self.max_position_embeddings = max_position_embeddings
|
||||||
|
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.use_sliding_window = use_sliding_window
|
||||||
|
self.sliding_window = sliding_window
|
||||||
|
|
||||||
|
# for backward compatibility
|
||||||
|
if num_key_value_heads is None:
|
||||||
|
num_key_value_heads = num_attention_heads
|
||||||
|
|
||||||
|
self.num_key_value_heads = num_key_value_heads
|
||||||
|
self.hidden_act = hidden_act
|
||||||
|
self.initializer_range = initializer_range
|
||||||
|
self.rms_norm_eps = rms_norm_eps
|
||||||
|
self.use_cache = use_cache
|
||||||
|
self.rope_theta = rope_theta
|
||||||
|
self.rope_scaling = rope_scaling
|
||||||
|
self.attention_bias = attention_bias
|
||||||
|
self.attention_dropout = attention_dropout
|
||||||
|
|
||||||
|
if no_rope_layers is None:
|
||||||
|
self.no_rope_layers = [
|
||||||
|
int((layer_idx + 1) % no_rope_layer_interval != 0) for layer_idx in range(num_hidden_layers)
|
||||||
|
]
|
||||||
|
else:
|
||||||
|
self.no_rope_layers = no_rope_layers
|
||||||
|
|
||||||
|
self.no_rope_layer_interval = no_rope_layer_interval
|
||||||
|
|
||||||
|
# Update layer_types based on sliding window and NoPE pattern
|
||||||
|
if layer_types is None:
|
||||||
|
layer_types = []
|
||||||
|
for layer_idx in range(num_hidden_layers):
|
||||||
|
has_rope = self.no_rope_layers[layer_idx]
|
||||||
|
if use_sliding_window and sliding_window is not None and not has_rope:
|
||||||
|
layer_types.append("sliding_attention")
|
||||||
|
else:
|
||||||
|
layer_types.append("full_attention")
|
||||||
|
|
||||||
|
self.layer_types = layer_types
|
||||||
|
layer_type_validation(self.layer_types)
|
||||||
|
|
||||||
|
# Validate the correctness of rotary position embeddings parameters
|
||||||
|
# BC: if there is a 'type' field, move it to 'rope_type'.
|
||||||
|
if self.rope_scaling is not None and "type" in self.rope_scaling:
|
||||||
|
self.rope_scaling["rope_type"] = self.rope_scaling["type"]
|
||||||
|
rope_config_validation(self)
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ["SmolLM3Config"]
|
||||||
845
src/transformers/models/smollm3/modeling_smollm3.py
Normal file
845
src/transformers/models/smollm3/modeling_smollm3.py
Normal file
@@ -0,0 +1,845 @@
|
|||||||
|
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
||||||
|
# This file was automatically generated from src/transformers/models/smollm3/modular_smollm3.py.
|
||||||
|
# Do NOT edit this file manually as any edits will be overwritten by the generation of
|
||||||
|
# the file from the modular. If any change should be done, please apply the change to the
|
||||||
|
# modular_smollm3.py file directly. One of our CI enforces this.
|
||||||
|
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
|
||||||
|
# coding=utf-8
|
||||||
|
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on 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.
|
||||||
|
|
||||||
|
from typing import Callable, Optional, Union
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from torch import nn
|
||||||
|
|
||||||
|
from ...activations import ACT2FN
|
||||||
|
from ...cache_utils import Cache, DynamicCache
|
||||||
|
from ...generation import GenerationMixin
|
||||||
|
from ...integrations import use_kernel_forward_from_hub
|
||||||
|
from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
|
||||||
|
from ...modeling_flash_attention_utils import FlashAttentionKwargs
|
||||||
|
from ...modeling_layers import GradientCheckpointingLayer
|
||||||
|
from ...modeling_outputs import (
|
||||||
|
BaseModelOutputWithPast,
|
||||||
|
CausalLMOutputWithPast,
|
||||||
|
QuestionAnsweringModelOutput,
|
||||||
|
SequenceClassifierOutputWithPast,
|
||||||
|
TokenClassifierOutput,
|
||||||
|
)
|
||||||
|
from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
|
||||||
|
from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
|
||||||
|
from ...processing_utils import Unpack
|
||||||
|
from ...utils import LossKwargs, auto_docstring, can_return_tuple, logging
|
||||||
|
from .configuration_smollm3 import SmolLM3Config
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def rotate_half(x):
|
||||||
|
"""Rotates half the hidden dims of the input."""
|
||||||
|
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, position_ids=None, unsqueeze_dim=1):
|
||||||
|
"""Applies Rotary Position Embedding to the query and key tensors.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
q (`torch.Tensor`): The query tensor.
|
||||||
|
k (`torch.Tensor`): The key tensor.
|
||||||
|
cos (`torch.Tensor`): The cosine part of the rotary embedding.
|
||||||
|
sin (`torch.Tensor`): The sine part of the rotary embedding.
|
||||||
|
position_ids (`torch.Tensor`, *optional*):
|
||||||
|
Deprecated and unused.
|
||||||
|
unsqueeze_dim (`int`, *optional*, defaults to 1):
|
||||||
|
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
|
||||||
|
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
|
||||||
|
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
|
||||||
|
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
|
||||||
|
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
|
||||||
|
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
|
||||||
|
Returns:
|
||||||
|
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
|
||||||
|
"""
|
||||||
|
cos = cos.unsqueeze(unsqueeze_dim)
|
||||||
|
sin = sin.unsqueeze(unsqueeze_dim)
|
||||||
|
q_embed = (q * cos) + (rotate_half(q) * sin)
|
||||||
|
k_embed = (k * cos) + (rotate_half(k) * sin)
|
||||||
|
return q_embed, k_embed
|
||||||
|
|
||||||
|
|
||||||
|
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
||||||
|
"""
|
||||||
|
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
|
||||||
|
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
|
||||||
|
"""
|
||||||
|
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
||||||
|
if n_rep == 1:
|
||||||
|
return hidden_states
|
||||||
|
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
||||||
|
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
||||||
|
|
||||||
|
|
||||||
|
def eager_attention_forward(
|
||||||
|
module: nn.Module,
|
||||||
|
query: torch.Tensor,
|
||||||
|
key: torch.Tensor,
|
||||||
|
value: torch.Tensor,
|
||||||
|
attention_mask: Optional[torch.Tensor],
|
||||||
|
scaling: float,
|
||||||
|
dropout: float = 0.0,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
key_states = repeat_kv(key, module.num_key_value_groups)
|
||||||
|
value_states = repeat_kv(value, module.num_key_value_groups)
|
||||||
|
|
||||||
|
attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
|
||||||
|
if attention_mask is not None:
|
||||||
|
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
|
||||||
|
attn_weights = attn_weights + causal_mask
|
||||||
|
|
||||||
|
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
|
||||||
|
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
|
||||||
|
attn_output = torch.matmul(attn_weights, value_states)
|
||||||
|
attn_output = attn_output.transpose(1, 2).contiguous()
|
||||||
|
|
||||||
|
return attn_output, attn_weights
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3Attention(nn.Module):
|
||||||
|
"""Multi-headed attention from 'Attention Is All You Need' paper"""
|
||||||
|
|
||||||
|
def __init__(self, config: SmolLM3Config, layer_idx: int):
|
||||||
|
super().__init__()
|
||||||
|
self.config = config
|
||||||
|
self.layer_idx = layer_idx
|
||||||
|
self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
|
||||||
|
self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
|
||||||
|
self.scaling = self.head_dim**-0.5
|
||||||
|
self.attention_dropout = config.attention_dropout
|
||||||
|
self.is_causal = True
|
||||||
|
|
||||||
|
self.q_proj = nn.Linear(
|
||||||
|
config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
|
||||||
|
)
|
||||||
|
self.k_proj = nn.Linear(
|
||||||
|
config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
|
||||||
|
)
|
||||||
|
self.v_proj = nn.Linear(
|
||||||
|
config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
|
||||||
|
)
|
||||||
|
self.o_proj = nn.Linear(
|
||||||
|
config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
|
||||||
|
)
|
||||||
|
|
||||||
|
self.use_rope = config.no_rope_layers[layer_idx]
|
||||||
|
self.sliding_window = (
|
||||||
|
config.sliding_window
|
||||||
|
if config.use_sliding_window and config.layer_types[layer_idx] == "sliding_attention"
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
||||||
|
attention_mask: Optional[torch.Tensor],
|
||||||
|
past_key_value: Optional[Cache] = None,
|
||||||
|
cache_position: Optional[torch.LongTensor] = None,
|
||||||
|
**kwargs: Unpack[FlashAttentionKwargs],
|
||||||
|
) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
|
||||||
|
input_shape = hidden_states.shape[:-1]
|
||||||
|
hidden_shape = (*input_shape, -1, self.head_dim)
|
||||||
|
|
||||||
|
query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||||
|
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||||
|
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||||
|
|
||||||
|
if self.use_rope:
|
||||||
|
cos, sin = position_embeddings
|
||||||
|
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
|
||||||
|
|
||||||
|
if past_key_value is not None:
|
||||||
|
cache_kwargs = {"cache_position": cache_position}
|
||||||
|
key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
|
||||||
|
|
||||||
|
attention_interface: Callable = eager_attention_forward
|
||||||
|
if self.config._attn_implementation != "eager":
|
||||||
|
attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
|
||||||
|
|
||||||
|
attn_output, attn_weights = attention_interface(
|
||||||
|
self,
|
||||||
|
query_states,
|
||||||
|
key_states,
|
||||||
|
value_states,
|
||||||
|
attention_mask,
|
||||||
|
dropout=0.0 if not self.training else self.attention_dropout,
|
||||||
|
scaling=self.scaling,
|
||||||
|
sliding_window=self.sliding_window,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
attn_output = attn_output.reshape(*input_shape, -1).contiguous()
|
||||||
|
attn_output = self.o_proj(attn_output)
|
||||||
|
return attn_output, attn_weights
|
||||||
|
|
||||||
|
|
||||||
|
@use_kernel_forward_from_hub("RMSNorm")
|
||||||
|
class SmolLM3RMSNorm(nn.Module):
|
||||||
|
def __init__(self, hidden_size, eps=1e-6):
|
||||||
|
"""
|
||||||
|
SmolLM3RMSNorm is equivalent to T5LayerNorm
|
||||||
|
"""
|
||||||
|
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)
|
||||||
|
|
||||||
|
def extra_repr(self):
|
||||||
|
return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
|
||||||
|
|
||||||
|
|
||||||
|
@auto_docstring
|
||||||
|
class SmolLM3PreTrainedModel(PreTrainedModel):
|
||||||
|
config_class = SmolLM3Config
|
||||||
|
base_model_prefix = "model"
|
||||||
|
supports_gradient_checkpointing = True
|
||||||
|
_no_split_modules = ["SmolLM3DecoderLayer"]
|
||||||
|
_skip_keys_device_placement = ["past_key_values"]
|
||||||
|
_supports_flash_attn_3 = True
|
||||||
|
_supports_flash_attn_2 = True
|
||||||
|
_supports_sdpa = True
|
||||||
|
_supports_flex_attn = True
|
||||||
|
_supports_cache_class = True
|
||||||
|
_supports_quantized_cache = True
|
||||||
|
_supports_static_cache = True
|
||||||
|
_supports_attention_backend = True
|
||||||
|
|
||||||
|
def _init_weights(self, module):
|
||||||
|
std = self.config.initializer_range
|
||||||
|
if isinstance(module, nn.Linear):
|
||||||
|
module.weight.data.normal_(mean=0.0, std=std)
|
||||||
|
if module.bias is not None:
|
||||||
|
module.bias.data.zero_()
|
||||||
|
elif isinstance(module, nn.Embedding):
|
||||||
|
module.weight.data.normal_(mean=0.0, std=std)
|
||||||
|
if module.padding_idx is not None:
|
||||||
|
module.weight.data[module.padding_idx].zero_()
|
||||||
|
elif isinstance(module, SmolLM3RMSNorm):
|
||||||
|
module.weight.data.fill_(1.0)
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3MLP(nn.Module):
|
||||||
|
def __init__(self, config):
|
||||||
|
super().__init__()
|
||||||
|
self.config = config
|
||||||
|
self.hidden_size = config.hidden_size
|
||||||
|
self.intermediate_size = config.intermediate_size
|
||||||
|
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||||
|
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||||
|
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
||||||
|
self.act_fn = ACT2FN[config.hidden_act]
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
|
||||||
|
return down_proj
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3DecoderLayer(GradientCheckpointingLayer):
|
||||||
|
def __init__(self, config: SmolLM3Config, layer_idx: int):
|
||||||
|
super().__init__()
|
||||||
|
self.hidden_size = config.hidden_size
|
||||||
|
|
||||||
|
self.self_attn = SmolLM3Attention(config=config, layer_idx=layer_idx)
|
||||||
|
|
||||||
|
self.mlp = SmolLM3MLP(config)
|
||||||
|
self.input_layernorm = SmolLM3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||||
|
self.post_attention_layernorm = SmolLM3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||||
|
self.attention_type = config.layer_types[layer_idx]
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
attention_mask: Optional[torch.Tensor] = None,
|
||||||
|
position_ids: Optional[torch.LongTensor] = None,
|
||||||
|
past_key_value: Optional[Cache] = None,
|
||||||
|
output_attentions: Optional[bool] = False,
|
||||||
|
use_cache: Optional[bool] = False,
|
||||||
|
cache_position: Optional[torch.LongTensor] = None,
|
||||||
|
position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
|
||||||
|
**kwargs: Unpack[FlashAttentionKwargs],
|
||||||
|
) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
|
||||||
|
residual = hidden_states
|
||||||
|
hidden_states = self.input_layernorm(hidden_states)
|
||||||
|
|
||||||
|
# Self Attention
|
||||||
|
hidden_states, self_attn_weights = self.self_attn(
|
||||||
|
hidden_states=hidden_states,
|
||||||
|
attention_mask=attention_mask,
|
||||||
|
position_ids=position_ids,
|
||||||
|
past_key_value=past_key_value,
|
||||||
|
output_attentions=output_attentions,
|
||||||
|
use_cache=use_cache,
|
||||||
|
cache_position=cache_position,
|
||||||
|
position_embeddings=position_embeddings,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
hidden_states = residual + hidden_states
|
||||||
|
|
||||||
|
# Fully Connected
|
||||||
|
residual = hidden_states
|
||||||
|
hidden_states = self.post_attention_layernorm(hidden_states)
|
||||||
|
hidden_states = self.mlp(hidden_states)
|
||||||
|
hidden_states = residual + hidden_states
|
||||||
|
|
||||||
|
outputs = (hidden_states,)
|
||||||
|
if output_attentions:
|
||||||
|
outputs += (self_attn_weights,)
|
||||||
|
|
||||||
|
return outputs
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3RotaryEmbedding(nn.Module):
|
||||||
|
def __init__(self, config: SmolLM3Config, device=None):
|
||||||
|
super().__init__()
|
||||||
|
# BC: "rope_type" was originally "type"
|
||||||
|
if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
|
||||||
|
self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
|
||||||
|
else:
|
||||||
|
self.rope_type = "default"
|
||||||
|
self.max_seq_len_cached = config.max_position_embeddings
|
||||||
|
self.original_max_seq_len = config.max_position_embeddings
|
||||||
|
|
||||||
|
self.config = config
|
||||||
|
self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
|
||||||
|
|
||||||
|
inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
|
||||||
|
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||||
|
self.original_inv_freq = self.inv_freq
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
@dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
|
||||||
|
def forward(self, x, position_ids):
|
||||||
|
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
|
||||||
|
position_ids_expanded = position_ids[:, None, :].float()
|
||||||
|
|
||||||
|
device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
|
||||||
|
with torch.autocast(device_type=device_type, enabled=False): # Force float32
|
||||||
|
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
|
||||||
|
emb = torch.cat((freqs, freqs), dim=-1)
|
||||||
|
cos = emb.cos() * self.attention_scaling
|
||||||
|
sin = emb.sin() * self.attention_scaling
|
||||||
|
|
||||||
|
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
||||||
|
|
||||||
|
|
||||||
|
@auto_docstring
|
||||||
|
class SmolLM3Model(SmolLM3PreTrainedModel):
|
||||||
|
def __init__(self, config: SmolLM3Config):
|
||||||
|
super().__init__(config)
|
||||||
|
self.padding_idx = config.pad_token_id
|
||||||
|
self.vocab_size = config.vocab_size
|
||||||
|
|
||||||
|
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
||||||
|
self.layers = nn.ModuleList(
|
||||||
|
[SmolLM3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
|
||||||
|
)
|
||||||
|
self.norm = SmolLM3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||||
|
self.rotary_emb = SmolLM3RotaryEmbedding(config=config)
|
||||||
|
self.gradient_checkpointing = False
|
||||||
|
self.has_sliding_layers = "sliding_attention" in self.config.layer_types
|
||||||
|
|
||||||
|
# Initialize weights and apply final processing
|
||||||
|
self.post_init()
|
||||||
|
|
||||||
|
def get_input_embeddings(self):
|
||||||
|
return self.embed_tokens
|
||||||
|
|
||||||
|
def set_input_embeddings(self, value):
|
||||||
|
self.embed_tokens = value
|
||||||
|
|
||||||
|
@can_return_tuple
|
||||||
|
@auto_docstring
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
input_ids: Optional[torch.LongTensor] = None,
|
||||||
|
attention_mask: Optional[torch.Tensor] = None,
|
||||||
|
position_ids: Optional[torch.LongTensor] = None,
|
||||||
|
past_key_values: Optional[Cache] = None,
|
||||||
|
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||||
|
use_cache: Optional[bool] = None,
|
||||||
|
output_attentions: Optional[bool] = None,
|
||||||
|
output_hidden_states: Optional[bool] = None,
|
||||||
|
cache_position: Optional[torch.LongTensor] = None,
|
||||||
|
**flash_attn_kwargs: Unpack[FlashAttentionKwargs],
|
||||||
|
) -> BaseModelOutputWithPast:
|
||||||
|
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||||
|
output_hidden_states = (
|
||||||
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||||
|
)
|
||||||
|
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
||||||
|
|
||||||
|
if (input_ids is None) ^ (inputs_embeds is not None):
|
||||||
|
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
|
||||||
|
|
||||||
|
if self.gradient_checkpointing and self.training and use_cache:
|
||||||
|
logger.warning_once(
|
||||||
|
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
|
||||||
|
)
|
||||||
|
use_cache = False
|
||||||
|
|
||||||
|
# TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache
|
||||||
|
if not isinstance(past_key_values, (type(None), Cache)):
|
||||||
|
raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")
|
||||||
|
|
||||||
|
if inputs_embeds is None:
|
||||||
|
inputs_embeds = self.embed_tokens(input_ids)
|
||||||
|
|
||||||
|
if use_cache and past_key_values is None:
|
||||||
|
past_key_values = DynamicCache()
|
||||||
|
|
||||||
|
if cache_position is None:
|
||||||
|
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
|
||||||
|
cache_position = torch.arange(
|
||||||
|
past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
|
||||||
|
)
|
||||||
|
|
||||||
|
if position_ids is None:
|
||||||
|
position_ids = cache_position.unsqueeze(0)
|
||||||
|
|
||||||
|
# It may already have been prepared by e.g. `generate`
|
||||||
|
if not isinstance(causal_mask_mapping := attention_mask, dict):
|
||||||
|
# Prepare mask arguments
|
||||||
|
mask_kwargs = {
|
||||||
|
"config": self.config,
|
||||||
|
"input_embeds": inputs_embeds,
|
||||||
|
"attention_mask": attention_mask,
|
||||||
|
"cache_position": cache_position,
|
||||||
|
"past_key_values": past_key_values,
|
||||||
|
}
|
||||||
|
# Create the masks
|
||||||
|
causal_mask_mapping = {
|
||||||
|
"full_attention": create_causal_mask(**mask_kwargs),
|
||||||
|
}
|
||||||
|
# The sliding window alternating layers are not always activated depending on the config
|
||||||
|
if self.has_sliding_layers:
|
||||||
|
causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
|
||||||
|
|
||||||
|
hidden_states = inputs_embeds
|
||||||
|
|
||||||
|
# create position embeddings to be shared across the decoder layers
|
||||||
|
position_embeddings = self.rotary_emb(hidden_states, position_ids)
|
||||||
|
|
||||||
|
# decoder layers
|
||||||
|
all_hidden_states = () if output_hidden_states else None
|
||||||
|
all_self_attns = () if output_attentions else None
|
||||||
|
|
||||||
|
for decoder_layer in self.layers[: self.config.num_hidden_layers]:
|
||||||
|
if output_hidden_states:
|
||||||
|
all_hidden_states += (hidden_states,)
|
||||||
|
|
||||||
|
layer_outputs = decoder_layer(
|
||||||
|
hidden_states,
|
||||||
|
attention_mask=causal_mask_mapping[decoder_layer.attention_type],
|
||||||
|
position_ids=position_ids,
|
||||||
|
past_key_value=past_key_values,
|
||||||
|
output_attentions=output_attentions,
|
||||||
|
use_cache=use_cache,
|
||||||
|
cache_position=cache_position,
|
||||||
|
position_embeddings=position_embeddings,
|
||||||
|
**flash_attn_kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
hidden_states = layer_outputs[0]
|
||||||
|
|
||||||
|
if output_attentions:
|
||||||
|
all_self_attns += (layer_outputs[1],)
|
||||||
|
|
||||||
|
hidden_states = self.norm(hidden_states)
|
||||||
|
|
||||||
|
# add hidden states from the last decoder layer
|
||||||
|
if output_hidden_states:
|
||||||
|
all_hidden_states += (hidden_states,)
|
||||||
|
|
||||||
|
return BaseModelOutputWithPast(
|
||||||
|
last_hidden_state=hidden_states,
|
||||||
|
past_key_values=past_key_values if use_cache else None,
|
||||||
|
hidden_states=all_hidden_states,
|
||||||
|
attentions=all_self_attns,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
|
||||||
|
|
||||||
|
|
||||||
|
@auto_docstring
|
||||||
|
class SmolLM3ForCausalLM(SmolLM3PreTrainedModel, GenerationMixin):
|
||||||
|
_tied_weights_keys = ["lm_head.weight"]
|
||||||
|
_tp_plan = {"lm_head": "colwise_rep"}
|
||||||
|
_pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
|
||||||
|
|
||||||
|
def __init__(self, config):
|
||||||
|
super().__init__(config)
|
||||||
|
self.model = SmolLM3Model(config)
|
||||||
|
self.vocab_size = config.vocab_size
|
||||||
|
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||||
|
|
||||||
|
# Initialize weights and apply final processing
|
||||||
|
self.post_init()
|
||||||
|
|
||||||
|
def get_input_embeddings(self):
|
||||||
|
return self.model.embed_tokens
|
||||||
|
|
||||||
|
def set_input_embeddings(self, value):
|
||||||
|
self.model.embed_tokens = value
|
||||||
|
|
||||||
|
def get_output_embeddings(self):
|
||||||
|
return self.lm_head
|
||||||
|
|
||||||
|
def set_output_embeddings(self, new_embeddings):
|
||||||
|
self.lm_head = new_embeddings
|
||||||
|
|
||||||
|
def set_decoder(self, decoder):
|
||||||
|
self.model = decoder
|
||||||
|
|
||||||
|
def get_decoder(self):
|
||||||
|
return self.model
|
||||||
|
|
||||||
|
@can_return_tuple
|
||||||
|
@auto_docstring
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
input_ids: Optional[torch.LongTensor] = None,
|
||||||
|
attention_mask: Optional[torch.Tensor] = None,
|
||||||
|
position_ids: Optional[torch.LongTensor] = None,
|
||||||
|
past_key_values: Optional[Cache] = None,
|
||||||
|
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||||
|
labels: Optional[torch.LongTensor] = None,
|
||||||
|
use_cache: Optional[bool] = None,
|
||||||
|
output_attentions: Optional[bool] = None,
|
||||||
|
output_hidden_states: Optional[bool] = None,
|
||||||
|
cache_position: Optional[torch.LongTensor] = None,
|
||||||
|
logits_to_keep: Union[int, torch.Tensor] = 0,
|
||||||
|
**kwargs: Unpack[KwargsForCausalLM],
|
||||||
|
) -> CausalLMOutputWithPast:
|
||||||
|
r"""
|
||||||
|
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
|
||||||
|
Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
|
||||||
|
config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
|
||||||
|
(masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
|
||||||
|
|
||||||
|
Example:
|
||||||
|
|
||||||
|
```python
|
||||||
|
>>> from transformers import AutoTokenizer, SmolLM3ForCausalLM
|
||||||
|
|
||||||
|
>>> model = SmolLM3ForCausalLM.from_pretrained("meta-smollm3/SmolLM3-2-7b-hf")
|
||||||
|
>>> tokenizer = AutoTokenizer.from_pretrained("meta-smollm3/SmolLM3-2-7b-hf")
|
||||||
|
|
||||||
|
>>> prompt = "Hey, are you conscious? Can you talk to me?"
|
||||||
|
>>> inputs = tokenizer(prompt, return_tensors="pt")
|
||||||
|
|
||||||
|
>>> # Generate
|
||||||
|
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
|
||||||
|
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
|
||||||
|
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
|
||||||
|
```"""
|
||||||
|
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||||
|
output_hidden_states = (
|
||||||
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||||
|
)
|
||||||
|
|
||||||
|
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||||
|
outputs: BaseModelOutputWithPast = self.model(
|
||||||
|
input_ids=input_ids,
|
||||||
|
attention_mask=attention_mask,
|
||||||
|
position_ids=position_ids,
|
||||||
|
past_key_values=past_key_values,
|
||||||
|
inputs_embeds=inputs_embeds,
|
||||||
|
use_cache=use_cache,
|
||||||
|
output_attentions=output_attentions,
|
||||||
|
output_hidden_states=output_hidden_states,
|
||||||
|
cache_position=cache_position,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
hidden_states = outputs.last_hidden_state
|
||||||
|
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
|
||||||
|
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
||||||
|
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
||||||
|
|
||||||
|
loss = None
|
||||||
|
if labels is not None:
|
||||||
|
loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
|
||||||
|
|
||||||
|
return CausalLMOutputWithPast(
|
||||||
|
loss=loss,
|
||||||
|
logits=logits,
|
||||||
|
past_key_values=outputs.past_key_values,
|
||||||
|
hidden_states=outputs.hidden_states,
|
||||||
|
attentions=outputs.attentions,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@auto_docstring(
|
||||||
|
custom_intro="""
|
||||||
|
The SmolLM3 Model transformer with a sequence classification head on top (linear layer).
|
||||||
|
|
||||||
|
[`SmolLM3ForSequenceClassification`] uses the last token in order to do the classification, as other causal models
|
||||||
|
(e.g. GPT-2) do.
|
||||||
|
|
||||||
|
Since it does classification on the last token, it requires to know the position of the last token. If a
|
||||||
|
`pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
|
||||||
|
no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
|
||||||
|
padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
|
||||||
|
each row of the batch).
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
class SmolLM3ForSequenceClassification(SmolLM3PreTrainedModel):
|
||||||
|
def __init__(self, config):
|
||||||
|
super().__init__(config)
|
||||||
|
self.num_labels = config.num_labels
|
||||||
|
self.model = SmolLM3Model(config)
|
||||||
|
self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
|
||||||
|
|
||||||
|
# Initialize weights and apply final processing
|
||||||
|
self.post_init()
|
||||||
|
|
||||||
|
def get_input_embeddings(self):
|
||||||
|
return self.model.embed_tokens
|
||||||
|
|
||||||
|
def set_input_embeddings(self, value):
|
||||||
|
self.model.embed_tokens = value
|
||||||
|
|
||||||
|
@can_return_tuple
|
||||||
|
@auto_docstring
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
input_ids: Optional[torch.LongTensor] = None,
|
||||||
|
attention_mask: Optional[torch.Tensor] = None,
|
||||||
|
position_ids: Optional[torch.LongTensor] = None,
|
||||||
|
past_key_values: Optional[Cache] = None,
|
||||||
|
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||||
|
labels: Optional[torch.LongTensor] = None,
|
||||||
|
use_cache: Optional[bool] = None,
|
||||||
|
output_attentions: Optional[bool] = None,
|
||||||
|
output_hidden_states: Optional[bool] = None,
|
||||||
|
) -> SequenceClassifierOutputWithPast:
|
||||||
|
r"""
|
||||||
|
labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
|
||||||
|
Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
|
||||||
|
config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
|
||||||
|
`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
|
||||||
|
"""
|
||||||
|
|
||||||
|
transformer_outputs: BaseModelOutputWithPast = self.model(
|
||||||
|
input_ids,
|
||||||
|
attention_mask=attention_mask,
|
||||||
|
position_ids=position_ids,
|
||||||
|
past_key_values=past_key_values,
|
||||||
|
inputs_embeds=inputs_embeds,
|
||||||
|
use_cache=use_cache,
|
||||||
|
output_attentions=output_attentions,
|
||||||
|
output_hidden_states=output_hidden_states,
|
||||||
|
)
|
||||||
|
hidden_states = transformer_outputs.last_hidden_state
|
||||||
|
logits = self.score(hidden_states)
|
||||||
|
|
||||||
|
if input_ids is not None:
|
||||||
|
batch_size = input_ids.shape[0]
|
||||||
|
else:
|
||||||
|
batch_size = inputs_embeds.shape[0]
|
||||||
|
|
||||||
|
if self.config.pad_token_id is None and batch_size != 1:
|
||||||
|
raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
|
||||||
|
if self.config.pad_token_id is None:
|
||||||
|
last_non_pad_token = -1
|
||||||
|
elif input_ids is not None:
|
||||||
|
# To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
|
||||||
|
non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
|
||||||
|
token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
|
||||||
|
last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
|
||||||
|
else:
|
||||||
|
last_non_pad_token = -1
|
||||||
|
logger.warning_once(
|
||||||
|
f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
|
||||||
|
"unexpected if using padding tokens in conjunction with `inputs_embeds.`"
|
||||||
|
)
|
||||||
|
|
||||||
|
pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]
|
||||||
|
|
||||||
|
loss = None
|
||||||
|
if labels is not None:
|
||||||
|
loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)
|
||||||
|
|
||||||
|
return SequenceClassifierOutputWithPast(
|
||||||
|
loss=loss,
|
||||||
|
logits=pooled_logits,
|
||||||
|
past_key_values=transformer_outputs.past_key_values,
|
||||||
|
hidden_states=transformer_outputs.hidden_states,
|
||||||
|
attentions=transformer_outputs.attentions,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@auto_docstring
|
||||||
|
class SmolLM3ForTokenClassification(SmolLM3PreTrainedModel):
|
||||||
|
def __init__(self, config):
|
||||||
|
super().__init__(config)
|
||||||
|
self.num_labels = config.num_labels
|
||||||
|
self.model = SmolLM3Model(config)
|
||||||
|
if getattr(config, "classifier_dropout", None) is not None:
|
||||||
|
classifier_dropout = config.classifier_dropout
|
||||||
|
elif getattr(config, "hidden_dropout", None) is not None:
|
||||||
|
classifier_dropout = config.hidden_dropout
|
||||||
|
else:
|
||||||
|
classifier_dropout = 0.1
|
||||||
|
self.dropout = nn.Dropout(classifier_dropout)
|
||||||
|
self.score = nn.Linear(config.hidden_size, config.num_labels)
|
||||||
|
|
||||||
|
# Initialize weights and apply final processing
|
||||||
|
self.post_init()
|
||||||
|
|
||||||
|
def get_input_embeddings(self):
|
||||||
|
return self.model.embed_tokens
|
||||||
|
|
||||||
|
def set_input_embeddings(self, value):
|
||||||
|
self.model.embed_tokens = value
|
||||||
|
|
||||||
|
@can_return_tuple
|
||||||
|
@auto_docstring
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
input_ids: Optional[torch.LongTensor] = None,
|
||||||
|
attention_mask: Optional[torch.Tensor] = None,
|
||||||
|
position_ids: Optional[torch.LongTensor] = None,
|
||||||
|
past_key_values: Optional[Cache] = None,
|
||||||
|
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||||
|
labels: Optional[torch.LongTensor] = None,
|
||||||
|
use_cache: Optional[bool] = None,
|
||||||
|
output_attentions: Optional[bool] = None,
|
||||||
|
output_hidden_states: Optional[bool] = None,
|
||||||
|
) -> TokenClassifierOutput:
|
||||||
|
r"""
|
||||||
|
labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
|
||||||
|
Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
|
||||||
|
config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
|
||||||
|
`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
|
||||||
|
"""
|
||||||
|
|
||||||
|
outputs: BaseModelOutputWithPast = self.model(
|
||||||
|
input_ids,
|
||||||
|
attention_mask=attention_mask,
|
||||||
|
position_ids=position_ids,
|
||||||
|
past_key_values=past_key_values,
|
||||||
|
inputs_embeds=inputs_embeds,
|
||||||
|
use_cache=use_cache,
|
||||||
|
output_attentions=output_attentions,
|
||||||
|
output_hidden_states=output_hidden_states,
|
||||||
|
)
|
||||||
|
sequence_output = outputs.last_hidden_state
|
||||||
|
sequence_output = self.dropout(sequence_output)
|
||||||
|
logits = self.score(sequence_output)
|
||||||
|
|
||||||
|
loss = None
|
||||||
|
if labels is not None:
|
||||||
|
loss = self.loss_function(logits, labels, self.config)
|
||||||
|
|
||||||
|
return TokenClassifierOutput(
|
||||||
|
loss=loss,
|
||||||
|
logits=logits,
|
||||||
|
hidden_states=outputs.hidden_states,
|
||||||
|
attentions=outputs.attentions,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@auto_docstring
|
||||||
|
class SmolLM3ForQuestionAnswering(SmolLM3PreTrainedModel):
|
||||||
|
base_model_prefix = "transformer"
|
||||||
|
|
||||||
|
def __init__(self, config):
|
||||||
|
super().__init__(config)
|
||||||
|
self.transformer = SmolLM3Model(config)
|
||||||
|
self.qa_outputs = nn.Linear(config.hidden_size, 2)
|
||||||
|
|
||||||
|
# Initialize weights and apply final processing
|
||||||
|
self.post_init()
|
||||||
|
|
||||||
|
def get_input_embeddings(self):
|
||||||
|
return self.transformer.embed_tokens
|
||||||
|
|
||||||
|
def set_input_embeddings(self, value):
|
||||||
|
self.transformer.embed_tokens = value
|
||||||
|
|
||||||
|
@can_return_tuple
|
||||||
|
@auto_docstring
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
input_ids: Optional[torch.LongTensor] = None,
|
||||||
|
attention_mask: Optional[torch.Tensor] = None,
|
||||||
|
position_ids: Optional[torch.LongTensor] = None,
|
||||||
|
past_key_values: Optional[Cache] = None,
|
||||||
|
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||||
|
start_positions: Optional[torch.LongTensor] = None,
|
||||||
|
end_positions: Optional[torch.LongTensor] = None,
|
||||||
|
output_attentions: Optional[bool] = None,
|
||||||
|
output_hidden_states: Optional[bool] = None,
|
||||||
|
**kwargs,
|
||||||
|
) -> QuestionAnsweringModelOutput:
|
||||||
|
outputs: BaseModelOutputWithPast = self.transformer(
|
||||||
|
input_ids,
|
||||||
|
attention_mask=attention_mask,
|
||||||
|
position_ids=position_ids,
|
||||||
|
past_key_values=past_key_values,
|
||||||
|
inputs_embeds=inputs_embeds,
|
||||||
|
output_attentions=output_attentions,
|
||||||
|
output_hidden_states=output_hidden_states,
|
||||||
|
)
|
||||||
|
|
||||||
|
sequence_output = outputs.last_hidden_state
|
||||||
|
|
||||||
|
logits = self.qa_outputs(sequence_output)
|
||||||
|
start_logits, end_logits = logits.split(1, dim=-1)
|
||||||
|
start_logits = start_logits.squeeze(-1).contiguous()
|
||||||
|
end_logits = end_logits.squeeze(-1).contiguous()
|
||||||
|
|
||||||
|
loss = None
|
||||||
|
if start_positions is not None and end_positions is not None:
|
||||||
|
loss = self.loss_function(start_logits, end_logits, start_positions, end_positions, **kwargs)
|
||||||
|
|
||||||
|
return QuestionAnsweringModelOutput(
|
||||||
|
loss=loss,
|
||||||
|
start_logits=start_logits,
|
||||||
|
end_logits=end_logits,
|
||||||
|
hidden_states=outputs.hidden_states,
|
||||||
|
attentions=outputs.attentions,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"SmolLM3PreTrainedModel",
|
||||||
|
"SmolLM3Model",
|
||||||
|
"SmolLM3ForCausalLM",
|
||||||
|
"SmolLM3ForSequenceClassification",
|
||||||
|
"SmolLM3ForTokenClassification",
|
||||||
|
"SmolLM3ForQuestionAnswering",
|
||||||
|
]
|
||||||
350
src/transformers/models/smollm3/modular_smollm3.py
Normal file
350
src/transformers/models/smollm3/modular_smollm3.py
Normal file
@@ -0,0 +1,350 @@
|
|||||||
|
# coding=utf-8
|
||||||
|
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on 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.
|
||||||
|
|
||||||
|
from typing import Callable, Optional
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from ...cache_utils import Cache
|
||||||
|
from ...configuration_utils import PretrainedConfig, layer_type_validation
|
||||||
|
from ...modeling_flash_attention_utils import FlashAttentionKwargs
|
||||||
|
from ...modeling_rope_utils import rope_config_validation
|
||||||
|
from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
|
||||||
|
from ...processing_utils import Unpack
|
||||||
|
from ...utils import logging
|
||||||
|
from ..llama.modeling_llama import (
|
||||||
|
LlamaAttention,
|
||||||
|
LlamaForCausalLM,
|
||||||
|
LlamaForQuestionAnswering,
|
||||||
|
LlamaForSequenceClassification,
|
||||||
|
LlamaForTokenClassification,
|
||||||
|
LlamaPreTrainedModel,
|
||||||
|
apply_rotary_pos_emb,
|
||||||
|
eager_attention_forward,
|
||||||
|
)
|
||||||
|
from ..qwen2.modeling_qwen2 import Qwen2Model
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3Config(PretrainedConfig):
|
||||||
|
r"""
|
||||||
|
This is the configuration class to store the configuration of a [`SmolLM3Model`]. It is used to instantiate a
|
||||||
|
SmolLM3 model according to the specified arguments, defining the model architecture. Instantiating a configuration
|
||||||
|
with the defaults will yield a similar configuration to that of the SmolLM3 3B.
|
||||||
|
e.g. [HuggingFaceTB/SmolLM3-3B](https://huggingface.co/HuggingFaceTB/SmolLM3-3B)
|
||||||
|
|
||||||
|
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
||||||
|
documentation from [`PretrainedConfig`] for more information.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
vocab_size (`int`, *optional*, defaults to 128256):
|
||||||
|
Vocabulary size of the SmolLM3 model. Defines the number of different tokens that can be represented by the
|
||||||
|
`inputs_ids` passed when calling [`SmolLM3Model`]
|
||||||
|
hidden_size (`int`, *optional*, defaults to 2048):
|
||||||
|
Dimension of the hidden representations.
|
||||||
|
intermediate_size (`int`, *optional*, defaults to 11008):
|
||||||
|
Dimension of the MLP representations.
|
||||||
|
num_hidden_layers (`int`, *optional*, defaults to 36):
|
||||||
|
Number of hidden layers in the Transformer encoder.
|
||||||
|
num_attention_heads (`int`, *optional*, defaults to 16):
|
||||||
|
Number of attention heads for each attention layer in the Transformer encoder.
|
||||||
|
num_key_value_heads (`int`, *optional*, defaults to 4):
|
||||||
|
This is the number of key_value heads that should be used to implement Grouped Query Attention. If
|
||||||
|
`num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
|
||||||
|
`num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
|
||||||
|
converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
|
||||||
|
by meanpooling all the original heads within that group. For more details checkout [this
|
||||||
|
paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `16`.
|
||||||
|
hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
|
||||||
|
The non-linear activation function (function or string) in the decoder.
|
||||||
|
max_position_embeddings (`int`, *optional*, defaults to 32768):
|
||||||
|
The maximum sequence length that this model might ever be used with.
|
||||||
|
initializer_range (`float`, *optional*, defaults to 0.02):
|
||||||
|
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
||||||
|
rms_norm_eps (`float`, *optional*, defaults to 1e-06):
|
||||||
|
The epsilon used by the rms normalization layers.
|
||||||
|
use_cache (`bool`, *optional*, defaults to `True`):
|
||||||
|
Whether or not the model should return the last key/values attentions (not used by all models). Only
|
||||||
|
relevant if `config.is_decoder=True`.
|
||||||
|
pad_token_id (`int`, *optional*, defaults to 128004):
|
||||||
|
The id of the padding token.
|
||||||
|
bos_token_id (`int`, *optional*, defaults to 128000):
|
||||||
|
The id of the beginning of sentence token.
|
||||||
|
eos_token_id (`int`, *optional*, defaults to 128001):
|
||||||
|
The id of the end of sentence token.
|
||||||
|
rope_theta (`float`, *optional*, defaults to 2000000.0):
|
||||||
|
The base period of the RoPE embeddings.
|
||||||
|
rope_scaling (`Dict`, *optional*):
|
||||||
|
Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
|
||||||
|
and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
|
||||||
|
accordingly.
|
||||||
|
Expected contents:
|
||||||
|
`rope_type` (`str`):
|
||||||
|
The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
|
||||||
|
'llama3'], with 'default' being the original RoPE implementation.
|
||||||
|
`factor` (`float`, *optional*):
|
||||||
|
Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
|
||||||
|
most scaling types, a `factor` of x will enable the model to handle sequences of length x *
|
||||||
|
original maximum pre-trained length.
|
||||||
|
`original_max_position_embeddings` (`int`, *optional*):
|
||||||
|
Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
|
||||||
|
pretraining.
|
||||||
|
`attention_factor` (`float`, *optional*):
|
||||||
|
Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
|
||||||
|
computation. If unspecified, it defaults to value recommended by the implementation, using the
|
||||||
|
`factor` field to infer the suggested value.
|
||||||
|
`beta_fast` (`float`, *optional*):
|
||||||
|
Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
|
||||||
|
ramp function. If unspecified, it defaults to 32.
|
||||||
|
`beta_slow` (`float`, *optional*):
|
||||||
|
Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
|
||||||
|
ramp function. If unspecified, it defaults to 1.
|
||||||
|
`short_factor` (`List[float]`, *optional*):
|
||||||
|
Only used with 'longrope'. The scaling factor to be applied to short contexts (<
|
||||||
|
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
||||||
|
size divided by the number of attention heads divided by 2
|
||||||
|
`long_factor` (`List[float]`, *optional*):
|
||||||
|
Only used with 'longrope'. The scaling factor to be applied to long contexts (<
|
||||||
|
`original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
|
||||||
|
size divided by the number of attention heads divided by 2
|
||||||
|
`low_freq_factor` (`float`, *optional*):
|
||||||
|
Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
|
||||||
|
`high_freq_factor` (`float`, *optional*):
|
||||||
|
Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
|
||||||
|
use_sliding_window (`bool`, *optional*, defaults to `False`):
|
||||||
|
Whether to use sliding window attention.
|
||||||
|
sliding_window (`int`, *optional*):
|
||||||
|
Sliding window attention (SWA) window size. If not specified, will default to `None`.
|
||||||
|
no_rope_layers (`List[int]`, *optional*):
|
||||||
|
List with at least the same length as the number of layers in the model.
|
||||||
|
A `1` at an index position indicates that the corresponding layer will use RoPE,
|
||||||
|
while a `0` indicates that it's a NoPE layer.
|
||||||
|
no_rope_layer_interval (`int`, *optional*, defaults to 4):
|
||||||
|
If `no_rope_layers` is `None`, it will be created using a NoPE layer every
|
||||||
|
`no_rope_layer_interval` layers.
|
||||||
|
layer_types (`list`, *optional*):
|
||||||
|
Attention pattern for each layer. Automatically computed based on sliding window and NoPE settings.
|
||||||
|
attention_bias (`bool`, *optional*, defaults to `False`):
|
||||||
|
Whether to use a bias in the query, key, value and output projection layers during self-attention.
|
||||||
|
attention_dropout (`float`, *optional*, defaults to 0.0):
|
||||||
|
The dropout ratio for the attention probabilities.
|
||||||
|
|
||||||
|
```python
|
||||||
|
>>> from transformers import SmolLM3Model, SmolLM3Config
|
||||||
|
|
||||||
|
>>> # Initializing a SmolLM3 style configuration
|
||||||
|
>>> configuration = SmolLM3Config()
|
||||||
|
|
||||||
|
>>> # Initializing a model from the SmolLM3 style configuration
|
||||||
|
>>> model = SmolLM3Model(configuration)
|
||||||
|
|
||||||
|
>>> # Accessing the model configuration
|
||||||
|
>>> configuration = model.config
|
||||||
|
```"""
|
||||||
|
|
||||||
|
model_type = "smollm3"
|
||||||
|
keys_to_ignore_at_inference = ["past_key_values"]
|
||||||
|
|
||||||
|
base_model_tp_plan = {
|
||||||
|
"layers.*.self_attn.q_proj": "colwise",
|
||||||
|
"layers.*.self_attn.k_proj": "colwise",
|
||||||
|
"layers.*.self_attn.v_proj": "colwise",
|
||||||
|
"layers.*.self_attn.o_proj": "rowwise",
|
||||||
|
"layers.*.mlp.gate_proj": "colwise",
|
||||||
|
"layers.*.mlp.up_proj": "colwise",
|
||||||
|
"layers.*.mlp.down_proj": "rowwise",
|
||||||
|
}
|
||||||
|
base_model_pp_plan = {
|
||||||
|
"embed_tokens": (["input_ids"], ["inputs_embeds"]),
|
||||||
|
"layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
|
||||||
|
"norm": (["hidden_states"], ["hidden_states"]),
|
||||||
|
}
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
vocab_size=128256,
|
||||||
|
hidden_size=2048,
|
||||||
|
intermediate_size=11008,
|
||||||
|
num_hidden_layers=36,
|
||||||
|
num_attention_heads=16,
|
||||||
|
num_key_value_heads=4,
|
||||||
|
hidden_act="silu",
|
||||||
|
max_position_embeddings=32768,
|
||||||
|
initializer_range=0.02,
|
||||||
|
rms_norm_eps=1e-6,
|
||||||
|
use_cache=True,
|
||||||
|
pad_token_id=128004,
|
||||||
|
bos_token_id=128000,
|
||||||
|
eos_token_id=128001,
|
||||||
|
rope_theta=2000000.0,
|
||||||
|
rope_scaling=None,
|
||||||
|
use_sliding_window=False,
|
||||||
|
sliding_window=None,
|
||||||
|
no_rope_layers=None,
|
||||||
|
no_rope_layer_interval=4,
|
||||||
|
layer_types=None,
|
||||||
|
attention_bias=False,
|
||||||
|
attention_dropout=0.0,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
super().__init__(
|
||||||
|
pad_token_id=pad_token_id,
|
||||||
|
bos_token_id=bos_token_id,
|
||||||
|
eos_token_id=eos_token_id,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
self.vocab_size = vocab_size
|
||||||
|
self.max_position_embeddings = max_position_embeddings
|
||||||
|
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.use_sliding_window = use_sliding_window
|
||||||
|
self.sliding_window = sliding_window
|
||||||
|
|
||||||
|
# for backward compatibility
|
||||||
|
if num_key_value_heads is None:
|
||||||
|
num_key_value_heads = num_attention_heads
|
||||||
|
|
||||||
|
self.num_key_value_heads = num_key_value_heads
|
||||||
|
self.hidden_act = hidden_act
|
||||||
|
self.initializer_range = initializer_range
|
||||||
|
self.rms_norm_eps = rms_norm_eps
|
||||||
|
self.use_cache = use_cache
|
||||||
|
self.rope_theta = rope_theta
|
||||||
|
self.rope_scaling = rope_scaling
|
||||||
|
self.attention_bias = attention_bias
|
||||||
|
self.attention_dropout = attention_dropout
|
||||||
|
|
||||||
|
if no_rope_layers is None:
|
||||||
|
self.no_rope_layers = [
|
||||||
|
int((layer_idx + 1) % no_rope_layer_interval != 0) for layer_idx in range(num_hidden_layers)
|
||||||
|
]
|
||||||
|
else:
|
||||||
|
self.no_rope_layers = no_rope_layers
|
||||||
|
|
||||||
|
self.no_rope_layer_interval = no_rope_layer_interval
|
||||||
|
|
||||||
|
# Update layer_types based on sliding window and NoPE pattern
|
||||||
|
if layer_types is None:
|
||||||
|
layer_types = []
|
||||||
|
for layer_idx in range(num_hidden_layers):
|
||||||
|
has_rope = self.no_rope_layers[layer_idx]
|
||||||
|
if use_sliding_window and sliding_window is not None and not has_rope:
|
||||||
|
layer_types.append("sliding_attention")
|
||||||
|
else:
|
||||||
|
layer_types.append("full_attention")
|
||||||
|
|
||||||
|
self.layer_types = layer_types
|
||||||
|
layer_type_validation(self.layer_types)
|
||||||
|
|
||||||
|
# Validate the correctness of rotary position embeddings parameters
|
||||||
|
# BC: if there is a 'type' field, move it to 'rope_type'.
|
||||||
|
if self.rope_scaling is not None and "type" in self.rope_scaling:
|
||||||
|
self.rope_scaling["rope_type"] = self.rope_scaling["type"]
|
||||||
|
rope_config_validation(self)
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3Attention(LlamaAttention):
|
||||||
|
def __init__(self, config: SmolLM3Config, layer_idx: int):
|
||||||
|
super().__init__(config, layer_idx)
|
||||||
|
|
||||||
|
self.use_rope = config.no_rope_layers[layer_idx]
|
||||||
|
self.sliding_window = (
|
||||||
|
config.sliding_window
|
||||||
|
if config.use_sliding_window and config.layer_types[layer_idx] == "sliding_attention"
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
||||||
|
attention_mask: Optional[torch.Tensor],
|
||||||
|
past_key_value: Optional[Cache] = None,
|
||||||
|
cache_position: Optional[torch.LongTensor] = None,
|
||||||
|
**kwargs: Unpack[FlashAttentionKwargs],
|
||||||
|
) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
|
||||||
|
input_shape = hidden_states.shape[:-1]
|
||||||
|
hidden_shape = (*input_shape, -1, self.head_dim)
|
||||||
|
|
||||||
|
query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||||
|
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||||
|
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||||
|
|
||||||
|
if self.use_rope:
|
||||||
|
cos, sin = position_embeddings
|
||||||
|
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
|
||||||
|
|
||||||
|
if past_key_value is not None:
|
||||||
|
cache_kwargs = {"cache_position": cache_position}
|
||||||
|
key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
|
||||||
|
|
||||||
|
attention_interface: Callable = eager_attention_forward
|
||||||
|
if self.config._attn_implementation != "eager":
|
||||||
|
attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
|
||||||
|
|
||||||
|
attn_output, attn_weights = attention_interface(
|
||||||
|
self,
|
||||||
|
query_states,
|
||||||
|
key_states,
|
||||||
|
value_states,
|
||||||
|
attention_mask,
|
||||||
|
dropout=0.0 if not self.training else self.attention_dropout,
|
||||||
|
scaling=self.scaling,
|
||||||
|
sliding_window=self.sliding_window,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
attn_output = attn_output.reshape(*input_shape, -1).contiguous()
|
||||||
|
attn_output = self.o_proj(attn_output)
|
||||||
|
return attn_output, attn_weights
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3PreTrainedModel(LlamaPreTrainedModel):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3Model(Qwen2Model):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3ForCausalLM(LlamaForCausalLM):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3ForSequenceClassification(LlamaForSequenceClassification):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3ForTokenClassification(LlamaForTokenClassification):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3ForQuestionAnswering(LlamaForQuestionAnswering):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"SmolLM3Config",
|
||||||
|
"SmolLM3PreTrainedModel",
|
||||||
|
"SmolLM3Model",
|
||||||
|
"SmolLM3ForCausalLM",
|
||||||
|
"SmolLM3ForSequenceClassification",
|
||||||
|
"SmolLM3ForTokenClassification",
|
||||||
|
"SmolLM3ForQuestionAnswering",
|
||||||
|
]
|
||||||
0
tests/models/smollm3/__init__.py
Normal file
0
tests/models/smollm3/__init__.py
Normal file
227
tests/models/smollm3/test_modeling_smollm3.py
Normal file
227
tests/models/smollm3/test_modeling_smollm3.py
Normal file
@@ -0,0 +1,227 @@
|
|||||||
|
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on 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.
|
||||||
|
"""Testing suite for the PyTorch SmolLM3 model."""
|
||||||
|
|
||||||
|
import gc
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from packaging import version
|
||||||
|
from parameterized import parameterized
|
||||||
|
|
||||||
|
from transformers import AutoTokenizer, SmolLM3Config, is_torch_available
|
||||||
|
from transformers.generation.configuration_utils import GenerationConfig
|
||||||
|
from transformers.testing_utils import (
|
||||||
|
backend_empty_cache,
|
||||||
|
is_flaky,
|
||||||
|
require_bitsandbytes,
|
||||||
|
require_flash_attn,
|
||||||
|
require_torch,
|
||||||
|
require_torch_sdpa,
|
||||||
|
slow,
|
||||||
|
torch_device,
|
||||||
|
)
|
||||||
|
from transformers.utils.import_utils import is_torch_greater_or_equal
|
||||||
|
|
||||||
|
|
||||||
|
if is_torch_available():
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from transformers import (
|
||||||
|
SmolLM3ForCausalLM,
|
||||||
|
SmolLM3ForQuestionAnswering,
|
||||||
|
SmolLM3ForSequenceClassification,
|
||||||
|
SmolLM3ForTokenClassification,
|
||||||
|
SmolLM3Model,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester
|
||||||
|
from ...test_modeling_common import (
|
||||||
|
TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION,
|
||||||
|
ModelTesterMixin,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class SmolLM3ModelTester(CausalLMModelTester):
|
||||||
|
config_class = SmolLM3Config
|
||||||
|
if is_torch_available():
|
||||||
|
base_model_class = SmolLM3Model
|
||||||
|
causal_lm_class = SmolLM3ForCausalLM
|
||||||
|
sequence_class = SmolLM3ForSequenceClassification
|
||||||
|
token_class = SmolLM3ForTokenClassification
|
||||||
|
question_answering_class = SmolLM3ForQuestionAnswering
|
||||||
|
|
||||||
|
|
||||||
|
@require_torch
|
||||||
|
class SmolLM3ModelTest(CausalLMModelTest, unittest.TestCase):
|
||||||
|
all_model_classes = (
|
||||||
|
(
|
||||||
|
SmolLM3Model,
|
||||||
|
SmolLM3ForCausalLM,
|
||||||
|
SmolLM3ForSequenceClassification,
|
||||||
|
SmolLM3ForTokenClassification,
|
||||||
|
SmolLM3ForQuestionAnswering,
|
||||||
|
)
|
||||||
|
if is_torch_available()
|
||||||
|
else ()
|
||||||
|
)
|
||||||
|
test_headmasking = False
|
||||||
|
test_pruning = False
|
||||||
|
model_tester_class = SmolLM3ModelTester
|
||||||
|
pipeline_model_mapping = (
|
||||||
|
{
|
||||||
|
"feature-extraction": SmolLM3Model,
|
||||||
|
"text-classification": SmolLM3ForSequenceClassification,
|
||||||
|
"token-classification": SmolLM3ForTokenClassification,
|
||||||
|
"text-generation": SmolLM3ForCausalLM,
|
||||||
|
"question-answering": SmolLM3ForQuestionAnswering,
|
||||||
|
}
|
||||||
|
if is_torch_available()
|
||||||
|
else {}
|
||||||
|
)
|
||||||
|
|
||||||
|
@parameterized.expand(TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION)
|
||||||
|
@require_torch_sdpa
|
||||||
|
@is_flaky()
|
||||||
|
def test_eager_matches_sdpa_inference(self, *args):
|
||||||
|
# flaky test_eager_matches_sdpa_inference_24_fp32_pad_left_output_attentions
|
||||||
|
return getattr(ModelTesterMixin, self._testMethodName)(self)
|
||||||
|
|
||||||
|
|
||||||
|
@require_torch
|
||||||
|
class SmolLM3IntegrationTest(unittest.TestCase):
|
||||||
|
model_id = "HuggingFaceTB/SmolLM3-3B"
|
||||||
|
|
||||||
|
@slow
|
||||||
|
def test_model_3b_logits(self):
|
||||||
|
input_ids = [1, 306, 4658, 278, 6593, 310, 2834, 338]
|
||||||
|
model = SmolLM3ForCausalLM.from_pretrained(self.model_id, device_map="auto")
|
||||||
|
input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
|
||||||
|
with torch.no_grad():
|
||||||
|
out = model(input_ids).logits.float().cpu()
|
||||||
|
# Expected mean on dim = -1
|
||||||
|
EXPECTED_MEAN = torch.tensor([[9.3306, 8.1721, 6.4764, 7.6011, 11.1218, 7.5343, 7.1195, 8.0956]])
|
||||||
|
torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, rtol=1e-2, atol=1e-2)
|
||||||
|
# slicing logits[0, 0, 0:30]
|
||||||
|
EXPECTED_SLICE = torch.tensor(
|
||||||
|
[15.7759, 17.6274, 16.3404, 14.5543, 13.1366, 14.2475, 15.8710, 15.6753, 12.3856, 13.0386, 14.0792, 12.7253,
|
||||||
|
13.9634, 12.1271, 12.4320, 16.0329, 17.3975, 17.1396, 17.8666, 17.0103, 17.2962, 16.8777, 16.7144, 16.3023,
|
||||||
|
16.6084, 12.4649, 12.0723, 14.1148, 14.8239, 15.2733]) # fmt: skip
|
||||||
|
torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE, rtol=1e-4, atol=1e-4)
|
||||||
|
|
||||||
|
del model
|
||||||
|
backend_empty_cache(torch_device)
|
||||||
|
gc.collect()
|
||||||
|
|
||||||
|
@slow
|
||||||
|
def test_model_3b_generation(self):
|
||||||
|
EXPECTED_TEXT_COMPLETION = """Gravity is the force that pulls objects toward the center of the Earth. It is a force that is always present, even"""
|
||||||
|
prompt = "Gravity is the force"
|
||||||
|
tokenizer = AutoTokenizer.from_pretrained(self.model_id)
|
||||||
|
model = SmolLM3ForCausalLM.from_pretrained(self.model_id, device_map="auto")
|
||||||
|
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.model.embed_tokens.weight.device)
|
||||||
|
|
||||||
|
# greedy generation outputs
|
||||||
|
generated_ids = model.generate(input_ids, max_new_tokens=20, temperature=0)
|
||||||
|
text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
|
||||||
|
self.assertEqual(EXPECTED_TEXT_COMPLETION, text)
|
||||||
|
|
||||||
|
del model
|
||||||
|
backend_empty_cache(torch_device)
|
||||||
|
gc.collect()
|
||||||
|
|
||||||
|
@require_bitsandbytes
|
||||||
|
@slow
|
||||||
|
@require_flash_attn
|
||||||
|
@pytest.mark.flash_attn_test
|
||||||
|
def test_model_3b_long_prompt(self):
|
||||||
|
EXPECTED_OUTPUT_TOKEN_IDS = [306, 338]
|
||||||
|
# An input with 4097 tokens that is above the size of the sliding window
|
||||||
|
input_ids = [1] + [306, 338] * 2048
|
||||||
|
model = SmolLM3ForCausalLM.from_pretrained(
|
||||||
|
self.model_id,
|
||||||
|
device_map="auto",
|
||||||
|
load_in_4bit=True,
|
||||||
|
attn_implementation="flash_attention_2",
|
||||||
|
)
|
||||||
|
input_ids = torch.tensor([input_ids]).to(model.model.embed_tokens.weight.device)
|
||||||
|
generated_ids = model.generate(input_ids, max_new_tokens=4, temperature=0)
|
||||||
|
self.assertEqual(EXPECTED_OUTPUT_TOKEN_IDS, generated_ids[0][-2:].tolist())
|
||||||
|
|
||||||
|
# Assisted generation
|
||||||
|
assistant_model = model
|
||||||
|
assistant_model.generation_config.num_assistant_tokens = 2
|
||||||
|
assistant_model.generation_config.num_assistant_tokens_schedule = "constant"
|
||||||
|
generated_ids = model.generate(input_ids, max_new_tokens=4, temperature=0)
|
||||||
|
self.assertEqual(EXPECTED_OUTPUT_TOKEN_IDS, generated_ids[0][-2:].tolist())
|
||||||
|
|
||||||
|
del assistant_model
|
||||||
|
del model
|
||||||
|
backend_empty_cache(torch_device)
|
||||||
|
gc.collect()
|
||||||
|
|
||||||
|
@slow
|
||||||
|
def test_export_static_cache(self):
|
||||||
|
if version.parse(torch.__version__) < version.parse("2.4.0"):
|
||||||
|
self.skipTest(reason="This test requires torch >= 2.4 to run.")
|
||||||
|
|
||||||
|
from transformers.integrations.executorch import (
|
||||||
|
TorchExportableModuleWithStaticCache,
|
||||||
|
convert_and_export_with_cache,
|
||||||
|
)
|
||||||
|
|
||||||
|
tokenizer = AutoTokenizer.from_pretrained(
|
||||||
|
self.model_id, pad_token="<|finetune_right_pad_id|>", padding_side="right"
|
||||||
|
)
|
||||||
|
EXPECTED_TEXT_COMPLETION = "Gravity is the force that pulls objects toward the center of the Earth. It is a force that is always present, and"
|
||||||
|
max_generation_length = tokenizer(EXPECTED_TEXT_COMPLETION, return_tensors="pt", padding=True)[
|
||||||
|
"input_ids"
|
||||||
|
].shape[-1]
|
||||||
|
|
||||||
|
# Load model
|
||||||
|
device = "cpu"
|
||||||
|
dtype = torch.bfloat16
|
||||||
|
cache_implementation = "static"
|
||||||
|
attn_implementation = "sdpa"
|
||||||
|
batch_size = 1
|
||||||
|
model = SmolLM3ForCausalLM.from_pretrained(
|
||||||
|
self.model_id,
|
||||||
|
device_map=device,
|
||||||
|
torch_dtype=dtype,
|
||||||
|
attn_implementation=attn_implementation,
|
||||||
|
generation_config=GenerationConfig(
|
||||||
|
use_cache=True,
|
||||||
|
cache_implementation=cache_implementation,
|
||||||
|
max_length=max_generation_length,
|
||||||
|
cache_config={
|
||||||
|
"batch_size": batch_size,
|
||||||
|
"max_cache_len": max_generation_length,
|
||||||
|
},
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
prompt = ["Gravity is the force"]
|
||||||
|
prompt_tokens = tokenizer(prompt, return_tensors="pt", padding=True).to(model.device)
|
||||||
|
prompt_token_ids = prompt_tokens["input_ids"]
|
||||||
|
max_new_tokens = max_generation_length - prompt_token_ids.shape[-1]
|
||||||
|
|
||||||
|
# Static Cache + export
|
||||||
|
strict = is_torch_greater_or_equal("2.7.0") # Due to https://github.com/pytorch/pytorch/issues/150994
|
||||||
|
exported_program = convert_and_export_with_cache(model, strict=strict)
|
||||||
|
ep_generated_ids = TorchExportableModuleWithStaticCache.generate(
|
||||||
|
exported_program=exported_program, prompt_token_ids=prompt_token_ids, max_new_tokens=max_new_tokens
|
||||||
|
)
|
||||||
|
ep_generated_text = tokenizer.batch_decode(ep_generated_ids, skip_special_tokens=True)
|
||||||
|
self.assertEqual(EXPECTED_TEXT_COMPLETION, ep_generated_text)
|
||||||
@@ -272,6 +272,7 @@ SPECIAL_CASES_TO_ALLOW = {
|
|||||||
"attention_chunk_size",
|
"attention_chunk_size",
|
||||||
],
|
],
|
||||||
"Llama4VisionConfig": ["multi_modal_projector_bias", "norm_eps"],
|
"Llama4VisionConfig": ["multi_modal_projector_bias", "norm_eps"],
|
||||||
|
"SmolLM3Config": ["no_rope_layer_interval"],
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user