Conformer OOM fix

[chandramouli-sastry]

Contains changes to fix #497

• I replaced the lazy linear layers with linear layers.
• Along with few minor changes/improvements throughout the conformer model.py, I configured the scaled-dot-product-attention backend to use the math backend and this helped fix the out of memory error.
• I updated the comparator and confirmed that the jax/pytorch implementations are identical:
  
• Running with torch.compile=True still gives OOM and this only fixes OOM for torch.compile=False (tested with NAdamW and AdamW) and I disabled compilation for librispeech_conformers until this can be fixed.
• For reference, I also profiled the run for 500 steps -- not sure if this is in the ballpark :)

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[msaroufim]

curious that math sdp was more memory efficient than memory efficient attention

[chandramouli-sastry]

It is likely more memory efficient but it seems like it requires a different memory configuration setting to sidestep the requirement of torch.cuda.empty_cache() after every update step.

[msaroufim]

yeah i feel like this is some corner case sdpa bug but should be fine to merge still

[msaroufim]

what's the reasoning for this?

[chandramouli-sastry]

Each of the two subsampling layers reduce the number mel-spectrogram features by half.

[msaroufim]

ok i assume that the new implementation is numerically equivalent, SDPA is probably the right bet

[chandramouli-sastry]

I also tried the default multihead self-attention without query scaler and it still couldn't run successfully without adjusting the attention backends. If this is useful for your debugging, I can create a separate branch with this setup?

[msaroufim]

so the previous implementation was quite long sdpa is something that's maintained in pytorch so this actually makes things better

[chandramouli-sastry]

The previous implementation was long because it was extending the nn.MultiheadAttention by changing the in-projection weights/biases in the forward pass but the attention was still managed entirely by pytorch.

[priyakasimbeg]

Confirmed that this works for 60K run:

(60000, {'train/ctc_loss': Array(0.09971161, dtype=float32), 'train/wer': 0.036610075914423744, 'validation/ctc_loss': 
Array(0.30221355, don/wer': type=float32), 'validation/wer': 0.08587637121438702, 'validation/num_examples': 5348,  'test/ctc_loss': Array(0.16302903, dtype=float32), 'teest/num_st/wer': 0.05093237974018627, 'test/num_examples': 2472,  'score': 46081.4025554657, 'total_duration': 50274.87093257904, 'accumulated_submissted_evalion_time':  46081.4025554657, 'accumulated_eval_time': 4190.642651796341, 'accumulated_logging_time': 1.6540186405181885,  'global_step': 60000c_loss':, 'preemption_count': 0})], 'global_step': 60000} 
 r': 0.08I1019 15:00:12.858658 139843127400256 submission_runner.py:550] Timing: 46081.4025554657                                                  um_exampI1019 15:00:12.858741 139843127400256 submission_runner.py:552] Total number of evals: 33                                                  l_time':I1019 15:00:12.858824 139843127400256 submission_runner.py:553] ====================