Skip to content

Commit 6284e1b

Browse files
lkm2835fynnsu
andauthored
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

File tree

1 file changed

+1
-0
lines changed

1 file changed

+1
-0
lines changed

src/llmcompressor/modifiers/awq/mappings.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,7 @@ class AWQMapping:
166166
"Llama4ForConditionalGeneration": _default_mappings,
167167
"Mistral3ForConditionalGeneration": _default_mappings,
168168
"MistralForCausalLM": _default_mappings,
169+
"Olmo3ForCausalLM": _exaone4_mappings,
169170
"Phi3ForCausalLM": _phi_mappings,
170171
"Phi3VForCausalLM": _phi_mappings,
171172
"Qwen2ForCausalLM": _default_mappings,

0 commit comments

Comments
 (0)