Merge pull request #2437 from AI-Hypercomputer:qinwen/add_grad_bf16

PiperOrigin-RevId: 815880783

Author: Google-ML-Automation

src/MaxText/configs/base.yml

@@ -83,6 +83,9 @@ gcs_metrics: False
 # If true save config to GCS in {base_output_directory}/{run_name}/
 save_config_to_gcs: False
 
+# Gradient dtype
+grad_dtype: "float32"
+
 # Activation dtypes.
 dtype: "bfloat16"
 # Used to configure quantization in the transformer layers, defaults to null implying bf16.

src/MaxText/pyconfig.py

@@ -691,6 +691,7 @@ def user_init(raw_keys):
 
     # Type conversions
     raw_keys["dtype"] = jax.numpy.dtype(raw_keys["dtype"])
+    raw_keys["grad_dtype"] = jax.numpy.dtype(raw_keys["grad_dtype"])
     raw_keys["weight_dtype"] = jax.numpy.dtype(raw_keys["weight_dtype"])
     raw_keys["mu_dtype"] = set_mu_dtype(raw_keys)
     raw_keys["logical_axis_rules"] = _lists_to_tuples(raw_keys["logical_axis_rules"])

src/MaxText/train.py

@@ -255,6 +255,8 @@ def reshape_to_microbatch_accumulations(batch_arr):
         extra_dpo_args = [reference_params]
     grad_func = jax.value_and_grad(_loss_fn, argnums=4, has_aux=True)
     (loss, aux), raw_grads = grad_func(model, config, data, dropout_rng, state.params, *extra_dpo_args, is_train=True)
+
+  raw_grads = jax.tree_util.tree_map(lambda x: x.astype(config.grad_dtype) if x.dtype == jnp.float32 else x, raw_grads)
   intermediate_outputs = aux["intermediate_outputs"]
   total_weights = aux["total_weights"]
   moe_lb_loss = aux["moe_lb_loss"]