Atharvsinh commited on
Commit
46cc6c9
·
verified ·
1 Parent(s): eaed585

Upload folder using huggingface_hub

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip 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
 
 
33
  *.zip 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
+ tokenizer.json filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ language: en
4
+ tags:
5
+ - 0labs
6
+ - sky
7
+ - crest
8
+ - adaptive-depth
9
+ ---
10
+
11
+ # Sky v2.0 — 6.5B (CREST K=2)
12
+
13
+ **Built by 0labs — Atharvsinh Jadav, Gujarat, India**
14
+
15
+ ## Architecture: CREST (Cognitively Recurrent Estimation of Step Termination)
16
+
17
+ This model uses the **CREST adaptive-depth architecture** — our proprietary
18
+ architecture where each FFN layer has 2 independent MLPs and a learned halting
19
+ gate. The model decides per-token, per-layer how many computational steps to take.
20
+
21
+ - Easy tokens ("the", "is") → 1 step (fast)
22
+ - Hard tokens (algorithm logic, math) → 2 steps (deeper thinking)
23
+
24
+ ## Specs
25
+
26
+ | Property | Value |
27
+ |---|---|
28
+ | Architecture | CREST (0labs proprietary) |
29
+ | Parameters | 6.5B |
30
+ | CREST Steps (K) | 2 |
31
+ | Base Model Size | 4B |
32
+ | Hidden Dimension | 2,560 |
33
+ | Layers | 32 |
34
+ | Attention Heads | 32 |
35
+ | Context Window | 32,768 tokens |
36
+ | Quantization | INT4 |
37
+
38
+ ## Usage
39
+
40
+ ```python
41
+ import torch
42
+ from transformers import AutoModelForCausalLM, AutoTokenizer
43
+
44
+ model = AutoModelForCausalLM.from_pretrained(
45
+ "0labs-in/Sky-v2.0",
46
+ trust_remote_code=True, # Required for CREST architecture
47
+ torch_dtype=torch.bfloat16,
48
+ device_map="auto",
49
+ )
50
+ tokenizer = AutoTokenizer.from_pretrained("0labs-in/Sky-v2.0")
51
+ ```
52
+
53
+ ## 0labs
54
+
55
+ Independent AI research lab. One researcher. One GPU. Adaptive-depth LLMs.
56
+ - Website: https://0labs.in
57
+ - HuggingFace: https://huggingface.co/0labs-in
chat_template.jinja ADDED
@@ -0,0 +1,154 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {%- set image_count = namespace(value=0) %}
2
+ {%- set video_count = namespace(value=0) %}
3
+ {%- macro render_content(content, do_vision_count, is_system_content=false) %}
4
+ {%- if content is string %}
5
+ {{- content }}
6
+ {%- elif content is iterable and content is not mapping %}
7
+ {%- for item in content %}
8
+ {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
9
+ {%- if is_system_content %}
10
+ {{- raise_exception('System message cannot contain images.') }}
11
+ {%- endif %}
12
+ {%- if do_vision_count %}
13
+ {%- set image_count.value = image_count.value + 1 %}
14
+ {%- endif %}
15
+ {%- if add_vision_id %}
16
+ {{- 'Picture ' ~ image_count.value ~ ': ' }}
17
+ {%- endif %}
18
+ {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
19
+ {%- elif 'video' in item or item.type == 'video' %}
20
+ {%- if is_system_content %}
21
+ {{- raise_exception('System message cannot contain videos.') }}
22
+ {%- endif %}
23
+ {%- if do_vision_count %}
24
+ {%- set video_count.value = video_count.value + 1 %}
25
+ {%- endif %}
26
+ {%- if add_vision_id %}
27
+ {{- 'Video ' ~ video_count.value ~ ': ' }}
28
+ {%- endif %}
29
+ {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
30
+ {%- elif 'text' in item %}
31
+ {{- item.text }}
32
+ {%- else %}
33
+ {{- raise_exception('Unexpected item type in content.') }}
34
+ {%- endif %}
35
+ {%- endfor %}
36
+ {%- elif content is none or content is undefined %}
37
+ {{- '' }}
38
+ {%- else %}
39
+ {{- raise_exception('Unexpected content type.') }}
40
+ {%- endif %}
41
+ {%- endmacro %}
42
+ {%- if not messages %}
43
+ {{- raise_exception('No messages provided.') }}
44
+ {%- endif %}
45
+ {%- if tools and tools is iterable and tools is not mapping %}
46
+ {{- '<|im_start|>system\n' }}
47
+ {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
48
+ {%- for tool in tools %}
49
+ {{- "\n" }}
50
+ {{- tool | tojson }}
51
+ {%- endfor %}
52
+ {{- "\n</tools>" }}
53
+ {{- '\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>' }}
54
+ {%- if messages[0].role == 'system' %}
55
+ {%- set content = render_content(messages[0].content, false, true)|trim %}
56
+ {%- if content %}
57
+ {{- '\n\n' + content }}
58
+ {%- endif %}
59
+ {%- endif %}
60
+ {{- '<|im_end|>\n' }}
61
+ {%- else %}
62
+ {%- if messages[0].role == 'system' %}
63
+ {%- set content = render_content(messages[0].content, false, true)|trim %}
64
+ {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
65
+ {%- endif %}
66
+ {%- endif %}
67
+ {%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
68
+ {%- for message in messages[::-1] %}
69
+ {%- set index = (messages|length - 1) - loop.index0 %}
70
+ {%- if ns.multi_step_tool and message.role == "user" %}
71
+ {%- set content = render_content(message.content, false)|trim %}
72
+ {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
73
+ {%- set ns.multi_step_tool = false %}
74
+ {%- set ns.last_query_index = index %}
75
+ {%- endif %}
76
+ {%- endif %}
77
+ {%- endfor %}
78
+ {%- if ns.multi_step_tool %}
79
+ {{- raise_exception('No user query found in messages.') }}
80
+ {%- endif %}
81
+ {%- for message in messages %}
82
+ {%- set content = render_content(message.content, true)|trim %}
83
+ {%- if message.role == "system" %}
84
+ {%- if not loop.first %}
85
+ {{- raise_exception('System message must be at the beginning.') }}
86
+ {%- endif %}
87
+ {%- elif message.role == "user" %}
88
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
89
+ {%- elif message.role == "assistant" %}
90
+ {%- set reasoning_content = '' %}
91
+ {%- if message.reasoning_content is string %}
92
+ {%- set reasoning_content = message.reasoning_content %}
93
+ {%- else %}
94
+ {%- if '</think>' in content %}
95
+ {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
96
+ {%- set content = content.split('</think>')[-1].lstrip('\n') %}
97
+ {%- endif %}
98
+ {%- endif %}
99
+ {%- set reasoning_content = reasoning_content|trim %}
100
+ {%- if loop.index0 > ns.last_query_index %}
101
+ {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
102
+ {%- else %}
103
+ {{- '<|im_start|>' + message.role + '\n' + content }}
104
+ {%- endif %}
105
+ {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
106
+ {%- for tool_call in message.tool_calls %}
107
+ {%- if tool_call.function is defined %}
108
+ {%- set tool_call = tool_call.function %}
109
+ {%- endif %}
110
+ {%- if loop.first %}
111
+ {%- if content|trim %}
112
+ {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
113
+ {%- else %}
114
+ {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
115
+ {%- endif %}
116
+ {%- else %}
117
+ {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
118
+ {%- endif %}
119
+ {%- if tool_call.arguments is defined %}
120
+ {%- for args_name, args_value in tool_call.arguments|items %}
121
+ {{- '<parameter=' + args_name + '>\n' }}
122
+ {%- 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 %}
123
+ {{- args_value }}
124
+ {{- '\n</parameter>\n' }}
125
+ {%- endfor %}
126
+ {%- endif %}
127
+ {{- '</function>\n</tool_call>' }}
128
+ {%- endfor %}
129
+ {%- endif %}
130
+ {{- '<|im_end|>\n' }}
131
+ {%- elif message.role == "tool" %}
132
+ {%- if loop.previtem and loop.previtem.role != "tool" %}
133
+ {{- '<|im_start|>user' }}
134
+ {%- endif %}
135
+ {{- '\n<tool_response>\n' }}
136
+ {{- content }}
137
+ {{- '\n</tool_response>' }}
138
+ {%- if not loop.last and loop.nextitem.role != "tool" %}
139
+ {{- '<|im_end|>\n' }}
140
+ {%- elif loop.last %}
141
+ {{- '<|im_end|>\n' }}
142
+ {%- endif %}
143
+ {%- else %}
144
+ {{- raise_exception('Unexpected message role.') }}
145
+ {%- endif %}
146
+ {%- endfor %}
147
+ {%- if add_generation_prompt %}
148
+ {{- '<|im_start|>assistant\n' }}
149
+ {%- if enable_thinking is defined and enable_thinking is false %}
150
+ {{- '<think>\n\n</think>\n\n' }}
151
+ {%- else %}
152
+ {{- '<think>\n' }}
153
+ {%- endif %}
154
+ {%- endif %}
config.json ADDED
@@ -0,0 +1,90 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "SkyCRESTForCausalLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "attn_output_gate": true,
8
+ "auto_map": {
9
+ "AutoConfig": "configuration_sky_crest.SkyCRESTConfig",
10
+ "AutoModelForCausalLM": "modeling_sky_crest.SkyCRESTForCausalLM"
11
+ },
12
+ "bos_token_id": null,
13
+ "creator": "Atharvsinh Jadav",
14
+ "crest_max_steps": 4,
15
+ "dtype": "bfloat16",
16
+ "eos_token_id": 248044,
17
+ "full_attention_interval": 4,
18
+ "head_dim": 256,
19
+ "hidden_act": "silu",
20
+ "hidden_size": 2560,
21
+ "initializer_range": 0.02,
22
+ "intermediate_size": 9216,
23
+ "layer_types": [
24
+ "linear_attention",
25
+ "linear_attention",
26
+ "linear_attention",
27
+ "full_attention",
28
+ "linear_attention",
29
+ "linear_attention",
30
+ "linear_attention",
31
+ "full_attention",
32
+ "linear_attention",
33
+ "linear_attention",
34
+ "linear_attention",
35
+ "full_attention",
36
+ "linear_attention",
37
+ "linear_attention",
38
+ "linear_attention",
39
+ "full_attention",
40
+ "linear_attention",
41
+ "linear_attention",
42
+ "linear_attention",
43
+ "full_attention",
44
+ "linear_attention",
45
+ "linear_attention",
46
+ "linear_attention",
47
+ "full_attention",
48
+ "linear_attention",
49
+ "linear_attention",
50
+ "linear_attention",
51
+ "full_attention",
52
+ "linear_attention",
53
+ "linear_attention",
54
+ "linear_attention",
55
+ "full_attention"
56
+ ],
57
+ "linear_conv_kernel_dim": 4,
58
+ "linear_key_head_dim": 128,
59
+ "linear_num_key_heads": 16,
60
+ "linear_num_value_heads": 32,
61
+ "linear_value_head_dim": 128,
62
+ "mamba_ssm_dtype": "float32",
63
+ "max_position_embeddings": 262144,
64
+ "mlp_only_layers": [],
65
+ "model_type": "sky-crest",
66
+ "mtp_num_hidden_layers": 1,
67
+ "mtp_use_dedicated_embeddings": false,
68
+ "num_attention_heads": 16,
69
+ "num_hidden_layers": 32,
70
+ "num_key_value_heads": 4,
71
+ "organization": "0labs",
72
+ "pad_token_id": null,
73
+ "partial_rotary_factor": 0.25,
74
+ "rms_norm_eps": 1e-06,
75
+ "rope_parameters": {
76
+ "mrope_interleaved": true,
77
+ "mrope_section": [
78
+ 11,
79
+ 11,
80
+ 10
81
+ ],
82
+ "partial_rotary_factor": 0.25,
83
+ "rope_theta": 10000000,
84
+ "rope_type": "default"
85
+ },
86
+ "sky_version": "2.0",
87
+ "tie_word_embeddings": true,
88
+ "transformers_version": "5.7.0",
89
+ "vocab_size": 248320
90
+ }
configuration_sky_crest.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ """Sky CREST Configuration — 0labs"""
2
+ from transformers import PretrainedConfig
3
+
4
+ class SkyCRESTConfig(PretrainedConfig):
5
+ model_type = "sky-crest"
6
+
7
+ def __init__(self, crest_max_steps=4, **kwargs):
8
+ self.crest_max_steps = crest_max_steps
9
+ super().__init__(**kwargs)
crest_block.py ADDED
@@ -0,0 +1,213 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ CREST Block v2 — Cognitively Recurrent Estimation of Step Termination
3
+ =====================================================================
4
+ Authors: ENI & LO (0labs, Gujarat, India)
5
+ Date: April 2026 — v2 (fixed initialization stability)
6
+
7
+ Key fix from v1: Step 1 now produces IDENTICAL output to original MLP.
8
+ No extra RMSNorm or h0 residual on step 1. Steps 2+ use lightweight
9
+ residual mixing instead of full norm, preventing gradient instability.
10
+ """
11
+
12
+ import math
13
+ import copy
14
+ import torch
15
+ import torch.nn as nn
16
+ import torch.nn.functional as F
17
+ from typing import Optional, Tuple
18
+
19
+
20
+ class CRESTBlock(nn.Module):
21
+ """
22
+ CREST Block v2 — replaces a standard FFN/MLP sublayer.
23
+
24
+ v2 fixes:
25
+ - Step 1 output is IDENTICAL to original MLP (no extra norm/residual)
26
+ - Steps 2+ use learned residual gate instead of RMSNorm
27
+ - Halting bias initialized higher (6.0) for safer start
28
+ """
29
+
30
+ def __init__(
31
+ self,
32
+ original_mlp: nn.Module,
33
+ hidden_size: int,
34
+ max_steps: int = 4,
35
+ eps: float = 1e-5,
36
+ ):
37
+ super().__init__()
38
+ self.max_steps = max_steps
39
+ self.hidden_size = hidden_size
40
+ self.eps = eps
41
+
42
+ # ── Independent FFN per step ────────────────────────────────
43
+ self.steps = nn.ModuleList()
44
+ self.steps.append(original_mlp) # Step 1 = original (untouched)
45
+ for _ in range(max_steps - 1):
46
+ self.steps.append(copy.deepcopy(original_mlp))
47
+
48
+ # ── Halting mechanism ───────────────────────────────────────
49
+ self.halt_linear = nn.Linear(hidden_size, 1, bias=True)
50
+ nn.init.zeros_(self.halt_linear.weight)
51
+ nn.init.constant_(self.halt_linear.bias, 6.0) # sigmoid(6)≈0.9975
52
+
53
+ # ── Residual gates for steps 2+ ─────────────────────────────
54
+ # Learned scalar that controls how much h0 mixes into step output
55
+ # Initialized to 0 → steps 2+ start identical to step 1
56
+ self.residual_gates = nn.ParameterList([
57
+ nn.Parameter(torch.zeros(1)) for _ in range(max_steps - 1)
58
+ ])
59
+
60
+ # ── Runtime storage ─────────────────────────────────────────
61
+ self._ponder_cost = 0.0
62
+ self._steps_used = 0.0
63
+
64
+ def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
65
+ h0 = hidden_states
66
+ B, S, D = h0.shape
67
+
68
+ total_prob = torch.zeros(B, S, 1, device=h0.device, dtype=h0.dtype)
69
+ accumulated = torch.zeros_like(h0)
70
+ h = h0
71
+
72
+ ponder_cost = torch.tensor(0.0, device=h0.device, dtype=h0.dtype)
73
+ steps_taken = torch.zeros(B, S, 1, device=h0.device, dtype=h0.dtype)
74
+
75
+ for i in range(self.max_steps):
76
+ # ── Halting probability ─────────────────────────────────
77
+ p_halt = torch.sigmoid(self.halt_linear(h))
78
+ remaining = 1.0 - total_prob
79
+ p_use = torch.min(p_halt, remaining)
80
+
81
+ # ── Compute step i ──────────────────────────────────────
82
+ if i == 0:
83
+ # Step 1: EXACT same as original MLP (no modifications)
84
+ h_new = self.steps[0](h)
85
+ else:
86
+ # Steps 2+: MLP output + learned residual gate from h0
87
+ mlp_out = self.steps[i](h)
88
+ gate = torch.sigmoid(self.residual_gates[i - 1])
89
+ h_new = mlp_out + gate * h0 # gate starts at 0.5, learned
90
+
91
+ # ── Weighted accumulation ───────────────────────────────
92
+ accumulated = accumulated + p_use * h_new
93
+ total_prob = total_prob + p_use
94
+ ponder_cost = ponder_cost + p_use.mean()
95
+ steps_taken = steps_taken + (p_use > self.eps).float()
96
+
97
+ # ── Early exit ──────────────────────────────────────────
98
+ if (total_prob >= (1.0 - self.eps)).all():
99
+ break
100
+
101
+ h = h_new
102
+
103
+ # ── Distribute remaining probability ────────────────────────
104
+ remainder = 1.0 - total_prob
105
+ if remainder.max() > self.eps:
106
+ accumulated = accumulated + remainder * h
107
+
108
+ self._ponder_cost = ponder_cost
109
+ self._steps_used = steps_taken.mean().item()
110
+
111
+ return accumulated
112
+
113
+
114
+ def retrofit_model_with_crest(
115
+ model: nn.Module,
116
+ max_steps: int = 4,
117
+ target_layers: Optional[list] = None,
118
+ ) -> Tuple[nn.Module, dict]:
119
+ """
120
+ Retrofit a pretrained transformer with CREST blocks.
121
+ Original MLP weights are preserved as Step 1.
122
+ """
123
+ layers = None
124
+ if hasattr(model, 'model') and hasattr(model.model, 'layers'):
125
+ layers = model.model.layers
126
+ elif hasattr(model, 'transformer') and hasattr(model.transformer, 'h'):
127
+ layers = model.transformer.h
128
+ else:
129
+ raise ValueError("Cannot find decoder layers")
130
+
131
+ hidden_size = getattr(model.config, 'hidden_size',
132
+ getattr(model.config, 'd_model', None))
133
+ if hidden_size is None:
134
+ for p in layers[0].parameters():
135
+ hidden_size = p.shape[-1]
136
+ break
137
+
138
+ n_layers = len(layers)
139
+ if target_layers is None:
140
+ target_layers = list(range(n_layers))
141
+
142
+ params_before = sum(p.numel() for p in model.parameters())
143
+ converted = 0
144
+ skipped = 0
145
+
146
+ for idx in target_layers:
147
+ layer = layers[idx]
148
+ mlp = None
149
+ mlp_attr = None
150
+ for attr_name in ['mlp', 'feed_forward', 'ffn']:
151
+ if hasattr(layer, attr_name):
152
+ mlp = getattr(layer, attr_name)
153
+ mlp_attr = attr_name
154
+ break
155
+ if mlp is None:
156
+ skipped += 1
157
+ continue
158
+
159
+ crest_block = CRESTBlock(
160
+ original_mlp=mlp,
161
+ hidden_size=hidden_size,
162
+ max_steps=max_steps,
163
+ )
164
+ setattr(layer, mlp_attr, crest_block)
165
+ converted += 1
166
+
167
+ params_after = sum(p.numel() for p in model.parameters())
168
+
169
+ stats = {
170
+ 'total_layers': n_layers,
171
+ 'converted': converted,
172
+ 'skipped': skipped,
173
+ 'max_steps': max_steps,
174
+ 'hidden_size': hidden_size,
175
+ 'params_before': params_before,
176
+ 'params_after': params_after,
177
+ 'params_added': params_after - params_before,
178
+ 'param_overhead_pct': round(100 * (params_after - params_before) / params_before, 1),
179
+ }
180
+ return model, stats
181
+
182
+
183
+ def collect_ponder_costs(model: nn.Module) -> torch.Tensor:
184
+ """Collect ponder costs from all CREST blocks."""
185
+ total = torch.tensor(0.0, device='cpu')
186
+ count = 0
187
+ for module in model.modules():
188
+ if isinstance(module, CRESTBlock):
189
+ if isinstance(module._ponder_cost, torch.Tensor):
190
+ total = total.to(module._ponder_cost.device)
191
+ total = total + module._ponder_cost
192
+ count += 1
193
+ return total / max(count, 1)
194
+
195
+
196
+ def get_crest_stats(model: nn.Module) -> dict:
197
+ """Get monitoring stats from all CREST blocks."""
198
+ stats = []
199
+ for name, module in model.named_modules():
200
+ if isinstance(module, CRESTBlock):
201
+ stats.append({
202
+ 'name': name,
203
+ 'avg_steps': module._steps_used,
204
+ 'ponder_cost': module._ponder_cost.item() if isinstance(module._ponder_cost, torch.Tensor) else module._ponder_cost,
205
+ })
206
+ if not stats:
207
+ return {'avg_steps': 0, 'avg_ponder': 0, 'n_blocks': 0}
208
+ return {
209
+ 'avg_steps': sum(s['avg_steps'] for s in stats) / len(stats),
210
+ 'avg_ponder': sum(s['ponder_cost'] for s in stats) / len(stats),
211
+ 'n_blocks': len(stats),
212
+ 'per_block': stats,
213
+ }
generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "eos_token_id": 248044,
4
+ "transformers_version": "5.7.0",
5
+ "use_cache": true
6
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:bb36953a0c8efc78a0c7bd3ba7d59300e0ac37ff92f4bcf5699be0a279c6278a
3
+ size 12941587560
modeling_sky_crest.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Sky CREST Model — 0labs
3
+ SkyCRESTForCausalLM: Adaptive-depth language model architecture.
4
+ """
5
+ import torch
6
+ import torch.nn as nn
7
+ import transformers
8
+ from .configuration_sky_crest import SkyCRESTConfig
9
+ from .crest_block import CRESTBlock
10
+
11
+ # Dynamically resolve base architecture
12
+ _BASE_CLASSES = [
13
+ "Qwen3_5ForCausalLM",
14
+ "Qwen2ForCausalLM",
15
+ "LlamaForCausalLM",
16
+ ]
17
+
18
+ _BaseClass = None
19
+ for _name in _BASE_CLASSES:
20
+ _BaseClass = getattr(transformers, _name, None)
21
+ if _BaseClass is not None:
22
+ break
23
+
24
+ if _BaseClass is None:
25
+ raise ImportError(
26
+ "Sky v2.0 requires transformers>=4.51.0. "
27
+ "Run: pip install --upgrade transformers"
28
+ )
29
+
30
+
31
+ class SkyCRESTForCausalLM(_BaseClass):
32
+ """Sky v2.0 — Adaptive-depth language model with CREST architecture by 0labs."""
33
+ config_class = SkyCRESTConfig
34
+
35
+ def __init__(self, config):
36
+ super().__init__(config)
37
+ max_steps = getattr(config, "crest_max_steps", 4)
38
+ hidden_size = config.hidden_size
39
+
40
+ # Find the layers — handle both model.layers and model.language_model.layers
41
+ if hasattr(self.model, "language_model"):
42
+ layers = self.model.language_model.layers
43
+ else:
44
+ layers = self.model.layers
45
+
46
+ for layer in layers:
47
+ orig_mlp = layer.mlp
48
+ layer.mlp = CRESTBlock(orig_mlp, hidden_size=hidden_size, max_steps=max_steps)
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523
3
+ size 19989325
tokenizer_config.json ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_prefix_space": false,
3
+ "audio_bos_token": "<|audio_start|>",
4
+ "audio_eos_token": "<|audio_end|>",
5
+ "audio_token": "<|audio_pad|>",
6
+ "backend": "tokenizers",
7
+ "bos_token": null,
8
+ "clean_up_tokenization_spaces": false,
9
+ "eos_token": "<|im_end|>",
10
+ "errors": "replace",
11
+ "image_token": "<|image_pad|>",
12
+ "is_local": true,
13
+ "local_files_only": false,
14
+ "model_max_length": 262144,
15
+ "model_specific_special_tokens": {
16
+ "audio_bos_token": "<|audio_start|>",
17
+ "audio_eos_token": "<|audio_end|>",
18
+ "audio_token": "<|audio_pad|>",
19
+ "image_token": "<|image_pad|>",
20
+ "video_token": "<|video_pad|>",
21
+ "vision_bos_token": "<|vision_start|>",
22
+ "vision_eos_token": "<|vision_end|>"
23
+ },
24
+ "pad_token": "<|endoftext|>",
25
+ "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
26
+ "split_special_tokens": false,
27
+ "tokenizer_class": "Qwen2Tokenizer",
28
+ "unk_token": null,
29
+ "video_token": "<|video_pad|>",
30
+ "vision_bos_token": "<|vision_start|>",
31
+ "vision_eos_token": "<|vision_end|>"
32
+ }