Add configurable beta parameters for adam

MaxText/configs/base.yml

@@ -109,7 +109,8 @@ prompt: "I love to "
 enable_profiler: False
 enable_checkpointing: True
 
+# Adam optimizer parameters
 adam_b1: 0.9 # Exponential decay rate to track the first moment of past gradients.
 adam_b2: 0.95 # Exponential decay rate to track the second moment of past gradients.
 adam_eps: 1.e-8 # A small constant applied to denominator outside of the square root.
-adam_eps_root: 0. # A small constant applied to denominator inside the square root
+adam_eps_root: 0. # A small constant applied to denominator inside the square root.