save intermediate

check_gradients_pt_flax.py

@@ -1,11 +1,11 @@
 #!/usr/bin/env python3
-from transformers import SpeechEncoderDecoderModel, FlaxSpeechEncoderDecoderModel
 import tempfile
 import random
 import numpy as np
 import torch
 import optax
 import jax
+import sys
 from flax.training.common_utils import onehot
 from flax.traverse_util import flatten_dict
 
@@ -63,7 +63,7 @@ def shift_tokens_right(input_ids: np.array, pad_token_id: int, decoder_start_tok
     return shifted_input_ids
 
 
-def assert_almost_equals(a: np.ndarray, b: np.ndarray, tol: float = 4e-2):
+def assert_almost_equals(a: np.ndarray, b: np.ndarray, tol: float = 1e-2):
     diff = np.abs((a - b)).max()
     if diff < tol:
         print(f"✅ Difference between Flax and PyTorch is {diff} (< {tol})")
@@ -71,46 +71,60 @@ def assert_almost_equals(a: np.ndarray, b: np.ndarray, tol: float = 4e-2):
         print(f"❌ Difference between Flax and PyTorch is {diff} (>= {tol})")
 
 
-def assert_dict_equal(a: dict, b: dict, tol: float = 4e-2):
+def assert_dict_equal(a: dict, b: dict, tol: float = 1e-2):
     if a.keys() != b.keys():
         print("❌ Dictionary keys for PyTorch and Flax do not match")
+    results_fail = []
+    results_correct = []
+
+    results_fail_rel = []
+    results_correct_rel = []
     for k in a:
-        diff = np.abs((a[k] - b[k])).max()
+        ak_norm = np.linalg.norm(a[k])
+        bk_norm = np.linalg.norm(b[k])
+        diff = np.abs(ak_norm - bk_norm)
+        diff_rel = np.abs(ak_norm - bk_norm) / np.abs(ak_norm)
         if diff < tol:
-            print(f"✅ Layer {k} diff is {diff} < {tol}).")
+            results_correct.append(f"✅ Layer {k} diff is {diff} < {tol}).")
+        else:
+            results_fail.append(f"❌ Layer {k} has PT grad norm {bk_norm} and flax grad norm {ak_norm}.")
+        if diff_rel < tol:
+            results_correct_rel.append(f"✅ Layer {k} rel diff is {diff} < {tol}).")
         else:
-            print(f"❌ Layer {k} diff is {diff} (>= {tol}).")
+            results_fail_rel.append(f"❌ Layer {k} has PT grad norm {bk_norm} and flax grad norm {ak_norm}.")
+    return results_fail_rel, results_correct_rel, results_fail, results_correct
 
 
-def main():
-    encoder_id = "hf-internal-testing/tiny-random-wav2vec2"
-    decoder_id = "hf-internal-testing/tiny-random-bart"
+def compare_grads(model_id, pt_architecture):
+    transformers_module = __import__("transformers", fromlist=[pt_architecture])
 
-    use_decoder_attention_mask = False
-    freeze_feature_encoder = False
+    model_cls = getattr(transformers_module, pt_architecture)
+    flax_model_cls = getattr(transformers_module, "Flax" + pt_architecture)
 
-    pt_model = SpeechEncoderDecoderModel.from_encoder_decoder_pretrained(encoder_id, decoder_id,
-                                                                         encoder_add_adapter=True)
+    pt_model, model_info = model_cls.from_pretrained(model_id, output_loading_info=True)
 
-    with tempfile.TemporaryDirectory() as tmpdirname:
-        pt_model.save_pretrained(tmpdirname)
-        fx_model = FlaxSpeechEncoderDecoderModel.from_pretrained(tmpdirname, from_pt=True)
+    if len(model_info["missing_keys"]) > 0:
+        raise ValueError(f"{model_id} with {pt_architecture} has missing keys: {model_info['missing_keys']}")
 
-    batch_size = 13
-    input_values = floats_tensor([batch_size, 512], fx_model.config.encoder.vocab_size)
-    attention_mask = random_attention_mask([batch_size, 512])
-    label_ids = ids_tensor([batch_size, 4], fx_model.config.decoder.vocab_size)
-    decoder_input_ids = shift_tokens_right(input_ids=label_ids, pad_token_id=fx_model.config.decoder.pad_token_id,
-                                           decoder_start_token_id=fx_model.config.decoder.decoder_start_token_id)
-    decoder_attention_mask = random_attention_mask([batch_size, 4])
+    fx_model = flax_model_cls.from_pretrained(model_id, from_pt=True)
+
+    batch_size = 2
+    seq_len = 64
+
+    input_ids = ids_tensor([batch_size, seq_len], fx_model.config.vocab_size)
+    label_ids = ids_tensor([batch_size, seq_len], fx_model.config.vocab_size)
+
+    attention_mask = random_attention_mask([batch_size, seq_len])
+    label_ids = ids_tensor([batch_size, seq_len], fx_model.config.vocab_size)
 
     fx_inputs = {
-        "inputs": input_values,
+        "input_ids": input_ids,
         "attention_mask": attention_mask,
-        "decoder_input_ids": decoder_input_ids,
     }
-    if use_decoder_attention_mask:
-        fx_inputs["decoder_attention_mask"] = decoder_attention_mask
+
+    if pt_model.config.is_encoder_decoder:
+        decoder_input_ids = shift_tokens_right(input_ids=label_ids, pad_token_id=fx_model.config.pad_token_id, decoder_start_token_id=fx_model.config.decoder_start_token_id)
+        fx_inputs["decoder_input_ids"] = decoder_input_ids
 
     pt_inputs = {k: torch.tensor(v.tolist()) for k, v in fx_inputs.items()}
     pt_inputs["labels"] = torch.tensor(label_ids.tolist())
@@ -118,9 +132,6 @@ def main():
     fx_outputs = fx_model(**fx_inputs)
     fx_logits = fx_outputs.logits
 
-    if freeze_feature_encoder:
-        pt_model.freeze_feature_encoder()
-
     pt_outputs = pt_model(**pt_inputs)
     pt_logits = pt_outputs.logits
     pt_loss = pt_outputs.loss
@@ -129,11 +140,10 @@ def main():
     print(f"Flax logits shape: {fx_logits.shape}, PyTorch logits shape: {pt_logits.shape}")
     assert_almost_equals(fx_logits, pt_logits.detach().numpy())
 
-    def fx_train_step(fx_model, batch, freeze_feature_encoder=False):
+    def fx_train_step(fx_model, batch):
         def compute_loss(params):
             label_ids = batch.pop('label_ids')
-            logits = fx_model(**batch, params=params,
-                              freeze_feature_encoder=freeze_feature_encoder).logits
+            logits = fx_model(**batch, params=params).logits
             vocab_size = logits.shape[-1]
             targets = onehot(label_ids, vocab_size)
             loss = optax.softmax_cross_entropy(logits, targets)
@@ -145,7 +155,7 @@ def main():
 
     fx_inputs["label_ids"] = label_ids
 
-    fx_loss, fx_grad = fx_train_step(fx_model, fx_inputs, freeze_feature_encoder=freeze_feature_encoder)
+    fx_loss, fx_grad = fx_train_step(fx_model, fx_inputs)
 
     print("--------------------------Checking losses match--------------------------")
     print(f"Flax loss: {fx_loss}, PyTorch loss: {pt_loss}")
@@ -166,13 +176,31 @@ def main():
 
     with tempfile.TemporaryDirectory() as tmpdirname:
         pt_model.save_pretrained(tmpdirname)
-        pt_grad_model_to_fx = FlaxSpeechEncoderDecoderModel.from_pretrained(tmpdirname, from_pt=True)
+        pt_grad_model_to_fx = flax_model_cls.from_pretrained(tmpdirname, from_pt=True)
 
     pt_grad_to_fx = pt_grad_model_to_fx.params
     fx_grad = flatten_dict(fx_grad)
     pt_grad_to_fx = flatten_dict(pt_grad_to_fx)
     print("--------------------------Checking gradients match--------------------------")
-    assert_dict_equal(fx_grad, pt_grad_to_fx)
+    results_fail_rel, results_correct_rel, results_fail, results_correct = assert_dict_equal(fx_grad, pt_grad_to_fx)
+
+    if len(results_fail) == 0:
+        print("✅ All grads pass")
+    else:
+        print("\n".join(results_fail))
+
+    print("--------------------------Checking rel gradients match--------------------------")
+
+    if len(results_fail_rel) == 0:
+        print("✅ All rel grads pass")
+    else:
+        print("\n".join(results_fail_rel))
+
+
+def main():
+    model_id = sys.argv[1]
+    pt_architecture_name = sys.argv[2]
+    compare_grads(model_id, pt_architecture_name)
 
 
 if __name__ == "__main__":

run_models.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+#model_ids=("hf-internal-testing/tiny-random-roberta" "hf-internal-testing/tiny-random-bert" "hf-internal-testing/tiny-random-bart" "tf-internal-testing/tiny-random-t5")
+#model_architectures=("RobertaForMaskedLM" "BertForMaskedLM" "BartForConditionalGeneration" "T5ForConditionalGeneration")
+model_ids=("hf-internal-testing/tiny-random-roberta")
+model_architectures=("RobertaForMaskedLM")
+
+rm -rf log.txt
+touch log.txt
+
+for model_idx in "${!model_ids[@]}"; do
+	model_id=${model_ids[model_idx]}
+	model_architecture=${model_architectures[model_idx]}
+
+	echo "Check ${model_id} ..." >> log.txt
+	./check_gradients_pt_flax.py "${model_id}" "${model_architecture}" >> log.txt
+	echo "=========================================" >> log.txt
+done