From fe73b86c614e030669f74651b6aaaa35ccd3c937 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 18 Mar 2025 02:23:03 -0700 Subject: [PATCH] Update rl_replacements.py --- unsloth/models/rl_replacements.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index b638dc6ccf..fe1a534e39 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -176,8 +176,9 @@ def grpo_trainer__prepare_inputs(function_name, function): "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():", + "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):", ) # Disable attaching a float32 conversion hook which upcasts logits to FP32