Skip to content

hcu: implement multi_tensor_scale_tensor using multi_tensor_scale in transformer_engine_hygon 2.13 #85

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhaoyinglia merged 3 commits into from
Jul 13, 2026
+14 −2
Merged

hcu: implement multi_tensor_scale_tensor using multi_tensor_scale in transformer_engine_hygon 2.13 #85

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8c4bd9e
fix(hygon/platform): Improve Hygon library path resolution, fall back…
wangyl166 Jun 23, 2026
3e2c9f2
fix: auto format code by pre-commit black & trim whitespace
wangyl166 Jun 27, 2026
b628315
fix(hygon): workaround missing multi_tensor_scale_tensor by using mul…
wangyl166 Jun 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 14 additions & 2 deletions transformer_engine/plugin/core/backends/vendor/hygon/hygon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,16 @@ def _get_sys_extension() -> str:
hygon_spec = importlib.util.find_spec("transformer_engine_hygon")
if hygon_spec is None:
return False
hygon_path = Path(hygon_spec.origin).parent
if hygon_spec.origin is not None:
hygon_path = Path(hygon_spec.origin).parent
elif hygon_spec.submodule_search_locations:
hygon_path = Path(hygon_spec.submodule_search_locations[0])
else:
print(
"[ERROR _load_hygon_libs] cannot determine package path, origin is None and"
" submodule_search_locations is empty"
)
return False
for file_path in hygon_path.iterdir():
if file_path.name.startswith(common_prefix) and file_path.suffix == ext:
common_files.append(file_path)
Expand Down Expand Up @@ -1628,8 +1637,11 @@ def multi_tensor_scale_tensor(
tensor_lists: List[List[torch.Tensor]],
scale: torch.Tensor,
) -> None:
# transform_engine_hygon does not support multi_tensor_scale_tensor
# (from upstream Nvidia TE v2.14). Use multi_tensor_scale as a workaround.
tex = self._get_tex()
return tex.multi_tensor_scale_tensor(chunk_size, noop_flag, tensor_lists, scale)
scale_value = scale.item()
return tex.multi_tensor_scale(chunk_size, noop_flag, tensor_lists, scale_value)

def multi_tensor_l2norm(
self,
Expand Down
Loading