[V0 deprecation] Deprecate V0 Neuron backend (#21159)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
This commit is contained in:
Woosuk Kwon
@@ -1,476 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Utilities for selecting and loading Neuron models in transformers-neuronx
|
||||
framework."""
|
||||
import ast
|
||||
import copy
|
||||
import importlib
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from vllm.config import (ModelConfig, ParallelConfig, SchedulerConfig,
|
||||
SpeculativeConfig)
|
||||
from vllm.logprobs import Logprob
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.layers.quantization import get_quantization_config
|
||||
from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
|
||||
from vllm.model_executor.sampling_metadata import SamplingMetadata
|
||||
from vllm.sequence import CompletionSequenceGroupOutput, SequenceOutput
|
||||
|
||||
TORCH_DTYPE_TO_NEURON_AMP = {
|
||||
"auto": "f32",
|
||||
"half": "f16",
|
||||
"float16": "f16",
|
||||
"bfloat16": "bf16",
|
||||
"float": "f32",
|
||||
"float32": "f32",
|
||||
torch.float16: "f16",
|
||||
torch.bfloat16: "bf16",
|
||||
torch.float32: "f32",
|
||||
}
|
||||
|
||||
# Models supported by Neuron.
|
||||
_NEURON_SUPPORTED_MODELS: dict[str, tuple[str, str, str]] = {
|
||||
"LlamaForCausalLM": ("transformers_neuronx.llama.model",
|
||||
"LlamaForSampling", "LlamaForCausalLM"),
|
||||
"MistralForCausalLM": ("transformers_neuronx.mistral.model",
|
||||
"MistralForSampling", "MistralForCausalLM")
|
||||
}
|
||||
|
||||
|
||||
class NeuronCausalLM(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
config: PretrainedConfig,
|
||||
on_device_sampling_disabled: bool = False) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.logits_processor = LogitsProcessor(config.vocab_size,
|
||||
logits_as_input=True)
|
||||
|
||||
self.on_device_sampling_disabled = on_device_sampling_disabled
|
||||
if self.on_device_sampling_disabled:
|
||||
# Use default sampler
|
||||
self.sampler = Sampler()
|
||||
|
||||
# Lazy initialized
|
||||
self.model: nn.Module
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
input_block_ids: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
logits = self.model(input_ids,
|
||||
cache_ids=positions,
|
||||
start_ids=input_block_ids)
|
||||
return logits
|
||||
|
||||
def compute_logits(self, hidden_states: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata) -> torch.Tensor:
|
||||
logits = self.logits_processor(None, hidden_states, sampling_metadata)
|
||||
return logits
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> Optional[SamplerOutput]:
|
||||
|
||||
if self.on_device_sampling_disabled:
|
||||
next_tokens = self.sampler(logits, sampling_metadata)
|
||||
return next_tokens
|
||||
|
||||
# On-device sampling outputs the token ids directly.
|
||||
sampled_token_ids = logits.flatten()
|
||||
next_tokens = []
|
||||
sample_idx = 0
|
||||
for seq_group in sampling_metadata.seq_groups:
|
||||
samples = []
|
||||
for seq_id in seq_group.seq_ids:
|
||||
token_id = sampled_token_ids[sample_idx].item()
|
||||
samples.append(
|
||||
SequenceOutput(parent_seq_id=seq_id,
|
||||
output_token=token_id,
|
||||
logprobs={token_id: Logprob(token_id)}))
|
||||
sample_idx += 1
|
||||
next_tokens.append(
|
||||
CompletionSequenceGroupOutput(samples=samples,
|
||||
prompt_logprobs=None))
|
||||
|
||||
return SamplerOutput(outputs=next_tokens)
|
||||
|
||||
def load_weights(self, model_name_or_path: str, **kwargs):
|
||||
arch = _get_model_architecture(self.config)
|
||||
neuronx_module_path, neuronx_model_cls_name, hf_model_cls_name = (
|
||||
_NEURON_SUPPORTED_MODELS[arch])
|
||||
neuronx_module = importlib.import_module(neuronx_module_path)
|
||||
neuronx_model_cls = getattr(neuronx_module, neuronx_model_cls_name)
|
||||
|
||||
self.model = neuronx_model_cls.from_pretrained(model_name_or_path,
|
||||
**kwargs)
|
||||
self.model.to_neuron()
|
||||
|
||||
|
||||
class NeuronSpeculationCausalLM(nn.Module):
|
||||
"""A Neuron-optimized causal language model with speculative decoding."""
|
||||
|
||||
SPECULATION_TERMINATION_ID = -1
|
||||
|
||||
def __init__(self, speculation_model) -> None:
|
||||
super().__init__()
|
||||
self.model = speculation_model
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
input_block_ids: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
tokens, counts = self.model.speculative_iteration(
|
||||
input_ids, positions, input_block_ids)
|
||||
|
||||
# Mark the end of accepted speculative tokens for each sequence with the
|
||||
# speculation termination id.
|
||||
batch_size, steps = tokens.shape
|
||||
mask = torch.arange(steps).expand(batch_size, -1) >= counts
|
||||
tokens[mask] = self.SPECULATION_TERMINATION_ID
|
||||
|
||||
return tokens
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> Optional[list[SamplerOutput]]:
|
||||
batch_size, num_steps = logits.shape
|
||||
seq_ids = [
|
||||
seq_id for sg in sampling_metadata.seq_groups
|
||||
for seq_id in sg.seq_ids
|
||||
]
|
||||
# Organize input tensors by step instead of by sequence.
|
||||
accepted_token_ids_by_step = logits.transpose(0, 1)
|
||||
accepted_token_ids_by_step = accepted_token_ids_by_step.tolist()
|
||||
|
||||
sampler_output_list = []
|
||||
for step_index in range(num_steps):
|
||||
if all(token_id == self.SPECULATION_TERMINATION_ID
|
||||
for token_id in accepted_token_ids_by_step[step_index]):
|
||||
break
|
||||
step_output_token_ids = []
|
||||
for sequence_index in range(batch_size):
|
||||
token_id = accepted_token_ids_by_step[step_index][
|
||||
sequence_index]
|
||||
step_output_token_ids.append(
|
||||
CompletionSequenceGroupOutput(samples=[
|
||||
SequenceOutput(parent_seq_id=seq_ids[sequence_index],
|
||||
output_token=token_id,
|
||||
logprobs={token_id: Logprob(token_id)})
|
||||
],
|
||||
prompt_logprobs=None))
|
||||
sampler_output_list.append(
|
||||
SamplerOutput(outputs=step_output_token_ids))
|
||||
return sampler_output_list
|
||||
|
||||
|
||||
def _get_model_architecture(config: PretrainedConfig) -> str:
|
||||
architectures = getattr(config, "architectures", [])
|
||||
for arch in architectures:
|
||||
if arch in _NEURON_SUPPORTED_MODELS:
|
||||
return arch
|
||||
raise ValueError(
|
||||
f"Model architectures {architectures} are not supported on Neuron "
|
||||
f"for now. Supported architectures: "
|
||||
f"{list(_NEURON_SUPPORTED_MODELS.keys())}")
|
||||
|
||||
|
||||
def _get_buckets(env: str, default_value: list[int]) -> list[int]:
|
||||
env_value = os.getenv(env)
|
||||
if env_value is None:
|
||||
return default_value
|
||||
buckets_remove_empty = filter(
|
||||
lambda x: x is not None and len(x.strip()) > 0, env_value.split(","))
|
||||
buckets_int = map(int, buckets_remove_empty)
|
||||
buckets_list = list(buckets_int)
|
||||
return buckets_list
|
||||
|
||||
|
||||
def _get_default_neuron_config(model_config: ModelConfig,
|
||||
parallel_config: ParallelConfig,
|
||||
scheduler_config: SchedulerConfig):
|
||||
"""Generate a neuron config based on vllm config args."""
|
||||
from transformers_neuronx.config import ContinuousBatchingConfig
|
||||
from transformers_neuronx.constants import LAYOUT_BSH
|
||||
|
||||
continuous_batching_config = ContinuousBatchingConfig(
|
||||
batch_size_for_shared_caches=scheduler_config.max_num_seqs)
|
||||
quant_config = dict(
|
||||
dequant_dtype=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
|
||||
quantize_method="vector_dynamic")
|
||||
neuron_quantization_config_builder = lambda quant: get_quantization_config(
|
||||
quant).from_config(quant_config).get_quant_method(None, "")
|
||||
# TODO: Add Paged attention config to the default neuron arguments.
|
||||
default_neuron_args = dict(
|
||||
collectives_layout=LAYOUT_BSH,
|
||||
attention_layout=LAYOUT_BSH,
|
||||
fuse_qkv=True,
|
||||
quant=neuron_quantization_config_builder(model_config.quantization)
|
||||
if model_config.quantization else None,
|
||||
continuous_batching=continuous_batching_config,
|
||||
weight_tiling=bool(model_config.quantization),
|
||||
on_device_generation=_get_neuron_on_device_generation_config(
|
||||
model_config))
|
||||
return default_neuron_args
|
||||
|
||||
|
||||
def _get_default_neuron_config_for_speculation(
|
||||
model_config: ModelConfig, parallel_config: ParallelConfig,
|
||||
scheduler_config: SchedulerConfig):
|
||||
"""Generate a neuron config for speculative decoding based on
|
||||
vllm config args."""
|
||||
from transformers_neuronx.config import ContinuousBatchingConfig
|
||||
from transformers_neuronx.constants import LAYOUT_BSH
|
||||
|
||||
continuous_batching_config = ContinuousBatchingConfig(
|
||||
batch_size_for_shared_caches=scheduler_config.max_num_seqs)
|
||||
|
||||
default_neuron_args = dict(collectives_layout=LAYOUT_BSH,
|
||||
attention_layout=LAYOUT_BSH,
|
||||
fuse_qkv=True,
|
||||
on_device_embedding=True,
|
||||
continuous_batching=continuous_batching_config,
|
||||
on_device_generation=copy.deepcopy(
|
||||
model_config.neuron_sampling_params))
|
||||
return default_neuron_args
|
||||
|
||||
|
||||
def _get_neuron_on_device_generation_config(model_config: ModelConfig):
|
||||
if not _is_neuron_on_device_sampling_disabled(model_config):
|
||||
return copy.deepcopy(model_config.neuron_sampling_params)
|
||||
return None
|
||||
|
||||
|
||||
def _is_neuron_on_device_sampling_disabled(model_config: ModelConfig) -> bool:
|
||||
return not getattr(model_config, "neuron_sampling_params", None)
|
||||
|
||||
|
||||
def _get_neuron_config_after_override(default_neuron_config,
|
||||
overridden_neuron_config):
|
||||
from transformers_neuronx.config import (ContinuousBatchingConfig,
|
||||
GenerationConfig,
|
||||
KVCacheQuantizationConfig,
|
||||
NeuronConfig, QuantizationConfig,
|
||||
SparseAttnConfig)
|
||||
|
||||
sparse_attn = overridden_neuron_config.pop("sparse_attn", {})
|
||||
if sparse_attn:
|
||||
overridden_neuron_config["sparse_attn"] = SparseAttnConfig(
|
||||
**sparse_attn)
|
||||
|
||||
kv_cache_quant = overridden_neuron_config.pop("kv_cache_quant", {})
|
||||
if kv_cache_quant:
|
||||
overridden_neuron_config["kv_cache_quant"] = KVCacheQuantizationConfig(
|
||||
**kv_cache_quant)
|
||||
|
||||
continuous_batching = overridden_neuron_config.pop("continuous_batching",
|
||||
{})
|
||||
if continuous_batching:
|
||||
overridden_neuron_config[
|
||||
"continuous_batching"] = ContinuousBatchingConfig(
|
||||
**continuous_batching)
|
||||
|
||||
quant = overridden_neuron_config.pop("quant", {})
|
||||
if quant:
|
||||
overridden_neuron_config["quant"] = QuantizationConfig(**quant)
|
||||
|
||||
on_device_generation = overridden_neuron_config.pop(
|
||||
"on_device_generation", {})
|
||||
if on_device_generation:
|
||||
overridden_neuron_config["on_device_generation"] = GenerationConfig(
|
||||
**on_device_generation)
|
||||
default_neuron_config.update(overridden_neuron_config)
|
||||
return NeuronConfig(**default_neuron_config)
|
||||
|
||||
|
||||
def get_neuron_model(model_config: ModelConfig,
|
||||
parallel_config: ParallelConfig,
|
||||
scheduler_config: SchedulerConfig) -> nn.Module:
|
||||
"""Initializes a neuron-optimized model for inference."""
|
||||
# Create a model instance.
|
||||
model = NeuronCausalLM(
|
||||
model_config.hf_config,
|
||||
_is_neuron_on_device_sampling_disabled(model_config))
|
||||
|
||||
default_neuron_config_args = _get_default_neuron_config(
|
||||
model_config, parallel_config, scheduler_config)
|
||||
|
||||
neuron_config = _get_neuron_config_after_override(
|
||||
default_neuron_config_args, model_config.override_neuron_config)
|
||||
|
||||
context_length_estimates = _get_buckets("NEURON_CONTEXT_LENGTH_BUCKETS",
|
||||
[scheduler_config.max_model_len])
|
||||
n_positions = _get_buckets("NEURON_TOKEN_GEN_BUCKETS",
|
||||
[scheduler_config.max_model_len])
|
||||
|
||||
model.load_weights(model_config.model,
|
||||
tp_degree=parallel_config.tensor_parallel_size,
|
||||
amp=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
|
||||
neuron_config=neuron_config,
|
||||
context_length_estimate=context_length_estimates,
|
||||
n_positions=n_positions,
|
||||
batch_size=scheduler_config.max_num_seqs)
|
||||
|
||||
return model.eval()
|
||||
|
||||
|
||||
def get_neuron_speculation_model(model_config: ModelConfig,
|
||||
parallel_config: ParallelConfig,
|
||||
scheduler_config: SchedulerConfig,
|
||||
speculation_config: SpeculativeConfig):
|
||||
"""Initializes a neuron-optimized speculation model for inference.
|
||||
|
||||
This method is only applicable for speculation with a standalone draft model
|
||||
"""
|
||||
from transformers_neuronx.fused_speculation import FusedSpeculativeDecoder
|
||||
|
||||
# For Eagle SD, we need to pass in additional parameters in neuron config.
|
||||
is_eagle = getattr(speculation_config.draft_model_config.hf_config,
|
||||
"is_eagle", False)
|
||||
|
||||
# Create target model instance.
|
||||
target_model = NeuronCausalLM(model_config.hf_config)
|
||||
|
||||
default_neuron_config_args = _get_default_neuron_config_for_speculation(
|
||||
model_config, parallel_config, scheduler_config)
|
||||
if is_eagle:
|
||||
default_neuron_config_args['is_eagle_target'] = True
|
||||
|
||||
neuron_config = _get_neuron_config_after_override(
|
||||
default_neuron_config_args, model_config.override_neuron_config)
|
||||
|
||||
context_length_estimates = _get_buckets("NEURON_CONTEXT_LENGTH_BUCKETS",
|
||||
[scheduler_config.max_model_len])
|
||||
n_positions = _get_buckets("NEURON_TOKEN_GEN_BUCKETS",
|
||||
[scheduler_config.max_model_len])
|
||||
|
||||
target_model.load_weights(
|
||||
model_config.model,
|
||||
tp_degree=parallel_config.tensor_parallel_size,
|
||||
amp=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
|
||||
neuron_config=neuron_config,
|
||||
context_length_estimate=context_length_estimates,
|
||||
n_positions=n_positions,
|
||||
batch_size=scheduler_config.max_num_seqs)
|
||||
|
||||
target_model.eval()
|
||||
|
||||
# Create draft model instance.
|
||||
draft_model = NeuronCausalLM(
|
||||
speculation_config.draft_model_config.hf_config)
|
||||
|
||||
default_draft_neuron_config_args = (
|
||||
_get_default_neuron_config_for_speculation(
|
||||
speculation_config.draft_model_config, parallel_config,
|
||||
scheduler_config))
|
||||
if is_eagle:
|
||||
default_draft_neuron_config_args['is_eagle_draft'] = True
|
||||
default_draft_neuron_config_args['has_pre_attention_norm'] = False
|
||||
|
||||
draft_neuron_config = _get_neuron_config_after_override(
|
||||
default_draft_neuron_config_args,
|
||||
speculation_config.draft_model_config.override_neuron_config)
|
||||
|
||||
draft_model.load_weights(speculation_config.draft_model_config.model,
|
||||
tp_degree=speculation_config.
|
||||
draft_parallel_config.tensor_parallel_size,
|
||||
amp=TORCH_DTYPE_TO_NEURON_AMP[
|
||||
speculation_config.draft_model_config.dtype],
|
||||
neuron_config=draft_neuron_config,
|
||||
context_length_estimate=context_length_estimates,
|
||||
n_positions=n_positions,
|
||||
batch_size=scheduler_config.max_num_seqs)
|
||||
|
||||
draft_model.eval()
|
||||
|
||||
num_speculative_tokens = speculation_config.num_speculative_tokens
|
||||
# Create speculation model instance.
|
||||
speculation_model = FusedSpeculativeDecoder(draft_model.model,
|
||||
target_model.model,
|
||||
num_speculative_tokens)
|
||||
speculation_model.to_neuron()
|
||||
|
||||
return NeuronSpeculationCausalLM(speculation_model)
|
||||
|
||||
|
||||
def get_neuron_eagle_speculation_model(model_config: ModelConfig,
|
||||
parallel_config: ParallelConfig,
|
||||
scheduler_config: SchedulerConfig,
|
||||
speculation_config: SpeculativeConfig):
|
||||
"""Initializes a neuron-optimized EAGLE speculation model for inference."""
|
||||
from transformers_neuronx.eagle_speculation import EagleSpeculativeDecoder
|
||||
|
||||
# Create target model instance.
|
||||
target_model = NeuronCausalLM(model_config.hf_config)
|
||||
|
||||
default_neuron_config_args = _get_default_neuron_config_for_speculation(
|
||||
model_config, parallel_config, scheduler_config)
|
||||
default_neuron_config_args['is_eagle_target'] = True
|
||||
neuron_config = _get_neuron_config_after_override(
|
||||
default_neuron_config_args, model_config.override_neuron_config)
|
||||
|
||||
context_length_estimates = _get_buckets("NEURON_CONTEXT_LENGTH_BUCKETS",
|
||||
[scheduler_config.max_model_len])
|
||||
n_positions = _get_buckets("NEURON_TOKEN_GEN_BUCKETS",
|
||||
[scheduler_config.max_model_len])
|
||||
|
||||
target_model.load_weights(
|
||||
model_config.model,
|
||||
tp_degree=parallel_config.tensor_parallel_size,
|
||||
amp=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
|
||||
neuron_config=neuron_config,
|
||||
context_length_estimate=context_length_estimates,
|
||||
n_positions=n_positions,
|
||||
batch_size=scheduler_config.max_num_seqs)
|
||||
|
||||
target_model.eval()
|
||||
|
||||
# Create draft model instance.
|
||||
draft_model = NeuronCausalLM(
|
||||
speculation_config.draft_model_config.hf_config)
|
||||
|
||||
default_draft_neuron_config_args = (
|
||||
_get_default_neuron_config_for_speculation(
|
||||
speculation_config.draft_model_config, parallel_config,
|
||||
scheduler_config))
|
||||
default_draft_neuron_config_args['is_eagle_draft'] = True
|
||||
default_draft_neuron_config_args['has_pre_attention_norm'] = False
|
||||
draft_neuron_config = _get_neuron_config_after_override(
|
||||
default_draft_neuron_config_args,
|
||||
speculation_config.draft_model_config.override_neuron_config)
|
||||
|
||||
draft_model.load_weights(speculation_config.draft_model_config.model,
|
||||
tp_degree=speculation_config.
|
||||
draft_parallel_config.tensor_parallel_size,
|
||||
amp=TORCH_DTYPE_TO_NEURON_AMP[
|
||||
speculation_config.draft_model_config.dtype],
|
||||
neuron_config=draft_neuron_config,
|
||||
context_length_estimate=context_length_estimates,
|
||||
n_positions=n_positions,
|
||||
batch_size=scheduler_config.max_num_seqs)
|
||||
|
||||
draft_model.eval()
|
||||
|
||||
token_tree: dict[int, list[int]] = ast.literal_eval(
|
||||
speculation_config.speculative_token_tree)
|
||||
|
||||
speculation_model = EagleSpeculativeDecoder(draft_model.model,
|
||||
target_model.model,
|
||||
token_tree=token_tree)
|
||||
speculation_model.to_neuron()
|
||||
|
||||
return NeuronSpeculationCausalLM(speculation_model)
|
||||
@@ -1,685 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Utilities for selecting and loading Neuron models in
|
||||
neuronx-distributed-inference framework."""
|
||||
# Disabling yapf because yapf and isort have conflicts for the below imports
|
||||
# yapf: disable
|
||||
import copy
|
||||
import hashlib
|
||||
import importlib
|
||||
import multiprocessing
|
||||
import os
|
||||
import shutil
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from neuronx_distributed_inference.models.config import (
|
||||
FusedSpecNeuronConfig, OnDeviceSamplingConfig)
|
||||
from neuronx_distributed_inference.models.mllama.utils import (
|
||||
create_vision_mask)
|
||||
from neuronx_distributed_inference.modules.lora_serving import (
|
||||
LoraServingConfig)
|
||||
from neuronx_distributed_inference.utils.hf_adapter import (
|
||||
load_pretrained_config)
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, PretrainedConfig
|
||||
|
||||
from vllm.config import (ModelConfig, ParallelConfig, SchedulerConfig,
|
||||
SpeculativeConfig)
|
||||
from vllm.logger import init_logger
|
||||
from vllm.logprobs import Logprob
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
|
||||
from vllm.model_executor.sampling_metadata import SamplingMetadata
|
||||
from vllm.sequence import CompletionSequenceGroupOutput, SequenceOutput
|
||||
|
||||
# yapf: enable
|
||||
logger = init_logger(__name__)
|
||||
|
||||
TORCH_DTYPE_TO_NEURON_AMP = {
|
||||
"auto": "float32",
|
||||
"half": "float16",
|
||||
"float16": "float16",
|
||||
"bfloat16": "bfloat16",
|
||||
"float": "float32",
|
||||
"float32": "float32",
|
||||
torch.float16: "float16",
|
||||
torch.bfloat16: "bfloat16",
|
||||
torch.float32: "float32",
|
||||
}
|
||||
|
||||
# Models supported by Neuronx distributed for inference.
|
||||
_NEURON_SUPPORTED_MODELS: dict[str, tuple[str, str]] = {
|
||||
"LlamaForCausalLM":
|
||||
("neuronx_distributed_inference.models.llama.modeling_llama",
|
||||
"NeuronLlamaForCausalLM"),
|
||||
"MistralForCausalLM":
|
||||
("neuronx_distributed_inference.models.llama.modeling_llama",
|
||||
"NeuronLlamaForCausalLM"),
|
||||
"DbrxForCausalLM":
|
||||
("neuronx_distributed_inference.models.dbrx.modeling_dbrx",
|
||||
"NeuronDbrxForCausalLM"),
|
||||
"MixtralForCausalLM":
|
||||
("neuronx_distributed_inference.models.mixtral.modeling_mixtral",
|
||||
"NeuronMixtralForCausalLM"),
|
||||
"MllamaForConditionalGeneration":
|
||||
("neuronx_distributed_inference.models.mllama.modeling_mllama",
|
||||
"NeuronMllamaForCausalLM"),
|
||||
}
|
||||
|
||||
|
||||
class NeuronCausalLM(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: PretrainedConfig,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.logits_processor = LogitsProcessor(config.vocab_size,
|
||||
logits_as_input=True)
|
||||
self.sampler = Sampler()
|
||||
|
||||
# Lazy initialized
|
||||
self.model: nn.Module
|
||||
|
||||
def forward(self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
input_block_ids: torch.Tensor,
|
||||
sampling_params: torch.Tensor,
|
||||
prev_hidden: Optional[torch.Tensor] = None,
|
||||
adapter_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
# sort block ids sequentially for perf/neuron support reasons
|
||||
sorted_input_block_ids, sorted_indices = torch.sort(input_block_ids)
|
||||
input_ids = torch.index_select(input_ids, 0, sorted_indices)
|
||||
positions = torch.index_select(positions, 0, sorted_indices)
|
||||
sampling_params = torch.index_select(sampling_params, 0,
|
||||
sorted_indices)
|
||||
output = self.model(input_ids,
|
||||
attention_mask=None,
|
||||
position_ids=positions,
|
||||
seq_ids=sorted_input_block_ids,
|
||||
sampling_params=sampling_params,
|
||||
prev_hidden=prev_hidden,
|
||||
adapter_ids=adapter_ids)
|
||||
# on-device sampling
|
||||
if self.config.neuron_config.on_device_sampling_config:
|
||||
output = output.hidden_states
|
||||
else:
|
||||
output = output.logits[:, -1, :]
|
||||
|
||||
restored_indices = torch.argsort(sorted_indices)
|
||||
if input_block_ids.shape[0] != 1:
|
||||
output = torch.index_select(output, 0, restored_indices)
|
||||
|
||||
return output
|
||||
|
||||
def compute_logits(self, hidden_states: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata) -> torch.Tensor:
|
||||
logits = self.logits_processor(None, hidden_states, sampling_metadata)
|
||||
return logits
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> Optional[SamplerOutput]:
|
||||
# on-device sampling
|
||||
if self.config.neuron_config.on_device_sampling_config:
|
||||
batch_size = logits.shape
|
||||
seq_ids = [
|
||||
seq_id for sg in sampling_metadata.seq_groups
|
||||
for seq_id in sg.seq_ids
|
||||
]
|
||||
assert len(seq_ids) == list(batch_size)[0], "batch size mismatch"
|
||||
# Organize input tensors by step instead of by sequence.
|
||||
accepted_token_ids_by_step = logits.flatten()
|
||||
accepted_token_ids_by_step = accepted_token_ids_by_step.tolist()
|
||||
|
||||
step_output_token_ids = []
|
||||
for i, seq_id in enumerate(seq_ids):
|
||||
token_id = accepted_token_ids_by_step[i]
|
||||
step_output_token_ids.append(
|
||||
CompletionSequenceGroupOutput(samples=[
|
||||
SequenceOutput(parent_seq_id=seq_id,
|
||||
output_token=token_id,
|
||||
logprobs={token_id: Logprob(token_id)})
|
||||
],
|
||||
prompt_logprobs=None))
|
||||
return SamplerOutput(outputs=step_output_token_ids)
|
||||
else:
|
||||
return self.sampler(logits, sampling_metadata)
|
||||
|
||||
def load_weights(self, model_name_or_path: str, **kwargs):
|
||||
arch = _get_model_architecture(self.config)
|
||||
neuronx_module_path, neuronx_model_cls_name = (
|
||||
_NEURON_SUPPORTED_MODELS[arch])
|
||||
neuronx_module = importlib.import_module(neuronx_module_path)
|
||||
neuronx_model_cls = getattr(neuronx_module, neuronx_model_cls_name)
|
||||
neuron_config = neuronx_model_cls.get_neuron_config_cls()(
|
||||
**kwargs['neuron_config'])
|
||||
self.config.neuron_config = neuron_config
|
||||
config = neuronx_model_cls.get_config_cls()(
|
||||
neuron_config,
|
||||
load_config=load_pretrained_config(model_name_or_path))
|
||||
hashed_config = hashlib.md5(config.to_json_string().encode('utf-8'),
|
||||
usedforsecurity=False).hexdigest()
|
||||
if os.getenv("NEURON_COMPILED_ARTIFACTS") is not None:
|
||||
compiled_model_path = os.getenv("NEURON_COMPILED_ARTIFACTS")
|
||||
elif os.path.exists(model_name_or_path):
|
||||
compiled_model_path = os.path.join(model_name_or_path,
|
||||
"neuron-compiled-artifacts",
|
||||
hashed_config)
|
||||
shutil.rmtree(compiled_model_path, ignore_errors=True)
|
||||
else:
|
||||
compiled_model_path = os.path.join("local-models",
|
||||
model_name_or_path,
|
||||
"neuron-compiled-artifacts",
|
||||
hashed_config)
|
||||
shutil.rmtree(compiled_model_path, ignore_errors=True)
|
||||
try:
|
||||
self.model = neuronx_model_cls(compiled_model_path)
|
||||
override_neuron_config = kwargs["override_neuron_config"]
|
||||
for k, v in override_neuron_config.items():
|
||||
setattr(self.model.config.neuron_config, k, v)
|
||||
self.model.load(compiled_model_path)
|
||||
return
|
||||
except (FileNotFoundError, ValueError) as e:
|
||||
logger.warning("Exception: %s", e)
|
||||
logger.warning("Failed to load the model from %s, Recompiling...",
|
||||
compiled_model_path)
|
||||
if not os.path.exists(model_name_or_path):
|
||||
hf_model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
|
||||
saved_path = os.path.join("local-models", model_name_or_path)
|
||||
hf_model.save_pretrained(saved_path)
|
||||
model_name_or_path = saved_path
|
||||
self.model = neuronx_model_cls(model_name_or_path, config)
|
||||
self.model.compile(compiled_model_path)
|
||||
self.model.load(compiled_model_path)
|
||||
|
||||
|
||||
class NeuronMllamaForCausalLM(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
config: PretrainedConfig,
|
||||
on_device_sampling_disabled: bool = False) -> None:
|
||||
super().__init__()
|
||||
# has_image is the only multimodal input that is used in
|
||||
# token-generation
|
||||
# This is a cache (on CPU) that saves has_image data per sequence id
|
||||
# The number of entries in this cache is <= Batch-Size
|
||||
self.has_image_cache: dict[int, torch.Tensor] = {}
|
||||
self.config = config
|
||||
self.logits_processor = LogitsProcessor(
|
||||
config.get_text_config().vocab_size, logits_as_input=True)
|
||||
|
||||
self.on_device_sampling_disabled = on_device_sampling_disabled
|
||||
if self.on_device_sampling_disabled:
|
||||
# Use default sampler
|
||||
self.sampler = Sampler()
|
||||
|
||||
# Lazy initialized
|
||||
self.model: nn.Module
|
||||
self.is_reorder_needed: bool = True
|
||||
|
||||
def read_from_has_image_cache(self, seq_ids: torch.Tensor):
|
||||
has_image_list = []
|
||||
for index in range(len(seq_ids)):
|
||||
seq_id = seq_ids[index].item()
|
||||
if seq_id in self.has_image_cache:
|
||||
has_image_list.append(self.has_image_cache[seq_id])
|
||||
else:
|
||||
has_image_list.append(torch.tensor([0]))
|
||||
return torch.tensor(has_image_list)
|
||||
|
||||
def write_to_has_image_cache(self, seq_ids: torch.Tensor,
|
||||
has_image: torch.Tensor):
|
||||
for index in range(len(seq_ids)):
|
||||
seq_id = seq_ids[index].item()
|
||||
if index < len(has_image):
|
||||
self.has_image_cache[seq_id] = has_image[index]
|
||||
else:
|
||||
self.has_image_cache[seq_id] = torch.zeros(1)
|
||||
|
||||
def forward(self, input_ids: torch.Tensor, positions: torch.Tensor,
|
||||
seq_ids: torch.Tensor, pixel_values: torch.Tensor,
|
||||
aspect_ratios: torch.Tensor, num_chunks: torch.Tensor,
|
||||
has_image: torch.Tensor, sampling_params) -> torch.Tensor:
|
||||
|
||||
# We update the has_image cache during prefill
|
||||
# and read the has_image cache during decode
|
||||
if input_ids.shape[-1] > 1: # prefill
|
||||
self.write_to_has_image_cache(seq_ids, has_image)
|
||||
else:
|
||||
has_image = self.read_from_has_image_cache(seq_ids)
|
||||
bs = input_ids.shape[0]
|
||||
num_chunks = torch.zeros((bs, 1))
|
||||
aspect_ratios = torch.zeros((bs, 1, 2))
|
||||
|
||||
input_block_ids = seq_ids
|
||||
origin_input_block_ids = seq_ids
|
||||
if self.is_reorder_needed:
|
||||
# sort block ids sequentially for perf/neuron support reasons
|
||||
input_block_ids, sorted_indices = torch.sort(input_block_ids)
|
||||
input_ids = torch.index_select(input_ids, 0, sorted_indices)
|
||||
positions = torch.index_select(positions, 0, sorted_indices)
|
||||
sampling_params = torch.index_select(sampling_params, 0,
|
||||
sorted_indices)
|
||||
pixel_values = torch.index_select(pixel_values, 0, sorted_indices)
|
||||
aspect_ratios = torch.index_select(aspect_ratios, 0,
|
||||
sorted_indices)
|
||||
num_chunks = torch.index_select(num_chunks, 0, sorted_indices)
|
||||
has_image = torch.index_select(has_image, 0, sorted_indices)
|
||||
|
||||
self.vision_mask = create_vision_mask(input_ids, self.vision_token_id)
|
||||
output = self.model(
|
||||
input_ids.to(torch.int32),
|
||||
attention_mask=None,
|
||||
position_ids=positions.to(torch.int32),
|
||||
seq_ids=seq_ids.flatten().to(torch.int32),
|
||||
pixel_values=pixel_values.to(
|
||||
self.config.vision_config.torch_dtype),
|
||||
aspect_ratios=aspect_ratios.to(torch.int32),
|
||||
vision_mask=self.vision_mask.to(torch.int32),
|
||||
sampling_params=sampling_params,
|
||||
num_chunks=num_chunks.to(torch.int32),
|
||||
has_image=has_image.to(torch.int32),
|
||||
)
|
||||
if self.config.neuron_config.on_device_sampling_config:
|
||||
output = output.hidden_states
|
||||
else:
|
||||
output = output.logits[:, -1, :]
|
||||
|
||||
if self.is_reorder_needed and origin_input_block_ids.shape[0] != 1:
|
||||
restored_indices = torch.argsort(sorted_indices)
|
||||
output = torch.index_select(output, 0, restored_indices)
|
||||
return output
|
||||
|
||||
def compute_logits(self, hidden_states: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata) -> torch.Tensor:
|
||||
logits = self.logits_processor(None, hidden_states, sampling_metadata)
|
||||
return logits
|
||||
|
||||
def sample(self, hidden_states, sampling_metadata):
|
||||
if not self.on_device_sampling_disabled:
|
||||
with torch.profiler.record_function("sample"):
|
||||
hidden_states = hidden_states.flatten()
|
||||
res = []
|
||||
sample_idx = 0
|
||||
for seq_group in sampling_metadata.seq_groups:
|
||||
seq_ids = seq_group.seq_ids
|
||||
samples = []
|
||||
for seq_id in seq_ids:
|
||||
token_id = hidden_states[sample_idx].item()
|
||||
samples.append(
|
||||
SequenceOutput(
|
||||
parent_seq_id=seq_id,
|
||||
output_token=token_id,
|
||||
logprobs={token_id: Logprob(token_id)}))
|
||||
sample_idx += 1
|
||||
res.append(
|
||||
CompletionSequenceGroupOutput(samples=samples,
|
||||
prompt_logprobs=None))
|
||||
next_tokens = SamplerOutput(outputs=res)
|
||||
else:
|
||||
next_tokens = self.sampler(None, hidden_states, sampling_metadata)
|
||||
return next_tokens
|
||||
|
||||
def load_weights(self, model_name_or_path: str, **kwargs):
|
||||
arch = _get_model_architecture(self.config)
|
||||
neuronx_module_path, neuronx_model_cls_name = (
|
||||
_NEURON_SUPPORTED_MODELS[arch])
|
||||
neuronx_module = importlib.import_module(neuronx_module_path)
|
||||
neuronx_model_cls = getattr(neuronx_module, neuronx_model_cls_name)
|
||||
neuron_config = neuronx_model_cls.get_neuron_config_cls()(
|
||||
**kwargs['neuron_config'])
|
||||
self.config.neuron_config = neuron_config
|
||||
logger.info("neuron_config buckets: %s",
|
||||
self.config.neuron_config.buckets)
|
||||
config = neuronx_model_cls.get_config_cls()(
|
||||
neuron_config,
|
||||
load_config=load_pretrained_config(model_name_or_path))
|
||||
hashed_config = hashlib.md5(config.to_json_string().encode('utf-8'),
|
||||
usedforsecurity=False).hexdigest()
|
||||
if os.getenv("NEURON_COMPILED_ARTIFACTS") is not None:
|
||||
compiled_model_path = os.getenv("NEURON_COMPILED_ARTIFACTS")
|
||||
elif os.path.exists(model_name_or_path):
|
||||
compiled_model_path = os.path.join(model_name_or_path,
|
||||
"neuron-compiled-artifacts",
|
||||
hashed_config)
|
||||
else:
|
||||
compiled_model_path = os.path.join("local-models",
|
||||
model_name_or_path,
|
||||
"neuron-compiled-artifacts",
|
||||
hashed_config)
|
||||
try:
|
||||
self.model = neuronx_model_cls(compiled_model_path)
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
|
||||
self.vision_token_id = tokenizer(
|
||||
"<|image|>", add_special_tokens=False).input_ids[0]
|
||||
self.model.load(compiled_model_path)
|
||||
return
|
||||
except (FileNotFoundError, ValueError):
|
||||
logger.warning("Failed to load the model from %s, Recompiling...",
|
||||
compiled_model_path)
|
||||
if not os.path.exists(model_name_or_path):
|
||||
hf_model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
|
||||
saved_path = os.path.join("local-models", model_name_or_path)
|
||||
hf_model.save_pretrained(saved_path)
|
||||
model_name_or_path = saved_path
|
||||
self.model = neuronx_model_cls(model_name_or_path, config)
|
||||
|
||||
logger.info("\nCompiling and saving model to %s", model_name_or_path)
|
||||
|
||||
p = multiprocessing.Process(target=compile_model,
|
||||
args=(self, compiled_model_path))
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
|
||||
tokenizer.save_pretrained(compiled_model_path)
|
||||
logger.info("Successfully compiled and saved the model in %s",
|
||||
compiled_model_path)
|
||||
|
||||
# Read "<|image|>" token_id from the tokenizer
|
||||
self.vision_token_id = tokenizer("<|image|>",
|
||||
add_special_tokens=False).input_ids[0]
|
||||
logger.info("\nLoading model from compiled checkpoint...")
|
||||
self.model.load(compiled_model_path)
|
||||
|
||||
|
||||
def compile_model(neuron_model, traced_model_path):
|
||||
neuron_model.model.compile(traced_model_path)
|
||||
|
||||
|
||||
class NeuronSpeculationCausalLM(nn.Module):
|
||||
"""A Neuron-optimized causal language model with speculative decoding."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: PretrainedConfig,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.logits_processor = LogitsProcessor(config.vocab_size,
|
||||
logits_as_input=True)
|
||||
# Lazy initialized
|
||||
self.model: nn.Module
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
input_block_ids: torch.Tensor,
|
||||
sampling_params: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
# sort block ids sequentially for perf/neuron support reasons
|
||||
sorted_input_block_ids, sorted_indices = torch.sort(input_block_ids)
|
||||
input_ids = torch.index_select(input_ids, 0, sorted_indices)
|
||||
positions = torch.index_select(positions, 0, sorted_indices)
|
||||
sampling_params = torch.index_select(sampling_params, 0,
|
||||
sorted_indices)
|
||||
|
||||
output = self.model(input_ids,
|
||||
attention_mask=None,
|
||||
position_ids=positions,
|
||||
seq_ids=sorted_input_block_ids,
|
||||
sampling_params=sampling_params)
|
||||
restored_indices = torch.argsort(sorted_indices)
|
||||
|
||||
# CTX encoding
|
||||
if (positions[:, 0]).sum().item() == 0:
|
||||
output = output.fused_outputs[0][:, 0:1]
|
||||
if input_block_ids.shape[0] != 1:
|
||||
output = torch.index_select(output, 0, restored_indices)
|
||||
return output
|
||||
|
||||
# Fused Spec (Generation)
|
||||
accepted_tokens_with_padding = output.fused_outputs[0]
|
||||
next_pos_ids = output.fused_outputs[-1]
|
||||
generated_token_counts = next_pos_ids - positions
|
||||
|
||||
assert torch.any(generated_token_counts == 0).item() is False, \
|
||||
"NxDI model generated no output for one or more sequences."
|
||||
|
||||
batch_size, steps = accepted_tokens_with_padding.shape
|
||||
mask = torch.arange(steps).expand(batch_size,
|
||||
-1) >= generated_token_counts
|
||||
accepted_tokens_with_padding[mask] = -1
|
||||
|
||||
if input_block_ids.shape[0] != 1:
|
||||
accepted_tokens_with_padding = torch.index_select(
|
||||
accepted_tokens_with_padding, 0, restored_indices)
|
||||
|
||||
return accepted_tokens_with_padding
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> Optional[list[SamplerOutput]]:
|
||||
batch_size, num_steps = logits.shape
|
||||
seq_ids = [
|
||||
seq_id for sg in sampling_metadata.seq_groups
|
||||
for seq_id in sg.seq_ids
|
||||
]
|
||||
# Organize input tensors by step instead of by sequence.
|
||||
accepted_token_ids_by_step = logits.transpose(0, 1)
|
||||
accepted_token_ids_by_step = accepted_token_ids_by_step.tolist()
|
||||
|
||||
sampler_output_list = []
|
||||
for step_index in range(num_steps):
|
||||
if all(token_id == -1
|
||||
for token_id in accepted_token_ids_by_step[step_index]):
|
||||
break
|
||||
step_output_token_ids = []
|
||||
for sequence_index in range(batch_size):
|
||||
token_id = accepted_token_ids_by_step[step_index][
|
||||
sequence_index]
|
||||
step_output_token_ids.append(
|
||||
CompletionSequenceGroupOutput(samples=[
|
||||
SequenceOutput(parent_seq_id=seq_ids[sequence_index],
|
||||
output_token=token_id,
|
||||
logprobs={token_id: Logprob(token_id)})
|
||||
],
|
||||
prompt_logprobs=None))
|
||||
sampler_output_list.append(
|
||||
SamplerOutput(outputs=step_output_token_ids))
|
||||
return sampler_output_list
|
||||
|
||||
def load_weights(self, model_name_or_path: str,
|
||||
draft_model_name_or_path: str, **kwargs):
|
||||
arch = _get_model_architecture(self.config)
|
||||
neuronx_module_path, neuronx_model_cls_name = (
|
||||
_NEURON_SUPPORTED_MODELS[arch])
|
||||
neuronx_module = importlib.import_module(neuronx_module_path)
|
||||
neuronx_model_cls = getattr(neuronx_module, neuronx_model_cls_name)
|
||||
neuron_config = neuronx_model_cls.get_neuron_config_cls()(
|
||||
**kwargs['neuron_config'])
|
||||
config = neuronx_model_cls.get_config_cls()(
|
||||
neuron_config,
|
||||
load_config=load_pretrained_config(model_name_or_path))
|
||||
|
||||
draft_neuron_config = copy.deepcopy(config.neuron_config)
|
||||
if not config.neuron_config.enable_eagle_speculation:
|
||||
draft_neuron_config.speculation_length = 0
|
||||
draft_neuron_config.trace_tokengen_model = True
|
||||
draft_neuron_config.enable_fused_speculation = False
|
||||
if getattr(config.neuron_config, "draft_model_modules_to_not_convert",
|
||||
None):
|
||||
draft_neuron_config.modules_to_not_convert = (
|
||||
draft_neuron_config.draft_model_modules_to_not_convert)
|
||||
if config.neuron_config.enable_eagle_speculation:
|
||||
draft_neuron_config.is_eagle_draft = True
|
||||
draft_neuron_config.sequence_parallel_enabled = False
|
||||
draft_config = neuronx_model_cls.get_config_cls()(
|
||||
draft_neuron_config,
|
||||
load_config=load_pretrained_config(draft_model_name_or_path))
|
||||
fused_spec_config = (FusedSpecNeuronConfig(
|
||||
neuronx_model_cls._model_cls,
|
||||
draft_config=draft_config,
|
||||
draft_model_path=draft_model_name_or_path))
|
||||
config.fused_spec_config = fused_spec_config
|
||||
self.config.neuron_config = neuron_config
|
||||
|
||||
hashed_config = hashlib.md5(config.to_json_string().encode('utf-8'),
|
||||
usedforsecurity=False).hexdigest()
|
||||
if os.getenv("NEURON_COMPILED_ARTIFACTS") is not None:
|
||||
compiled_model_path = os.getenv("NEURON_COMPILED_ARTIFACTS")
|
||||
elif os.path.exists(model_name_or_path):
|
||||
compiled_model_path = os.path.join(model_name_or_path,
|
||||
"neuron-compiled-artifacts",
|
||||
hashed_config)
|
||||
shutil.rmtree(compiled_model_path, ignore_errors=True)
|
||||
else:
|
||||
compiled_model_path = os.path.join("local-models",
|
||||
model_name_or_path,
|
||||
"neuron-compiled-artifacts",
|
||||
hashed_config)
|
||||
shutil.rmtree(compiled_model_path, ignore_errors=True)
|
||||
try:
|
||||
self.model = neuronx_model_cls(compiled_model_path)
|
||||
override_neuron_config = kwargs["override_neuron_config"]
|
||||
for k, v in override_neuron_config.items():
|
||||
setattr(self.model.config.neuron_config, k, v)
|
||||
self.model.load(compiled_model_path)
|
||||
return
|
||||
except (FileNotFoundError, ValueError) as e:
|
||||
logger.warning("Exception: %s", e)
|
||||
logger.warning("Failed to load the model from %s Recompiling...",
|
||||
compiled_model_path)
|
||||
if not os.path.exists(model_name_or_path):
|
||||
hf_model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
|
||||
saved_path = os.path.join("local-models", model_name_or_path)
|
||||
hf_model.save_pretrained(saved_path)
|
||||
model_name_or_path = saved_path
|
||||
if not os.path.exists(draft_model_name_or_path):
|
||||
if draft_model_name_or_path != model_name_or_path:
|
||||
hf_model = AutoModelForCausalLM.from_pretrained(
|
||||
draft_model_name_or_path)
|
||||
saved_path = os.path.join("local-models",
|
||||
draft_model_name_or_path)
|
||||
hf_model.save_pretrained(saved_path)
|
||||
draft_model_name_or_path = saved_path
|
||||
else:
|
||||
draft_model_name_or_path = model_name_or_path
|
||||
config.fused_spec_config.draft_model_path = draft_model_name_or_path
|
||||
self.model = neuronx_model_cls(model_name_or_path, config)
|
||||
self.model.compile(compiled_model_path)
|
||||
self.model.load(compiled_model_path)
|
||||
|
||||
|
||||
def _get_model_architecture(config: PretrainedConfig) -> str:
|
||||
architectures = getattr(config, "architectures", [])
|
||||
for arch in architectures:
|
||||
if arch in _NEURON_SUPPORTED_MODELS:
|
||||
return arch
|
||||
raise ValueError(
|
||||
f"Model architectures {architectures} are not supported on Neuron "
|
||||
f"for now. Supported architectures: "
|
||||
f"{list(_NEURON_SUPPORTED_MODELS.keys())}")
|
||||
|
||||
|
||||
def _get_default_neuron_config(model_config: ModelConfig,
|
||||
parallel_config: ParallelConfig,
|
||||
scheduler_config: SchedulerConfig,
|
||||
lora_serving_config: LoraServingConfig):
|
||||
"""Generate a neuron config based on vllm config args."""
|
||||
on_device_sampling_config = OnDeviceSamplingConfig(dynamic=True,
|
||||
deterministic=False)
|
||||
batch_size = scheduler_config.max_num_seqs
|
||||
|
||||
neuron_config = dict(
|
||||
tp_degree=parallel_config.tensor_parallel_size,
|
||||
ctx_batch_size=1,
|
||||
batch_size=batch_size,
|
||||
max_context_length=scheduler_config.max_model_len,
|
||||
seq_len=scheduler_config.max_model_len,
|
||||
enable_bucketing=True,
|
||||
is_continuous_batching=True,
|
||||
quantized=False,
|
||||
torch_dtype=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
|
||||
padding_side="right",
|
||||
on_device_sampling_config=on_device_sampling_config,
|
||||
sequence_parallel_enabled=True,
|
||||
lora_serving_config=lora_serving_config)
|
||||
return neuron_config
|
||||
|
||||
|
||||
def _get_default_speculation_config(model_config: ModelConfig,
|
||||
parallel_config: ParallelConfig,
|
||||
scheduler_config: SchedulerConfig,
|
||||
speculation_config: SpeculativeConfig):
|
||||
"""Generate a neuron config for speculative decoding based on vllm config
|
||||
args."""
|
||||
neuron_config = dict(
|
||||
tp_degree=parallel_config.tensor_parallel_size,
|
||||
ctx_batch_size=1,
|
||||
batch_size=scheduler_config.max_num_seqs,
|
||||
max_context_length=scheduler_config.max_model_len,
|
||||
seq_len=scheduler_config.max_model_len,
|
||||
speculation_length=speculation_config.num_speculative_tokens,
|
||||
trace_tokengen_model=False,
|
||||
enable_fused_speculation=True,
|
||||
enable_bucketing=True,
|
||||
is_continuous_batching=True,
|
||||
quantized=False,
|
||||
torch_dtype=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
|
||||
on_device_sampling_config=dict(
|
||||
top_k=1,
|
||||
do_sample=False,
|
||||
))
|
||||
return neuron_config
|
||||
|
||||
|
||||
def _get_neuron_config_after_override(default_neuron_config,
|
||||
overridden_neuron_config):
|
||||
"""Update default neuron config values with override args"""
|
||||
overridden_neuron_config = overridden_neuron_config or {}
|
||||
default_neuron_config.update(overridden_neuron_config)
|
||||
return default_neuron_config
|
||||
|
||||
|
||||
def get_neuron_model(model_config: ModelConfig,
|
||||
parallel_config: ParallelConfig,
|
||||
scheduler_config: SchedulerConfig,
|
||||
lora_serving_config: LoraServingConfig) -> nn.Module:
|
||||
"""Initializes a neuron-optimized model for inference."""
|
||||
model_arch = _get_model_architecture(model_config.hf_config)
|
||||
if model_arch == "MllamaForConditionalGeneration":
|
||||
model = NeuronMllamaForCausalLM(model_config.hf_config)
|
||||
else:
|
||||
model = NeuronCausalLM(model_config.hf_config)
|
||||
default_neuron_config_args = _get_default_neuron_config(
|
||||
model_config, parallel_config, scheduler_config, lora_serving_config)
|
||||
neuron_config = _get_neuron_config_after_override(
|
||||
default_neuron_config_args, model_config.override_neuron_config)
|
||||
|
||||
override_neuron_config = model_config.override_neuron_config
|
||||
model.load_weights(model_config.model,
|
||||
neuron_config=neuron_config,
|
||||
override_neuron_config=override_neuron_config)
|
||||
return model.eval()
|
||||
|
||||
|
||||
def get_neuron_speculation_model(model_config: ModelConfig,
|
||||
parallel_config: ParallelConfig,
|
||||
scheduler_config: SchedulerConfig,
|
||||
speculation_config: SpeculativeConfig):
|
||||
"""Initializes a neuron-optimized speculation model for inference.
|
||||
|
||||
This model handles speculation using both a draft model and an EAGLE draft.
|
||||
"""
|
||||
model = NeuronSpeculationCausalLM(model_config.hf_config)
|
||||
default_neuron_config_args = _get_default_speculation_config(
|
||||
model_config, parallel_config, scheduler_config, speculation_config)
|
||||
neuron_config = _get_neuron_config_after_override(
|
||||
default_neuron_config_args, model_config.override_neuron_config)
|
||||
|
||||
override_neuron_config = model_config.override_neuron_config
|
||||
model.load_weights(model_config.model,
|
||||
speculation_config.draft_model_config.model,
|
||||
neuron_config=neuron_config,
|
||||
override_neuron_config=override_neuron_config)
|
||||
return model.eval()
|
||||
Reference in New Issue
Block a user