Skip to content

FP8 training enhancements #3496

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Oct 27, 2025
Merged

FP8 training enhancements #3496

Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
78f7c79
Fix FP8 for models with non 8 multiple weights
Datta0 Oct 23, 2025
c315cba
patch fp8 forward methods for compiled models
Datta0 Oct 23, 2025
87cff2e
patch hf quantizer for fp8
Datta0 Oct 23, 2025
37a019a
Failsafe import of fbgemmfp8linear and fp8linear
Datta0 Oct 27, 2025
0d61b24
Beautify
Datta0 Oct 27, 2025
File filter

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Failsafe import of fbgemmfp8linear and fp8linear
  • Loading branch information
@Datta0
Datta0 committed Oct 27, 2025
commit 37a019a333a5a1de00cf6884e5089e2c8dff3379
13 changes: 9 additions & 4 deletions unsloth/kernels/fp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,16 +17,19 @@
import triton.language as tl
from torch.nn import functional as F
import math
from unsloth_zoo.log import logger

try:
from transformers.integrations.finegrained_fp8 import FP8Linear
except ImportError:
raise ImportError("Unsloth: FP8 models need importing FP8Linear from `transformers.integrations.finegrained_fp8` but we don't see it.")
FP8Linear = None
logger.log("Unsloth: FP8 models need importing FP8Linear from `transformers.integrations.finegrained_fp8` but we don't see it.")

try:
from transformers.integrations.fbgemm_fp8 import FbgemmFp8Linear
except ImportError:
raise ImportError("Unsloth: FP8 models need importing FbgemmFP8Linear from `transformers.integrations.fbgemm_fp8` but we don't see it.")
FbgemmFp8Linear = None
logger.log("Unsloth: FP8 models need importing FbgemmFP8Linear from `transformers.integrations.fbgemm_fp8` but we don't see it.")

try:
from fbgemm_gpu.experimental.gemm.triton_gemm.fp8_gemm import triton_quantize_fp8_block
Expand Down Expand Up @@ -486,5 +489,7 @@ def patched_forward(self, X):


# Patch the forward functions of the layers (for compiled models)
FbgemmFp8Linear.forward = module_forward_patch(fbgemm_fp8_linear, 'weight_scale')
FP8Linear.forward = module_forward_patch(fp8_block_quant_forward, 'weight_scale_inv')
if FbgemmFp8Linear is not None:
FbgemmFp8Linear.forward = module_forward_patch(fbgemm_fp8_linear, 'weight_scale')
if FP8Linear is not None:
FP8Linear.forward = module_forward_patch(fp8_block_quant_forward, 'weight_scale_inv')