Skip to content

Commit 1f98511

Browse files
lxd-cumtlxd-cumt
authored
Fix quantizer dtype conversion errors (#54)
- Fix quantizer dtype attr conversion errors for vendor backends - Polish logger for vendor backend
1 parent 7f788a3 commit 1f98511

7 files changed

Lines changed: 98 additions & 7 deletions

File tree

‎transformer_engine/__init__.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,7 @@
2121
from .plugin.core.backends.vendor.musa.patches import apply_patch as _musa_apply_patch
2222

2323
_musa_apply_patch()
24-
print("[TE-FL] MUSA patches applied")
2524
except Exception as e:
26-
print(f"[TE-FL] MUSA patches not applied: {e}")
2725
pass
2826

2927

‎transformer_engine/plugin/core/backends/vendor/cuda/cuda.py‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,6 @@ def try_load_lib(name, search_patterns):
7979
return True
8080
return False
8181
except Exception as e:
82-
print(f"[CUDA] Failed to load CUDA libs: {e}")
8382
return False
8483

8584

@@ -90,6 +89,8 @@ def _ensure_cuda_libs():
9089
global _cuda_libs_loaded
9190
if not _cuda_libs_loaded:
9291
_cuda_libs_loaded = _load_cuda_libs()
92+
if _cuda_libs_loaded:
93+
print(f"[CUDA] Successfully loaded CUDA libs")
9394
return _cuda_libs_loaded
9495

9596

@@ -166,6 +167,15 @@ def quantize(
166167
noop: Optional[torch.Tensor] = None,
167168
) -> Any:
168169
tex = self._get_tex()
170+
# Normalize quantizer.dtype to this backend's `tex.DType`.
171+
try:
172+
if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
173+
qdtype = quantizer.dtype
174+
if qdtype is not None:
175+
quantizer.dtype = tex.DType(int(qdtype))
176+
except Exception:
177+
pass
178+
169179
return tex.quantize(tensor, quantizer, output, noop)
170180

171181
def dequantize(
@@ -183,6 +193,16 @@ def bgrad_quantize(
183193
quantizer: Any,
184194
) -> List[Any]:
185195
tex = self._get_tex()
196+
197+
# Normalize quantizer.dtype to this backend's `tex.DType`.
198+
try:
199+
if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
200+
qdtype = quantizer.dtype
201+
if qdtype is not None:
202+
quantizer.dtype = tex.DType(int(qdtype))
203+
except Exception:
204+
pass
205+
186206
return tex.bgrad_quantize(input, quantizer)
187207

188208
def generic_gemm(

‎transformer_engine/plugin/core/backends/vendor/hygon/hygon.py‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,6 @@ def _get_sys_extension() -> str:
5454
spec.loader.exec_module(solib)
5555
return True
5656
except Exception as e:
57-
print(f"[HYGON] Failed to load hygon libs: {e}")
5857
return False
5958

6059

@@ -65,6 +64,8 @@ def _ensure_hygon_libs():
6564
global _hygon_libs_loaded
6665
if not _hygon_libs_loaded:
6766
_hygon_libs_loaded = _load_hygon_libs()
67+
if _hygon_libs_loaded:
68+
print(f"[HYGON] Successfully loaded HYGON libs")
6869
return _hygon_libs_loaded
6970

7071

@@ -145,6 +146,13 @@ def quantize(
145146
noop: Optional[torch.Tensor] = None,
146147
) -> Any:
147148
tex = self._get_tex()
149+
try:
150+
if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
151+
qdtype = quantizer.dtype
152+
if qdtype is not None:
153+
quantizer.dtype = tex.DType(int(qdtype))
154+
except Exception:
155+
pass
148156
return tex.quantize(tensor, quantizer, output, noop)
149157

150158
def dequantize(
@@ -162,6 +170,16 @@ def bgrad_quantize(
162170
quantizer: Any,
163171
) -> List[Any]:
164172
tex = self._get_tex()
173+
174+
# Normalize quantizer.dtype to this backend's `tex.DType`.
175+
try:
176+
if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
177+
qdtype = quantizer.dtype
178+
if qdtype is not None:
179+
quantizer.dtype = tex.DType(int(qdtype))
180+
except Exception:
181+
pass
182+
165183
return tex.bgrad_quantize(input, quantizer)
166184

167185
def generic_gemm(

‎transformer_engine/plugin/core/backends/vendor/iluvatar/iluvatar.py‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,6 @@ def try_load_lib(name, search_patterns):
7777
return True
7878
return False
7979
except Exception as e:
80-
print(f"[ILUVATAR] Failed to load ILUVATAR libs: {e}")
8180
return False
8281

8382

@@ -88,6 +87,8 @@ def _ensure_iluvatar_libs():
8887
global _iluvatar_libs_loaded
8988
if not _iluvatar_libs_loaded:
9089
_iluvatar_libs_loaded = _load_iluvatar_libs()
90+
if _iluvatar_libs_loaded:
91+
print(f"[ILUVATAR] Successfully loaded ILUVATAR libs")
9192
return _iluvatar_libs_loaded
9293

9394

@@ -171,6 +172,13 @@ def quantize(
171172
noop: Optional[torch.Tensor] = None,
172173
) -> Any:
173174
tex = self._get_tex()
175+
try:
176+
if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
177+
qdtype = quantizer.dtype
178+
if qdtype is not None:
179+
quantizer.dtype = tex.DType(int(qdtype))
180+
except Exception:
181+
pass
174182
return tex.quantize(tensor, quantizer, output, noop)
175183

176184
def dequantize(
@@ -188,6 +196,16 @@ def bgrad_quantize(
188196
quantizer: Any,
189197
) -> List[Any]:
190198
tex = self._get_tex()
199+
200+
# Normalize quantizer.dtype to this backend's `tex.DType`.
201+
try:
202+
if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
203+
qdtype = quantizer.dtype
204+
if qdtype is not None:
205+
quantizer.dtype = tex.DType(int(qdtype))
206+
except Exception:
207+
pass
208+
191209
return tex.bgrad_quantize(input, quantizer)
192210

193211
def generic_gemm(

‎transformer_engine/plugin/core/backends/vendor/metax/metax.py‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,6 @@ def get_ext():
3737
return True
3838
return False
3939
except Exception as e:
40-
print(f"[Metax] Failed to load Metax libs: {e}")
4140
return False
4241

4342

@@ -48,6 +47,8 @@ def _ensure_metax_libs():
4847
global _metax_libs_loaded
4948
if not _metax_libs_loaded:
5049
_metax_libs_loaded = _load_metax_libs()
50+
if _metax_libs_loaded:
51+
print(f"[Metax] Successfully loaded Metax libs")
5152
return _metax_libs_loaded
5253

5354

@@ -126,6 +127,13 @@ def quantize(
126127
noop: Optional[torch.Tensor] = None,
127128
) -> Any:
128129
tex = self._get_tex()
130+
try:
131+
if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
132+
qdtype = quantizer.dtype
133+
if qdtype is not None:
134+
quantizer.dtype = tex.DType(int(qdtype))
135+
except Exception:
136+
pass
129137
return tex.quantize(tensor, quantizer, output, noop)
130138

131139
def dequantize(
@@ -143,6 +151,16 @@ def bgrad_quantize(
143151
quantizer: Any,
144152
) -> List[Any]:
145153
tex = self._get_tex()
154+
155+
# Normalize quantizer.dtype to this backend's `tex.DType`.
156+
try:
157+
if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
158+
qdtype = quantizer.dtype
159+
if qdtype is not None:
160+
quantizer.dtype = tex.DType(int(qdtype))
161+
except Exception:
162+
pass
163+
146164
return tex.bgrad_quantize(input, quantizer)
147165

148166
def generic_gemm(

‎transformer_engine/plugin/core/backends/vendor/musa/musa.py‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,6 @@ def try_load_lib(name, search_patterns):
6666

6767
return True
6868
except Exception as e:
69-
print(f"[MUSA] Failed to load MUSA libs: {e}")
7069
return False
7170

7271

@@ -77,6 +76,8 @@ def _ensure_musa_libs():
7776
global _musa_libs_loaded
7877
if not _musa_libs_loaded:
7978
_musa_libs_loaded = _load_musa_libs()
79+
if _musa_libs_loaded:
80+
print(f"[MUSA] Successfully loaded MUSA libs")
8081
return _musa_libs_loaded
8182

8283

@@ -138,6 +139,13 @@ def quantize(
138139
noop: Optional[torch.Tensor] = None,
139140
) -> Any:
140141
tex = self._get_tex()
142+
try:
143+
if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
144+
qdtype = quantizer.dtype
145+
if qdtype is not None:
146+
quantizer.dtype = tex.DType(int(qdtype))
147+
except Exception:
148+
pass
141149
return tex.quantize(tensor, quantizer, output, noop)
142150

143151
def dequantize(
@@ -155,6 +163,16 @@ def bgrad_quantize(
155163
quantizer: Any,
156164
) -> List[Any]:
157165
tex = self._get_tex()
166+
167+
# Normalize quantizer.dtype to this backend's `tex.DType`.
168+
try:
169+
if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
170+
qdtype = quantizer.dtype
171+
if qdtype is not None:
172+
quantizer.dtype = tex.DType(int(qdtype))
173+
except Exception:
174+
pass
175+
158176
return tex.bgrad_quantize(input, quantizer)
159177

160178
def generic_gemm(

‎transformer_engine/plugin/core/backends/vendor/musa/patches.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,3 +70,4 @@ def apply_patch() -> None:
7070
except Exception:
7171
# Best-effort: patching should never crash import/initialization.
7272
continue
73+
print(f"[TE-FL] MUSA backend patches applied")

0 commit comments

Comments
 (0)