slow training of Squeezeformer with bf16 precision
Hi, all!
Trying to learn a Squeezeformer-S using bf16 precision, but I got **~15 s/it**. But everything is fine learning Conformer with the same number of params with bf16 (got **~3 s/it**) . Also everything goes well using fp16 training with both models (**~3 s/it**). So with all above, I expect Squeezeformer-S to be learnt faster.

**Environment details**
torch==1.10.2
pytorch-lightning==1.7.7
nemo-toolkit==1.11.0
Debian GNU/Linux 10
Python 3.9.7



    name: "Squeezeformer-CTC-ru"
    model:
      sample_rate: 16000
      labels:  [' ', 'а', 'б', 'в', 'г', 'д', 'е', 'ж', 'з', 'и', 'й', 'к', 'л', 'м', 'н', 'о', 'п', 'р', 'с', 'т', 'у', 'ф', 'х', 'ц', 'ч', 'ш', 'щ', 'ъ', 'ы', 'ь', 'э', 'ю', 'я']
    
      log_prediction: false # enables logging sample predictions in the output during training
      ctc_reduction: 'mean_batch'
      skip_nan_grad: false
    
      train_ds:
        manifest_filepath: ???
        labels: ${model.labels}
        sample_rate: ${model.sample_rate}
        batch_size: 128 # you may increase batch_size if your memory allows
        shuffle: true
        num_workers: 32
        pin_memory: true
        trim_silence: false
        max_duration: 15.0 # it is set for LibriSpeech, you may need to update it for your dataset
        min_duration: 0.1
        # tarred datasets
        is_tarred: false
        tarred_audio_filepaths: null
        tarred_shard_strategy: "scatter"
        use_start_end_token: false
        # bucketing params
        bucketing_strategy: "synced_randomized"
        bucketing_batch_size: null
    
      preprocessor:
        _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
        sample_rate: ${model.sample_rate}
        normalize: "per_feature"
        window_size: 0.025
        window_stride: 0.01
        window: "hann"
        features: 80
        n_fft: 512
        log: true
        frame_splicing: 1
        dither: 0.00001
        pad_to: 0
        pad_value: 0.0
    
      spec_augment:
        _target_: nemo.collections.asr.modules.SpectrogramAugmentation
        freq_masks: 2 # set to zero to disable it
        # you may use lower time_masks for smaller models to have a faster convergence
        time_masks: 5 # set to zero to disable it
        freq_width: 27
        time_width: 0.05
    
      encoder:
        _target_: nemo.collections.asr.modules.SqueezeformerEncoder
        feat_in: ${model.preprocessor.features}
        feat_out: -1 # you may set it if you need different output size other than the default d_model
        n_layers: 18
        d_model: 196
    
        # Squeezeformer params
        adaptive_scale: true
        time_reduce_idx: 8
        time_recovery_idx: null
    
        # Sub-sampling params
        subsampling: dw_striding # dw_striding, vggnet, striding or stacking, vggnet may give better results but needs more memory
        subsampling_factor: 4 # must be power of 2
        subsampling_conv_channels: -1 # -1 sets it to d_model
    
        # Feed forward module's params
        ff_expansion_factor: 4
    
        # Multi-headed Attention Module's params
        self_attention_model: rel_pos # rel_pos or abs_pos
        n_heads: 4 # may need to be lower for smaller d_models
        # [left, right] specifies the number of steps to be seen from left and right of each step in self-attention
        att_context_size: [-1, -1] # -1 means unlimited context
        xscaling: true # scales up the input embeddings by sqrt(d_model)
        untie_biases: true # unties the biases of the TransformerXL layers
        pos_emb_max_len: 5000
    
        # Convolution module's params
        conv_kernel_size: 31
        conv_norm_type: 'batch_norm' # batch_norm or layer_norm
    
        ### regularization
        dropout: 0.1 # The dropout used in most of the Conformer Modules
        dropout_emb: 0.0 # The dropout used for embeddings
        dropout_att: 0.1 # The dropout for multi-headed attention modules
    
      decoder:
        _target_: nemo.collections.asr.modules.ConvASRDecoder
        feat_in: null
        num_classes: -1
        vocabulary: ${model.labels}
    
      optim:
        name: adamw
        lr: 0.002
        # optimizer arguments
        betas: [0.9, 0.98]
        # less necessity for weight_decay as we already have large augmentations with SpecAug
        # you may need weight_decay for large models, stable AMP training, small datasets, or when lower augmentations are used
        # weight decay of 0.0 with lr of 2.0 also works fine
        weight_decay: 5e-4
    
        # scheduler setup
        sched:
          name: NoamHoldAnnealing
          # scheduler config override
          warmup_steps: 5000 # paper uses ~ 6500 steps (20 epochs) out of 500 epochs.
          warmup_ratio: null
          hold_steps: 40000 
          hold_ratio: null  # paper uses ~ 40000 steps (160 epochs) out of 500 epochs.
          decay_rate: 1.0  # Noam decay = 0.5 and no hold steps. For Squeezeformer, use hold ~ 10-30% of training, then faster decay.
          min_lr: 1e-7
    
    trainer:
      devices: 8 # number of GPUs, -1 would use all available GPUs
      num_nodes: 1
      max_epochs: 500 
      max_steps: -1 # computed at runtime if not set
      val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
      accelerator: auto
      strategy: ddp
      accumulate_grad_batches: 1
      gradient_clip_val: 1.0
      precision: bf16 # Should be set to 16 for O1 and O2 to enable the AMP.
      log_every_n_steps: 10  # Interval of logging.
      enable_progress_bar: True
      num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
      check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
      sync_batchnorm: true
      enable_checkpointing: False  # Provided by exp_manager
      logger: false  # Provided by exp_manager
      benchmark: false
      resume_from_checkpoint: null 
                                                                                        
    exp_manager:
      exp_dir: ???
      name: ${name} 
      create_tensorboard_logger: true
      create_checkpoint_callback: true
      checkpoint_callback_params:
        # in case of multiple validation sets, first one is used
        monitor: "val_wer"
        mode: "min"
        save_top_k: 5
        always_save_nemo: True # saves the checkpoints as nemo files instead of PTL checkpoints
    
      # you need to set these two to True to continue the training
      #resume_if_exists: false
        #resume_ignore_no_checkpoint: false
    
      # You may use this section to create a W&B logger
      create_wandb_logger: false
      wandb_logger_kwargs:
        name: null
        project: null`