bobber commited on
Commit
1a96a52
·
verified ·
1 Parent(s): aff278c

Add merged weights: GRPO v12 (score=0.760)

Browse files
.gitattributes CHANGED
@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
  adapter/tokenizer.json filter=lfs diff=lfs merge=lfs -text
 
 
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
  adapter/tokenizer.json filter=lfs diff=lfs merge=lfs -text
37
+ tokenizer.json filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ base_model: nvidia/Nemotron-3-4B-Base
4
+ library_name: transformers
5
+ pipeline_tag: text-generation
6
+ tags:
7
+ - grpo
8
+ - interview
9
+ - lex-fridman
10
+ - nemotron
11
+ - mamba
12
+ ---
13
+
14
+ # Lex Fridman Interviewer — Nemotron-3-Nano-4B GRPO v12 (merged weights)
15
+
16
+ Full merged weights of Nemotron-3-Nano-4B fine-tuned to ask Lex Fridman-style interview questions.
17
+
18
+ ## Training
19
+ 1. **LoRA v1** (r=64, LR=2e-4, 1 epoch, 4,772 pairs) → score 0.733
20
+ 2. **GRPO v12** (reward_v12, 200 steps, LR=5e-6) → score **0.760**
21
+
22
+ ## Eval (functional judge: on_topic × uses_guest × probing)
23
+
24
+ | Model | Score | uses_guest | probing |
25
+ |-------|-------|------------|---------|
26
+ | Base | 0.653 | 48% | 84% |
27
+ | LoRA v1 | 0.733 | 56% | 92% |
28
+ | **This model** | **0.760** | **60%** | **96%** |
29
+
30
+ ## Note on ONNX
31
+ NemotronH (Mamba-2 hybrid) cannot be exported to ONNX — the 38 SSM layers
32
+ use compiled CUDA kernels with no ONNX equivalent.
33
+ Use this model with `vllm` or `llama.cpp`.
34
+
35
+ ## Reward Design (reward_v12)
36
+ `uses_guest_logit^0.67 × probing_logit^0.33 + lexical_bonus`
37
+ Continuous judge logits from Qwen3.5-4B.
chat_template.jinja ADDED
@@ -0,0 +1,204 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {% macro render_extra_keys(json_dict, handled_keys) %}
2
+ {%- if json_dict is mapping %}
3
+ {%- for json_key in json_dict if json_key not in handled_keys %}
4
+ {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
5
+ {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson ) ~ '</' ~ json_key ~ '>' }}
6
+ {%- else %}
7
+ {{-'\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
8
+ {%- endif %}
9
+ {%- endfor %}
10
+ {%- endif %}
11
+ {% endmacro %}
12
+ {%- set enable_thinking = enable_thinking if enable_thinking is defined else True %}
13
+ {%- set truncate_history_thinking = truncate_history_thinking if truncate_history_thinking is defined else True %}
14
+
15
+ {%- set ns = namespace(last_user_idx = -1) %}
16
+ {%- set loop_messages = messages %}
17
+ {%- for m in loop_messages %}
18
+ {%- if m["role"] == "user" %}
19
+ {%- set ns.last_user_idx = loop.index0 %}
20
+ {%- endif %}
21
+ {%- endfor %}
22
+
23
+ {%- if messages[0]["role"] == "system" %}
24
+ {%- set system_message = messages[0]["content"] %}
25
+ {%- set loop_messages = messages[1:] %}
26
+ {%- else %}
27
+ {%- set system_message = "" %}
28
+ {%- set loop_messages = messages %}
29
+ {%- endif %}
30
+ {%- if not tools is defined %}
31
+ {%- set tools = [] %}
32
+ {%- endif %}
33
+ {# Recompute last_user_idx relative to loop_messages after handling system #}
34
+ {%- set ns = namespace(last_user_idx = -1) %}
35
+ {%- for m in loop_messages %}
36
+ {%- if m["role"] == "user" %}
37
+ {%- set ns.last_user_idx = loop.index0 %}
38
+ {%- endif %}
39
+ {%- endfor %}
40
+ {%- if system_message is defined %}
41
+ {{- "<|im_start|>system\n" + system_message }}
42
+ {%- else %}
43
+ {%- if tools is iterable and tools | length > 0 %}
44
+ {{- "<|im_start|>system\n" }}
45
+ {%- endif %}
46
+ {%- endif %}
47
+ {%- if tools is iterable and tools | length > 0 %}
48
+ {%- if system_message is defined and system_message | length > 0 %}
49
+ {{- "\n\n" }}
50
+ {%- endif %}
51
+ {{- "# Tools\n\nYou have access to the following functions:\n\n" }}
52
+ {{- "<tools>" }}
53
+ {%- for tool in tools %}
54
+ {%- if tool.function is defined %}
55
+ {%- set tool = tool.function %}
56
+ {%- endif %}
57
+ {{- "\n<function>\n<name>" ~ tool.name ~ "</name>" }}
58
+ {%- if tool.description is defined %}
59
+ {{- '\n<description>' ~ (tool.description | trim) ~ '</description>' }}
60
+ {%- endif %}
61
+ {{- '\n<parameters>' }}
62
+ {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
63
+ {%- for param_name, param_fields in tool.parameters.properties|items %}
64
+ {{- '\n<parameter>' }}
65
+ {{- '\n<name>' ~ param_name ~ '</name>' }}
66
+ {%- if param_fields.type is defined %}
67
+ {{- '\n<type>' ~ (param_fields.type | string) ~ '</type>' }}
68
+ {%- endif %}
69
+ {%- if param_fields.description is defined %}
70
+ {{- '\n<description>' ~ (param_fields.description | trim) ~ '</description>' }}
71
+ {%- endif %}
72
+ {%- if param_fields.enum is defined %}
73
+ {{- '\n<enum>' ~ (param_fields.enum | tojson ) ~ '</enum>' }}
74
+ {%- endif %}
75
+ {%- set handled_keys = ['name', 'type', 'description', 'enum'] %}
76
+ {{- render_extra_keys(param_fields, handled_keys) }}
77
+ {{- '\n</parameter>' }}
78
+ {%- endfor %}
79
+ {%- endif %}
80
+ {% set handled_keys = ['type', 'properties', 'required'] %}
81
+ {{- render_extra_keys(tool.parameters, handled_keys) }}
82
+ {%- if tool.parameters is defined and tool.parameters.required is defined %}
83
+ {{- '\n<required>' ~ (tool.parameters.required | tojson ) ~ '</required>' }}
84
+ {%- endif %}
85
+ {{- '\n</parameters>' }}
86
+ {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
87
+ {{- render_extra_keys(tool, handled_keys) }}
88
+ {{- '\n</function>' }}
89
+ {%- endfor %}
90
+ {{- "\n</tools>" }}
91
+
92
+ {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
93
+ {%- endif %}
94
+
95
+
96
+ {%- if system_message is defined %}
97
+ {{- '<|im_end|>\n' }}
98
+ {%- else %}
99
+ {%- if tools is iterable and tools | length > 0 %}
100
+ {{- '<|im_end|>\n' }}
101
+ {%- endif %}
102
+ {%- endif %}
103
+
104
+ {%- for message in loop_messages %}
105
+ {%- if message.role == "assistant" %}
106
+ {# Add reasoning content in to content field for unified processing below. #}
107
+ {%- if message.reasoning_content is defined and message.reasoning_content is string and message.reasoning_content | trim | length > 0 %}
108
+ {%- set content = "<think>\n" ~ message.reasoning_content ~ "\n</think>\n" ~ (message.content | default('', true)) %}
109
+ {%- else %}
110
+ {%- set content = message.content | default('', true) %}
111
+ {%- if content is string -%}
112
+ {# Allow downstream logic to to take care of broken thought, only handle coherent reasoning here. #}
113
+ {%- if '<think>' not in content and '</think>' not in content -%}
114
+ {%- set content = "<think></think>" ~ content -%}
115
+ {%- endif -%}
116
+ {%- else -%}
117
+ {%- set content = content -%}
118
+ {%- endif -%}
119
+ {%- endif %}
120
+ {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
121
+ {# Assistant message has tool calls. #}
122
+ {{- '<|im_start|>assistant\n' }}
123
+ {%- set include_content = not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
124
+ {%- if content is string and content | trim | length > 0 %}
125
+ {%- if include_content %}
126
+ {{- (content | trim) ~ '\n' -}}
127
+ {%- else %}
128
+ {%- set c = (content | string) %}
129
+ {%- if '</think>' in c %}
130
+ {# Keep only content after the last closing think. Also generation prompt causes this. #}
131
+ {%- set c = c.split('</think>')[-1] %}
132
+ {%- elif '<think>' in c %}
133
+ {# If <think> was opened but never closed, drop the trailing think segment #}
134
+ {%- set c = c.split('<think>')[0] %}
135
+ {%- endif %}
136
+ {%- set c = "<think></think>" ~ c | trim %}
137
+ {%- if c | length > 0 %}
138
+ {{- c ~ '\n' -}}
139
+ {%- endif %}
140
+ {%- endif %}
141
+ {%- else %}
142
+ {{- "<think></think>" -}}
143
+ {%- endif %}
144
+ {%- for tool_call in message.tool_calls %}
145
+ {%- if tool_call.function is defined %}
146
+ {%- set tool_call = tool_call.function %}
147
+ {%- endif %}
148
+ {{- '<tool_call>\n<function=' ~ tool_call.name ~ '>\n' -}}
149
+ {%- if tool_call.arguments is defined %}
150
+ {%- for args_name, args_value in tool_call.arguments|items %}
151
+ {{- '<parameter=' ~ args_name ~ '>\n' -}}
152
+ {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
153
+ {{- args_value ~ '\n</parameter>\n' -}}
154
+ {%- endfor %}
155
+ {%- endif %}
156
+ {{- '</function>\n</tool_call>\n' -}}
157
+ {%- endfor %}
158
+ {{- '<|im_end|>\n' }}
159
+ {%- else %}
160
+ {# Assistant message doesn't have tool calls. #}
161
+ {%- if not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
162
+ {{- '<|im_start|>assistant\n' ~ (content | default('', true) | string | trim) ~ '<|im_end|>\n' }}
163
+ {%- else %}
164
+ {%- set c = (content | default('', true) | string) %}
165
+ {%- if '<think>' in c and '</think>' in c %}
166
+ {%- set c = "<think></think>" ~ c.split('</think>')[-1] %}
167
+ {%- endif %}
168
+ {%- set c = c | trim %}
169
+ {%- if c | length > 0 %}
170
+ {{- '<|im_start|>assistant\n' ~ c ~ '<|im_end|>\n' }}
171
+ {%- else %}
172
+ {{- '<|im_start|>assistant\n<|im_end|>\n' }}
173
+ {%- endif %}
174
+ {%- endif %}
175
+ {%- endif %}
176
+ {%- elif message.role == "user" or message.role == "system" %}
177
+ {{- '<|im_start|>' + message.role + '\n' }}
178
+ {%- set content = message.content | string %}
179
+ {{- content }}
180
+ {{- '<|im_end|>\n' }}
181
+ {%- elif message.role == "tool" %}
182
+ {%- if loop.previtem and loop.previtem.role != "tool" %}
183
+ {{- '<|im_start|>user\n' }}
184
+ {%- endif %}
185
+ {{- '<tool_response>\n' }}
186
+ {{- message.content }}
187
+ {{- '\n</tool_response>\n' }}
188
+ {%- if not loop.last and loop.nextitem.role != "tool" %}
189
+ {{- '<|im_end|>\n' }}
190
+ {%- elif loop.last %}
191
+ {{- '<|im_end|>\n' }}
192
+ {%- endif %}
193
+ {%- else %}
194
+ {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' }}
195
+ {%- endif %}
196
+ {%- endfor %}
197
+
198
+ {%- if add_generation_prompt %}
199
+ {%- if enable_thinking %}
200
+ {{- '<|im_start|>assistant\n<think>\n' }}
201
+ {%- else %}
202
+ {{- '<|im_start|>assistant\n<think></think>' }}
203
+ {%- endif %}
204
+ {%- endif %}
config.json ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "NemotronHForCausalLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "attention_head_dim": 128,
8
+ "auto_map": {
9
+ "AutoConfig": "configuration_nemotron_h.NemotronHConfig",
10
+ "AutoModelForCausalLM": "modeling_nemotron_h.NemotronHForCausalLM"
11
+ },
12
+ "bos_token_id": 1,
13
+ "chunk_size": 256,
14
+ "conv_kernel": 4,
15
+ "dtype": "bfloat16",
16
+ "eos_token_id": 2,
17
+ "expand": 2,
18
+ "head_dim": 128,
19
+ "hidden_dropout": 0.0,
20
+ "hidden_size": 3136,
21
+ "hybrid_override_pattern": "M-M-M-MM-M-M*-M-M*-M-M-M*-M-M-MM*-MMM-M-M-",
22
+ "initializer_range": 0.02,
23
+ "intermediate_size": 12544,
24
+ "layer_norm_epsilon": 1e-05,
25
+ "mamba_head_dim": 80,
26
+ "mamba_hidden_act": "silu",
27
+ "mamba_num_heads": 96,
28
+ "mamba_proj_bias": false,
29
+ "max_position_embeddings": 262144,
30
+ "mlp_bias": false,
31
+ "mlp_hidden_act": "relu2",
32
+ "model_name": "models/NVIDIA-Nemotron-3-Nano-4B",
33
+ "model_type": "nemotron_h",
34
+ "n_groups": 8,
35
+ "num_attention_heads": 40,
36
+ "num_hidden_layers": 42,
37
+ "num_key_value_heads": 8,
38
+ "num_logits_to_keep": 1,
39
+ "pad_token_id": 999,
40
+ "rescale_prenorm_residual": true,
41
+ "residual_in_fp32": false,
42
+ "rms_norm_eps": 1e-05,
43
+ "sliding_window": null,
44
+ "ssm_state_size": 128,
45
+ "tie_word_embeddings": false,
46
+ "time_step_floor": 0.0001,
47
+ "time_step_limit": [
48
+ 0.0,
49
+ {
50
+ "__float__": "Infinity"
51
+ }
52
+ ],
53
+ "time_step_max": 0.1,
54
+ "time_step_min": 0.001,
55
+ "time_step_rank": 256,
56
+ "transformers_version": "5.3.0",
57
+ "unsloth_fixed": true,
58
+ "unsloth_version": "2026.3.17",
59
+ "use_bias": false,
60
+ "use_cache": true,
61
+ "use_conv_bias": true,
62
+ "use_mamba_kernels": true,
63
+ "vocab_size": 131072
64
+ }
configuration_nemotron_h.py ADDED
@@ -0,0 +1,243 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2024 AI21 Labs Ltd. and the HuggingFace Inc. team. All rights reserved.
3
+ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+ """NemotronH model configuration"""
17
+
18
+ import re
19
+
20
+ from transformers.configuration_utils import PretrainedConfig
21
+ from transformers.utils import logging
22
+
23
+
24
+ logger = logging.get_logger(__name__)
25
+
26
+
27
+ class NemotronHConfig(PretrainedConfig):
28
+ r"""
29
+ This is the configuration class to store the configuration of a [`NemotronHModel`]. It is used to instantiate a
30
+ NemotronH model according to the specified arguments, defining the model architecture. Instantiating a configuration
31
+ with the defaults will yield a similar configuration to that of the NemotronH-v0.1 model.
32
+
33
+ [todo](todo)
34
+
35
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
36
+ documentation from [`PretrainedConfig`] for more information.
37
+
38
+
39
+ Args:
40
+ vocab_size (`int`, *optional*, defaults to 131072):
41
+ Vocabulary size of the NemotronH model. Defines the number of different tokens that can be represented by the
42
+ `inputs_ids` passed when calling [`NemotronHModel`]
43
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
44
+ Whether the model's input and output word embeddings should be tied. Note that this is only relevant if the
45
+ model has a output word embedding layer.
46
+ hidden_size (`int`, *optional*, defaults to 4096):
47
+ Dimension of the hidden representations.
48
+ intermediate_size (`int`, *optional*, defaults to 21504):
49
+ Dimension of the MLP representations.
50
+ num_hidden_layers (`int`, *optional*, defaults to 52):
51
+ Number of hidden layers in the Transformer encoder.
52
+ hybrid_override_pattern (`str`, *optional*, defaults to `"M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-"`):
53
+ The pattern of the hybrid model. The pattern is a string of characters where each character represents M: Mamba2, *: Attention, -: MLP
54
+ num_attention_heads (`int`, *optional*, defaults to 32):
55
+ Number of attention heads for each attention layer in the Transformer encoder.
56
+ attention_head_dim (`int`, *optional*, defaults to 128):
57
+ Dimension of each attention head.
58
+ num_key_value_heads (`int`, *optional*, defaults to 8):
59
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
60
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
61
+ `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used.
62
+ mlp_hidden_act (`str`, *optional*, defaults to "relu2"):
63
+ The non-linear activation function in the MLP layers.
64
+ attention_bias (`bool`, *optional*, defaults to `False`):
65
+ Whether to use bias in attention layers.
66
+ mlp_bias (`bool`, *optional*, defaults to `False`):
67
+ Whether to use bias in MLP layers.
68
+ use_bias (`bool`, *optional*, defaults to `False`):
69
+ Whether to use bias in the model.
70
+ initializer_range (`float`, *optional*, defaults to 0.02):
71
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
72
+ layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
73
+ The epsilon used by the layer normalization layers.
74
+ residual_in_fp32 (`bool`, *optional*, defaults to `False`):
75
+ Whether or not residuals should be in `float32`. If set to `False` residuals will keep the same `dtype` as the rest of the model.
76
+ use_cache (`bool`, *optional*, defaults to `True`):
77
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
78
+ relevant if `config.is_decoder=True`.
79
+ num_logits_to_keep (`int` or `None`, *optional*, defaults to 1):
80
+ Number of prompt logits to calculate during generation. If `None`, all logits will be calculated. If an
81
+ integer value, only last `num_logits_to_keep` logits will be calculated.
82
+ pad_token_id (`int`, *optional*, defaults to 0):
83
+ The id of the padding token.
84
+ bos_token_id (`int`, *optional*, defaults to 1):
85
+ The id of the "beginning-of-sequence" token.
86
+ eos_token_id (`int`, *optional*, defaults to 2):
87
+ The id of the "end-of-sequence" token.
88
+ sliding_window (`int`, *optional*, defaults to None):
89
+ Sliding window attention window size.
90
+ max_position_embeddings (`int`, *optional*, defaults to 4096):
91
+ The maximum sequence length that this model might ever be used with.
92
+ attention_dropout (`float`, *optional*, defaults to 0.0):
93
+ The dropout ratio for the attention probabilities.
94
+ hidden_dropout (`float`, *optional*, defaults to 0.0):
95
+ The dropout ratio for the hidden states.
96
+ use_mamba_kernels (`bool`, *optional*, defaults to `True`):
97
+ Flag indicating whether or not to use the fast mamba kernels. These are available only if `mamba-ssm` and
98
+ `causal-conv1d` are installed, and the mamba modules are running on a CUDA device.
99
+ ssm_state_size (`int`, *optional*, defaults to 128):
100
+ The dimension of the mamba state space latents.
101
+ mamba_num_heads (`int`, *optional*, defaults to 128):
102
+ Number of heads in Mamba layers.
103
+ mamba_n_groups (`int`, *optional*, defaults to 8):
104
+ Number of groups in Mamba layers.
105
+ mamba_head_dim (`int`, *optional*, defaults to 64):
106
+ Dimension of each Mamba head.
107
+ mamba_d_conv (`int`, *optional*, defaults to 4):
108
+ The size of the mamba convolution kernel.
109
+ mamba_expand (`int`, *optional*, defaults to 2):
110
+ Expanding factor used to determine the mamba intermediate size.
111
+ mamba_hidden_act (`str`, *optional*, defaults to "silu"):
112
+ The non-linear activation function in the Mamba layers.
113
+ mamba_dt_min (`float`, *optional*, defaults to 0.001):
114
+ Minimum value for the time step in Mamba.
115
+ mamba_dt_max (`float`, *optional*, defaults to 0.1):
116
+ Maximum value for the time step in Mamba.
117
+ mamba_dt_limit (`tuple`, *optional*, defaults to (0.0, float("inf"))):
118
+ Limits for the time step in Mamba.
119
+ mamba_dt_init_floor (`float`, *optional*, defaults to 1e-4):
120
+ Floor value for time step initialization in Mamba.
121
+ mamba_conv_bias (`bool`, *optional*, defaults to `True`):
122
+ Whether to use bias in the convolution layer of the mamba mixer block.
123
+ mamba_proj_bias (`bool`, *optional*, defaults to `False`):
124
+ Whether to use bias in the input and output projections of the mamba mixer block.
125
+ mamba_chunk_size (`int`, *optional*, defaults to 256):
126
+ Size of chunks for Mamba processing.
127
+ rescale_prenorm_residual (`bool`, *optional*, defaults to `True`):
128
+ Whether to rescale the pre-normalization residual connections.
129
+ """
130
+
131
+ model_type = "nemotron_h"
132
+ keys_to_ignore_at_inference = ["past_key_values"]
133
+
134
+ def __init__(
135
+ self,
136
+ vocab_size=131072,
137
+ tie_word_embeddings=False,
138
+ hidden_size=4096,
139
+ intermediate_size=21504,
140
+ num_hidden_layers=52,
141
+ hybrid_override_pattern="M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-",
142
+ num_attention_heads=32,
143
+ attention_head_dim=128,
144
+ num_key_value_heads=8, # nemo: num_query_groups
145
+ mlp_hidden_act="relu2",
146
+ attention_bias=False,
147
+ mlp_bias=False,
148
+ use_bias=False,
149
+ initializer_range=0.02, # nemo: init_method_std
150
+ layer_norm_epsilon=1e-5, # nemo: layernorm_epsilon
151
+ residual_in_fp32=False, # Megatron Core default value
152
+ use_cache=True,
153
+ num_logits_to_keep=1,
154
+ pad_token_id=0,
155
+ bos_token_id=1,
156
+ eos_token_id=2,
157
+ sliding_window=None,
158
+ max_position_embeddings=4096,
159
+ attention_dropout=0.0,
160
+ hidden_dropout=0.0, # * ADDED
161
+ use_mamba_kernels=True,
162
+ ssm_state_size=128, # mamba_state_size
163
+ mamba_num_heads=128,
164
+ mamba_n_groups=8, # nemo: mamba_ssm_ngroups = num_heads
165
+ mamba_head_dim=64,
166
+ mamba_d_conv=4,
167
+ mamba_expand=2,
168
+ mamba_hidden_act="silu",
169
+ mamba_dt_min=0.001,
170
+ mamba_dt_max=0.1,
171
+ mamba_dt_limit=(0.0, float("inf")),
172
+ mamba_dt_init_floor=1e-4,
173
+ mamba_conv_bias=True,
174
+ mamba_proj_bias=False,
175
+ mamba_chunk_size=256,
176
+ rescale_prenorm_residual=True,
177
+ **kwargs,
178
+ ):
179
+ self.vocab_size = vocab_size
180
+ self.tie_word_embeddings = tie_word_embeddings
181
+ self.hidden_size = hidden_size
182
+ self.intermediate_size = intermediate_size
183
+ self.num_hidden_layers = num_hidden_layers
184
+ self.hybrid_override_pattern = hybrid_override_pattern
185
+ self.num_attention_heads = num_attention_heads
186
+ self.attention_head_dim = attention_head_dim
187
+ self.sliding_window = sliding_window
188
+ self.max_position_embeddings = max_position_embeddings
189
+ self.attention_dropout = attention_dropout
190
+ self.hidden_dropout = hidden_dropout
191
+
192
+ # Validate hybrid_override_pattern
193
+ # M: Mamba2, *: Attention, -: MLP
194
+ assert len(self.hybrid_override_pattern) == self.num_hidden_layers, "hybrid_override_pattern must have the same length as num_hidden_layers"
195
+ assert re.match(r"^[*-M]+$", self.hybrid_override_pattern), "hybrid_override_pattern must only contain characters 'M', '*', or '-'"
196
+
197
+ # for backward compatibility
198
+ if num_key_value_heads is None:
199
+ num_key_value_heads = num_attention_heads
200
+
201
+ self.num_key_value_heads = num_key_value_heads
202
+ self.mlp_hidden_act = mlp_hidden_act
203
+ self.attention_bias = attention_bias
204
+ self.mlp_bias = mlp_bias
205
+ self.use_bias = use_bias
206
+ self.initializer_range = initializer_range
207
+ self.layer_norm_epsilon = layer_norm_epsilon
208
+ self.residual_in_fp32 = residual_in_fp32
209
+
210
+ self.use_cache = use_cache
211
+ self.num_logits_to_keep = num_logits_to_keep
212
+
213
+ self.use_mamba_kernels = use_mamba_kernels
214
+ self.n_groups = mamba_n_groups
215
+ self.mamba_head_dim = mamba_head_dim
216
+ self.ssm_state_size = ssm_state_size
217
+ self.mamba_num_heads = mamba_num_heads
218
+ self.conv_kernel = mamba_d_conv
219
+ self.expand = mamba_expand
220
+ self.mamba_hidden_act = mamba_hidden_act
221
+ self.time_step_min = mamba_dt_min
222
+ self.time_step_max = mamba_dt_max
223
+ self.time_step_limit = mamba_dt_limit
224
+ self.time_step_floor = mamba_dt_init_floor
225
+ self.use_conv_bias = mamba_conv_bias
226
+ self.mamba_proj_bias = mamba_proj_bias
227
+ self.chunk_size = mamba_chunk_size
228
+ self.rescale_prenorm_residual = rescale_prenorm_residual
229
+
230
+ super().__init__(
231
+ pad_token_id=pad_token_id,
232
+ bos_token_id=bos_token_id,
233
+ eos_token_id=eos_token_id,
234
+ tie_word_embeddings=tie_word_embeddings,
235
+ **kwargs,
236
+ )
237
+
238
+ @property
239
+ def layers_block_type(self):
240
+ return [
241
+ "mamba" if self.hybrid_override_pattern[i] == "M" else
242
+ "attention" if self.hybrid_override_pattern[i] == "*" else "mlp"
243
+ for i in range(self.num_hidden_layers)]
generation_config.json ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 1,
4
+ "eos_token_id": 2,
5
+ "max_length": 262144,
6
+ "pad_token_id": 0,
7
+ "transformers_version": "5.3.0"
8
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:aa25f0f7722900e9472a0e9056ed68cd98eb2d54e36323277b9a836688a5e03e
3
+ size 7947142664
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:623c34567aebb18582765289fbe23d901c62704d6518d71866e0e58db892b5b7
3
+ size 17077484
tokenizer_config.json ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_prefix_space": false,
3
+ "backend": "tokenizers",
4
+ "bos_token": "<s>",
5
+ "clean_up_tokenization_spaces": false,
6
+ "eos_token": "<|im_end|>",
7
+ "is_local": true,
8
+ "model_input_names": [
9
+ "input_ids",
10
+ "attention_mask"
11
+ ],
12
+ "model_max_length": 262144,
13
+ "pad_token": "<SPECIAL_999>",
14
+ "padding_side": "left",
15
+ "tokenizer_class": "TokenizersBackend",
16
+ "unk_token": "<unk>"
17
+ }