optimum-rbln 0.2.1a4__tar.gz → 0.2.1a5__tar.gz

{optimum_rbln-0.2.1a4 → optimum_rbln-0.2.1a5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.2.1a4
+Version: 0.2.1a5
 Summary: Optimum RBLN is the interface between the Hugging Face Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai

{optimum_rbln-0.2.1a4 → optimum_rbln-0.2.1a5}/src/optimum/rbln/__version__.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1a4'
+__version__ = version = '0.2.1a5'
 __version_tuple__ = version_tuple = (0, 2, 1)

{optimum_rbln-0.2.1a4 → optimum_rbln-0.2.1a5}/src/optimum/rbln/modeling_base.py

@@ -442,8 +442,9 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
             return
 
-        real_save_dir = self.model_save_dir / self.subfolder
-        save_directory_path = Path(save_directory)
+        # Normalize paths to handle relative paths and symlinks
+        real_save_dir = Path(self.model_save_dir).resolve() / self.subfolder
+        save_directory_path = Path(save_directory).resolve()
 
         if not os.path.exists(real_save_dir) or not os.path.isdir(real_save_dir):
             raise FileNotFoundError(
@@ -452,13 +453,13 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                 f"Please ensure the model directory exists and you have the necessary permissions to access it."
             )
 
-        if save_directory_path.absolute() == real_save_dir.absolute():
+        if save_directory_path == real_save_dir:
             raise FileExistsError(
                 f"Cannot save model to '{save_directory}'. This directory already exists and contains the model files."
             )
 
-        # Create a temporary directory next to the target directory
-        tmp_dir = save_directory + ".tmp"
+        # Create a temporary directory with normalized path
+        tmp_dir = str(save_directory_path) + ".tmp"
         try:
             # Remove temporary directory if it exists from a previous failed attempt
             if os.path.exists(tmp_dir):
@@ -473,9 +474,9 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                 self.generation_config.save_pretrained(tmp_dir)
 
             # If everything succeeded, atomically replace the target directory
-            if os.path.exists(save_directory):
-                shutil.rmtree(save_directory)
-            os.rename(tmp_dir, save_directory)
+            if os.path.exists(save_directory_path):
+                shutil.rmtree(save_directory_path)
+            os.rename(tmp_dir, save_directory_path)
 
         except Exception as e:
             # Clean up the temporary directory if anything fails
@@ -484,7 +485,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             raise e  # Re-raise the exception after cleanup
 
         if push_to_hub:
-            return super().push_to_hub(save_directory, **kwargs)
+            return super().push_to_hub(str(save_directory_path), **kwargs)
 
     @staticmethod
     def _raise_missing_compiled_file_error(missing_files: List[str]):

{optimum_rbln-0.2.1a4 → optimum_rbln-0.2.1a5}/src/optimum/rbln/transformers/models/decoderonly/decoderonly_architecture.py

@@ -427,12 +427,14 @@ class DecoderOnlyModel(nn.Module):
             cos, sin = None, None
 
         # (batch, seq_len) -> (batch,)
-        seq_positions = cache_position[:, 0]
         if self.attn_impl == "flash_attn":
+            seq_positions = cache_position[:, 0]
             max_seq_len = past_key_values[0][0].shape[-2]
             seq_positions = self.convert_sequence_positions_for_flash_attn(
                 seq_positions=seq_positions, max_seq_len=max_seq_len
             )
+        else:
+            seq_positions = cache_position[:, :1]
 
         present_key_values = past_key_values
         for layer in self.layers:

{optimum_rbln-0.2.1a4 → optimum_rbln-0.2.1a5}/src/optimum/rbln/transformers/models/seq2seq/seq2seq_architecture.py

@@ -459,7 +459,7 @@ class Seq2SeqSelfAttention(nn.Module):
             ),  # Unsqueeze group axis since CustomKernel expects it for group query attention
             past_key_value[0].view(bsz, self.num_heads, 1, -1, self.head_dim),
             past_key_value[1].view(bsz, self.num_heads, 1, -1, self.head_dim),
-            cache_position.squeeze(1),
+            cache_position,
             torch.tensor(1.0, dtype=torch.float32),  # scale
         )