Upload modeling_nemotron_h.py

modeling_nemotron_h.py

@@ -24,21 +24,21 @@ import torch.utils.checkpoint
 from torch import nn
 from torch.nn import CrossEntropyLoss
 
-from transformers.activations import ACT2FN
-from transformers.cache_utils import DynamicCache  # we need __iter__ and __len__ of pkv
-from transformers.generation import GenerationMixin
-from transformers.modeling_attn_mask_utils import (
+from ...activations import ACT2FN
+from ...cache_utils import DynamicCache  # we need __iter__ and __len__ of pkv
+from ...generation import GenerationMixin
+from ...modeling_attn_mask_utils import (
     AttentionMaskConverter,
 )
-from transformers.modeling_utils import PreTrainedModel
-from transformers.utils import (
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
     ModelOutput,
     add_code_sample_docstrings,
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
     logging,
 )
-from transformers.utils.import_utils import (
+from ...utils.import_utils import (
     is_causal_conv1d_available,
     is_flash_attn_2_available,
     is_flash_attn_greater_or_equal_2_10,
@@ -70,7 +70,7 @@ else:
     causal_conv1d_update, causal_conv1d_fn = None, None
 
 if is_flash_attn_2_available():
-    from transformers.modeling_flash_attention_utils import _flash_attention_forward
+    from ...modeling_flash_attention_utils import _flash_attention_forward
 
 is_fast_path_available = all(
     (
@@ -844,8 +844,8 @@ class NemotronHAttention(nn.Module):
         self.attention_dropout = config.attention_dropout
         self.hidden_size = config.hidden_size
         self.num_heads = config.num_attention_heads
-        if config.head_dim is not None:
-            self.head_dim = config.head_dim
+        if config.attention_head_dim is not None:
+            self.head_dim = config.attention_head_dim
         else:
             self.head_dim = config.hidden_size // config.num_attention_heads
         self.num_key_value_heads = config.num_key_value_heads
@@ -1542,6 +1542,11 @@ class NemotronHForCausalLM(NemotronHPreTrainedModel, GenerationMixin):
 
         # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
         if inputs_embeds is not None and empty_past_kv:
+            # TODO(pjin): workaround fix for properly extending inputs_embeds;
+            # longer term, may be better handled elsewhere in .generate().
+            if input_ids is not None and inputs_embeds.shape[1] < input_ids.shape[1]:
+                new_token_embeds = self.get_input_embeddings()(input_ids[:,inputs_embeds.shape[1]:])
+                inputs_embeds = torch.cat([inputs_embeds, new_token_embeds], dim=1)
             model_inputs = {"inputs_embeds": inputs_embeds}
         else:
             model_inputs = {"input_ids": input_ids.contiguous()}  # `contiguous()` needed for compilation use cases