Upload folder using huggingface_hub

.gitattributes

@@ -63,3 +63,4 @@ stage2_v3pf_qwen3_5_9b_b300/latest/__7_0.distcp filter=lfs diff=lfs merge=lfs -t
 qwen3_5_9b_stage2_v3pf_b300/vllm_compat/tokenizer.json filter=lfs diff=lfs merge=lfs -text
 qwen3_6_35b_a3b_allretr_stage1/tokenizer.json filter=lfs diff=lfs merge=lfs -text
 qwen3_6_35b_a3b_allretr_stage1/vllm_compat/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

chat_template.jinja

@@ -0,0 +1,154 @@
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\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>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- 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 %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n' }}
+    {%- endif %}
+{%- endif %}

ckpt.pt

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

loss_curve.csv

@@ -0,0 +1,121 @@
+step,loss,distill_kl,agree_kl,lr,tok_per_s
+1,38.789196,38.789196,0.000000,2.000e-05,15216.2
+7,37.839962,37.839962,0.000000,8.000e-05,7445.3
+10,37.053440,37.053440,0.000000,1.100e-04,39654.8
+17,39.608994,39.608994,0.000000,1.800e-04,7494.5
+23,40.573605,40.573605,0.000000,2.400e-04,5058.3
+26,38.859566,38.859566,0.000000,2.700e-04,18584.9
+30,30.389381,30.389381,0.000000,3.100e-04,42489.1
+38,36.713814,36.713814,0.000000,3.900e-04,4424.7
+40,35.190712,35.190712,0.000000,4.100e-04,68411.4
+46,35.047192,35.047192,0.000000,4.700e-04,17683.3
+51,35.109936,35.109936,0.000000,5.200e-04,15793.2
+57,31.554787,31.554787,0.000000,5.800e-04,11639.5
+61,29.545170,29.545170,0.000000,6.200e-04,27695.4
+66,30.060705,30.060705,0.000000,6.700e-04,13943.4
+72,26.739662,26.739662,0.000000,7.300e-04,6812.6
+77,25.327383,25.327383,0.000000,7.800e-04,6649.2
+81,25.097601,25.097601,0.000000,8.200e-04,13981.2
+86,21.446486,21.446486,0.000000,8.700e-04,12738.9
+91,22.065109,22.065109,0.000000,9.200e-04,17261.7
+98,20.093021,20.093021,0.000000,9.900e-04,4244.4
+102,20.581947,20.581947,0.000000,1.000e-03,7410.7
+107,15.226495,15.226495,0.000000,9.995e-04,7855.8
+113,13.259332,13.259332,0.000000,9.983e-04,4252.2
+118,11.778306,11.778306,0.000000,9.968e-04,4244.0
+123,10.824446,10.824446,0.000000,9.948e-04,4237.4
+128,11.554870,11.554870,0.000000,9.923e-04,6084.8
+130,9.023664,9.023664,0.000000,9.911e-04,38039.2
+136,8.738180,8.738180,0.000000,9.873e-04,16595.4
+141,6.847273,6.847273,0.000000,9.835e-04,20890.9
+148,7.188715,7.188715,0.000000,9.774e-04,4558.4
+150,7.429653,7.429653,0.000000,9.755e-04,67745.7
+155,5.146590,5.146590,0.000000,9.704e-04,43204.8
+161,6.125732,6.125732,0.000000,9.637e-04,19938.2
+165,4.732550,4.732550,0.000000,9.589e-04,48386.1
+171,4.825805,4.825805,0.000000,9.511e-04,27331.1
+178,6.297479,6.297479,0.000000,9.411e-04,4337.8
+181,4.748247,4.748247,0.000000,9.366e-04,13750.5
+186,4.462758,4.462758,0.000000,9.288e-04,20161.5
+193,5.449700,5.449700,0.000000,9.170e-04,4248.7
+198,4.322855,4.322855,0.000000,9.082e-04,4238.2
+203,4.213812,4.213812,0.000000,8.989e-04,4238.0
+207,3.799648,3.799648,0.000000,8.912e-04,7402.3
+213,3.698510,3.698510,0.000000,8.792e-04,4141.1
+216,3.461662,3.461662,0.000000,8.730e-04,21524.2
+220,3.459363,3.459363,0.000000,8.645e-04,38688.9
+226,3.435270,3.435270,0.000000,8.513e-04,31606.8
+232,4.484162,4.484162,0.000000,8.377e-04,8023.5
+238,4.498869,4.498869,0.000000,8.235e-04,4578.3
+243,4.061327,4.061327,0.000000,8.114e-04,4768.6
+248,3.419380,3.419380,0.000000,7.990e-04,4244.5
+251,3.500835,3.500835,0.000000,7.913e-04,23952.4
+256,2.942934,2.942934,0.000000,7.784e-04,13313.2
+263,3.399804,3.399804,0.000000,7.599e-04,4228.8
+268,3.223961,3.223961,0.000000,7.464e-04,4223.0
+272,3.108742,3.108742,0.000000,7.354e-04,8646.8
+278,3.738605,3.738605,0.000000,7.186e-04,4239.0
+283,2.970390,2.970390,0.000000,7.043e-04,5598.0
+286,2.732635,2.732635,0.000000,6.957e-04,27332.0
+293,3.430861,3.430861,0.000000,6.753e-04,4226.6
+296,2.607919,2.607919,0.000000,6.664e-04,15970.0
+303,2.842535,2.842535,0.000000,6.455e-04,5802.0
+308,3.380185,3.380185,0.000000,6.304e-04,5188.5
+311,2.785544,2.785544,0.000000,6.213e-04,15352.5
+316,2.626540,2.626540,0.000000,6.060e-04,13395.3
+322,2.905012,2.905012,0.000000,5.875e-04,7611.9
+326,2.531753,2.531753,0.000000,5.751e-04,23859.9
+333,2.860965,2.860965,0.000000,5.533e-04,4530.5
+337,2.684601,2.684601,0.000000,5.408e-04,7667.4
+343,2.976162,2.976162,0.000000,5.220e-04,4239.1
+346,3.125975,3.125975,0.000000,5.126e-04,31237.3
+353,3.071309,3.071309,0.000000,4.906e-04,4228.0
+358,3.240630,3.240630,0.000000,4.749e-04,4227.4
+363,3.261945,3.261945,0.000000,4.592e-04,5125.1
+367,3.260708,3.260708,0.000000,4.467e-04,11573.6
+373,2.963808,2.963808,0.000000,4.280e-04,4394.7
+376,2.565841,2.565841,0.000000,4.187e-04,13676.1
+381,2.394880,2.394880,0.000000,4.032e-04,12541.4
+386,2.355356,2.355356,0.000000,3.879e-04,17428.2
+393,2.955215,2.955215,0.000000,3.665e-04,4401.0
+397,2.899159,2.899159,0.000000,3.545e-04,8377.9
+403,3.015910,3.015910,0.000000,3.366e-04,4296.0
+407,2.538303,2.538303,0.000000,3.247e-04,8043.7
+413,3.127864,3.127864,0.000000,3.072e-04,4298.1
+418,2.998759,2.998759,0.000000,2.928e-04,4798.8
+423,3.147634,3.147634,0.000000,2.786e-04,4244.1
+428,3.075955,3.075955,0.000000,2.646e-04,4927.6
+433,2.858573,2.858573,0.000000,2.509e-04,5693.6
+435,1.998581,1.998581,0.000000,2.455e-04,36566.3
+442,3.579702,3.579702,0.000000,2.268e-04,12106.0
+445,2.221598,2.221598,0.000000,2.190e-04,55121.3
+453,2.799558,2.799558,0.000000,1.985e-04,4292.6
+458,2.736555,2.736555,0.000000,1.862e-04,4260.5
+463,2.676741,2.676741,0.000000,1.741e-04,4284.6
+467,2.582446,2.582446,0.000000,1.647e-04,11607.7
+473,2.945476,2.945476,0.000000,1.509e-04,4287.6
+476,2.599123,2.599123,0.000000,1.442e-04,24064.8
+483,3.524433,3.524433,0.000000,1.291e-04,4280.2
+488,2.868298,2.868298,0.000000,1.188e-04,4551.5
+491,2.431671,2.431671,0.000000,1.127e-04,28452.6
+496,2.452458,2.452458,0.000000,1.030e-04,18471.7
+503,2.572148,2.572148,0.000000,9.002e-05,4284.7
+508,2.780742,2.780742,0.000000,8.124e-05,5854.9
+511,2.827664,2.827664,0.000000,7.616e-05,23118.3
+517,2.799655,2.799655,0.000000,6.646e-05,11999.8
+521,2.634980,2.634980,0.000000,6.034e-05,31520.5
+528,3.433747,3.433747,0.000000,5.030e-05,5737.2
+533,2.586520,2.586520,0.000000,4.365e-05,4283.5
+536,2.465654,2.465654,0.000000,3.988e-05,12816.6
+540,2.351502,2.351502,0.000000,3.511e-05,37026.2
+546,2.347608,2.347608,0.000000,2.850e-05,26273.9
+553,2.781969,2.781969,0.000000,2.164e-05,4250.9
+558,2.570587,2.570587,0.000000,1.731e-05,4256.8
+560,2.042787,2.042787,0.000000,1.571e-05,48395.1
+568,2.783283,2.783283,0.000000,1.007e-05,4967.6
+573,2.862680,2.862680,0.000000,7.178e-06,4233.9
+577,2.632743,2.632743,0.000000,5.212e-06,6726.3
+583,2.652764,2.652764,0.000000,2.850e-06,4228.9
+586,2.522513,2.522513,0.000000,1.933e-06,17763.5
+591,2.335135,2.335135,0.000000,7.992e-07,16899.3
+596,2.357494,2.357494,0.000000,1.579e-07,13673.2

stage1_summary.json

@@ -0,0 +1,162 @@
+{
+  "model": "Qwen/Qwen3.5-9B",
+  "probe": "probe_results/mi300x_9b_allretr.json",
+  "steps": 600,
+  "lr": 0.001,
+  "seq_len": 131072,
+  "batch_size": 1,
+  "d_idx": 16,
+  "group_agree_loss": 0.0,
+  "init_from": null,
+  "chunked_teacher": true,
+  "cq": 4096,
+  "elapsed_min": 174.15655491059977,
+  "final_per_head_kl": {
+    "3": [
+      4.593803726739482,
+      3.410243195254239,
+      3.5888593552131702,
+      4.6822851397609755,
+      0.8305473322269198,
+      1.1673379760931124,
+      1.5762652097513894,
+      1.0128812176677704,
+      3.3167540997212255,
+      1.799044251458729,
+      1.7019176598556,
+      2.2712551432016848,
+      1.5196908843723618,
+      2.024793378777086,
+      0.6604239991299291,
+      2.072595749019453
+    ],
+    "7": [
+      5.075489091960391,
+      4.540986425685064,
+      4.601246872581609,
+      3.2168032301450156,
+      3.2345920711424476,
+      2.9049126370237968,
+      1.3735442423897655,
+      4.060737014068436,
+      3.7775711611351483,
+      1.4299579887127127,
+      2.266692423897654,
+      2.159260964651928,
+      3.1257064173086393,
+      3.891731070049775,
+      3.299265084159983,
+      3.5663830966608625
+    ],
+    "11": [
+      2.8785906569337447,
+      3.115237226227223,
+      2.5582499881828653,
+      2.9077316762962573,
+      2.446724132292546,
+      3.140160868301059,
+      3.124129585597826,
+      2.898979817158699,
+      3.175635906115889,
+      1.4351075455426867,
+      3.2905795589955327,
+      2.2393125698035385,
+      3.3401959624423982,
+      3.11980162878588,
+      2.9127986162410298,
+      3.0960572955141057
+    ],
+    "15": [
+      1.529403555555922,
+      2.7887688235960146,
+      2.5803037498241173,
+      2.01445731939571,
+      1.7256206159805034,
+      1.8691750801916245,
+      2.399475312357095,
+      2.83946223518626,
+      2.396328656771053,
+      2.177124830712678,
+      2.4256641435710393,
+      2.4409239803736633,
+      2.36247424963944,
+      2.753803468411425,
+      2.031825037567495,
+      2.1247666631215263
+    ],
+    "19": [
+      2.460418164419762,
+      2.0898832549555015,
+      2.4441304773792565,
+      2.5005181676076402,
+      2.0902674148836535,
+      1.8442892426425972,
+      1.5240928872117718,
+      2.337066242459019,
+      1.8528365176388595,
+      2.329931544439021,
+      2.356586808944966,
+      2.132457247393635,
+      2.6984642053499193,
+      3.890836643505083,
+      2.3112759303657486,
+      2.4218452854030357
+    ],
+    "23": [
+      2.4067274946795414,
+      2.042321143546855,
+      2.9173936952564374,
+      2.2469329039591246,
+      1.5838642627592074,
+      1.5878576641262576,
+      1.72485078463025,
+      1.7928821997953162,
+      2.2872549593490574,
+      2.1080817966936207,
+      2.5997792806607922,
+      2.074354939040075,
+      1.8014794571483221,
+      1.8537860450084425,
+      2.1193633814208264,
+      2.286988489836165
+    ],
+    "27": [
+      1.370882810713689,
+      2.225345129032116,
+      3.11398069382519,
+      1.9831379109065006,
+      1.4017502138076896,
+      1.3009589364194192,
+      1.3352709655203496,
+      1.7207389886291684,
+      1.395687748984276,
+      1.9375296630687615,
+      1.46587096472338,
+      1.380170975386283,
+      1.2181396140853649,
+      2.0743062448884024,
+      1.3117177847050003,
+      1.139523241143195
+    ],
+    "31": [
+      1.5005077245585339,
+      1.3716400004781817,
+      1.9651338881340579,
+      1.302182336187034,
+      1.0163908550906018,
+      1.0520073120580942,
+      0.8861319897663615,
+      1.7970688490990396,
+      2.420727604331557,
+      4.982886384572429,
+      2.9604456461279374,
+      1.8760843595003165,
+      2.8927373471467392,
+      3.1637061309492402,
+      2.557191186671152,
+      2.0281517293964137
+    ]
+  },
+  "final_mean_kl": 2.3574943174134333,
+  "final_max_kl": 5.075489091960391
+}

state_dict.pt

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

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,31 @@
+{
+  "add_prefix_space": false,
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "audio_token": "<|audio_pad|>",
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "image_token": "<|image_pad|>",
+  "is_local": false,
+  "model_max_length": 262144,
+  "model_specific_special_tokens": {
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "audio_token": "<|audio_pad|>",
+    "image_token": "<|image_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>"
+  },
+  "pad_token": "<|endoftext|>",
+  "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+",
+  "split_special_tokens": false,
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": null,
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>"
+}