Add 128x128 PerBlock FP8 + RL

[andrewor14]

Summary: Following #3440, this PR extends torchao FP8 + RL support to also handle 128x128 PerBlock granularity (in addition to PerRow).

Example usage:

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/Qwen3-8B-Base",
    max_seq_length = 2048,
    load_in_4bit = False,
    fast_inference = True,
    max_lora_rank = 32,
    load_in_fp8 = "block",  # or "row" or True
)

Initial results: TBD

Note:

• Requires Add PerBlock to safe globals pytorch/ao#3370

[gemini-code-assist]

Summary of Changes

Hello @andrewor14, I'm Gemini Code Assist¹! 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 significantly enhances the FP8 quantization capabilities within the unsloth library by integrating 128x128 PerBlock granularity. This new option provides users with more control over the quantization process, potentially leading to improved model performance or memory efficiency for specific use cases. The changes involve updating the core quantization logic and refactoring related utility functions to support the new granularity mode.

Highlights

• Expanded FP8 Granularity: Introduced support for 128x128 PerBlock granularity for FP8 quantization, complementing the existing PerRow method.
• Flexible FP8 Loading: The load_in_fp8 parameter now accepts "block" or "row" strings, allowing users to specify the desired quantization granularity.
• Refactored FP8 Configuration: Centralized FP8 mode determination and validation into a new utility function, _get_fp8_mode_and_check_settings.

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. ↩

[gemini-code-assist]

Code Review

This pull request successfully extends FP8 support to include 128x128 PerBlock granularity, in addition to the existing PerRow method. The changes are well-implemented across unsloth/kernels/utils.py, unsloth/models/loader.py, and unsloth/models/loader_utils.py. The logic for selecting and configuring the FP8 mode is clear, and the necessary adjustments to the kernel-level functions are correctly handled. I have one suggestion regarding code maintainability by removing a redundant validation check. Overall, the changes are solid and align well with the PR's objective.

[gemini-code-assist]

This validation is redundant, as fp8_mode is already validated in _get_fp8_mode_and_check_settings before being passed to this function. For internal functions, it's better to rely on assertions for contract checking rather than raising user-facing ValueErrors. This avoids duplicated validation logic and makes the code cleaner.

Consider removing this else block. If you want to keep a check for robustness, an assert would be more appropriate, for example:

assert fp8_mode in ['row', 'block']

However, given the call chain, even an assert is likely unnecessary.

[andrewor14]

@danielhanchen I had to change this back for this to work. When I tried AutoModel it removed the lm_head from the state dict, which later caused an out of bounds exception on this line: https://github.com/unslothai/unsloth-zoo/blob/54dce973426ee61670e15b720619c8539bf05104/unsloth_zoo/vllm_utils.py#L1110. Maybe we can generalize this to multimodal models later

[danielhanchen]

Oh ok oh wait I can make this work for vision models without using AutoModel