[Bugfix] Fix NemotronLayerNorm1P hardcoded cuda device type (#46495)
Signed-off-by: <Michal Ganczarenko> <michal.ganczarenko@intel.com>
Michał Ganczarenko committed
Jun 24, 2026 at 08:21 UTC
ac1fa74616feb95f5db873af724f340993922234
1 file changed
+10
-8
vllm/model_executor/models/nemotron.py
+10
-8
@@ -71,13 +71,14 @@ from .utils import (
71
# - Adds a partial_rotary_factor to RoPE
72
73
74
-def _cast_if_autocast_enabled(*args):
75
- if not torch.is_autocast_enabled():
74
+def _cast_if_autocast_enabled(device_type: str, *args):
75
+ if not torch.is_autocast_enabled(device_type):
76
return args
77
- else:
78
- return torch.amp.autocast_mode._cast(
79
- args, device_type="cuda", dtype=torch.get_autocast_gpu_dtype()
80
- )
77
+ return torch.amp.autocast_mode._cast(
78
+ args,
79
+ device_type=device_type,
80
+ dtype=torch.get_autocast_dtype(device_type),
81
+ )
82
83
84
class NemotronLayerNorm1P(nn.LayerNorm):
@@ -100,10 +101,11 @@ class NemotronLayerNorm1P(nn.LayerNorm):
101
if residual is not None:
102
x = x + residual
103
residual = x
104
+ device_type = x.device.type
105
args = _cast_if_autocast_enabled(
104
- x, self.normalized_shape, self.weight + 1, self.bias, self.eps
106
+ device_type, x, self.normalized_shape, self.weight + 1, self.bias, self.eps
107
)
106
- with torch.amp.autocast("cuda", enabled=False):
108
+ with torch.amp.autocast(device_type, enabled=False):
109
x = torch.nn.functional.layer_norm(*args)
110
return x if residual is None else (x, residual)
111