Commit 6284e1b
Add Olmo3 AWQ mapping (#2063)
Huggingface: [Olmo 3
Collection](https://huggingface.co/collections/allenai/olmo-3)
Transformers:
[Olmo3DecoderLayer](https://github.com/huggingface/transformers/blob/main/src/transformers/models/olmo3/modeling_olmo3.py#L245-L275)
```python
class Olmo3DecoderLayer(GradientCheckpointingLayer):
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[Cache] = None,
use_cache: Optional[bool] = False,
cache_position: Optional[torch.LongTensor] = None,
position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
**kwargs: Unpack[TransformersKwargs],
) -> torch.Tensor:
residual = hidden_states
hidden_states, _ = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
use_cache=use_cache,
cache_position=cache_position,
position_embeddings=position_embeddings,
**kwargs,
)
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = residual + hidden_states
# Fully Connected
residual = hidden_states
hidden_states = self.mlp(hidden_states)
hidden_states = self.post_feedforward_layernorm(hidden_states)
hidden_states = residual + hidden_states
return hidden_states
```
Olmo3 uses the same decoder layer structure as EXAONE4. Related to
#2046.
### Test
Tested it with `examples/awq/llama_example.py` and modified
`MODEL_ID="allenai/Olmo-3-1125-32B"`.
with `_exaone4_mappings`
```
$ python3 examples/awq/llama_example.py
========== SAMPLE GENERATION ==============
Hello my name is John and I am a 25 year old male. I have been having a problem with my left testicle for about 2 years now. I have been to the doctor and they have done a ultrasound and a CT scan and they have found nothing wrong. I have been having pain in my left testicle for about 2 years now. It is a dull pain that is always there. It is not a sharp pain. It is more of a dull ache. I have also been having a
==========================================
```
with `_default_mappings`
```
$ python3 examples/awq/llama_example.py
========== SAMPLE GENERATION ==============
Hello my name is
•
==========================================
```
Signed-off-by: lkm2835 <[email protected]>
Co-authored-by: Fynn Schmitt-Ulms <[email protected]>1 parent 0340724 commit 6284e1b
1 file changed
+1
-0
lines changed| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
166 | 166 | | |
167 | 167 | | |
168 | 168 | | |
| 169 | + | |
169 | 170 | | |
170 | 171 | | |
171 | 172 | | |
| |||
0 commit comments