Skip to content

feat: Allow serializing and deserializing with Tensorizer without passing --model-loader-extra-config#8

Merged
sangstar merged 11 commits intosangstar/tensorizer-aws-update-and-any-kwargsfrom
sangstar/tensorizer-fix-aws-and-make-cli-streamlined
Jun 12, 2025
Merged

feat: Allow serializing and deserializing with Tensorizer without passing --model-loader-extra-config#8
sangstar merged 11 commits intosangstar/tensorizer-aws-update-and-any-kwargsfrom
sangstar/tensorizer-fix-aws-and-make-cli-streamlined

Conversation

@sangstar
Copy link
Collaborator

@sangstar sangstar commented Jun 5, 2025

Serializing and deserializing with Tensorizer just with AWS config/credential files or env vars

This PR streamlines Tensorizer usage patterns with vLLM by allowing Tensorizer to handle model loading (and saving) from vLLM just by passing load_format=tensorizer.

As a basic demonstrating of what this looks like, the targets being run in this Makefile demonstrate the ways Tensorizer can be used to now save and load models without passing things like --model-loader-extra-config JSON strings.

SHELL := /bin/bash
PYTHON := $(shell python3 -c 'import sys; print(sys.executable)')
SERIALIZATION_FP := /vllm-workspace/examples/others/tensorize_vllm_model.py
WITH_PDB := -m pdb
AWS_ENDPOINT_URL := https://cwobject.com/
ENV := AWS_ACCESS_KEY_ID=$(AWS_ACCESS_KEY_ID) AWS_SECRET_ACCESS_KEY=$(AWS_SECRET_ACCESS_KEY) AWS_ENDPOINT_URL=$(AWS_ENDPOINT_URL)

.PHONY: test_serialize_env_vars test_serialize_aws_cfg test_vllm_serve_with_aws_cfg test_vllm_serve_with_aws_env_vars

test_serialize_env_vars:
	@$(ENV) $(PYTHON) $(SERIALIZATION_FP) \
		--model facebook/opt-125m \
		serialize \
		--serialized-directory s3://ml-dev-sandbox-bucket \
		--suffix v1

# Assumes aws config and creds files are in ~/.aws/* as they are conventionally
test_serialize_aws_cfg:
	$(PYTHON) $(SERIALIZATION_FP) \
			--model facebook/opt-125m \
			serialize \
			--serialized-directory s3://ml-dev-sandbox-bucket \
			--suffix v1

test_vllm_serve_with_aws_cfg:
	@$(ENV) AWS_S3_ADDRESSING_STYLE=virtual $(PYTHON) -m vllm.entrypoints.cli.main serve s3://ml-dev-sandbox-bucket/vllm/facebook/opt-125m/v1 \
		--load-format=tensorizer

test_vllm_serve_with_aws_env_vars:
	@$(ENV) AWS_S3_ADDRESSING_STYLE=virtual $(PYTHON) -m vllm.entrypoints.cli.main serve s3://ml-dev-sandbox-bucket/vllm/facebook/opt-125m/v1 \
		--load-format=tensorizer

Please hold off on any requests for changes relating to formatting. This will all be done in a later stage with vLLM's formatter.

sangstar added 3 commits June 2, 2025 16:08
…el files

Signed-off-by: Sanger Steel <sangersteel@gmail.com>
Some changes to `TensorizerConfig` have added a few
parameters that are used for convenience internally,
but are exposed as public parameters. This unnecessarily
complicates `TensorizerConfig` as it makes it seem like
these are important parameters users need to understand
and contend with to use `TensorizerConfig` with
the public-facing API. They have been made private,
so users can disregard them and have less parameters to
wrap their heads around.
Signed-off-by: Sanger Steel <sangersteel@gmail.com>
Signed-off-by: Sanger Steel <sangersteel@gmail.com>
@sangstar sangstar requested review from Eta0 and arsenetar June 5, 2025 19:28
@sangstar sangstar self-assigned this Jun 5, 2025
sangstar added 3 commits June 6, 2025 13:37
Simply call `snapshot_download` to a tempdir and serialize that to
S3 for model artifacts, completely decoupling Tensorizer from the
original machinery needed to load specific files.

Signed-off-by: Sanger Steel <sangersteel@gmail.com>
Signed-off-by: Sanger Steel <sangersteel@gmail.com>
Signed-off-by: Sanger Steel <sangersteel@gmail.com>
sangstar added 3 commits June 9, 2025 10:31
Signed-off-by: Sanger Steel <sangersteel@gmail.com>
Signed-off-by: Sanger Steel <sangersteel@gmail.com>
Since `model_loader_extra_config` can be a `TensorizerConfig` instance
as well as a dict, add a `__getitem__` method to `TensorizerConfig` and
fix checker function to work without importing `TensorizerConfig` (that
would've caused a circular import)

Signed-off-by: Sanger Steel <sangersteel@gmail.com>
@sangstar sangstar changed the base branch from main to sangstar/tensorizer-aws-update-and-any-kwargs June 11, 2025 21:49
Copy link

@Eta0 Eta0 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Requested a few changes when reviewing this in a Zoom call, will approve when those are fixed.

Signed-off-by: Sanger Steel <sangersteel@gmail.com>
Signed-off-by: Sanger Steel <sangersteel@gmail.com>
@sangstar sangstar merged commit 0a7020b into sangstar/tensorizer-aws-update-and-any-kwargs Jun 12, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants