File size: 13,573 Bytes
b386992 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 | # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
from typing import Any, List, Optional
import numpy as np
import torch
from peft import PeftModel
from pytriton.decorators import batch
from pytriton.model_config import Tensor
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
from nemo.deploy import ITritonDeployable
from nemo.deploy.utils import broadcast_list, cast_output, str_ndarray2list
LOGGER = logging.getLogger("NeMo")
SUPPORTED_TASKS = ["text-generation"]
class HuggingFaceLLMDeploy(ITritonDeployable):
"""A Triton inference server compatible wrapper for HuggingFace models.
This class provides a standardized interface for deploying HuggingFace models
in Triton inference server. It supports various NLP tasks and handles model
loading, inference, and deployment configurations.
Args:
hf_model_id_path (Optional[str]): Path to the HuggingFace model or model identifier.
Can be a local path or a model ID from HuggingFace Hub.
hf_peft_model_id_path (Optional[str]): Path to the PEFT model or model identifier.
Can be a local path or a model ID from HuggingFace Hub.
tokenizer_id_path (Optional[str]): Path to the tokenizer or tokenizer identifier.
If None, will use the same path as hf_model_id_path.
model (Optional[AutoModel]): Pre-loaded HuggingFace model.
tokenizer (Optional[AutoTokenizer]): Pre-loaded HuggingFace tokenizer.
tokenizer_padding (bool): Whether to enable padding in tokenizer. Defaults to True.
tokenizer_truncation (bool): Whether to enable truncation in tokenizer. Defaults to True.
tokenizer_padding_side (str): Which side to pad on ('left' or 'right'). Defaults to 'left'.
task (str): HuggingFace task type (e.g., "text-generation"). Defaults to "text-generation".
**hf_kwargs: Additional keyword arguments to pass to HuggingFace model loading.
"""
def __init__(
self,
hf_model_id_path: Optional[str] = None,
hf_peft_model_id_path: Optional[str] = None,
tokenizer_id_path: Optional[str] = None,
model: Optional[AutoModel] = None,
tokenizer: Optional[AutoTokenizer] = None,
tokenizer_padding=True,
tokenizer_truncation=True,
tokenizer_padding_side="left",
task: Optional[str] = "text-generation",
**hf_kwargs,
):
if hf_model_id_path is None and model is None:
raise ValueError("hf_model_id_path or model parameters has to be passed.")
elif hf_model_id_path is not None and model is not None:
LOGGER.warning(
"hf_model_id_path will be ignored and the HuggingFace model " "set with model parameter will be used."
)
assert task in SUPPORTED_TASKS, "Task {0} is not a support task.".format(task)
self.hf_model_id_path = hf_model_id_path
self.hf_peft_model_id_path = hf_peft_model_id_path
self.task = task
self.model = model
self.tokenizer = tokenizer
self.tokenizer_padding = tokenizer_padding
self.tokenizer_truncation = tokenizer_truncation
self.tokenizer_padding_side = tokenizer_padding_side
if tokenizer_id_path is None:
self.tokenizer_id_path = hf_model_id_path
else:
self.tokenizer_id_path = tokenizer_id_path
if model is None:
self._load(**hf_kwargs)
def _load(self, **hf_kwargs) -> None:
"""
Load the HuggingFace pipeline with the specified model and task.
This method initializes the HuggingFace AutoModel classes using the provided model
configuration and task type. It handles the model and tokenizer loading
process.
Raises:
AssertionError: If task is not specified.
"""
assert self.task is not None, "A task has to be given for the generation task."
if self.task == "text-generation":
self.model = AutoModelForCausalLM.from_pretrained(self.hf_model_id_path, **hf_kwargs)
if self.hf_peft_model_id_path is not None:
self.model = PeftModel.from_pretrained(self.model, self.hf_peft_model_id_path)
else:
raise ValueError("Task {0} is not supported.".format(self.task))
self.model.cuda()
self.tokenizer = AutoTokenizer.from_pretrained(
self.tokenizer_id_path,
trust_remote_code=hf_kwargs.pop("trust_remote_code", False),
padding=self.tokenizer_padding,
truncation=self.tokenizer_truncation,
padding_side=self.tokenizer_padding_side,
)
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
def generate(
self,
**kwargs: Any,
) -> List[str]:
"""Generate text based on the provided input prompts.
This method processes input prompts through the loaded pipeline and
generates text according to the specified parameters.
Args:
**kwargs: Generation parameters including:
- text_inputs: List of input prompts
- max_length: Maximum number of tokens to generate
- num_return_sequences: Number of sequences to generate per prompt
- temperature: Sampling temperature
- top_k: Number of highest probability tokens to consider
- top_p: Cumulative probability threshold for token sampling
- do_sample: Whether to use sampling
- return_full_text: Whether to return full text or only generated part
Returns:
If output logits and output scores are False:
List[str]: A list of generated texts, one for each input prompt.
If output logits and output scores are True:
Dict: A dictionary containing:
- sentences: List of generated texts
- logits: List of logits
- scores: List of scores
Raises:
RuntimeError: If the pipeline is not initialized.
"""
if not self.model:
raise RuntimeError("Model is not initialized")
inputs = self.tokenizer(
kwargs["text_inputs"],
return_tensors="pt",
padding=self.tokenizer_padding,
truncation=self.tokenizer_truncation,
)
kwargs = {**inputs, **kwargs}
kwargs.pop("text_inputs")
for key, val in kwargs.items():
if torch.is_tensor(val):
kwargs[key] = val.cuda()
with torch.no_grad():
generated_ids = self.model.generate(**kwargs)
return_dict_in_generate = kwargs.get("return_dict_in_generate", False)
if return_dict_in_generate:
output = {"sentences": self.tokenizer.batch_decode(generated_ids["sequences"], skip_special_tokens=True)}
if kwargs.get("output_logits", False):
output["logits"] = generated_ids["logits"]
if kwargs.get("output_scores", False):
output["scores"] = generated_ids["scores"]
else:
output = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
return output
def generate_other_ranks(self):
"""
Generate function for ranks other than the rank 0.
"""
while True:
message = torch.empty(1, dtype=torch.long, device="cuda")
torch.distributed.broadcast(message, src=0)
if message == 0:
prompts = broadcast_list(data=[None], src=0)
temperature, top_k, top_p, num_tokens_to_generate, output_logits, output_scores = broadcast_list(
data=[None], src=0
)
return_dict_in_generate = False
if output_logits or output_scores:
return_dict_in_generate = True
self.generate(
text_inputs=prompts,
do_sample=True,
top_k=top_k,
top_p=top_p,
temperature=temperature,
max_new_tokens=num_tokens_to_generate,
output_logits=output_logits,
output_scores=output_scores,
return_dict_in_generate=return_dict_in_generate,
)
else:
return
@property
def get_triton_input(self):
inputs = (
Tensor(name="prompts", shape=(-1,), dtype=bytes),
Tensor(name="max_length", shape=(-1,), dtype=np.int_, optional=True),
Tensor(name="max_batch_size", shape=(-1,), dtype=np.int_, optional=True),
Tensor(name="top_k", shape=(-1,), dtype=np.int_, optional=True),
Tensor(name="top_p", shape=(-1,), dtype=np.single, optional=True),
Tensor(name="temperature", shape=(-1,), dtype=np.single, optional=True),
Tensor(name="random_seed", shape=(-1,), dtype=np.int_, optional=True),
Tensor(name="max_length", shape=(-1,), dtype=np.int_, optional=True),
Tensor(name="output_logits", shape=(-1,), dtype=np.bool_, optional=True),
Tensor(name="output_scores", shape=(-1,), dtype=np.bool_, optional=True),
)
return inputs
@property
def get_triton_output(self):
return (
Tensor(name="sentences", shape=(-1,), dtype=bytes),
Tensor(name="logits", shape=(-1,), dtype=np.single),
Tensor(name="scores", shape=(-1,), dtype=np.single),
)
@batch
def triton_infer_fn(self, **inputs: np.ndarray):
output_infer = {}
try:
prompts = str_ndarray2list(inputs.pop("prompts"))
temperature = inputs.pop("temperature")[0][0] if "temperature" in inputs else 1.0
top_k = int(inputs.pop("top_k")[0][0] if "top_k" in inputs else 1)
top_p = inputs.pop("top_p")[0][0] if "top_k" in inputs else 0.0
num_tokens_to_generate = inputs.pop("max_length")[0][0] if "max_length" in inputs else 256
output_logits = inputs.pop("output_logits")[0][0] if "output_logits" in inputs else False
output_scores = inputs.pop("output_scores")[0][0] if "output_scores" in inputs else False
return_dict_in_generate = False
if output_logits or output_scores:
return_dict_in_generate = True
if torch.distributed.is_initialized():
if torch.distributed.get_world_size() > 1:
torch.distributed.broadcast(torch.tensor([0], dtype=torch.long, device="cuda"), src=0)
broadcast_list(prompts, src=0)
broadcast_list(
data=[
temperature,
top_k,
top_p,
num_tokens_to_generate,
output_logits,
output_scores,
],
src=0,
)
output = self.generate(
text_inputs=prompts,
do_sample=True,
top_k=top_k,
top_p=top_p,
temperature=temperature,
max_new_tokens=num_tokens_to_generate,
output_logits=output_logits,
output_scores=output_scores,
return_dict_in_generate=return_dict_in_generate,
)
if isinstance(output, dict):
output_infer = {"sentences": cast_output(output["sentences"], np.bytes_)}
if "scores" in output.keys():
output_scores = []
for r in output["scores"]:
lp = torch.tensor(r).cpu().detach().numpy()
if len(lp) == 0:
output_scores.append([0])
else:
output_scores.append(lp)
output_infer["scores"] = np.array(output_scores).transpose(1, 0, 2)
if "logits" in output.keys():
output_logits = []
for r in output["logits"]:
lp = torch.tensor(r).cpu().detach().numpy()
if len(lp) == 0:
output_logits.append([0])
else:
output_logits.append(lp)
output_infer["logits"] = np.array(output_logits).transpose(1, 0, 2)
else:
output_infer = {"sentences": cast_output(output, np.bytes_)}
except Exception as error:
err_msg = "An error occurred: {0}".format(str(error))
output_infer["sentences"] = cast_output([err_msg], np.bytes_)
return output_infer
|