【vLLM 学习】Load Sharded State

vLLM 是一款专为大语言模型推理加速而设计的框架,实现了 KV 缓存内存几乎零浪费,解决了内存管理瓶颈问题。

更多 vLLM 中文文档及教程可访问 →https://vllm.hyper.ai/

*在线运行 vLLM 入门教程:零基础分步指南

源码 examples/offline_inference/load_sharded_state.py.

复制代码
# SPDX-License-Identifier: Apache-2.0
"""
Validates the loading of a model saved with the sharded_state format.
This script demonstrates how to load a model that was previously saved
using save_sharded_state.py and validates it by running inference.
Example usage:
(First need to save a sharded_state mode)

python save_sharded_state.py \
    --model /path/to/load \
    --quantization deepspeedfp \
    --tensor-parallel-size 8 \
    --output /path/to/save/sharded/modele

python load_sharded_state.py \
    --model /path/to/saved/sharded/model \
    --load-format sharded_state \
    --quantization deepspeedfp \
    --tensor-parallel-size 8 \
    --prompt "Hello, my name is" \
    --max-tokens 50
"""

import dataclasses

from vllm import LLM, EngineArgs, SamplingParams
from vllm.utils import FlexibleArgumentParser


def parse_args():
    parser = FlexibleArgumentParser()
    # Add engine arguments
    EngineArgs.add_cli_args(parser)

    # Override default load_format for clarity
    parser.set_defaults(load_format="sharded_state")

    # Add validation arguments
    parser.add_argument("--prompt",
                        type=str,
                        default="Hello, world!",
                        help="Prompt for validation")
    parser.add_argument("--max-tokens",
                        type=int,
                        default=100,
                        help="Maximum number of tokens to generate")
    parser.add_argument("--temperature",
                        type=float,
                        default=0.7,
                        help="Sampling temperature")
    parser.add_argument("--top-p",
                        type=float,
                        default=1.0,
                        help="Top-p sampling parameter")

    return parser.parse_args()


def main():
    args = parse_args()
    engine_args = EngineArgs.from_cli_args(args)

    print(f"Loading model from {engine_args.model} "
          f"using format {engine_args.load_format}")
    print(f"Tensor parallel size: {engine_args.tensor_parallel_size}")

    # Load the model using engine args
    llm = LLM(**dataclasses.asdict(engine_args))

    # Prepare sampling parameters
    sampling_params = SamplingParams(
        temperature=args.temperature,
        top_p=args.top_p,
        max_tokens=args.max_tokens,
    )

    print("\nRunning inference:")
    print(f"Prompt: {args.prompt}")

    # Generate completion
    outputs = llm.generate(args.prompt, sampling_params)

    # Display generated text
    print("\nGenerated outputs:")
    for output in outputs:
        generated_text = output.outputs[0].text
        print("-" * 50)
        print(f"Full output: {args.prompt}{generated_text}")
        print("-" * 50)


if __name__ == "__main__":
    main()
相关推荐
Baihai_IDP18 小时前
GPU 网络通信基础,Part 1(横向与纵向扩展;训练中的通信环节...)
人工智能·llm·gpu
镰刀韭菜18 小时前
【大语言模型】大模型后训练入门指南
人工智能·自然语言处理·大语言模型·强化学习·ppo·后训练·grpo
bylander1 天前
【论文阅读】A Survey of Reinforcement Learning for Large Reasoning Models
论文阅读·大模型·llm·强化学习
Tadas-Gao1 天前
华为OmniPlacement技术深度解析:突破超大规模MoE模型推理瓶颈的创新设计
人工智能·架构·大模型·llm
AI大模型2 天前
大模型提示词开发:从“简单指令”到“工程化架构”的技术跃迁
程序员·llm·agent
AI大模型2 天前
手把手教:LangChain+Qwen3搭建本地RAG问答系统,从0到1全流程
程序员·langchain·llm
大模型教程2 天前
深入剖析大模型Embedding技术:原理、应用与实践
程序员·llm·agent
智泊AI2 天前
2025最新大语言模型(LLM)推理框架分析与选型建议,附大模型智能体(LLM Agent)学习笔记
llm
Baihai_IDP2 天前
上下文工程实施过程中会遇到什么挑战?有哪些优化策略?
人工智能·llm·aigc
阿里云云原生2 天前
基于 AI 网关提升大模型应用可用性的实践
llm