Skip to content

Conversation

@lkm2835
Copy link
Contributor

@lkm2835 lkm2835 commented Nov 22, 2025

Huggingface: Olmo 3 Collection
Transformers: Olmo3DecoderLayer

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]>
@github-actions
Copy link

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello @lkm2835, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces support for quantizing Olmo3 language models using Activation-aware Weight Quantization (AWQ). By leveraging an existing mapping configuration that aligns with Olmo3's architectural design, this change allows for improved inference efficiency for Olmo3 models. The integration has been validated through testing, demonstrating successful text generation post-quantization.

Highlights

  • Olmo3 AWQ Mapping: Added Olmo3ForCausalLM to the AWQ mapping configurations, enabling support for Activation-aware Weight Quantization for Olmo3 models.
  • Architectural Alignment: The Olmo3 model is mapped to the _exaone4_mappings due to its decoder layer structure being similar to EXAONE4, ensuring correct quantization behavior.
  • Verification: The integration was tested by modifying an existing AWQ example script to use allenai/Olmo-3-1125-32B, which successfully produced coherent text generation, confirming the mapping's effectiveness.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request adds an AWQ (Activation-aware Weight Quantization) mapping for the Olmo3 model architecture. The change correctly identifies that Olmo3 shares a similar decoder structure with EXAONE4, particularly its pre-normalization scheme, and thus reuses the _exaone4_mappings. The new mapping is added to AWQ_MAPPING_REGISTRY in the correct alphabetical position, maintaining code organization. The provided test results confirm that this change is necessary and effective for generating coherent output with the quantized Olmo3 model. The change is straightforward, well-justified, and I have no comments.

Copy link
Collaborator

@kylesayrs kylesayrs left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants