[Feature]: Support loading of sharded vLLM serialized models with Tensorizer

[tjohnson31415]

🚀 The feature, motivation and pitch

PR #3476 added support for loading models with Tensorizer, but has the limitation that it does not support loading a sharded vllm-serialized model to multiple GPUs (see this verification check). Use of sharded models would also benefit from the faster loading and encryption provided by Tensorizer.

This issue open with Tensorizer suggests a couple of approaches to support sharding. With tensor-parallel models, the model is split across the GPUs and the suggestion is to serialize each shard separately.

I have prototyped this approach of splitting the vllm-tensorized model into multiple shards and am working on a PR.

Alternatives

The alternative given in the Tensorizer issue is to deserializing Tensors to CPU memory and then send the tensors to the GPUs. This would decouple the serialization of the model from the sharding configuration, but would also be less efficient.

Additional context

No response

[tjohnson31415]

cc: @sangstar I'm curious what your thoughts are on supporting sharded vllm tensorized models and if this is the right approach.

[sangstar]

cc: @sangstar I'm curious what your thoughts are on supporting sharded vllm tensorized models and if this is the right approach.

Eager to see your PR on this! I'll give you a substantive reply on this soon.

[sangstar]

Working with serialized shards will be performant, but there are a couple of considerations one should make if trying to implement this with tensorizer:

• Deserialization incurs a performance cost as the number of shards increase, as apart from increased peak CPU RAM usage, the buffer setup phase for multiple deserializers can not occur perfectly concurrently (some part of the process are simply not parallelizable). As such num_shards = num_devices would be ideal, and the fewer shards the better.

• Sharded serialization to a certain num_devices expects that same num_devices during deserialization as the device count is baked into the serialization, so if num_devices changes, you'd need to reserialize.

We have plans to implement this natively with tensorizer, but can't speak as to when that would be, although conversations on this matter have been leaning towards a feature that allows direct tensor-device mapping that we think could be more robust than a multiple file approach. This is due to tensorizer's native parallelism when reading from a single file: all metadata is loaded up front, so we'd be able to know what devices to load tensors to.

[tjohnson31415]

Thank you for the additional context!
Yeah, my current approach requires num_shards == num_devices and assumes that the serialization and loading use the same number of devices, but it would be nice to lift that restriction in the future.

Initial PR has been pushed up: #4990

@sangstar I'd appreciate any feedback you have on it. Thanks!