Compare commits
3 commits
main
...
dh/combine
| Author | SHA1 | Date | |
|---|---|---|---|
|
de28b8e7db |
||
|
|
a68590e6c2 | ||
|
|
d031935bbd |
3 changed files with 97 additions and 43 deletions
|
|
@ -1331,8 +1331,18 @@ class FastModel(FastBaseModel):
|
|||
redirector = contextlib.redirect_stdout(open(os.devnull, "w"))
|
||||
|
||||
model_types = ["siglip"] + model_types
|
||||
# Set forced float32 env flag
|
||||
os.environ["UNSLOTH_FORCE_FLOAT32"] = "0"
|
||||
# Preserve explicit user force-float32 requests without leaking prior
|
||||
# auto-detected force state into later model loads in the same process.
|
||||
user_forced_float32 = (
|
||||
type(float32_mixed_precision) is bool and float32_mixed_precision
|
||||
)
|
||||
if not user_forced_float32:
|
||||
user_forced_float32 = (
|
||||
os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1"
|
||||
and os.environ.get("UNSLOTH_AUTO_FORCE_FLOAT32", "0") != "1"
|
||||
)
|
||||
os.environ["UNSLOTH_FORCE_FLOAT32"] = "1" if user_forced_float32 else "0"
|
||||
os.environ["UNSLOTH_AUTO_FORCE_FLOAT32"] = "0"
|
||||
do_forced_float32 = False
|
||||
for model_type_arch in model_types:
|
||||
if model_type_arch != "siglip":
|
||||
|
|
@ -1346,6 +1356,7 @@ class FastModel(FastBaseModel):
|
|||
or disable_name.lower() in model_types_all
|
||||
) and ((dtype == torch.float16) or not SUPPORTS_BFLOAT16):
|
||||
os.environ["UNSLOTH_FORCE_FLOAT32"] = "1"
|
||||
os.environ["UNSLOTH_AUTO_FORCE_FLOAT32"] = "1"
|
||||
dtype = torch.bfloat16 # Change to bfloat16 loading
|
||||
break
|
||||
# Apply gradient checkpointing with smart heuristics
|
||||
|
|
|
|||
|
|
@ -715,39 +715,71 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
"if type(use_bf16) is not bool: use_bf16 = False\n"
|
||||
"use_fp16 = getattr(args, 'fp16', False)\n"
|
||||
"if type(use_fp16) is not bool: use_fp16 = False\n"
|
||||
"force_float32 = False\n"
|
||||
"full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1'\n"
|
||||
"if not full_finetuning and (os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'):\n"
|
||||
"force_float32 = os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'\n"
|
||||
"if force_float32:\n"
|
||||
" print('Unsloth: Switching to float32 training since model cannot work with float16')\n"
|
||||
" force_float32 = True\n"
|
||||
"mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32')\n"
|
||||
"dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None)\n"
|
||||
"if dtype is None: dtype = model.get_input_embeddings().weight.dtype\n"
|
||||
"from unsloth_zoo.utils import _get_dtype\n"
|
||||
"dtype = _get_dtype(dtype)\n"
|
||||
"float16 = dtype == torch.float16\n"
|
||||
"if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`')\n"
|
||||
"if not force_float32 and (not float16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')\n"
|
||||
"original_dtype = getattr(model, '_unsloth_original_dtype', None)\n"
|
||||
"if original_dtype is None and hasattr(model, 'model'): original_dtype = getattr(model.model, '_unsloth_original_dtype', None)\n"
|
||||
"if original_dtype is None: original_dtype = dtype\n"
|
||||
"else: original_dtype = _get_dtype(original_dtype)\n"
|
||||
"is_float16 = original_dtype == torch.float16\n"
|
||||
"is_bfloat16 = original_dtype == torch.bfloat16\n"
|
||||
"is_float32 = original_dtype == torch.float32\n"
|
||||
"explicit_mixed_precision = None\n"
|
||||
"if use_bf16: explicit_mixed_precision = 'bf16'\n"
|
||||
"elif use_fp16: explicit_mixed_precision = 'fp16'\n"
|
||||
"if not force_float32 and (is_float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`')\n"
|
||||
"if not force_float32 and (is_bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')\n"
|
||||
"if not force_float32 and (dtype == torch.float32) and (explicit_mixed_precision == 'fp16') and not is_float16:\n"
|
||||
" print('Unsloth: Model is float32 but fp16 was requested. Switching to float32 training.')\n"
|
||||
" force_float32 = True\n"
|
||||
"if not force_float32 and (dtype == torch.float32) and (explicit_mixed_precision == 'bf16') and not is_bfloat16:\n"
|
||||
" print('Unsloth: Model is float32 but bf16 was requested. Switching to float32 training.')\n"
|
||||
" force_float32 = True\n"
|
||||
"if force_float32:\n"
|
||||
" # Forced float32 training\n"
|
||||
" os.environ['UNSLOTH_FORCE_FLOAT32'] = '1'\n"
|
||||
" args.fp16 = False\n"
|
||||
" args.bf16 = False\n"
|
||||
" os.environ['ACCELERATE_MIXED_PRECISION'] = 'no'\n"
|
||||
" if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no'\n"
|
||||
" # args.mixed_precision is a new argument which needs to be set now\n"
|
||||
"elif explicit_mixed_precision is not None:\n"
|
||||
" args.fp16 = explicit_mixed_precision == 'fp16'\n"
|
||||
" args.bf16 = explicit_mixed_precision == 'bf16'\n"
|
||||
" os.environ['ACCELERATE_MIXED_PRECISION'] = explicit_mixed_precision\n"
|
||||
" if hasattr(args, 'mixed_precision'): args.mixed_precision = explicit_mixed_precision\n"
|
||||
" # args.mixed_precision is a new argument which needs to be set now\n"
|
||||
"elif mixed_precision_dtype == 'bfloat16':\n"
|
||||
" args.fp16 = False\n"
|
||||
" args.bf16 = True\n"
|
||||
" os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16'\n"
|
||||
" if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16'\n"
|
||||
" # args.mixed_precision is a new argument which needs to be set now\n"
|
||||
"elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32':\n"
|
||||
" # Mixed precision training\n"
|
||||
" args.fp16 = float16\n"
|
||||
" args.bf16 = not float16\n"
|
||||
" os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16'\n"
|
||||
" if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16'\n"
|
||||
" # args.mixed_precision is a new argument which needs to be set now\n"
|
||||
"elif mixed_precision_dtype == 'bfloat16':\n"
|
||||
" # Both False since bfloat16 full finetuning doesn't do any autocasting.\n"
|
||||
" args.fp16 = False\n"
|
||||
" args.bf16 = False\n"
|
||||
" os.environ['ACCELERATE_MIXED_PRECISION'] = 'no'\n"
|
||||
" if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no'\n"
|
||||
" if is_float16:\n"
|
||||
" args.fp16 = True\n"
|
||||
" args.bf16 = False\n"
|
||||
" os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16'\n"
|
||||
" if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16'\n"
|
||||
" elif is_bfloat16:\n"
|
||||
" args.fp16 = False\n"
|
||||
" args.bf16 = True\n"
|
||||
" os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16'\n"
|
||||
" if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16'\n"
|
||||
" else:\n"
|
||||
" # Float32 params should not force bf16/fp16 mixed precision.\n"
|
||||
" args.fp16 = False\n"
|
||||
" args.bf16 = False\n"
|
||||
" os.environ['ACCELERATE_MIXED_PRECISION'] = 'no'\n"
|
||||
" if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no'\n"
|
||||
" # args.mixed_precision is a new argument which needs to be set now\n"
|
||||
"\n"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ import re
|
|||
import torch
|
||||
import inspect
|
||||
import linecache
|
||||
from contextlib import nullcontext
|
||||
from collections import defaultdict
|
||||
from unsloth_zoo.rl_replacements import (
|
||||
RL_REPLACEMENTS,
|
||||
|
|
@ -244,10 +245,10 @@ def grpo_trainer__prepare_inputs(function_name, function):
|
|||
function = function.replace(
|
||||
"with torch.inference_mode():",
|
||||
"with torch.inference_mode(), "
|
||||
"torch.amp.autocast(device_type = 'cuda', "
|
||||
"dtype = ((torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16) "
|
||||
"if not torch.is_autocast_enabled('cuda') else nullcontext())"
|
||||
"if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '0' else torch.float16):",
|
||||
"(nullcontext() if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' else "
|
||||
"(torch.amp.autocast(device_type = 'cuda', "
|
||||
"dtype = (torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16)) "
|
||||
"if not torch.is_autocast_enabled('cuda') else nullcontext())):",
|
||||
)
|
||||
function = function.replace(
|
||||
"self.accelerator.unwrap_model(self.model)",
|
||||
|
|
@ -624,17 +625,21 @@ def grpo_trainer__get_per_token_logps(function_name, function):
|
|||
if True: # os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0':
|
||||
return None # Unsloth efficient GRPO
|
||||
# Otherwise, calculate normally:
|
||||
if not hasattr(self, "_autocast_dtype"):
|
||||
self._autocast_dtype = (
|
||||
torch.float16
|
||||
if os.environ.get("ACCELERATE_MIXED_PRECISION", "fp16") == "fp16"
|
||||
else torch.bfloat16
|
||||
if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
|
||||
autocaster = nullcontext()
|
||||
else:
|
||||
if not hasattr(self, "_autocast_dtype"):
|
||||
self._autocast_dtype = (
|
||||
torch.float16
|
||||
if os.environ.get("ACCELERATE_MIXED_PRECISION", "fp16") == "fp16"
|
||||
else torch.bfloat16
|
||||
)
|
||||
autocaster = torch.amp.autocast(
|
||||
device_type = DEVICE_TYPE, dtype = self._autocast_dtype
|
||||
)
|
||||
if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
|
||||
self._autocast_dtype = torch.float16
|
||||
|
||||
os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
|
||||
with torch.amp.autocast(device_type = DEVICE_TYPE, dtype = self._autocast_dtype):
|
||||
with autocaster:
|
||||
# We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded
|
||||
logits = model(
|
||||
input_ids = input_ids,
|
||||
|
|
@ -690,14 +695,25 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
|
|||
if compute_efficient:
|
||||
return None, None
|
||||
else:
|
||||
if not hasattr(self, "_autocast_dtype"):
|
||||
self._autocast_dtype = (
|
||||
torch.float16
|
||||
if os.environ.get("ACCELERATE_MIXED_PRECISION", "fp16") == "fp16"
|
||||
else torch.bfloat16
|
||||
if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
|
||||
autocaster = nullcontext()
|
||||
dtype_bytes = 32
|
||||
else:
|
||||
if not hasattr(self, "_autocast_dtype"):
|
||||
self._autocast_dtype = (
|
||||
torch.float16
|
||||
if os.environ.get("ACCELERATE_MIXED_PRECISION", "fp16")
|
||||
== "fp16"
|
||||
else torch.bfloat16
|
||||
)
|
||||
autocaster = torch.amp.autocast(
|
||||
device_type = DEVICE_TYPE, dtype = self._autocast_dtype
|
||||
)
|
||||
dtype_bytes = (
|
||||
16
|
||||
if self._autocast_dtype in [torch.float16, torch.bfloat16]
|
||||
else 32
|
||||
)
|
||||
if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
|
||||
self._autocast_dtype = torch.float16
|
||||
|
||||
pixel_values, image_grid_thw = (
|
||||
kwargs.get("pixel_values", None),
|
||||
|
|
@ -714,9 +730,6 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
|
|||
|
||||
lm_head = self.model.get_output_embeddings().weight
|
||||
|
||||
dtype_bytes = (
|
||||
16 if self._autocast_dtype in [torch.float16, torch.bfloat16] else 32
|
||||
)
|
||||
total_rows = input_ids.shape[0]
|
||||
seq_len = input_ids.shape[1]
|
||||
hidden_dim = lm_head.shape[1]
|
||||
|
|
@ -843,9 +856,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
|
|||
pixel_attention_mask_chunk,
|
||||
image_sizes_chunk,
|
||||
) in zipped_inputs:
|
||||
with torch.amp.autocast(
|
||||
device_type = "cuda", dtype = self._autocast_dtype
|
||||
):
|
||||
with autocaster:
|
||||
if pixel_values is None:
|
||||
logits_chunk = unwrapped_model(
|
||||
input_ids = input_ids_chunk,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue