Skip to content

fix(encoders): preserve BF16 constants - #1937

Draft
JiaxinD wants to merge 1 commit into
NVIDIA:mainfrom
JiaxinD:fix/encoders-bf16-constants
Draft

JiaxinD wants to merge 1 commit into
NVIDIA:mainfrom
JiaxinD:fix/encoders-bf16-constants

Conversation

@JiaxinD

@JiaxinD JiaxinD commented Nov 2, 2026 •

Copy link
Copy Markdown
Contributor

Background

DistilBERT, ELECTRA and MPNet accept BF16 builds, but their family-owned encoder builders store embedding and mask constants in FP16. This loses BF16-range values before TensorRT sees them and mixes FP16 constants with BF16 tensors in strongly typed mask arithmetic. Reproduced against current upstream; no originating issue.

Exit Criteria

Preserve BF16 constant range and requested graph types, keep the FP32 output boundary, and cover each family's existing official checkpoint with an unchanged FP32 CLS oracle. The scoped official FP16/BF16 engine/reference checks now pass in public Dev CI.

Implementation

  • Use FP32 host storage for BF16 in the three encoder builders.
  • Cast embedding tables and mask constants through each family's existing helper; FP16/FP32 type choices remain unchanged.
  • Add family-local tests executing production builder and graph helpers with a typed CPU interpreter: sequence lengths 1/3, padding, BF16-range values, FP32-layer overrides, ELECTRA embedding projection and MPNet relative bias.
  • Add one BF16 variant of each existing single-GPU manifest, retaining its checkpoint, prompt, FP32 reference and 0.999 CLS cosine threshold.

Changes are confined to families/distilbert/**, families/electra/** and families/mpnet/**; no cross-family imports or shared implementation. No public API, ABI, dependency, bundle-schema or runtime-source change.

Change categories

  • Model or runtime behavior
  • Public API
  • ABI
  • Bundle or artifact format
  • Dependencies
  • Documentation only
  • CI or developer tooling

Validation

Commands and Results

Command Result
python -m pytest families/{distilbert,electra,mpnet}/tests/test_bf16_constants.py -q -p no:cacheprovider Baseline: 12 BF16 failures, 24 controls pass; fixed: 36 pass. Independent reviewer reproduced both.
python -m pytest families/distilbert/tests families/electra/tests families/mpnet/tests tools/tests/test_architecture.py tools/tests/test_family_impact.py tools/tests/test_pr_metadata.py tools/tests/test_public_source_hygiene.py -q -p no:cacheprovider Linux: 163 passed, 9 GPU E2E skipped; one existing Dash deprecation warning
Public Dev GPU run41376122006 Six requested/executed/passed FP16/BF16 cases; zero failed/skipped; family E2E pytest totals 2/2/3 passed
Ruff check on six changed Python files; format check on three new tests; python -m tools.model_ci validate; python tools/test_impact.py --validate; python -m tools.legal_headers --check; git diff --check Passed

Hardware, Environment, and Revisions

Head e51a9041d33cd9e613dabb61659fac1f92dd158a; upstream 2674c813ad927e7cccdeb00347d4bae61cb4c668. Linux CPU/Python3.12, NumPy/ml_dtypes; focused tests also ran on Windows/Python3.13. Public Dev ran merge ca4b6be7, whose second parent is this head.

Existing official checkpoint Pinned revision
distilbert/distilbert-base-uncased 12040accade4e8a0f71eabdb258fecc2e7e948be
google/electra-base-discriminator 1ae76a97c7e84a4e640876a07453fccd636f0667
sentence-transformers/all-mpnet-base-v2 e8c3b32edf5434bc2275fc9bab85f82640a19130

Not Run / Remaining Gaps

No local GPU execution, TP/multicard, performance or peak-memory proof. Radix authentication remains expired. Public Dev validates the scoped engine/reference cases; the CPU interpreter alone does not. Existing builder formatting fails on both baseline and head; required Ruff checks pass, without unrelated reformatting.

Contributor Self-Review

  • I have completed a self-review of this change.

Independent review found no code blocker.

Notes For Future Readers

Rebuild BF16 bundles. Official model configs and Apache-2.0 model metadata were inspected; no weights or third-party implementation were added. Raw CLS features are not normalized sentence embeddings. Existing GELU dispatch is unchanged, tracked separately in #1801. Draft now remains only for the existing Ready-capacity restriction.

Risk level

  • Low
  • Medium
  • High

BF16 constant storage and graph casts change across three independently owned families; scoped public Dev qualification passes, broader hardware and performance remain unqualified.

Signed-off-by: JiaxinD <djx2048@gmail.com>
@coderabbitai

coderabbitai Bot commented Nov 2, 2026

Copy link
Copy Markdown

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

  • Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

@JiaxinD

JiaxinD commented Nov 2, 2026

Copy link
Copy Markdown
Contributor Author

@yifeif-nv Community CPU, Stable and Dev pass for e51a904, including all six official FP16/BF16 cases. Please trigger protected CI for this head.

This branch has not been deployed

No deployments
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.

1 participant