Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

support for concurrency in llm models #519

Closed
wants to merge 5 commits into from
Closed

support for concurrency in llm models #519

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9b55100
support for concurrency in llm models
dtrawins Jan 17, 2024
9e4ab17
style fixes
dtrawins Jan 18, 2024
dcb2a8f
concurrency in seq2seq and stable diffusion classes
dtrawins Jan 19, 2024
03b797f
merge from main
dtrawins Jan 19, 2024
eb5db08
merge from upstream
dtrawins Jan 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 82 additions & 24 deletions optimum/intel/openvino/modeling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,7 @@ def to(self, device: str):
be in upper or lower case. To speed up first inference, call `.compile()` after `.to()`.
"""
self._device = device.upper()
self.compiled_model = None
self.request = None
return self

Expand Down Expand Up @@ -197,8 +198,14 @@ def forward(
inputs["token_type_ids"] = token_type_ids

# Run inference
outputs = self.request(inputs)
logits = torch.from_numpy(outputs["logits"]).to(self.device) if not np_inputs else outputs["logits"]
infer_request = self.compiled_model.create_infer_request()
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we need to create an inference request for each prediction ? (from my understanding it's needed for stateful models at each generation steps, but might not be the case here)

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

actually new infer_requests are needed for stateless models. It ensure each prediction from each execution thread is independent. Otherwise only one generate operation would be possible on a model. In stateful LLM models there is different approach - each generate operation has a single infer request so the state is preserved between integrations.

infer_request.start_async(inputs)
infer_request.wait()
logits = (
torch.from_numpy(infer_request.get_tensor("logits").data).to(self.device)
if not np_inputs
else infer_request.get_tensor("logits").data
)
return SequenceClassifierOutput(logits=logits)


Expand Down Expand Up @@ -263,12 +270,18 @@ def forward(
inputs["token_type_ids"] = token_type_ids

# Run inference
outputs = self.request(inputs)
infer_request = self.compiled_model.create_infer_request()
infer_request.start_async(inputs)
infer_request.wait()
start_logits = (
torch.from_numpy(outputs["start_logits"]).to(self.device) if not np_inputs else outputs["start_logits"]
torch.from_numpy(infer_request.get_tensor("start_logits").data).to(self.device)
if not np_inputs
else infer_request.get_tensor("start_logits").data
)
end_logits = (
torch.from_numpy(outputs["end_logits"]).to(self.device) if not np_inputs else outputs["end_logits"]
torch.from_numpy(infer_request.get_tensor("end_logits").data).to(self.device)
if not np_inputs
else infer_request.get_tensor("end_logits").data
)
return QuestionAnsweringModelOutput(start_logits=start_logits, end_logits=end_logits)

Expand Down Expand Up @@ -333,8 +346,14 @@ def forward(
inputs["token_type_ids"] = token_type_ids

# Run inference
outputs = self.request(inputs)
logits = torch.from_numpy(outputs["logits"]).to(self.device) if not np_inputs else outputs["logits"]
infer_request = self.compiled_model.create_infer_request()
infer_request.start_async(inputs)
infer_request.wait()
logits = (
torch.from_numpy(infer_request.get_tensor("logits").data).to(self.device)
if not np_inputs
else infer_request.get_tensor("logits").data
)
return TokenClassifierOutput(logits=logits)


Expand Down Expand Up @@ -398,11 +417,13 @@ def forward(
inputs["token_type_ids"] = token_type_ids

# Run inference
outputs = self.request(inputs)
infer_request = self.compiled_model.create_infer_request()
infer_request.start_async(inputs)
infer_request.wait()
last_hidden_state = (
torch.from_numpy(outputs["last_hidden_state"]).to(self.device)
torch.from_numpy(infer_request.get_tensor("last_hidden_state").data).to(self.device)
if not np_inputs
else outputs["last_hidden_state"]
else infer_request.get_tensor("last_hidden_state").data
)
return BaseModelOutput(last_hidden_state=last_hidden_state)

Expand Down Expand Up @@ -468,8 +489,14 @@ def forward(
inputs["token_type_ids"] = token_type_ids

# Run inference
outputs = self.request(inputs)
logits = torch.from_numpy(outputs["logits"]).to(self.device) if not np_inputs else outputs["logits"]
infer_request = self.compiled_model.create_infer_request()
infer_request.start_async(inputs)
infer_request.wait()
logits = (
torch.from_numpy(infer_request.get_tensor("logits").data).to(self.device)
if not np_inputs
else infer_request.get_tensor("logits").data
)
return MaskedLMOutput(logits=logits)


Expand Down Expand Up @@ -595,8 +622,14 @@ def forward(
}

# Run inference
outputs = self.request(inputs)
logits = torch.from_numpy(outputs["logits"]).to(self.device) if not np_inputs else outputs["logits"]
infer_request = self.compiled_model.create_infer_request()
infer_request.start_async(inputs)
infer_request.wait()
logits = (
torch.from_numpy(infer_request.get_tensor("logits").data).to(self.device)
if not np_inputs
else infer_request.get_tensor("logits").data
)
return ImageClassifierOutput(logits=logits)


Expand Down Expand Up @@ -660,8 +693,14 @@ def forward(
inputs["attention_mask"] = attention_mask

# Run inference
outputs = self.request(inputs)
logits = torch.from_numpy(outputs["logits"]).to(self.device) if not np_inputs else outputs["logits"]
infer_request = self.compiled_model.create_infer_request()
infer_request.start_async(inputs)
infer_request.wait()
logits = (
torch.from_numpy(infer_request.get_tensor("logits").data).to(self.device)
if not np_inputs
else infer_request.get_tensor("logits").data
)
return SequenceClassifierOutput(logits=logits)


Expand Down Expand Up @@ -732,8 +771,14 @@ def forward(
inputs["attention_mask"] = attention_mask

# Run inference
outputs = self.request(inputs)
logits = torch.from_numpy(outputs["logits"]).to(self.device) if not np_inputs else outputs["logits"]
infer_request = self.compiled_model.create_infer_request()
infer_request.start_async(inputs)
infer_request.wait()
logits = (
torch.from_numpy(infer_request.get_tensor("logits").data).to(self.device)
if not np_inputs
else infer_request.get_tensor("logits").data
)
return CausalLMOutput(logits=logits)


Expand Down Expand Up @@ -813,12 +858,19 @@ def forward(
inputs["attention_mask"] = attention_mask

# Run inference
outputs = self.request(inputs)
logits = torch.from_numpy(outputs["logits"]).to(self.device) if not np_inputs else outputs["logits"]
infer_request = self.compiled_model.create_infer_request()
infer_request.start_async(inputs)
infer_request.wait()
logits = (
torch.from_numpy(infer_request.get_tensor("logits").data).to(self.device)
if not np_inputs
else infer_request.get_tensor("logits").data
)
embeddings = (
torch.from_numpy(outputs["embeddings"]).to(self.device) if not np_inputs else outputs["embeddings"]
torch.from_numpy(infer_request.get_tensor("embeddings").data).to(self.device)
if not np_inputs
else infer_request.get_tensor("embeddings").data
)

return XVectorOutput(logits=logits, embeddings=embeddings)


Expand Down Expand Up @@ -890,7 +942,13 @@ def forward(
inputs["attention_mask"] = attention_mask

# Run inference
outputs = self.request(inputs)
logits = torch.from_numpy(outputs["logits"]).to(self.device) if not np_inputs else outputs["logits"]
infer_request = self.compiled_model.create_infer_request()
infer_request.start_async(inputs)
infer_request.wait()
logits = (
torch.from_numpy(infer_request.get_tensor("logits").data).to(self.device)
if not np_inputs
else infer_request.get_tensor("logits").data
)

return TokenClassifierOutput(logits=logits)
6 changes: 4 additions & 2 deletions optimum/intel/openvino/modeling_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@ def __init__(

self.model = model
self.request = None
self.compiled_model = None
if enable_compilation:
self.compile()

Expand Down Expand Up @@ -343,15 +344,15 @@ def _to_load(
)

def compile(self):
if self.request is None:
if self.compiled_model is None:
logger.info(f"Compiling the model to {self._device} ...")
ov_config = {**self.ov_config}
if "CACHE_DIR" not in self.ov_config.keys() and not str(self.model_save_dir).startswith(gettempdir()):
# Set default CACHE_DIR only if it is not set, and if the model is not in a temporary directory
cache_dir = Path(self.model_save_dir).joinpath("model_cache")
ov_config["CACHE_DIR"] = str(cache_dir)
logger.info(f"Setting OpenVINO CACHE_DIR to {str(cache_dir)}")
self.request = core.compile_model(self.model, self._device, ov_config)
self.compiled_model = core.compile_model(self.model, self._device, ov_config)
# OPENVINO_LOG_LEVEL can be found in https://docs.openvino.ai/2023.2/openvino_docs_OV_UG_supported_plugins_AUTO_debugging.html
if "OPENVINO_LOG_LEVEL" in os.environ and int(os.environ["OPENVINO_LOG_LEVEL"]) > 2:
logger.info(f"{self._device} SUPPORTED_PROPERTIES:")
Expand Down Expand Up @@ -403,6 +404,7 @@ def half(self):
apply_moc_transformations(self.model, cf=False)
compress_model_transformation(self.model)
self.request = None
self.compiled_model = None
return self

def forward(self, *args, **kwargs):
Expand Down
43 changes: 29 additions & 14 deletions optimum/intel/openvino/modeling_decoder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,6 @@ def __init__(
"`dynamic_shapes` was set to `False` but static shapes are not supported for causal language model. Please set `dynamic_shapes=True`."
)

enable_compilation = kwargs.get("compile", True)
kwargs["compile"] = False # avoid extra compilation in the base class

super().__init__(
Expand All @@ -135,21 +134,13 @@ def __init__(
self.normalized_config = NormalizedConfigManager.get_normalized_config_class(config.model_type)(config)
self.key_value_input_names = [key for key in self.input_names if "key_values" in key]
self.key_value_output_names = [key for key in self.output_names if "present" in key]
self._original_model = self.model.clone() # keep original model for serialization
self._pkv_precision = Type.f32
self.next_beam_idx = None
self.update_pkv_precision()
if self.is_dynamic:
self.model = self._reshape(self.model, -1, -1)
is_stateful_supported = ensure_stateful_is_available(warn=False)

if self.use_cache and not self.stateful:
logger.warn(
"Provided model does not contain state. It may lead to sub-optimal performance."
"Please reexport model with updated OpenVINO version >= 2023.3.0 calling the `from_pretrained` method with original model "
"and `export=True` parameter"
)

if self.stateful:
if stateful is None:
stateful = is_stateful_supported
Expand All @@ -176,7 +167,13 @@ def raise_error(model_prop, user_prop, name):
if use_cache ^ self.use_cache:
raise_error(self.use_cache, use_cache, "use_cache")

if enable_compilation:
def init_ov_model(self, compile=True):
self._pkv_precision = Type.f32
self.update_pkv_precision(force_fp32=False)
if self.is_dynamic:
self.model = self._reshape(self.model, -1, -1)
self._original_model = self.model.clone() # keep original model for serialization
if compile:
self.compile()

def update_pkv_precision(self, force_fp32=False):
Expand Down Expand Up @@ -282,9 +279,10 @@ def _from_transformers(
config.is_decoder = True
config.is_encoder_decoder = False
config.save_pretrained(save_dir_path)
return cls._from_pretrained(
model_instance = cls._from_pretrained(
model_id=save_dir_path, config=config, use_cache=use_cache, load_in_8bit=False, stateful=None, **kwargs
)
return model_instance

def _reshape(
self,
Expand Down Expand Up @@ -322,14 +320,19 @@ def reshape(self, batch_size: int, sequence_length: int):
return self

def compile(self):
if self.request is None:
if self.compiled_model is None:
super().compile()
self.request = self.request.create_infer_request()

def _make_stateful(self):
patch_stateful(self.config, self.model)
self.stateful = True

def create_infer_request(self):
if self.compiled_model is None:
self.compile()
if self.request is None:
self.request = self.compiled_model.create_infer_request()


@add_start_docstrings(
"""
Expand Down Expand Up @@ -359,6 +362,8 @@ def forward(
**kwargs,
) -> CausalLMOutputWithPast:
self.compile()
self.create_infer_request()

if self.use_cache and past_key_values is not None:
input_ids = input_ids[:, -1:]

Expand Down Expand Up @@ -556,7 +561,17 @@ def _from_pretrained(
else:
init_cls = cls

return init_cls(model=model, config=config, model_save_dir=model_cache_path.parent, **kwargs)
model_instance = init_cls(model=model, config=config, model_save_dir=model_cache_path.parent, **kwargs)
model_instance.init_ov_model(compile=kwargs.get("compile", True))
model_instance.request = None
return model_instance

def clone(self):
model_instance = self.__class__(model=self.model, config=self.config, compile=False)
model_instance.compiled_model = self.compiled_model
model_instance._pkv_precision = self._pkv_precision
model_instance.request = None
return model_instance


class OVBloomForCausalLM(OVModelForCausalLM):
Expand Down
Loading
Loading