feat(train): use MultiSteps for gradient accumulation

tools/train/train.py

@@ -647,9 +647,7 @@ def main():
 
     # add gradient accumulation
     if training_args.gradient_accumulation_steps > 1:
-        optimizer = optax.chain(
-            optax.apply_every(training_args.gradient_accumulation_steps), optimizer
-        )
+        optimizer = optax.MultiSteps(optimizer, training_args.gradient_accumulation_steps)
 
     # Setup train state
     state = TrainState.create(
@@ -693,7 +691,7 @@ def main():
 
         metrics = {
             "loss": loss,
-            "learning_rate": learning_rate_fn(state.step),
+            "learning_rate": learning_rate_fn(state.step // training_args.gradient_accumulation_steps),
         }
         metrics = jax.lax.pmean(metrics, axis_name="batch")