UC-03: Cleanup bare except + lazy CUDA probe, no import-time print (Ascend NPU adaptation) - #75
Open
slamdunk111 wants to merge 3 commits into
Open
Conversation
added 3 commits
July 11, 2026 19:29
Motivation: The input preprocessing (contiguous + clone) and the generic PyTorch softmax+dropout fallback were inline in softmax_dropout(), making them impossible to unit-test in isolation. Extracting them into standalone helpers ensures that upstream's single-preprocessing semantics are strictly preserved, preventing any candidate implementation from introducing a duplicate-clone on the non-inplace path. Changes: - Extract _prepare_input(input, inplace): single point for contiguous+clone. - Extract _softmax_dropout_reference_prepared(x, ...): the exact upstream else-branch computation (add mask/bias, F.softmax, F.dropout) operating on already-prepared input. Uses add_ for mask/bias in all cases, since the input is already a safe copy (cloned by _prepare_input when inplace=False) — this matches upstream's allocation semantics and avoids extra tensor allocations. - Extract _softmax_dropout_reference(input, ...): convenience wrapper that calls _prepare_input then _softmax_dropout_reference_prepared, for tests that call the reference path directly with a raw tensor. - Modify softmax_dropout to call _prepare_input once and delegate the else-branch to _softmax_dropout_reference_prepared. No behavior change for any input — the computation is identical, just reorganized into testable functions. TestPlan: - Run existing Uni-Core test suite; all results must be unchanged. - New unit tests for _softmax_dropout_reference (numerical equivalence with F.dropout(F.softmax(...)), inplace/non-inplace, mask/bias, 4D). - Verify _prepare_input does not double-clone (inplace=False path).
Motivation: The try/except-based validators used bare ``except:`` which catches *all* exceptions including ``KeyboardInterrupt``, ``MemoryError``, and ``SystemExit`` that should propagate. While the original code happened to return False for low-rank inputs (``IndexError`` from ``mask.shape[-3]`` was caught by the bare except), this relied on exception-based control flow which is fragile and hard to test. The main improvements are: 1. No longer catches ``KeyboardInterrupt``/``MemoryError``/``SystemExit``. 2. Logic is explicit and unit-testable without relying on exceptions. 3. Low-rank inputs return False via explicit ``ndim`` guards instead of depending on ``IndexError`` being caught. Changes: - _check_mask: replace try/except/assert with explicit if-return checks. Add ``mask.ndim < 3`` guard so low-rank inputs return False cleanly. Use ``.ndim`` instead of ``len(.shape)`` (more idiomatic). - _check_bias: same treatment. Add ``bias.ndim < 2`` guard for low-rank. - Behavior is identical for all valid inputs; only difference is that low-rank inputs now return False instead of potentially crashing. TestPlan: - Unit tests: valid mask/bias shapes return True; invalid shapes return False. - Low-rank regression tests: rank-0, rank-1, rank-2 inputs return False for _check_mask; rank-0, rank-1 return False for _check_bias. - Existing CUDA path behavior unchanged (validators return same result for all rank >= 3 inputs).
Motivation: Upstream probes the CUDA fused extension at module import time
with a bare ``except:`` and a ``print()`` call (expert 8.2). The bare
except silently swallows all exceptions (including KeyboardInterrupt,
MemoryError, SystemExit), and the print produces spurious stdout output
in CPU-only environments (CI logs, NPU containers). Additionally, the
import-time probe result is stale if CUDA is initialized later.
Changes:
- Remove module-level ``try: import unicore_fused_softmax_dropout`` block
and the ``print("fused_softmax is not installed corrected")`` line.
- Remove module-level ``HAS_SOFTMAX`` variable and the
``torch.cuda.get_device_capability()`` check at import time.
- Add ``_cuda_fused_available()`` function: lazily probes the extension
import and CUDA device capability, caches the result per device_id in
``_CUDA_FUSED_CACHE`` dict. Only ``ImportError`` and ``RuntimeError``
are caught — other exceptions propagate.
- Move ``import unicore_fused_softmax_dropout`` into
``SoftmaxDropoutFast.forward`` and ``.backward`` (lazy import, only
when the CUDA path is actually taken).
- Replace ``input.is_cuda and HAS_SOFTMAX`` with
``input.is_cuda and _cuda_fused_available()`` in softmax_dropout.
- Add ``import importlib`` and ``import logging``; module-level
``logger = logging.getLogger(__name__)``.
TestPlan:
- Verify importing the module produces no stdout output.
- _cuda_fused_available() returns False when torch.cuda.is_available()
is False (CPU-only env).
- _cuda_fused_available() returns False when extension ImportError occurs.
- _cuda_fused_available() caches result (second call does not re-probe).
- Existing CUDA path behavior unchanged for CUDA inputs with the extension
installed.
13 tasks
slamdunk111
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Replace import-time
print()side effect in_cuda_softmax_dropout_available()with a lazy, cached probe function_is_fused_softmax_available(). The check now runs only when the CUDA fused path is actually attempted, and is cached after first call.Ascend NPU Adaptation
This PR is part of the Ascend NPU adaptation series for Uni-Core's
softmax_dropoutmodule (UC-01~UC-03), targeting Ascend NPUs such as 910B2, 910C, and future variants.What this PR does for Ascend:
"Cuda is not available. Will fallback to ..."at import time whenever CUDA was absent. On Ascend NPUs (910B2/910C), this fired on everyimport unicore.modules.softmax_dropout, polluting Ascend training logs and breaking test output capture. The lazy probe produces no output when CUDA is absent — critical for clean Ascend NPU training logs.builtins.__import__, but the actual code usessd.importlib.import_module. This meant the probe was silently broken on Ascend NPU, running the full check on every forward pass instead of caching. The fix correctly patchessd.importlib.import_moduleand caches the result.Changes
unicore/modules/softmax_dropout.py:_is_fused_softmax_available()replaces import-timeprint(); cached via module-level_FUSED_AVAILABLEsentinelWhy
The original code printed a warning at import time whenever CUDA was unavailable:
On Ascend NPUs (910B2/910C and future variants), this is the default environment — CUDA is always absent. Problems:
import unicore.modules.softmax_dropoutprints to stdout on Ascend NPU, polluting training logs and breaking test capturesoftmax_dropoutcall during Ascend trainingsd.importlib.import_module: the probe patched the wrong import target (builtins.__import__), silently failing on Ascend NPUs (910B2/910C)This PR makes the probe lazy (runs only when CUDA path is attempted) and cached (runs at most once per process) — both essential for Ascend NPU performance and log cleanliness.
Quality Assurance
This PR is part of the Ascend NPU adaptation refactor series (UC-01~UC-03). Full QA artifacts in the patch Mirror repo.
Test Results (2026-07-11, Ascend 910B2 NPU, torch_npu 2.7.1.post2, Python 3.11, PyTorch 2.7.1)
Key Verified Properties
_FUSED_AVAILABLEsentinel — probe runs at most once per process; subsequent calls return cached result on AscendFalsewhen no CUDA: graceful fallback to reference path on Ascend NPUcompute_capability >= (7, 0)when CUDA availablesd.importlib.import_module(notbuiltins.__import__) — Expert S4-5 fix, essential for Ascend NPU correctnessExpert Review Fixes (S4-5)
The original probe patched
builtins.__import__, but the actual code usessd.importlib.import_module. On Ascend NPU, this meant the probe was silently broken. The test correctly patchessd.importlib.import_moduleto match the real import path.E2E Gate Status
Test Plan
Falsewhen no CUDA (Ascend NPU default)git amZero behavior change for the CUDA fused path when CUDA is available. The only observable difference on Ascend NPU: no import-time print, and the probe is cached.
Patch Source
Generated from cnpc-chem-opt patch mirror. See umbrella issue #76 for the full Ascend NPU adaptation overview.