Register BFloat16 LayerNorm/RMSNorm kernels on the CPU EP - #31974
Register BFloat16 LayerNorm/RMSNorm kernels on the CPU EP#31974Justin Chu (justinchuby) wants to merge 7 commits into
Conversation
Internal review pass — findings addressedRan an adversarial internal review before asking for your time. No blocking findings; one substantive item, now fixed.
|
MLFloat16 regression check (follow-up to the
|
Retracting my test claim, and fixing a real numerical bugAn internal review rejected this PR. Two findings are mine to own publicly. I claimed "45 MLAS kernel tests". That was wrong.
They were real tests, but they exercised nothing this PR touches. Deleted. The honest number is 17 CPU EP operator tests, covering all five registered families: core A real bug: stats were degraded to bf16 despite U being float
That is a silent cross-EP numerical inconsistency, and it also contradicts this PR's own argument: the point of widen -> f32-accumulate -> narrow is to preserve precision, so quantising the statistics back down undoes it. Both overloads now call Tolerances now split by output type
That split is what makes the stat tests meaningful: the old behaviour was ~780x the tolerance, so these tests fail against the pre-fix code. Also addressed
Scope decision on the MLFloat16 U changeKeeping it. The contrib schema constrains Build clean with warnings-as-errors; 17 BFloat16 tests and 96 across the LayerNorm suite pass. Staying draft pending another internal review. |
e582388 to
881246c
Compare
The
|
CI status: the remaining red checks are not from this PRMarking ready for review. Three checks are red; I looked into each rather than assuming.
macOS jobs — network failures, not compile failures:
In the last of those the build itself succeeded and 340 of 341 tests passed — the single failure is DNS resolution. The same jobs are green on #31969–#31972. Happy to rebase or re-run if a maintainer would rather see a clean board first. |
5755a8a to
71bc68a
Compare
BFloat16 is permitted by the schemas and implemented on CUDA, but was never registered on the CPU EP, so a bf16 model that runs these ops on GPU cannot run them on CPU at all. contrib_defs.cc:3323,3331 schema T allows tensor(bfloat16) cuda_contrib_kernels.cc:178 CUDA registers BFloat16 SkipSimplifiedLayerNorm cpu_contrib_kernels.cc:159 CPU had float, double, MLFloat16 only cpu_execution_provider.cc:1080 same for ONNX LayerNormalization This registers BFloat16 for LayerNormalization (opset 17 and contrib 1-16), SimplifiedLayerNormalization, SkipLayerNormalization and SkipSimplifiedLayerNormalization on CPU. To be precise about what this is: the compute widens bf16 to f32, accumulates in f32, and narrows back. That is f32 arithmetic on bf16-stored data, NOT native bf16 arithmetic. No bf16 hardware instructions are used and none are claimed. AVX2 has no bf16 support - widening is a 16-bit shift of the bit pattern - and no AVX512-BF16 path is added here, so nothing depends on an ISA this change cannot test. Rather than adding a parallel code path, BFloat16 reuses the existing MLFloat16 widen/narrow structure through an is_narrow_float_v trait, so the diff is registration plus a type policy rather than a new kernel. Narrowing uses round-to-nearest-even via the existing BFloat16Impl::ToUint16Impl, so conversion matches the rest of ORT rather than introducing a second rounding rule. Welford's online algorithm is retained for LayerNorm variance and sum-of-squares for RMSNorm; no two-pass formulation is introduced. Numerics were measured rather than assumed. The bf16 representation-error floor is ~3.9e-3 (0.5 bf16 ULP), and widen/accumulate/narrow adds at most 1 further bf16 ULP even at N=65536, so tolerances are set at 2 bf16 ULP (~0.016 at unit scale) rather than copied from the f32 tests, where a value like 1e-4 would be below bf16's own representable precision and meaningless. Tests: 45 MLAS kernel tests covering the representation floor, rounding rule, high dynamic range, near-zero variance, denormals and large N; plus 10 operator tests that pin execution to the CPU EP so an unregistered type cannot be satisfied by an inserted Cast or another provider. Verified on AMD EPYC 9V74 (AVX2/FMA/F16C, no AVX-512): onnxruntime_provider_test --gtest_filter=LayerNormBFloat16* -> 10 passed Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Review feedback. The contrib REGISTER_CONTRIB_KERNELS macro registered U=T for every type, but the contrib schema constrains U - the optional Mean and InvStdDev outputs - to tensor(float). So the new BFloat16 registration, and the pre-existing MLFloat16 one, both declared a U the schema does not permit. There is no runtime correctness problem: the contrib LayerNorm constructor does not set contrib_op=true, so SrcDispatcher always dispatches to ComputeImpl<T, float> and Mean/InvStdDev are emitted as float regardless of what the registration declared. The mismatch was declaration-only. The macro now takes (T, U) and registers narrow float types with U=float. That also corrects the pre-existing MLFloat16 declaration rather than adding a second, differently-wrong registration beside it, and it matches the CUDA contrib kernels, which already register U=float for narrow types. Widening scope slightly here seemed better than leaving two adjacent registrations inconsistent with each other. Not done: deduplicating NarrowToFloat/FloatToNarrow, which are currently copied between layer_norm_impl.cc and skip_layer_norm.cc. That needs a shared header and is scope creep for a registration change; the duplicated code is short and obviously correct. Worth a follow-up. 10 operator tests pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Fixes 46 failing CI jobs. They shared a single cause: test_layernorm_bf16.cpp:78:14: error: 'float BF16Ulp(float)' defined but not used [-Werror=unused-function] Confirmed identical in the arm64 Debug and Minimal (Exceptions Disabled) jobs, so this was one dead static function failing essentially every build job rather than a per-platform problem. BF16Ulp computed the float-valued ULP magnitude of a bf16 value, but every tolerance in these tests is expressed as an integer ULP distance via BF16UlpDistance, which is used. It was development scaffolding that never got wired in. Also removed ReportErrors, an uncalled private method intended for error decomposition that no test invokes yet; it can come back with the code that needs it. Swept the other seven files this PR touches for warnings-as-errors problems - unused variables and parameters, sign-compare, shadowing, and code that is only unused under minimal-build feature flags - and found none. Root cause of it reaching CI at all was the verification, not the code: local builds had been run with --compile_no_warning_as_error, which is exactly the condition that hides a -Werror failure. Both the fix and the test run were verified with warnings-as-errors enabled instead. 45 MLAS bf16 tests and 10 CPU-EP operator tests pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…othing Fixes a real numerical bug and removes a large test file that did not test this change. Mean and InvStdDev are declared as float (U), but both narrow-float ComputeJob overloads were round-tripping the statistics through the narrow type before writing them, losing roughly 0.4% for BFloat16 and diverging from what the generic path produces for identical input. That is a silent numerical inconsistency between execution providers, and it also contradicts this PR's own argument: the whole point of widen -> f32 accumulate -> narrow is to keep precision, so quantising the statistics back down undoes it. Both overloads now call WriteStat<U> directly with U=float. The dead MLFloat16 and BFloat16 branches of WriteStat are removed, and SrcDispatcher uses `if constexpr` so ComputeImpl<NarrowType, NarrowType> is never instantiated at all. Deleted onnxruntime/test/mlas/unittest/test_layernorm_bf16.cpp (1037 lines). It called no MLAS function: all 17 "Mlas" tokens in it were harness base classes, and its own header referred to a "kernel hook (MlasLayerNormBF16)" that does not exist and is not part of this PR. It tested standalone BFloat16 rounding arithmetic while the change here is CPU EP registration, with compute going through layer_norm_impl.cc. Those tests were real C++ tests, but they exercised nothing this PR touches, so describing them as coverage for it was wrong. Coverage now matches what is actually registered: 17 CPU EP operator tests spanning core LayerNormalization opset 17, contrib LayerNormalization opset 1-16, SimplifiedLayerNormalization, SkipLayerNormalization and SkipSimplifiedLayerNormalization. Tolerances are split by output type rather than applying one blanket value: BFloat16 Y is checked at 2 bf16 ULP, while the float Mean and InvStdDev are checked at 1e-5. That distinction is what makes the stat tests meaningful - the previous behaviour was off by about 780x that tolerance, so these tests fail against the pre-fix code. Regenerating docs/OperatorKernels.md requires built Python bindings, which was not practical here, so the five affected rows were hand-edited to match the existing format. Flagging that explicitly rather than passing it off as generated. Also removed internal working notes that had been committed by mistake: comments naming internal reviewers, replaced with the substance they were citing. Build is clean with warnings as errors. 17 BFloat16 operator tests pass, and 96 across the whole LayerNorm suite. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The two narrow-float conversion helpers were duplicated across layer_norm_impl.cc and skip_layer_norm.cc. Move them into a new header core/util/narrow_float_utils.h and include it from both sites. No logic change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
71bc68a to
7cf4ac9
Compare
Coverage:
- LayerNorm17_PrePack_ScaleBiasInitializers: scale/bias as constant
initializers, exercising the PrePack bf16→f32 conversion at session init.
- SkipLayerNorm_PrePack_GammaBetaInitializers: gamma/beta as initializers
for the skip variant.
- LayerNorm17_GenericBroadcast: X={2,2,2} with scale/bias={2,2} triggers
use_generic_broadcast=true, exercising ComputeJobGeneric / BFloat16Math.
Hygiene:
- Remove internal 'B5' labels from test comments (2 occurrences).
- Fix SrcDispatcher comment to accurately describe the if-constexpr
behaviour that prevents ComputeImpl<NarrowType, NarrowType> instantiation.
- Align tolerance comments with what the checker actually applies:
tolerance = absolute + relative * |expected| (numpy.isclose semantics),
where the relative component is the framework default.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…entical results
Each PrePack test now loops over is_initializer={false, true} against the
same reference output, directly proving that the prepacked code path does not
change results. SCOPED_TRACE labels failures by configuration.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 participant
What this adds
BFloat16 registration on the CPU EP for five LayerNorm/RMSNorm op families. The schema already permits
bfloat16and the CUDA EP already registers it; the CPU EP did not, so a bf16 model fell back or failed.Registered:
LayerNormalization(core, opset 17)LayerNormalization(contrib, opset 1)SimplifiedLayerNormalizationSkipLayerNormalizationSkipSimplifiedLayerNormalizationCompute widens to f32, accumulates, and narrows back — no native bf16 arithmetic is introduced, and no AVX512-BF16 or other ISA path is added.
Two changes to pre-existing behaviour, called out explicitly
1.
Mean/InvStdDevare now written at float precision.Both narrow-float
ComputeJoboverloads previously round-tripped the statistics through the narrow type:But
U— the type ofMeanandInvStdDev— is float. Quantising the stats down and back cost roughly 0.4% and, more importantly, made these overloads disagree with the generic path for identical input. Both now callWriteStat<U>directly.This changes existing MLFloat16 output, not just the new bf16 path:
Mean/InvStdDevfrom an fp16 LayerNorm are now more accurate and now match the generic path. The deadMLFloat16/BFloat16WriteStatbranches are removed, andSrcDispatcherusesif constexprsoComputeImpl<NarrowType, NarrowType>is never instantiated.2. The contrib MLFloat16 registration now declares
U = float.REGISTER_CONTRIB_KERNELS(MLFloat16)becomesREGISTER_CONTRIB_KERNELS(MLFloat16, float). The contrib schema constrainsUto float and the CUDA contrib kernels already registerU=floatfor narrow types, so this makes the CPU registration consistent rather than leaving two adjacent registrations disagreeing. It is declaration-only at runtime —SrcDispatcheralways callsComputeImpl<T, float>.If you would prefer this PR stay strictly additive, I am happy to split that registration change out.
Tests
20 CPU EP tests, covering all five registered families and:
Mean/InvStdDevprecision. Tolerances are split by output type — 2 bf16 ULP (0.016) forY, but 1e-5 for the float stats. That split is what makes them meaningful: the pre-fix error was roughly 780x the stat tolerance, so these tests fail against the pre-fix code.scale/bias(andgamma/beta) as constant initializers, which routes throughPrePack— bf16 weights converted once at session init rather than per run. Each case runs withis_initializerbothfalseandtrueagainst the same reference, so the prepacked and non-prepacked paths are asserted to agree. The implementation genuinely branches here (layer_norm_impl.cc:622).ComputeJobGeneric/BFloat16Mathrather than the fast path.106 tests pass across the wider LayerNorm suite, and upstream #31676's prepacked-length validation stays green.
Other cleanup
NarrowToFloat/FloatToNarrowwere duplicated acrosslayer_norm_impl.ccandskip_layer_norm.cc; they now live once incore/util/narrow_float_utils.h.docs/OperatorKernels.mdis updated for the new registrations. Regenerating it needs built Python bindings, which was not practical here, so the affected rows were hand-edited to match the existing format — flagging that rather than presenting them as generated.Verification
Built and run with warnings-as-errors (no
--compile_no_warning_as_error), from a clean rebuild. No performance claims are made — this PR adds no kernel and nothing was measured.