we should not autocast

ServiceNow · Mar 3, 2025 · edf0a26 · edf0a26
1 parent 7333ddd
commit edf0a26
Showing 1 changed file with 9 additions and 10 deletions.
diff --git a/tapeagents/finetune/finetune.py b/tapeagents/finetune/finetune.py
@@ -203,16 +203,15 @@ def toggle_sync(sync: bool):
                 torch.cuda.empty_cache()
 
             do_optimizer_step = training_metrics.passes % args.gradient_accumulation_passes == 0
-            with torch.autocast("cuda"):
-                with toggle_sync(do_optimizer_step):
-                    loss, this_step_rl_metrics = forward(model, batch)
-                    for k, v in this_step_rl_metrics.items():
-                        rl_metrics[k].append(v)
-                    training_metrics.train_loss = loss.item()
-                    training_metrics.lr = optimizer.param_groups[0]["lr"]
-                    training_metrics.max_batch_len = max(batch["input_ids"].shape[1], training_metrics.max_batch_len)
-                    training_metrics.min_batch_len = min(batch["input_ids"].shape[1], training_metrics.min_batch_len)
-                    accelerator.backward(loss / args.gradient_accumulation_passes)
+            with toggle_sync(do_optimizer_step):
+                loss, this_step_rl_metrics = forward(model, batch)
+                for k, v in this_step_rl_metrics.items():
+                    rl_metrics[k].append(v)
+                training_metrics.train_loss = loss.item()
+                training_metrics.lr = optimizer.param_groups[0]["lr"]
+                training_metrics.max_batch_len = max(batch["input_ids"].shape[1], training_metrics.max_batch_len)
+                training_metrics.min_batch_len = min(batch["input_ids"].shape[1], training_metrics.min_batch_len)
+                accelerator.backward(loss / args.gradient_accumulation_passes)
 
             if not do_optimizer_step:
                 continue