UC-02: Rewrite _check_mask/_check_bias with explicit shape checks (Ascend NPU adaptation) - #74
Open
slamdunk111 wants to merge 2 commits into
Open
Conversation
added 2 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).
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
Extract
_check_maskand_check_biasfrom inlinetry/exceptblocks in the reference path into standalone, unit-testable validators. Replace bareexcept:with explicit exception types.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:
except:silently swallowed mask/bias validation errors during Ascend training runs, making it impossible to diagnose shape mismatches. With explicitRuntimeErrorcatching, mask/bias issues on Ascend NPU now surface immediately with clear error messages._check_maskand_check_biascan now be unit-tested in isolation on Ascend NPUs (910B2/910C) without constructing the fullsoftmax_dropoutpipeline.KeyboardInterrupt/MemoryErrorswallowing — on Ascend NPU long-running training jobs, a bareexcept:could interceptKeyboardInterrupt(preventing clean shutdown) orMemoryError(hiding Ascend NPU out-of-memory conditions). OnlyRuntimeErroris now caught.Changes
unicore/modules/softmax_dropout.py:_check_mask(tensor, mask, inplace)and_check_bias(tensor, bias, inplace)as standalone functions; reference path calls themWhy
The original reference path (which is the path Ascend NPUs — 910B2, 910C, and future variants — always use, since no CUDA fused operator exists on Ascend) used bare
except: print(...)which:KeyboardInterruptandMemoryError— dangerous on Ascend NPU training jobsThis PR makes the validators independently testable and replaces bare
exceptwith explicitRuntimeErrorcatching.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
RuntimeErrorcaught (no moreKeyboardInterrupt/MemoryErrorswallowing on Ascend NPU)inplacesafety — verified on Ascend 910B2inplacesafety — verified on Ascend 910B2softmax_dropoutpipeline — works on Ascend NPUMotivation Correction (Expert Review S2-8)
The original bare
exceptdid not crash on low-rank input (historical fact corrected). The real value for Ascend NPU adaptation is: no longer catchesKeyboardInterrupt/MemoryError(critical for Ascend training jobs), and the validation logic is now explicit and unit-testable on Ascend.E2E Gate Status
Test Plan
git amZero behavior change for valid inputs on Ascend NPU. Invalid inputs now raise
RuntimeErrorexplicitly instead of being silently swallowed.Patch Source
Generated from cnpc-chem-opt patch mirror. See umbrella issue #76 for the full Ascend NPU adaptation overview.