Package: tensorizer
tensorizerModule, Model, and Tensor Serialization/DeserializationTLDRExtremely fast model loads from HTTP/HTTPS, Redis, and S3 endpoints.
GPT-J (20GB) loads at wire-speed (~5GB/s) on a 40GbE network, and
is only bottlenecked by the Linux kernel TCP stack.RationaleCoreWeave and our customers use KNative to deploy models as serverless
functions. How long a model takes to load is a major factor in the latency
of KNative scale-up.tensorizeris a tool to serialize models and their
associated tensors into a single file that can be loaded quickly and
efficiently off an HTTP/HTTPS or S3 endpoint.By not embedding the model in the container image, we can reduce the
container image size and the time it takes to load the model. This is
especially important for models that are large in size, such asEleutherAI/gpt-neox-20Bthat weighs in at~40GB.This decoupling of the model from the container image also allows us to
update the model without having to rebuild the container image. This allows
us to quickly iterate on the model and deploy new versions without having
to wait for the container image to build or for the container image cache
to be populated.tensorizerhas S3 support, so we can store the serialized model in S3
object storage, and perform streaming loads from S3. This allows us to
stream the model directly from S3 into the container without having to
download the model to the container's local filesystem. This also
pertains to HTTP/HTTPS endpoints, as S3 is just an HTTP/HTTPS endpoint.tensorizeralso has support for loading models from a local filesystem,
so you can use it to serialize models locally and load them locally. This
is extremely fast, as the same principles that make it fast for HTTP/HTTPS
and S3 endpoints also apply to local filesystems.tensorizerhas preliminary support for Redis, but it is not recommended
for model deployment due to the lack of distributed caching. It is intended
for sharing state between inference pods, or for loading data on a per-request
basis from a Redis cache.Speedtensorizer's deserialization speed is primarily network-bound.The following graph presents data collected from the scripts and Kubernetes
manifests inexamples/benchmark_buffer_sizecomparing the various deserialization modes available intensorizerrelease
2.5.0—along with the raw network speed, and the speed oftorch.load().InstallationFrom PyPItensorizercan be installed from PyPI withpip:python-mpipinstalltensorizerFrom SourceYou can also installtensorizerfrom source usingpip.To clone the repository and installtensorizerineditable mode,
run:gitclonehttps://github.com/coreweave/tensorizercdtensorizer
python-mpipinstall-e.Or, run the following forpipto installtensorizerdirectly from GitHub:python-mpipinstallgit+https://github.com/coreweave/tensorizerBasic UsageSerialization is done with theTensorSerializerclass. It takes apath_uriargument that can be a local filesystem path, an HTTP/HTTPS
endpoint, or an S3 endpoint.write_moduleis the main method of theTensorSerializerclass. It
takes atorch.nn.Moduleand serializes the tensors to thepath_uriendpoint.The below example serializes theEleutherAI/gpt-j-6Bmodel to an S3
endpoint. It assumes that you have already configured your S3
credentials in~/.s3cfg.NOTE:Loading and serializinggpt-j-6Bwill take a lot of CPU RAM,
up to~20GB. Additionally, when loadinggpt-j-6Binto a GPU, you
will need about~16GBof VRAM. If you don't have that much RAM or VRAM,
you can use the smallergpt-neo-125Mmodel instead.NOTE2:The below examples require thetransformersandacceleratelibraries. You can install them withpip:python-mpipinstalltransformersaccelerateserialize.pyimporttorchfromtensorizerimportTensorSerializerfromtransformersimportAutoModelForCausalLMmodel_ref="EleutherAI/gpt-j-6B"# For less intensive requirements, swap above with the line below:# model_ref = "EleutherAI/gpt-neo-125M"model_name=model_ref.split("/")[-1]# Change this to your S3 bucket.s3_bucket="bucket"s3_uri=f"s3://{s3_bucket}/{model_name}.tensors"model=AutoModelForCausalLM.from_pretrained(model_ref,revision="float16",torch_dtype=torch.float16,low_cpu_mem_usage=True,)serializer=TensorSerializer(s3_uri)serializer.write_module(model)serializer.close()Conversely, deserialization is done with theTensorDeserializerclass.
It takes apath_uriargument that can be a local filesystem path, an
HTTP/HTTPS endpoint, or an S3 endpoint.load_into_moduleis the main method of theTensorDeserializerclass.
It takes atorch.nn.Moduleand loads the tensors from thepath_uriendpoint into thetorch.nn.Module.The below example loads theEleutherAI/gpt-j-6Bmodel from an S3
endpoint.deserialize.pyimporttimeimporttorchfromtensorizerimportTensorDeserializerfromtensorizer.utilsimportno_init_or_tensor,convert_bytes,get_mem_usagefromtransformersimportAutoModelForCausalLM,AutoTokenizer,AutoConfigmodel_ref="EleutherAI/gpt-j-6B"# To run this at home, swap this with the line below for a smaller example:# model_ref = "EleutherAI/gpt-neo-125M"model_name=model_ref.split("/")[-1]# Change this to your S3 bucket.s3_bucket="bucket"s3_uri=f"s3://{s3_bucket}/{model_name}.tensors"config=AutoConfig.from_pretrained(model_ref)# This ensures that the model is not initialized.withno_init_or_tensor():model=AutoModelForCausalLM.from_config(config)before_mem=get_mem_usage()# Lazy load the tensors from S3 into the model.start=time.time()deserializer=TensorDeserializer(s3_uri)deserializer.load_into_module(model)end=time.time()# Brag about how fast we are.total_bytes_str=convert_bytes(deserializer.total_tensor_bytes)duration=end-startper_second=convert_bytes(deserializer.total_tensor_bytes/duration)after_mem=get_mem_usage()deserializer.close()print(f"Deserialized{total_bytes_str}in{end-start:0.2f}s,{per_second}/s")print(f"Memory usage before:{before_mem}")print(f"Memory usage after:{after_mem}")# Tokenize and generatemodel.eval()tokenizer=AutoTokenizer.from_pretrained(model_ref)eos=tokenizer.eos_token_idinput_ids=tokenizer.encode("¡Hola! Encantado de conocerte. hoy voy a",return_tensors="pt").to("cuda")withtorch.no_grad():output=model.generate(input_ids,max_new_tokens=50,do_sample=True,pad_token_id=eos)print(f"Output:{tokenizer.decode(output[0],skip_special_tokens=True)}")It should produce output similar to the following, with GPT-J-6B:Deserialized model in 6.25 seconds
Test Output: ¡Hola! Encantado de conocerte. hoy voy a comentar por primera
vez una teoría de trineo, que quizá te parezca
algo desconocido, ya que en este mundo han
llegado a dominar tantosMore practical examples for the usage oftensorizercan be found inexamples/hf_serialization.py,
wheredf_main()serializes models fromHuggingFace Diffusersandhf_main()serializesHuggingFace Transformersmodels.Tensor Weight Encryptiontensorizersupports fast tensor weight encryption and decryption during
serialization and deserialization, respectively.Be aware that metadata (tensor names, dtypes, shapes, etc.) are not encrypted,
only the weights themselves.[!NOTE]Refer todocs/encryption.mdfor details, instructions,
and warnings on usingtensorizerencryption correctly and safely.To usetensorizerencryption, a recent version oflibsodiummust be
installed. Installlibsodiumwithapt-get install libsodium23on Ubuntu or Debian, or followthe instructions inlibsodium's documentationfor other platforms.Quick Encryption ExampleThe following outline demonstrates how to encrypt and decrypt a tensorized model
with a randomly-generated encryption key:fromtensorizerimport(EncryptionParams,DecryptionParams,TensorDeserializer,TensorSerializer)# Serialize and encrypt a model:encryption_params=EncryptionParams.random()serializer=TensorSerializer("model.tensors",encryption=encryption_params)serializer.write_module(...)# or write_state_dict(), etc.serializer.close()# Save the randomly-generated encryption key somewherewithopen("tensor.key","wb")askey_file:key_file.write(encryption_params.key)# Then decrypt it again:# Load the randomly-generated key from where it was savedwithopen("tensor.key","rb")askey_file:key:bytes=key_file.read()decryption_params=DecryptionParams.from_key(key)deserializer=TensorDeserializer("model.tensors",encryption=decryption_params)deserializer.load_into_module(...)deserializer.close()For more detail, refer todocs/encryption.md.
A complete example is also available asexamples/encryption.py.
TheEncryptionParamsandDecryptionParamsclass docstrings additionally
contain some usage information for quick reference from an IDE.An example command line tool to add or remove encryption from existing
serialized models is also available asexamples/encryption.py.BenchmarksYou can run your own benchmarks on CoreWeave or your own Kubernetes cluster
by using thebenchmark.yamlfile in theexamples/benchmark_buffer_sizedirectory. Please see theREADME.Available Pre-Tensorized Models on the CoreWeave CloudThe following models are available on the CoreWeave Cloud for free, and can be
used with theTensorDeserializerclass. The S3 support defaults to theaccel-object.ord1.coreweave.comendpoint, and the bucket to use astensorized.We name the keys in the S3 bucket after the HuggingFace model identifier, and
append the/fp16suffix for the half-precision version.For example, the S3 URI for theEleutherAI/gpt-j-6Bmodel is:s3://tensorized/EleutherAI/gpt-j-6B/fp16/model.tensorsThe below table shows the available models and their S3 URIs.Large Language ModelsModelPrecisionS3 URIEleutherAI/gpt-neo-125Mfp32s3://tensorized/EleutherAI/gpt-neo-125M/model.tensorsEleutherAI/gpt-neo-125Mfp16s3://tensorized/EleutherAI/gpt-neo-125M/fp16/model.tensorsEleutherAI/gpt-neo-1.3Bfp32s3://tensorized/EleutherAI/gpt-neo-1.3B/model.tensorsEleutherAI/gpt-neo-1.3Bfp16s3://tensorized/EleutherAI/gpt-neo-1.3B/fp16/model.tensorsEleutherAI/gpt-neo-2.7Bfp32s3://tensorized/EleutherAI/gpt-neo-2.7B/model.tensorsEleutherAI/gpt-neo-2.7Bfp16s3://tensorized/EleutherAI/gpt-neo-2.7B/fp16/model.tensorsEleutherAI/gpt-j-6Bfp32s3://tensorized/EleutherAI/gpt-j-6B/model.tensorsEleutherAI/gpt-j-6Bfp16s3://tensorized/EleutherAI/gpt-j-6B/fp16/model.tensorsEleutherAI/gpt-neox-20bfp32s3://tensorized/EleutherAI/gpt-neox-20b/model.tensorsEleutherAI/gpt-neox-20bfp16s3://tensorized/EleutherAI/gpt-neox-20b/fp16/model.tensorsEleutherAI/pythia-70mfp32s3://tensorized/EleutherAI/pythia-70m/model.tensorsEleutherAI/pythia-70mfp16s3://tensorized/EleutherAI/pythia-70m/fp16/model.tensorsEleutherAI/pythia-1.4bfp32s3://tensorized/EleutherAI/pythia-1.4b/model.tensorsEleutherAI/pythia-1.4bfp16s3://tensorized/EleutherAI/pythia-1.4b/fp16/model.tensorsEleutherAI/pythia-2.8bfp32s3://tensorized/EleutherAI/pythia-2.8b/model.tensorsEleutherAI/pythia-2.8bfp16s3://tensorized/EleutherAI/pythia-2.8b/fp16/model.tensorsEleutherAI/pythia-6.9bfp32s3://tensorized/EleutherAI/pythia-6.9b/model.tensorsEleutherAI/pythia-6.9bfp16s3://tensorized/EleutherAI/pythia-6.9b/fp16/model.tensorsEleutherAI/pythia-12bfp32s3://tensorized/EleutherAI/pythia-12b/model.tensorsEleutherAI/pythia-12bfp16s3://tensorized/EleutherAI/pythia-12b/fp16/model.tensorsEleutherAI/pythia-70m-dedupedfp32s3://tensorized/EleutherAI/pythia-70m-deduped/model.tensorsEleutherAI/pythia-70m-dedupedfp16s3://tensorized/EleutherAI/pythia-70m-deduped/fp16/model.tensorsEleutherAI/pythia-1.4b-dedupedfp32s3://tensorized/EleutherAI/pythia-1.4b-deduped/model.tensorsEleutherAI/pythia-1.4b-dedupedfp16s3://tensorized/EleutherAI/pythia-1.4b-deduped/fp16/model.tensorsEleutherAI/pythia-2.8b-dedupedfp32s3://tensorized/EleutherAI/pythia-2.8b-deduped/model.tensorsEleutherAI/pythia-2.8b-dedupedfp16s3://tensorized/EleutherAI/pythia-2.8b-deduped/fp16/model.tensorsEleutherAI/pythia-6.9b-dedupedfp32s3://tensorized/EleutherAI/pythia-6.9b-deduped/model.tensorsEleutherAI/pythia-6.9b-dedupedfp16s3://tensorized/EleutherAI/pythia-6.9b-deduped/fp16/model.tensorsEleutherAI/pythia-12b-dedupedfp32s3://tensorized/EleutherAI/pythia-12b-deduped/model.tensorsEleutherAI/pythia-12b-dedupedfp16s3://tensorized/EleutherAI/pythia-12b-deduped/fp16/model.tensorsKoboldAI/fairseq-dense-125Mfp32s3://tensorized/KoboldAI/fairseq-dense-125M/model.tensorsKoboldAI/fairseq-dense-125Mfp16s3://tensorized/KoboldAI/fairseq-dense-125M/fp16/model.tensorsKoboldAI/fairseq-dense-355Mfp32s3://tensorized/KoboldAI/fairseq-dense-355M/model.tensorsKoboldAI/fairseq-dense-355Mfp16s3://tensorized/KoboldAI/fairseq-dense-355M/fp16/model.tensorsKoboldAI/fairseq-dense-2.7Bfp32s3://tensorized/KoboldAI/fairseq-dense-2.7B/model.tensorsKoboldAI/fairseq-dense-2.7Bfp16s3://tensorized/KoboldAI/fairseq-dense-2.7B/fp16/model.tensorsKoboldAI/fairseq-dense-6.7Bfp32s3://tensorized/KoboldAI/fairseq-dense-6.7B/model.tensorsKoboldAI/fairseq-dense-6.7Bfp16s3://tensorized/KoboldAI/fairseq-dense-6.7B/fp16/model.tensorsKoboldAI/fairseq-dense-13Bfp32s3://tensorized/KoboldAI/fairseq-dense-13B/model.tensorsKoboldAI/fairseq-dense-13Bfp16s3://tensorized/KoboldAI/fairseq-dense-13B/fp16/model.tensorsSalesforce/codegen-350M-monofp32s3://tensorized/Salesforce/codegen-350M-mono/model.tensorsSalesforce/codegen-350M-monofp16s3://tensorized/Salesforce/codegen-350M-mono/fp16/model.tensorsSalesforce/codegen-350M-multifp32s3://tensorized/Salesforce/codegen-350M-multi/model.tensorsSalesforce/codegen-350M-multifp16s3://tensorized/Salesforce/codegen-350M-multi/fp16/model.tensorsSalesforce/codegen-2B-multifp32s3://tensorized/Salesforce/codegen-2B-multi/model.tensorsSalesforce/codegen-2B-multifp16s3://tensorized/Salesforce/codegen-2B-multi/fp16/model.tensorsSalesforce/codegen-6B-monofp32s3://tensorized/Salesforce/codegen-6B-mono/model.tensorsSalesforce/codegen-6B-monofp16s3://tensorized/Salesforce/codegen-6B-mono/fp16/model.tensorsSalesforce/codegen-6B-multifp32s3://tensorized/Salesforce/codegen-6B-multi/model.tensorsSalesforce/codegen-6B-multifp16s3://tensorized/Salesforce/codegen-6B-multi/fp16/model.tensorsSalesforce/codegen-16B-monofp32s3://tensorized/Salesforce/codegen-16B-mono/model.tensorsSalesforce/codegen-16B-monofp16s3://tensorized/Salesforce/codegen-16B-mono/fp16/model.tensorsSalesforce/codegen-16B-multifp32s3://tensorized/Salesforce/codegen-16B-multi/model.tensorsSalesforce/codegen-16B-multifp16s3://tensorized/Salesforce/codegen-16B-multi/fp16/model.tensorsGenerative Diffusion ModelsModelComponentPrecisionS3 URIRunwayML/stable-diffusion-v1-5VAEfp32s3://tensorized/runwayml/stable-diffusion-v1-5/vae.tensorsRunwayML/stable-diffusion-v1-5UNetfp32s3://tensorized/runwayml/stable-diffusion-v1-5/unet.tensorsRunwayML/stable-diffusion-v1-5TextEncfp32s3://tensorized/runwayml/stable-diffusion-v1-5/text_encoder.tensorsRunwayML/stable-diffusion-v1-5VAEfp16s3://tensorized/runwayml/stable-diffusion-v1-5/fp16/vae.tensorsRunwayML/stable-diffusion-v1-5UNetfp16s3://tensorized/runwayml/stable-diffusion-v1-5/fp16/unet.tensorsRunwayML/stable-diffusion-v1-5TextEncfp16s3://tensorized/runwayml/stable-diffusion-v1-5/fp16/text_encoder.tensorsStabilityAI/stable-diffusion-2-1VAEfp32s3://tensorized/stabilityai/stable-diffusion-2-1/vae.tensorsStabilityAI/stable-diffusion-2-1UNetfp32s3://tensorized/stabilityai/stable-diffusion-2-1/unet.tensorsStabilityAI/stable-diffusion-2-1TextEncfp32s3://tensorized/stabilityai/stable-diffusion-2-1/text_encoder.tensorsStabilityAI/stable-diffusion-2-1VAEfp16s3://tensorized/stabilityai/stable-diffusion-2-1/fp16/vae.tensorsStabilityAI/stable-diffusion-2-1UNetfp16s3://tensorized/stabilityai/stable-diffusion-2-1/fp16/unet.tensorsStabilityAI/stable-diffusion-2-1TextEncfp16s3://tensorized/stabilityai/stable-diffusion-2-1/fp16/text_encoder.tensorsStabilityAI/stable-diffusion-xl-base-1.0VAEfp32s3://tensorized/stabilityai/stable-diffusion-xl-base-1.0/vae.tensorsStabilityAI/stable-diffusion-xl-base-1.0UNetfp32s3://tensorized/stabilityai/stable-diffusion-xl-base-1.0/unet.tensorsStabilityAI/stable-diffusion-xl-base-1.0TextEncfp32s3://tensorized/stabilityai/stable-diffusion-xl-base-1.0/text_encoder.tensorsStabilityAI/stable-diffusion-xl-base-1.0TextEnc2fp32s3://tensorized/stabilityai/stable-diffusion-xl-base-1.0/text_encoder_2.tensorsStabilityAI/stable-diffusion-xl-base-1.0VAEfp16s3://tensorized/stabilityai/stable-diffusion-xl-base-1.0/fp16/vae.tensorsStabilityAI/stable-diffusion-xl-base-1.0UNetfp16s3://tensorized/stabilityai/stable-diffusion-xl-base-1.0/fp16/unet.tensorsStabilityAI/stable-diffusion-xl-base-1.0TextEncfp16s3://tensorized/stabilityai/stable-diffusion-xl-base-1.0/fp16/text_encoder.tensorsStabilityAI/stable-diffusion-xl-base-1.0TextEnc2fp16s3://tensorized/stabilityai/stable-diffusion-xl-base-1.0/fp16/text_encoder_2.tensorsS3 Usage Notestensorizeruses theboto3library to interact with S3. The easiest way
to usetensorizerwith S3 is to configure your S3 credentials in~/.s3cfg.If you don't want to use~/.s3cfg, or wish to use a.s3cfgconfig file
saved at a nonstandard location (e.g. under/var/run), you can also specify
your S3 credentials using thetensorizer.stream_io.open_stream()function,
and then pass that into theTensorSerializerorTensorDeserializerconstructor.Thestream_io.open_stream()function takes apath_uriargument, which can
be ans3://URI, and accepts the following keyword arguments:s3_access_key_id: S3 access key IDs3_secret_access_key: S3 secret access keys3_endpoint: S3 endpointOr,s3_config_path: Alternative filesystem path to a.s3cfgconfig fileFor example:TensorSerializer(open_stream(s3_uri,"wb",s3_access_key_id=ACCESS_KEY,s3_secret_access_key=SECRET_KEY,s3_endpoint="object.ord1.coreweave.com"))and...TensorDeserializer(open_stream(s3_uri,"rb",s3_access_key_id=ACCESS_KEY,s3_secret_access_key=SECRET_KEY,s3_endpoint="object.ord1.coreweave.com"))NOTE:For faster object downloads in the CoreWeave Cloud, you can use
theaccel-object.ord1.coreweave.comendpoint. This endpoint is optimized
for object downloads, and will be faster than theobject.ord1.coreweave.comendpoint once the object is cached.NOTE2:The cache above does not get invalidated when the object is updated
in S3. If you update an object in S3, you will need to wait for the cache to
expire before you can download the updated object. This takes 24 hours since
the last download.For this reason, it is recommended to use a unique S3 key for each version
of a model if you use theaccel-object.ord1.coreweave.comendpoint.Additional Featurestensorizerhas a few additional features that make it more useful than
just a serialization/deserialization tool.Plaid Modetensorizerhas aplaid_modeargument that can be passed to theTensorDeserializerclass when deserializing tensors to the GPU.
Whenplaid_modeisTrue,tensorizerwill load tensors extremely fast
by reusing a single small CPU buffer as a staging area to load tensors before
transferring them to the GPU through DMA.
This can also greatly decrease CPU RAM usage during deserialization.plaid_modeis enabled by default when available.state_dictSupportTheTensorDeserializerobject can be used as-is as astate_dictfortorch.nn.Module.load_state_dict. This is useful for loading the tensors
into atorch.nn.Modulethat is already initialized, or for inspection.Keep in mind thatload_state_dictis not a fast operation, and will
likely be much slower thanload_into_module.Thestate_dictcan also be used to initialize a HuggingFace Transformers
AutoModel. But HuggingFace Transformers performs three or more copies of
the data, so memory use will explode.bfloat16SupportTensorizer supports models using thebfloat16data type. However, tensorizer
uses numpy to save the tensors as binary and numpy doesn't supportbfloat16.
This means that special conversions need to be applied.To be saved, the torch tensor is cast toint16before being converted to
numpy, which doesn't change any of the underlying data. When serialized, the
originalbfloat16datatype string is also saved so that it will be cast back
tobfloat16during the deserialization process.Thecomplex32datatype is supported in a similar way, by casting toint32.
The quantized datatypes (qint8,qint32, etc.) are not currently supported
by tensorizer as they would require supplemental quantization parameters to be
deserialized correctly.NOTE:The exact choice of intermediate types asint16andint32is
considered an implementation detail, and is subject to change,
so they should not be relied upon.NOTE2:This does not interfere with storing actualintdatatypes
used in tensors in tensorized files.Numpy SupportTensorizer can be used withnumpydirectly to read and writenumpy.ndarrays.The serializer'swrite_tensorfunction handles supplying bothtorch.Tensors andnumpy.ndarrays.The deserializer has a separate functionread_numpy_arraysthat will return
the data asnumpy.ndarrays.As explained above inbfloat16 support, tensorizer uses
special conversions to write "opaque" datatypes, those not supported by numpy.
Therefore, special considerations need to be taken when loading such data asnumpy.ndarrays.By default, theTensorDeserializer.read_numpy_arraysfunction sets itsallow_raw_dataparameter toFalse. This means that if a file contains
opaque datatypes, aValueErrorwill be raised during deserialization.If you want to return the raw data regardless, setallow_raw_datatoTrue.
Otherwise, the file may be read withTensorDeserializer.read_tensorsinstead, which yieldstorch.Tensorobjects of the correct datatype.A fifth and sixth variable are also returned by theread_numpy_arraysgenerator. The fifth is aboolthat indicates whether the returned array
has an opaque datatype and requires special handling (only legal whenallow_raw_data=True). The sixth is a string describing the true, non-numpy
datatype that the raw data should be interpreted as in such cases.
For all other datatypes that require no special handling, these are returned asFalseandNone, respectively.
The exact numpy datatypes used by the returned opaquenumpy.ndarrayobjects
is not guaranteed, and should not be relied upon.Running Teststensorizerusesunittestfor testing.
The tests have their own set of dependencies, which can be installed withpip install -r tests/requirements.txt.Some tests require a GPU, and will be skipped if no GPU is available.
To run the tests, run the following in the root of the repository:python-mpipinstall-e.
python-mpipinstall-rtests/requirements.txt
python-munittestdiscovertests/--verbose