Enable FP8 + RL training for bf16 models (#3440)
* Enable FP8 + RL training for bf16 models **Summary:** Enable FP8 + RL training using TorchAO for 1.33x faster training and 42% less model memory usage: - We quantize the frozen LoRA weights into fp8 and keep the LoRA adapters in bf16 - We leverage TorchAO's `Float8Tensor`, which calls into fbgemm's fp8 x fp8 rowwise matmul kernel - For now, we need to do an offline quantization first, because vllm doesn't support on-the-fly quantization for torchao yet (this is in progress: https://github.com/vllm-project/vllm/pull/26327) **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 = True, # set this to True ) \# the rest is the same as before model = FastLanguageModel.get_peft_model(...) ``` **Initial results:** ``` \# fp8 {'train_runtime': 1725.4337, 'train_samples_per_second': 0.232, 'train_steps_per_second': 0.058, 'train_loss': 0.00015715716748673002, 'epoch': 0.01} \# bf16 {'train_runtime': 2297.8145, 'train_samples_per_second': 0.174, 'train_steps_per_second': 0.044, 'train_loss': 0.00016081033063528594, 'epoch': 0.01} ``` <img width="1199" height="448" alt="Screenshot 2025-11-11 at 4 10 50 PM" src="https://github.com/user-attachments/assets/b6304afd-89e9-42b1-8064-775807e17b23" /> Test script: https://gist.github.com/andrewor14/5b85119fae46845d07b608d420907423 **Requires:** - https://github.com/pytorch/ao/pull/3158 (torchao nightly or 0.15.0+) - https://github.com/unslothai/unsloth-zoo/pull/351 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update utils.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * _get_inference_mode_context_manager * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update utils.py * Update utils.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
andrewor14
GitHub
parent
ecacc17f7e
commit
89d8677c0c
6 changed files with 257 additions and 12 deletions
|
|
@ -12,6 +12,7 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import importlib
|
||||
import triton
|
||||
import ctypes
|
||||
|
||||
|
|
@ -35,7 +36,7 @@ import functools
|
|||
import torch
|
||||
|
||||
torch_Tensor = torch.Tensor
|
||||
from packaging.version import Version
|
||||
from unsloth_zoo.utils import Version
|
||||
|
||||
if DEVICE_TYPE == "xpu" and Version(torch.__version__) < Version("2.6.0"):
|
||||
raise RuntimeError(
|
||||
|
|
@ -55,7 +56,6 @@ if DEVICE_TYPE == "xpu":
|
|||
|
||||
|
||||
# tl.math.tanh now is libdevice.tanh
|
||||
from packaging.version import Version
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
|
|
@ -211,6 +211,22 @@ torch_float16 = torch.float16
|
|||
torch_bfloat16 = torch.bfloat16
|
||||
|
||||
|
||||
# Check whether torchao can be imported to get Float8Tensor
|
||||
if importlib.util.find_spec("torchao") is not None:
|
||||
try:
|
||||
from torchao.quantization import Float8Tensor
|
||||
except:
|
||||
import torchao
|
||||
|
||||
if Version(torchao.__version__) >= Version("0.15.0"):
|
||||
print(
|
||||
f"Unsloth: `from torchao.quantization import Float8Tensor` failed on version={torchao.__version__}"
|
||||
)
|
||||
Float8Tensor = type(None)
|
||||
else:
|
||||
Float8Tensor = type(None)
|
||||
|
||||
|
||||
def QUANT_STATE(W):
|
||||
return getattr(W, "quant_state", None)
|
||||
|
||||
|
|
@ -335,6 +351,13 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM:
|
|||
@torch.inference_mode
|
||||
def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False):
|
||||
# TODO: After adding XPU BNB support, check this function
|
||||
if isinstance(W, Float8Tensor):
|
||||
# TorchAO Float8Tensor
|
||||
# In the backward pass, rowwise scaled becomes colwise scaled after we
|
||||
# transpose the weight tensor. Use this case to detect backward
|
||||
assert W.ndim == 2
|
||||
if W.block_size[0] == W.shape[0] and W.block_size[1] == 1:
|
||||
return W.dequantize()
|
||||
if quant_state is None:
|
||||
return W
|
||||
if W.dtype == torch.float8_e4m3fn:
|
||||
|
|
@ -441,6 +464,13 @@ elif DEVICE_TYPE in ("cuda", "hip") and HAS_CUDA_STREAM:
|
|||
|
||||
@torch.inference_mode
|
||||
def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False):
|
||||
if isinstance(W, Float8Tensor):
|
||||
# TorchAO Float8Tensor
|
||||
# In the backward pass, rowwise scaled becomes colwise scaled after we
|
||||
# transpose the weight tensor. Use this case to detect backward
|
||||
assert W.ndim == 2
|
||||
if W.block_size[0] == W.shape[0] and W.block_size[1] == 1:
|
||||
return W.dequantize()
|
||||
if quant_state is None:
|
||||
return W
|
||||
if W.dtype == torch.float8_e4m3fn:
|
||||
|
|
@ -551,6 +581,13 @@ else:
|
|||
|
||||
@torch.inference_mode
|
||||
def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False):
|
||||
if isinstance(W, Float8Tensor):
|
||||
# TorchAO Float8Tensor
|
||||
# In the backward pass, rowwise scaled becomes colwise scaled after we
|
||||
# transpose the weight tensor. Use this case to detect backward
|
||||
assert W.ndim == 2
|
||||
if W.block_size[0] == W.shape[0] and W.block_size[1] == 1:
|
||||
return W.dequantize()
|
||||
if quant_state is None:
|
||||
return W
|
||||
if W.dtype == torch.float8_e4m3fn:
|
||||
|
|
@ -987,8 +1024,8 @@ def matmul_lora(X, W, W_quant, A, B, s, out = None):
|
|||
if W.dtype == torch.float8_e4m3fn:
|
||||
out = fp8_linear(X, W, W_quant)
|
||||
else:
|
||||
W = fast_dequantize(W.t(), W_quant, use_global_buffer = True)
|
||||
out = torch_matmul(X, W, out = out)
|
||||
W = fast_dequantize(W, W_quant, use_global_buffer = True)
|
||||
out = torch_matmul(X, W.t(), out = out)
|
||||
if W_quant is not None:
|
||||
del W
|
||||
|
||||
|
|
|
|||
|
|
@ -71,6 +71,7 @@ __all__ = [
|
|||
"dequantize_module_weight",
|
||||
"patch_hf_quantizer",
|
||||
"verify_fp8_support_if_applicable",
|
||||
"_get_inference_mode_context_manager",
|
||||
]
|
||||
|
||||
import torch
|
||||
|
|
@ -2056,7 +2057,7 @@ except:
|
|||
|
||||
@dataclass
|
||||
class TorchAOConfig:
|
||||
qat_scheme: str = "int4"
|
||||
qat_scheme: Optional[str] = "int4"
|
||||
|
||||
# Each (config, filter_fn) pair defines a quantization rule
|
||||
base_config_and_filter_fns: List[
|
||||
|
|
@ -2306,3 +2307,22 @@ def verify_fp8_support_if_applicable(model_config):
|
|||
raise ValueError(
|
||||
f"Unsloth: FP8 quantization is only supported on L4 and higher GPUs with compute capability 8.9 or higher. You are using {torch.cuda.get_device_name()}. Refer to https://developer.nvidia.com/cuda-gpus for more details."
|
||||
)
|
||||
|
||||
|
||||
def _get_inference_mode_context_manager(model: torch.nn.Module):
|
||||
"""
|
||||
If the state dict was quantized using torchao, we will run into
|
||||
the following error when calling ops like aten.t() in inference mode.
|
||||
This is a bug in PyTorch that affects all tensor subclasses.
|
||||
|
||||
Cannot set version_counter for inference tensor
|
||||
|
||||
For now, we work around this issue by using `torch.no_grad()` in this case.
|
||||
See https://github.com/pytorch/pytorch/issues/164872 for more details.
|
||||
Otherwise, just return `torch.inference_mode()`.
|
||||
"""
|
||||
torchao_config = getattr(model, "torchao_config", None)
|
||||
if torchao_config is not None and torchao_config.qat_scheme is None:
|
||||
return torch.no_grad()
|
||||
else:
|
||||
return torch.inference_mode()
|
||||
|
|
|
|||
|
|
@ -21,7 +21,10 @@ from ._utils import *
|
|||
from ._utils import patch_unsloth_smart_gradient_checkpointing
|
||||
from ._utils import __version__, importlib_version
|
||||
from ._utils import move_to_device
|
||||
from ._utils import _prepare_model_for_qat
|
||||
from ._utils import (
|
||||
_get_inference_mode_context_manager,
|
||||
_prepare_model_for_qat,
|
||||
)
|
||||
from torch.nn.functional import scaled_dot_product_attention
|
||||
from transformers import __version__ as transformers_version
|
||||
from unsloth_zoo.utils import Version, _get_dtype
|
||||
|
|
@ -2030,7 +2033,7 @@ def unsloth_fast_generate(
|
|||
|
||||
# Mixed precision autocast
|
||||
with (
|
||||
torch.inference_mode(),
|
||||
_get_inference_mode_context_manager(self),
|
||||
torch.autocast(device_type = DEVICE_TYPE_TORCH, dtype = dtype),
|
||||
):
|
||||
output = self._old_generate(*args, **kwargs)
|
||||
|
|
|
|||
|
|
@ -31,7 +31,12 @@ from .cohere import FastCohereModel
|
|||
from transformers import AutoConfig
|
||||
from transformers import __version__ as transformers_version
|
||||
from peft import PeftConfig, PeftModel
|
||||
from .loader_utils import get_model_name
|
||||
from .loader_utils import (
|
||||
_check_load_in_fp8_settings,
|
||||
_offline_quantize_to_fp8,
|
||||
_tag_model_with_fp8_torchao_config,
|
||||
get_model_name,
|
||||
)
|
||||
import os, contextlib, sys
|
||||
|
||||
try:
|
||||
|
|
@ -140,6 +145,7 @@ class FastLanguageModel(FastLlamaModel):
|
|||
max_lora_rank = 64,
|
||||
disable_log_stats = True,
|
||||
qat_scheme = None,
|
||||
load_in_fp8 = False, # fp8 LoRA
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
|
|
@ -183,6 +189,7 @@ class FastLanguageModel(FastLlamaModel):
|
|||
max_lora_rank = max_lora_rank,
|
||||
disable_log_stats = disable_log_stats,
|
||||
qat_scheme = qat_scheme,
|
||||
load_in_fp8 = load_in_fp8,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
|
|
@ -212,9 +219,23 @@ class FastLanguageModel(FastLlamaModel):
|
|||
)
|
||||
load_in_4bit = False
|
||||
|
||||
if load_in_fp8:
|
||||
_check_load_in_fp8_settings(
|
||||
fast_inference,
|
||||
full_finetuning,
|
||||
load_in_4bit,
|
||||
load_in_8bit,
|
||||
load_in_16bit,
|
||||
use_exact_model_name,
|
||||
)
|
||||
|
||||
old_model_name = model_name
|
||||
if not use_exact_model_name:
|
||||
model_name = get_model_name(model_name, load_in_4bit)
|
||||
if load_in_fp8:
|
||||
model_name = _offline_quantize_to_fp8(model_name)
|
||||
else:
|
||||
model_name = get_model_name(model_name, load_in_4bit)
|
||||
|
||||
# Check if pre-quantized models are allowed
|
||||
# For eg AMD GPUs need blocksize = 128, but our pre-quants are blocksize = 64
|
||||
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
|
||||
|
|
@ -476,6 +497,8 @@ class FastLanguageModel(FastLlamaModel):
|
|||
random_state = random_state,
|
||||
max_lora_rank = max_lora_rank,
|
||||
disable_log_stats = disable_log_stats,
|
||||
qat_scheme = qat_scheme,
|
||||
load_in_fp8 = load_in_fp8,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
|
|
@ -554,6 +577,9 @@ class FastLanguageModel(FastLlamaModel):
|
|||
}
|
||||
model.config.update({"quantization_config": quantization_config})
|
||||
|
||||
if load_in_fp8:
|
||||
_tag_model_with_fp8_torchao_config(model)
|
||||
|
||||
if is_peft:
|
||||
# From https://github.com/huggingface/peft/issues/184
|
||||
# Now add PEFT adapters
|
||||
|
|
@ -634,6 +660,7 @@ class FastModel(FastBaseModel):
|
|||
max_lora_rank = 64,
|
||||
disable_log_stats = True,
|
||||
qat_scheme = None,
|
||||
load_in_fp8 = False, # fp8 LoRA
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
|
|
@ -694,9 +721,23 @@ class FastModel(FastBaseModel):
|
|||
)
|
||||
load_in_4bit = False
|
||||
|
||||
if load_in_fp8:
|
||||
_check_load_in_fp8_settings(
|
||||
fast_inference,
|
||||
full_finetuning,
|
||||
load_in_4bit,
|
||||
load_in_8bit,
|
||||
load_in_16bit,
|
||||
use_exact_model_name,
|
||||
)
|
||||
|
||||
old_model_name = model_name
|
||||
if not use_exact_model_name:
|
||||
model_name = get_model_name(model_name, load_in_4bit)
|
||||
if load_in_fp8:
|
||||
model_name = _offline_quantize_to_fp8(model_name)
|
||||
else:
|
||||
model_name = get_model_name(model_name, load_in_4bit)
|
||||
|
||||
# Check if pre-quantized models are allowed
|
||||
# For eg AMD GPUs need blocksize = 128, but our pre-quants are blocksize = 64
|
||||
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
|
||||
|
|
@ -1130,6 +1171,9 @@ class FastModel(FastBaseModel):
|
|||
}
|
||||
model.config.update({"quantization_config": quantization_config})
|
||||
|
||||
if load_in_fp8:
|
||||
_tag_model_with_fp8_torchao_config(model)
|
||||
|
||||
if is_peft:
|
||||
# From https://github.com/huggingface/peft/issues/184
|
||||
# Now add PEFT adapters
|
||||
|
|
|
|||
|
|
@ -12,11 +12,23 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import importlib
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
from .mapper import INT_TO_FLOAT_MAPPER, FLOAT_TO_INT_MAPPER, MAP_TO_UNSLOTH_16bit
|
||||
|
||||
# https://github.com/huggingface/transformers/pull/26037 allows 4 bit loading!
|
||||
from packaging.version import Version
|
||||
from transformers import __version__ as transformers_version
|
||||
from transformers import (
|
||||
AutoModelForCausalLM,
|
||||
AutoTokenizer,
|
||||
TorchAoConfig,
|
||||
__version__ as transformers_version,
|
||||
)
|
||||
from unsloth.models._utils import TorchAOConfig
|
||||
from unsloth_zoo.utils import Version
|
||||
import torch
|
||||
|
||||
transformers_version = Version(transformers_version)
|
||||
SUPPORTS_FOURBIT = transformers_version >= Version("4.37")
|
||||
|
|
@ -144,3 +156,128 @@ def get_model_name(model_name, load_in_4bit = True):
|
|||
'pip install --upgrade --no-cache-dir "git+https://github.com/unslothai/unsloth-zoo.git"\n'
|
||||
)
|
||||
return new_model_name if new_model_name is not None else model_name
|
||||
|
||||
|
||||
def _get_torchao_fp8_config():
|
||||
"""
|
||||
Return a `torchao.quantization.Float8DynamicActivationFloat8WeightConfig`
|
||||
to be used for `load_in_fp8=True`.
|
||||
"""
|
||||
from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, PerRow
|
||||
|
||||
return Float8DynamicActivationFloat8WeightConfig(
|
||||
granularity = PerRow(),
|
||||
activation_value_lb = 1e-12,
|
||||
)
|
||||
|
||||
|
||||
def _offline_quantize_to_fp8(model_name: str) -> str:
|
||||
"""
|
||||
Quantizes the model to fp8 using torchao and saving the quantized model to a
|
||||
temporary location. Return the path to the quantized model.
|
||||
|
||||
Note: Once on-the-fly quantization is added in vllm in
|
||||
https://github.com/vllm-project/vllm/pull/26327, we should
|
||||
dynamically quantize the model there instead:
|
||||
|
||||
llm = LLM(
|
||||
...
|
||||
hf_overrides={"quantization_config_file": "torchao_config.json"},
|
||||
)
|
||||
"""
|
||||
temp_dir = tempfile.gettempdir()
|
||||
new_model_name = model_name.split("/")[-1] + "-fp8"
|
||||
new_model_name = os.path.join(temp_dir, new_model_name)
|
||||
print(
|
||||
f"Quantizing '{model_name}' to fp8, using model_name='{new_model_name}' instead"
|
||||
)
|
||||
if not os.path.isdir(new_model_name):
|
||||
qconfig = _get_torchao_fp8_config()
|
||||
qconfig = TorchAoConfig(qconfig)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
torch_dtype = "auto",
|
||||
device_map = "auto",
|
||||
quantization_config = qconfig,
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
model.save_pretrained(new_model_name, safe_serialization = False)
|
||||
tokenizer.save_pretrained(new_model_name)
|
||||
return new_model_name
|
||||
|
||||
|
||||
def _tag_model_with_fp8_torchao_config(model: torch.nn.Module):
|
||||
"""
|
||||
Tag a model with a `TorchAOConfig` so downstream callers will know what to do with it.
|
||||
"""
|
||||
base_config = _get_torchao_fp8_config()
|
||||
model.torchao_config = TorchAOConfig(
|
||||
qat_scheme = None,
|
||||
base_config_and_filter_fns = [(base_config, None)],
|
||||
)
|
||||
|
||||
|
||||
def _check_load_in_fp8_settings(
|
||||
fast_inference: bool,
|
||||
full_finetuning: bool,
|
||||
load_in_4bit: bool,
|
||||
load_in_8bit: bool,
|
||||
load_in_16bit: bool,
|
||||
use_exact_model_name: bool,
|
||||
):
|
||||
"""
|
||||
Assuming `load_in_fp8=True`, raise appropriate errors on incompatible settings
|
||||
and environment. Currently this feature requires:
|
||||
1. H100 GPUs or after
|
||||
2. torchao 0.15.0+ (or nightly)
|
||||
3. torch 2.9.0+
|
||||
4. If fbgemm_gpu_genai is installed, require 1.4.1+
|
||||
"""
|
||||
if not fast_inference:
|
||||
raise ValueError(
|
||||
"Unsloth: `load_in_fp8` is only supported for `fast_inference` for now"
|
||||
)
|
||||
if full_finetuning:
|
||||
raise ValueError(
|
||||
"Unsloth: `load_in_fp8` is not compatible with full finetuning"
|
||||
)
|
||||
if load_in_4bit or load_in_8bit or load_in_16bit:
|
||||
raise ValueError(
|
||||
"Unsloth: `load_in_fp8` is not compatible with `load_in_4bit`, `load_in_8bit` or `load_in_16bit`",
|
||||
)
|
||||
if use_exact_model_name:
|
||||
raise ValueError("Unsloth: `load_in_fp8` requires `use_exact_model_name=False`")
|
||||
|
||||
# Check if this is Hopper or above
|
||||
if not (
|
||||
torch.cuda.is_available()
|
||||
and torch.version.cuda
|
||||
and torch.cuda.get_device_capability() >= (9, 0)
|
||||
):
|
||||
raise ValueError("Unsloth: `load_in_fp8` requires H100 GPUs or after")
|
||||
|
||||
# Check if torch >= 2.9.0
|
||||
if Version(torch.__version__) < Version("2.9.0"):
|
||||
raise ValueError("Unsloth: `load_in_fp8` requires torch 2.9.0+")
|
||||
|
||||
# Check if torchao has this PR: https://github.com/pytorch/ao/pull/3158,
|
||||
# which will be released in 0.15.0.
|
||||
error_message = "Unsloth: `load_in_fp8` requires torchao 0.15.0+ (or nightly)"
|
||||
if importlib.util.find_spec("torchao") is None:
|
||||
raise ValueError(error_message)
|
||||
import torchao
|
||||
|
||||
if Version(torchao.__version__) < Version("0.15.0"):
|
||||
raise ValueError(error_message)
|
||||
|
||||
# If fbgemm_gpu_genai is installed, check if it's >= 1.4.1
|
||||
if (
|
||||
importlib.util.find_spec("fbgemm_gpu") is not None
|
||||
and importlib.util.find_spec("fbgemm_gpu.experimental") is not None
|
||||
):
|
||||
import fbgemm_gpu.experimental.gen_ai
|
||||
|
||||
if Version(fbgemm_gpu.__version__) < Version("1.4.1"):
|
||||
raise ValueError(
|
||||
"Unsloth: `load_in_fp8` is only compatible with fbgemm_gpu_genai 1.4.1+"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ from ..device_type import (
|
|||
ALLOW_PREQUANTIZED_MODELS,
|
||||
)
|
||||
import textwrap
|
||||
from ._utils import _get_inference_mode_context_manager
|
||||
|
||||
RL_EXTRA_ARGS = defaultdict(list)
|
||||
RL_FUNCTIONS = defaultdict(list)
|
||||
|
|
@ -536,7 +537,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
|
|||
)
|
||||
|
||||
with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype):
|
||||
with torch.inference_mode():
|
||||
with _get_inference_mode_context_manager(model):
|
||||
if pixel_values is None:
|
||||
attention_mask = input_ids != self.processing_class.pad_token_id
|
||||
attention_mask = attention_mask.to(attention_mask.dtype)
|
||||
|
|
@ -603,6 +604,9 @@ RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(UnslothEfficientGRPO))
|
|||
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_accumulated_loss))
|
||||
RL_PRE_ITEMS["grpo_trainer"].append(grpo_compute_loss_slow)
|
||||
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_update_SamplingParams))
|
||||
RL_PRE_ITEMS["grpo_trainer"].append(
|
||||
inspect.getsource(_get_inference_mode_context_manager)
|
||||
)
|
||||
|
||||
|
||||
# Edit _get_per_token_logps to handle mixed precision
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue