Merge pull request #2942 from ChingTsai:enable-grad-clipping-in-sft

PiperOrigin-RevId: 857296866

Author: Google-ML-Automation

src/MaxText/sft/sft_trainer.py

@@ -40,6 +40,7 @@
 from absl import app
 import os
 import jax
+import optax
 import pathwaysutils
 
 from flax.linen import partitioning as nn_partitioning
@@ -147,6 +148,12 @@ def setup_trainer_state(mt_config, goodput_recorder=None):
     # pass in model for muon
     optimizer = optimizers.get_optimizer(mt_config, learning_rate_schedule, model)
 
+    if mt_config.gradient_clipping_threshold > 0:
+      optimizer = optax.chain(
+          optax.clip_by_global_norm(max_norm=mt_config.gradient_clipping_threshold),
+          optimizer,
+      )
+
   with maybe_record_goodput(goodput_recorder, GoodputEvent.TRAINING_PREPARATION):
     training_hooks = hooks.SFTTrainingHooks(mt_config, mesh, learning_rate_schedule, goodput_recorder)
     data_hooks = hooks.SFTDataHooks(mt_config, mesh, goodput_recorder)