diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index b8dc10c16..706d448b5 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -214,6 +214,14 @@ 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": + 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