diff --git a/.gitattributes b/.gitattributes
index 7e35f53ef70bc50484bed9b3e3ff12d30108e836..75b75828c460ae2aaf9c65ab574c5c48376a0d10 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -106,3 +106,8 @@ qwen3/srkl/checkpoint-628/tokenizer.json filter=lfs diff=lfs merge=lfs -text
qwen3/srkl/checkpoint-942/tokenizer.json filter=lfs diff=lfs merge=lfs -text
qwen3/srkl/tokenizer.json filter=lfs diff=lfs merge=lfs -text
qwen3/teacher_lora/checkpoint-1256/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+qwen3/teacher_lora/checkpoint-1570/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+qwen3/teacher_lora/checkpoint-314/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+qwen3/teacher_lora/checkpoint-628/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+qwen3/teacher_lora/checkpoint-942/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+qwen3/teacher_lora/tokenizer.json filter=lfs diff=lfs merge=lfs -text
diff --git a/qwen3/teacher_lora/checkpoint-1570/adapter_config.json b/qwen3/teacher_lora/checkpoint-1570/adapter_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..376dee7aec128c9bdb50c2cc1ddcf6a68a37ab9b
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/adapter_config.json
@@ -0,0 +1,46 @@
+{
+ "alora_invocation_tokens": null,
+ "alpha_pattern": {},
+ "arrow_config": null,
+ "auto_mapping": null,
+ "base_model_name_or_path": "Qwen/Qwen3-4B-Instruct-2507",
+ "bias": "none",
+ "corda_config": null,
+ "ensure_weight_tying": false,
+ "eva_config": null,
+ "exclude_modules": null,
+ "fan_in_fan_out": false,
+ "inference_mode": true,
+ "init_lora_weights": true,
+ "layer_replication": null,
+ "layers_pattern": null,
+ "layers_to_transform": null,
+ "loftq_config": {},
+ "lora_alpha": 64,
+ "lora_bias": false,
+ "lora_dropout": 0.1,
+ "megatron_config": null,
+ "megatron_core": "megatron.core",
+ "modules_to_save": null,
+ "peft_type": "LORA",
+ "peft_version": "0.18.1",
+ "qalora_group_size": 16,
+ "r": 32,
+ "rank_pattern": {},
+ "revision": null,
+ "target_modules": [
+ "down_proj",
+ "k_proj",
+ "up_proj",
+ "o_proj",
+ "q_proj",
+ "v_proj",
+ "gate_proj"
+ ],
+ "target_parameters": null,
+ "task_type": "CAUSAL_LM",
+ "trainable_token_indices": null,
+ "use_dora": false,
+ "use_qalora": false,
+ "use_rslora": false
+}
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-1570/adapter_model.safetensors b/qwen3/teacher_lora/checkpoint-1570/adapter_model.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..9c06f5d11bbc6b39df9ee0caf6b30c7edbdafb0e
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/adapter_model.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cefa38e0ed411f3656d30fffca3cad60870d758d78c6b5cc3d93a1d4494dc6c
+size 132188392
diff --git a/qwen3/teacher_lora/checkpoint-1570/chat_template.jinja b/qwen3/teacher_lora/checkpoint-1570/chat_template.jinja
new file mode 100644
index 0000000000000000000000000000000000000000..70adff8a08fb31e0636f618564838d4bf3c05286
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/chat_template.jinja
@@ -0,0 +1,61 @@
+{%- if tools %}
+ {{- '<|im_start|>system\n' }}
+ {%- if messages[0].role == 'system' %}
+ {{- messages[0].content + '\n\n' }}
+ {%- endif %}
+ {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }}
+ {%- for tool in tools %}
+ {{- "\n" }}
+ {{- tool | tojson }}
+ {%- endfor %}
+ {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }}
+{%- else %}
+ {%- if messages[0].role == 'system' %}
+ {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+ {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+ {%- if message.content is string %}
+ {%- set content = message.content %}
+ {%- else %}
+ {%- set content = '' %}
+ {%- endif %}
+ {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+ {%- elif message.role == "assistant" %}
+ {{- '<|im_start|>' + message.role + '\n' + content }}
+ {%- if message.tool_calls %}
+ {%- for tool_call in message.tool_calls %}
+ {%- if (loop.first and content) or (not loop.first) %}
+ {{- '\n' }}
+ {%- endif %}
+ {%- if tool_call.function %}
+ {%- set tool_call = tool_call.function %}
+ {%- endif %}
+ {{- '\n{"name": "' }}
+ {{- tool_call.name }}
+ {{- '", "arguments": ' }}
+ {%- if tool_call.arguments is string %}
+ {{- tool_call.arguments }}
+ {%- else %}
+ {{- tool_call.arguments | tojson }}
+ {%- endif %}
+ {{- '}\n' }}
+ {%- endfor %}
+ {%- endif %}
+ {{- '<|im_end|>\n' }}
+ {%- elif message.role == "tool" %}
+ {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+ {{- '<|im_start|>user' }}
+ {%- endif %}
+ {{- '\n\n' }}
+ {{- content }}
+ {{- '\n' }}
+ {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+ {{- '<|im_end|>\n' }}
+ {%- endif %}
+ {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+ {{- '<|im_start|>assistant\n' }}
+{%- endif %}
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-1570/global_step1570/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt b/qwen3/teacher_lora/checkpoint-1570/global_step1570/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt
new file mode 100644
index 0000000000000000000000000000000000000000..a2ecf9a20ab7c864528c323322480364d1691860
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/global_step1570/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7512b4dc25e66e83b9b0b1a543ef89eedc0f785b43f9a7518a33fede0b910cb5
+size 396393221
diff --git a/qwen3/teacher_lora/checkpoint-1570/global_step1570/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt b/qwen3/teacher_lora/checkpoint-1570/global_step1570/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt
new file mode 100644
index 0000000000000000000000000000000000000000..3e6d2230ca4df3a75f222316a24afce53136fa3e
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/global_step1570/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4d7e7f3d4709f8688e5a0784e30aad927fe50c664ce810b50f4f5f8c81ac64c
+size 396393349
diff --git a/qwen3/teacher_lora/checkpoint-1570/global_step1570/mp_rank_00_model_states.pt b/qwen3/teacher_lora/checkpoint-1570/global_step1570/mp_rank_00_model_states.pt
new file mode 100644
index 0000000000000000000000000000000000000000..64e302634e8ebfe6588964d708335079ec97f8f2
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/global_step1570/mp_rank_00_model_states.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe0ea5f78da12b9d509e39ad2d2d1c2b4ad2b6a0f9b448c8193ec49fe08bb204
+size 910426661
diff --git a/qwen3/teacher_lora/checkpoint-1570/latest b/qwen3/teacher_lora/checkpoint-1570/latest
new file mode 100644
index 0000000000000000000000000000000000000000..18c2fcff58fa6daa0a2e7c3a015552a69085778b
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/latest
@@ -0,0 +1 @@
+global_step1570
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-1570/rng_state_0.pth b/qwen3/teacher_lora/checkpoint-1570/rng_state_0.pth
new file mode 100644
index 0000000000000000000000000000000000000000..ff6b9233eaa520933ba2dbc76b95626ddeb034a9
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/rng_state_0.pth
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52230a3515621808db71eefe0a2b6ac9fd0cbaecf9b8c1e7c906c59fbd3868f2
+size 14917
diff --git a/qwen3/teacher_lora/checkpoint-1570/rng_state_1.pth b/qwen3/teacher_lora/checkpoint-1570/rng_state_1.pth
new file mode 100644
index 0000000000000000000000000000000000000000..4a3148fc45a494e2895748d1861cd808d37f55a4
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/rng_state_1.pth
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b7634d5573a5f553b496b96a782ce44fed076b8442474278e12c926c6ad293c6
+size 14917
diff --git a/qwen3/teacher_lora/checkpoint-1570/scheduler.pt b/qwen3/teacher_lora/checkpoint-1570/scheduler.pt
new file mode 100644
index 0000000000000000000000000000000000000000..e22954463492d8a070ae427c95418df61893c54c
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/scheduler.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6332fac17ae99a07597512c705658da2fd8429d1c39bb776b30e107c2cb7aae
+size 1465
diff --git a/qwen3/teacher_lora/checkpoint-1570/tokenizer.json b/qwen3/teacher_lora/checkpoint-1570/tokenizer.json
new file mode 100644
index 0000000000000000000000000000000000000000..c7afbed2efcdf019f88ab0572ec29d3bf595dfe2
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/tokenizer.json
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506
+size 11422650
diff --git a/qwen3/teacher_lora/checkpoint-1570/tokenizer_config.json b/qwen3/teacher_lora/checkpoint-1570/tokenizer_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..907e863a5a4109fb0e373ccdeb4d657bbeca9416
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/tokenizer_config.json
@@ -0,0 +1,31 @@
+{
+ "add_prefix_space": false,
+ "backend": "tokenizers",
+ "bos_token": null,
+ "clean_up_tokenization_spaces": false,
+ "eos_token": "<|im_end|>",
+ "errors": "replace",
+ "extra_special_tokens": [
+ "<|im_start|>",
+ "<|im_end|>",
+ "<|object_ref_start|>",
+ "<|object_ref_end|>",
+ "<|box_start|>",
+ "<|box_end|>",
+ "<|quad_start|>",
+ "<|quad_end|>",
+ "<|vision_start|>",
+ "<|vision_end|>",
+ "<|vision_pad|>",
+ "<|image_pad|>",
+ "<|video_pad|>"
+ ],
+ "is_local": false,
+ "local_files_only": false,
+ "model_max_length": 1010000,
+ "pad_token": "<|endoftext|>",
+ "padding_side": "right",
+ "split_special_tokens": false,
+ "tokenizer_class": "Qwen2Tokenizer",
+ "unk_token": null
+}
diff --git a/qwen3/teacher_lora/checkpoint-1570/trainer_state.json b/qwen3/teacher_lora/checkpoint-1570/trainer_state.json
new file mode 100644
index 0000000000000000000000000000000000000000..2930eaef2b59193a3a4240a51c1e85c32b84d181
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/trainer_state.json
@@ -0,0 +1,655 @@
+{
+ "best_global_step": null,
+ "best_metric": null,
+ "best_model_checkpoint": null,
+ "epoch": 5.0,
+ "eval_steps": 500,
+ "global_step": 1570,
+ "is_hyper_param_search": false,
+ "is_local_process_zero": true,
+ "is_world_process_zero": true,
+ "log_history": [
+ {
+ "epoch": 0.06383403151805306,
+ "grad_norm": 0.4070272743701935,
+ "learning_rate": 9.996390385474569e-05,
+ "loss": 0.37070536613464355,
+ "step": 20
+ },
+ {
+ "epoch": 0.12766806303610612,
+ "grad_norm": 0.5005771517753601,
+ "learning_rate": 9.984797509888031e-05,
+ "loss": 0.03893736004829407,
+ "step": 40
+ },
+ {
+ "epoch": 0.19150209455415917,
+ "grad_norm": 0.17237664759159088,
+ "learning_rate": 9.965229938050925e-05,
+ "loss": 0.02901872992515564,
+ "step": 60
+ },
+ {
+ "epoch": 0.25533612607221223,
+ "grad_norm": 0.17700566351413727,
+ "learning_rate": 9.937719005663432e-05,
+ "loss": 0.012938253581523895,
+ "step": 80
+ },
+ {
+ "epoch": 0.3191701575902653,
+ "grad_norm": 0.29617995023727417,
+ "learning_rate": 9.902308769000724e-05,
+ "loss": 0.01476612538099289,
+ "step": 100
+ },
+ {
+ "epoch": 0.38300418910831835,
+ "grad_norm": 0.09139569103717804,
+ "learning_rate": 9.859055934360789e-05,
+ "loss": 0.013924291729927063,
+ "step": 120
+ },
+ {
+ "epoch": 0.4468382206263714,
+ "grad_norm": 0.6957159042358398,
+ "learning_rate": 9.808029767254437e-05,
+ "loss": 0.015140479803085327,
+ "step": 140
+ },
+ {
+ "epoch": 0.5106722521444245,
+ "grad_norm": 0.19447490572929382,
+ "learning_rate": 9.749311981482817e-05,
+ "loss": 0.016634005308151244,
+ "step": 160
+ },
+ {
+ "epoch": 0.5745062836624776,
+ "grad_norm": 0.13639570772647858,
+ "learning_rate": 9.682996608280154e-05,
+ "loss": 0.008574086427688598,
+ "step": 180
+ },
+ {
+ "epoch": 0.6383403151805306,
+ "grad_norm": 0.05407094210386276,
+ "learning_rate": 9.609189845731207e-05,
+ "loss": 0.00992567539215088,
+ "step": 200
+ },
+ {
+ "epoch": 0.7021743466985837,
+ "grad_norm": 0.08482213318347931,
+ "learning_rate": 9.52800988870462e-05,
+ "loss": 0.009285379946231843,
+ "step": 220
+ },
+ {
+ "epoch": 0.7660083782166367,
+ "grad_norm": 0.07621290534734726,
+ "learning_rate": 9.439586739574513e-05,
+ "loss": 0.004985855519771576,
+ "step": 240
+ },
+ {
+ "epoch": 0.8298424097346898,
+ "grad_norm": 0.3552233576774597,
+ "learning_rate": 9.344062000033414e-05,
+ "loss": 0.014777910709381104,
+ "step": 260
+ },
+ {
+ "epoch": 0.8936764412527428,
+ "grad_norm": 0.07681854814291,
+ "learning_rate": 9.241588644329919e-05,
+ "loss": 0.0040129125118255615,
+ "step": 280
+ },
+ {
+ "epoch": 0.9575104727707959,
+ "grad_norm": 0.08489467203617096,
+ "learning_rate": 9.132330774294252e-05,
+ "loss": 0.012399718165397644,
+ "step": 300
+ },
+ {
+ "epoch": 1.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 94.6262,
+ "eval_generator_rouge1": 99.8758,
+ "eval_generator_rouge2": 99.7812,
+ "eval_generator_rougeL": 99.5997,
+ "eval_loss": 0.004354734905064106,
+ "eval_runtime": 1522.2841,
+ "eval_samples_per_second": 2.243,
+ "eval_selector_accuracy": 98.014,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.141,
+ "step": 314
+ },
+ {
+ "epoch": 1.019150209455416,
+ "grad_norm": 0.23293638229370117,
+ "learning_rate": 9.01646335654399e-05,
+ "loss": 0.011312296986579895,
+ "step": 320
+ },
+ {
+ "epoch": 1.082984240973469,
+ "grad_norm": 0.07232551276683807,
+ "learning_rate": 8.894171942290817e-05,
+ "loss": 0.006549905985593796,
+ "step": 340
+ },
+ {
+ "epoch": 1.1468182724915221,
+ "grad_norm": 0.05371030047535896,
+ "learning_rate": 8.765652370197003e-05,
+ "loss": 0.003336009383201599,
+ "step": 360
+ },
+ {
+ "epoch": 1.2106523040095751,
+ "grad_norm": 0.1448812186717987,
+ "learning_rate": 8.631110452757467e-05,
+ "loss": 0.008410868048667908,
+ "step": 380
+ },
+ {
+ "epoch": 1.2744863355276281,
+ "grad_norm": 0.0636834055185318,
+ "learning_rate": 8.490761646709642e-05,
+ "loss": 0.006242972612380981,
+ "step": 400
+ },
+ {
+ "epoch": 1.3383203670456814,
+ "grad_norm": 0.05049936845898628,
+ "learning_rate": 8.344830707998953e-05,
+ "loss": 0.00744413286447525,
+ "step": 420
+ },
+ {
+ "epoch": 1.4021543985637344,
+ "grad_norm": 0.032405462116003036,
+ "learning_rate": 8.193551331852458e-05,
+ "loss": 0.0031117375940084456,
+ "step": 440
+ },
+ {
+ "epoch": 1.4659884300817874,
+ "grad_norm": 0.13639135658740997,
+ "learning_rate": 8.037165778537033e-05,
+ "loss": 0.005259513109922409,
+ "step": 460
+ },
+ {
+ "epoch": 1.5298224615998404,
+ "grad_norm": 0.021995026618242264,
+ "learning_rate": 7.875924485401427e-05,
+ "loss": 0.004164557158946991,
+ "step": 480
+ },
+ {
+ "epoch": 1.5936564931178934,
+ "grad_norm": 0.01941438391804695,
+ "learning_rate": 7.710085665823444e-05,
+ "loss": 0.0015524964779615401,
+ "step": 500
+ },
+ {
+ "epoch": 1.6574905246359466,
+ "grad_norm": 0.1271296739578247,
+ "learning_rate": 7.539914895704521e-05,
+ "loss": 0.004233330115675926,
+ "step": 520
+ },
+ {
+ "epoch": 1.7213245561539996,
+ "grad_norm": 0.22600136697292328,
+ "learning_rate": 7.365684688173908e-05,
+ "loss": 0.0029299482703208923,
+ "step": 540
+ },
+ {
+ "epoch": 1.7851585876720528,
+ "grad_norm": 0.0755968764424324,
+ "learning_rate": 7.18767405718347e-05,
+ "loss": 0.0044687166810035706,
+ "step": 560
+ },
+ {
+ "epoch": 1.8489926191901058,
+ "grad_norm": 0.08520206063985825,
+ "learning_rate": 7.006168070692036e-05,
+ "loss": 0.005520488321781159,
+ "step": 580
+ },
+ {
+ "epoch": 1.9128266507081588,
+ "grad_norm": 0.008491385728120804,
+ "learning_rate": 6.821457394154786e-05,
+ "loss": 0.0015786593779921533,
+ "step": 600
+ },
+ {
+ "epoch": 1.9766606822262118,
+ "grad_norm": 0.009171651676297188,
+ "learning_rate": 6.633837825048768e-05,
+ "loss": 0.007294930517673492,
+ "step": 620
+ },
+ {
+ "epoch": 2.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 96.2617,
+ "eval_generator_rouge1": 99.9321,
+ "eval_generator_rouge2": 99.8752,
+ "eval_generator_rougeL": 99.7349,
+ "eval_loss": 0.0033019811380654573,
+ "eval_runtime": 1482.8547,
+ "eval_samples_per_second": 2.302,
+ "eval_selector_accuracy": 98.4813,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.144,
+ "step": 628
+ },
+ {
+ "epoch": 2.038300418910832,
+ "grad_norm": 0.0519501231610775,
+ "learning_rate": 6.443609819179926e-05,
+ "loss": 0.003084934316575527,
+ "step": 640
+ },
+ {
+ "epoch": 2.102134450428885,
+ "grad_norm": 0.051104024052619934,
+ "learning_rate": 6.251078009530256e-05,
+ "loss": 0.0026566479355096817,
+ "step": 660
+ },
+ {
+ "epoch": 2.165968481946938,
+ "grad_norm": 0.011012597940862179,
+ "learning_rate": 6.0565507184155677e-05,
+ "loss": 0.0016018826514482497,
+ "step": 680
+ },
+ {
+ "epoch": 2.229802513464991,
+ "grad_norm": 0.02783738076686859,
+ "learning_rate": 5.860339463735136e-05,
+ "loss": 0.0031667064875364305,
+ "step": 700
+ },
+ {
+ "epoch": 2.2936365449830443,
+ "grad_norm": 0.008262189105153084,
+ "learning_rate": 5.662758460103895e-05,
+ "loss": 0.0029385803267359734,
+ "step": 720
+ },
+ {
+ "epoch": 2.3574705765010973,
+ "grad_norm": 0.03108038939535618,
+ "learning_rate": 5.464124115666095e-05,
+ "loss": 0.0051560476422309875,
+ "step": 740
+ },
+ {
+ "epoch": 2.4213046080191503,
+ "grad_norm": 0.04321969673037529,
+ "learning_rate": 5.264754525396212e-05,
+ "loss": 0.0011130817234516143,
+ "step": 760
+ },
+ {
+ "epoch": 2.4851386395372033,
+ "grad_norm": 0.022080617025494576,
+ "learning_rate": 5.06496896169857e-05,
+ "loss": 0.002006329409778118,
+ "step": 780
+ },
+ {
+ "epoch": 2.5489726710552563,
+ "grad_norm": 0.04725557193160057,
+ "learning_rate": 4.865087363121384e-05,
+ "loss": 0.0021070817485451697,
+ "step": 800
+ },
+ {
+ "epoch": 2.6128067025733093,
+ "grad_norm": 0.27594882249832153,
+ "learning_rate": 4.6654298220040673e-05,
+ "loss": 0.0016123956069350243,
+ "step": 820
+ },
+ {
+ "epoch": 2.6766407340913627,
+ "grad_norm": 0.0028557360637933016,
+ "learning_rate": 4.466316071878229e-05,
+ "loss": 0.00044896998442709444,
+ "step": 840
+ },
+ {
+ "epoch": 2.7404747656094157,
+ "grad_norm": 0.03328096121549606,
+ "learning_rate": 4.2680649754432913e-05,
+ "loss": 0.001640898734331131,
+ "step": 860
+ },
+ {
+ "epoch": 2.8043087971274687,
+ "grad_norm": 0.03642182797193527,
+ "learning_rate": 4.070994013936646e-05,
+ "loss": 0.0008663022890686989,
+ "step": 880
+ },
+ {
+ "epoch": 2.8681428286455217,
+ "grad_norm": 0.06211363896727562,
+ "learning_rate": 3.8754187787161155e-05,
+ "loss": 0.000627265078946948,
+ "step": 900
+ },
+ {
+ "epoch": 2.9319768601635747,
+ "grad_norm": 0.0025077553000301123,
+ "learning_rate": 3.681652465868875e-05,
+ "loss": 0.0009966723620891571,
+ "step": 920
+ },
+ {
+ "epoch": 2.9958108916816277,
+ "grad_norm": 0.08202111721038818,
+ "learning_rate": 3.490005374656184e-05,
+ "loss": 0.0005241389386355877,
+ "step": 940
+ },
+ {
+ "epoch": 3.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 96.8458,
+ "eval_generator_rouge1": 99.9512,
+ "eval_generator_rouge2": 99.9138,
+ "eval_generator_rougeL": 99.7589,
+ "eval_loss": 0.002848332514986396,
+ "eval_runtime": 1491.6076,
+ "eval_samples_per_second": 2.289,
+ "eval_selector_accuracy": 99.1822,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.143,
+ "step": 942
+ },
+ {
+ "epoch": 3.0574506283662477,
+ "grad_norm": 0.05988731235265732,
+ "learning_rate": 3.3007844105971336e-05,
+ "loss": 0.003412053734064102,
+ "step": 960
+ },
+ {
+ "epoch": 3.1212846598843007,
+ "grad_norm": 0.0027547082863748074,
+ "learning_rate": 3.114292593987149e-05,
+ "loss": 0.0007339773699641227,
+ "step": 980
+ },
+ {
+ "epoch": 3.1851186914023537,
+ "grad_norm": 0.004056593868881464,
+ "learning_rate": 2.9308285746383408e-05,
+ "loss": 0.0007870178669691086,
+ "step": 1000
+ },
+ {
+ "epoch": 3.248952722920407,
+ "grad_norm": 0.02156253345310688,
+ "learning_rate": 2.7506861536187867e-05,
+ "loss": 0.0004239068366587162,
+ "step": 1020
+ },
+ {
+ "epoch": 3.31278675443846,
+ "grad_norm": 0.01963026262819767,
+ "learning_rate": 2.5741538127566396e-05,
+ "loss": 0.0004422076977789402,
+ "step": 1040
+ },
+ {
+ "epoch": 3.376620785956513,
+ "grad_norm": 0.0010617657098919153,
+ "learning_rate": 2.4015142526625115e-05,
+ "loss": 0.001177133247256279,
+ "step": 1060
+ },
+ {
+ "epoch": 3.440454817474566,
+ "grad_norm": 0.04265115037560463,
+ "learning_rate": 2.233043940009967e-05,
+ "loss": 0.00048412443138659,
+ "step": 1080
+ },
+ {
+ "epoch": 3.504288848992619,
+ "grad_norm": 0.0026085481513291597,
+ "learning_rate": 2.0690126647990974e-05,
+ "loss": 0.0007659945171326399,
+ "step": 1100
+ },
+ {
+ "epoch": 3.568122880510672,
+ "grad_norm": 0.008790579624474049,
+ "learning_rate": 1.9096831083121828e-05,
+ "loss": 0.000472747627645731,
+ "step": 1120
+ },
+ {
+ "epoch": 3.631956912028725,
+ "grad_norm": 0.010375800542533398,
+ "learning_rate": 1.7553104224533324e-05,
+ "loss": 0.0004488838836550713,
+ "step": 1140
+ },
+ {
+ "epoch": 3.695790943546778,
+ "grad_norm": 0.005610933993011713,
+ "learning_rate": 1.606141821145735e-05,
+ "loss": 0.0006687684915959835,
+ "step": 1160
+ },
+ {
+ "epoch": 3.759624975064831,
+ "grad_norm": 0.01078067161142826,
+ "learning_rate": 1.4624161844408763e-05,
+ "loss": 0.0008440790697932243,
+ "step": 1180
+ },
+ {
+ "epoch": 3.8234590065828846,
+ "grad_norm": 0.009699724614620209,
+ "learning_rate": 1.324363675973698e-05,
+ "loss": 0.0008271764032542706,
+ "step": 1200
+ },
+ {
+ "epoch": 3.8872930381009376,
+ "grad_norm": 0.018169160932302475,
+ "learning_rate": 1.1922053743763052e-05,
+ "loss": 0.0003760921768844128,
+ "step": 1220
+ },
+ {
+ "epoch": 3.9511270696189906,
+ "grad_norm": 0.010761883109807968,
+ "learning_rate": 1.0661529192404973e-05,
+ "loss": 0.0003943405114114285,
+ "step": 1240
+ },
+ {
+ "epoch": 4.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 96.7874,
+ "eval_generator_rouge1": 99.9518,
+ "eval_generator_rouge2": 99.9142,
+ "eval_generator_rougeL": 99.7472,
+ "eval_loss": 0.0028539097402244806,
+ "eval_runtime": 1492.5205,
+ "eval_samples_per_second": 2.287,
+ "eval_selector_accuracy": 99.1822,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.143,
+ "step": 1256
+ },
+ {
+ "epoch": 4.012766806303611,
+ "grad_norm": 0.0212729349732399,
+ "learning_rate": 9.464081721960592e-06,
+ "loss": 0.0008496158756315708,
+ "step": 1260
+ },
+ {
+ "epoch": 4.076600837821664,
+ "grad_norm": 0.04593608155846596,
+ "learning_rate": 8.331628936475818e-06,
+ "loss": 0.00023648960050195455,
+ "step": 1280
+ },
+ {
+ "epoch": 4.140434869339717,
+ "grad_norm": 0.00037404155591502786,
+ "learning_rate": 7.2659843568748385e-06,
+ "loss": 0.00027941586449742317,
+ "step": 1300
+ },
+ {
+ "epoch": 4.20426890085777,
+ "grad_norm": 0.0028630236629396677,
+ "learning_rate": 6.268854516770008e-06,
+ "loss": 0.00032952893525362015,
+ "step": 1320
+ },
+ {
+ "epoch": 4.268102932375823,
+ "grad_norm": 0.03912129998207092,
+ "learning_rate": 5.341836229602331e-06,
+ "loss": 0.0003575342707335949,
+ "step": 1340
+ },
+ {
+ "epoch": 4.331936963893876,
+ "grad_norm": 0.0010673116194084287,
+ "learning_rate": 4.4864140314888216e-06,
+ "loss": 0.0003374688792973757,
+ "step": 1360
+ },
+ {
+ "epoch": 4.395770995411929,
+ "grad_norm": 0.0022425639908760786,
+ "learning_rate": 3.703957803871932e-06,
+ "loss": 0.0008516624569892883,
+ "step": 1380
+ },
+ {
+ "epoch": 4.459605026929982,
+ "grad_norm": 0.01993773877620697,
+ "learning_rate": 2.9957205797780495e-06,
+ "loss": 0.000931937713176012,
+ "step": 1400
+ },
+ {
+ "epoch": 4.5234390584480355,
+ "grad_norm": 0.027315137907862663,
+ "learning_rate": 2.362836537198199e-06,
+ "loss": 0.0003644102718681097,
+ "step": 1420
+ },
+ {
+ "epoch": 4.5872730899660885,
+ "grad_norm": 0.007133164908736944,
+ "learning_rate": 1.8063191828043703e-06,
+ "loss": 0.00021162929479032755,
+ "step": 1440
+ },
+ {
+ "epoch": 4.6511071214841415,
+ "grad_norm": 0.02018481120467186,
+ "learning_rate": 1.3270597289100437e-06,
+ "loss": 0.00022698421962559224,
+ "step": 1460
+ },
+ {
+ "epoch": 4.7149411530021945,
+ "grad_norm": 0.014296010136604309,
+ "learning_rate": 9.258256662741205e-07,
+ "loss": 0.00022367050405591726,
+ "step": 1480
+ },
+ {
+ "epoch": 4.7787751845202475,
+ "grad_norm": 0.007890012115240097,
+ "learning_rate": 6.032595350337341e-07,
+ "loss": 0.00043175043538212776,
+ "step": 1500
+ },
+ {
+ "epoch": 4.8426092160383005,
+ "grad_norm": 0.022568073123693466,
+ "learning_rate": 3.5987789573421356e-07,
+ "loss": 0.00024233334697782994,
+ "step": 1520
+ },
+ {
+ "epoch": 4.9064432475563535,
+ "grad_norm": 0.0283685103058815,
+ "learning_rate": 1.9607050210401192e-07,
+ "loss": 0.0002710998523980379,
+ "step": 1540
+ },
+ {
+ "epoch": 4.9702772790744065,
+ "grad_norm": 0.019454164430499077,
+ "learning_rate": 1.1209967689925832e-07,
+ "loss": 0.00033448461908847094,
+ "step": 1560
+ },
+ {
+ "epoch": 5.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 97.0794,
+ "eval_generator_rouge1": 99.9483,
+ "eval_generator_rouge2": 99.9036,
+ "eval_generator_rougeL": 99.7893,
+ "eval_loss": 0.00291790091432631,
+ "eval_runtime": 1494.8716,
+ "eval_samples_per_second": 2.284,
+ "eval_selector_accuracy": 99.2407,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.143,
+ "step": 1570
+ }
+ ],
+ "logging_steps": 20,
+ "max_steps": 1570,
+ "num_input_tokens_seen": 0,
+ "num_train_epochs": 5,
+ "save_steps": 500,
+ "stateful_callbacks": {
+ "TrainerControl": {
+ "args": {
+ "should_epoch_stop": false,
+ "should_evaluate": false,
+ "should_log": false,
+ "should_save": true,
+ "should_training_stop": true
+ },
+ "attributes": {}
+ }
+ },
+ "total_flos": 3.981241098143007e+17,
+ "train_batch_size": 1,
+ "trial_name": null,
+ "trial_params": null
+}
diff --git a/qwen3/teacher_lora/checkpoint-1570/training_args.bin b/qwen3/teacher_lora/checkpoint-1570/training_args.bin
new file mode 100644
index 0000000000000000000000000000000000000000..1ef98d1a171b439f2f8968f31ccbc846a421fe48
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/training_args.bin
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:84fda57546f783ec338bd3a292c7dba3d917a8a52ac59ca7cebdd546da52b586
+size 7185
diff --git a/qwen3/teacher_lora/checkpoint-1570/zero_to_fp32.py b/qwen3/teacher_lora/checkpoint-1570/zero_to_fp32.py
new file mode 100644
index 0000000000000000000000000000000000000000..f40f1f0acbbff30e8e8d29906b56aae4e8b8a138
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-1570/zero_to_fp32.py
@@ -0,0 +1,790 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example:
+# python zero_to_fp32.py . output_dir/
+# or
+# python zero_to_fp32.py . output_dir/ --safe_serialization
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+import gc
+import json
+import numpy as np
+from tqdm import tqdm
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+ FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+ FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS, AUTOEP_LAYERS_KEY,
+ AUTOEP_LAYERS_KEY_LEGACY, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY,
+ AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT)
+
+
+@dataclass
+class zero_model_state:
+ buffers: dict()
+ param_shapes: dict()
+ shared_params: list
+ ds_version: int
+ frozen_param_shapes: dict()
+ frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+ return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+ '''
+ alist.sort(key=natural_keys) sorts in human order
+ http://nedbatchelder.com/blog/200712/human_sorting.html
+ (See Toothy's implementation in the comments)
+ '''
+ return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+ if not os.path.isdir(checkpoint_dir):
+ raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+ # there should be only one file
+ if zero_stage <= 2:
+ file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+ elif zero_stage == 3:
+ file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+ if not os.path.exists(file):
+ raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+ return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+ # XXX: need to test that this simple glob rule works for multi-node setup too
+ ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+ if len(ckpt_files) == 0:
+ raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+ return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+ return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+ return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def _has_autoep_zero3_partitioned_metadata(state_dict):
+ autoep_layers = state_dict.get(AUTOEP_LAYERS_KEY)
+ if autoep_layers is None:
+ autoep_layers = state_dict.get(AUTOEP_LAYERS_KEY_LEGACY)
+ if not isinstance(autoep_layers, list):
+ return False
+ return any(
+ isinstance(entry, dict)
+ and entry.get(AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY) == AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT
+ for entry in autoep_layers)
+
+
+def _raise_if_autoep_zero3_partitioned_state(state_dict):
+ if _has_autoep_zero3_partitioned_metadata(state_dict):
+ raise NotImplementedError("zero_to_fp32 does not support AutoEP ZeRO-3 partition-native checkpoints. "
+ "AutoEP expert parameters are partitioned over expert replica groups, so "
+ "global data-parallel consolidation would produce incomplete expert tensors. "
+ "Use ds_to_universal.py for expert-aware conversion.")
+
+
+def _raise_if_autoep_zero3_partitioned_checkpoint(model_files):
+ for file in model_files:
+ state_dict = torch.load(file, map_location=device, weights_only=False)
+ _raise_if_autoep_zero3_partitioned_state(state_dict)
+
+
+def parse_model_states(files):
+ zero_model_states = []
+ for file in files:
+ state_dict = torch.load(file, map_location=device, weights_only=False)
+ _raise_if_autoep_zero3_partitioned_state(state_dict)
+
+ if BUFFER_NAMES not in state_dict:
+ raise ValueError(f"{file} is not a model state checkpoint")
+ buffer_names = state_dict[BUFFER_NAMES]
+ if debug:
+ print("Found buffers:", buffer_names)
+
+ # recover just the buffers while restoring them to fp32 if they were saved in fp16
+ buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+ param_shapes = state_dict[PARAM_SHAPES]
+
+ # collect parameters that are included in param_shapes
+ param_names = []
+ for s in param_shapes:
+ for name in s.keys():
+ param_names.append(name)
+
+ # update with frozen parameters
+ frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+ if frozen_param_shapes is not None:
+ if debug:
+ print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+ param_names += list(frozen_param_shapes.keys())
+
+ # handle shared params
+ shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+ ds_version = state_dict.get(DS_VERSION, None)
+
+ frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+ z_model_state = zero_model_state(buffers=buffers,
+ param_shapes=param_shapes,
+ shared_params=shared_params,
+ ds_version=ds_version,
+ frozen_param_shapes=frozen_param_shapes,
+ frozen_param_fragments=frozen_param_fragments)
+ zero_model_states.append(z_model_state)
+
+ return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+ total_files = len(files)
+ state_dicts = []
+ for f in tqdm(files, desc='Loading checkpoint shards'):
+ state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
+ # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
+ # and also handle the case where it was already removed by another helper script
+ state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
+ state_dicts.append(state_dict)
+
+ if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
+ raise ValueError(f"{files[0]} is not a zero checkpoint")
+ zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+ world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+ # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+ # parameters can be different from data parallelism for non-expert parameters. So we can just
+ # use the max of the partition_count to get the dp world_size.
+
+ if type(world_size) is list:
+ world_size = max(world_size)
+
+ if world_size != total_files:
+ raise ValueError(
+ f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+ "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+ )
+
+ # the groups are named differently in each stage
+ if zero_stage <= 2:
+ fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+ elif zero_stage == 3:
+ fp32_groups_key = FP32_FLAT_GROUPS
+ else:
+ raise ValueError(f"unknown zero stage {zero_stage}")
+
+ fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+ return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
+ """
+ Returns fp32 state_dict reconstructed from ds checkpoint
+
+ Args:
+ - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+ """
+ print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+ # parse_model_states rejects AutoEP ZeRO-3 partition-native checkpoints
+ # before the expensive optimizer-shard load below.
+ model_files = get_model_state_files(ds_checkpoint_dir)
+ zero_model_states = parse_model_states(model_files)
+ print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+ optim_files = get_optim_files(ds_checkpoint_dir)
+ zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+ print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+ if zero_stage <= 2:
+ return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters)
+ elif zero_stage == 3:
+ return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+ return
+
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+ frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+ if debug:
+ num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+ print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+ wanted_params = len(frozen_param_shapes)
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+ avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+ print(f'Frozen params: Have {avail_numel} numels to process.')
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+ total_params = 0
+ total_numel = 0
+ for name, shape in frozen_param_shapes.items():
+ total_params += 1
+ unpartitioned_numel = shape.numel()
+ total_numel += unpartitioned_numel
+
+ state_dict[name] = frozen_param_fragments[name]
+
+ if debug:
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _has_callable(obj, fn):
+ attr = getattr(obj, fn, None)
+ return callable(attr)
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+ param_shapes = zero_model_states[0].param_shapes
+
+ # Reconstruction protocol:
+ #
+ # XXX: document this
+
+ if debug:
+ for i in range(world_size):
+ for j in range(len(fp32_flat_groups[0])):
+ print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+ # XXX: memory usage doubles here (zero2)
+ num_param_groups = len(fp32_flat_groups[0])
+ merged_single_partition_of_fp32_groups = []
+ for i in range(num_param_groups):
+ merged_partitions = [sd[i] for sd in fp32_flat_groups]
+ full_single_fp32_vector = torch.cat(merged_partitions, 0)
+ merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+ avail_numel = sum(
+ [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+ if debug:
+ wanted_params = sum([len(shapes) for shapes in param_shapes])
+ wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+ # not asserting if there is a mismatch due to possible padding
+ print(f"Have {avail_numel} numels to process.")
+ print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+ # params
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+ # out-of-core computing solution
+ total_numel = 0
+ total_params = 0
+ for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+ offset = 0
+ avail_numel = full_single_fp32_vector.numel()
+ for name, shape in shapes.items():
+
+ unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
+ total_numel += unpartitioned_numel
+ total_params += 1
+
+ if debug:
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+ state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+ offset += unpartitioned_numel
+
+ # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+ # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+ # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+ # live optimizer object, so we are checking that the numbers are within the right range
+ align_to = 2 * world_size
+
+ def zero2_align(x):
+ return align_to * math.ceil(x / align_to)
+
+ if debug:
+ print(f"original offset={offset}, avail_numel={avail_numel}")
+
+ offset = zero2_align(offset)
+ avail_numel = zero2_align(avail_numel)
+
+ if debug:
+ print(f"aligned offset={offset}, avail_numel={avail_numel}")
+
+ # Sanity check
+ if offset != avail_numel:
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+ print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters):
+ state_dict = OrderedDict()
+
+ # buffers
+ buffers = zero_model_states[0].buffers
+ state_dict.update(buffers)
+ if debug:
+ print(f"added {len(buffers)} buffers")
+
+ if not exclude_frozen_parameters:
+ _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+ _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+ # recover shared parameters
+ for pair in zero_model_states[0].shared_params:
+ if pair[1] in state_dict:
+ state_dict[pair[0]] = state_dict[pair[1]]
+
+ return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+ remainder = unpartitioned_numel % world_size
+ padding_numel = (world_size - remainder) if remainder else 0
+ partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+ return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+ return
+
+ if debug:
+ for i in range(world_size):
+ num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+ print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+ wanted_params = len(frozen_param_shapes)
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+ avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+ print(f'Frozen params: Have {avail_numel} numels to process.')
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+ total_params = 0
+ total_numel = 0
+ for name, shape in zero_model_states[0].frozen_param_shapes.items():
+ total_params += 1
+ unpartitioned_numel = shape.numel()
+ total_numel += unpartitioned_numel
+
+ param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+ state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+ if debug:
+ print(
+ f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+ )
+
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+class GatheredTensor:
+ """
+ A pseudo tensor that collects partitioned weights.
+ It is more memory efficient when there are multiple groups.
+ """
+
+ def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
+ self.flat_groups = flat_groups
+ self.flat_groups_offset = flat_groups_offset
+ self.offset = offset
+ self.partitioned_numel = partitioned_numel
+ self.shape = shape
+ self.dtype = self.flat_groups[0][0].dtype
+
+ def contiguous(self):
+ """
+ Merge partitioned weights from flat_groups into a single tensor.
+ """
+ end_idx = self.offset + self.partitioned_numel
+ world_size = len(self.flat_groups)
+ pad_flat_param_chunks = []
+
+ for rank_i in range(world_size):
+ # for each rank, we need to collect weights from related group/groups
+ flat_groups_at_rank_i = self.flat_groups[rank_i]
+ start_group_id = None
+ end_group_id = None
+ for group_id in range(len(self.flat_groups_offset)):
+ if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
+ start_group_id = group_id
+ if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
+ end_group_id = group_id
+ break
+ # collect weights from related group/groups
+ for group_id in range(start_group_id, end_group_id + 1):
+ flat_tensor = flat_groups_at_rank_i[group_id]
+ start_offset = self.offset - self.flat_groups_offset[group_id]
+ end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
+ pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
+
+ # collect weights from all ranks
+ pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
+ param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
+ return param
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+ param_shapes = zero_model_states[0].param_shapes
+ avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
+
+ # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+ # param, re-consolidating each param, while dealing with padding if any
+
+ # merge list of dicts, preserving order
+ param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+ if debug:
+ for i in range(world_size):
+ print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+ wanted_params = len(param_shapes)
+ wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+ # not asserting if there is a mismatch due to possible padding
+ avail_numel = fp32_flat_groups[0].numel() * world_size
+ print(f"Trainable params: Have {avail_numel} numels to process.")
+ print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+ # params
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+ # out-of-core computing solution
+ offset = 0
+ total_numel = 0
+ total_params = 0
+ flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
+ for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
+ unpartitioned_numel = shape.numel()
+ total_numel += unpartitioned_numel
+ total_params += 1
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+ if debug:
+ print(
+ f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+ )
+
+ # memory efficient tensor
+ tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
+ state_dict[name] = tensor
+ offset += partitioned_numel
+
+ offset *= world_size
+
+ # Sanity check
+ if offset != avail_numel:
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+ print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters):
+ state_dict = OrderedDict()
+
+ # buffers
+ buffers = zero_model_states[0].buffers
+ state_dict.update(buffers)
+ if debug:
+ print(f"added {len(buffers)} buffers")
+
+ if not exclude_frozen_parameters:
+ _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+ _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+ # recover shared parameters
+ for pair in zero_model_states[0].shared_params:
+ if pair[1] in state_dict:
+ state_dict[pair[0]] = state_dict[pair[1]]
+
+ return state_dict
+
+
+def to_torch_tensor(state_dict, return_empty_tensor=False):
+ """
+ Convert state_dict of GatheredTensor to torch tensor
+ """
+ torch_state_dict = {}
+ converted_tensors = {}
+ for name, tensor in state_dict.items():
+ tensor_id = id(tensor)
+ if tensor_id in converted_tensors: # shared tensors
+ shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
+ torch_state_dict[name] = shared_tensor
+ else:
+ converted_tensors[tensor_id] = name
+ if return_empty_tensor:
+ torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
+ else:
+ torch_state_dict[name] = tensor.contiguous()
+ return torch_state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+ tag=None,
+ exclude_frozen_parameters=False,
+ lazy_mode=False):
+ """
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+ ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+ via a model hub.
+
+ Args:
+ - ``checkpoint_dir``: path to the desired checkpoint folder
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+ - ``exclude_frozen_parameters``: exclude frozen parameters
+ - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
+ Convert the pesduo tensor to torch tensor by ``.contiguous()``
+
+ Returns:
+ - pytorch ``state_dict``
+
+ A typical usage might be ::
+
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+ # do the training and checkpoint saving
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+ model = model.cpu() # move to cpu
+ model.load_state_dict(state_dict)
+ # submit to model hub or save the model to share with others
+
+ In this example the ``model`` will no longer be usable in the deepspeed context of the same
+ application. i.e. you will need to re-initialize the deepspeed engine, since
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+ If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+ Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
+ You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+ the checkpoint. Or you can load state_dict in lazy mode ::
+
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
+ for name, lazy_tensor in state_dict.item():
+ tensor = lazy_tensor.contiguous() # to cpu
+ print(name, tensor)
+ # del tensor to release memory if it no longer in use
+ """
+ if tag is None:
+ latest_path = os.path.join(checkpoint_dir, 'latest')
+ if os.path.isfile(latest_path):
+ with open(latest_path, 'r') as fd:
+ tag = fd.read().strip()
+ else:
+ raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+ ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+ if not os.path.isdir(ds_checkpoint_dir):
+ raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+ state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
+ if lazy_mode:
+ return state_dict
+ else:
+ return to_torch_tensor(state_dict)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
+ output_dir,
+ max_shard_size="5GB",
+ safe_serialization=False,
+ tag=None,
+ exclude_frozen_parameters=False):
+ """
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+ loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+ Args:
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+ - ``output_dir``: directory to the pytorch fp32 state_dict output files
+ - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
+ - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+ - ``exclude_frozen_parameters``: exclude frozen parameters
+ """
+
+ # Dependency pre-check
+ if safe_serialization:
+ try:
+ from safetensors.torch import save_file
+ except ImportError:
+ print('If you want to use `safe_serialization`, please `pip install safetensors`')
+ raise
+ if max_shard_size is not None:
+ try:
+ from huggingface_hub import split_torch_state_dict_into_shards
+ except ImportError:
+ print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
+ raise
+
+ # Convert zero checkpoint to state_dict
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+ tag,
+ exclude_frozen_parameters,
+ lazy_mode=True)
+
+ # Shard the model if it is too big.
+ weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
+ if max_shard_size is not None:
+ filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
+ # an memory-efficient approach for sharding
+ empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
+ state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
+ filename_pattern=filename_pattern,
+ max_shard_size=max_shard_size)
+ else:
+ from collections import namedtuple
+ StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
+ state_dict_split = StateDictSplit(is_sharded=False,
+ filename_to_tensors={weights_name: list(state_dict.keys())})
+
+ # Save the model by shard
+ os.makedirs(output_dir, exist_ok=True)
+ filename_to_tensors = state_dict_split.filename_to_tensors.items()
+ for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
+ shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
+ shard_state_dict = to_torch_tensor(shard_state_dict)
+ output_path = os.path.join(output_dir, shard_file)
+ if safe_serialization:
+ save_file(shard_state_dict, output_path, metadata={"format": "pt"})
+ else:
+ torch.save(shard_state_dict, output_path)
+ # release the memory of current shard
+ for tensor_name in list(shard_state_dict.keys()):
+ del state_dict[tensor_name]
+ del shard_state_dict[tensor_name]
+ del shard_state_dict
+ gc.collect()
+
+ # Save index if sharded
+ if state_dict_split.is_sharded:
+ index = {
+ "metadata": state_dict_split.metadata,
+ "weight_map": state_dict_split.tensor_to_filename,
+ }
+ save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
+ save_index_file = os.path.join(output_dir, save_index_file)
+ with open(save_index_file, "w", encoding="utf-8") as f:
+ content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+ f.write(content)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+ """
+ 1. Put the provided model to cpu
+ 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+ 3. Load it into the provided model
+
+ Args:
+ - ``model``: the model object to update
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+ Returns:
+ - ``model`: modified model
+
+ Make sure you have plenty of CPU memory available before you call this function. If you don't
+ have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+ conveniently placed for you in the checkpoint folder.
+
+ A typical usage might be ::
+
+ from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+ model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+ # submit to model hub or save the model to share with others
+
+ Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+ of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+ """
+ logger.info("Extracting fp32 weights")
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+ logger.info("Overwriting model with fp32 weights")
+ model = model.cpu()
+ model.load_state_dict(state_dict, strict=False)
+
+ return model
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("checkpoint_dir",
+ type=str,
+ help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+ parser.add_argument("output_dir",
+ type=str,
+ help="directory to the pytorch fp32 state_dict output files"
+ "(e.g. path/checkpoint-12-output/)")
+ parser.add_argument(
+ "--max_shard_size",
+ type=str,
+ default="5GB",
+ help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
+ "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
+ "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
+ "without CPU OOM issues.")
+ parser.add_argument(
+ "--safe_serialization",
+ default=False,
+ action='store_true',
+ help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
+ parser.add_argument("-t",
+ "--tag",
+ type=str,
+ default=None,
+ help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
+ parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
+ parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+ args = parser.parse_args()
+
+ debug = args.debug
+
+ convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
+ args.output_dir,
+ max_shard_size=args.max_shard_size,
+ safe_serialization=args.safe_serialization,
+ tag=args.tag,
+ exclude_frozen_parameters=args.exclude_frozen_parameters)
diff --git a/qwen3/teacher_lora/checkpoint-314/README.md b/qwen3/teacher_lora/checkpoint-314/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..4fabf4f222ce5d86630bcfbcae8e90fb08d87580
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/README.md
@@ -0,0 +1,208 @@
+---
+base_model: Qwen/Qwen3-4B-Instruct-2507
+library_name: peft
+pipeline_tag: text-generation
+tags:
+- base_model:adapter:Qwen/Qwen3-4B-Instruct-2507
+- llama-factory
+- lora
+- transformers
+---
+
+# Model Card for Model ID
+
+
+
+
+
+## Model Details
+
+### Model Description
+
+
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+
+
+### Direct Use
+
+
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+
+
+[More Information Needed]
+
+### Recommendations
+
+
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+
+
+[More Information Needed]
+
+### Training Procedure
+
+
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed]
+
+#### Speeds, Sizes, Times [optional]
+
+
+
+[More Information Needed]
+
+## Evaluation
+
+
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+
+
+[More Information Needed]
+
+#### Factors
+
+
+
+[More Information Needed]
+
+#### Metrics
+
+
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+
+
+[More Information Needed]
+
+## Environmental Impact
+
+
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.18.1
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-314/adapter_config.json b/qwen3/teacher_lora/checkpoint-314/adapter_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..376dee7aec128c9bdb50c2cc1ddcf6a68a37ab9b
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/adapter_config.json
@@ -0,0 +1,46 @@
+{
+ "alora_invocation_tokens": null,
+ "alpha_pattern": {},
+ "arrow_config": null,
+ "auto_mapping": null,
+ "base_model_name_or_path": "Qwen/Qwen3-4B-Instruct-2507",
+ "bias": "none",
+ "corda_config": null,
+ "ensure_weight_tying": false,
+ "eva_config": null,
+ "exclude_modules": null,
+ "fan_in_fan_out": false,
+ "inference_mode": true,
+ "init_lora_weights": true,
+ "layer_replication": null,
+ "layers_pattern": null,
+ "layers_to_transform": null,
+ "loftq_config": {},
+ "lora_alpha": 64,
+ "lora_bias": false,
+ "lora_dropout": 0.1,
+ "megatron_config": null,
+ "megatron_core": "megatron.core",
+ "modules_to_save": null,
+ "peft_type": "LORA",
+ "peft_version": "0.18.1",
+ "qalora_group_size": 16,
+ "r": 32,
+ "rank_pattern": {},
+ "revision": null,
+ "target_modules": [
+ "down_proj",
+ "k_proj",
+ "up_proj",
+ "o_proj",
+ "q_proj",
+ "v_proj",
+ "gate_proj"
+ ],
+ "target_parameters": null,
+ "task_type": "CAUSAL_LM",
+ "trainable_token_indices": null,
+ "use_dora": false,
+ "use_qalora": false,
+ "use_rslora": false
+}
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-314/adapter_model.safetensors b/qwen3/teacher_lora/checkpoint-314/adapter_model.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..dd42041682e5ffca18fe0e8da652c2ee4dfc6253
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/adapter_model.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0010cb9c35bff8086f79e3aaeb297822f66779826389e7dac5f44f25fc44f940
+size 132188392
diff --git a/qwen3/teacher_lora/checkpoint-314/chat_template.jinja b/qwen3/teacher_lora/checkpoint-314/chat_template.jinja
new file mode 100644
index 0000000000000000000000000000000000000000..70adff8a08fb31e0636f618564838d4bf3c05286
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/chat_template.jinja
@@ -0,0 +1,61 @@
+{%- if tools %}
+ {{- '<|im_start|>system\n' }}
+ {%- if messages[0].role == 'system' %}
+ {{- messages[0].content + '\n\n' }}
+ {%- endif %}
+ {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }}
+ {%- for tool in tools %}
+ {{- "\n" }}
+ {{- tool | tojson }}
+ {%- endfor %}
+ {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }}
+{%- else %}
+ {%- if messages[0].role == 'system' %}
+ {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+ {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+ {%- if message.content is string %}
+ {%- set content = message.content %}
+ {%- else %}
+ {%- set content = '' %}
+ {%- endif %}
+ {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+ {%- elif message.role == "assistant" %}
+ {{- '<|im_start|>' + message.role + '\n' + content }}
+ {%- if message.tool_calls %}
+ {%- for tool_call in message.tool_calls %}
+ {%- if (loop.first and content) or (not loop.first) %}
+ {{- '\n' }}
+ {%- endif %}
+ {%- if tool_call.function %}
+ {%- set tool_call = tool_call.function %}
+ {%- endif %}
+ {{- '\n{"name": "' }}
+ {{- tool_call.name }}
+ {{- '", "arguments": ' }}
+ {%- if tool_call.arguments is string %}
+ {{- tool_call.arguments }}
+ {%- else %}
+ {{- tool_call.arguments | tojson }}
+ {%- endif %}
+ {{- '}\n' }}
+ {%- endfor %}
+ {%- endif %}
+ {{- '<|im_end|>\n' }}
+ {%- elif message.role == "tool" %}
+ {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+ {{- '<|im_start|>user' }}
+ {%- endif %}
+ {{- '\n\n' }}
+ {{- content }}
+ {{- '\n' }}
+ {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+ {{- '<|im_end|>\n' }}
+ {%- endif %}
+ {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+ {{- '<|im_start|>assistant\n' }}
+{%- endif %}
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-314/global_step314/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt b/qwen3/teacher_lora/checkpoint-314/global_step314/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt
new file mode 100644
index 0000000000000000000000000000000000000000..212ed4bd25241d1ebf3d0388221976a92b50eff5
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/global_step314/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:631d21e0022f5d0b44e2c7214c58adc2af980b28f9ce3cc0318c3ff963e56c87
+size 396393221
diff --git a/qwen3/teacher_lora/checkpoint-314/global_step314/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt b/qwen3/teacher_lora/checkpoint-314/global_step314/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt
new file mode 100644
index 0000000000000000000000000000000000000000..3711766cff969d9339e4c005c1804b8af72f7429
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/global_step314/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f59fffe84be95e12a757b2a33699fd6d1843a6ccf88c6dbf73f706dd69170d5a
+size 396393349
diff --git a/qwen3/teacher_lora/checkpoint-314/global_step314/mp_rank_00_model_states.pt b/qwen3/teacher_lora/checkpoint-314/global_step314/mp_rank_00_model_states.pt
new file mode 100644
index 0000000000000000000000000000000000000000..468ec8a1599f3cac21f2b6d3156bd93b8ec344cd
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/global_step314/mp_rank_00_model_states.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b9020ef7756dc5dbab37599230bd3d6d2ef747f7628dfb111ea8cca090c0918c
+size 910426661
diff --git a/qwen3/teacher_lora/checkpoint-314/latest b/qwen3/teacher_lora/checkpoint-314/latest
new file mode 100644
index 0000000000000000000000000000000000000000..a9786513b1a514db2c3a9da7408844635663d55d
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/latest
@@ -0,0 +1 @@
+global_step314
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-314/rng_state_0.pth b/qwen3/teacher_lora/checkpoint-314/rng_state_0.pth
new file mode 100644
index 0000000000000000000000000000000000000000..cce25075ed3e63bc1f6659a81a0df8826c86f938
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/rng_state_0.pth
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:30eb6aea8f635464d0f11026c8e459d262ee20ddfeb98e3c344f7e5c74394c1d
+size 14917
diff --git a/qwen3/teacher_lora/checkpoint-314/rng_state_1.pth b/qwen3/teacher_lora/checkpoint-314/rng_state_1.pth
new file mode 100644
index 0000000000000000000000000000000000000000..7c3ccd064e81df3c5f852d3bfe571cb1210f5503
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/rng_state_1.pth
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c987cd378df86ec81de5dbf5ad531cda963a7b2ddf82632503a6eaca00033cce
+size 14917
diff --git a/qwen3/teacher_lora/checkpoint-314/scheduler.pt b/qwen3/teacher_lora/checkpoint-314/scheduler.pt
new file mode 100644
index 0000000000000000000000000000000000000000..a3cb460802cda03bf0833a98a6f1428a3fd4f585
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/scheduler.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:17e254b5642bb75983fd5105bec79b35376cfbc99d29702287668b8767d46149
+size 1465
diff --git a/qwen3/teacher_lora/checkpoint-314/tokenizer.json b/qwen3/teacher_lora/checkpoint-314/tokenizer.json
new file mode 100644
index 0000000000000000000000000000000000000000..c7afbed2efcdf019f88ab0572ec29d3bf595dfe2
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/tokenizer.json
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506
+size 11422650
diff --git a/qwen3/teacher_lora/checkpoint-314/tokenizer_config.json b/qwen3/teacher_lora/checkpoint-314/tokenizer_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..907e863a5a4109fb0e373ccdeb4d657bbeca9416
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/tokenizer_config.json
@@ -0,0 +1,31 @@
+{
+ "add_prefix_space": false,
+ "backend": "tokenizers",
+ "bos_token": null,
+ "clean_up_tokenization_spaces": false,
+ "eos_token": "<|im_end|>",
+ "errors": "replace",
+ "extra_special_tokens": [
+ "<|im_start|>",
+ "<|im_end|>",
+ "<|object_ref_start|>",
+ "<|object_ref_end|>",
+ "<|box_start|>",
+ "<|box_end|>",
+ "<|quad_start|>",
+ "<|quad_end|>",
+ "<|vision_start|>",
+ "<|vision_end|>",
+ "<|vision_pad|>",
+ "<|image_pad|>",
+ "<|video_pad|>"
+ ],
+ "is_local": false,
+ "local_files_only": false,
+ "model_max_length": 1010000,
+ "pad_token": "<|endoftext|>",
+ "padding_side": "right",
+ "split_special_tokens": false,
+ "tokenizer_class": "Qwen2Tokenizer",
+ "unk_token": null
+}
diff --git a/qwen3/teacher_lora/checkpoint-314/trainer_state.json b/qwen3/teacher_lora/checkpoint-314/trainer_state.json
new file mode 100644
index 0000000000000000000000000000000000000000..19dffe5fdae4faf52050b48214ccc68024ffe177
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/trainer_state.json
@@ -0,0 +1,154 @@
+{
+ "best_global_step": null,
+ "best_metric": null,
+ "best_model_checkpoint": null,
+ "epoch": 1.0,
+ "eval_steps": 500,
+ "global_step": 314,
+ "is_hyper_param_search": false,
+ "is_local_process_zero": true,
+ "is_world_process_zero": true,
+ "log_history": [
+ {
+ "epoch": 0.06383403151805306,
+ "grad_norm": 0.4070272743701935,
+ "learning_rate": 9.996390385474569e-05,
+ "loss": 0.37070536613464355,
+ "step": 20
+ },
+ {
+ "epoch": 0.12766806303610612,
+ "grad_norm": 0.5005771517753601,
+ "learning_rate": 9.984797509888031e-05,
+ "loss": 0.03893736004829407,
+ "step": 40
+ },
+ {
+ "epoch": 0.19150209455415917,
+ "grad_norm": 0.17237664759159088,
+ "learning_rate": 9.965229938050925e-05,
+ "loss": 0.02901872992515564,
+ "step": 60
+ },
+ {
+ "epoch": 0.25533612607221223,
+ "grad_norm": 0.17700566351413727,
+ "learning_rate": 9.937719005663432e-05,
+ "loss": 0.012938253581523895,
+ "step": 80
+ },
+ {
+ "epoch": 0.3191701575902653,
+ "grad_norm": 0.29617995023727417,
+ "learning_rate": 9.902308769000724e-05,
+ "loss": 0.01476612538099289,
+ "step": 100
+ },
+ {
+ "epoch": 0.38300418910831835,
+ "grad_norm": 0.09139569103717804,
+ "learning_rate": 9.859055934360789e-05,
+ "loss": 0.013924291729927063,
+ "step": 120
+ },
+ {
+ "epoch": 0.4468382206263714,
+ "grad_norm": 0.6957159042358398,
+ "learning_rate": 9.808029767254437e-05,
+ "loss": 0.015140479803085327,
+ "step": 140
+ },
+ {
+ "epoch": 0.5106722521444245,
+ "grad_norm": 0.19447490572929382,
+ "learning_rate": 9.749311981482817e-05,
+ "loss": 0.016634005308151244,
+ "step": 160
+ },
+ {
+ "epoch": 0.5745062836624776,
+ "grad_norm": 0.13639570772647858,
+ "learning_rate": 9.682996608280154e-05,
+ "loss": 0.008574086427688598,
+ "step": 180
+ },
+ {
+ "epoch": 0.6383403151805306,
+ "grad_norm": 0.05407094210386276,
+ "learning_rate": 9.609189845731207e-05,
+ "loss": 0.00992567539215088,
+ "step": 200
+ },
+ {
+ "epoch": 0.7021743466985837,
+ "grad_norm": 0.08482213318347931,
+ "learning_rate": 9.52800988870462e-05,
+ "loss": 0.009285379946231843,
+ "step": 220
+ },
+ {
+ "epoch": 0.7660083782166367,
+ "grad_norm": 0.07621290534734726,
+ "learning_rate": 9.439586739574513e-05,
+ "loss": 0.004985855519771576,
+ "step": 240
+ },
+ {
+ "epoch": 0.8298424097346898,
+ "grad_norm": 0.3552233576774597,
+ "learning_rate": 9.344062000033414e-05,
+ "loss": 0.014777910709381104,
+ "step": 260
+ },
+ {
+ "epoch": 0.8936764412527428,
+ "grad_norm": 0.07681854814291,
+ "learning_rate": 9.241588644329919e-05,
+ "loss": 0.0040129125118255615,
+ "step": 280
+ },
+ {
+ "epoch": 0.9575104727707959,
+ "grad_norm": 0.08489467203617096,
+ "learning_rate": 9.132330774294252e-05,
+ "loss": 0.012399718165397644,
+ "step": 300
+ },
+ {
+ "epoch": 1.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 94.6262,
+ "eval_generator_rouge1": 99.8758,
+ "eval_generator_rouge2": 99.7812,
+ "eval_generator_rougeL": 99.5997,
+ "eval_loss": 0.004354734905064106,
+ "eval_runtime": 1522.2841,
+ "eval_samples_per_second": 2.243,
+ "eval_selector_accuracy": 98.014,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.141,
+ "step": 314
+ }
+ ],
+ "logging_steps": 20,
+ "max_steps": 1570,
+ "num_input_tokens_seen": 0,
+ "num_train_epochs": 5,
+ "save_steps": 500,
+ "stateful_callbacks": {
+ "TrainerControl": {
+ "args": {
+ "should_epoch_stop": false,
+ "should_evaluate": false,
+ "should_log": false,
+ "should_save": true,
+ "should_training_stop": false
+ },
+ "attributes": {}
+ }
+ },
+ "total_flos": 7.962272500547584e+16,
+ "train_batch_size": 1,
+ "trial_name": null,
+ "trial_params": null
+}
diff --git a/qwen3/teacher_lora/checkpoint-314/training_args.bin b/qwen3/teacher_lora/checkpoint-314/training_args.bin
new file mode 100644
index 0000000000000000000000000000000000000000..1ef98d1a171b439f2f8968f31ccbc846a421fe48
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/training_args.bin
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:84fda57546f783ec338bd3a292c7dba3d917a8a52ac59ca7cebdd546da52b586
+size 7185
diff --git a/qwen3/teacher_lora/checkpoint-314/zero_to_fp32.py b/qwen3/teacher_lora/checkpoint-314/zero_to_fp32.py
new file mode 100644
index 0000000000000000000000000000000000000000..f40f1f0acbbff30e8e8d29906b56aae4e8b8a138
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-314/zero_to_fp32.py
@@ -0,0 +1,790 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example:
+# python zero_to_fp32.py . output_dir/
+# or
+# python zero_to_fp32.py . output_dir/ --safe_serialization
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+import gc
+import json
+import numpy as np
+from tqdm import tqdm
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+ FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+ FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS, AUTOEP_LAYERS_KEY,
+ AUTOEP_LAYERS_KEY_LEGACY, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY,
+ AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT)
+
+
+@dataclass
+class zero_model_state:
+ buffers: dict()
+ param_shapes: dict()
+ shared_params: list
+ ds_version: int
+ frozen_param_shapes: dict()
+ frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+ return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+ '''
+ alist.sort(key=natural_keys) sorts in human order
+ http://nedbatchelder.com/blog/200712/human_sorting.html
+ (See Toothy's implementation in the comments)
+ '''
+ return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+ if not os.path.isdir(checkpoint_dir):
+ raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+ # there should be only one file
+ if zero_stage <= 2:
+ file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+ elif zero_stage == 3:
+ file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+ if not os.path.exists(file):
+ raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+ return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+ # XXX: need to test that this simple glob rule works for multi-node setup too
+ ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+ if len(ckpt_files) == 0:
+ raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+ return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+ return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+ return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def _has_autoep_zero3_partitioned_metadata(state_dict):
+ autoep_layers = state_dict.get(AUTOEP_LAYERS_KEY)
+ if autoep_layers is None:
+ autoep_layers = state_dict.get(AUTOEP_LAYERS_KEY_LEGACY)
+ if not isinstance(autoep_layers, list):
+ return False
+ return any(
+ isinstance(entry, dict)
+ and entry.get(AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY) == AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT
+ for entry in autoep_layers)
+
+
+def _raise_if_autoep_zero3_partitioned_state(state_dict):
+ if _has_autoep_zero3_partitioned_metadata(state_dict):
+ raise NotImplementedError("zero_to_fp32 does not support AutoEP ZeRO-3 partition-native checkpoints. "
+ "AutoEP expert parameters are partitioned over expert replica groups, so "
+ "global data-parallel consolidation would produce incomplete expert tensors. "
+ "Use ds_to_universal.py for expert-aware conversion.")
+
+
+def _raise_if_autoep_zero3_partitioned_checkpoint(model_files):
+ for file in model_files:
+ state_dict = torch.load(file, map_location=device, weights_only=False)
+ _raise_if_autoep_zero3_partitioned_state(state_dict)
+
+
+def parse_model_states(files):
+ zero_model_states = []
+ for file in files:
+ state_dict = torch.load(file, map_location=device, weights_only=False)
+ _raise_if_autoep_zero3_partitioned_state(state_dict)
+
+ if BUFFER_NAMES not in state_dict:
+ raise ValueError(f"{file} is not a model state checkpoint")
+ buffer_names = state_dict[BUFFER_NAMES]
+ if debug:
+ print("Found buffers:", buffer_names)
+
+ # recover just the buffers while restoring them to fp32 if they were saved in fp16
+ buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+ param_shapes = state_dict[PARAM_SHAPES]
+
+ # collect parameters that are included in param_shapes
+ param_names = []
+ for s in param_shapes:
+ for name in s.keys():
+ param_names.append(name)
+
+ # update with frozen parameters
+ frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+ if frozen_param_shapes is not None:
+ if debug:
+ print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+ param_names += list(frozen_param_shapes.keys())
+
+ # handle shared params
+ shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+ ds_version = state_dict.get(DS_VERSION, None)
+
+ frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+ z_model_state = zero_model_state(buffers=buffers,
+ param_shapes=param_shapes,
+ shared_params=shared_params,
+ ds_version=ds_version,
+ frozen_param_shapes=frozen_param_shapes,
+ frozen_param_fragments=frozen_param_fragments)
+ zero_model_states.append(z_model_state)
+
+ return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+ total_files = len(files)
+ state_dicts = []
+ for f in tqdm(files, desc='Loading checkpoint shards'):
+ state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
+ # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
+ # and also handle the case where it was already removed by another helper script
+ state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
+ state_dicts.append(state_dict)
+
+ if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
+ raise ValueError(f"{files[0]} is not a zero checkpoint")
+ zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+ world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+ # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+ # parameters can be different from data parallelism for non-expert parameters. So we can just
+ # use the max of the partition_count to get the dp world_size.
+
+ if type(world_size) is list:
+ world_size = max(world_size)
+
+ if world_size != total_files:
+ raise ValueError(
+ f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+ "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+ )
+
+ # the groups are named differently in each stage
+ if zero_stage <= 2:
+ fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+ elif zero_stage == 3:
+ fp32_groups_key = FP32_FLAT_GROUPS
+ else:
+ raise ValueError(f"unknown zero stage {zero_stage}")
+
+ fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+ return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
+ """
+ Returns fp32 state_dict reconstructed from ds checkpoint
+
+ Args:
+ - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+ """
+ print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+ # parse_model_states rejects AutoEP ZeRO-3 partition-native checkpoints
+ # before the expensive optimizer-shard load below.
+ model_files = get_model_state_files(ds_checkpoint_dir)
+ zero_model_states = parse_model_states(model_files)
+ print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+ optim_files = get_optim_files(ds_checkpoint_dir)
+ zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+ print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+ if zero_stage <= 2:
+ return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters)
+ elif zero_stage == 3:
+ return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+ return
+
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+ frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+ if debug:
+ num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+ print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+ wanted_params = len(frozen_param_shapes)
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+ avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+ print(f'Frozen params: Have {avail_numel} numels to process.')
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+ total_params = 0
+ total_numel = 0
+ for name, shape in frozen_param_shapes.items():
+ total_params += 1
+ unpartitioned_numel = shape.numel()
+ total_numel += unpartitioned_numel
+
+ state_dict[name] = frozen_param_fragments[name]
+
+ if debug:
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _has_callable(obj, fn):
+ attr = getattr(obj, fn, None)
+ return callable(attr)
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+ param_shapes = zero_model_states[0].param_shapes
+
+ # Reconstruction protocol:
+ #
+ # XXX: document this
+
+ if debug:
+ for i in range(world_size):
+ for j in range(len(fp32_flat_groups[0])):
+ print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+ # XXX: memory usage doubles here (zero2)
+ num_param_groups = len(fp32_flat_groups[0])
+ merged_single_partition_of_fp32_groups = []
+ for i in range(num_param_groups):
+ merged_partitions = [sd[i] for sd in fp32_flat_groups]
+ full_single_fp32_vector = torch.cat(merged_partitions, 0)
+ merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+ avail_numel = sum(
+ [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+ if debug:
+ wanted_params = sum([len(shapes) for shapes in param_shapes])
+ wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+ # not asserting if there is a mismatch due to possible padding
+ print(f"Have {avail_numel} numels to process.")
+ print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+ # params
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+ # out-of-core computing solution
+ total_numel = 0
+ total_params = 0
+ for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+ offset = 0
+ avail_numel = full_single_fp32_vector.numel()
+ for name, shape in shapes.items():
+
+ unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
+ total_numel += unpartitioned_numel
+ total_params += 1
+
+ if debug:
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+ state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+ offset += unpartitioned_numel
+
+ # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+ # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+ # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+ # live optimizer object, so we are checking that the numbers are within the right range
+ align_to = 2 * world_size
+
+ def zero2_align(x):
+ return align_to * math.ceil(x / align_to)
+
+ if debug:
+ print(f"original offset={offset}, avail_numel={avail_numel}")
+
+ offset = zero2_align(offset)
+ avail_numel = zero2_align(avail_numel)
+
+ if debug:
+ print(f"aligned offset={offset}, avail_numel={avail_numel}")
+
+ # Sanity check
+ if offset != avail_numel:
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+ print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters):
+ state_dict = OrderedDict()
+
+ # buffers
+ buffers = zero_model_states[0].buffers
+ state_dict.update(buffers)
+ if debug:
+ print(f"added {len(buffers)} buffers")
+
+ if not exclude_frozen_parameters:
+ _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+ _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+ # recover shared parameters
+ for pair in zero_model_states[0].shared_params:
+ if pair[1] in state_dict:
+ state_dict[pair[0]] = state_dict[pair[1]]
+
+ return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+ remainder = unpartitioned_numel % world_size
+ padding_numel = (world_size - remainder) if remainder else 0
+ partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+ return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+ return
+
+ if debug:
+ for i in range(world_size):
+ num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+ print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+ wanted_params = len(frozen_param_shapes)
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+ avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+ print(f'Frozen params: Have {avail_numel} numels to process.')
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+ total_params = 0
+ total_numel = 0
+ for name, shape in zero_model_states[0].frozen_param_shapes.items():
+ total_params += 1
+ unpartitioned_numel = shape.numel()
+ total_numel += unpartitioned_numel
+
+ param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+ state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+ if debug:
+ print(
+ f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+ )
+
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+class GatheredTensor:
+ """
+ A pseudo tensor that collects partitioned weights.
+ It is more memory efficient when there are multiple groups.
+ """
+
+ def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
+ self.flat_groups = flat_groups
+ self.flat_groups_offset = flat_groups_offset
+ self.offset = offset
+ self.partitioned_numel = partitioned_numel
+ self.shape = shape
+ self.dtype = self.flat_groups[0][0].dtype
+
+ def contiguous(self):
+ """
+ Merge partitioned weights from flat_groups into a single tensor.
+ """
+ end_idx = self.offset + self.partitioned_numel
+ world_size = len(self.flat_groups)
+ pad_flat_param_chunks = []
+
+ for rank_i in range(world_size):
+ # for each rank, we need to collect weights from related group/groups
+ flat_groups_at_rank_i = self.flat_groups[rank_i]
+ start_group_id = None
+ end_group_id = None
+ for group_id in range(len(self.flat_groups_offset)):
+ if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
+ start_group_id = group_id
+ if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
+ end_group_id = group_id
+ break
+ # collect weights from related group/groups
+ for group_id in range(start_group_id, end_group_id + 1):
+ flat_tensor = flat_groups_at_rank_i[group_id]
+ start_offset = self.offset - self.flat_groups_offset[group_id]
+ end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
+ pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
+
+ # collect weights from all ranks
+ pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
+ param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
+ return param
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+ param_shapes = zero_model_states[0].param_shapes
+ avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
+
+ # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+ # param, re-consolidating each param, while dealing with padding if any
+
+ # merge list of dicts, preserving order
+ param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+ if debug:
+ for i in range(world_size):
+ print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+ wanted_params = len(param_shapes)
+ wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+ # not asserting if there is a mismatch due to possible padding
+ avail_numel = fp32_flat_groups[0].numel() * world_size
+ print(f"Trainable params: Have {avail_numel} numels to process.")
+ print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+ # params
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+ # out-of-core computing solution
+ offset = 0
+ total_numel = 0
+ total_params = 0
+ flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
+ for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
+ unpartitioned_numel = shape.numel()
+ total_numel += unpartitioned_numel
+ total_params += 1
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+ if debug:
+ print(
+ f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+ )
+
+ # memory efficient tensor
+ tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
+ state_dict[name] = tensor
+ offset += partitioned_numel
+
+ offset *= world_size
+
+ # Sanity check
+ if offset != avail_numel:
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+ print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters):
+ state_dict = OrderedDict()
+
+ # buffers
+ buffers = zero_model_states[0].buffers
+ state_dict.update(buffers)
+ if debug:
+ print(f"added {len(buffers)} buffers")
+
+ if not exclude_frozen_parameters:
+ _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+ _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+ # recover shared parameters
+ for pair in zero_model_states[0].shared_params:
+ if pair[1] in state_dict:
+ state_dict[pair[0]] = state_dict[pair[1]]
+
+ return state_dict
+
+
+def to_torch_tensor(state_dict, return_empty_tensor=False):
+ """
+ Convert state_dict of GatheredTensor to torch tensor
+ """
+ torch_state_dict = {}
+ converted_tensors = {}
+ for name, tensor in state_dict.items():
+ tensor_id = id(tensor)
+ if tensor_id in converted_tensors: # shared tensors
+ shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
+ torch_state_dict[name] = shared_tensor
+ else:
+ converted_tensors[tensor_id] = name
+ if return_empty_tensor:
+ torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
+ else:
+ torch_state_dict[name] = tensor.contiguous()
+ return torch_state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+ tag=None,
+ exclude_frozen_parameters=False,
+ lazy_mode=False):
+ """
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+ ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+ via a model hub.
+
+ Args:
+ - ``checkpoint_dir``: path to the desired checkpoint folder
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+ - ``exclude_frozen_parameters``: exclude frozen parameters
+ - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
+ Convert the pesduo tensor to torch tensor by ``.contiguous()``
+
+ Returns:
+ - pytorch ``state_dict``
+
+ A typical usage might be ::
+
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+ # do the training and checkpoint saving
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+ model = model.cpu() # move to cpu
+ model.load_state_dict(state_dict)
+ # submit to model hub or save the model to share with others
+
+ In this example the ``model`` will no longer be usable in the deepspeed context of the same
+ application. i.e. you will need to re-initialize the deepspeed engine, since
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+ If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+ Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
+ You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+ the checkpoint. Or you can load state_dict in lazy mode ::
+
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
+ for name, lazy_tensor in state_dict.item():
+ tensor = lazy_tensor.contiguous() # to cpu
+ print(name, tensor)
+ # del tensor to release memory if it no longer in use
+ """
+ if tag is None:
+ latest_path = os.path.join(checkpoint_dir, 'latest')
+ if os.path.isfile(latest_path):
+ with open(latest_path, 'r') as fd:
+ tag = fd.read().strip()
+ else:
+ raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+ ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+ if not os.path.isdir(ds_checkpoint_dir):
+ raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+ state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
+ if lazy_mode:
+ return state_dict
+ else:
+ return to_torch_tensor(state_dict)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
+ output_dir,
+ max_shard_size="5GB",
+ safe_serialization=False,
+ tag=None,
+ exclude_frozen_parameters=False):
+ """
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+ loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+ Args:
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+ - ``output_dir``: directory to the pytorch fp32 state_dict output files
+ - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
+ - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+ - ``exclude_frozen_parameters``: exclude frozen parameters
+ """
+
+ # Dependency pre-check
+ if safe_serialization:
+ try:
+ from safetensors.torch import save_file
+ except ImportError:
+ print('If you want to use `safe_serialization`, please `pip install safetensors`')
+ raise
+ if max_shard_size is not None:
+ try:
+ from huggingface_hub import split_torch_state_dict_into_shards
+ except ImportError:
+ print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
+ raise
+
+ # Convert zero checkpoint to state_dict
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+ tag,
+ exclude_frozen_parameters,
+ lazy_mode=True)
+
+ # Shard the model if it is too big.
+ weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
+ if max_shard_size is not None:
+ filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
+ # an memory-efficient approach for sharding
+ empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
+ state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
+ filename_pattern=filename_pattern,
+ max_shard_size=max_shard_size)
+ else:
+ from collections import namedtuple
+ StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
+ state_dict_split = StateDictSplit(is_sharded=False,
+ filename_to_tensors={weights_name: list(state_dict.keys())})
+
+ # Save the model by shard
+ os.makedirs(output_dir, exist_ok=True)
+ filename_to_tensors = state_dict_split.filename_to_tensors.items()
+ for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
+ shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
+ shard_state_dict = to_torch_tensor(shard_state_dict)
+ output_path = os.path.join(output_dir, shard_file)
+ if safe_serialization:
+ save_file(shard_state_dict, output_path, metadata={"format": "pt"})
+ else:
+ torch.save(shard_state_dict, output_path)
+ # release the memory of current shard
+ for tensor_name in list(shard_state_dict.keys()):
+ del state_dict[tensor_name]
+ del shard_state_dict[tensor_name]
+ del shard_state_dict
+ gc.collect()
+
+ # Save index if sharded
+ if state_dict_split.is_sharded:
+ index = {
+ "metadata": state_dict_split.metadata,
+ "weight_map": state_dict_split.tensor_to_filename,
+ }
+ save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
+ save_index_file = os.path.join(output_dir, save_index_file)
+ with open(save_index_file, "w", encoding="utf-8") as f:
+ content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+ f.write(content)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+ """
+ 1. Put the provided model to cpu
+ 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+ 3. Load it into the provided model
+
+ Args:
+ - ``model``: the model object to update
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+ Returns:
+ - ``model`: modified model
+
+ Make sure you have plenty of CPU memory available before you call this function. If you don't
+ have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+ conveniently placed for you in the checkpoint folder.
+
+ A typical usage might be ::
+
+ from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+ model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+ # submit to model hub or save the model to share with others
+
+ Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+ of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+ """
+ logger.info("Extracting fp32 weights")
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+ logger.info("Overwriting model with fp32 weights")
+ model = model.cpu()
+ model.load_state_dict(state_dict, strict=False)
+
+ return model
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("checkpoint_dir",
+ type=str,
+ help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+ parser.add_argument("output_dir",
+ type=str,
+ help="directory to the pytorch fp32 state_dict output files"
+ "(e.g. path/checkpoint-12-output/)")
+ parser.add_argument(
+ "--max_shard_size",
+ type=str,
+ default="5GB",
+ help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
+ "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
+ "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
+ "without CPU OOM issues.")
+ parser.add_argument(
+ "--safe_serialization",
+ default=False,
+ action='store_true',
+ help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
+ parser.add_argument("-t",
+ "--tag",
+ type=str,
+ default=None,
+ help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
+ parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
+ parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+ args = parser.parse_args()
+
+ debug = args.debug
+
+ convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
+ args.output_dir,
+ max_shard_size=args.max_shard_size,
+ safe_serialization=args.safe_serialization,
+ tag=args.tag,
+ exclude_frozen_parameters=args.exclude_frozen_parameters)
diff --git a/qwen3/teacher_lora/checkpoint-628/README.md b/qwen3/teacher_lora/checkpoint-628/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..4fabf4f222ce5d86630bcfbcae8e90fb08d87580
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/README.md
@@ -0,0 +1,208 @@
+---
+base_model: Qwen/Qwen3-4B-Instruct-2507
+library_name: peft
+pipeline_tag: text-generation
+tags:
+- base_model:adapter:Qwen/Qwen3-4B-Instruct-2507
+- llama-factory
+- lora
+- transformers
+---
+
+# Model Card for Model ID
+
+
+
+
+
+## Model Details
+
+### Model Description
+
+
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+
+
+### Direct Use
+
+
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+
+
+[More Information Needed]
+
+### Recommendations
+
+
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+
+
+[More Information Needed]
+
+### Training Procedure
+
+
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed]
+
+#### Speeds, Sizes, Times [optional]
+
+
+
+[More Information Needed]
+
+## Evaluation
+
+
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+
+
+[More Information Needed]
+
+#### Factors
+
+
+
+[More Information Needed]
+
+#### Metrics
+
+
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+
+
+[More Information Needed]
+
+## Environmental Impact
+
+
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.18.1
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-628/adapter_config.json b/qwen3/teacher_lora/checkpoint-628/adapter_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..376dee7aec128c9bdb50c2cc1ddcf6a68a37ab9b
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/adapter_config.json
@@ -0,0 +1,46 @@
+{
+ "alora_invocation_tokens": null,
+ "alpha_pattern": {},
+ "arrow_config": null,
+ "auto_mapping": null,
+ "base_model_name_or_path": "Qwen/Qwen3-4B-Instruct-2507",
+ "bias": "none",
+ "corda_config": null,
+ "ensure_weight_tying": false,
+ "eva_config": null,
+ "exclude_modules": null,
+ "fan_in_fan_out": false,
+ "inference_mode": true,
+ "init_lora_weights": true,
+ "layer_replication": null,
+ "layers_pattern": null,
+ "layers_to_transform": null,
+ "loftq_config": {},
+ "lora_alpha": 64,
+ "lora_bias": false,
+ "lora_dropout": 0.1,
+ "megatron_config": null,
+ "megatron_core": "megatron.core",
+ "modules_to_save": null,
+ "peft_type": "LORA",
+ "peft_version": "0.18.1",
+ "qalora_group_size": 16,
+ "r": 32,
+ "rank_pattern": {},
+ "revision": null,
+ "target_modules": [
+ "down_proj",
+ "k_proj",
+ "up_proj",
+ "o_proj",
+ "q_proj",
+ "v_proj",
+ "gate_proj"
+ ],
+ "target_parameters": null,
+ "task_type": "CAUSAL_LM",
+ "trainable_token_indices": null,
+ "use_dora": false,
+ "use_qalora": false,
+ "use_rslora": false
+}
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-628/adapter_model.safetensors b/qwen3/teacher_lora/checkpoint-628/adapter_model.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..da0d61c309facdc1dd70f4e8604ef8a4f25d4c0e
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/adapter_model.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06034ccc5a4fd0a5508f3d32881d3048dd62ac84997177a641e709ceb40a20ce
+size 132188392
diff --git a/qwen3/teacher_lora/checkpoint-628/chat_template.jinja b/qwen3/teacher_lora/checkpoint-628/chat_template.jinja
new file mode 100644
index 0000000000000000000000000000000000000000..70adff8a08fb31e0636f618564838d4bf3c05286
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/chat_template.jinja
@@ -0,0 +1,61 @@
+{%- if tools %}
+ {{- '<|im_start|>system\n' }}
+ {%- if messages[0].role == 'system' %}
+ {{- messages[0].content + '\n\n' }}
+ {%- endif %}
+ {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }}
+ {%- for tool in tools %}
+ {{- "\n" }}
+ {{- tool | tojson }}
+ {%- endfor %}
+ {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }}
+{%- else %}
+ {%- if messages[0].role == 'system' %}
+ {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+ {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+ {%- if message.content is string %}
+ {%- set content = message.content %}
+ {%- else %}
+ {%- set content = '' %}
+ {%- endif %}
+ {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+ {%- elif message.role == "assistant" %}
+ {{- '<|im_start|>' + message.role + '\n' + content }}
+ {%- if message.tool_calls %}
+ {%- for tool_call in message.tool_calls %}
+ {%- if (loop.first and content) or (not loop.first) %}
+ {{- '\n' }}
+ {%- endif %}
+ {%- if tool_call.function %}
+ {%- set tool_call = tool_call.function %}
+ {%- endif %}
+ {{- '\n{"name": "' }}
+ {{- tool_call.name }}
+ {{- '", "arguments": ' }}
+ {%- if tool_call.arguments is string %}
+ {{- tool_call.arguments }}
+ {%- else %}
+ {{- tool_call.arguments | tojson }}
+ {%- endif %}
+ {{- '}\n' }}
+ {%- endfor %}
+ {%- endif %}
+ {{- '<|im_end|>\n' }}
+ {%- elif message.role == "tool" %}
+ {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+ {{- '<|im_start|>user' }}
+ {%- endif %}
+ {{- '\n\n' }}
+ {{- content }}
+ {{- '\n' }}
+ {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+ {{- '<|im_end|>\n' }}
+ {%- endif %}
+ {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+ {{- '<|im_start|>assistant\n' }}
+{%- endif %}
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-628/global_step628/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt b/qwen3/teacher_lora/checkpoint-628/global_step628/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt
new file mode 100644
index 0000000000000000000000000000000000000000..10a6f5de55244b6ff3f9cbfc7e969338e07bc3dd
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/global_step628/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47a9a979990e5c53d73e3e90ee5d02c45f3eab850f726761552c7053cf81596f
+size 396393221
diff --git a/qwen3/teacher_lora/checkpoint-628/global_step628/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt b/qwen3/teacher_lora/checkpoint-628/global_step628/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt
new file mode 100644
index 0000000000000000000000000000000000000000..76d531c48517b21591ce117d7f9de179b01c3e10
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/global_step628/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee92d7958ea8bec13bb26836486a69dbfe105aa6dc732804eb3e6d68854ded69
+size 396393349
diff --git a/qwen3/teacher_lora/checkpoint-628/global_step628/mp_rank_00_model_states.pt b/qwen3/teacher_lora/checkpoint-628/global_step628/mp_rank_00_model_states.pt
new file mode 100644
index 0000000000000000000000000000000000000000..a369c8c675693ab97350f1e8207a50fb954d286e
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/global_step628/mp_rank_00_model_states.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:45363c4fdb5c8ad54c0c71ea1315b2f16b42512eb6bde72dd585857852d3ad15
+size 910426661
diff --git a/qwen3/teacher_lora/checkpoint-628/latest b/qwen3/teacher_lora/checkpoint-628/latest
new file mode 100644
index 0000000000000000000000000000000000000000..af055c3db05dcb2af541005f22308dbd946aba65
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/latest
@@ -0,0 +1 @@
+global_step628
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-628/rng_state_0.pth b/qwen3/teacher_lora/checkpoint-628/rng_state_0.pth
new file mode 100644
index 0000000000000000000000000000000000000000..97591724fcb98be5d95333175c99b35d0b373d46
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/rng_state_0.pth
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b72b373cce1dd790cbb47b37388702c8d18eaaa5c908ae380202fbb69a257482
+size 14917
diff --git a/qwen3/teacher_lora/checkpoint-628/rng_state_1.pth b/qwen3/teacher_lora/checkpoint-628/rng_state_1.pth
new file mode 100644
index 0000000000000000000000000000000000000000..dddd4b4171ef73300e4ac2ae220ba877062f768f
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/rng_state_1.pth
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82d5b1b1dde04e3237fcfa724ab05bdc0555b795efdd6c89bacb8eac42e00a06
+size 14917
diff --git a/qwen3/teacher_lora/checkpoint-628/scheduler.pt b/qwen3/teacher_lora/checkpoint-628/scheduler.pt
new file mode 100644
index 0000000000000000000000000000000000000000..96c0cabcd99be4945ffbd7897ecc5367ac4beb28
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/scheduler.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb823a2c51d7fb42dfe3b693a7159784285dbf38c5a797e6e73292096123ab7a
+size 1465
diff --git a/qwen3/teacher_lora/checkpoint-628/tokenizer.json b/qwen3/teacher_lora/checkpoint-628/tokenizer.json
new file mode 100644
index 0000000000000000000000000000000000000000..c7afbed2efcdf019f88ab0572ec29d3bf595dfe2
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/tokenizer.json
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506
+size 11422650
diff --git a/qwen3/teacher_lora/checkpoint-628/tokenizer_config.json b/qwen3/teacher_lora/checkpoint-628/tokenizer_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..907e863a5a4109fb0e373ccdeb4d657bbeca9416
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/tokenizer_config.json
@@ -0,0 +1,31 @@
+{
+ "add_prefix_space": false,
+ "backend": "tokenizers",
+ "bos_token": null,
+ "clean_up_tokenization_spaces": false,
+ "eos_token": "<|im_end|>",
+ "errors": "replace",
+ "extra_special_tokens": [
+ "<|im_start|>",
+ "<|im_end|>",
+ "<|object_ref_start|>",
+ "<|object_ref_end|>",
+ "<|box_start|>",
+ "<|box_end|>",
+ "<|quad_start|>",
+ "<|quad_end|>",
+ "<|vision_start|>",
+ "<|vision_end|>",
+ "<|vision_pad|>",
+ "<|image_pad|>",
+ "<|video_pad|>"
+ ],
+ "is_local": false,
+ "local_files_only": false,
+ "model_max_length": 1010000,
+ "pad_token": "<|endoftext|>",
+ "padding_side": "right",
+ "split_special_tokens": false,
+ "tokenizer_class": "Qwen2Tokenizer",
+ "unk_token": null
+}
diff --git a/qwen3/teacher_lora/checkpoint-628/trainer_state.json b/qwen3/teacher_lora/checkpoint-628/trainer_state.json
new file mode 100644
index 0000000000000000000000000000000000000000..c0dfb22a78786f929bce317244a787f367a4769c
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/trainer_state.json
@@ -0,0 +1,281 @@
+{
+ "best_global_step": null,
+ "best_metric": null,
+ "best_model_checkpoint": null,
+ "epoch": 2.0,
+ "eval_steps": 500,
+ "global_step": 628,
+ "is_hyper_param_search": false,
+ "is_local_process_zero": true,
+ "is_world_process_zero": true,
+ "log_history": [
+ {
+ "epoch": 0.06383403151805306,
+ "grad_norm": 0.4070272743701935,
+ "learning_rate": 9.996390385474569e-05,
+ "loss": 0.37070536613464355,
+ "step": 20
+ },
+ {
+ "epoch": 0.12766806303610612,
+ "grad_norm": 0.5005771517753601,
+ "learning_rate": 9.984797509888031e-05,
+ "loss": 0.03893736004829407,
+ "step": 40
+ },
+ {
+ "epoch": 0.19150209455415917,
+ "grad_norm": 0.17237664759159088,
+ "learning_rate": 9.965229938050925e-05,
+ "loss": 0.02901872992515564,
+ "step": 60
+ },
+ {
+ "epoch": 0.25533612607221223,
+ "grad_norm": 0.17700566351413727,
+ "learning_rate": 9.937719005663432e-05,
+ "loss": 0.012938253581523895,
+ "step": 80
+ },
+ {
+ "epoch": 0.3191701575902653,
+ "grad_norm": 0.29617995023727417,
+ "learning_rate": 9.902308769000724e-05,
+ "loss": 0.01476612538099289,
+ "step": 100
+ },
+ {
+ "epoch": 0.38300418910831835,
+ "grad_norm": 0.09139569103717804,
+ "learning_rate": 9.859055934360789e-05,
+ "loss": 0.013924291729927063,
+ "step": 120
+ },
+ {
+ "epoch": 0.4468382206263714,
+ "grad_norm": 0.6957159042358398,
+ "learning_rate": 9.808029767254437e-05,
+ "loss": 0.015140479803085327,
+ "step": 140
+ },
+ {
+ "epoch": 0.5106722521444245,
+ "grad_norm": 0.19447490572929382,
+ "learning_rate": 9.749311981482817e-05,
+ "loss": 0.016634005308151244,
+ "step": 160
+ },
+ {
+ "epoch": 0.5745062836624776,
+ "grad_norm": 0.13639570772647858,
+ "learning_rate": 9.682996608280154e-05,
+ "loss": 0.008574086427688598,
+ "step": 180
+ },
+ {
+ "epoch": 0.6383403151805306,
+ "grad_norm": 0.05407094210386276,
+ "learning_rate": 9.609189845731207e-05,
+ "loss": 0.00992567539215088,
+ "step": 200
+ },
+ {
+ "epoch": 0.7021743466985837,
+ "grad_norm": 0.08482213318347931,
+ "learning_rate": 9.52800988870462e-05,
+ "loss": 0.009285379946231843,
+ "step": 220
+ },
+ {
+ "epoch": 0.7660083782166367,
+ "grad_norm": 0.07621290534734726,
+ "learning_rate": 9.439586739574513e-05,
+ "loss": 0.004985855519771576,
+ "step": 240
+ },
+ {
+ "epoch": 0.8298424097346898,
+ "grad_norm": 0.3552233576774597,
+ "learning_rate": 9.344062000033414e-05,
+ "loss": 0.014777910709381104,
+ "step": 260
+ },
+ {
+ "epoch": 0.8936764412527428,
+ "grad_norm": 0.07681854814291,
+ "learning_rate": 9.241588644329919e-05,
+ "loss": 0.0040129125118255615,
+ "step": 280
+ },
+ {
+ "epoch": 0.9575104727707959,
+ "grad_norm": 0.08489467203617096,
+ "learning_rate": 9.132330774294252e-05,
+ "loss": 0.012399718165397644,
+ "step": 300
+ },
+ {
+ "epoch": 1.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 94.6262,
+ "eval_generator_rouge1": 99.8758,
+ "eval_generator_rouge2": 99.7812,
+ "eval_generator_rougeL": 99.5997,
+ "eval_loss": 0.004354734905064106,
+ "eval_runtime": 1522.2841,
+ "eval_samples_per_second": 2.243,
+ "eval_selector_accuracy": 98.014,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.141,
+ "step": 314
+ },
+ {
+ "epoch": 1.019150209455416,
+ "grad_norm": 0.23293638229370117,
+ "learning_rate": 9.01646335654399e-05,
+ "loss": 0.011312296986579895,
+ "step": 320
+ },
+ {
+ "epoch": 1.082984240973469,
+ "grad_norm": 0.07232551276683807,
+ "learning_rate": 8.894171942290817e-05,
+ "loss": 0.006549905985593796,
+ "step": 340
+ },
+ {
+ "epoch": 1.1468182724915221,
+ "grad_norm": 0.05371030047535896,
+ "learning_rate": 8.765652370197003e-05,
+ "loss": 0.003336009383201599,
+ "step": 360
+ },
+ {
+ "epoch": 1.2106523040095751,
+ "grad_norm": 0.1448812186717987,
+ "learning_rate": 8.631110452757467e-05,
+ "loss": 0.008410868048667908,
+ "step": 380
+ },
+ {
+ "epoch": 1.2744863355276281,
+ "grad_norm": 0.0636834055185318,
+ "learning_rate": 8.490761646709642e-05,
+ "loss": 0.006242972612380981,
+ "step": 400
+ },
+ {
+ "epoch": 1.3383203670456814,
+ "grad_norm": 0.05049936845898628,
+ "learning_rate": 8.344830707998953e-05,
+ "loss": 0.00744413286447525,
+ "step": 420
+ },
+ {
+ "epoch": 1.4021543985637344,
+ "grad_norm": 0.032405462116003036,
+ "learning_rate": 8.193551331852458e-05,
+ "loss": 0.0031117375940084456,
+ "step": 440
+ },
+ {
+ "epoch": 1.4659884300817874,
+ "grad_norm": 0.13639135658740997,
+ "learning_rate": 8.037165778537033e-05,
+ "loss": 0.005259513109922409,
+ "step": 460
+ },
+ {
+ "epoch": 1.5298224615998404,
+ "grad_norm": 0.021995026618242264,
+ "learning_rate": 7.875924485401427e-05,
+ "loss": 0.004164557158946991,
+ "step": 480
+ },
+ {
+ "epoch": 1.5936564931178934,
+ "grad_norm": 0.01941438391804695,
+ "learning_rate": 7.710085665823444e-05,
+ "loss": 0.0015524964779615401,
+ "step": 500
+ },
+ {
+ "epoch": 1.6574905246359466,
+ "grad_norm": 0.1271296739578247,
+ "learning_rate": 7.539914895704521e-05,
+ "loss": 0.004233330115675926,
+ "step": 520
+ },
+ {
+ "epoch": 1.7213245561539996,
+ "grad_norm": 0.22600136697292328,
+ "learning_rate": 7.365684688173908e-05,
+ "loss": 0.0029299482703208923,
+ "step": 540
+ },
+ {
+ "epoch": 1.7851585876720528,
+ "grad_norm": 0.0755968764424324,
+ "learning_rate": 7.18767405718347e-05,
+ "loss": 0.0044687166810035706,
+ "step": 560
+ },
+ {
+ "epoch": 1.8489926191901058,
+ "grad_norm": 0.08520206063985825,
+ "learning_rate": 7.006168070692036e-05,
+ "loss": 0.005520488321781159,
+ "step": 580
+ },
+ {
+ "epoch": 1.9128266507081588,
+ "grad_norm": 0.008491385728120804,
+ "learning_rate": 6.821457394154786e-05,
+ "loss": 0.0015786593779921533,
+ "step": 600
+ },
+ {
+ "epoch": 1.9766606822262118,
+ "grad_norm": 0.009171651676297188,
+ "learning_rate": 6.633837825048768e-05,
+ "loss": 0.007294930517673492,
+ "step": 620
+ },
+ {
+ "epoch": 2.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 96.2617,
+ "eval_generator_rouge1": 99.9321,
+ "eval_generator_rouge2": 99.8752,
+ "eval_generator_rougeL": 99.7349,
+ "eval_loss": 0.0033019811380654573,
+ "eval_runtime": 1482.8547,
+ "eval_samples_per_second": 2.302,
+ "eval_selector_accuracy": 98.4813,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.144,
+ "step": 628
+ }
+ ],
+ "logging_steps": 20,
+ "max_steps": 1570,
+ "num_input_tokens_seen": 0,
+ "num_train_epochs": 5,
+ "save_steps": 500,
+ "stateful_callbacks": {
+ "TrainerControl": {
+ "args": {
+ "should_epoch_stop": false,
+ "should_evaluate": false,
+ "should_log": false,
+ "should_save": true,
+ "should_training_stop": false
+ },
+ "attributes": {}
+ }
+ },
+ "total_flos": 1.592447410729124e+17,
+ "train_batch_size": 1,
+ "trial_name": null,
+ "trial_params": null
+}
diff --git a/qwen3/teacher_lora/checkpoint-628/training_args.bin b/qwen3/teacher_lora/checkpoint-628/training_args.bin
new file mode 100644
index 0000000000000000000000000000000000000000..1ef98d1a171b439f2f8968f31ccbc846a421fe48
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/training_args.bin
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:84fda57546f783ec338bd3a292c7dba3d917a8a52ac59ca7cebdd546da52b586
+size 7185
diff --git a/qwen3/teacher_lora/checkpoint-628/zero_to_fp32.py b/qwen3/teacher_lora/checkpoint-628/zero_to_fp32.py
new file mode 100644
index 0000000000000000000000000000000000000000..f40f1f0acbbff30e8e8d29906b56aae4e8b8a138
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-628/zero_to_fp32.py
@@ -0,0 +1,790 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example:
+# python zero_to_fp32.py . output_dir/
+# or
+# python zero_to_fp32.py . output_dir/ --safe_serialization
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+import gc
+import json
+import numpy as np
+from tqdm import tqdm
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+ FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+ FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS, AUTOEP_LAYERS_KEY,
+ AUTOEP_LAYERS_KEY_LEGACY, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY,
+ AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT)
+
+
+@dataclass
+class zero_model_state:
+ buffers: dict()
+ param_shapes: dict()
+ shared_params: list
+ ds_version: int
+ frozen_param_shapes: dict()
+ frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+ return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+ '''
+ alist.sort(key=natural_keys) sorts in human order
+ http://nedbatchelder.com/blog/200712/human_sorting.html
+ (See Toothy's implementation in the comments)
+ '''
+ return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+ if not os.path.isdir(checkpoint_dir):
+ raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+ # there should be only one file
+ if zero_stage <= 2:
+ file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+ elif zero_stage == 3:
+ file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+ if not os.path.exists(file):
+ raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+ return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+ # XXX: need to test that this simple glob rule works for multi-node setup too
+ ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+ if len(ckpt_files) == 0:
+ raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+ return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+ return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+ return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def _has_autoep_zero3_partitioned_metadata(state_dict):
+ autoep_layers = state_dict.get(AUTOEP_LAYERS_KEY)
+ if autoep_layers is None:
+ autoep_layers = state_dict.get(AUTOEP_LAYERS_KEY_LEGACY)
+ if not isinstance(autoep_layers, list):
+ return False
+ return any(
+ isinstance(entry, dict)
+ and entry.get(AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY) == AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT
+ for entry in autoep_layers)
+
+
+def _raise_if_autoep_zero3_partitioned_state(state_dict):
+ if _has_autoep_zero3_partitioned_metadata(state_dict):
+ raise NotImplementedError("zero_to_fp32 does not support AutoEP ZeRO-3 partition-native checkpoints. "
+ "AutoEP expert parameters are partitioned over expert replica groups, so "
+ "global data-parallel consolidation would produce incomplete expert tensors. "
+ "Use ds_to_universal.py for expert-aware conversion.")
+
+
+def _raise_if_autoep_zero3_partitioned_checkpoint(model_files):
+ for file in model_files:
+ state_dict = torch.load(file, map_location=device, weights_only=False)
+ _raise_if_autoep_zero3_partitioned_state(state_dict)
+
+
+def parse_model_states(files):
+ zero_model_states = []
+ for file in files:
+ state_dict = torch.load(file, map_location=device, weights_only=False)
+ _raise_if_autoep_zero3_partitioned_state(state_dict)
+
+ if BUFFER_NAMES not in state_dict:
+ raise ValueError(f"{file} is not a model state checkpoint")
+ buffer_names = state_dict[BUFFER_NAMES]
+ if debug:
+ print("Found buffers:", buffer_names)
+
+ # recover just the buffers while restoring them to fp32 if they were saved in fp16
+ buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+ param_shapes = state_dict[PARAM_SHAPES]
+
+ # collect parameters that are included in param_shapes
+ param_names = []
+ for s in param_shapes:
+ for name in s.keys():
+ param_names.append(name)
+
+ # update with frozen parameters
+ frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+ if frozen_param_shapes is not None:
+ if debug:
+ print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+ param_names += list(frozen_param_shapes.keys())
+
+ # handle shared params
+ shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+ ds_version = state_dict.get(DS_VERSION, None)
+
+ frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+ z_model_state = zero_model_state(buffers=buffers,
+ param_shapes=param_shapes,
+ shared_params=shared_params,
+ ds_version=ds_version,
+ frozen_param_shapes=frozen_param_shapes,
+ frozen_param_fragments=frozen_param_fragments)
+ zero_model_states.append(z_model_state)
+
+ return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+ total_files = len(files)
+ state_dicts = []
+ for f in tqdm(files, desc='Loading checkpoint shards'):
+ state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
+ # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
+ # and also handle the case where it was already removed by another helper script
+ state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
+ state_dicts.append(state_dict)
+
+ if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
+ raise ValueError(f"{files[0]} is not a zero checkpoint")
+ zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+ world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+ # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+ # parameters can be different from data parallelism for non-expert parameters. So we can just
+ # use the max of the partition_count to get the dp world_size.
+
+ if type(world_size) is list:
+ world_size = max(world_size)
+
+ if world_size != total_files:
+ raise ValueError(
+ f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+ "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+ )
+
+ # the groups are named differently in each stage
+ if zero_stage <= 2:
+ fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+ elif zero_stage == 3:
+ fp32_groups_key = FP32_FLAT_GROUPS
+ else:
+ raise ValueError(f"unknown zero stage {zero_stage}")
+
+ fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+ return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
+ """
+ Returns fp32 state_dict reconstructed from ds checkpoint
+
+ Args:
+ - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+ """
+ print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+ # parse_model_states rejects AutoEP ZeRO-3 partition-native checkpoints
+ # before the expensive optimizer-shard load below.
+ model_files = get_model_state_files(ds_checkpoint_dir)
+ zero_model_states = parse_model_states(model_files)
+ print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+ optim_files = get_optim_files(ds_checkpoint_dir)
+ zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+ print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+ if zero_stage <= 2:
+ return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters)
+ elif zero_stage == 3:
+ return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+ return
+
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+ frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+ if debug:
+ num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+ print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+ wanted_params = len(frozen_param_shapes)
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+ avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+ print(f'Frozen params: Have {avail_numel} numels to process.')
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+ total_params = 0
+ total_numel = 0
+ for name, shape in frozen_param_shapes.items():
+ total_params += 1
+ unpartitioned_numel = shape.numel()
+ total_numel += unpartitioned_numel
+
+ state_dict[name] = frozen_param_fragments[name]
+
+ if debug:
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _has_callable(obj, fn):
+ attr = getattr(obj, fn, None)
+ return callable(attr)
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+ param_shapes = zero_model_states[0].param_shapes
+
+ # Reconstruction protocol:
+ #
+ # XXX: document this
+
+ if debug:
+ for i in range(world_size):
+ for j in range(len(fp32_flat_groups[0])):
+ print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+ # XXX: memory usage doubles here (zero2)
+ num_param_groups = len(fp32_flat_groups[0])
+ merged_single_partition_of_fp32_groups = []
+ for i in range(num_param_groups):
+ merged_partitions = [sd[i] for sd in fp32_flat_groups]
+ full_single_fp32_vector = torch.cat(merged_partitions, 0)
+ merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+ avail_numel = sum(
+ [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+ if debug:
+ wanted_params = sum([len(shapes) for shapes in param_shapes])
+ wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+ # not asserting if there is a mismatch due to possible padding
+ print(f"Have {avail_numel} numels to process.")
+ print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+ # params
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+ # out-of-core computing solution
+ total_numel = 0
+ total_params = 0
+ for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+ offset = 0
+ avail_numel = full_single_fp32_vector.numel()
+ for name, shape in shapes.items():
+
+ unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
+ total_numel += unpartitioned_numel
+ total_params += 1
+
+ if debug:
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+ state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+ offset += unpartitioned_numel
+
+ # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+ # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+ # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+ # live optimizer object, so we are checking that the numbers are within the right range
+ align_to = 2 * world_size
+
+ def zero2_align(x):
+ return align_to * math.ceil(x / align_to)
+
+ if debug:
+ print(f"original offset={offset}, avail_numel={avail_numel}")
+
+ offset = zero2_align(offset)
+ avail_numel = zero2_align(avail_numel)
+
+ if debug:
+ print(f"aligned offset={offset}, avail_numel={avail_numel}")
+
+ # Sanity check
+ if offset != avail_numel:
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+ print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters):
+ state_dict = OrderedDict()
+
+ # buffers
+ buffers = zero_model_states[0].buffers
+ state_dict.update(buffers)
+ if debug:
+ print(f"added {len(buffers)} buffers")
+
+ if not exclude_frozen_parameters:
+ _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+ _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+ # recover shared parameters
+ for pair in zero_model_states[0].shared_params:
+ if pair[1] in state_dict:
+ state_dict[pair[0]] = state_dict[pair[1]]
+
+ return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+ remainder = unpartitioned_numel % world_size
+ padding_numel = (world_size - remainder) if remainder else 0
+ partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+ return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+ return
+
+ if debug:
+ for i in range(world_size):
+ num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+ print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+ wanted_params = len(frozen_param_shapes)
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+ avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+ print(f'Frozen params: Have {avail_numel} numels to process.')
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+ total_params = 0
+ total_numel = 0
+ for name, shape in zero_model_states[0].frozen_param_shapes.items():
+ total_params += 1
+ unpartitioned_numel = shape.numel()
+ total_numel += unpartitioned_numel
+
+ param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+ state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+ if debug:
+ print(
+ f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+ )
+
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+class GatheredTensor:
+ """
+ A pseudo tensor that collects partitioned weights.
+ It is more memory efficient when there are multiple groups.
+ """
+
+ def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
+ self.flat_groups = flat_groups
+ self.flat_groups_offset = flat_groups_offset
+ self.offset = offset
+ self.partitioned_numel = partitioned_numel
+ self.shape = shape
+ self.dtype = self.flat_groups[0][0].dtype
+
+ def contiguous(self):
+ """
+ Merge partitioned weights from flat_groups into a single tensor.
+ """
+ end_idx = self.offset + self.partitioned_numel
+ world_size = len(self.flat_groups)
+ pad_flat_param_chunks = []
+
+ for rank_i in range(world_size):
+ # for each rank, we need to collect weights from related group/groups
+ flat_groups_at_rank_i = self.flat_groups[rank_i]
+ start_group_id = None
+ end_group_id = None
+ for group_id in range(len(self.flat_groups_offset)):
+ if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
+ start_group_id = group_id
+ if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
+ end_group_id = group_id
+ break
+ # collect weights from related group/groups
+ for group_id in range(start_group_id, end_group_id + 1):
+ flat_tensor = flat_groups_at_rank_i[group_id]
+ start_offset = self.offset - self.flat_groups_offset[group_id]
+ end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
+ pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
+
+ # collect weights from all ranks
+ pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
+ param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
+ return param
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+ param_shapes = zero_model_states[0].param_shapes
+ avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
+
+ # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+ # param, re-consolidating each param, while dealing with padding if any
+
+ # merge list of dicts, preserving order
+ param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+ if debug:
+ for i in range(world_size):
+ print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+ wanted_params = len(param_shapes)
+ wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+ # not asserting if there is a mismatch due to possible padding
+ avail_numel = fp32_flat_groups[0].numel() * world_size
+ print(f"Trainable params: Have {avail_numel} numels to process.")
+ print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+ # params
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+ # out-of-core computing solution
+ offset = 0
+ total_numel = 0
+ total_params = 0
+ flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
+ for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
+ unpartitioned_numel = shape.numel()
+ total_numel += unpartitioned_numel
+ total_params += 1
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+ if debug:
+ print(
+ f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+ )
+
+ # memory efficient tensor
+ tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
+ state_dict[name] = tensor
+ offset += partitioned_numel
+
+ offset *= world_size
+
+ # Sanity check
+ if offset != avail_numel:
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+ print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters):
+ state_dict = OrderedDict()
+
+ # buffers
+ buffers = zero_model_states[0].buffers
+ state_dict.update(buffers)
+ if debug:
+ print(f"added {len(buffers)} buffers")
+
+ if not exclude_frozen_parameters:
+ _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+ _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+ # recover shared parameters
+ for pair in zero_model_states[0].shared_params:
+ if pair[1] in state_dict:
+ state_dict[pair[0]] = state_dict[pair[1]]
+
+ return state_dict
+
+
+def to_torch_tensor(state_dict, return_empty_tensor=False):
+ """
+ Convert state_dict of GatheredTensor to torch tensor
+ """
+ torch_state_dict = {}
+ converted_tensors = {}
+ for name, tensor in state_dict.items():
+ tensor_id = id(tensor)
+ if tensor_id in converted_tensors: # shared tensors
+ shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
+ torch_state_dict[name] = shared_tensor
+ else:
+ converted_tensors[tensor_id] = name
+ if return_empty_tensor:
+ torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
+ else:
+ torch_state_dict[name] = tensor.contiguous()
+ return torch_state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+ tag=None,
+ exclude_frozen_parameters=False,
+ lazy_mode=False):
+ """
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+ ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+ via a model hub.
+
+ Args:
+ - ``checkpoint_dir``: path to the desired checkpoint folder
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+ - ``exclude_frozen_parameters``: exclude frozen parameters
+ - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
+ Convert the pesduo tensor to torch tensor by ``.contiguous()``
+
+ Returns:
+ - pytorch ``state_dict``
+
+ A typical usage might be ::
+
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+ # do the training and checkpoint saving
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+ model = model.cpu() # move to cpu
+ model.load_state_dict(state_dict)
+ # submit to model hub or save the model to share with others
+
+ In this example the ``model`` will no longer be usable in the deepspeed context of the same
+ application. i.e. you will need to re-initialize the deepspeed engine, since
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+ If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+ Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
+ You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+ the checkpoint. Or you can load state_dict in lazy mode ::
+
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
+ for name, lazy_tensor in state_dict.item():
+ tensor = lazy_tensor.contiguous() # to cpu
+ print(name, tensor)
+ # del tensor to release memory if it no longer in use
+ """
+ if tag is None:
+ latest_path = os.path.join(checkpoint_dir, 'latest')
+ if os.path.isfile(latest_path):
+ with open(latest_path, 'r') as fd:
+ tag = fd.read().strip()
+ else:
+ raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+ ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+ if not os.path.isdir(ds_checkpoint_dir):
+ raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+ state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
+ if lazy_mode:
+ return state_dict
+ else:
+ return to_torch_tensor(state_dict)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
+ output_dir,
+ max_shard_size="5GB",
+ safe_serialization=False,
+ tag=None,
+ exclude_frozen_parameters=False):
+ """
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+ loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+ Args:
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+ - ``output_dir``: directory to the pytorch fp32 state_dict output files
+ - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
+ - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+ - ``exclude_frozen_parameters``: exclude frozen parameters
+ """
+
+ # Dependency pre-check
+ if safe_serialization:
+ try:
+ from safetensors.torch import save_file
+ except ImportError:
+ print('If you want to use `safe_serialization`, please `pip install safetensors`')
+ raise
+ if max_shard_size is not None:
+ try:
+ from huggingface_hub import split_torch_state_dict_into_shards
+ except ImportError:
+ print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
+ raise
+
+ # Convert zero checkpoint to state_dict
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+ tag,
+ exclude_frozen_parameters,
+ lazy_mode=True)
+
+ # Shard the model if it is too big.
+ weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
+ if max_shard_size is not None:
+ filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
+ # an memory-efficient approach for sharding
+ empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
+ state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
+ filename_pattern=filename_pattern,
+ max_shard_size=max_shard_size)
+ else:
+ from collections import namedtuple
+ StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
+ state_dict_split = StateDictSplit(is_sharded=False,
+ filename_to_tensors={weights_name: list(state_dict.keys())})
+
+ # Save the model by shard
+ os.makedirs(output_dir, exist_ok=True)
+ filename_to_tensors = state_dict_split.filename_to_tensors.items()
+ for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
+ shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
+ shard_state_dict = to_torch_tensor(shard_state_dict)
+ output_path = os.path.join(output_dir, shard_file)
+ if safe_serialization:
+ save_file(shard_state_dict, output_path, metadata={"format": "pt"})
+ else:
+ torch.save(shard_state_dict, output_path)
+ # release the memory of current shard
+ for tensor_name in list(shard_state_dict.keys()):
+ del state_dict[tensor_name]
+ del shard_state_dict[tensor_name]
+ del shard_state_dict
+ gc.collect()
+
+ # Save index if sharded
+ if state_dict_split.is_sharded:
+ index = {
+ "metadata": state_dict_split.metadata,
+ "weight_map": state_dict_split.tensor_to_filename,
+ }
+ save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
+ save_index_file = os.path.join(output_dir, save_index_file)
+ with open(save_index_file, "w", encoding="utf-8") as f:
+ content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+ f.write(content)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+ """
+ 1. Put the provided model to cpu
+ 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+ 3. Load it into the provided model
+
+ Args:
+ - ``model``: the model object to update
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+ Returns:
+ - ``model`: modified model
+
+ Make sure you have plenty of CPU memory available before you call this function. If you don't
+ have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+ conveniently placed for you in the checkpoint folder.
+
+ A typical usage might be ::
+
+ from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+ model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+ # submit to model hub or save the model to share with others
+
+ Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+ of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+ """
+ logger.info("Extracting fp32 weights")
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+ logger.info("Overwriting model with fp32 weights")
+ model = model.cpu()
+ model.load_state_dict(state_dict, strict=False)
+
+ return model
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("checkpoint_dir",
+ type=str,
+ help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+ parser.add_argument("output_dir",
+ type=str,
+ help="directory to the pytorch fp32 state_dict output files"
+ "(e.g. path/checkpoint-12-output/)")
+ parser.add_argument(
+ "--max_shard_size",
+ type=str,
+ default="5GB",
+ help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
+ "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
+ "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
+ "without CPU OOM issues.")
+ parser.add_argument(
+ "--safe_serialization",
+ default=False,
+ action='store_true',
+ help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
+ parser.add_argument("-t",
+ "--tag",
+ type=str,
+ default=None,
+ help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
+ parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
+ parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+ args = parser.parse_args()
+
+ debug = args.debug
+
+ convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
+ args.output_dir,
+ max_shard_size=args.max_shard_size,
+ safe_serialization=args.safe_serialization,
+ tag=args.tag,
+ exclude_frozen_parameters=args.exclude_frozen_parameters)
diff --git a/qwen3/teacher_lora/checkpoint-942/README.md b/qwen3/teacher_lora/checkpoint-942/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..4fabf4f222ce5d86630bcfbcae8e90fb08d87580
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/README.md
@@ -0,0 +1,208 @@
+---
+base_model: Qwen/Qwen3-4B-Instruct-2507
+library_name: peft
+pipeline_tag: text-generation
+tags:
+- base_model:adapter:Qwen/Qwen3-4B-Instruct-2507
+- llama-factory
+- lora
+- transformers
+---
+
+# Model Card for Model ID
+
+
+
+
+
+## Model Details
+
+### Model Description
+
+
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+
+
+### Direct Use
+
+
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+
+
+[More Information Needed]
+
+### Recommendations
+
+
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+
+
+[More Information Needed]
+
+### Training Procedure
+
+
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed]
+
+#### Speeds, Sizes, Times [optional]
+
+
+
+[More Information Needed]
+
+## Evaluation
+
+
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+
+
+[More Information Needed]
+
+#### Factors
+
+
+
+[More Information Needed]
+
+#### Metrics
+
+
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+
+
+[More Information Needed]
+
+## Environmental Impact
+
+
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.18.1
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-942/adapter_config.json b/qwen3/teacher_lora/checkpoint-942/adapter_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..376dee7aec128c9bdb50c2cc1ddcf6a68a37ab9b
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/adapter_config.json
@@ -0,0 +1,46 @@
+{
+ "alora_invocation_tokens": null,
+ "alpha_pattern": {},
+ "arrow_config": null,
+ "auto_mapping": null,
+ "base_model_name_or_path": "Qwen/Qwen3-4B-Instruct-2507",
+ "bias": "none",
+ "corda_config": null,
+ "ensure_weight_tying": false,
+ "eva_config": null,
+ "exclude_modules": null,
+ "fan_in_fan_out": false,
+ "inference_mode": true,
+ "init_lora_weights": true,
+ "layer_replication": null,
+ "layers_pattern": null,
+ "layers_to_transform": null,
+ "loftq_config": {},
+ "lora_alpha": 64,
+ "lora_bias": false,
+ "lora_dropout": 0.1,
+ "megatron_config": null,
+ "megatron_core": "megatron.core",
+ "modules_to_save": null,
+ "peft_type": "LORA",
+ "peft_version": "0.18.1",
+ "qalora_group_size": 16,
+ "r": 32,
+ "rank_pattern": {},
+ "revision": null,
+ "target_modules": [
+ "down_proj",
+ "k_proj",
+ "up_proj",
+ "o_proj",
+ "q_proj",
+ "v_proj",
+ "gate_proj"
+ ],
+ "target_parameters": null,
+ "task_type": "CAUSAL_LM",
+ "trainable_token_indices": null,
+ "use_dora": false,
+ "use_qalora": false,
+ "use_rslora": false
+}
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-942/adapter_model.safetensors b/qwen3/teacher_lora/checkpoint-942/adapter_model.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..9cbc789ab45547fa81213b1f3d6c937b4da0819f
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/adapter_model.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:311656d93a2665768b9ea7b0bda10b4bf921ab80376ff736ee4a10ce43a17ebd
+size 132188392
diff --git a/qwen3/teacher_lora/checkpoint-942/chat_template.jinja b/qwen3/teacher_lora/checkpoint-942/chat_template.jinja
new file mode 100644
index 0000000000000000000000000000000000000000..70adff8a08fb31e0636f618564838d4bf3c05286
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/chat_template.jinja
@@ -0,0 +1,61 @@
+{%- if tools %}
+ {{- '<|im_start|>system\n' }}
+ {%- if messages[0].role == 'system' %}
+ {{- messages[0].content + '\n\n' }}
+ {%- endif %}
+ {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }}
+ {%- for tool in tools %}
+ {{- "\n" }}
+ {{- tool | tojson }}
+ {%- endfor %}
+ {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }}
+{%- else %}
+ {%- if messages[0].role == 'system' %}
+ {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+ {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+ {%- if message.content is string %}
+ {%- set content = message.content %}
+ {%- else %}
+ {%- set content = '' %}
+ {%- endif %}
+ {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+ {%- elif message.role == "assistant" %}
+ {{- '<|im_start|>' + message.role + '\n' + content }}
+ {%- if message.tool_calls %}
+ {%- for tool_call in message.tool_calls %}
+ {%- if (loop.first and content) or (not loop.first) %}
+ {{- '\n' }}
+ {%- endif %}
+ {%- if tool_call.function %}
+ {%- set tool_call = tool_call.function %}
+ {%- endif %}
+ {{- '\n{"name": "' }}
+ {{- tool_call.name }}
+ {{- '", "arguments": ' }}
+ {%- if tool_call.arguments is string %}
+ {{- tool_call.arguments }}
+ {%- else %}
+ {{- tool_call.arguments | tojson }}
+ {%- endif %}
+ {{- '}\n' }}
+ {%- endfor %}
+ {%- endif %}
+ {{- '<|im_end|>\n' }}
+ {%- elif message.role == "tool" %}
+ {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+ {{- '<|im_start|>user' }}
+ {%- endif %}
+ {{- '\n\n' }}
+ {{- content }}
+ {{- '\n' }}
+ {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+ {{- '<|im_end|>\n' }}
+ {%- endif %}
+ {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+ {{- '<|im_start|>assistant\n' }}
+{%- endif %}
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-942/global_step942/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt b/qwen3/teacher_lora/checkpoint-942/global_step942/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt
new file mode 100644
index 0000000000000000000000000000000000000000..803e708c1e5cca24574c2b1e0f88b8d23a8bdc74
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/global_step942/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:93cb7590542114ed6a4855825b4472a09e51f46eb645402db03264d0edb0feb7
+size 396393221
diff --git a/qwen3/teacher_lora/checkpoint-942/global_step942/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt b/qwen3/teacher_lora/checkpoint-942/global_step942/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt
new file mode 100644
index 0000000000000000000000000000000000000000..0c9e1118c6728fcb0823dcbd99efc937f58b9048
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/global_step942/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:44adfc6300f02eb7e3a926b0f0ad166ecb7f2aa5735a722cf603a8e48f79f71f
+size 396393349
diff --git a/qwen3/teacher_lora/checkpoint-942/global_step942/mp_rank_00_model_states.pt b/qwen3/teacher_lora/checkpoint-942/global_step942/mp_rank_00_model_states.pt
new file mode 100644
index 0000000000000000000000000000000000000000..4e6e8f5811b331ab3675af0ded13d1d09bb9964b
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/global_step942/mp_rank_00_model_states.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b5d2aa3f924560db8e088bb1fe1c2b18ce31cc3d0e84e7d2285bd405b69b1511
+size 910426661
diff --git a/qwen3/teacher_lora/checkpoint-942/latest b/qwen3/teacher_lora/checkpoint-942/latest
new file mode 100644
index 0000000000000000000000000000000000000000..4e9a68e4460aca593c480611a486d2a4cd73ef42
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/latest
@@ -0,0 +1 @@
+global_step942
\ No newline at end of file
diff --git a/qwen3/teacher_lora/checkpoint-942/rng_state_0.pth b/qwen3/teacher_lora/checkpoint-942/rng_state_0.pth
new file mode 100644
index 0000000000000000000000000000000000000000..b45b66f2dee40dd64bc5f07a41848dbadfe48083
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/rng_state_0.pth
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8d0b32e27e4120b86020f95399815db4bf7a06f8a084c53510982d5c7eb065d5
+size 14917
diff --git a/qwen3/teacher_lora/checkpoint-942/rng_state_1.pth b/qwen3/teacher_lora/checkpoint-942/rng_state_1.pth
new file mode 100644
index 0000000000000000000000000000000000000000..f1c2d688702c11f4def880d2f43fe8a05bb46591
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/rng_state_1.pth
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19355b2e0f80851843bbb83db57325b1c2057b4ac24ca8a2e62f0bd99ae6501b
+size 14917
diff --git a/qwen3/teacher_lora/checkpoint-942/scheduler.pt b/qwen3/teacher_lora/checkpoint-942/scheduler.pt
new file mode 100644
index 0000000000000000000000000000000000000000..6b1f8af0480bf83ab728abfeff11bc3288214349
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/scheduler.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8076da6f897951c5bdf8649e4fba596e8c49819927746017442ffd85f40b18a3
+size 1465
diff --git a/qwen3/teacher_lora/checkpoint-942/tokenizer.json b/qwen3/teacher_lora/checkpoint-942/tokenizer.json
new file mode 100644
index 0000000000000000000000000000000000000000..c7afbed2efcdf019f88ab0572ec29d3bf595dfe2
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/tokenizer.json
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506
+size 11422650
diff --git a/qwen3/teacher_lora/checkpoint-942/tokenizer_config.json b/qwen3/teacher_lora/checkpoint-942/tokenizer_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..907e863a5a4109fb0e373ccdeb4d657bbeca9416
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/tokenizer_config.json
@@ -0,0 +1,31 @@
+{
+ "add_prefix_space": false,
+ "backend": "tokenizers",
+ "bos_token": null,
+ "clean_up_tokenization_spaces": false,
+ "eos_token": "<|im_end|>",
+ "errors": "replace",
+ "extra_special_tokens": [
+ "<|im_start|>",
+ "<|im_end|>",
+ "<|object_ref_start|>",
+ "<|object_ref_end|>",
+ "<|box_start|>",
+ "<|box_end|>",
+ "<|quad_start|>",
+ "<|quad_end|>",
+ "<|vision_start|>",
+ "<|vision_end|>",
+ "<|vision_pad|>",
+ "<|image_pad|>",
+ "<|video_pad|>"
+ ],
+ "is_local": false,
+ "local_files_only": false,
+ "model_max_length": 1010000,
+ "pad_token": "<|endoftext|>",
+ "padding_side": "right",
+ "split_special_tokens": false,
+ "tokenizer_class": "Qwen2Tokenizer",
+ "unk_token": null
+}
diff --git a/qwen3/teacher_lora/checkpoint-942/trainer_state.json b/qwen3/teacher_lora/checkpoint-942/trainer_state.json
new file mode 100644
index 0000000000000000000000000000000000000000..f738f068f2bb9f0a09a64d4451e656853ab24014
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/trainer_state.json
@@ -0,0 +1,408 @@
+{
+ "best_global_step": null,
+ "best_metric": null,
+ "best_model_checkpoint": null,
+ "epoch": 3.0,
+ "eval_steps": 500,
+ "global_step": 942,
+ "is_hyper_param_search": false,
+ "is_local_process_zero": true,
+ "is_world_process_zero": true,
+ "log_history": [
+ {
+ "epoch": 0.06383403151805306,
+ "grad_norm": 0.4070272743701935,
+ "learning_rate": 9.996390385474569e-05,
+ "loss": 0.37070536613464355,
+ "step": 20
+ },
+ {
+ "epoch": 0.12766806303610612,
+ "grad_norm": 0.5005771517753601,
+ "learning_rate": 9.984797509888031e-05,
+ "loss": 0.03893736004829407,
+ "step": 40
+ },
+ {
+ "epoch": 0.19150209455415917,
+ "grad_norm": 0.17237664759159088,
+ "learning_rate": 9.965229938050925e-05,
+ "loss": 0.02901872992515564,
+ "step": 60
+ },
+ {
+ "epoch": 0.25533612607221223,
+ "grad_norm": 0.17700566351413727,
+ "learning_rate": 9.937719005663432e-05,
+ "loss": 0.012938253581523895,
+ "step": 80
+ },
+ {
+ "epoch": 0.3191701575902653,
+ "grad_norm": 0.29617995023727417,
+ "learning_rate": 9.902308769000724e-05,
+ "loss": 0.01476612538099289,
+ "step": 100
+ },
+ {
+ "epoch": 0.38300418910831835,
+ "grad_norm": 0.09139569103717804,
+ "learning_rate": 9.859055934360789e-05,
+ "loss": 0.013924291729927063,
+ "step": 120
+ },
+ {
+ "epoch": 0.4468382206263714,
+ "grad_norm": 0.6957159042358398,
+ "learning_rate": 9.808029767254437e-05,
+ "loss": 0.015140479803085327,
+ "step": 140
+ },
+ {
+ "epoch": 0.5106722521444245,
+ "grad_norm": 0.19447490572929382,
+ "learning_rate": 9.749311981482817e-05,
+ "loss": 0.016634005308151244,
+ "step": 160
+ },
+ {
+ "epoch": 0.5745062836624776,
+ "grad_norm": 0.13639570772647858,
+ "learning_rate": 9.682996608280154e-05,
+ "loss": 0.008574086427688598,
+ "step": 180
+ },
+ {
+ "epoch": 0.6383403151805306,
+ "grad_norm": 0.05407094210386276,
+ "learning_rate": 9.609189845731207e-05,
+ "loss": 0.00992567539215088,
+ "step": 200
+ },
+ {
+ "epoch": 0.7021743466985837,
+ "grad_norm": 0.08482213318347931,
+ "learning_rate": 9.52800988870462e-05,
+ "loss": 0.009285379946231843,
+ "step": 220
+ },
+ {
+ "epoch": 0.7660083782166367,
+ "grad_norm": 0.07621290534734726,
+ "learning_rate": 9.439586739574513e-05,
+ "loss": 0.004985855519771576,
+ "step": 240
+ },
+ {
+ "epoch": 0.8298424097346898,
+ "grad_norm": 0.3552233576774597,
+ "learning_rate": 9.344062000033414e-05,
+ "loss": 0.014777910709381104,
+ "step": 260
+ },
+ {
+ "epoch": 0.8936764412527428,
+ "grad_norm": 0.07681854814291,
+ "learning_rate": 9.241588644329919e-05,
+ "loss": 0.0040129125118255615,
+ "step": 280
+ },
+ {
+ "epoch": 0.9575104727707959,
+ "grad_norm": 0.08489467203617096,
+ "learning_rate": 9.132330774294252e-05,
+ "loss": 0.012399718165397644,
+ "step": 300
+ },
+ {
+ "epoch": 1.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 94.6262,
+ "eval_generator_rouge1": 99.8758,
+ "eval_generator_rouge2": 99.7812,
+ "eval_generator_rougeL": 99.5997,
+ "eval_loss": 0.004354734905064106,
+ "eval_runtime": 1522.2841,
+ "eval_samples_per_second": 2.243,
+ "eval_selector_accuracy": 98.014,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.141,
+ "step": 314
+ },
+ {
+ "epoch": 1.019150209455416,
+ "grad_norm": 0.23293638229370117,
+ "learning_rate": 9.01646335654399e-05,
+ "loss": 0.011312296986579895,
+ "step": 320
+ },
+ {
+ "epoch": 1.082984240973469,
+ "grad_norm": 0.07232551276683807,
+ "learning_rate": 8.894171942290817e-05,
+ "loss": 0.006549905985593796,
+ "step": 340
+ },
+ {
+ "epoch": 1.1468182724915221,
+ "grad_norm": 0.05371030047535896,
+ "learning_rate": 8.765652370197003e-05,
+ "loss": 0.003336009383201599,
+ "step": 360
+ },
+ {
+ "epoch": 1.2106523040095751,
+ "grad_norm": 0.1448812186717987,
+ "learning_rate": 8.631110452757467e-05,
+ "loss": 0.008410868048667908,
+ "step": 380
+ },
+ {
+ "epoch": 1.2744863355276281,
+ "grad_norm": 0.0636834055185318,
+ "learning_rate": 8.490761646709642e-05,
+ "loss": 0.006242972612380981,
+ "step": 400
+ },
+ {
+ "epoch": 1.3383203670456814,
+ "grad_norm": 0.05049936845898628,
+ "learning_rate": 8.344830707998953e-05,
+ "loss": 0.00744413286447525,
+ "step": 420
+ },
+ {
+ "epoch": 1.4021543985637344,
+ "grad_norm": 0.032405462116003036,
+ "learning_rate": 8.193551331852458e-05,
+ "loss": 0.0031117375940084456,
+ "step": 440
+ },
+ {
+ "epoch": 1.4659884300817874,
+ "grad_norm": 0.13639135658740997,
+ "learning_rate": 8.037165778537033e-05,
+ "loss": 0.005259513109922409,
+ "step": 460
+ },
+ {
+ "epoch": 1.5298224615998404,
+ "grad_norm": 0.021995026618242264,
+ "learning_rate": 7.875924485401427e-05,
+ "loss": 0.004164557158946991,
+ "step": 480
+ },
+ {
+ "epoch": 1.5936564931178934,
+ "grad_norm": 0.01941438391804695,
+ "learning_rate": 7.710085665823444e-05,
+ "loss": 0.0015524964779615401,
+ "step": 500
+ },
+ {
+ "epoch": 1.6574905246359466,
+ "grad_norm": 0.1271296739578247,
+ "learning_rate": 7.539914895704521e-05,
+ "loss": 0.004233330115675926,
+ "step": 520
+ },
+ {
+ "epoch": 1.7213245561539996,
+ "grad_norm": 0.22600136697292328,
+ "learning_rate": 7.365684688173908e-05,
+ "loss": 0.0029299482703208923,
+ "step": 540
+ },
+ {
+ "epoch": 1.7851585876720528,
+ "grad_norm": 0.0755968764424324,
+ "learning_rate": 7.18767405718347e-05,
+ "loss": 0.0044687166810035706,
+ "step": 560
+ },
+ {
+ "epoch": 1.8489926191901058,
+ "grad_norm": 0.08520206063985825,
+ "learning_rate": 7.006168070692036e-05,
+ "loss": 0.005520488321781159,
+ "step": 580
+ },
+ {
+ "epoch": 1.9128266507081588,
+ "grad_norm": 0.008491385728120804,
+ "learning_rate": 6.821457394154786e-05,
+ "loss": 0.0015786593779921533,
+ "step": 600
+ },
+ {
+ "epoch": 1.9766606822262118,
+ "grad_norm": 0.009171651676297188,
+ "learning_rate": 6.633837825048768e-05,
+ "loss": 0.007294930517673492,
+ "step": 620
+ },
+ {
+ "epoch": 2.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 96.2617,
+ "eval_generator_rouge1": 99.9321,
+ "eval_generator_rouge2": 99.8752,
+ "eval_generator_rougeL": 99.7349,
+ "eval_loss": 0.0033019811380654573,
+ "eval_runtime": 1482.8547,
+ "eval_samples_per_second": 2.302,
+ "eval_selector_accuracy": 98.4813,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.144,
+ "step": 628
+ },
+ {
+ "epoch": 2.038300418910832,
+ "grad_norm": 0.0519501231610775,
+ "learning_rate": 6.443609819179926e-05,
+ "loss": 0.003084934316575527,
+ "step": 640
+ },
+ {
+ "epoch": 2.102134450428885,
+ "grad_norm": 0.051104024052619934,
+ "learning_rate": 6.251078009530256e-05,
+ "loss": 0.0026566479355096817,
+ "step": 660
+ },
+ {
+ "epoch": 2.165968481946938,
+ "grad_norm": 0.011012597940862179,
+ "learning_rate": 6.0565507184155677e-05,
+ "loss": 0.0016018826514482497,
+ "step": 680
+ },
+ {
+ "epoch": 2.229802513464991,
+ "grad_norm": 0.02783738076686859,
+ "learning_rate": 5.860339463735136e-05,
+ "loss": 0.0031667064875364305,
+ "step": 700
+ },
+ {
+ "epoch": 2.2936365449830443,
+ "grad_norm": 0.008262189105153084,
+ "learning_rate": 5.662758460103895e-05,
+ "loss": 0.0029385803267359734,
+ "step": 720
+ },
+ {
+ "epoch": 2.3574705765010973,
+ "grad_norm": 0.03108038939535618,
+ "learning_rate": 5.464124115666095e-05,
+ "loss": 0.0051560476422309875,
+ "step": 740
+ },
+ {
+ "epoch": 2.4213046080191503,
+ "grad_norm": 0.04321969673037529,
+ "learning_rate": 5.264754525396212e-05,
+ "loss": 0.0011130817234516143,
+ "step": 760
+ },
+ {
+ "epoch": 2.4851386395372033,
+ "grad_norm": 0.022080617025494576,
+ "learning_rate": 5.06496896169857e-05,
+ "loss": 0.002006329409778118,
+ "step": 780
+ },
+ {
+ "epoch": 2.5489726710552563,
+ "grad_norm": 0.04725557193160057,
+ "learning_rate": 4.865087363121384e-05,
+ "loss": 0.0021070817485451697,
+ "step": 800
+ },
+ {
+ "epoch": 2.6128067025733093,
+ "grad_norm": 0.27594882249832153,
+ "learning_rate": 4.6654298220040673e-05,
+ "loss": 0.0016123956069350243,
+ "step": 820
+ },
+ {
+ "epoch": 2.6766407340913627,
+ "grad_norm": 0.0028557360637933016,
+ "learning_rate": 4.466316071878229e-05,
+ "loss": 0.00044896998442709444,
+ "step": 840
+ },
+ {
+ "epoch": 2.7404747656094157,
+ "grad_norm": 0.03328096121549606,
+ "learning_rate": 4.2680649754432913e-05,
+ "loss": 0.001640898734331131,
+ "step": 860
+ },
+ {
+ "epoch": 2.8043087971274687,
+ "grad_norm": 0.03642182797193527,
+ "learning_rate": 4.070994013936646e-05,
+ "loss": 0.0008663022890686989,
+ "step": 880
+ },
+ {
+ "epoch": 2.8681428286455217,
+ "grad_norm": 0.06211363896727562,
+ "learning_rate": 3.8754187787161155e-05,
+ "loss": 0.000627265078946948,
+ "step": 900
+ },
+ {
+ "epoch": 2.9319768601635747,
+ "grad_norm": 0.0025077553000301123,
+ "learning_rate": 3.681652465868875e-05,
+ "loss": 0.0009966723620891571,
+ "step": 920
+ },
+ {
+ "epoch": 2.9958108916816277,
+ "grad_norm": 0.08202111721038818,
+ "learning_rate": 3.490005374656184e-05,
+ "loss": 0.0005241389386355877,
+ "step": 940
+ },
+ {
+ "epoch": 3.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 96.8458,
+ "eval_generator_rouge1": 99.9512,
+ "eval_generator_rouge2": 99.9138,
+ "eval_generator_rougeL": 99.7589,
+ "eval_loss": 0.002848332514986396,
+ "eval_runtime": 1491.6076,
+ "eval_samples_per_second": 2.289,
+ "eval_selector_accuracy": 99.1822,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.143,
+ "step": 942
+ }
+ ],
+ "logging_steps": 20,
+ "max_steps": 1570,
+ "num_input_tokens_seen": 0,
+ "num_train_epochs": 5,
+ "save_steps": 500,
+ "stateful_callbacks": {
+ "TrainerControl": {
+ "args": {
+ "should_epoch_stop": false,
+ "should_evaluate": false,
+ "should_log": false,
+ "should_save": true,
+ "should_training_stop": false
+ },
+ "attributes": {}
+ }
+ },
+ "total_flos": 2.388724412275753e+17,
+ "train_batch_size": 1,
+ "trial_name": null,
+ "trial_params": null
+}
diff --git a/qwen3/teacher_lora/checkpoint-942/training_args.bin b/qwen3/teacher_lora/checkpoint-942/training_args.bin
new file mode 100644
index 0000000000000000000000000000000000000000..1ef98d1a171b439f2f8968f31ccbc846a421fe48
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/training_args.bin
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:84fda57546f783ec338bd3a292c7dba3d917a8a52ac59ca7cebdd546da52b586
+size 7185
diff --git a/qwen3/teacher_lora/checkpoint-942/zero_to_fp32.py b/qwen3/teacher_lora/checkpoint-942/zero_to_fp32.py
new file mode 100644
index 0000000000000000000000000000000000000000..f40f1f0acbbff30e8e8d29906b56aae4e8b8a138
--- /dev/null
+++ b/qwen3/teacher_lora/checkpoint-942/zero_to_fp32.py
@@ -0,0 +1,790 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example:
+# python zero_to_fp32.py . output_dir/
+# or
+# python zero_to_fp32.py . output_dir/ --safe_serialization
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+import gc
+import json
+import numpy as np
+from tqdm import tqdm
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+ FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+ FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS, AUTOEP_LAYERS_KEY,
+ AUTOEP_LAYERS_KEY_LEGACY, AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY,
+ AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT)
+
+
+@dataclass
+class zero_model_state:
+ buffers: dict()
+ param_shapes: dict()
+ shared_params: list
+ ds_version: int
+ frozen_param_shapes: dict()
+ frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+ return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+ '''
+ alist.sort(key=natural_keys) sorts in human order
+ http://nedbatchelder.com/blog/200712/human_sorting.html
+ (See Toothy's implementation in the comments)
+ '''
+ return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+ if not os.path.isdir(checkpoint_dir):
+ raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+ # there should be only one file
+ if zero_stage <= 2:
+ file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+ elif zero_stage == 3:
+ file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+ if not os.path.exists(file):
+ raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+ return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+ # XXX: need to test that this simple glob rule works for multi-node setup too
+ ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+ if len(ckpt_files) == 0:
+ raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+ return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+ return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+ return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def _has_autoep_zero3_partitioned_metadata(state_dict):
+ autoep_layers = state_dict.get(AUTOEP_LAYERS_KEY)
+ if autoep_layers is None:
+ autoep_layers = state_dict.get(AUTOEP_LAYERS_KEY_LEGACY)
+ if not isinstance(autoep_layers, list):
+ return False
+ return any(
+ isinstance(entry, dict)
+ and entry.get(AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY) == AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT
+ for entry in autoep_layers)
+
+
+def _raise_if_autoep_zero3_partitioned_state(state_dict):
+ if _has_autoep_zero3_partitioned_metadata(state_dict):
+ raise NotImplementedError("zero_to_fp32 does not support AutoEP ZeRO-3 partition-native checkpoints. "
+ "AutoEP expert parameters are partitioned over expert replica groups, so "
+ "global data-parallel consolidation would produce incomplete expert tensors. "
+ "Use ds_to_universal.py for expert-aware conversion.")
+
+
+def _raise_if_autoep_zero3_partitioned_checkpoint(model_files):
+ for file in model_files:
+ state_dict = torch.load(file, map_location=device, weights_only=False)
+ _raise_if_autoep_zero3_partitioned_state(state_dict)
+
+
+def parse_model_states(files):
+ zero_model_states = []
+ for file in files:
+ state_dict = torch.load(file, map_location=device, weights_only=False)
+ _raise_if_autoep_zero3_partitioned_state(state_dict)
+
+ if BUFFER_NAMES not in state_dict:
+ raise ValueError(f"{file} is not a model state checkpoint")
+ buffer_names = state_dict[BUFFER_NAMES]
+ if debug:
+ print("Found buffers:", buffer_names)
+
+ # recover just the buffers while restoring them to fp32 if they were saved in fp16
+ buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+ param_shapes = state_dict[PARAM_SHAPES]
+
+ # collect parameters that are included in param_shapes
+ param_names = []
+ for s in param_shapes:
+ for name in s.keys():
+ param_names.append(name)
+
+ # update with frozen parameters
+ frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+ if frozen_param_shapes is not None:
+ if debug:
+ print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+ param_names += list(frozen_param_shapes.keys())
+
+ # handle shared params
+ shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+ ds_version = state_dict.get(DS_VERSION, None)
+
+ frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+ z_model_state = zero_model_state(buffers=buffers,
+ param_shapes=param_shapes,
+ shared_params=shared_params,
+ ds_version=ds_version,
+ frozen_param_shapes=frozen_param_shapes,
+ frozen_param_fragments=frozen_param_fragments)
+ zero_model_states.append(z_model_state)
+
+ return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+ total_files = len(files)
+ state_dicts = []
+ for f in tqdm(files, desc='Loading checkpoint shards'):
+ state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
+ # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
+ # and also handle the case where it was already removed by another helper script
+ state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
+ state_dicts.append(state_dict)
+
+ if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
+ raise ValueError(f"{files[0]} is not a zero checkpoint")
+ zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+ world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+ # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+ # parameters can be different from data parallelism for non-expert parameters. So we can just
+ # use the max of the partition_count to get the dp world_size.
+
+ if type(world_size) is list:
+ world_size = max(world_size)
+
+ if world_size != total_files:
+ raise ValueError(
+ f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+ "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+ )
+
+ # the groups are named differently in each stage
+ if zero_stage <= 2:
+ fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+ elif zero_stage == 3:
+ fp32_groups_key = FP32_FLAT_GROUPS
+ else:
+ raise ValueError(f"unknown zero stage {zero_stage}")
+
+ fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+ return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
+ """
+ Returns fp32 state_dict reconstructed from ds checkpoint
+
+ Args:
+ - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+ """
+ print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+ # parse_model_states rejects AutoEP ZeRO-3 partition-native checkpoints
+ # before the expensive optimizer-shard load below.
+ model_files = get_model_state_files(ds_checkpoint_dir)
+ zero_model_states = parse_model_states(model_files)
+ print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+ optim_files = get_optim_files(ds_checkpoint_dir)
+ zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+ print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+ if zero_stage <= 2:
+ return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters)
+ elif zero_stage == 3:
+ return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+ return
+
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+ frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+ if debug:
+ num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+ print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+ wanted_params = len(frozen_param_shapes)
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+ avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+ print(f'Frozen params: Have {avail_numel} numels to process.')
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+ total_params = 0
+ total_numel = 0
+ for name, shape in frozen_param_shapes.items():
+ total_params += 1
+ unpartitioned_numel = shape.numel()
+ total_numel += unpartitioned_numel
+
+ state_dict[name] = frozen_param_fragments[name]
+
+ if debug:
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _has_callable(obj, fn):
+ attr = getattr(obj, fn, None)
+ return callable(attr)
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+ param_shapes = zero_model_states[0].param_shapes
+
+ # Reconstruction protocol:
+ #
+ # XXX: document this
+
+ if debug:
+ for i in range(world_size):
+ for j in range(len(fp32_flat_groups[0])):
+ print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+ # XXX: memory usage doubles here (zero2)
+ num_param_groups = len(fp32_flat_groups[0])
+ merged_single_partition_of_fp32_groups = []
+ for i in range(num_param_groups):
+ merged_partitions = [sd[i] for sd in fp32_flat_groups]
+ full_single_fp32_vector = torch.cat(merged_partitions, 0)
+ merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+ avail_numel = sum(
+ [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+ if debug:
+ wanted_params = sum([len(shapes) for shapes in param_shapes])
+ wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+ # not asserting if there is a mismatch due to possible padding
+ print(f"Have {avail_numel} numels to process.")
+ print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+ # params
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+ # out-of-core computing solution
+ total_numel = 0
+ total_params = 0
+ for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+ offset = 0
+ avail_numel = full_single_fp32_vector.numel()
+ for name, shape in shapes.items():
+
+ unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
+ total_numel += unpartitioned_numel
+ total_params += 1
+
+ if debug:
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+ state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+ offset += unpartitioned_numel
+
+ # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+ # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+ # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+ # live optimizer object, so we are checking that the numbers are within the right range
+ align_to = 2 * world_size
+
+ def zero2_align(x):
+ return align_to * math.ceil(x / align_to)
+
+ if debug:
+ print(f"original offset={offset}, avail_numel={avail_numel}")
+
+ offset = zero2_align(offset)
+ avail_numel = zero2_align(avail_numel)
+
+ if debug:
+ print(f"aligned offset={offset}, avail_numel={avail_numel}")
+
+ # Sanity check
+ if offset != avail_numel:
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+ print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters):
+ state_dict = OrderedDict()
+
+ # buffers
+ buffers = zero_model_states[0].buffers
+ state_dict.update(buffers)
+ if debug:
+ print(f"added {len(buffers)} buffers")
+
+ if not exclude_frozen_parameters:
+ _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+ _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+ # recover shared parameters
+ for pair in zero_model_states[0].shared_params:
+ if pair[1] in state_dict:
+ state_dict[pair[0]] = state_dict[pair[1]]
+
+ return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+ remainder = unpartitioned_numel % world_size
+ padding_numel = (world_size - remainder) if remainder else 0
+ partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+ return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+ return
+
+ if debug:
+ for i in range(world_size):
+ num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+ print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+ wanted_params = len(frozen_param_shapes)
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+ avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+ print(f'Frozen params: Have {avail_numel} numels to process.')
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+ total_params = 0
+ total_numel = 0
+ for name, shape in zero_model_states[0].frozen_param_shapes.items():
+ total_params += 1
+ unpartitioned_numel = shape.numel()
+ total_numel += unpartitioned_numel
+
+ param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+ state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+ if debug:
+ print(
+ f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+ )
+
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+class GatheredTensor:
+ """
+ A pseudo tensor that collects partitioned weights.
+ It is more memory efficient when there are multiple groups.
+ """
+
+ def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
+ self.flat_groups = flat_groups
+ self.flat_groups_offset = flat_groups_offset
+ self.offset = offset
+ self.partitioned_numel = partitioned_numel
+ self.shape = shape
+ self.dtype = self.flat_groups[0][0].dtype
+
+ def contiguous(self):
+ """
+ Merge partitioned weights from flat_groups into a single tensor.
+ """
+ end_idx = self.offset + self.partitioned_numel
+ world_size = len(self.flat_groups)
+ pad_flat_param_chunks = []
+
+ for rank_i in range(world_size):
+ # for each rank, we need to collect weights from related group/groups
+ flat_groups_at_rank_i = self.flat_groups[rank_i]
+ start_group_id = None
+ end_group_id = None
+ for group_id in range(len(self.flat_groups_offset)):
+ if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
+ start_group_id = group_id
+ if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
+ end_group_id = group_id
+ break
+ # collect weights from related group/groups
+ for group_id in range(start_group_id, end_group_id + 1):
+ flat_tensor = flat_groups_at_rank_i[group_id]
+ start_offset = self.offset - self.flat_groups_offset[group_id]
+ end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
+ pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
+
+ # collect weights from all ranks
+ pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
+ param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
+ return param
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+ param_shapes = zero_model_states[0].param_shapes
+ avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
+
+ # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+ # param, re-consolidating each param, while dealing with padding if any
+
+ # merge list of dicts, preserving order
+ param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+ if debug:
+ for i in range(world_size):
+ print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+ wanted_params = len(param_shapes)
+ wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+ # not asserting if there is a mismatch due to possible padding
+ avail_numel = fp32_flat_groups[0].numel() * world_size
+ print(f"Trainable params: Have {avail_numel} numels to process.")
+ print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+ # params
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+ # out-of-core computing solution
+ offset = 0
+ total_numel = 0
+ total_params = 0
+ flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
+ for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
+ unpartitioned_numel = shape.numel()
+ total_numel += unpartitioned_numel
+ total_params += 1
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+ if debug:
+ print(
+ f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+ )
+
+ # memory efficient tensor
+ tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
+ state_dict[name] = tensor
+ offset += partitioned_numel
+
+ offset *= world_size
+
+ # Sanity check
+ if offset != avail_numel:
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+ print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+ exclude_frozen_parameters):
+ state_dict = OrderedDict()
+
+ # buffers
+ buffers = zero_model_states[0].buffers
+ state_dict.update(buffers)
+ if debug:
+ print(f"added {len(buffers)} buffers")
+
+ if not exclude_frozen_parameters:
+ _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+ _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+ # recover shared parameters
+ for pair in zero_model_states[0].shared_params:
+ if pair[1] in state_dict:
+ state_dict[pair[0]] = state_dict[pair[1]]
+
+ return state_dict
+
+
+def to_torch_tensor(state_dict, return_empty_tensor=False):
+ """
+ Convert state_dict of GatheredTensor to torch tensor
+ """
+ torch_state_dict = {}
+ converted_tensors = {}
+ for name, tensor in state_dict.items():
+ tensor_id = id(tensor)
+ if tensor_id in converted_tensors: # shared tensors
+ shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
+ torch_state_dict[name] = shared_tensor
+ else:
+ converted_tensors[tensor_id] = name
+ if return_empty_tensor:
+ torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
+ else:
+ torch_state_dict[name] = tensor.contiguous()
+ return torch_state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+ tag=None,
+ exclude_frozen_parameters=False,
+ lazy_mode=False):
+ """
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+ ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+ via a model hub.
+
+ Args:
+ - ``checkpoint_dir``: path to the desired checkpoint folder
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+ - ``exclude_frozen_parameters``: exclude frozen parameters
+ - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
+ Convert the pesduo tensor to torch tensor by ``.contiguous()``
+
+ Returns:
+ - pytorch ``state_dict``
+
+ A typical usage might be ::
+
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+ # do the training and checkpoint saving
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+ model = model.cpu() # move to cpu
+ model.load_state_dict(state_dict)
+ # submit to model hub or save the model to share with others
+
+ In this example the ``model`` will no longer be usable in the deepspeed context of the same
+ application. i.e. you will need to re-initialize the deepspeed engine, since
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+ If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+ Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
+ You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+ the checkpoint. Or you can load state_dict in lazy mode ::
+
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
+ for name, lazy_tensor in state_dict.item():
+ tensor = lazy_tensor.contiguous() # to cpu
+ print(name, tensor)
+ # del tensor to release memory if it no longer in use
+ """
+ if tag is None:
+ latest_path = os.path.join(checkpoint_dir, 'latest')
+ if os.path.isfile(latest_path):
+ with open(latest_path, 'r') as fd:
+ tag = fd.read().strip()
+ else:
+ raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+ ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+ if not os.path.isdir(ds_checkpoint_dir):
+ raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+ state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
+ if lazy_mode:
+ return state_dict
+ else:
+ return to_torch_tensor(state_dict)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
+ output_dir,
+ max_shard_size="5GB",
+ safe_serialization=False,
+ tag=None,
+ exclude_frozen_parameters=False):
+ """
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+ loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+ Args:
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+ - ``output_dir``: directory to the pytorch fp32 state_dict output files
+ - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
+ - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+ - ``exclude_frozen_parameters``: exclude frozen parameters
+ """
+
+ # Dependency pre-check
+ if safe_serialization:
+ try:
+ from safetensors.torch import save_file
+ except ImportError:
+ print('If you want to use `safe_serialization`, please `pip install safetensors`')
+ raise
+ if max_shard_size is not None:
+ try:
+ from huggingface_hub import split_torch_state_dict_into_shards
+ except ImportError:
+ print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
+ raise
+
+ # Convert zero checkpoint to state_dict
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+ tag,
+ exclude_frozen_parameters,
+ lazy_mode=True)
+
+ # Shard the model if it is too big.
+ weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
+ if max_shard_size is not None:
+ filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
+ # an memory-efficient approach for sharding
+ empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
+ state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
+ filename_pattern=filename_pattern,
+ max_shard_size=max_shard_size)
+ else:
+ from collections import namedtuple
+ StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
+ state_dict_split = StateDictSplit(is_sharded=False,
+ filename_to_tensors={weights_name: list(state_dict.keys())})
+
+ # Save the model by shard
+ os.makedirs(output_dir, exist_ok=True)
+ filename_to_tensors = state_dict_split.filename_to_tensors.items()
+ for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
+ shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
+ shard_state_dict = to_torch_tensor(shard_state_dict)
+ output_path = os.path.join(output_dir, shard_file)
+ if safe_serialization:
+ save_file(shard_state_dict, output_path, metadata={"format": "pt"})
+ else:
+ torch.save(shard_state_dict, output_path)
+ # release the memory of current shard
+ for tensor_name in list(shard_state_dict.keys()):
+ del state_dict[tensor_name]
+ del shard_state_dict[tensor_name]
+ del shard_state_dict
+ gc.collect()
+
+ # Save index if sharded
+ if state_dict_split.is_sharded:
+ index = {
+ "metadata": state_dict_split.metadata,
+ "weight_map": state_dict_split.tensor_to_filename,
+ }
+ save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
+ save_index_file = os.path.join(output_dir, save_index_file)
+ with open(save_index_file, "w", encoding="utf-8") as f:
+ content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+ f.write(content)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+ """
+ 1. Put the provided model to cpu
+ 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+ 3. Load it into the provided model
+
+ Args:
+ - ``model``: the model object to update
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+ Returns:
+ - ``model`: modified model
+
+ Make sure you have plenty of CPU memory available before you call this function. If you don't
+ have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+ conveniently placed for you in the checkpoint folder.
+
+ A typical usage might be ::
+
+ from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+ model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+ # submit to model hub or save the model to share with others
+
+ Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+ of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+ """
+ logger.info("Extracting fp32 weights")
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+ logger.info("Overwriting model with fp32 weights")
+ model = model.cpu()
+ model.load_state_dict(state_dict, strict=False)
+
+ return model
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("checkpoint_dir",
+ type=str,
+ help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+ parser.add_argument("output_dir",
+ type=str,
+ help="directory to the pytorch fp32 state_dict output files"
+ "(e.g. path/checkpoint-12-output/)")
+ parser.add_argument(
+ "--max_shard_size",
+ type=str,
+ default="5GB",
+ help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
+ "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
+ "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
+ "without CPU OOM issues.")
+ parser.add_argument(
+ "--safe_serialization",
+ default=False,
+ action='store_true',
+ help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
+ parser.add_argument("-t",
+ "--tag",
+ type=str,
+ default=None,
+ help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
+ parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
+ parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+ args = parser.parse_args()
+
+ debug = args.debug
+
+ convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
+ args.output_dir,
+ max_shard_size=args.max_shard_size,
+ safe_serialization=args.safe_serialization,
+ tag=args.tag,
+ exclude_frozen_parameters=args.exclude_frozen_parameters)
diff --git a/qwen3/teacher_lora/eval_results.json b/qwen3/teacher_lora/eval_results.json
new file mode 100644
index 0000000000000000000000000000000000000000..269181d4d2f2b41eac908b2e53c89ec79884553a
--- /dev/null
+++ b/qwen3/teacher_lora/eval_results.json
@@ -0,0 +1,14 @@
+{
+ "epoch": 5.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 97.0794,
+ "eval_generator_rouge1": 99.9517,
+ "eval_generator_rouge2": 99.914,
+ "eval_generator_rougeL": 99.7744,
+ "eval_loss": 0.00291790091432631,
+ "eval_runtime": 1472.9037,
+ "eval_samples_per_second": 2.318,
+ "eval_selector_accuracy": 99.2407,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.145
+}
\ No newline at end of file
diff --git a/qwen3/teacher_lora/tokenizer.json b/qwen3/teacher_lora/tokenizer.json
new file mode 100644
index 0000000000000000000000000000000000000000..c7afbed2efcdf019f88ab0572ec29d3bf595dfe2
--- /dev/null
+++ b/qwen3/teacher_lora/tokenizer.json
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506
+size 11422650
diff --git a/qwen3/teacher_lora/tokenizer_config.json b/qwen3/teacher_lora/tokenizer_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..907e863a5a4109fb0e373ccdeb4d657bbeca9416
--- /dev/null
+++ b/qwen3/teacher_lora/tokenizer_config.json
@@ -0,0 +1,31 @@
+{
+ "add_prefix_space": false,
+ "backend": "tokenizers",
+ "bos_token": null,
+ "clean_up_tokenization_spaces": false,
+ "eos_token": "<|im_end|>",
+ "errors": "replace",
+ "extra_special_tokens": [
+ "<|im_start|>",
+ "<|im_end|>",
+ "<|object_ref_start|>",
+ "<|object_ref_end|>",
+ "<|box_start|>",
+ "<|box_end|>",
+ "<|quad_start|>",
+ "<|quad_end|>",
+ "<|vision_start|>",
+ "<|vision_end|>",
+ "<|vision_pad|>",
+ "<|image_pad|>",
+ "<|video_pad|>"
+ ],
+ "is_local": false,
+ "local_files_only": false,
+ "model_max_length": 1010000,
+ "pad_token": "<|endoftext|>",
+ "padding_side": "right",
+ "split_special_tokens": false,
+ "tokenizer_class": "Qwen2Tokenizer",
+ "unk_token": null
+}
diff --git a/qwen3/teacher_lora/train_results.json b/qwen3/teacher_lora/train_results.json
new file mode 100644
index 0000000000000000000000000000000000000000..9a8351fde29920c793d91870479cdc25710797ba
--- /dev/null
+++ b/qwen3/teacher_lora/train_results.json
@@ -0,0 +1,8 @@
+{
+ "epoch": 5.0,
+ "total_flos": 3.981241098143007e+17,
+ "train_loss": 0.009030620363736704,
+ "train_runtime": 16300.8877,
+ "train_samples_per_second": 3.076,
+ "train_steps_per_second": 0.096
+}
\ No newline at end of file
diff --git a/qwen3/teacher_lora/trainer_state.json b/qwen3/teacher_lora/trainer_state.json
new file mode 100644
index 0000000000000000000000000000000000000000..4415da34c9f9fd1cc5c07545fae64232f6c28914
--- /dev/null
+++ b/qwen3/teacher_lora/trainer_state.json
@@ -0,0 +1,664 @@
+{
+ "best_global_step": null,
+ "best_metric": null,
+ "best_model_checkpoint": null,
+ "epoch": 5.0,
+ "eval_steps": 500,
+ "global_step": 1570,
+ "is_hyper_param_search": false,
+ "is_local_process_zero": true,
+ "is_world_process_zero": true,
+ "log_history": [
+ {
+ "epoch": 0.06383403151805306,
+ "grad_norm": 0.4070272743701935,
+ "learning_rate": 9.996390385474569e-05,
+ "loss": 0.37070536613464355,
+ "step": 20
+ },
+ {
+ "epoch": 0.12766806303610612,
+ "grad_norm": 0.5005771517753601,
+ "learning_rate": 9.984797509888031e-05,
+ "loss": 0.03893736004829407,
+ "step": 40
+ },
+ {
+ "epoch": 0.19150209455415917,
+ "grad_norm": 0.17237664759159088,
+ "learning_rate": 9.965229938050925e-05,
+ "loss": 0.02901872992515564,
+ "step": 60
+ },
+ {
+ "epoch": 0.25533612607221223,
+ "grad_norm": 0.17700566351413727,
+ "learning_rate": 9.937719005663432e-05,
+ "loss": 0.012938253581523895,
+ "step": 80
+ },
+ {
+ "epoch": 0.3191701575902653,
+ "grad_norm": 0.29617995023727417,
+ "learning_rate": 9.902308769000724e-05,
+ "loss": 0.01476612538099289,
+ "step": 100
+ },
+ {
+ "epoch": 0.38300418910831835,
+ "grad_norm": 0.09139569103717804,
+ "learning_rate": 9.859055934360789e-05,
+ "loss": 0.013924291729927063,
+ "step": 120
+ },
+ {
+ "epoch": 0.4468382206263714,
+ "grad_norm": 0.6957159042358398,
+ "learning_rate": 9.808029767254437e-05,
+ "loss": 0.015140479803085327,
+ "step": 140
+ },
+ {
+ "epoch": 0.5106722521444245,
+ "grad_norm": 0.19447490572929382,
+ "learning_rate": 9.749311981482817e-05,
+ "loss": 0.016634005308151244,
+ "step": 160
+ },
+ {
+ "epoch": 0.5745062836624776,
+ "grad_norm": 0.13639570772647858,
+ "learning_rate": 9.682996608280154e-05,
+ "loss": 0.008574086427688598,
+ "step": 180
+ },
+ {
+ "epoch": 0.6383403151805306,
+ "grad_norm": 0.05407094210386276,
+ "learning_rate": 9.609189845731207e-05,
+ "loss": 0.00992567539215088,
+ "step": 200
+ },
+ {
+ "epoch": 0.7021743466985837,
+ "grad_norm": 0.08482213318347931,
+ "learning_rate": 9.52800988870462e-05,
+ "loss": 0.009285379946231843,
+ "step": 220
+ },
+ {
+ "epoch": 0.7660083782166367,
+ "grad_norm": 0.07621290534734726,
+ "learning_rate": 9.439586739574513e-05,
+ "loss": 0.004985855519771576,
+ "step": 240
+ },
+ {
+ "epoch": 0.8298424097346898,
+ "grad_norm": 0.3552233576774597,
+ "learning_rate": 9.344062000033414e-05,
+ "loss": 0.014777910709381104,
+ "step": 260
+ },
+ {
+ "epoch": 0.8936764412527428,
+ "grad_norm": 0.07681854814291,
+ "learning_rate": 9.241588644329919e-05,
+ "loss": 0.0040129125118255615,
+ "step": 280
+ },
+ {
+ "epoch": 0.9575104727707959,
+ "grad_norm": 0.08489467203617096,
+ "learning_rate": 9.132330774294252e-05,
+ "loss": 0.012399718165397644,
+ "step": 300
+ },
+ {
+ "epoch": 1.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 94.6262,
+ "eval_generator_rouge1": 99.8758,
+ "eval_generator_rouge2": 99.7812,
+ "eval_generator_rougeL": 99.5997,
+ "eval_loss": 0.004354734905064106,
+ "eval_runtime": 1522.2841,
+ "eval_samples_per_second": 2.243,
+ "eval_selector_accuracy": 98.014,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.141,
+ "step": 314
+ },
+ {
+ "epoch": 1.019150209455416,
+ "grad_norm": 0.23293638229370117,
+ "learning_rate": 9.01646335654399e-05,
+ "loss": 0.011312296986579895,
+ "step": 320
+ },
+ {
+ "epoch": 1.082984240973469,
+ "grad_norm": 0.07232551276683807,
+ "learning_rate": 8.894171942290817e-05,
+ "loss": 0.006549905985593796,
+ "step": 340
+ },
+ {
+ "epoch": 1.1468182724915221,
+ "grad_norm": 0.05371030047535896,
+ "learning_rate": 8.765652370197003e-05,
+ "loss": 0.003336009383201599,
+ "step": 360
+ },
+ {
+ "epoch": 1.2106523040095751,
+ "grad_norm": 0.1448812186717987,
+ "learning_rate": 8.631110452757467e-05,
+ "loss": 0.008410868048667908,
+ "step": 380
+ },
+ {
+ "epoch": 1.2744863355276281,
+ "grad_norm": 0.0636834055185318,
+ "learning_rate": 8.490761646709642e-05,
+ "loss": 0.006242972612380981,
+ "step": 400
+ },
+ {
+ "epoch": 1.3383203670456814,
+ "grad_norm": 0.05049936845898628,
+ "learning_rate": 8.344830707998953e-05,
+ "loss": 0.00744413286447525,
+ "step": 420
+ },
+ {
+ "epoch": 1.4021543985637344,
+ "grad_norm": 0.032405462116003036,
+ "learning_rate": 8.193551331852458e-05,
+ "loss": 0.0031117375940084456,
+ "step": 440
+ },
+ {
+ "epoch": 1.4659884300817874,
+ "grad_norm": 0.13639135658740997,
+ "learning_rate": 8.037165778537033e-05,
+ "loss": 0.005259513109922409,
+ "step": 460
+ },
+ {
+ "epoch": 1.5298224615998404,
+ "grad_norm": 0.021995026618242264,
+ "learning_rate": 7.875924485401427e-05,
+ "loss": 0.004164557158946991,
+ "step": 480
+ },
+ {
+ "epoch": 1.5936564931178934,
+ "grad_norm": 0.01941438391804695,
+ "learning_rate": 7.710085665823444e-05,
+ "loss": 0.0015524964779615401,
+ "step": 500
+ },
+ {
+ "epoch": 1.6574905246359466,
+ "grad_norm": 0.1271296739578247,
+ "learning_rate": 7.539914895704521e-05,
+ "loss": 0.004233330115675926,
+ "step": 520
+ },
+ {
+ "epoch": 1.7213245561539996,
+ "grad_norm": 0.22600136697292328,
+ "learning_rate": 7.365684688173908e-05,
+ "loss": 0.0029299482703208923,
+ "step": 540
+ },
+ {
+ "epoch": 1.7851585876720528,
+ "grad_norm": 0.0755968764424324,
+ "learning_rate": 7.18767405718347e-05,
+ "loss": 0.0044687166810035706,
+ "step": 560
+ },
+ {
+ "epoch": 1.8489926191901058,
+ "grad_norm": 0.08520206063985825,
+ "learning_rate": 7.006168070692036e-05,
+ "loss": 0.005520488321781159,
+ "step": 580
+ },
+ {
+ "epoch": 1.9128266507081588,
+ "grad_norm": 0.008491385728120804,
+ "learning_rate": 6.821457394154786e-05,
+ "loss": 0.0015786593779921533,
+ "step": 600
+ },
+ {
+ "epoch": 1.9766606822262118,
+ "grad_norm": 0.009171651676297188,
+ "learning_rate": 6.633837825048768e-05,
+ "loss": 0.007294930517673492,
+ "step": 620
+ },
+ {
+ "epoch": 2.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 96.2617,
+ "eval_generator_rouge1": 99.9321,
+ "eval_generator_rouge2": 99.8752,
+ "eval_generator_rougeL": 99.7349,
+ "eval_loss": 0.0033019811380654573,
+ "eval_runtime": 1482.8547,
+ "eval_samples_per_second": 2.302,
+ "eval_selector_accuracy": 98.4813,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.144,
+ "step": 628
+ },
+ {
+ "epoch": 2.038300418910832,
+ "grad_norm": 0.0519501231610775,
+ "learning_rate": 6.443609819179926e-05,
+ "loss": 0.003084934316575527,
+ "step": 640
+ },
+ {
+ "epoch": 2.102134450428885,
+ "grad_norm": 0.051104024052619934,
+ "learning_rate": 6.251078009530256e-05,
+ "loss": 0.0026566479355096817,
+ "step": 660
+ },
+ {
+ "epoch": 2.165968481946938,
+ "grad_norm": 0.011012597940862179,
+ "learning_rate": 6.0565507184155677e-05,
+ "loss": 0.0016018826514482497,
+ "step": 680
+ },
+ {
+ "epoch": 2.229802513464991,
+ "grad_norm": 0.02783738076686859,
+ "learning_rate": 5.860339463735136e-05,
+ "loss": 0.0031667064875364305,
+ "step": 700
+ },
+ {
+ "epoch": 2.2936365449830443,
+ "grad_norm": 0.008262189105153084,
+ "learning_rate": 5.662758460103895e-05,
+ "loss": 0.0029385803267359734,
+ "step": 720
+ },
+ {
+ "epoch": 2.3574705765010973,
+ "grad_norm": 0.03108038939535618,
+ "learning_rate": 5.464124115666095e-05,
+ "loss": 0.0051560476422309875,
+ "step": 740
+ },
+ {
+ "epoch": 2.4213046080191503,
+ "grad_norm": 0.04321969673037529,
+ "learning_rate": 5.264754525396212e-05,
+ "loss": 0.0011130817234516143,
+ "step": 760
+ },
+ {
+ "epoch": 2.4851386395372033,
+ "grad_norm": 0.022080617025494576,
+ "learning_rate": 5.06496896169857e-05,
+ "loss": 0.002006329409778118,
+ "step": 780
+ },
+ {
+ "epoch": 2.5489726710552563,
+ "grad_norm": 0.04725557193160057,
+ "learning_rate": 4.865087363121384e-05,
+ "loss": 0.0021070817485451697,
+ "step": 800
+ },
+ {
+ "epoch": 2.6128067025733093,
+ "grad_norm": 0.27594882249832153,
+ "learning_rate": 4.6654298220040673e-05,
+ "loss": 0.0016123956069350243,
+ "step": 820
+ },
+ {
+ "epoch": 2.6766407340913627,
+ "grad_norm": 0.0028557360637933016,
+ "learning_rate": 4.466316071878229e-05,
+ "loss": 0.00044896998442709444,
+ "step": 840
+ },
+ {
+ "epoch": 2.7404747656094157,
+ "grad_norm": 0.03328096121549606,
+ "learning_rate": 4.2680649754432913e-05,
+ "loss": 0.001640898734331131,
+ "step": 860
+ },
+ {
+ "epoch": 2.8043087971274687,
+ "grad_norm": 0.03642182797193527,
+ "learning_rate": 4.070994013936646e-05,
+ "loss": 0.0008663022890686989,
+ "step": 880
+ },
+ {
+ "epoch": 2.8681428286455217,
+ "grad_norm": 0.06211363896727562,
+ "learning_rate": 3.8754187787161155e-05,
+ "loss": 0.000627265078946948,
+ "step": 900
+ },
+ {
+ "epoch": 2.9319768601635747,
+ "grad_norm": 0.0025077553000301123,
+ "learning_rate": 3.681652465868875e-05,
+ "loss": 0.0009966723620891571,
+ "step": 920
+ },
+ {
+ "epoch": 2.9958108916816277,
+ "grad_norm": 0.08202111721038818,
+ "learning_rate": 3.490005374656184e-05,
+ "loss": 0.0005241389386355877,
+ "step": 940
+ },
+ {
+ "epoch": 3.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 96.8458,
+ "eval_generator_rouge1": 99.9512,
+ "eval_generator_rouge2": 99.9138,
+ "eval_generator_rougeL": 99.7589,
+ "eval_loss": 0.002848332514986396,
+ "eval_runtime": 1491.6076,
+ "eval_samples_per_second": 2.289,
+ "eval_selector_accuracy": 99.1822,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.143,
+ "step": 942
+ },
+ {
+ "epoch": 3.0574506283662477,
+ "grad_norm": 0.05988731235265732,
+ "learning_rate": 3.3007844105971336e-05,
+ "loss": 0.003412053734064102,
+ "step": 960
+ },
+ {
+ "epoch": 3.1212846598843007,
+ "grad_norm": 0.0027547082863748074,
+ "learning_rate": 3.114292593987149e-05,
+ "loss": 0.0007339773699641227,
+ "step": 980
+ },
+ {
+ "epoch": 3.1851186914023537,
+ "grad_norm": 0.004056593868881464,
+ "learning_rate": 2.9308285746383408e-05,
+ "loss": 0.0007870178669691086,
+ "step": 1000
+ },
+ {
+ "epoch": 3.248952722920407,
+ "grad_norm": 0.02156253345310688,
+ "learning_rate": 2.7506861536187867e-05,
+ "loss": 0.0004239068366587162,
+ "step": 1020
+ },
+ {
+ "epoch": 3.31278675443846,
+ "grad_norm": 0.01963026262819767,
+ "learning_rate": 2.5741538127566396e-05,
+ "loss": 0.0004422076977789402,
+ "step": 1040
+ },
+ {
+ "epoch": 3.376620785956513,
+ "grad_norm": 0.0010617657098919153,
+ "learning_rate": 2.4015142526625115e-05,
+ "loss": 0.001177133247256279,
+ "step": 1060
+ },
+ {
+ "epoch": 3.440454817474566,
+ "grad_norm": 0.04265115037560463,
+ "learning_rate": 2.233043940009967e-05,
+ "loss": 0.00048412443138659,
+ "step": 1080
+ },
+ {
+ "epoch": 3.504288848992619,
+ "grad_norm": 0.0026085481513291597,
+ "learning_rate": 2.0690126647990974e-05,
+ "loss": 0.0007659945171326399,
+ "step": 1100
+ },
+ {
+ "epoch": 3.568122880510672,
+ "grad_norm": 0.008790579624474049,
+ "learning_rate": 1.9096831083121828e-05,
+ "loss": 0.000472747627645731,
+ "step": 1120
+ },
+ {
+ "epoch": 3.631956912028725,
+ "grad_norm": 0.010375800542533398,
+ "learning_rate": 1.7553104224533324e-05,
+ "loss": 0.0004488838836550713,
+ "step": 1140
+ },
+ {
+ "epoch": 3.695790943546778,
+ "grad_norm": 0.005610933993011713,
+ "learning_rate": 1.606141821145735e-05,
+ "loss": 0.0006687684915959835,
+ "step": 1160
+ },
+ {
+ "epoch": 3.759624975064831,
+ "grad_norm": 0.01078067161142826,
+ "learning_rate": 1.4624161844408763e-05,
+ "loss": 0.0008440790697932243,
+ "step": 1180
+ },
+ {
+ "epoch": 3.8234590065828846,
+ "grad_norm": 0.009699724614620209,
+ "learning_rate": 1.324363675973698e-05,
+ "loss": 0.0008271764032542706,
+ "step": 1200
+ },
+ {
+ "epoch": 3.8872930381009376,
+ "grad_norm": 0.018169160932302475,
+ "learning_rate": 1.1922053743763052e-05,
+ "loss": 0.0003760921768844128,
+ "step": 1220
+ },
+ {
+ "epoch": 3.9511270696189906,
+ "grad_norm": 0.010761883109807968,
+ "learning_rate": 1.0661529192404973e-05,
+ "loss": 0.0003943405114114285,
+ "step": 1240
+ },
+ {
+ "epoch": 4.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 96.7874,
+ "eval_generator_rouge1": 99.9518,
+ "eval_generator_rouge2": 99.9142,
+ "eval_generator_rougeL": 99.7472,
+ "eval_loss": 0.0028539097402244806,
+ "eval_runtime": 1492.5205,
+ "eval_samples_per_second": 2.287,
+ "eval_selector_accuracy": 99.1822,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.143,
+ "step": 1256
+ },
+ {
+ "epoch": 4.012766806303611,
+ "grad_norm": 0.0212729349732399,
+ "learning_rate": 9.464081721960592e-06,
+ "loss": 0.0008496158756315708,
+ "step": 1260
+ },
+ {
+ "epoch": 4.076600837821664,
+ "grad_norm": 0.04593608155846596,
+ "learning_rate": 8.331628936475818e-06,
+ "loss": 0.00023648960050195455,
+ "step": 1280
+ },
+ {
+ "epoch": 4.140434869339717,
+ "grad_norm": 0.00037404155591502786,
+ "learning_rate": 7.2659843568748385e-06,
+ "loss": 0.00027941586449742317,
+ "step": 1300
+ },
+ {
+ "epoch": 4.20426890085777,
+ "grad_norm": 0.0028630236629396677,
+ "learning_rate": 6.268854516770008e-06,
+ "loss": 0.00032952893525362015,
+ "step": 1320
+ },
+ {
+ "epoch": 4.268102932375823,
+ "grad_norm": 0.03912129998207092,
+ "learning_rate": 5.341836229602331e-06,
+ "loss": 0.0003575342707335949,
+ "step": 1340
+ },
+ {
+ "epoch": 4.331936963893876,
+ "grad_norm": 0.0010673116194084287,
+ "learning_rate": 4.4864140314888216e-06,
+ "loss": 0.0003374688792973757,
+ "step": 1360
+ },
+ {
+ "epoch": 4.395770995411929,
+ "grad_norm": 0.0022425639908760786,
+ "learning_rate": 3.703957803871932e-06,
+ "loss": 0.0008516624569892883,
+ "step": 1380
+ },
+ {
+ "epoch": 4.459605026929982,
+ "grad_norm": 0.01993773877620697,
+ "learning_rate": 2.9957205797780495e-06,
+ "loss": 0.000931937713176012,
+ "step": 1400
+ },
+ {
+ "epoch": 4.5234390584480355,
+ "grad_norm": 0.027315137907862663,
+ "learning_rate": 2.362836537198199e-06,
+ "loss": 0.0003644102718681097,
+ "step": 1420
+ },
+ {
+ "epoch": 4.5872730899660885,
+ "grad_norm": 0.007133164908736944,
+ "learning_rate": 1.8063191828043703e-06,
+ "loss": 0.00021162929479032755,
+ "step": 1440
+ },
+ {
+ "epoch": 4.6511071214841415,
+ "grad_norm": 0.02018481120467186,
+ "learning_rate": 1.3270597289100437e-06,
+ "loss": 0.00022698421962559224,
+ "step": 1460
+ },
+ {
+ "epoch": 4.7149411530021945,
+ "grad_norm": 0.014296010136604309,
+ "learning_rate": 9.258256662741205e-07,
+ "loss": 0.00022367050405591726,
+ "step": 1480
+ },
+ {
+ "epoch": 4.7787751845202475,
+ "grad_norm": 0.007890012115240097,
+ "learning_rate": 6.032595350337341e-07,
+ "loss": 0.00043175043538212776,
+ "step": 1500
+ },
+ {
+ "epoch": 4.8426092160383005,
+ "grad_norm": 0.022568073123693466,
+ "learning_rate": 3.5987789573421356e-07,
+ "loss": 0.00024233334697782994,
+ "step": 1520
+ },
+ {
+ "epoch": 4.9064432475563535,
+ "grad_norm": 0.0283685103058815,
+ "learning_rate": 1.9607050210401192e-07,
+ "loss": 0.0002710998523980379,
+ "step": 1540
+ },
+ {
+ "epoch": 4.9702772790744065,
+ "grad_norm": 0.019454164430499077,
+ "learning_rate": 1.1209967689925832e-07,
+ "loss": 0.00033448461908847094,
+ "step": 1560
+ },
+ {
+ "epoch": 5.0,
+ "eval_generator_count": 1712.0,
+ "eval_generator_exact_match": 97.0794,
+ "eval_generator_rouge1": 99.9483,
+ "eval_generator_rouge2": 99.9036,
+ "eval_generator_rougeL": 99.7893,
+ "eval_loss": 0.00291790091432631,
+ "eval_runtime": 1494.8716,
+ "eval_samples_per_second": 2.284,
+ "eval_selector_accuracy": 99.2407,
+ "eval_selector_count": 1712.0,
+ "eval_steps_per_second": 0.143,
+ "step": 1570
+ },
+ {
+ "epoch": 5.0,
+ "step": 1570,
+ "total_flos": 3.981241098143007e+17,
+ "train_loss": 0.009030620363736704,
+ "train_runtime": 16300.8877,
+ "train_samples_per_second": 3.076,
+ "train_steps_per_second": 0.096
+ }
+ ],
+ "logging_steps": 20,
+ "max_steps": 1570,
+ "num_input_tokens_seen": 0,
+ "num_train_epochs": 5,
+ "save_steps": 500,
+ "stateful_callbacks": {
+ "TrainerControl": {
+ "args": {
+ "should_epoch_stop": false,
+ "should_evaluate": false,
+ "should_log": false,
+ "should_save": true,
+ "should_training_stop": true
+ },
+ "attributes": {}
+ }
+ },
+ "total_flos": 3.981241098143007e+17,
+ "train_batch_size": 1,
+ "trial_name": null,
+ "trial_params": null
+}
diff --git a/qwen3/teacher_lora/training_args.bin b/qwen3/teacher_lora/training_args.bin
new file mode 100644
index 0000000000000000000000000000000000000000..1ef98d1a171b439f2f8968f31ccbc846a421fe48
--- /dev/null
+++ b/qwen3/teacher_lora/training_args.bin
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:84fda57546f783ec338bd3a292c7dba3d917a8a52ac59ca7cebdd546da52b586
+size 7185