Skip to content

Commit

Permalink
fix model patcher
Browse files Browse the repository at this point in the history
  • Loading branch information
eaidova committed Jun 26, 2024
1 parent aa38288 commit 1785f6d
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions optimum/exporters/openvino/model_patcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -528,8 +528,8 @@ def __enter__(self):

# llama/gemma has some accuracy issues with bf16 with transformers >= 4.39
# fill causal mask in slightly different way for avoid overflow on some platforms
max_positions = self._model.config.max_position_embeddings
if is_transformers_version(">=", "4.39.0"):
max_positions = self._model.config.max_position_embeddings
self._model.model.register_buffer(
"bias",
torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
Expand All @@ -552,9 +552,10 @@ def create_sinusoidal_positions(num_pos: int, dim: int, base: int = 10000) -> to
return torch.cat((torch.sin(emb), torch.cos(emb)), dim=1)

base = self._model.model.layers[0].self_attn.rotary_emb.base
dim = self._model.model.layers[0].self_attn.rotary_emb.dim
self._model.model.register_buffer(
"embed_positions",
create_sinusoidal_positions(max_positions, self._model.config.head_dim, self._model.config, base),
create_sinusoidal_positions(max_positions, dim, base),
)

# init cos/sin for torchscript tracing
Expand Down

0 comments on commit 1785f6d

Please sign in to comment.