Skip to content

hcu: implement multi_tensor_scale_tensor using multi_tensor_scale in transformer_engine_hygon 2.13#85

Merged
zhaoyinglia merged 3 commits into
flagos-ai:mainfrom
wangxshuai:main
Jul 13, 2026
Merged

hcu: implement multi_tensor_scale_tensor using multi_tensor_scale in transformer_engine_hygon 2.13#85
zhaoyinglia merged 3 commits into
flagos-ai:mainfrom
wangxshuai:main

Conversation

@wangxshuai

Copy link
Copy Markdown

Description

This PR fixes two compatibility issues in the Hygon backend.

First, it improves the package path resolution for transformer_engine_hygon. In some environments, importlib.util.find_spec() may return a ModuleSpec with origin=None, which causes the previous implementation to fail when resolving the package directory. This PR adds a fallback to submodule_search_locations and reports an error if neither source is available.

Second, transformer_engine_hygon currently does not implement the multi_tensor_scale_tensor API (available in upstream NVIDIA TransformerEngine v2.14). This PR replaces the unsupported call with the existing multi_tensor_scale implementation by extracting the scalar value from the input tensor, preserving equivalent functionality on the Hygon backend.

Fixes # (issue)

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Improve transformer_engine_hygon package path resolution by handling the case where ModuleSpec.origin is None.
  • Add a fallback to submodule_search_locations when locating Hygon backend libraries.
  • Return a descriptive error when the package path cannot be determined.
  • Replace the unsupported multi_tensor_scale_tensor call with multi_tensor_scale by converting the scale tensor to a scalar value.
  • Add comments explaining the compatibility workaround for the Hygon backend.

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@lxd-cumt lxd-cumt left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@zhaoyinglia zhaoyinglia left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@zhaoyinglia zhaoyinglia merged commit 3b5fbb5 into flagos-ai:main Jul 13, 2026
31 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants