cjiao/OpenThoughts3-random-groups-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-random-groups-openthinker3-1.5B-checkpoint-375-length-filtered
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-04-27 22:17:19 +08:00
commit 39e6309ba9
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:34a193e33b797fbcfda37684b7c24dd7055c8efead3d339f5c13adfefb124e92
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.0897409915924072,
"learning_rate": 0.00015996052482925854,
"loss": 1.2438,
"step": 1
},
{
"epoch": 0.04,
"grad_norm": 28.813220977783203,
"learning_rate": 0.00015984213827426174,
"loss": 3.7464,
"step": 2
},
{
"epoch": 0.06,
"grad_norm": 16.25991439819336,
"learning_rate": 0.0001596449571682464,
"loss": 2.9219,
"step": 3
},
{
"epoch": 0.08,
"grad_norm": 5.324892520904541,
"learning_rate": 0.00015936917610515826,
"loss": 2.1242,
"step": 4
},
{
"epoch": 0.1,
"grad_norm": 2.8375790119171143,
"learning_rate": 0.00015901506724761103,
"loss": 1.7799,
"step": 5
},
{
"epoch": 0.12,
"grad_norm": 4.01439905166626,
"learning_rate": 0.00015858298005829512,
"loss": 1.6514,
"step": 6
},
{
"epoch": 0.14,
"grad_norm": 2.095858335494995,
"learning_rate": 0.0001580733409550998,
"loss": 1.5758,
"step": 7
},
{
"epoch": 0.16,
"grad_norm": 1.694789171218872,
"learning_rate": 0.0001574866528902905,
"loss": 1.4651,
"step": 8
},
{
"epoch": 0.18,
"grad_norm": 1.4686285257339478,
"learning_rate": 0.00015682349485415545,
"loss": 1.3782,
"step": 9
},
{
"epoch": 0.2,
"grad_norm": 0.9562129378318787,
"learning_rate": 0.0001560845213036123,
"loss": 1.353,
"step": 10
},
{
"epoch": 0.22,
"grad_norm": 1.4398468732833862,
"learning_rate": 0.00015527046151633805,
"loss": 1.3325,
"step": 11
},
{
"epoch": 0.24,
"grad_norm": 0.827878475189209,
"learning_rate": 0.00015438211887106013,
"loss": 1.2997,
"step": 12
},
{
"epoch": 0.26,
"grad_norm": 1.015064001083374,
"learning_rate": 0.0001534203700547185,
"loss": 1.2664,
"step": 13
},
{
"epoch": 0.28,
"grad_norm": 0.680793821811676,
"learning_rate": 0.00015238616419728157,
"loss": 1.239,
"step": 14
},
{
"epoch": 0.3,
"grad_norm": 0.9654543995857239,
"learning_rate": 0.00015128052193506944,
"loss": 1.2528,
"step": 15
},
{
"epoch": 0.32,
"grad_norm": 0.7509903907775879,
"learning_rate": 0.0001501045344035091,
"loss": 1.2447,
"step": 16
},
{
"epoch": 0.34,
"grad_norm": 0.6938178539276123,
"learning_rate": 0.0001488593621603155,
"loss": 1.2231,
"step": 17
},
{
"epoch": 0.36,
"grad_norm": 0.6297674179077148,
"learning_rate": 0.00014754623404016122,
"loss": 1.2305,
"step": 18
},
{
"epoch": 0.38,
"grad_norm": 0.535707950592041,
"learning_rate": 0.00014616644594196495,
"loss": 1.2339,
"step": 19
},
{
"epoch": 0.4,
"grad_norm": 0.5825291275978088,
"learning_rate": 0.00014472135954999581,
"loss": 1.1903,
"step": 20
},
{
"epoch": 0.42,
"grad_norm": 0.36997389793395996,
"learning_rate": 0.00014321240099005524,
"loss": 1.1979,
"step": 21
},
{
"epoch": 0.44,
"grad_norm": 0.4852772355079651,
"learning_rate": 0.00014164105942206316,
"loss": 1.1951,
"step": 22
},
{
"epoch": 0.46,
"grad_norm": 0.43263790011405945,
"learning_rate": 0.00014000888557043678,
"loss": 1.1632,
"step": 23
},
{
"epoch": 0.48,
"grad_norm": 0.39633846282958984,
"learning_rate": 0.00013831749019371293,
"loss": 1.1631,
"step": 24
},
{
"epoch": 0.5,
"grad_norm": 0.41002726554870605,
"learning_rate": 0.00013656854249492382,
"loss": 1.1886,
"step": 25
},
{
"epoch": 0.52,
"grad_norm": 0.39650410413742065,
"learning_rate": 0.00013476376847429511,
"loss": 1.202,
"step": 26
},
{
"epoch": 0.54,
"grad_norm": 0.3824397921562195,
"learning_rate": 0.00013290494922589216,
"loss": 1.1832,
"step": 27
},
{
"epoch": 0.56,
"grad_norm": 0.3212684094905853,
"learning_rate": 0.0001309939191798952,
"loss": 1.1466,
"step": 28
},
{
"epoch": 0.58,
"grad_norm": 0.32189854979515076,
"learning_rate": 0.00012903256429223813,
"loss": 1.1975,
"step": 29
},
{
"epoch": 0.6,
"grad_norm": 0.3070578873157501,
"learning_rate": 0.00012702282018339786,
"loss": 1.154,
"step": 30
},
{
"epoch": 0.62,
"grad_norm": 0.25685596466064453,
"learning_rate": 0.00012496667022817044,
"loss": 1.1357,
"step": 31
},
{
"epoch": 0.64,
"grad_norm": 0.2539421021938324,
"learning_rate": 0.00012286614359831974,
"loss": 1.158,
"step": 32
},
{
"epoch": 0.66,
"grad_norm": 0.2875991761684418,
"learning_rate": 0.00012072331326002972,
"loss": 1.1849,
"step": 33
},
{
"epoch": 0.68,
"grad_norm": 0.2683248221874237,
"learning_rate": 0.00011854029392813723,
"loss": 1.1935,
"step": 34
},
{
"epoch": 0.7,
"grad_norm": 0.2622566819190979,
"learning_rate": 0.00011631923997916375,
"loss": 1.1624,
"step": 35
},
{
"epoch": 0.72,
"grad_norm": 0.2862881124019623,
"learning_rate": 0.00011406234332520582,
"loss": 1.1854,
"step": 36
},
{
"epoch": 0.74,
"grad_norm": 0.24599239230155945,
"learning_rate": 0.00011177183125078245,
"loss": 1.1597,
"step": 37
},
{
"epoch": 0.76,
"grad_norm": 0.23713542520999908,
"learning_rate": 0.00010944996421477426,
"loss": 1.125,
"step": 38
},
{
"epoch": 0.78,
"grad_norm": 0.26411688327789307,
"learning_rate": 0.00010709903361962333,
"loss": 1.1144,
"step": 39
},
{
"epoch": 0.8,
"grad_norm": 0.24425670504570007,
"learning_rate": 0.0001047213595499958,
"loss": 1.1443,
"step": 40
},
{
"epoch": 0.82,
"grad_norm": 0.26452597975730896,
"learning_rate": 0.00010231928848313836,
"loss": 1.1376,
"step": 41
},
{
"epoch": 0.84,
"grad_norm": 0.21576781570911407,
"learning_rate": 9.989519097318841e-05,
"loss": 1.1616,
"step": 42
},
{
"epoch": 0.86,
"grad_norm": 0.2854762375354767,
"learning_rate": 9.745145931172342e-05,
"loss": 1.1511,
"step": 43
},
{
"epoch": 0.88,
"grad_norm": 0.2194770872592926,
"learning_rate": 9.4990505166858e-05,
"loss": 1.1385,
"step": 44
},
{
"epoch": 0.9,
"grad_norm": 0.22979624569416046,
"learning_rate": 9.251475720321848e-05,
"loss": 1.1429,
"step": 45
},
{
"epoch": 0.92,
"grad_norm": 0.22235210239887238,
"learning_rate": 9.002665868514435e-05,
"loss": 1.119,
"step": 46
},
{
"epoch": 0.94,
"grad_norm": 0.22438690066337585,
"learning_rate": 8.752866506548117e-05,
"loss": 1.1642,
"step": 47
},
{
"epoch": 0.96,
"grad_norm": 0.18937288224697113,
"learning_rate": 8.502324156234508e-05,
"loss": 1.1281,
"step": 48
},
{
"epoch": 0.98,
"grad_norm": 0.22655418515205383,
"learning_rate": 8.251286072625027e-05,
"loss": 1.1399,
"step": 49
},
{
"epoch": 1.0,
"grad_norm": 0.21689347922801971,
"learning_rate": 8e-05,
"loss": 1.0937,
"step": 50
},
{
"epoch": 1.02,
"grad_norm": 0.370595782995224,
"learning_rate": 7.748713927374974e-05,
"loss": 1.0383,
"step": 51
},
{
"epoch": 1.04,
"grad_norm": 0.45078134536743164,
"learning_rate": 7.497675843765493e-05,
"loss": 1.1376,
"step": 52
},
{
"epoch": 1.06,
"grad_norm": 0.3670865297317505,
"learning_rate": 7.247133493451886e-05,
"loss": 1.1206,
"step": 53
},
{
"epoch": 1.08,
"grad_norm": 0.3233189582824707,
"learning_rate": 6.997334131485565e-05,
"loss": 1.1502,
"step": 54
},
{
"epoch": 1.1,
"grad_norm": 0.399816632270813,
"learning_rate": 6.748524279678152e-05,
"loss": 1.1279,
"step": 55
},
{
"epoch": 1.12,
"grad_norm": 0.33596694469451904,
"learning_rate": 6.500949483314202e-05,
"loss": 1.1218,
"step": 56
},
{
"epoch": 1.1400000000000001,
"grad_norm": 0.5186057090759277,
"learning_rate": 6.254854068827662e-05,
"loss": 1.1378,
"step": 57
},
{
"epoch": 1.16,
"grad_norm": 0.2955310344696045,
"learning_rate": 6.0104809026811634e-05,
"loss": 1.1156,
"step": 58
},
{
"epoch": 1.18,
"grad_norm": 0.5075625777244568,
"learning_rate": 5.7680711516861674e-05,
"loss": 1.079,
"step": 59
},
{
"epoch": 1.2,
"grad_norm": 0.25788846611976624,
"learning_rate": 5.5278640450004216e-05,
"loss": 1.078,
"step": 60
},
{
"epoch": 1.22,
"grad_norm": 0.4179430603981018,
"learning_rate": 5.2900966380376693e-05,
"loss": 1.0959,
"step": 61
},
{
"epoch": 1.24,
"grad_norm": 0.34371474385261536,
"learning_rate": 5.055003578522577e-05,
"loss": 1.0807,
"step": 62
},
{
"epoch": 1.26,
"grad_norm": 0.2947935461997986,
"learning_rate": 4.822816874921756e-05,
"loss": 1.0645,
"step": 63
},
{
"epoch": 1.28,
"grad_norm": 0.27775266766548157,
"learning_rate": 4.593765667479419e-05,
"loss": 1.0591,
"step": 64
},
{
"epoch": 1.3,
"grad_norm": 0.23427054286003113,
"learning_rate": 4.3680760020836266e-05,
"loss": 1.0848,
"step": 65
},
{
"epoch": 1.32,
"grad_norm": 0.26551228761672974,
"learning_rate": 4.145970607186277e-05,
"loss": 1.0722,
"step": 66
},
{
"epoch": 1.34,
"grad_norm": 0.20378205180168152,
"learning_rate": 3.92766867399703e-05,
"loss": 1.0654,
"step": 67
},
{
"epoch": 1.3599999999999999,
"grad_norm": 0.3097435235977173,
"learning_rate": 3.7133856401680256e-05,
"loss": 1.0786,
"step": 68
},
{
"epoch": 1.38,
"grad_norm": 0.1919064074754715,
"learning_rate": 3.5033329771829576e-05,
"loss": 1.0804,
"step": 69
},
{
"epoch": 1.4,
"grad_norm": 0.19627317786216736,
"learning_rate": 3.297717981660216e-05,
"loss": 1.0532,
"step": 70
},
{
"epoch": 1.42,
"grad_norm": 0.22595791518688202,
"learning_rate": 3.09674357077619e-05,
"loss": 1.0669,
"step": 71
},
{
"epoch": 1.44,
"grad_norm": 0.20693820714950562,
"learning_rate": 2.9006080820104823e-05,
"loss": 1.0586,
"step": 72
},
{
"epoch": 1.46,
"grad_norm": 0.21305669844150543,
"learning_rate": 2.7095050774107867e-05,
"loss": 1.0425,
"step": 73
},
{
"epoch": 1.48,
"grad_norm": 0.19866228103637695,
"learning_rate": 2.5236231525704902e-05,
"loss": 1.0439,
"step": 74
},
{
"epoch": 1.5,
"grad_norm": 0.1801978200674057,
"learning_rate": 2.3431457505076205e-05,
"loss": 1.067,
"step": 75
},
{
"epoch": 1.52,
"grad_norm": 0.19466304779052734,
"learning_rate": 2.1682509806287094e-05,
"loss": 1.0777,
"step": 76
},
{
"epoch": 1.54,
"grad_norm": 0.23117047548294067,
"learning_rate": 1.9991114429563236e-05,
"loss": 1.064,
"step": 77
},
{
"epoch": 1.56,
"grad_norm": 0.1913105994462967,
"learning_rate": 1.835894057793687e-05,
"loss": 1.0343,
"step": 78
},
{
"epoch": 1.58,
"grad_norm": 0.17815832793712616,
"learning_rate": 1.678759900994477e-05,
"loss": 1.0886,
"step": 79
},
{
"epoch": 1.6,
"grad_norm": 0.20096604526042938,
"learning_rate": 1.5278640450004213e-05,
"loss": 1.0535,
"step": 80
},
{
"epoch": 1.62,
"grad_norm": 0.18867887556552887,
"learning_rate": 1.3833554058035045e-05,
"loss": 1.0352,
"step": 81
},
{
"epoch": 1.6400000000000001,
"grad_norm": 0.1709662228822708,
"learning_rate": 1.2453765959838813e-05,
"loss": 1.0571,
"step": 82
},
{
"epoch": 1.6600000000000001,
"grad_norm": 0.17397433519363403,
"learning_rate": 1.1140637839684519e-05,
"loss": 1.0827,
"step": 83
},
{
"epoch": 1.6800000000000002,
"grad_norm": 0.1669914424419403,
"learning_rate": 9.895465596490931e-06,
"loss": 1.0925,
"step": 84
},
{
"epoch": 1.7,
"grad_norm": 0.1735496073961258,
"learning_rate": 8.719478064930578e-06,
"loss": 1.0733,
"step": 85
},
{
"epoch": 1.72,
"grad_norm": 0.17620506882667542,
"learning_rate": 7.613835802718452e-06,
"loss": 1.0804,
"step": 86
},
{
"epoch": 1.74,
"grad_norm": 0.14163199067115784,
"learning_rate": 6.579629945281509e-06,
"loss": 1.0658,
"step": 87
},
{
"epoch": 1.76,
"grad_norm": 0.13924406468868256,
"learning_rate": 5.6178811289398925e-06,
"loss": 1.0417,
"step": 88
},
{
"epoch": 1.78,
"grad_norm": 0.1430363655090332,
"learning_rate": 4.729538483661964e-06,
"loss": 1.0325,
"step": 89
},
{
"epoch": 1.8,
"grad_norm": 0.1582486629486084,
"learning_rate": 3.915478696387718e-06,
"loss": 1.0589,
"step": 90
},
{
"epoch": 1.8199999999999998,
"grad_norm": 0.1521899402141571,
"learning_rate": 3.176505145844555e-06,
"loss": 1.0562,
"step": 91
},
{
"epoch": 1.8399999999999999,
"grad_norm": 0.1510230004787445,
"learning_rate": 2.513347109709514e-06,
"loss": 1.0795,
"step": 92
},
{
"epoch": 1.8599999999999999,
"grad_norm": 0.14750948548316956,
"learning_rate": 1.9266590449002052e-06,
"loss": 1.0656,
"step": 93
},
{
"epoch": 1.88,
"grad_norm": 0.13227519392967224,
"learning_rate": 1.4170199417049114e-06,
"loss": 1.0675,
"step": 94
},
{
"epoch": 1.9,
"grad_norm": 0.1416577845811844,
"learning_rate": 9.849327523889873e-07,
"loss": 1.065,
"step": 95
},
{
"epoch": 1.92,
"grad_norm": 0.1315535306930542,
"learning_rate": 6.308238948417788e-07,
"loss": 1.0449,
"step": 96
},
{
"epoch": 1.94,
"grad_norm": 0.1372697800397873,
"learning_rate": 3.550428317536003e-07,
"loss": 1.0836,
"step": 97
},
{
"epoch": 1.96,
"grad_norm": 0.1258520931005478,
"learning_rate": 1.578617257382753e-07,
"loss": 1.0578,
"step": 98
},
{
"epoch": 1.98,
"grad_norm": 0.12827666103839874,
"learning_rate": 3.9475170741472005e-08,
"loss": 1.0651,
"step": 99
},
{
"epoch": 2.0,
"grad_norm": 0.13206015527248383,
"learning_rate": 0.0,
"loss": 1.0265,
"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.263464103047594e+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:999b270af0c9033d1db6bc5a5d9570df8b783fa5fbd8ce2cd5e5f99509c3967a
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)