From ff3375ee8e06d7954de284ddc915704f49cdff21 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 12 Feb 2025 03:50:07 -0800 Subject: [PATCH] autocast --- unsloth/models/rl.py | 1 + unsloth/models/rl_replacements.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index cf351ebf3f..466101d16c 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -80,6 +80,7 @@ from typing import * from dataclasses import dataclass, field from packaging.version import Version import torch +from contextlib import nullcontext @dataclass class Unsloth{RLConfig_name}({RLConfig_name}): diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 65138feb13..2ea12f69c4 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -116,3 +116,22 @@ def sft_trainer_compute_loss(function_name, function): return function pass RL_FUNCTIONS["sft_trainer"].append(sft_trainer_compute_loss) + + +# Autocast precision for GRPO +def grpo_trainer__prepare_inputs(function_name, function): + if function_name != "_prepare_inputs": return function + + if "with torch.inference_mode()" not in function: return 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()", + ) + return function +pass +RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__prepare_inputs)