Update run_mlm_flax.py

run_mlm_flax.py

@@ -760,7 +760,7 @@ def main():
             logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
 
             # compute loss, ignore padded input tokens
-            label_mask = jnp.where(labels > -100, 1.0, 0.0)
+            label_mask = jnp.where(labels!=-100, 1.0, 0.0)
             loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])) * label_mask
 
             # take average
@@ -816,6 +816,11 @@ def main():
     train_time = 0
     epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
     for epoch in epochs:
+        # Debug
+        testsample = [tokenized_datasets["train"][0]]  # Just one sample
+        testbatch = data_collator(testsample, pad_to_multiple_of=16)
+        print("Labels in a sample batch:", testbatch["labels"])
+
         # ======================== Training ================================
         train_start = time.time()
         train_metrics = []
@@ -837,6 +842,9 @@ def main():
             samples = [tokenized_datasets["train"][int(idx)] for idx in batch_idx]
             model_inputs = data_collator(samples, pad_to_multiple_of=16)
 
+
+
+            
             # Model forward
             model_inputs = shard(model_inputs.data)
             state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)
@@ -845,6 +853,17 @@ def main():
             cur_step = epoch * (num_train_samples // (train_batch_size * jax.process_count())) + step
 
             if cur_step % training_args.logging_steps == 0 and cur_step > 0:
+
+
+                # Debug
+                if cur_step % training_args.logging_steps == 0 and cur_step > 0:
+                    params = jax.device_get(jax.tree_util.tree_map(lambda x: x[0], state.params))
+                    total_norm = 0
+                    for p in jax.tree_util.tree_leaves(params):
+                        total_norm += np.sum(np.square(p))
+                    total_norm = np.sqrt(total_norm)
+                    print("Parameter norm:", total_norm)
+     
                 # Save metrics
                 train_metric = jax_utils.unreplicate(train_metric)
                 train_time += time.time() - train_start