with grad_accum

run_mlm_flax.py

@@ -288,7 +288,8 @@ def rotate_checkpoints(ckpt_dir: str, save_total_limit: int):
 
 
     
-
+class TrainState(train_state.TrainState):
+    grad_accum: jnp.ndarray
 
 
     
@@ -396,10 +397,10 @@ if __name__ == "__main__":
             return train, val
         train, val = train_val_files()
         datasets = load_dataset('json', data_files={'train': train, 'validation': val})
-        #datasets["train"] = datasets["train"].select(range(int(0.8*len(datasets["train"]))))
-        #datasets["validation"] = datasets["validation"].select(range(int(0.8*len(datasets["validation"]))))
-        datasets["train"] = datasets["train"].select(range(10000))
-        datasets["validation"] = datasets["validation"].select(range(10000))
+        datasets["train"] = datasets["train"].select(range(int(0.8*len(datasets["train"]))))
+        datasets["validation"] = datasets["validation"].select(range(int(0.8*len(datasets["validation"]))))
+        #datasets["train"] = datasets["train"].select(range(10000))
+        #datasets["validation"] = datasets["validation"].select(range(10000))
 
 
 
@@ -566,7 +567,7 @@ if __name__ == "__main__":
 
     # Store some constant
     num_epochs = int(training_args.num_train_epochs)
-    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count() #* training_args.gradient_accumulation_steps
+    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
     eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
 
     num_train_steps = len(train_dataset) // train_batch_size * num_epochs
@@ -613,14 +614,14 @@ if __name__ == "__main__":
             mask=decay_mask_fn,
         )
 
-    if training_args.gradient_accumulation_steps > 1:
-        optimizer = optax.MultiSteps(optimizer, training_args.gradient_accumulation_steps)
-    grad_accum_steps = training_args.gradient_accumulation_steps
+    #if training_args.gradient_accumulation_steps > 1:
+    #    optimizer = optax.MultiSteps(optimizer, training_args.gradient_accumulation_steps)
+    #grad_accum_steps = training_args.gradient_accumulation_steps
 
     # Setup train state
     
     
-    state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer)
+    state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer,grad_accum=jax.tree_map(jnp.zeros_like, model.params))
         
     if training_args.resume_from_checkpoint:
         state = restore_checkpoint(training_args.resume_from_checkpoint, state)
@@ -645,17 +646,30 @@ if __name__ == "__main__":
             # take average
             loss = loss.sum() / label_mask.sum()
 
-            return loss
+            return loss / training_args.gradient_accumulation_steps
 
         grad_fn = jax.value_and_grad(loss_fn)
-        loss, grad = grad_fn(state.params)
-        grad = jax.lax.pmean(grad, "batch")
-        new_state = state.apply_gradients(grads=grad)
-
+        loss, grads = grad_fn(state.params)
+        grad_accum = jax.tree_multimap(lambda x, y: x + y, grads, state.grad_accum)
+        
+        def update_fn():
+            grads = jax.tree_map(lambda x: x / training_args.gradient_accumulation_steps, grad_accum)
+            grads = jax.lax.pmean(grad_accum, "batch")
+            new_state = state.apply_gradients(grads=grads,grad_accum=jax.tree_map(jnp.zeros_like, grads))
+            return new_state
+
+        new_state = jax.lax.cond(
+            state.step % training_args.gradient_accumulation_steps == 0,
+            lambda _: update_fn(),
+            lambda _: state.replace(grad_accum=grad_accum, step=state.step + 1),
+            None,
+        )
+        
         metrics = jax.lax.pmean(
-            {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step // grad_accum_steps)}, axis_name="batch" #
+            {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}, axis_name="batch" #
         )
 
+        #return new_state.replace(new_dropout_rng=new_dropout_rng), metrics
         return new_state, metrics, new_dropout_rng
 
     # Create parallel version of the train step
@@ -699,10 +713,10 @@ if __name__ == "__main__":
         # Generate an epoch by shuffling sampling indices from the train dataset
         num_train_samples = len(train_dataset)
         train_samples_idx = jax.random.permutation(input_rng, jnp.arange(num_train_samples))
-        train_batch_idx = generate_batch_splits(train_samples_idx, train_batch_size // grad_accum_steps) #
+        train_batch_idx = generate_batch_splits(train_samples_idx, train_batch_size) #// grad_accum_steps
 
         # Gather the indexes for creating the batch and do a training step
-        for step, batch_idx in enumerate(tqdm(train_batch_idx, desc="Training...", position=1,initial=resume_step // grad_accum_steps)): #
+        for step, batch_idx in enumerate(tqdm(train_batch_idx, desc="Training...", position=1,initial=resume_step)): #grad_accum
             samples = [train_dataset[int(idx)] for idx in batch_idx]
             model_inputs = data_collator(samples, pad_to_multiple_of=16)
             
@@ -716,7 +730,7 @@ if __name__ == "__main__":
             if cur_step < resume_step:
                 continue
 
-            if (cur_step % training_args.logging_steps * grad_accum_steps) == 0 and cur_step > 0: #
+            if (cur_step % training_args.logging_steps) == 0 and cur_step > 0: # * grad_accum_steps
                 # Save metrics
                 train_metric = jax_utils.unreplicate(train_metric)
                 train_time += time.time() - train_start
@@ -733,7 +747,7 @@ if __name__ == "__main__":
 
                 train_metrics = []
 
-            if cur_step % (training_args.eval_steps * grad_accum_steps) == 0 and cur_step > 0: #
+            if cur_step % (training_args.eval_steps) == 0 and cur_step > 0: #* grad_accum_steps
                 # ======================== Evaluating ==============================
                 num_eval_samples = len(eval_dataset)
                 eval_samples_idx = jnp.arange(num_eval_samples)