PCC drop in Mixer_L16_224_In21k , PCC=0.9843129116423043

[kamalrajkannan78]

Describe the bug

• Mixer_L16_224_In21k having PCC=0.9843129116423043 - feb16_mlp_mixer.log
• Drop comes after the Layernorm operation
  • PCC of the model while returning output before norm - 0.9987465364058089 - feb16_mlp_mixer_d1.log, mlp_mixer.py
  • PCC of the model while returning output after norm - 0.984995753301239 - feb16_mlp_mixer_d2.log, mlp_mixer.py
• This indicates that issue is either from norm alone or norm + few ops from last block.
• Norm sanity passed with PCC=0.9999935384890567 - feb16_mlp_mixer_s1.log, loader.py
• Expanded the sanity by adding preceding add op, but this sanity also passed - 0.9998903181887254 - feb16_mlp_mixer_s2.log, loader.py
• Further expanded it by adding entire MixerBlock with it. but this also passed - 0.9996960392476073 - feb16_mlp_mixer_s3.log, loader.py
• Note : Saved inputs are passed to all sanity experiments
• In sum, I'm unable to replicate the drop with op-level/ block-level sanity

Repro

git checkout kkannan/feb16_Mixer_L16_224_In21k_pcc_drop
git submodule update --init --recursive
git lfs pull --include '*.pt'

# whole model run 
pytest tests/runner/test_models.py::test_all_models_torch[mlp_mixer/pytorch-Mixer_L16_224_In21k-single_device-inference] -svv

#  For each exp, place the attached version of files in original files & rerun above command