finish

after_fix_log.txt

@@ -0,0 +1,62 @@
+Check hf-internal-testing/tiny-random-roberta ...
+--------------------------Checking logits match--------------------------
+Flax logits shape: (2, 64, 1000), PyTorch logits shape: torch.Size([2, 64, 1000])
+✅ Difference between Flax and PyTorch is 1.7881393432617188e-07 (< 0.01)
+--------------------------Checking losses match--------------------------
+Flax loss: 6.887884140014648, PyTorch loss: 6.887884616851807
+✅ Difference between Flax and PyTorch is 4.76837158203125e-07 (< 0.01)
+--------------------------Checking gradients match--------------------------
+✅ All grads pass
+--------------------------Checking rel gradients match--------------------------
+❌ Layer ('roberta', 'encoder', 'layer', '0', 'attention', 'self', 'key', 'bias') has PT grad norm 7.584575871001642e-13 and flax grad norm 6.388195094436666e-13.
+❌ Layer ('roberta', 'encoder', 'layer', '1', 'attention', 'self', 'key', 'bias') has PT grad norm 7.811836030477415e-13 and flax grad norm 6.42668156105447e-13.
+❌ Layer ('roberta', 'encoder', 'layer', '2', 'attention', 'self', 'key', 'bias') has PT grad norm 8.422985074175993e-13 and flax grad norm 6.414080963405844e-13.
+❌ Layer ('roberta', 'encoder', 'layer', '3', 'attention', 'self', 'key', 'bias') has PT grad norm 8.625919531608794e-13 and flax grad norm 7.699825477734679e-13.
+❌ Layer ('roberta', 'encoder', 'layer', '4', 'attention', 'self', 'key', 'bias') has PT grad norm 1.0383360837806777e-12 and flax grad norm 6.049140680551568e-13.
+=========================================
+Check hf-internal-testing/tiny-random-bert ...
+--------------------------Checking logits match--------------------------
+Flax logits shape: (2, 64, 1124), PyTorch logits shape: torch.Size([2, 64, 1124])
+✅ Difference between Flax and PyTorch is 1.7881393432617188e-07 (< 0.01)
+--------------------------Checking losses match--------------------------
+Flax loss: 7.036032199859619, PyTorch loss: 7.036032676696777
+✅ Difference between Flax and PyTorch is 4.76837158203125e-07 (< 0.01)
+--------------------------Checking gradients match--------------------------
+✅ All grads pass
+--------------------------Checking rel gradients match--------------------------
+❌ Layer ('bert', 'encoder', 'layer', '0', 'attention', 'self', 'key', 'bias') has PT grad norm 5.234438642080785e-13 and flax grad norm 4.935363641205004e-13.
+❌ Layer ('bert', 'encoder', 'layer', '1', 'attention', 'self', 'key', 'bias') has PT grad norm 9.028551018787356e-13 and flax grad norm 6.16219206737989e-13.
+❌ Layer ('bert', 'encoder', 'layer', '2', 'attention', 'self', 'key', 'bias') has PT grad norm 8.728350616056535e-13 and flax grad norm 6.037235598596591e-13.
+❌ Layer ('bert', 'encoder', 'layer', '3', 'attention', 'self', 'key', 'bias') has PT grad norm 8.327751465850297e-13 and flax grad norm 7.390156737431541e-13.
+❌ Layer ('bert', 'encoder', 'layer', '4', 'attention', 'self', 'key', 'bias') has PT grad norm 7.404479048130075e-13 and flax grad norm 7.178592030705755e-13.
+=========================================
+Check hf-internal-testing/tiny-random-t5 ...
+--------------------------Checking logits match--------------------------
+Flax logits shape: (2, 64, 1103), PyTorch logits shape: torch.Size([2, 64, 1103])
+✅ Difference between Flax and PyTorch is 3.725290298461914e-09 (< 0.01)
+--------------------------Checking losses match--------------------------
+Flax loss: 7.006012916564941, PyTorch loss: 7.006012916564941
+✅ Difference between Flax and PyTorch is 0.0 (< 0.01)
+--------------------------Checking gradients match--------------------------
+✅ All grads pass
+--------------------------Checking rel gradients match--------------------------
+✅ All rel grads pass
+=========================================
+Check hf-internal-testing/tiny-random-bart ...
+--------------------------Checking logits match--------------------------
+Flax logits shape: (2, 64, 1000), PyTorch logits shape: torch.Size([2, 64, 1000])
+✅ Difference between Flax and PyTorch is 8.940696716308594e-08 (< 0.01)
+--------------------------Checking losses match--------------------------
+Flax loss: 6.919522285461426, PyTorch loss: 6.919522285461426
+✅ Difference between Flax and PyTorch is 0.0 (< 0.01)
+--------------------------Checking gradients match--------------------------
+✅ All grads pass
+--------------------------Checking rel gradients match--------------------------
+❌ Layer ('final_logits_bias',) has PT grad norm 0.0 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '0', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 1.1293364247239035e-13 and flax grad norm 7.444291358479557e-14.
+❌ Layer ('model', 'decoder', 'layers', '0', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.9028742882613858e-13 and flax grad norm 1.0847509820726894e-13.
+❌ Layer ('model', 'decoder', 'layers', '1', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 1.0747876384459981e-13 and flax grad norm 1.1924105637346055e-13.
+❌ Layer ('model', 'decoder', 'layers', '1', 'self_attn', 'k_proj', 'bias') has PT grad norm 2.0553104032074165e-13 and flax grad norm 2.416926793251395e-13.
+❌ Layer ('model', 'encoder', 'layers', '0', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.6505221806215232e-16 and flax grad norm 8.704786207109524e-17.
+❌ Layer ('model', 'encoder', 'layers', '1', 'self_attn', 'k_proj', 'bias') has PT grad norm 2.145616114665214e-16 and flax grad norm 1.6750639014770325e-16.
+=========================================

after_fix_pretrained_log.txt

@@ -0,0 +1,212 @@
+Check roberta-base ...
+--------------------------Checking logits match--------------------------
+Flax logits shape: (2, 64, 50265), PyTorch logits shape: torch.Size([2, 64, 50265])
+✅ Difference between Flax and PyTorch is 0.00013017654418945312 (< 0.01)
+--------------------------Checking losses match--------------------------
+Flax loss: 14.801228523254395, PyTorch loss: 14.801219940185547
+✅ Difference between Flax and PyTorch is 8.58306884765625e-06 (< 0.01)
+--------------------------Checking gradients match--------------------------
+✅ All grads pass
+--------------------------Checking rel gradients match--------------------------
+❌ Layer ('roberta', 'encoder', 'layer', '0', 'attention', 'self', 'key', 'bias') has PT grad norm 6.889232651019483e-08 and flax grad norm 5.7956174970286156e-08.
+❌ Layer ('roberta', 'encoder', 'layer', '1', 'attention', 'self', 'key', 'bias') has PT grad norm 7.026115156349988e-08 and flax grad norm 6.282134989987753e-08.
+❌ Layer ('roberta', 'encoder', 'layer', '10', 'attention', 'self', 'key', 'bias') has PT grad norm 2.273949206710313e-08 and flax grad norm 1.8748883334751554e-08.
+❌ Layer ('roberta', 'encoder', 'layer', '11', 'attention', 'self', 'key', 'bias') has PT grad norm 2.9379741306456708e-08 and flax grad norm 2.6026933497291793e-08.
+❌ Layer ('roberta', 'encoder', 'layer', '2', 'attention', 'self', 'key', 'bias') has PT grad norm 6.197853963385569e-08 and flax grad norm 5.317058082709991e-08.
+❌ Layer ('roberta', 'encoder', 'layer', '3', 'attention', 'self', 'key', 'bias') has PT grad norm 7.359258802352997e-08 and flax grad norm 8.573702814373974e-08.
+❌ Layer ('roberta', 'encoder', 'layer', '4', 'attention', 'self', 'key', 'bias') has PT grad norm 5.1634213349416314e-08 and flax grad norm 5.744939457485998e-08.
+❌ Layer ('roberta', 'encoder', 'layer', '5', 'attention', 'self', 'key', 'bias') has PT grad norm 4.652720519970899e-08 and flax grad norm 6.121346984855336e-08.
+❌ Layer ('roberta', 'encoder', 'layer', '6', 'attention', 'self', 'key', 'bias') has PT grad norm 3.8810604507943935e-08 and flax grad norm 4.2490388096894094e-08.
+❌ Layer ('roberta', 'encoder', 'layer', '7', 'attention', 'self', 'key', 'bias') has PT grad norm 3.7450202938771326e-08 and flax grad norm 3.219445687818734e-08.
+❌ Layer ('roberta', 'encoder', 'layer', '8', 'attention', 'self', 'key', 'bias') has PT grad norm 3.3088259243641005e-08 and flax grad norm 2.6118801343955056e-08.
+❌ Layer ('roberta', 'encoder', 'layer', '9', 'attention', 'self', 'key', 'bias') has PT grad norm 2.6417508180998084e-08 and flax grad norm 2.415968225477627e-08.
+=========================================
+Check bert-base-cased ...
+--------------------------Checking logits match--------------------------
+Flax logits shape: (2, 64, 28996), PyTorch logits shape: torch.Size([2, 64, 28996])
+✅ Difference between Flax and PyTorch is 5.4836273193359375e-05 (< 0.01)
+--------------------------Checking losses match--------------------------
+Flax loss: 13.967159271240234, PyTorch loss: 13.967162132263184
+✅ Difference between Flax and PyTorch is 2.86102294921875e-06 (< 0.01)
+--------------------------Checking gradients match--------------------------
+✅ All grads pass
+--------------------------Checking rel gradients match--------------------------
+❌ Layer ('bert', 'encoder', 'layer', '0', 'attention', 'self', 'key', 'bias') has PT grad norm 8.025740783068613e-08 and flax grad norm 8.381563532111613e-08.
+❌ Layer ('bert', 'encoder', 'layer', '1', 'attention', 'self', 'key', 'bias') has PT grad norm 7.262840284738559e-08 and flax grad norm 5.0372555904232286e-08.
+❌ Layer ('bert', 'encoder', 'layer', '10', 'attention', 'self', 'key', 'bias') has PT grad norm 2.6523425233904163e-08 and flax grad norm 2.7082945663892133e-08.
+❌ Layer ('bert', 'encoder', 'layer', '11', 'attention', 'self', 'key', 'bias') has PT grad norm 2.9038789151059063e-08 and flax grad norm 3.3138192634396546e-08.
+❌ Layer ('bert', 'encoder', 'layer', '2', 'attention', 'self', 'key', 'bias') has PT grad norm 5.880680831182872e-08 and flax grad norm 5.04786008548308e-08.
+❌ Layer ('bert', 'encoder', 'layer', '3', 'attention', 'self', 'key', 'bias') has PT grad norm 4.705585965325554e-08 and flax grad norm 4.983893475696277e-08.
+❌ Layer ('bert', 'encoder', 'layer', '4', 'attention', 'self', 'key', 'bias') has PT grad norm 6.595875134962625e-08 and flax grad norm 5.823812543326312e-08.
+❌ Layer ('bert', 'encoder', 'layer', '5', 'attention', 'self', 'key', 'bias') has PT grad norm 4.716540402682767e-08 and flax grad norm 6.053270595884896e-08.
+❌ Layer ('bert', 'encoder', 'layer', '6', 'attention', 'self', 'key', 'bias') has PT grad norm 5.4432636176215965e-08 and flax grad norm 4.0700697923057305e-08.
+❌ Layer ('bert', 'encoder', 'layer', '8', 'attention', 'self', 'key', 'bias') has PT grad norm 4.059621971919114e-08 and flax grad norm 4.575255374561493e-08.
+❌ Layer ('bert', 'encoder', 'layer', '9', 'attention', 'self', 'key', 'bias') has PT grad norm 2.9032529269557017e-08 and flax grad norm 2.659336217902819e-08.
+=========================================
+Check t5-small ...
+--------------------------Checking logits match--------------------------
+Flax logits shape: (2, 64, 32128), PyTorch logits shape: torch.Size([2, 64, 32128])
+✅ Difference between Flax and PyTorch is 7.62939453125e-05 (< 0.01)
+--------------------------Checking losses match--------------------------
+Flax loss: 20.534835815429688, PyTorch loss: 20.534835815429688
+✅ Difference between Flax and PyTorch is 0.0 (< 0.01)
+--------------------------Checking gradients match--------------------------
+✅ All grads pass
+--------------------------Checking rel gradients match--------------------------
+✅ All rel grads pass
+=========================================
+Check facebook/bart-large ...
+--------------------------Checking logits match--------------------------
+Flax logits shape: (2, 64, 50265), PyTorch logits shape: torch.Size([2, 64, 50265])
+✅ Difference between Flax and PyTorch is 0.0004191398620605469 (< 0.01)
+--------------------------Checking losses match--------------------------
+Flax loss: 13.993148803710938, PyTorch loss: 13.993138313293457
+✅ Difference between Flax and PyTorch is 1.049041748046875e-05 (< 0.01)
+--------------------------Checking gradients match--------------------------
+❌ Layer ('model', 'decoder', 'layers', '0', 'fc1', 'kernel') has PT grad norm 11.655710220336914 and flax grad norm 11.6015625.
+❌ Layer ('model', 'decoder', 'layers', '0', 'fc2', 'kernel') has PT grad norm 7.740886211395264 and flax grad norm 7.71484375.
+❌ Layer ('model', 'decoder', 'layers', '1', 'fc1', 'kernel') has PT grad norm 6.798306465148926 and flax grad norm 6.765625.
+❌ Layer ('model', 'decoder', 'layers', '1', 'fc2', 'kernel') has PT grad norm 7.071859836578369 and flax grad norm 7.05078125.
+❌ Layer ('model', 'decoder', 'layers', '10', 'fc1', 'kernel') has PT grad norm 16.904926300048828 and flax grad norm 16.859375.
+❌ Layer ('model', 'decoder', 'layers', '10', 'fc2', 'kernel') has PT grad norm 6.783661842346191 and flax grad norm 6.765625.
+❌ Layer ('model', 'decoder', 'layers', '10', 'self_attn', 'v_proj', 'kernel') has PT grad norm 6.97633171081543 and flax grad norm 6.96484375.
+❌ Layer ('model', 'decoder', 'layers', '11', 'fc1', 'kernel') has PT grad norm 13.733073234558105 and flax grad norm 13.6875.
+❌ Layer ('model', 'decoder', 'layers', '11', 'fc2', 'kernel') has PT grad norm 6.311193466186523 and flax grad norm 6.29296875.
+❌ Layer ('model', 'decoder', 'layers', '2', 'fc1', 'kernel') has PT grad norm 6.043461322784424 and flax grad norm 6.01171875.
+❌ Layer ('model', 'decoder', 'layers', '2', 'fc2', 'kernel') has PT grad norm 8.091109275817871 and flax grad norm 8.0703125.
+❌ Layer ('model', 'decoder', 'layers', '3', 'fc1', 'kernel') has PT grad norm 6.561250686645508 and flax grad norm 6.52734375.
+❌ Layer ('model', 'decoder', 'layers', '3', 'fc2', 'kernel') has PT grad norm 8.535536766052246 and flax grad norm 8.5.
+❌ Layer ('model', 'decoder', 'layers', '4', 'fc1', 'kernel') has PT grad norm 5.882389545440674 and flax grad norm 5.859375.
+❌ Layer ('model', 'decoder', 'layers', '4', 'fc2', 'kernel') has PT grad norm 8.772762298583984 and flax grad norm 8.75.
+❌ Layer ('model', 'decoder', 'layers', '5', 'fc1', 'kernel') has PT grad norm 4.559173107147217 and flax grad norm 4.5390625.
+❌ Layer ('model', 'decoder', 'layers', '5', 'fc2', 'kernel') has PT grad norm 7.053295612335205 and flax grad norm 7.03515625.
+❌ Layer ('model', 'decoder', 'layers', '6', 'fc1', 'kernel') has PT grad norm 4.724750995635986 and flax grad norm 4.703125.
+❌ Layer ('model', 'decoder', 'layers', '6', 'fc2', 'kernel') has PT grad norm 6.6051740646362305 and flax grad norm 6.5859375.
+❌ Layer ('model', 'decoder', 'layers', '7', 'fc1', 'kernel') has PT grad norm 3.9028773307800293 and flax grad norm 3.884765625.
+❌ Layer ('model', 'decoder', 'layers', '7', 'fc2', 'kernel') has PT grad norm 6.16121244430542 and flax grad norm 6.14453125.
+❌ Layer ('model', 'decoder', 'layers', '8', 'fc1', 'kernel') has PT grad norm 3.8737242221832275 and flax grad norm 3.85546875.
+❌ Layer ('model', 'decoder', 'layers', '8', 'fc2', 'kernel') has PT grad norm 6.476221084594727 and flax grad norm 6.45703125.
+❌ Layer ('model', 'decoder', 'layers', '9', 'fc1', 'kernel') has PT grad norm 6.240624904632568 and flax grad norm 6.203125.
+❌ Layer ('model', 'decoder', 'layers', '9', 'fc2', 'kernel') has PT grad norm 6.872060775756836 and flax grad norm 6.8515625.
+❌ Layer ('model', 'encoder', 'layers', '0', 'fc1', 'kernel') has PT grad norm 18.140913009643555 and flax grad norm 18.0625.
+❌ Layer ('model', 'encoder', 'layers', '0', 'fc2', 'kernel') has PT grad norm 13.278300285339355 and flax grad norm 13.2265625.
+❌ Layer ('model', 'encoder', 'layers', '0', 'self_attn', 'v_proj', 'kernel') has PT grad norm 13.971784591674805 and flax grad norm 13.9609375.
+❌ Layer ('model', 'encoder', 'layers', '1', 'fc1', 'kernel') has PT grad norm 9.453910827636719 and flax grad norm 9.4140625.
+❌ Layer ('model', 'encoder', 'layers', '1', 'fc2', 'kernel') has PT grad norm 5.271415710449219 and flax grad norm 5.25.
+❌ Layer ('model', 'encoder', 'layers', '10', 'fc1', 'kernel') has PT grad norm 13.269391059875488 and flax grad norm 13.2109375.
+❌ Layer ('model', 'encoder', 'layers', '10', 'fc2', 'kernel') has PT grad norm 15.780173301696777 and flax grad norm 15.734375.
+❌ Layer ('model', 'encoder', 'layers', '10', 'self_attn', 'v_proj', 'kernel') has PT grad norm 21.07574462890625 and flax grad norm 21.0625.
+❌ Layer ('model', 'encoder', 'layers', '11', 'fc1', 'kernel') has PT grad norm 16.847095489501953 and flax grad norm 16.765625.
+❌ Layer ('model', 'encoder', 'layers', '11', 'fc2', 'kernel') has PT grad norm 17.480010986328125 and flax grad norm 17.421875.
+❌ Layer ('model', 'encoder', 'layers', '11', 'self_attn', 'out_proj', 'kernel') has PT grad norm 26.91538429260254 and flax grad norm 26.890625.
+❌ Layer ('model', 'encoder', 'layers', '11', 'self_attn', 'v_proj', 'kernel') has PT grad norm 26.706096649169922 and flax grad norm 26.6875.
+❌ Layer ('model', 'encoder', 'layers', '2', 'fc1', 'kernel') has PT grad norm 6.756587505340576 and flax grad norm 6.7265625.
+❌ Layer ('model', 'encoder', 'layers', '2', 'fc2', 'kernel') has PT grad norm 5.000077724456787 and flax grad norm 4.98046875.
+❌ Layer ('model', 'encoder', 'layers', '2', 'self_attn', 'out_proj', 'kernel') has PT grad norm 18.72007942199707 and flax grad norm 18.703125.
+❌ Layer ('model', 'encoder', 'layers', '3', 'fc1', 'kernel') has PT grad norm 9.560458183288574 and flax grad norm 9.515625.
+❌ Layer ('model', 'encoder', 'layers', '3', 'fc2', 'kernel') has PT grad norm 8.892805099487305 and flax grad norm 8.859375.
+❌ Layer ('model', 'encoder', 'layers', '4', 'fc1', 'kernel') has PT grad norm 8.845908164978027 and flax grad norm 8.8046875.
+❌ Layer ('model', 'encoder', 'layers', '4', 'fc2', 'kernel') has PT grad norm 8.670329093933105 and flax grad norm 8.640625.
+❌ Layer ('model', 'encoder', 'layers', '4', 'self_attn', 'out_proj', 'kernel') has PT grad norm 18.135663986206055 and flax grad norm 18.125.
+❌ Layer ('model', 'encoder', 'layers', '5', 'fc1', 'kernel') has PT grad norm 10.071086883544922 and flax grad norm 10.03125.
+❌ Layer ('model', 'encoder', 'layers', '5', 'fc2', 'kernel') has PT grad norm 9.528592109680176 and flax grad norm 9.5.
+❌ Layer ('model', 'encoder', 'layers', '5', 'self_attn', 'out_proj', 'kernel') has PT grad norm 16.202266693115234 and flax grad norm 16.1875.
+❌ Layer ('model', 'encoder', 'layers', '5', 'self_attn', 'v_proj', 'kernel') has PT grad norm 17.79180908203125 and flax grad norm 17.78125.
+❌ Layer ('model', 'encoder', 'layers', '6', 'fc1', 'kernel') has PT grad norm 12.527167320251465 and flax grad norm 12.46875.
+❌ Layer ('model', 'encoder', 'layers', '6', 'fc2', 'kernel') has PT grad norm 11.495430946350098 and flax grad norm 11.453125.
+❌ Layer ('model', 'encoder', 'layers', '6', 'self_attn', 'v_proj', 'kernel') has PT grad norm 18.312782287597656 and flax grad norm 18.296875.
+❌ Layer ('model', 'encoder', 'layers', '7', 'fc1', 'kernel') has PT grad norm 11.963201522827148 and flax grad norm 11.9140625.
+❌ Layer ('model', 'encoder', 'layers', '7', 'fc2', 'kernel') has PT grad norm 13.052857398986816 and flax grad norm 13.0078125.
+❌ Layer ('model', 'encoder', 'layers', '7', 'self_attn', 'out_proj', 'kernel') has PT grad norm 17.2364501953125 and flax grad norm 17.21875.
+❌ Layer ('model', 'encoder', 'layers', '7', 'self_attn', 'v_proj', 'kernel') has PT grad norm 18.88938331604004 and flax grad norm 18.875.
+❌ Layer ('model', 'encoder', 'layers', '8', 'fc1', 'kernel') has PT grad norm 11.773221969604492 and flax grad norm 11.71875.
+❌ Layer ('model', 'encoder', 'layers', '8', 'fc2', 'kernel') has PT grad norm 14.441213607788086 and flax grad norm 14.3984375.
+❌ Layer ('model', 'encoder', 'layers', '8', 'self_attn', 'out_proj', 'kernel') has PT grad norm 19.045316696166992 and flax grad norm 19.03125.
+❌ Layer ('model', 'encoder', 'layers', '8', 'self_attn', 'v_proj', 'kernel') has PT grad norm 18.37466049194336 and flax grad norm 18.359375.
+❌ Layer ('model', 'encoder', 'layers', '9', 'fc1', 'kernel') has PT grad norm 12.223063468933105 and flax grad norm 12.1640625.
+❌ Layer ('model', 'encoder', 'layers', '9', 'fc2', 'kernel') has PT grad norm 15.896522521972656 and flax grad norm 15.8515625.
+--------------------------Checking rel gradients match--------------------------
+❌ Layer ('final_logits_bias',) has PT grad norm 0.0 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '0', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 8.274865592738934e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '0', 'self_attn', 'k_proj', 'bias') has PT grad norm 3.530680814378684e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '1', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 9.156278935051887e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '1', 'self_attn', 'k_proj', 'bias') has PT grad norm 2.7762926180230352e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '10', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 2.314587632668008e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '10', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.0350275303494527e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '11', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 2.4003222520718737e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '11', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.0850777165671843e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '2', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 8.753153935003866e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '2', 'self_attn', 'k_proj', 'bias') has PT grad norm 3.340263532436438e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '3', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 7.864576190286243e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '3', 'self_attn', 'k_proj', 'bias') has PT grad norm 3.161582284860742e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '4', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 3.1698593971896116e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '4', 'self_attn', 'k_proj', 'bias') has PT grad norm 3.397210690536667e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '5', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 4.269594100492213e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '5', 'self_attn', 'k_proj', 'bias') has PT grad norm 2.3111518032692402e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '6', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 3.222339373110117e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '6', 'self_attn', 'k_proj', 'bias') has PT grad norm 2.2391466458770992e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '7', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 2.4886496419185278e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '7', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.1538945798861278e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '8', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 2.3322632713984603e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '8', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.525707205018989e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '9', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 2.3752901867624132e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '9', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.4804857784156411e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'encoder', 'layers', '0', 'self_attn', 'k_proj', 'bias') has PT grad norm 8.215602775862862e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'encoder', 'layers', '1', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.3884370275718538e-07 and flax grad norm 0.0.
+❌ Layer ('model', 'encoder', 'layers', '10', 'self_attn', 'k_proj', 'bias') has PT grad norm 6.668184937552724e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'encoder', 'layers', '11', 'self_attn', 'k_proj', 'bias') has PT grad norm 8.486110658623147e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'encoder', 'layers', '2', 'self_attn', 'k_proj', 'bias') has PT grad norm 7.912614563565512e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'encoder', 'layers', '3', 'self_attn', 'k_proj', 'bias') has PT grad norm 7.959246062227976e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'encoder', 'layers', '4', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.1111747255654336e-07 and flax grad norm 0.0.
+❌ Layer ('model', 'encoder', 'layers', '5', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.0121711824240265e-07 and flax grad norm 0.0.
+❌ Layer ('model', 'encoder', 'layers', '6', 'self_attn', 'k_proj', 'bias') has PT grad norm 7.735735607639072e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'encoder', 'layers', '7', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.0352330548357713e-07 and flax grad norm 0.0.
+❌ Layer ('model', 'encoder', 'layers', '8', 'self_attn', 'k_proj', 'bias') has PT grad norm 8.3155640595578e-08 and flax grad norm 0.0.
+❌ Layer ('model', 'encoder', 'layers', '9', 'self_attn', 'k_proj', 'bias') has PT grad norm 6.08824493042448e-08 and flax grad norm 0.0.
+=========================================
+Check facebook/bart-large-cnn ...
+--------------------------Checking logits match--------------------------
+Flax logits shape: (2, 64, 50264), PyTorch logits shape: torch.Size([2, 64, 50264])
+✅ Difference between Flax and PyTorch is 0.0003502368927001953 (< 0.01)
+--------------------------Checking losses match--------------------------
+Flax loss: 13.418181419372559, PyTorch loss: 13.418176651000977
+✅ Difference between Flax and PyTorch is 4.76837158203125e-06 (< 0.01)
+--------------------------Checking gradients match--------------------------
+✅ All grads pass
+--------------------------Checking rel gradients match--------------------------
+❌ Layer ('final_logits_bias',) has PT grad norm 0.0 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '0', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 3.5387660091146245e-07 and flax grad norm 4.874667069998395e-07.
+❌ Layer ('model', 'decoder', 'layers', '0', 'self_attn', 'k_proj', 'bias') has PT grad norm 6.254911966152576e-08 and flax grad norm 6.927437112835833e-08.
+❌ Layer ('model', 'decoder', 'layers', '1', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 2.0864062832970376e-07 and flax grad norm 1.8239754240312323e-07.
+❌ Layer ('model', 'decoder', 'layers', '1', 'self_attn', 'k_proj', 'bias') has PT grad norm 6.265895535761956e-08 and flax grad norm 6.79184637419894e-08.
+❌ Layer ('model', 'decoder', 'layers', '10', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 2.3348372124587513e-08 and flax grad norm 1.9192864186834413e-08.
+❌ Layer ('model', 'decoder', 'layers', '10', 'self_attn', 'k_proj', 'bias') has PT grad norm 2.108701835368265e-08 and flax grad norm 1.938536442480654e-08.
+❌ Layer ('model', 'decoder', 'layers', '11', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 2.4383044916476138e-08 and flax grad norm 1.890670375814807e-08.
+❌ Layer ('model', 'decoder', 'layers', '11', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.4334801790027996e-08 and flax grad norm 1.3059753278810149e-08.
+❌ Layer ('model', 'decoder', 'layers', '2', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 2.0973679681901558e-07 and flax grad norm 2.6336286396144715e-07.
+❌ Layer ('model', 'decoder', 'layers', '2', 'self_attn', 'k_proj', 'bias') has PT grad norm 8.528008521579977e-08 and flax grad norm 1.286241939624233e-07.
+❌ Layer ('model', 'decoder', 'layers', '3', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 1.4659757141544105e-07 and flax grad norm 1.4895935862568876e-07.
+❌ Layer ('model', 'decoder', 'layers', '3', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.2112769809391466e-07 and flax grad norm 1.382354497536653e-07.
+❌ Layer ('model', 'decoder', 'layers', '4', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 6.575358924010288e-08 and flax grad norm 7.240216604031957e-08.
+❌ Layer ('model', 'decoder', 'layers', '4', 'self_attn', 'k_proj', 'bias') has PT grad norm 5.0567738441031906e-08 and flax grad norm 5.817578241362753e-08.
+❌ Layer ('model', 'decoder', 'layers', '5', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 2.5854723162410664e-07 and flax grad norm 2.9657505251634575e-07.
+❌ Layer ('model', 'decoder', 'layers', '5', 'self_attn', 'k_proj', 'bias') has PT grad norm 8.977868048987148e-08 and flax grad norm 8.695727160556999e-08.
+❌ Layer ('model', 'decoder', 'layers', '6', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 1.271907592581556e-07 and flax grad norm 1.5420768306739774e-07.
+❌ Layer ('model', 'decoder', 'layers', '6', 'self_attn', 'k_proj', 'bias') has PT grad norm 3.32223777377294e-08 and flax grad norm 2.9034252335691235e-08.
+❌ Layer ('model', 'decoder', 'layers', '7', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 2.932817011469524e-08 and flax grad norm 2.990202219166349e-08.
+❌ Layer ('model', 'decoder', 'layers', '7', 'self_attn', 'k_proj', 'bias') has PT grad norm 2.2563149570942187e-08 and flax grad norm 1.903124946522894e-08.
+❌ Layer ('model', 'decoder', 'layers', '8', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 2.1031839381180362e-08 and flax grad norm 1.922012415889185e-08.
+❌ Layer ('model', 'decoder', 'layers', '8', 'self_attn', 'k_proj', 'bias') has PT grad norm 1.9617312219111227e-08 and flax grad norm 1.901776514046105e-08.
+❌ Layer ('model', 'encoder', 'layers', '0', 'self_attn', 'k_proj', 'bias') has PT grad norm 9.498415209918676e-08 and flax grad norm 1.0714285281210323e-07.
+❌ Layer ('model', 'encoder', 'layers', '1', 'self_attn', 'k_proj', 'bias') has PT grad norm 4.2883741002697207e-08 and flax grad norm 3.708849050099161e-08.
+❌ Layer ('model', 'encoder', 'layers', '10', 'self_attn', 'k_proj', 'bias') has PT grad norm 6.720534884152585e-08 and flax grad norm 5.363398614122161e-08.
+❌ Layer ('model', 'encoder', 'layers', '11', 'self_attn', 'k_proj', 'bias') has PT grad norm 7.474503149751399e-08 and flax grad norm 7.271213320336756e-08.
+❌ Layer ('model', 'encoder', 'layers', '2', 'self_attn', 'k_proj', 'bias') has PT grad norm 2.8607770374833308e-08 and flax grad norm 2.428515344377047e-08.
+❌ Layer ('model', 'encoder', 'layers', '3', 'self_attn', 'k_proj', 'bias') has PT grad norm 6.400713203902342e-08 and flax grad norm 5.387828849734433e-08.
+❌ Layer ('model', 'encoder', 'layers', '4', 'self_attn', 'k_proj', 'bias') has PT grad norm 4.873627545975978e-08 and flax grad norm 4.757723104376055e-08.
+❌ Layer ('model', 'encoder', 'layers', '5', 'self_attn', 'k_proj', 'bias') has PT grad norm 5.0619281211083944e-08 and flax grad norm 4.67279193117065e-08.
+❌ Layer ('model', 'encoder', 'layers', '6', 'self_attn', 'k_proj', 'bias') has PT grad norm 5.6844903895125753e-08 and flax grad norm 6.739185920423552e-08.
+❌ Layer ('model', 'encoder', 'layers', '7', 'self_attn', 'k_proj', 'bias') has PT grad norm 5.603576624935158e-08 and flax grad norm 5.457893337279529e-08.
+❌ Layer ('model', 'encoder', 'layers', '8', 'self_attn', 'k_proj', 'bias') has PT grad norm 5.864935914701164e-08 and flax grad norm 6.345069891722233e-08.
+❌ Layer ('model', 'encoder', 'layers', '9', 'self_attn', 'k_proj', 'bias') has PT grad norm 6.470781244161117e-08 and flax grad norm 6.696199505995537e-08.
+=========================================

before_fix_log.txt

@@ -0,0 +1,15 @@
+=========================================
+Check hf-internal-testing/tiny-random-bart ...
+--------------------------Checking logits match--------------------------
+Flax logits shape: (2, 64, 1000), PyTorch logits shape: torch.Size([2, 64, 1000])
+✅ Difference between Flax and PyTorch is 8.940696716308594e-08 (< 0.01)
+--------------------------Checking losses match--------------------------
+Flax loss: 6.923163414001465, PyTorch loss: 6.923163414001465
+✅ Difference between Flax and PyTorch is 0.0 (< 0.01)
+--------------------------Checking gradients match--------------------------
+❌ Layer ('final_logits_bias',) has PT grad norm 0.0 and flax grad norm 0.09163407981395721.
+--------------------------Checking rel gradients match--------------------------
+❌ Layer ('final_logits_bias',) has PT grad norm 0.0 and flax grad norm 0.09163407981395721.
+❌ Layer ('model', 'decoder', 'layers', '0', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 9.3028212357852e-14 and flax grad norm 1.6552796459901042e-13.
+...
+=========================================

before_fix_pretrained_log.txt

@@ -0,0 +1,40 @@
+=========================================
+Check facebook/bart-large ...
+--------------------------Checking logits match--------------------------
+Flax logits shape: (2, 64, 50265), PyTorch logits shape: torch.Size([2, 64, 50265])
+✅ Difference between Flax and PyTorch is 0.00039315223693847656 (< 0.01)
+--------------------------Checking losses match--------------------------
+Flax loss: 15.027304649353027, PyTorch loss: 15.027304649353027
+✅ Difference between Flax and PyTorch is 0.0 (< 0.01)
+--------------------------Checking gradients match--------------------------
+❌ Layer ('final_logits_bias',) has PT grad norm 0.0 and flax grad norm 0.09944064915180206.
+❌ Layer ('model', 'decoder', 'layers', '0', 'fc1', 'kernel') has PT grad norm 13.111018180847168 and flax grad norm 13.0546875.
+❌ Layer ('model', 'decoder', 'layers', '0', 'fc2', 'kernel') has PT grad norm 8.751346588134766 and flax grad norm 8.71875.
+...
+❌ Layer ('model', 'encoder', 'layers', '0', 'self_attn', 'k_proj', 'kernel') has PT grad norm 18.60892105102539 and flax grad norm 18.59375.
+...
+❌ Layer ('model', 'encoder', 'layers', '0', 'self_attn', 'v_proj', 'kernel') has PT grad norm 96.85579681396484 and flax grad norm 96.8125.
+...
+❌ Layer ('model', 'encoder', 'layers', '1', 'self_attn', 'out_proj', 'kernel') has PT grad norm 199.41278076171875 and flax grad norm 199.25.
+...
+--------------------------Checking rel gradients match--------------------------
+❌ Layer ('final_logits_bias',) has PT grad norm 0.0 and flax grad norm 0.09944064915180206.
+❌ Layer ('model', 'decoder', 'layers', '0', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 1.4212106691502413e-07 and flax grad norm 0.0.
+❌ Layer ('model', 'decoder', 'layers', '0', 'self_attn', 'k_proj', 'bias') has PT grad norm 2.0100719311244575e-08 and flax grad norm 0.0.
+...
+=========================================
+Check facebook/bart-large-cnn ...
+--------------------------Checking logits match--------------------------
+Flax logits shape: (2, 64, 50264), PyTorch logits shape: torch.Size([2, 64, 50264])
+✅ Difference between Flax and PyTorch is 0.0001919269561767578 (< 0.01)
+--------------------------Checking losses match--------------------------
+Flax loss: 13.262251853942871, PyTorch loss: 13.262249946594238
+✅ Difference between Flax and PyTorch is 1.9073486328125e-06 (< 0.01)
+--------------------------Checking gradients match--------------------------
+❌ Layer ('final_logits_bias',) has PT grad norm 0.0 and flax grad norm 0.09764379262924194.
+--------------------------Checking rel gradients match--------------------------
+❌ Layer ('final_logits_bias',) has PT grad norm 0.0 and flax grad norm 0.09764379262924194.
+❌ Layer ('model', 'decoder', 'layers', '0', 'encoder_attn', 'k_proj', 'bias') has PT grad norm 2.1513474735002092e-07 and flax grad norm 1.5481474235912174e-07.
+❌ Layer ('model', 'decoder', 'layers', '0', 'self_attn', 'k_proj', 'bias') has PT grad norm 3.8047311079481005e-08 and flax grad norm 3.508952062247772e-08.
+...
+=========================================

check_gradients_pt_flax.py

@@ -163,16 +163,12 @@ def compare_grads(model_id, pt_architecture):
 
     pt_loss.backward()
 
-    pt_grad_dict = {k: v.grad if v.grad is not None else torch.zeros_like(v) for k, v in pt_model.named_parameters()}
+    pt_grad_dict = {k: v.grad for k, v in pt_model.named_parameters()}
+    missing_grads = [k for k in pt_model.state_dict().keys() if k not in pt_grad_dict]
 
-    for k in pt_model.state_dict():
-        if k not in pt_grad_dict:
-            # set any unused parameters to zero in the grad-dict
-            # these won't be compared to the Flax model, but required for loading the PT model from state-dict
-            pt_grad_dict[k] = torch.zeros_like(pt_model.state_dict()[k])
-        pt_model.state_dict()[k] = pt_grad_dict[k]
+    missing_keys, unexpected_keys = pt_model.load_state_dict(pt_grad_dict, strict=False)
 
-    pt_model.load_state_dict(pt_grad_dict)
+    assert missing_grads == missing_keys, f"Error with either grads {missing_keys} or keys {unexpected_keys}"
 
     with tempfile.TemporaryDirectory() as tmpdirname:
         pt_model.save_pretrained(tmpdirname)

run_models.sh

@@ -1,8 +1,8 @@
 #!/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")
+model_ids=("roberta-base" "bert-base-cased" "t5-small" "facebook/bart-large" "facebook/bart-large-cnn")
+model_architectures=("RobertaForMaskedLM" "BertForMaskedLM" "T5ForConditionalGeneration" "BartForConditionalGeneration" "BartForConditionalGeneration" )
+#model_ids=("hf-internal-testing/tiny-random-roberta" "hf-internal-testing/tiny-random-bert" "hf-internal-testing/tiny-random-t5" "hf-internal-testing/tiny-random-bart")
+#model_architectures=("RobertaForMaskedLM" "BertForMaskedLM" "T5ForConditionalGeneration" "BartForConditionalGeneration")
 
 rm -rf log.txt
 touch log.txt