Warn about safetensors CUDA alignment issue - #248
Merged
Conversation
Safetensors format does not guarantee 256-byte offset alignment for
tensor data within the file. The header size (variable, typically
~34KB) plus the 8-byte length prefix creates a data start offset
that is not aligned to 256 bytes, causing ALL subsequent tensor
offsets to be misaligned. This can trigger CUBLAS_STATUS_INVALID_VALUE
('misaligned address') errors on CUDA when weights are loaded via
memory-mapped I/O.
ONNX external data format (.onnx.data) uses onnx_ir's alignment
logic which ensures 256-byte alignment for all tensors.
Changes:
- Add warning in ModelPackage.save() docstring about safetensors
alignment limitation
- Add runtime warning in CLI when --external-data safetensors is
used with --ep cuda
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
…once' Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
37 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add documentation and runtime warning about safetensors external data format causing CUBLAS misaligned address errors on CUDA.
Root Cause
Safetensors header size (variable, ~34KB) + 8-byte length prefix creates a data start offset that is not 256-byte aligned. ALL subsequent tensor offsets inherit this misalignment. When ORT memory-maps the safetensors file and passes the pointer to cuBLAS, the misaligned address triggers
CUBLAS_STATUS_INVALID_VALUE.Verified on gemma-4-e2b-it:
Changes
ModelPackage.save()docstring about safetensors alignment--external-data safetensors+--ep cudaRecommendation
Use
--external-data onnx(the default) for CUDA builds. The safetensors format is fine for CPU and for inference viacudaMemcpy(which handles alignment internally), but can fail with mmap-based loading paths.Ref: onnxruntime/mobius#2120