Skip to content

fix(checkpoint): Preserve destination tensor dtype - #149

Open
suharvest wants to merge 1 commit into
NVIDIA:mainfrom
suharvest:pr/fix-checkpoint-destination-dtype
Open

fix(checkpoint): Preserve destination tensor dtype#149
suharvest wants to merge 1 commit into
NVIDIA:mainfrom
suharvest:pr/fix-checkpoint-destination-dtype

Conversation

@suharvest

Copy link
Copy Markdown

What does this PR do?

Preserves the declared half-precision dtype of an existing checkpoint
destination when loading a floating-point source tensor.

Previously _set_tensor() only converted BF16 sources through its legacy FP16
fallback. Assigning an FP32 checkpoint tensor could replace an existing FP16
or BF16 parameter with an FP32 parameter, leaving the exported model in an
unexpected mixed-dtype state.

FP32, BF16, and FP16 sources now use the existing destination dtype when that
destination is FP16 or BF16. Existing behavior remains unchanged for FP32
destinations, non-floating tensors, and attributes without a declared
parameter/buffer dtype; the legacy unknown-destination BF16-to-FP16 fallback is
retained.

Scope

  • One checkpoint-loader function.
  • One CPU-only regression-test file.
  • No model classes, quantization formats, runtime, or product-specific logic.

Validation

  • Before the fix: 3 failed, 4 passed.
  • After the fix: 7 passed.
  • The matrix covers FP32 → FP16, FP32 → BF16, matching half dtype,
    BF16 → declared FP16, FP32 destination, unknown destination legacy fallback,
    and non-floating input.
  • Tests import the production _set_tensor() function and use only standard
    PyTorch modules/parameters/buffers.
  • YAPF, isort, autoflake, codespell, compileall, git diff --check, and
    merge-tree validation pass.

The CPU test environment used PyTorch 2.11 rather than the repository-pinned
2.12 release. A complete model export → build → inference run requires the
official GPU environment and is still requested from CI/device testing.

Closes #144.

Signed-off-by: suharvest <suharvest@gmail.com>
@suharvest
suharvest requested a review from a team July 25, 2026 12:09
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.

[Bug][Export]: checkpoint loader can replace declared FP16/BF16 parameters with FP32 tensors

1 participant