Stream model weights

flyte.extras.model_streamer loads a model’s safetensors weights from object storage tensor by tensor, inside a task. Use it when a GPU task loads a large model, for example in an alru_cached bootstrap on reusable containers (see Batch inference). The safetensors weights are not written to local disk. load_hf_model() and the vLLM engine_args() download every other file under the model’s object-store prefix (config, tokenizer, and any non-safetensors files) to local disk.

Requires flyte 2.11.1 or later. The streamer needs torch, load_hf_model() also needs transformers, and the vLLM load format needs vllm and flyteplugins-vllm 2.11.1 or later.

Model location

Every entry point takes path, the object-store prefix of a Hugging Face–style model directory, such as s3://bucket/models/qwen2.5-7b. The prefix holds the *.safetensors files, optionally a model.safetensors.index.json, and the config and tokenizer files. To copy a Hub model there, use flyte.prefetch.hf_model() (see Prefetching models). A task that takes the model as a flyte.io.Dir passes model.path.

A local filesystem path, with or without file://, is not supported.

Load a Hugging Face model

flyte.extras.model_streamer.load_hf_model() is async. It builds a transformers model with its parameters on the meta device, then streams each weight onto device as it arrives. It returns the model in eval mode and the local directory holding the config and tokenizer files, so you can load the tokenizer with AutoTokenizer.from_pretrained(local_dir).

Checkpoint tensors that the model has no parameter for are logged and ignored. A parameter that no tensor fills raises KeyError.

This task embeds text on a reusable GPU container, and alru_cache loads the model once per container:

import flyte
from async_lru import alru_cache
from flyte.io import Dir

env = flyte.TaskEnvironment(
    name="streamed-embedder",
    image=flyte.Image.from_debian_base().with_pip_packages("async-lru", "torch", "transformers", "unionai-reuse"),
    resources=flyte.Resources(cpu=4, memory="8Gi", gpu="L4:1"),
    reusable=flyte.ReusePolicy(replicas=1, concurrency=8),
)

@alru_cache(maxsize=1)
async def get_model(model_path: str):
    """Load the model once per container."""
    import torch
    from flyte.extras.model_streamer import load_hf_model
    from transformers import AutoModel, AutoTokenizer

    model, local_dir = await load_hf_model(model_path, model_class=AutoModel, device="cuda", dtype=torch.float16)
    return model, AutoTokenizer.from_pretrained(local_dir)

@env.task
async def embed(model: Dir, texts: list[str]) -> list[list[float]]:
    import torch

    hf_model, tokenizer = await get_model(model.path)
    batch = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt").to("cuda")
    with torch.inference_mode():
        cls = hf_model(**batch).last_hidden_state[:, 0]
    return torch.nn.functional.normalize(cls, dim=-1).float().cpu().tolist()

Stream into any module

ModelStreamer(path, chunk_size=None, max_concurrency=None) is the streamer underneath, for models that are not transformers models. It uses the same chunk_size and max_concurrency defaults as load_hf_model().

• stream(device=None, dtype=None): an async iterator of (name, tensor) pairs, in the order the tensors finish downloading. With device, each tensor is moved there as it arrives. With dtype, floating-point tensors are cast.
• stream_sync(device=None, dtype=None): the same, as a regular iterator. It works from any thread, including one with a running event loop.
• state_dict(device=None, dtype=None): collects every tensor into a dict.
• load_into(module, device=None, dtype=None, strict=True, key_mapping=None): streams the weights into a module’s parameters and buffers, replacing each entry as it arrives. Build the module with with empty_weights(): so its parameters are on meta and never allocated twice. Then pass device: a meta entry with no device raises ValueError. After the stream ends, load_into() calls the module’s tie_weights() if it has one. With strict=True, a parameter left on meta or an unmatched checkpoint tensor raises KeyError. With strict=False, it returns a LoadResult listing missing_keys and unexpected_keys.
• download_metadata(local_dir): downloads everything under the prefix except the *.safetensors files.

This task streams a checkpoint into a PyTorch module:

import flyte
from flyte.io import Dir

env = flyte.TaskEnvironment(
    name="streamed-classifier",
    image=flyte.Image.from_debian_base().with_pip_packages("torch"),
    resources=flyte.Resources(cpu=4, memory="8Gi", gpu="L4:1"),
)

@env.task
async def predict(model: Dir, features: list[float]) -> list[float]:
    import torch
    from flyte.extras.model_streamer import ModelStreamer, empty_weights

    # The parameters stay on the meta device until their tensors arrive.
    with empty_weights():
        classifier = torch.nn.Sequential(torch.nn.Linear(4, 16), torch.nn.ReLU(), torch.nn.Linear(16, 2))

    await ModelStreamer(model.path).load_into(classifier, device="cuda")

    with torch.inference_mode():
        return classifier(torch.tensor([features], device="cuda")).softmax(dim=-1)[0].tolist()

Load weights into vLLM

flyteplugins-vllm registers a vLLM load format, flyte-streaming, through vLLM’s vllm.general_plugins entry point. flyteplugins.vllm.model_streamer.engine_args(path) downloads the non-safetensors files under path and returns the LLM keyword arguments that select the load format: model, load_format, and model_loader_extra_config. Pass any other engine arguments alongside them. It takes the same local_dir, chunk_size, and max_concurrency arguments as load_hf_model().

In an async task, build the engine off the event loop, for example await asyncio.to_thread(LLM, **args, max_model_len=4096), because LLM() blocks for the whole load.

This task builds the engine once per reusable container and shares it between calls:

import asyncio

import flyte
from async_lru import alru_cache
from flyte.io import Dir
from flyteplugins.vllm import DEFAULT_VLLM_IMAGE

env = flyte.TaskEnvironment(
    name="streamed-vllm",
    image=DEFAULT_VLLM_IMAGE.clone(name="streamed-vllm").with_pip_packages("async-lru", "unionai-reuse"),
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu="L4:1"),
    reusable=flyte.ReusePolicy(replicas=1, concurrency=4),
)

@alru_cache(maxsize=1)
async def get_llm(model_path: str):
    """Build the engine once per container."""
    from flyteplugins.vllm.model_streamer import engine_args
    from vllm import LLM

    args = await engine_args(model_path)
    # LLM() blocks for the whole load, so run it off the event loop.
    return await asyncio.to_thread(LLM, **args, max_model_len=4096)

# Concurrent calls on a replica share the engine, and the offline LLM isn't re-entrant.
generate_lock = asyncio.Lock()

@env.task
async def generate(model: Dir, prompts: list[str]) -> list[str]:
    from vllm import SamplingParams

    llm = await get_llm(model.path)
    messages = [[{"role": "user", "content": prompt}] for prompt in prompts]
    async with generate_lock:
        outputs = await asyncio.to_thread(llm.chat, messages, SamplingParams(max_tokens=128))
    return [output.outputs[0].text for output in outputs]

engine_args() also sets VLLM_WORKER_MULTIPROC_METHOD=spawn unless it is already set. Keep it at spawn. obstore’s runtime does not survive fork(), and vLLM forks its engine-core process by default. ModelStreamer detects a process forked from one that has already streamed or downloaded metadata and raises RuntimeError instead of hanging on the first read.

Gotchas

• A prefix with no *.safetensors files raises ValueError: No files found matching pattern: *.safetensors.
• Weights in other formats under the prefix, such as pytorch_model.bin or an original/*.pth directory, are downloaded to local disk with the config and tokenizer. Keep weights in other formats out of the prefix.

Related pages

To serve a model behind an endpoint instead of loading it inside a task, use stream_model=True on a vLLM endpoint or an SGLang endpoint.