From 71b9e487c4949f327e593f5930a7a15475415cfc Mon Sep 17 00:00:00 2001 From: suharvest Date: Sat, 25 Jul 2026 20:04:40 +0800 Subject: [PATCH] fix(checkpoint): preserve destination tensor dtype Signed-off-by: suharvest --- tensorrt_edgellm/checkpoint/loader.py | 18 ++-- .../test_loader_dtype_cast.py | 93 +++++++++++++++++++ 2 files changed, 105 insertions(+), 6 deletions(-) create mode 100644 tests/python-unittests/test_loader_dtype_cast.py diff --git a/tensorrt_edgellm/checkpoint/loader.py b/tensorrt_edgellm/checkpoint/loader.py index c7d5c6409..e02a4261b 100644 --- a/tensorrt_edgellm/checkpoint/loader.py +++ b/tensorrt_edgellm/checkpoint/loader.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -406,9 +406,9 @@ def _set_tensor(model: nn.Module, mapping: Optional[Mapping] = None) -> bool: """Assign *tensor* to the buffer or parameter at *key* inside *model*. - Bfloat16 tensors are cast to float16 on the fly. The export pipeline - assumes FP16 activations and the C++ runtime requires FP16 (or FP8) - weight files. Doing the cast here avoids a separate post-loading sweep. + Floating-point checkpoint tensors are cast to an existing destination's + declared half-precision dtype. Bfloat16 tensors retain the legacy float16 + fallback when the destination dtype is unavailable. Returns True on success, False if the key does not resolve to a known buffer or parameter. @@ -420,8 +420,14 @@ def _set_tensor(model: nn.Module, except (AttributeError, IndexError, TypeError): return False - if tensor.dtype == torch.bfloat16: - tensor = tensor.to(torch.float16) + if tensor.dtype in (torch.float32, torch.bfloat16, torch.float16): + destination = module._parameters.get( + attr) if attr in module._parameters else module._buffers.get(attr) + destination_dtype = getattr(destination, "dtype", None) + if destination_dtype in (torch.float16, torch.bfloat16): + tensor = tensor.to(destination_dtype) + elif tensor.dtype == torch.bfloat16: + tensor = tensor.to(torch.float16) tensor = _shard_for_module(module, attr, tensor, mapping) diff --git a/tests/python-unittests/test_loader_dtype_cast.py b/tests/python-unittests/test_loader_dtype_cast.py new file mode 100644 index 000000000..3448ad40d --- /dev/null +++ b/tests/python-unittests/test_loader_dtype_cast.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for checkpoint destination dtype preservation.""" + +import torch +import torch.nn as nn + +from tensorrt_edgellm.checkpoint.loader import _set_tensor + + +class _TinyModel(nn.Module): + + def __init__(self): + super().__init__() + self.fp16 = nn.Linear(3, 2, dtype=torch.float16) + self.bf16 = nn.Linear(3, 2, dtype=torch.bfloat16) + self.fp32 = nn.Linear(3, 2, dtype=torch.float32) + self.register_buffer("integer", torch.zeros(2, dtype=torch.int32)) + self.attribute = torch.zeros(2, dtype=torch.float32) + + +def test_fp32_source_cast_to_declared_fp16_destination(): + model = _TinyModel() + + assert _set_tensor(model, "fp16.weight", + torch.ones_like(model.fp16.weight, dtype=torch.float32)) + + assert model.fp16.weight.dtype == torch.float16 + + +def test_fp32_source_cast_to_declared_bf16_destination(): + model = _TinyModel() + + assert _set_tensor(model, "bf16.weight", + torch.ones_like(model.bf16.weight, dtype=torch.float32)) + + assert model.bf16.weight.dtype == torch.bfloat16 + + +def test_matching_half_dtype_is_preserved(): + model = _TinyModel() + + assert _set_tensor(model, "bf16.weight", + torch.ones_like(model.bf16.weight)) + + assert model.bf16.weight.dtype == torch.bfloat16 + + +def test_bf16_source_cast_to_declared_fp16_destination(): + model = _TinyModel() + + assert _set_tensor( + model, "fp16.weight", + torch.ones_like(model.fp16.weight, dtype=torch.bfloat16)) + + assert model.fp16.weight.dtype == torch.float16 + + +def test_fp32_destination_keeps_fp32_source(): + model = _TinyModel() + + assert _set_tensor(model, "fp32.weight", + torch.ones_like(model.fp32.weight)) + + assert model.fp32.weight.dtype == torch.float32 + + +def test_unknown_destination_retains_legacy_bf16_fallback(): + model = _TinyModel() + + assert _set_tensor(model, "attribute", torch.ones(2, dtype=torch.bfloat16)) + + assert model.attribute.dtype == torch.float16 + + +def test_non_floating_source_is_not_cast(): + model = _TinyModel() + + assert _set_tensor(model, "integer", torch.ones(2, dtype=torch.int64)) + + assert model.integer.dtype == torch.int64