cjiao/OpenThoughts3-greedy-groups-top-openthinker3-1.5B-checkpoint-375-length-filtered
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: cjiao/OpenThoughts3-greedy-groups-top-openthinker3-1.5B-checkpoint-375-length-filtered
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-04-27 01:24:02 +08:00
commit 85ac608c11
49 changed files with 457662 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
checkpoint-100/added_tokens.json Normal file
View File

@@ -0,0 +1,24 @@
{
"</tool_call>": 151658,
"<tool_call>": 151657,
"<|box_end|>": 151649,
"<|box_start|>": 151648,
"<|endoftext|>": 151643,
"<|file_sep|>": 151664,
"<|fim_middle|>": 151660,
"<|fim_pad|>": 151662,
"<|fim_prefix|>": 151659,
"<|fim_suffix|>": 151661,
"<|im_end|>": 151645,
"<|im_start|>": 151644,
"<|image_pad|>": 151655,
"<|object_ref_end|>": 151647,
"<|object_ref_start|>": 151646,
"<|quad_end|>": 151651,
"<|quad_start|>": 151650,
"<|repo_name|>": 151663,
"<|video_pad|>": 151656,
"<|vision_end|>": 151653,
"<|vision_pad|>": 151654,
"<|vision_start|>": 151652
}

29
checkpoint-100/config.json Normal file
View File

@@ -0,0 +1,29 @@
{
"_name_or_path": "/data/hf_cache/openthoughts_outputs/OpenThinker3-1.5B/checkpoint-375",
"architectures": [
"Qwen2ForCausalLM"
],
"attention_dropout": 0.0,
"bos_token_id": 151643,
"eos_token_id": 151645,
"hidden_act": "silu",
"hidden_size": 1536,
"initializer_range": 0.02,
"intermediate_size": 8960,
"max_position_embeddings": 32768,
"max_window_layers": 21,
"model_type": "qwen2",
"num_attention_heads": 12,
"num_hidden_layers": 28,
"num_key_value_heads": 2,
"rms_norm_eps": 1e-06,
"rope_scaling": null,
"rope_theta": 1000000.0,
"sliding_window": null,
"tie_word_embeddings": true,
"torch_dtype": "bfloat16",
"transformers_version": "4.46.1",
"use_cache": false,
"use_sliding_window": false,
"vocab_size": 151936
}

14
checkpoint-100/generation_config.json Normal file
View File

@@ -0,0 +1,14 @@
{
"bos_token_id": 151643,
"do_sample": true,
"eos_token_id": [
151645,
151643
],
"pad_token_id": 151643,
"repetition_penalty": 1.1,
"temperature": 0.7,
"top_k": 20,
"top_p": 0.8,
"transformers_version": "4.46.1"
}

1
checkpoint-100/latest Normal file
View File

@@ -0,0 +1 @@
global_step100

151388
checkpoint-100/merges.txt Normal file
View File

File diff suppressed because it is too large Load Diff

3
checkpoint-100/model.safetensors Normal file
View File

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

3
checkpoint-100/rng_state_0.pth Normal file
View File

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

3
checkpoint-100/rng_state_1.pth Normal file
View File

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

3
checkpoint-100/rng_state_2.pth Normal file
View File

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

3
checkpoint-100/rng_state_3.pth Normal file
View File

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

3
checkpoint-100/scheduler.pt Normal file
View File

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

31
checkpoint-100/special_tokens_map.json Normal file
View File

@@ -0,0 +1,31 @@
{
"additional_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|>"
],
"eos_token": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
},
"pad_token": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false
}
}

3
checkpoint-100/tokenizer.json Normal file
View File

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

208
checkpoint-100/tokenizer_config.json Normal file
View File

@@ -0,0 +1,208 @@
{
"add_bos_token": false,
"add_prefix_space": false,
"added_tokens_decoder": {
"151643": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151644": {
"content": "<|im_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151645": {
"content": "<|im_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151646": {
"content": "<|object_ref_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151647": {
"content": "<|object_ref_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151648": {
"content": "<|box_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151649": {
"content": "<|box_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151650": {
"content": "<|quad_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151651": {
"content": "<|quad_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151652": {
"content": "<|vision_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151653": {
"content": "<|vision_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151654": {
"content": "<|vision_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151655": {
"content": "<|image_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151656": {
"content": "<|video_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"151657": {
"content": "<tool_call>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151658": {
"content": "</tool_call>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151659": {
"content": "<|fim_prefix|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151660": {
"content": "<|fim_middle|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151661": {
"content": "<|fim_suffix|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151662": {
"content": "<|fim_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151663": {
"content": "<|repo_name|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"151664": {
"content": "<|file_sep|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
}
},
"additional_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|>"
],
"bos_token": null,
"chat_template": "{%- if tools %}\n {{- '<|im_start|>system\\n' }}\n {%- if messages[0]['role'] == 'system' %}\n {{- messages[0]['content'] }}\n {%- else %}\n {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}\n {%- endif %}\n {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n {%- for tool in tools %}\n {{- \"\\n\" }}\n {{- tool | tojson }}\n {%- endfor %}\n {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n {%- if messages[0]['role'] == 'system' %}\n {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n {%- else %}\n {{- '<|im_start|>system\\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n {%- elif message.role == \"assistant\" %}\n {{- '<|im_start|>' + message.role }}\n {%- if message.content %}\n {{- '\\n' + message.content }}\n {%- endif %}\n {%- for tool_call in message.tool_calls %}\n {%- if tool_call.function is defined %}\n {%- set tool_call = tool_call.function %}\n {%- endif %}\n {{- '\\n<tool_call>\\n{\"name\": \"' }}\n {{- tool_call.name }}\n {{- '\", \"arguments\": ' }}\n {{- tool_call.arguments | tojson }}\n {{- '}\\n</tool_call>' }}\n {%- endfor %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"tool\" %}\n {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n {{- '<|im_start|>user' }}\n {%- endif %}\n {{- '\\n<tool_response>\\n' }}\n {{- message.content }}\n {{- '\\n</tool_response>' }}\n {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
"clean_up_tokenization_spaces": false,
"eos_token": "<|endoftext|>",
"errors": "replace",
"model_max_length": 131072,
"pad_token": "<|endoftext|>",
"padding_side": "right",
"split_special_tokens": false,
"tokenizer_class": "Qwen2Tokenizer",
"unk_token": null
}

733
checkpoint-100/trainer_state.json Normal file
View File

@@ -0,0 +1,733 @@
{
"best_metric": null,
"best_model_checkpoint": null,
"epoch": 2.0,
"eval_steps": 500,
"global_step": 100,
"is_hyper_param_search": false,
"is_local_process_zero": true,
"is_world_process_zero": true,
"log_history": [
{
"epoch": 0.02,
"grad_norm": 2.5414493083953857,
"learning_rate": 0.00015996052482925854,
"loss": 1.3158,
"step": 1
},
{
"epoch": 0.04,
"grad_norm": 22.558874130249023,
"learning_rate": 0.00015984213827426174,
"loss": 3.9659,
"step": 2
},
{
"epoch": 0.06,
"grad_norm": 13.792680740356445,
"learning_rate": 0.0001596449571682464,
"loss": 2.6879,
"step": 3
},
{
"epoch": 0.08,
"grad_norm": 4.81754207611084,
"learning_rate": 0.00015936917610515826,
"loss": 2.0717,
"step": 4
},
{
"epoch": 0.1,
"grad_norm": 6.582490921020508,
"learning_rate": 0.00015901506724761103,
"loss": 1.9221,
"step": 5
},
{
"epoch": 0.12,
"grad_norm": 3.573866844177246,
"learning_rate": 0.00015858298005829512,
"loss": 1.7919,
"step": 6
},
{
"epoch": 0.14,
"grad_norm": 2.014141082763672,
"learning_rate": 0.0001580733409550998,
"loss": 1.6514,
"step": 7
},
{
"epoch": 0.16,
"grad_norm": 2.1408352851867676,
"learning_rate": 0.0001574866528902905,
"loss": 1.561,
"step": 8
},
{
"epoch": 0.18,
"grad_norm": 1.8259636163711548,
"learning_rate": 0.00015682349485415545,
"loss": 1.4794,
"step": 9
},
{
"epoch": 0.2,
"grad_norm": 1.1060450077056885,
"learning_rate": 0.0001560845213036123,
"loss": 1.4122,
"step": 10
},
{
"epoch": 0.22,
"grad_norm": 1.6639591455459595,
"learning_rate": 0.00015527046151633805,
"loss": 1.4144,
"step": 11
},
{
"epoch": 0.24,
"grad_norm": 1.007230281829834,
"learning_rate": 0.00015438211887106013,
"loss": 1.3658,
"step": 12
},
{
"epoch": 0.26,
"grad_norm": 0.9042938947677612,
"learning_rate": 0.0001534203700547185,
"loss": 1.3555,
"step": 13
},
{
"epoch": 0.28,
"grad_norm": 1.5797706842422485,
"learning_rate": 0.00015238616419728157,
"loss": 1.3031,
"step": 14
},
{
"epoch": 0.3,
"grad_norm": 0.8265913128852844,
"learning_rate": 0.00015128052193506944,
"loss": 1.2992,
"step": 15
},
{
"epoch": 0.32,
"grad_norm": 0.7161855697631836,
"learning_rate": 0.0001501045344035091,
"loss": 1.3053,
"step": 16
},
{
"epoch": 0.34,
"grad_norm": 0.655879557132721,
"learning_rate": 0.0001488593621603155,
"loss": 1.2677,
"step": 17
},
{
"epoch": 0.36,
"grad_norm": 0.7548253536224365,
"learning_rate": 0.00014754623404016122,
"loss": 1.243,
"step": 18
},
{
"epoch": 0.38,
"grad_norm": 0.7004591822624207,
"learning_rate": 0.00014616644594196495,
"loss": 1.2565,
"step": 19
},
{
"epoch": 0.4,
"grad_norm": 0.5514001250267029,
"learning_rate": 0.00014472135954999581,
"loss": 1.2198,
"step": 20
},
{
"epoch": 0.42,
"grad_norm": 0.45575690269470215,
"learning_rate": 0.00014321240099005524,
"loss": 1.2149,
"step": 21
},
{
"epoch": 0.44,
"grad_norm": 0.45689329504966736,
"learning_rate": 0.00014164105942206316,
"loss": 1.2114,
"step": 22
},
{
"epoch": 0.46,
"grad_norm": 0.47528719902038574,
"learning_rate": 0.00014000888557043678,
"loss": 1.2158,
"step": 23
},
{
"epoch": 0.48,
"grad_norm": 0.4457974135875702,
"learning_rate": 0.00013831749019371293,
"loss": 1.2238,
"step": 24
},
{
"epoch": 0.5,
"grad_norm": 0.35548001527786255,
"learning_rate": 0.00013656854249492382,
"loss": 1.2145,
"step": 25
},
{
"epoch": 0.52,
"grad_norm": 0.5427165627479553,
"learning_rate": 0.00013476376847429511,
"loss": 1.2081,
"step": 26
},
{
"epoch": 0.54,
"grad_norm": 0.44462138414382935,
"learning_rate": 0.00013290494922589216,
"loss": 1.1975,
"step": 27
},
{
"epoch": 0.56,
"grad_norm": 0.43967288732528687,
"learning_rate": 0.0001309939191798952,
"loss": 1.1867,
"step": 28
},
{
"epoch": 0.58,
"grad_norm": 0.3888394236564636,
"learning_rate": 0.00012903256429223813,
"loss": 1.1566,
"step": 29
},
{
"epoch": 0.6,
"grad_norm": 0.3624572455883026,
"learning_rate": 0.00012702282018339786,
"loss": 1.1663,
"step": 30
},
{
"epoch": 0.62,
"grad_norm": 0.37822291254997253,
"learning_rate": 0.00012496667022817044,
"loss": 1.1654,
"step": 31
},
{
"epoch": 0.64,
"grad_norm": 0.312124103307724,
"learning_rate": 0.00012286614359831974,
"loss": 1.1489,
"step": 32
},
{
"epoch": 0.66,
"grad_norm": 0.3584006428718567,
"learning_rate": 0.00012072331326002972,
"loss": 1.18,
"step": 33
},
{
"epoch": 0.68,
"grad_norm": 0.2872241735458374,
"learning_rate": 0.00011854029392813723,
"loss": 1.139,
"step": 34
},
{
"epoch": 0.7,
"grad_norm": 0.2621050179004669,
"learning_rate": 0.00011631923997916375,
"loss": 1.1694,
"step": 35
},
{
"epoch": 0.72,
"grad_norm": 0.2508472502231598,
"learning_rate": 0.00011406234332520582,
"loss": 1.1404,
"step": 36
},
{
"epoch": 0.74,
"grad_norm": 0.2924502491950989,
"learning_rate": 0.00011177183125078245,
"loss": 1.1311,
"step": 37
},
{
"epoch": 0.76,
"grad_norm": 0.3054457902908325,
"learning_rate": 0.00010944996421477426,
"loss": 1.1536,
"step": 38
},
{
"epoch": 0.78,
"grad_norm": 0.255205363035202,
"learning_rate": 0.00010709903361962333,
"loss": 1.137,
"step": 39
},
{
"epoch": 0.8,
"grad_norm": 0.2247438281774521,
"learning_rate": 0.0001047213595499958,
"loss": 1.1247,
"step": 40
},
{
"epoch": 0.82,
"grad_norm": 0.22704657912254333,
"learning_rate": 0.00010231928848313836,
"loss": 1.1189,
"step": 41
},
{
"epoch": 0.84,
"grad_norm": 0.22271329164505005,
"learning_rate": 9.989519097318841e-05,
"loss": 1.1254,
"step": 42
},
{
"epoch": 0.86,
"grad_norm": 0.21321555972099304,
"learning_rate": 9.745145931172342e-05,
"loss": 1.0805,
"step": 43
},
{
"epoch": 0.88,
"grad_norm": 0.2108481377363205,
"learning_rate": 9.4990505166858e-05,
"loss": 1.1303,
"step": 44
},
{
"epoch": 0.9,
"grad_norm": 0.2166047990322113,
"learning_rate": 9.251475720321848e-05,
"loss": 1.1149,
"step": 45
},
{
"epoch": 0.92,
"grad_norm": 0.20189222693443298,
"learning_rate": 9.002665868514435e-05,
"loss": 1.1049,
"step": 46
},
{
"epoch": 0.94,
"grad_norm": 0.23353922367095947,
"learning_rate": 8.752866506548117e-05,
"loss": 1.0833,
"step": 47
},
{
"epoch": 0.96,
"grad_norm": 0.20002281665802002,
"learning_rate": 8.502324156234508e-05,
"loss": 1.0629,
"step": 48
},
{
"epoch": 0.98,
"grad_norm": 0.22263740003108978,
"learning_rate": 8.251286072625027e-05,
"loss": 1.1009,
"step": 49
},
{
"epoch": 1.0,
"grad_norm": 0.2224554568529129,
"learning_rate": 8e-05,
"loss": 1.07,
"step": 50
},
{
"epoch": 1.02,
"grad_norm": 0.25747236609458923,
"learning_rate": 7.748713927374974e-05,
"loss": 1.113,
"step": 51
},
{
"epoch": 1.04,
"grad_norm": 0.3350413143634796,
"learning_rate": 7.497675843765493e-05,
"loss": 1.2082,
"step": 52
},
{
"epoch": 1.06,
"grad_norm": 0.25098180770874023,
"learning_rate": 7.247133493451886e-05,
"loss": 1.1663,
"step": 53
},
{
"epoch": 1.08,
"grad_norm": 0.2518307864665985,
"learning_rate": 6.997334131485565e-05,
"loss": 1.195,
"step": 54
},
{
"epoch": 1.1,
"grad_norm": 0.27086207270622253,
"learning_rate": 6.748524279678152e-05,
"loss": 1.1827,
"step": 55
},
{
"epoch": 1.12,
"grad_norm": 0.2570662498474121,
"learning_rate": 6.500949483314202e-05,
"loss": 1.2075,
"step": 56
},
{
"epoch": 1.1400000000000001,
"grad_norm": 0.2899071276187897,
"learning_rate": 6.254854068827662e-05,
"loss": 1.1787,
"step": 57
},
{
"epoch": 1.16,
"grad_norm": 0.2995426654815674,
"learning_rate": 6.0104809026811634e-05,
"loss": 1.1946,
"step": 58
},
{
"epoch": 1.18,
"grad_norm": 0.2579825818538666,
"learning_rate": 5.7680711516861674e-05,
"loss": 1.1632,
"step": 59
},
{
"epoch": 1.2,
"grad_norm": 0.2803911864757538,
"learning_rate": 5.5278640450004216e-05,
"loss": 1.15,
"step": 60
},
{
"epoch": 1.22,
"grad_norm": 0.27153438329696655,
"learning_rate": 5.2900966380376693e-05,
"loss": 1.1817,
"step": 61
},
{
"epoch": 1.24,
"grad_norm": 0.26352250576019287,
"learning_rate": 5.055003578522577e-05,
"loss": 1.1487,
"step": 62
},
{
"epoch": 1.26,
"grad_norm": 0.30853864550590515,
"learning_rate": 4.822816874921756e-05,
"loss": 1.1649,
"step": 63
},
{
"epoch": 1.28,
"grad_norm": 0.2495918869972229,
"learning_rate": 4.593765667479419e-05,
"loss": 1.1325,
"step": 64
},
{
"epoch": 1.3,
"grad_norm": 0.2213287055492401,
"learning_rate": 4.3680760020836266e-05,
"loss": 1.1281,
"step": 65
},
{
"epoch": 1.32,
"grad_norm": 0.18903322517871857,
"learning_rate": 4.145970607186277e-05,
"loss": 1.1476,
"step": 66
},
{
"epoch": 1.34,
"grad_norm": 0.20846880972385406,
"learning_rate": 3.92766867399703e-05,
"loss": 1.1233,
"step": 67
},
{
"epoch": 1.3599999999999999,
"grad_norm": 0.17987675964832306,
"learning_rate": 3.7133856401680256e-05,
"loss": 1.1051,
"step": 68
},
{
"epoch": 1.38,
"grad_norm": 0.17960353195667267,
"learning_rate": 3.5033329771829576e-05,
"loss": 1.1266,
"step": 69
},
{
"epoch": 1.4,
"grad_norm": 0.19571684300899506,
"learning_rate": 3.297717981660216e-05,
"loss": 1.0983,
"step": 70
},
{
"epoch": 1.42,
"grad_norm": 0.17779013514518738,
"learning_rate": 3.09674357077619e-05,
"loss": 1.0976,
"step": 71
},
{
"epoch": 1.44,
"grad_norm": 0.16692589223384857,
"learning_rate": 2.9006080820104823e-05,
"loss": 1.0988,
"step": 72
},
{
"epoch": 1.46,
"grad_norm": 0.14753758907318115,
"learning_rate": 2.7095050774107867e-05,
"loss": 1.1055,
"step": 73
},
{
"epoch": 1.48,
"grad_norm": 0.1639815717935562,
"learning_rate": 2.5236231525704902e-05,
"loss": 1.117,
"step": 74
},
{
"epoch": 1.5,
"grad_norm": 0.16390062868595123,
"learning_rate": 2.3431457505076205e-05,
"loss": 1.1131,
"step": 75
},
{
"epoch": 1.52,
"grad_norm": 0.15375538170337677,
"learning_rate": 2.1682509806287094e-05,
"loss": 1.1111,
"step": 76
},
{
"epoch": 1.54,
"grad_norm": 0.13153688609600067,
"learning_rate": 1.9991114429563236e-05,
"loss": 1.1045,
"step": 77
},
{
"epoch": 1.56,
"grad_norm": 0.14417824149131775,
"learning_rate": 1.835894057793687e-05,
"loss": 1.0936,
"step": 78
},
{
"epoch": 1.58,
"grad_norm": 0.1515759974718094,
"learning_rate": 1.678759900994477e-05,
"loss": 1.0665,
"step": 79
},
{
"epoch": 1.6,
"grad_norm": 0.15299999713897705,
"learning_rate": 1.5278640450004213e-05,
"loss": 1.0794,
"step": 80
},
{
"epoch": 1.62,
"grad_norm": 0.1259460747241974,
"learning_rate": 1.3833554058035045e-05,
"loss": 1.084,
"step": 81
},
{
"epoch": 1.6400000000000001,
"grad_norm": 0.13053837418556213,
"learning_rate": 1.2453765959838813e-05,
"loss": 1.0694,
"step": 82
},
{
"epoch": 1.6600000000000001,
"grad_norm": 0.1260591298341751,
"learning_rate": 1.1140637839684519e-05,
"loss": 1.0983,
"step": 83
},
{
"epoch": 1.6800000000000002,
"grad_norm": 0.11727996915578842,
"learning_rate": 9.895465596490931e-06,
"loss": 1.0624,
"step": 84
},
{
"epoch": 1.7,
"grad_norm": 0.12494169920682907,
"learning_rate": 8.719478064930578e-06,
"loss": 1.0937,
"step": 85
},
{
"epoch": 1.72,
"grad_norm": 0.13201913237571716,
"learning_rate": 7.613835802718452e-06,
"loss": 1.07,
"step": 86
},
{
"epoch": 1.74,
"grad_norm": 0.12661895155906677,
"learning_rate": 6.579629945281509e-06,
"loss": 1.0622,
"step": 87
},
{
"epoch": 1.76,
"grad_norm": 0.11581739038228989,
"learning_rate": 5.6178811289398925e-06,
"loss": 1.0805,
"step": 88
},
{
"epoch": 1.78,
"grad_norm": 0.1270962953567505,
"learning_rate": 4.729538483661964e-06,
"loss": 1.0685,
"step": 89
},
{
"epoch": 1.8,
"grad_norm": 0.1226486787199974,
"learning_rate": 3.915478696387718e-06,
"loss": 1.0612,
"step": 90
},
{
"epoch": 1.8199999999999998,
"grad_norm": 0.12050648778676987,
"learning_rate": 3.176505145844555e-06,
"loss": 1.053,
"step": 91
},
{
"epoch": 1.8399999999999999,
"grad_norm": 0.12416510283946991,
"learning_rate": 2.513347109709514e-06,
"loss": 1.0649,
"step": 92
},
{
"epoch": 1.8599999999999999,
"grad_norm": 0.10619990527629852,
"learning_rate": 1.9266590449002052e-06,
"loss": 1.024,
"step": 93
},
{
"epoch": 1.88,
"grad_norm": 0.10766904056072235,
"learning_rate": 1.4170199417049114e-06,
"loss": 1.0724,
"step": 94
},
{
"epoch": 1.9,
"grad_norm": 0.10820022970438004,
"learning_rate": 9.849327523889873e-07,
"loss": 1.0594,
"step": 95
},
{
"epoch": 1.92,
"grad_norm": 0.10673188418149948,
"learning_rate": 6.308238948417788e-07,
"loss": 1.0471,
"step": 96
},
{
"epoch": 1.94,
"grad_norm": 0.10520946979522705,
"learning_rate": 3.550428317536003e-07,
"loss": 1.0315,
"step": 97
},
{
"epoch": 1.96,
"grad_norm": 0.10756369680166245,
"learning_rate": 1.578617257382753e-07,
"loss": 1.0129,
"step": 98
},
{
"epoch": 1.98,
"grad_norm": 0.11246038973331451,
"learning_rate": 3.9475170741472005e-08,
"loss": 1.0481,
"step": 99
},
{
"epoch": 2.0,
"grad_norm": 0.11008872836828232,
"learning_rate": 0.0,
"loss": 1.0204,
"step": 100
}
],
"logging_steps": 1,
"max_steps": 100,
"num_input_tokens_seen": 0,
"num_train_epochs": 2,
"save_steps": 50,
"stateful_callbacks": {
"TrainerControl": {
"args": {
"should_epoch_stop": false,
"should_evaluate": false,
"should_log": false,
"should_save": true,
"should_training_stop": true
},
"attributes": {}
}
},
"total_flos": 2.671112626070618e+18,
"train_batch_size": 1,
"trial_name": null,
"trial_params": null
}

3
checkpoint-100/training_args.bin Normal file
View File

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

1
checkpoint-100/vocab.json Normal file
View File

File diff suppressed because one or more lines are too long

604
checkpoint-100/zero_to_fp32.py Normal file
View File

@@ -0,0 +1,604 @@
#!/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 . pytorch_model.bin
import argparse
import torch
import glob
import math
import os
import re
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)
@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 parse_model_states(files):
zero_model_states = []
for file in files:
state_dict = torch.load(file, map_location=device)
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 files:
state_dict = torch.load(f, map_location=device)
# 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 not ZERO_STAGE 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}")
if zero_stage <= 2:
fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
elif zero_stage == 3:
# if there is more than one param group, there will be multiple flattened tensors - one
# flattened tensor per group - for simplicity merge them into a single tensor
#
# XXX: could make the script more memory efficient for when there are multiple groups - it
# will require matching the sub-lists of param_shapes for each param group flattened tensor
fp32_flat_groups = [
torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key], 0) 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}'")
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}")
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}')
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")
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 = fp32_flat_groups[0].numel() * 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
for name, shape in param_shapes.items():
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}"
)
# XXX: memory usage doubles here
state_dict[name] = torch.cat(
tuple(fp32_flat_groups[i].narrow(0, offset, partitioned_numel) for i in range(world_size)),
0).narrow(0, 0, unpartitioned_numel).view(shape)
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 get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None, exclude_frozen_parameters=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
Returns:
- pytorch ``state_dict``
Note: this approach may not work if your application doesn't have sufficient free CPU memory and
you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
the checkpoint.
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.
"""
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")
return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, 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_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)
- ``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
"""
state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag, exclude_frozen_parameters)
print(f"Saving fp32 state dict to {output_file}")
torch.save(state_dict, output_file)
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(f"Extracting fp32 weights")
state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
logger.info(f"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_file",
type=str,
help="path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)")
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_file,
tag=args.tag,
exclude_frozen_parameters=args.exclude_frozen_parameters)