Force activations and weights cast to FP32 Jasper Encoder Squeeze-Excite (#14715)

* comment out x = x.float()

Signed-off-by: Elena Rastorgueva <erastorgueva@nvidia.com>

* Update Reasoning-SFT.ipynb (#14716) (#14717)

Signed-off-by: Chen Cui <chcui@nvidia.com>
Co-authored-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Elena Rastorgueva <erastorgueva@nvidia.com>

* cast SE weights and activations to fp32

Signed-off-by: Elena Rastorgueva <erastorgueva@nvidia.com>

---------

Signed-off-by: Elena Rastorgueva <erastorgueva@nvidia.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
Co-authored-by: Chen Cui <chcui@nvidia.com>

Author: 3 people

nemo/collections/asr/parts/submodules/jasper.py

@@ -477,7 +477,13 @@ def forward_for_export(self, x, lengths):
             # Create sample mask - 1 represents value, 0 represents pad
             mask = self.make_pad_mask(lengths, max_audio_length=max_len, device=x.device)
             mask = ~mask  # 0 represents value, 1 represents pad
-            x = x.float()  # For stable AMP, SE must be computed at fp32.
+
+            # Ensure SE runs in FP32: cast fc weights and activations to float32
+            if self.fc[0].weight.dtype != torch.float32:
+                self.fc.float()
+            if x.dtype != torch.float32:
+                x = x.float()
+
             x = x.masked_fill(mask, 0.0)  # mask padded values explicitly to 0
             y = self._se_pool_step(x, mask)  # [B, C, 1]
             y = y.transpose(1, -1)  # [B, 1, C]
@@ -490,6 +496,8 @@ def forward_for_export(self, x, lengths):
 
             y = torch.sigmoid(y)
             y = x * y
+            # Cast back to original dtype for downstream consistency
+            y = y.to(dtype)
         return y, lengths
 
     def _se_pool_step(self, x, mask):