Inc on vLLM -Change RMS norm to BF16

HabanaAI · Aug 13, 2024 · 3b34893 · 3b34893
1 parent ed81599
commit 3b34893
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py
@@ -79,18 +79,16 @@ def forward_hpu(
         if HPUFusedRMSNorm is None:
             return self.forward_native(x, residual)
         if residual is not None:
-            orig_dtype = x.dtype
             orig_shape = x.shape
             residual += x.view(residual.shape)
             # Note: HPUFusedRMSNorm requires 3D tensors as inputs
-            x = HPUFusedRMSNorm.apply(residual.float(), self.weight.float(),
+            x = HPUFusedRMSNorm.apply(residual, self.weight,
                                       self.variance_epsilon)
-            return x.to(orig_dtype).view(orig_shape), residual
+            return x.view(orig_shape), residual
 
-        orig_dtype = x.dtype
-        x = HPUFusedRMSNorm.apply(x.float(), self.weight.float(),
+        x = HPUFusedRMSNorm.apply(x, self.weight,
                                   self.variance_epsilon)
-        return x.to(orig_dtype)
+        return x
 
     def forward_xpu(
         self,