From f3e882b3842df1aa0ea7610c0b5e2116430d5523 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 4 May 2026 21:41:42 +0000 Subject: [PATCH 1/2] Warn about safetensors CUDA alignment issue (#2120) 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 --- src/mobius/__main__.py | 10 ++++++++++ src/mobius/_model_package.py | 8 ++++++++ 2 files changed, 18 insertions(+) diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index b8dc10c16..d4dda2f4f 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -214,6 +214,16 @@ def _save_package( _apply_optimize(model, optimize) max_shard_size_bytes = _parse_size(args.max_shard_size) if args.max_shard_size else None + + if args.external_data == "safetensors" and getattr(args, "ep", "default") == "cuda": + import logging + + logging.getLogger(__name__).warning( + "Safetensors external data does not guarantee 256-byte offset " + "alignment, which can cause CUBLAS misaligned address errors on " + "CUDA. Consider using --external-data onnx for CUDA builds." + ) + pkg.save( output_dir, external_data=args.external_data, diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 85274eb97..7c099296f 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -78,6 +78,14 @@ def save( external_data: External data format. ``"onnx"`` (default) saves weights to ``model.onnx.data``. ``"safetensors"`` saves weights in safetensors format. + + .. warning:: + The safetensors format does not guarantee 256-byte offset + alignment for tensor data within the file. This can cause + ``CUBLAS_STATUS_INVALID_VALUE`` ("misaligned address") + errors on some CUDA/cuBLAS versions when loading weights + via memory-mapped I/O. Use ``"onnx"`` (the default) for + models targeting CUDA execution. max_shard_size_bytes: Maximum shard size in bytes for safetensors format. Only used when *external_data* is ``"safetensors"``. components: Optional predicate ``(name) -> bool`` that selects From 1f70e26c4b919cde4b59e95fbf42ae97e4c207ea Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Mon, 4 May 2026 15:41:55 -0700 Subject: [PATCH 2/2] Potential fix for pull request finding 'Module is imported more than once' Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> Signed-off-by: Justin Chu --- src/mobius/__main__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index d4dda2f4f..706d448b5 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -216,8 +216,6 @@ def _save_package( max_shard_size_bytes = _parse_size(args.max_shard_size) if args.max_shard_size else None if args.external_data == "safetensors" and getattr(args, "ep", "default") == "cuda": - import logging - logging.getLogger(__name__).warning( "Safetensors external data does not guarantee 256-byte offset " "alignment, which can cause CUBLAS misaligned address errors on "