diff --git a/cosmos_framework/inference/offloading_storage.py b/cosmos_framework/inference/offloading_storage.py index 51d81942..7dc3a7e3 100644 --- a/cosmos_framework/inference/offloading_storage.py +++ b/cosmos_framework/inference/offloading_storage.py @@ -90,6 +90,9 @@ def __init__(self, tensors: list[torch.Tensor]): for tensor in physical: entries = self.views.setdefault(tensor.dtype, []) offset = entries[-1].offset + entries[-1].span if entries else 0 + # Work around cuBLAS 12.9/13.0 requiring 16-byte-aligned FP8 scalar scales. + alignment = max(1, 16 // tensor.element_size()) + offset = (offset + alignment - 1) // alignment * alignment span = 1 + sum((size - 1) * stride for size, stride in zip(tensor.shape, tensor.stride(), strict=True)) entries.append(_WeightView(tensor, offset, span if tensor.numel() else 0, tensor.shape, tensor.stride())) # Complete the copies before changing any module storage. diff --git a/cosmos_framework/inference/offloading_test.py b/cosmos_framework/inference/offloading_test.py index 2960c4d1..153de9b4 100644 --- a/cosmos_framework/inference/offloading_test.py +++ b/cosmos_framework/inference/offloading_test.py @@ -174,6 +174,23 @@ def init_weights(self, buffer_device): assert getattr(weight, part).device.type == "cpu" +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float8_e4m3fn]) +def test_weight_binding_preserves_alignment_and_values(dtype): + # Scalar scales and odd-sized payloads must stay aligned after packing. + tensors = [torch.arange(size, dtype=torch.float32).to(dtype) for size in (1, 1, 3, 5)] + tensors.append(torch.arange(6, dtype=torch.float32).reshape(2, 3).t().to(dtype)) + expected = [tensor.clone() for tensor in tensors] + strides = [tensor.stride() for tensor in tensors] + weights = _LayerWeights(tensors) + for device in ("cpu", "cuda", "cpu"): + backing = {dtype: host.to(device) for dtype, host in weights.host.items()} + weights.bind(backing) + for tensor, reference, stride in zip(tensors, expected, strides, strict=True): + assert tensor.data_ptr() % 16 == 0 + assert tensor.stride() == stride + torch.testing.assert_close(tensor.cpu().float(), reference.float(), rtol=0, atol=0) + + @pytest.mark.parametrize("inference_parameter", [False, True]) def test_weight_binding_preserves_inference_status(inference_parameter): with torch.inference_mode(inference_parameter):