Upload modeling_nemotron_h.py

modeling_nemotron_h.py

@@ -24,25 +24,25 @@ import torch.utils.checkpoint
 from torch import nn
 from torch.nn import CrossEntropyLoss
 
-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 (
+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 (
     AttentionMaskConverter,
 )
-from ...modeling_utils import PreTrainedModel
-from ...utils import (
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (
     ModelOutput,
     add_code_sample_docstrings,
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
     logging,
 )
-from ...utils.import_utils import (
+from transformers.utils.import_utils import (
     is_causal_conv1d_available,
     is_flash_attn_2_available,
     is_flash_attn_greater_or_equal_2_10,
-    is_mamba_2_ssm_available,    
+    is_mamba_2_ssm_available,
 )
 from .configuration_nemotron_h import NemotronHConfig
 
@@ -70,7 +70,7 @@ else:
     causal_conv1d_update, causal_conv1d_fn = None, None
 
 if is_flash_attn_2_available():
-    from ...modeling_flash_attention_utils import _flash_attention_forward
+    from transformers.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.attention_head_dim is not None:
-            self.head_dim = config.attention_head_dim
+        if config.head_dim is not None:
+            self.head_dim = config.head_dim
         else:
             self.head_dim = config.hidden_size // config.num_attention_heads
         self.num_key_value_heads = config.num_key_value_heads