flagos-ai / flagos-ai/FlagScale
FlagGems import fails on Kunlunxin P800 because current Triton lacks `triton.language.math.asin`
- Dominant language
- Python
- Stars
- 537
- Forks
- 174
- Avg merge
- 4d 1h
- Merged PRs (30d)
- 12
Description
## Background
When bringing up the Kunlunxin P800 CI path, the Qwen3 functional training smoke test fails when `enable_flag_gems: true` is enabled.
Affected CI scope:
- `Kunlunxin functional train`
- `functional_test_train (p800, train, qwen3, 0_6b_kunlunxin)`
- Downstream `Kunlunxin all-in-one tests`
Affected config:
```text
tests/functional_tests/train/qwen3/conf/train/0_6b_kunlunxin.yaml
```
The case was originally configured with:
```yaml
enable_flag_gems: true
```
## Environment
- Platform: Kunlunxin P800
- Python: 3.10.18
- Task: Qwen3 functional training smoke test
- Case: `0_6b_kunlunxin`
- Megatron-LM-FL path in the fixed image: `/opt/flagscale/deps/Megatron-LM-FL`
- Distributed backend: NCCL-compatible Kunlunxin/XCCL runtime
## Symptom
Training fails during initialization before any `lm loss:` value is produced.
The failure is triggered while importing `flag_gems`:
```text
import flag_gems
...
from flag_gems.fused import *
from flag_gems.fused.outer import outer
from flag_gems.ops import mul, mv
from flag_gems.ops.arcsin import arcsin, arcsin_, arcsin_out
...
_ASIN = tl_extra_shim.asin
AttributeError: module 'triton.language.math' has no attribute 'asin'. Did you mean: 'sin'?
```
A related warning also appears earlier during operator registration:
```text
[WARNING] Failed to register FlagOS operators: module 'triton.language.math' has no attribute 'asin'
```
The later validation error:
```text
AssertionError: Expected at least 10 values for 'lm loss:', but extracted 0
```
is downstream of the initialization failure. Since training exits before the loop starts, no loss values are written.
## Root cause hypothesis
The current Kunlunxin P800 runtime image appears to provide a Triton version whose `triton.language.math` module does not expose `asin`.
However, the installed `flag_gems` package imports `flag_gems.ops.arcsin` unconditionally during package initialization. That module expects:
```python
tl_extra_shim.asin
```
to exist, so `import flag_gems` fails before training starts.
This looks like a FlagGems/Triton compatibility issue, or a runtime image dependency version mismatch.
It may also indicate that FlagGems imports optional ops too eagerly: even if the training smoke case does not explicitly need the arcsin op, the whole `flag_gems` import fails because one optional op is incompatible with the current Triton API.
## Temporary workaround
To keep the Kunlunxin P800 CI smoke path green while tracking this compatibility issue separately, FlagGems was temporarily disabled only in the Kunlunxin CI branch for this specific Kunlunxin Qwen3 smoke case:
```diff
# tests/functional_tests/train/qwen3/conf/train/0_6b_kunlunxin.yaml
- enable_flag_gems: true
+ enable_flag_gems: false
```
This keeps the core smoke coverage for:
- Megatron-LM-FL
- Transformer Engine FL vendor path
- Kunlunxin/XCCL distributed runtime
- P800 8-card training initialization
- loss logging through `lm loss:`
while avoiding a hard dependency on an incompatible optional FlagGems/Triton path.
## Expected behavior
`enable_flag_gems: true` should not make training fail at `import flag_gems` because a single optional op requires a Triton API that is unavailable in the runtime image.
Possible fixes could include one of the following:
1. provide a FlagGems version compatible with the Triton version in the Kunlunxin P800 image;
2. provide a Triton version exposing `triton.language.math.asin`;
3. guard `tl_extra_shim.asin` usage in `flag_gems.ops.arcsin`;
4. lazily import optional FlagGems ops instead of importing all ops during package initialization;
5. allow disabling only the incompatible `arcsin` op instead of disabling all FlagGems.
## Impact
Without the workaround, the Kunlunxin P800 Qwen3 functional training smoke test fails before producing training loss, which blocks the all-in-one CI validation path.
With the workaround, the CI can still validate the baseline Kunlunxin training stack while FlagGems/Triton compatibility is tracked as a separate issue.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with tests/functional_tests/train/qwen3/conf/train/0_6b_kunlunxin.yaml and reproduce the Qwen3 P800 smoke test with enable_flag_gems enabled. Trace the flag_gems import through flag_gems.ops.arcsin and the missing triton.language.math.asin entry point. Done means the test initializes and produces lm loss values without disabling the intended FlagGems path.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- ci-cd, testing-qa
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100