add custom ops

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+lib/liboptimus_ths-torch2.2-cu121.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+lib/liboptimus_ths-torch2.3-cu121.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text

config.json

@@ -6,7 +6,7 @@
     "AutoConfig": "configuration_step1.Step1Config",
     "AutoModelForCausalLM": "modeling_step1.Step1ForCausalLM"
   },
-  "model_type": "step_audio",
+  "model_type": "step1",
   "bos_token_id": 1,
   "pad_token_id": 0,
   "eos_token_id": 3,

lib/liboptimus_ths-torch2.2-cu121.cpython-310-x86_64-linux-gnu.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e018916e5e93fb904be6b34af32e71d03ba9e888d8c086a43a5c9fcacda661a1
+size 31250408

lib/liboptimus_ths-torch2.3-cu121.cpython-310-x86_64-linux-gnu.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee23bba95f7806364e101e285720892b755a176d603842fb4646822800ac2344
+size 31250472

lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6fa1a77f035203ff90a071218f775381f705269ef454163474d22501684b7e1f
+size 31258792

modeling_step1.py

@@ -76,6 +76,11 @@ class StepAttention(torch.nn.Module):
 
         self.layer_idx = layer_idx
 
+    def flash_attn_func(self, q, k, v, dropout_p=0.0, softmax_scale=None, causal=True,
+                    return_attn_probs=False, tp_group_rank=0, tp_group_size=1):
+        softmax_scale = q.size(-1) ** (-0.5) if softmax_scale is None else softmax_scale
+        return torch.ops.Optimus.fwd(q, k, v, None, dropout_p, softmax_scale, causal, return_attn_probs, None, tp_group_rank, tp_group_size)[0]
+
     def forward(
         self,
         x: torch.Tensor,
@@ -95,24 +100,31 @@ class StepAttention(torch.nn.Module):
         k = rearrange(k, "b s (g d) -> b s g d", g=self.num_groups)
         v = rearrange(v, "b s (g d) -> b s g d", g=self.num_groups)
 
-        k = k.repeat_interleave(self.num_heads // self.num_groups, dim=-2)
-        v = v.repeat_interleave(self.num_heads // self.num_groups, dim=-2)
-
-        attention_mask = build_alibi_cache(
-            k.size(1), self.num_heads, dtype=q.dtype, device=q.device
-        )[:, :, -q.size(1) :, :].contiguous()
-
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
+        try:
+            if self.head_dim not in (64, 128):
+                raise ValueError("head_dim must be 64 or 128")
+            attn_output = self.flash_attn_func(q, k, v)
+            attn_output = attn_output.flatten(-2, -1)
+        except:
+            k = k.repeat_interleave(self.num_heads // self.num_groups, dim=-2)
+            v = v.repeat_interleave(self.num_heads // self.num_groups, dim=-2)
+
+            attention_mask = build_alibi_cache(
+                k.size(1), self.num_heads, dtype=q.dtype, device=q.device
+            )[:, :, -q.size(1) :, :].contiguous()
+
+            q = q.transpose(1, 2)
+            k = k.transpose(1, 2)
+            v = v.transpose(1, 2)
+
+            attn_output: torch.Tensor = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=attention_mask
+            )
 
-        o: torch.Tensor = torch.nn.functional.scaled_dot_product_attention(
-            q, k, v, attn_mask=attention_mask
-        )
-        o = o.transpose(1, 2).flatten(-2, -1)
+            attn_output = attn_output.transpose(1, 2).flatten(-2, -1)
 
-        o = self.o_proj(o)
-        return o
+        out = self.o_proj(attn_output)
+        return out, None  # attn weights are not returned
 
 
 class StepMLP(torch.nn.Module):
@@ -153,26 +165,26 @@ class StepLayer(torch.nn.Module):
 
     def forward(
         self,
-        x,
+        hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
         past_key_value: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
         cache_position: Optional[torch.LongTensor] = None,
     ):
-        def f(x):
-            x = self.input_layernorm(x)
-            x = self.self_attn(x, past_key_value, attention_mask, cache_position)
-            return x
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, self_attn_weights = self.self_attn(hidden_states, past_key_value, attention_mask, cache_position)
+        hidden_states = residual + hidden_states
 
-        x = x + f(x)
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
 
-        def f(x):
-            x = self.post_attention_layernorm(x)
-            x = self.mlp(x)
-            return x
-
-        x = x + f(x)
-
-        return x
+        outputs = (hidden_states, )
+        if output_attentions:
+            outputs += (self_attn_weights,)
+        return outputs
 
 
 class StepPreTrainedModel(PreTrainedModel):
@@ -241,9 +253,16 @@ class Step1Model(StepPreTrainedModel):
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.LongTensor] = None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
-        output_attentions = False
-        output_hidden_states = False
-
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
         use_cache = use_cache if use_cache is not None else self.config.use_cache
         return_dict = (
             return_dict if return_dict is not None else self.config.use_return_dict
@@ -274,22 +293,37 @@ class Step1Model(StepPreTrainedModel):
 
         hidden_states = inputs_embeds
 
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+
         for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
             layer_outputs = decoder_layer(
                 hidden_states,
                 attention_mask=causal_mask,
                 past_key_value=past_key_values,
                 cache_position=cache_position,
+                output_attentions=output_attentions,
             )
 
-            hidden_states = layer_outputs
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
 
         hidden_states = self.norm(hidden_states)
 
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
         output = BaseModelOutputWithPast(
             last_hidden_state=hidden_states,
             past_key_values=past_key_values if use_cache else None,
-            hidden_states=hidden_states,
+            hidden_states=all_hidden_states,
             attentions=None,
         )
         return output if return_dict else output.to_tuple()
@@ -313,12 +347,6 @@ class Step1ForCausalLM(StepPreTrainedModel, GenerationMixin):
     def set_input_embeddings(self, value):
         self.model.embed_tokens = value
 
-    # def get_output_embeddings(self):
-    #     return self.lm_head
-
-    # def set_output_embeddings(self, new_embeddings):
-    #     self.lm_head = new_embeddings
-
     def set_decoder(self, decoder):
         self.model = decoder
 
@@ -338,14 +366,11 @@ class Step1ForCausalLM(StepPreTrainedModel, GenerationMixin):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.LongTensor] = None,
-        num_logits_to_keep: int = 0,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
-        # output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_attentions = False
-        output_hidden_states = False
-        # output_hidden_states = (
-        #     output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        # )
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
         return_dict = (
             return_dict if return_dict is not None else self.config.use_return_dict
         )
@@ -368,15 +393,12 @@ class Step1ForCausalLM(StepPreTrainedModel, GenerationMixin):
 
         logits = self.lm_head(hidden_states)
 
-        # logits = torch.matmul(hidden_states, lm_stat)
-
         loss = None
         if labels is not None:
             loss = self.loss_function(
                 logits=logits,
                 labels=labels,
                 vocab_size=self.config.vocab_size,
-                **kwargs
             )
 
         if not return_dict:

tokenizer_config.json

@@ -1,4 +1,5 @@
 {
+  "add_bos_token": true,
   "bos_token": "<s>",
   "clean_up_tokenization_spaces": false,
   "eos_token": "</s>",
@@ -9,6 +10,7 @@
   "sp_model_kwargs": {},
   "tokenizer_class": "LlamaTokenizer",
   "unk_token": "<unk>",
-  "use_default_system_prompt": false
+  "use_default_system_prompt": false,
+  "chat_template": "{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|BOT|>system\nYou are a helpful assistant.<|EOT|>' }}{% endif %}{{'<|BOT|>' + (message['role'] if message['role'] != 'user' else 'human') + '\n' + message['content'] + '<|EOT|>'}}{% endfor %}{% if add_generation_prompt %}{{ '<|BOT|>assistant\n' }}{% endif %}"
 }